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.