From f9b774c30700010b5d46cd96e39b695646cc0d5a Mon Sep 17 00:00:00 2001 From: 0z5a Date: Sat, 7 Nov 2026 12:48:25 +0100 Subject: [PATCH 1/2] [Fix] Use the retained dimensions when slicing a register SliceRegisterInst and SliceAssignInst keep the dimensions in inst.dims, but the emitters indexed the indices of the sliced tensor by source dimension. Slice reads also took those indices from the source layout, so r[:, 1] on local(4, 4) read regs[1] for every element, and r[1, :] = v raised IndexError. Read the indices from the sliced tensor's layout and map them by their position in dims. Signed-off-by: 0z5a --- python/tilus/backends/emitters/regs.py | 18 +-- tests/instructions/test_slice_register.py | 155 ++++++++++++++++++++++ 2 files changed, 164 insertions(+), 9 deletions(-) create mode 100644 tests/instructions/test_slice_register.py diff --git a/python/tilus/backends/emitters/regs.py b/python/tilus/backends/emitters/regs.py index 238c8865..aa0664f1 100644 --- a/python/tilus/backends/emitters/regs.py +++ b/python/tilus/backends/emitters/regs.py @@ -40,11 +40,11 @@ def _emit_default(self, inst, dst_tensor, dst_layout, src_tensor, src_layout): src_var = self.tensor2var[src_tensor] with self.for_range(extent=dst_layout.local_size) as dst_local: - dst_indices = src_layout.get_global(spatial_index=self.current_thread, local_index=dst_local) + dst_indices = dst_layout.get_global(spatial_index=self.current_thread, local_index=dst_local) src_indices = list(inst.offsets) dims = range(len(src_layout.shape)) if inst.dims is None else inst.dims - for dim in dims: - src_indices[dim] = src_indices[dim] + dst_indices[dim] + for dst_dim, dim in enumerate(dims): + src_indices[dim] = src_indices[dim] + dst_indices[dst_dim] src_local = src_layout.get_local(global_indices=src_indices) self.buffer_store(buf=dst_var, indices=[dst_local], value=src_var[src_local]) @@ -61,8 +61,8 @@ def _emit_const_reg(self, inst, dst_tensor, dst_layout, src_tensor, src_layout, dst_global = dst_layout.get_global(spatial_index=self.current_thread, local_index=dst_local) src_logical_indices = list(inst.offsets) dims = range(len(src_layout.shape)) if inst.dims is None else inst.dims - for dim in dims: - src_logical_indices[dim] = src_logical_indices[dim] + dst_global[dim] + for dst_dim, dim in enumerate(dims): + src_logical_indices[dim] = src_logical_indices[dim] + dst_global[dst_dim] value = const_reg_ctx.get_value(src_tensor, src_logical_indices) self.buffer_store(buf=dst_var, indices=[dst_local], value=value) @@ -74,8 +74,8 @@ def _emit_const_reg(self, inst, dst_tensor, dst_layout, src_tensor, src_layout, ) src_logical_indices = list(inst.offsets) dims = range(len(src_layout.shape)) if inst.dims is None else inst.dims - for dim in dims: - src_logical_indices[dim] = src_logical_indices[dim] + dst_global[dim] + for dst_dim, dim in enumerate(dims): + src_logical_indices[dim] = src_logical_indices[dim] + dst_global[dst_dim] dst_expr = const_reg_ctx.get_value(src_tensor, src_logical_indices) const_reg_ctx.register(dst_tensor, axes=dst_axes, expr=dst_expr) @@ -95,8 +95,8 @@ def emit(self, inst: SliceAssignInst) -> None: # type: ignore src_indices = src_layout.get_global(spatial_index=self.current_thread, local_index=src_local) dst_indices = list(inst.offsets) dims = range(len(src_layout.shape)) if inst.dims is None else inst.dims - for dim in dims: - dst_indices[dim] = dst_indices[dim] + src_indices[dim] + for src_dim, dim in enumerate(dims): + dst_indices[dim] = dst_indices[dim] + src_indices[src_dim] dst_local = dst_layout.get_local(global_indices=dst_indices) self.buffer_store(buf=dst_var, indices=[dst_local], value=src_var[src_local]) diff --git a/tests/instructions/test_slice_register.py b/tests/instructions/test_slice_register.py new file mode 100644 index 00000000..0c6d5fd5 --- /dev/null +++ b/tests/instructions/test_slice_register.py @@ -0,0 +1,155 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import re +from pathlib import Path + +import pytest +import tilus +import torch +from tilus import int32 +from tilus.backends.codegen import generate_ir_module +from tilus.drivers import BuildOptions, optimize_ir_module, optimize_program +from tilus.hidet.backend.codegen import codegen +from tilus.ir.layout.ops import local, spatial +from tilus.target import nvgpu_sm80, scope +from tilus.testing import requires + + +class ReadSlice(tilus.Script): + def __init__(self, dim): + super().__init__() + self.dim = dim + + def __call__(self, x: ~int32, y: ~int32): + self.attrs.blocks = 1 + self.attrs.warps = 1 + + g_x = self.global_view(x, dtype=int32, shape=[4, 4]) + g_y = self.global_view(y, dtype=int32, shape=[4]) + with self.thread_group(0, num_threads=1): + r = self.load_global(g_x, offsets=[0, 0], shape=[4, 4]) + self.annotate_layout(r, local(4, 4)) + if self.dim == 0: + s = r[1, :] + else: + s = r[:, 1] + self.store_global(g_y, s, offsets=[0]) + + +class AssignSlice(tilus.Script): + def __init__(self, dim): + super().__init__() + self.dim = dim + + def __call__(self, x: ~int32, y: ~int32): + self.attrs.blocks = 1 + self.attrs.warps = 1 + + g_x = self.global_view(x, dtype=int32, shape=[4, 4]) + g_y = self.global_view(y, dtype=int32, shape=[4]) + with self.thread_group(0, num_threads=1): + r = self.load_global(g_x, offsets=[0, 0], shape=[4, 4]) + self.annotate_layout(r, local(4, 4)) + v = self.load_global(g_y, offsets=[0], shape=[4]) + self.annotate_layout(v, local(4)) + if self.dim == 0: + r[1, :] = v + else: + r[:, 1] = v + self.store_global(g_x, r, offsets=[0, 0]) + + +class ReadColumnAcrossThreads(tilus.Script): + def __call__(self, x: ~int32, y: ~int32): + self.attrs.blocks = 1 + self.attrs.warps = 1 + + g_x = self.global_view(x, dtype=int32, shape=[4, 4]) + g_y = self.global_view(y, dtype=int32, shape=[4]) + with self.thread_group(0, num_threads=4): + r = self.load_global(g_x, offsets=[0, 0], shape=[4, 4]) + self.annotate_layout(r, spatial(4, 1).local(1, 4)) + self.store_global(g_y, r[:, 1], offsets=[0]) + + +def generated_source(script, tmp_path: Path) -> str: + with scope(nvgpu_sm80): + program = script._jit_instance_for(None, None).transpiled_programs[0] # type: ignore + ir_module = generate_ir_module(optimize_program(program, BuildOptions(), tmp_path)) + ir_module = optimize_ir_module(ir_module, tmp_path) + codegen(ir_module, src_out_path=str(tmp_path / "source.cu"), target="cuda") + return (tmp_path / "source.cu").read_text() + + +def copied_indices(source: str, pattern: str) -> list[int]: + """Evaluate the register index of the matching element copy for each iteration of its loop.""" + match = re.search(pattern, source) + return [eval(match["expr"], {match["var"]: i}) for i in range(4)] + + +@pytest.mark.parametrize("dim, expected", [(0, [4, 5, 6, 7]), (1, [1, 5, 9, 13])]) +def test_read_slice_codegen(dim, expected, tmp_path: Path): + source = generated_source(ReadSlice(dim), tmp_path) + assert copied_indices(source, r"slice_regs\[(?P\w+)\] = regs\[(?P.+)\];") == expected + + +@pytest.mark.parametrize("dim, expected", [(0, [4, 5, 6, 7]), (1, [1, 5, 9, 13])]) +def test_assign_slice_codegen(dim, expected, tmp_path: Path): + source = generated_source(AssignSlice(dim), tmp_path) + assert copied_indices(source, r"regs\[(?P.+)\] = regs_1\[(?P\w+)\];") == expected + + +def matrix(device: str) -> torch.Tensor: + return torch.arange(16, dtype=torch.int32, device=device).reshape(4, 4) + + +@requires.nvgpu_sm80 +@pytest.mark.parametrize("dim", [0, 1]) +def test_read_slice(dim): + device = "cuda" if torch.cuda.is_available() else "cpu" + x = matrix(device) + y = torch.full((4,), -1, dtype=torch.int32, device=device) + + ReadSlice(dim)(x, y) + + torch.testing.assert_close(y, x[1, :] if dim == 0 else x[:, 1]) + + +@requires.nvgpu_sm80 +def test_read_column_across_threads(): + device = "cuda" if torch.cuda.is_available() else "cpu" + x = matrix(device) + y = torch.full((4,), -1, dtype=torch.int32, device=device) + + ReadColumnAcrossThreads()(x, y) + + torch.testing.assert_close(y, x[:, 1]) + + +@requires.nvgpu_sm80 +@pytest.mark.parametrize("dim", [0, 1]) +def test_assign_slice(dim): + device = "cuda" if torch.cuda.is_available() else "cpu" + x = matrix(device) + y = torch.arange(100, 104, dtype=torch.int32, device=device) + expected = x.clone() + if dim == 0: + expected[1, :] = y + else: + expected[:, 1] = y + + AssignSlice(dim)(x, y) + + torch.testing.assert_close(x, expected) From dd43c370a0cb6e690ae74c80d45a0397a6b3dae5 Mon Sep 17 00:00:00 2001 From: Yaoyao Ding Date: Mon, 9 Nov 2026 12:24:34 -0500 Subject: [PATCH 2/2] Address review: guard slice assigns along thread-distributed dropped dims When a slice assign drops a destination dim that is distributed across threads (e.g. r[1, :] = v with r in spatial(4, 1).local(1, 4)), every thread wrote its local copy and overwrote the other rows. Only threads that own the selected slice now write. Tests: cover multi-thread assigns and a 3-D slice with non-contiguous retained dims, assert regex matches before use, and drop the dead CPU device fallback. Co-Authored-By: Claude Opus 5.5 --- python/tilus/backends/emitters/regs.py | 20 ++++- tests/instructions/test_slice_register.py | 98 ++++++++++++++++++++--- 2 files changed, 104 insertions(+), 14 deletions(-) diff --git a/python/tilus/backends/emitters/regs.py b/python/tilus/backends/emitters/regs.py index aa0664f1..60101228 100644 --- a/python/tilus/backends/emitters/regs.py +++ b/python/tilus/backends/emitters/regs.py @@ -14,6 +14,7 @@ # limitations under the License. from tilus.backends.emitter import BaseInstEmitter, register_emitter +from tilus.hidet.ir.expr import logical_and from tilus.ir import RegisterTensor from tilus.ir.instructions import SliceAssignInst, SliceRegisterInst from tilus.ir.layout import RegisterLayout @@ -91,13 +92,28 @@ def emit(self, inst: SliceAssignInst) -> None: # type: ignore dst_var = self.tensor2var[dst_tensor] src_var = self.tensor2var[src_tensor] + dims = range(len(src_layout.shape)) if inst.dims is None else inst.dims + + # A dropped destination dim that is distributed across threads selects the threads that own the + # assigned slice; get_local() ignores spatial modes, so the other threads must not write. + spatial_modes = {mode for mode in dst_layout.spatial_modes if mode >= 0} + guarded_dims = [ + dim + for dim in range(len(dst_layout.shape)) + if dim not in dims and any(mode in spatial_modes for mode in dst_layout.grouped_modes[dim]) + ] + with self.for_range(extent=src_layout.local_size) as src_local: src_indices = src_layout.get_global(spatial_index=self.current_thread, local_index=src_local) dst_indices = list(inst.offsets) - dims = range(len(src_layout.shape)) if inst.dims is None else inst.dims for src_dim, dim in enumerate(dims): dst_indices[dim] = dst_indices[dim] + src_indices[src_dim] dst_local = dst_layout.get_local(global_indices=dst_indices) - self.buffer_store(buf=dst_var, indices=[dst_local], value=src_var[src_local]) + if guarded_dims: + owned = dst_layout.get_global(spatial_index=self.current_thread, local_index=dst_local) + with self.if_then(logical_and(*[owned[dim] == dst_indices[dim] for dim in guarded_dims])): + self.buffer_store(buf=dst_var, indices=[dst_local], value=src_var[src_local]) + else: + self.buffer_store(buf=dst_var, indices=[dst_local], value=src_var[src_local]) self.tensor2var[dst_tensor] = dst_var diff --git a/tests/instructions/test_slice_register.py b/tests/instructions/test_slice_register.py index 0c6d5fd5..305d27dd 100644 --- a/tests/instructions/test_slice_register.py +++ b/tests/instructions/test_slice_register.py @@ -22,7 +22,7 @@ from tilus.backends.codegen import generate_ir_module from tilus.drivers import BuildOptions, optimize_ir_module, optimize_program from tilus.hidet.backend.codegen import codegen -from tilus.ir.layout.ops import local, spatial +from tilus.ir.layout.ops import local, replicated, spatial from tilus.target import nvgpu_sm80, scope from tilus.testing import requires @@ -84,6 +84,51 @@ def __call__(self, x: ~int32, y: ~int32): self.store_global(g_y, r[:, 1], offsets=[0]) +class AssignAcrossThreads(tilus.Script): + def __init__(self, dim): + super().__init__() + self.dim = dim + + def __call__(self, x: ~int32, y: ~int32): + self.attrs.blocks = 1 + self.attrs.warps = 1 + + g_x = self.global_view(x, dtype=int32, shape=[4, 4]) + g_y = self.global_view(y, dtype=int32, shape=[4]) + with self.thread_group(0, num_threads=4): + r = self.load_global(g_x, offsets=[0, 0], shape=[4, 4]) + self.annotate_layout(r, spatial(4, 1).local(1, 4)) + v = self.load_global(g_y, offsets=[0], shape=[4]) + if self.dim == 0: + # the dropped row dim is distributed across threads: only the owner of row 1 may write + self.annotate_layout(v, replicated(4, num_workers=4)) + r[1, :] = v + else: + self.annotate_layout(v, spatial(4)) + r[:, 1] = v + self.store_global(g_x, r, offsets=[0, 0]) + + +class Slice3D(tilus.Script): + """Read and assign r[:, 1, :], whose retained dims (0 and 2) are not contiguous.""" + + def __call__(self, x: ~int32, y: ~int32, z: ~int32): + self.attrs.blocks = 1 + self.attrs.warps = 1 + + g_x = self.global_view(x, dtype=int32, shape=[2, 3, 4]) + g_y = self.global_view(y, dtype=int32, shape=[2, 4]) + g_z = self.global_view(z, dtype=int32, shape=[2, 4]) + with self.thread_group(0, num_threads=1): + r = self.load_global(g_x, offsets=[0, 0, 0], shape=[2, 3, 4]) + self.annotate_layout(r, local(2, 3, 4)) + self.store_global(g_z, r[:, 1, :], offsets=[0, 0]) + v = self.load_global(g_y, offsets=[0, 0], shape=[2, 4]) + self.annotate_layout(v, local(2, 4)) + r[:, 1, :] = v + self.store_global(g_x, r, offsets=[0, 0, 0]) + + def generated_source(script, tmp_path: Path) -> str: with scope(nvgpu_sm80): program = script._jit_instance_for(None, None).transpiled_programs[0] # type: ignore @@ -96,6 +141,7 @@ def generated_source(script, tmp_path: Path) -> str: def copied_indices(source: str, pattern: str) -> list[int]: """Evaluate the register index of the matching element copy for each iteration of its loop.""" match = re.search(pattern, source) + assert match is not None, f"no element copy matching {pattern!r} in the generated source" return [eval(match["expr"], {match["var"]: i}) for i in range(4)] @@ -111,16 +157,15 @@ def test_assign_slice_codegen(dim, expected, tmp_path: Path): assert copied_indices(source, r"regs\[(?P.+)\] = regs_1\[(?P\w+)\];") == expected -def matrix(device: str) -> torch.Tensor: - return torch.arange(16, dtype=torch.int32, device=device).reshape(4, 4) +def matrix() -> torch.Tensor: + return torch.arange(16, dtype=torch.int32, device="cuda").reshape(4, 4) @requires.nvgpu_sm80 @pytest.mark.parametrize("dim", [0, 1]) def test_read_slice(dim): - device = "cuda" if torch.cuda.is_available() else "cpu" - x = matrix(device) - y = torch.full((4,), -1, dtype=torch.int32, device=device) + x = matrix() + y = torch.full((4,), -1, dtype=torch.int32, device="cuda") ReadSlice(dim)(x, y) @@ -129,9 +174,8 @@ def test_read_slice(dim): @requires.nvgpu_sm80 def test_read_column_across_threads(): - device = "cuda" if torch.cuda.is_available() else "cpu" - x = matrix(device) - y = torch.full((4,), -1, dtype=torch.int32, device=device) + x = matrix() + y = torch.full((4,), -1, dtype=torch.int32, device="cuda") ReadColumnAcrossThreads()(x, y) @@ -141,9 +185,8 @@ def test_read_column_across_threads(): @requires.nvgpu_sm80 @pytest.mark.parametrize("dim", [0, 1]) def test_assign_slice(dim): - device = "cuda" if torch.cuda.is_available() else "cpu" - x = matrix(device) - y = torch.arange(100, 104, dtype=torch.int32, device=device) + x = matrix() + y = torch.arange(100, 104, dtype=torch.int32, device="cuda") expected = x.clone() if dim == 0: expected[1, :] = y @@ -153,3 +196,34 @@ def test_assign_slice(dim): AssignSlice(dim)(x, y) torch.testing.assert_close(x, expected) + + +@requires.nvgpu_sm80 +@pytest.mark.parametrize("dim", [0, 1]) +def test_assign_slice_across_threads(dim): + x = matrix() + y = torch.arange(100, 104, dtype=torch.int32, device="cuda") + expected = x.clone() + if dim == 0: + expected[1, :] = y + else: + expected[:, 1] = y + + AssignAcrossThreads(dim)(x, y) + + torch.testing.assert_close(x, expected) + + +@requires.nvgpu_sm80 +def test_slice_3d_non_contiguous_dims(): + x = torch.arange(24, dtype=torch.int32, device="cuda").reshape(2, 3, 4) + y = torch.arange(100, 108, dtype=torch.int32, device="cuda").reshape(2, 4) + z = torch.full((2, 4), -1, dtype=torch.int32, device="cuda") + expected_x = x.clone() + expected_x[:, 1, :] = y + expected_z = x[:, 1, :].clone() + + Slice3D()(x, y, z) + + torch.testing.assert_close(z, expected_z) + torch.testing.assert_close(x, expected_x)