diff --git a/python/tilus/backends/emitters/regs.py b/python/tilus/backends/emitters/regs.py index 238c8865..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 @@ -40,11 +41,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 +62,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 +75,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) @@ -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 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]) + 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 new file mode 100644 index 00000000..305d27dd --- /dev/null +++ b/tests/instructions/test_slice_register.py @@ -0,0 +1,229 @@ +# 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, replicated, 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]) + + +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 + 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) + 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)] + + +@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() -> 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): + x = matrix() + y = torch.full((4,), -1, dtype=torch.int32, device="cuda") + + 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(): + x = matrix() + y = torch.full((4,), -1, dtype=torch.int32, device="cuda") + + ReadColumnAcrossThreads()(x, y) + + torch.testing.assert_close(y, x[:, 1]) + + +@requires.nvgpu_sm80 +@pytest.mark.parametrize("dim", [0, 1]) +def test_assign_slice(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 + + 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)