[0z5a]

What

Slicing a register along anything but its trailing dimensions is broken:

• Reading r[:, 1] on a local(4, 4) register emits slice_regs[i] = regs[1] for every element, instead of regs[i * 4 + 1]. No error is raised, and the kernel gets four copies of one value.
• Writing r[1, :] = v raises IndexError in the emitter.

Why

inst.dims lists the source dimensions the slice keeps, so the sliced tensor has len(dims) dimensions. The emitters got this wrong in two ways:

• SliceRegisterInstEmitter computed the sliced tensor's indices from the source layout (src_layout.get_global(...)) with the sliced tensor's local index.
• Both emitters then indexed that list by source dimension (dst_indices[dim]), not by position in dims.

The two mistakes cancel when the kept dimensions are the trailing ones, for example r[1, :] or the philox output slices. So those cases worked and the bug stayed hidden.

How

Take the indices from the sliced tensor's layout, and map them with enumerate(dims). The same mapping is applied to SliceAssignInstEmitter and to the const-register path. Tracked const registers are 1-D today, so dim and its position are the same there and that path behaves as before.

Testing

Host tests ran on the Grace CPU of a GH200: GCC 13.3, CUDA 13.2, PyTorch 2.14.1+cu130, tilus main at f23018a.

tests/instructions/test_slice_register.py is new. The codegen tests evaluate the register index in the generated source.cu for each loop iteration:

test	needs	main	this PR
test_read_slice_codegen[1]: r[:, 1] reads 1, 5, 9, 13	no GPU, no nvcc	fails: 1, 1, 1, 1	passes
test_assign_slice_codegen[0]: r[1, :] = v writes 4, 5, 6, 7	no GPU, no nvcc	fails: IndexError	passes
test_read_slice_codegen[0], test_assign_slice_codegen[1] (controls)	no GPU, no nvcc	pass	pass
test_read_slice[1], test_assign_slice[0]	GPU	2 fail	2 pass
test_read_slice[0], test_assign_slice[1], test_read_column_across_threads (controls)	GPU	pass	pass

GPU tests ran on GH200 (sm_90, CUDA 13.2). Not run on B300.

To check that working slices do not change, I compared the generated source.cu of main and this PR for 38 programs: tests/instructions scripts, hopper_matmul v0 to v5, blackwell_matmul v0 to v3 and the three philox kernels with two layouts each. 21 of them contain register slices. All 38 are byte-identical.

The rest of the CPU-runnable suite (tests/lang, tests/ir, tests/drivers, tests/transforms and the top-level tests) gives the same result on main and here: 192 passed, and the same 5 failed because they need a GPU.

ruff check and ruff format --check (0.11.0 from pyproject.toml, 0.12.3 from pre-commit), mypy (1.15.0) and the copyright check pass on the changed files.

Not in this PR

These behave the same on main and with this PR:

• r[1, :] += v fails in the transpiler.
• r[1, :] = 5 fails layout inference.
• r[1, :] on a register whose dimension 0 is spatial (spatial(4, 1).local(1, 4)) emits slice_regs[i] = regs[i], so each thread reads its own row and not row 1. Fixing that needs communication between threads, or an error.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[yaoyaoding]

/ok to test dd43c37