From 9d4c909dc3c891c6e2126149634a5be5fe6da7aa Mon Sep 17 00:00:00 2001
From: johnnynunez <johnnynuca14@gmail.com>
Date: Wed, 24 Jun 2026 05:21:58 +0200
Subject: [PATCH] DivergentHF: fix off-by-one in sequence selector (0-based
 kernel)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

The 3 failing DivergentHF tests were an off-by-one: plane 0 always ran the
wrong sequence. DivergentBatchTransformDPP is 0-BASED — exec() calls
divergent_operate<0>(z, seqs...) and runs the sequence whose 0-based position
equals at(z) (data_parallel_patterns.h; the upstream regression test selector
returns index==0?0u:1u, i.e. 0 picks the FIRST sequence). _selector_cpp was
emitting the 1-based plane_map verbatim, so at(0)=1 selected the SECOND
sequence for plane 0.

compose_divergent keeps the readable 1-based plane_map at the API; the
selector now emits (s-1). The stale comment citing circular_tensor.h's
SequenceSelectorType as '1-based convention' was wrong for this kernel (that
is a different selector contract) — corrected.

Added sv=2 to the divergent cache signature so stale .so files from the
previous (buggy) selector are not reused.

test_batch_divergent_hf now 12/12 (was 9/12). circular/HF/e2e still green.
---
 fkl/jit.py                           | 19 ++++++++++++-------
 skills/fkl-python-extending/SKILL.md |  9 ++++++++-
 2 files changed, 20 insertions(+), 8 deletions(-)

diff --git a/fkl/jit.py b/fkl/jit.py
index 59e9583..38779a3 100644
--- a/fkl/jit.py
+++ b/fkl/jit.py
@@ -269,12 +269,15 @@ def __init__(self, plane_map, chains):
         self._variants = {}
 
     def _selector_cpp(self):
-        # FKL convention (see SequenceSelectorType in circular_tensor.h):
-        # at(z) is uint, 1-based sequence index. Generate a chain of ternaries.
-        terms = []
-        for z, s in enumerate(self.plane_map):
-            terms.append((z, s))
-        expr = f"{self.plane_map[-1]}u"
+        # DivergentBatchTransformDPP is 0-BASED: exec() calls
+        # divergent_operate<0>(z, seqs...) and runs the sequence whose
+        # 0-based position matches at(z) (see data_parallel_patterns.h; the
+        # upstream regression test uses at(index)=index==0?0u:1u, i.e. 0 picks
+        # the FIRST sequence). plane_map is given 1-based at the API for
+        # readability, so emit (s-1) here. (NOTE: this is a DIFFERENT selector
+        # contract than circular_tensor.h's SequenceSelectorType.)
+        terms = [(z, s - 1) for z, s in enumerate(self.plane_map)]
+        expr = f"{terms[-1][1]}u"
         for z, s in reversed(terms[:-1]):
             expr = f"(index == {z}u ? {s}u : {expr})"
         return (
@@ -384,7 +387,9 @@ def _get_variant(self, dt, shape, B):
             "|".join(op.token(st) for op, st in
                      zip(ch, _plan(ch, dt, shape, B)[0]))
             for ch in self.chains)
-        sig = (f"divergent;arch={_ARCH};in={dt}x{B};map={tuple(self.plane_map)};"
+        # sv (selector version): bump when _selector_cpp's emitted C++ changes
+        # for the same plane_map, so stale .so files are not reused.
+        sig = (f"divergent;arch={_ARCH};sv=2;in={dt}x{B};map={tuple(self.plane_map)};"
                f"chains={chain_toks}")
         so = get_backend().compile(src, sig)
         lib = ctypes.CDLL(str(so))
diff --git a/skills/fkl-python-extending/SKILL.md b/skills/fkl-python-extending/SKILL.md
index 695e2a9..9173835 100644
--- a/skills/fkl-python-extending/SKILL.md
+++ b/skills/fkl-python-extending/SKILL.md
@@ -124,7 +124,14 @@ def t_myop():
    over a const tuple -> rvalue-ref binding error (upstream bug). The
    divergent codegen launches launchDivergentBatchTransformDPP_Kernel
    directly with buildOperationSequence(...) lvalues and a generated
-   PySequenceSelector (1-based at(z), FKL convention).
+   PySequenceSelector. The selector is 0-BASED: exec() calls
+   divergent_operate<0>(z, seqs...) and runs the sequence at the 0-based
+   position equal to at(z) (data_parallel_patterns.h; upstream regression
+   test MySelector::at returns index==0?0u:1u). compose_divergent takes a
+   1-based plane_map at the API for readability, so _selector_cpp emits
+   (s-1). Do NOT confuse this with circular_tensor.h's SequenceSelectorType
+   (a different contract). Bump the `sv=` token in the divergent signature
+   when the emitted selector C++ changes, or stale .so files get reused.
 8. IOpSequences passed to the divergent kernel must be LVALUES (const auto
    seqN = buildOperationSequence(...)), not temporaries inlined in the
    launch expression.