[BUG] domain_offset overflows runtime Int32 coordinate with static stride

Which component has the problem?

CuTe DSL

Bug Report

Describe the bug

cute.domain_offset can overflow a runtime Int32 coordinate when the tensor has a static stride. The element offset is calculated in 32 bits before conversion to a 64-bit byte offset, which can cause CUDA_ERROR_ILLEGAL_ADDRESS. The equivalent dynamic-stride layout uses 64-bit arithmetic and succeeds.

Steps/Code to reproduce bug

import argparse
import os

os.environ.setdefault("CUDA_LAUNCH_BLOCKING", "1")

import torch
import cutlass
import cutlass.cute as cute
from cutlass.cute.runtime import from_dlpack


ROWS = 2**29 + 1
COLS = 4
LAST_ROW = ROWS - 1


# IMA: from_dlpack(out) preserves the fully static compact layout:
#   !cute.layout<"(536870913,4):(4,1)">
@cute.kernel
def no_make_layout_static_stride(tensor: cute.Tensor, row: cutlass.Int32):
    # IR emitted by domain_offset:
    #   %idx = cute.crd2idx(%coord, %lay)
    #       : (...) -> !cute.int_tuple<"?{div=4}">
    #   %ptr = cute.add_offset(%iter, %int_tuple)
    shifted = cute.domain_offset((row, 0), tensor)
    shifted[(0, 0)] = cutlass.BFloat16(1.0)


# Works: the runtime row stride remains i64.
@cute.kernel
def make_layout_dynamic_stride(tensor: cute.Tensor, row: cutlass.Int32):
    layout = cute.make_layout(
        (tensor.shape[0], COLS),
        stride=(tensor.stride[0], 1),
    )
    # IR for layout:
    #   %lay_3 = cute.make_layout(%shape, %stride)
    #       : !cute.layout<"(?,4):(?{i64},1)">
    rebuilt = cute.make_tensor(tensor.iterator, layout)
    # IR emitted by domain_offset:
    #   %idx = cute.crd2idx(%coord, %lay_3)
    #       : (...) -> !cute.int_tuple<"?{i64}">
    #   %ptr = cute.add_offset(%iter, %int_tuple)
    shifted = cute.domain_offset((row, 0), rebuilt)
    shifted[(0, 0)] = cutlass.BFloat16(1.0)


# IMA: make_layout replaces the dynamic stride with the static stride (4, 1).
@cute.kernel
def make_layout_static_stride(tensor: cute.Tensor, row: cutlass.Int32):
    layout = cute.make_layout(
        (tensor.shape[0], COLS),
        stride=(COLS, 1),
    )
    # IR for layout:
    #   %lay_0 = cute.make_layout(%shape, %stride)
    #       : !cute.layout<"(?,4):(4,1)">
    rebuilt = cute.make_tensor(tensor.iterator, layout)
    # IR emitted by domain_offset:
    #   %idx = cute.crd2idx(%coord, %lay_0)
    #       : (...) -> !cute.int_tuple<"?{div=4}">
    #   %ptr = cute.add_offset(%iter, %int_tuple)
    shifted = cute.domain_offset((row, 0), rebuilt)
    shifted[(0, 0)] = cutlass.BFloat16(1.0)


@cute.jit
def launch_no_make_layout_static_stride(tensor: cute.Tensor, row: cutlass.Int32):
    no_make_layout_static_stride(tensor, row).launch(grid=[1, 1, 1], block=[1, 1, 1])


@cute.jit
def launch_make_layout_dynamic_stride(tensor: cute.Tensor, row: cutlass.Int32):
    make_layout_dynamic_stride(tensor, row).launch(grid=[1, 1, 1], block=[1, 1, 1])


@cute.jit
def launch_make_layout_static_stride(tensor: cute.Tensor, row: cutlass.Int32):
    make_layout_static_stride(tensor, row).launch(grid=[1, 1, 1], block=[1, 1, 1])


cases = (
    "no_make_layout_static_stride",
    "make_layout_dynamic_stride",
    "make_layout_static_stride",
)
parser = argparse.ArgumentParser()
parser.add_argument("case", choices=cases)
args = parser.parse_args()

out = torch.zeros((ROWS, COLS), device="cuda", dtype=torch.bfloat16)
row = cutlass.Int32(LAST_ROW)

if args.case == "no_make_layout_static_stride":
    launch_no_make_layout_static_stride(from_dlpack(out), row)
elif args.case == "make_layout_dynamic_stride":
    launch_make_layout_dynamic_stride(from_dlpack(out).mark_layout_dynamic(), row)
else:
    launch_make_layout_static_stride(from_dlpack(out).mark_layout_dynamic(), row)

torch.cuda.synchronize()
assert float(out[0, 0]) == 0.0
assert float(out[-1, 0]) == 1.0

$ python repro.py no_make_layout_static_stride
CUDA_ERROR_ILLEGAL_ADDRESS (error code: 700)

$ python repro.py make_layout_dynamic_stride
# succeeds

$ python repro.py make_layout_static_stride
CUDA_ERROR_ILLEGAL_ADDRESS (error code: 700)

Expected behavior

All three cases should write out[-1, 0] without an illegal memory access.

Environment details

• Docker
• nvidia-cutlass-dsl==4.8.0a0+20260828210702.38664cd
• CUDA 13.4
• GB300 (sm_103a)

Additional context

The static-stride path lowers the address calculation as:

shl.b32      %r2, %r1, 2;
mul.wide.s32 %rd2, %r2, 2;
add.s64      %rd3, %rd1, %rd2;

For row = 2**29, the first instruction produces 0x80000000; signed widening then generates a -4 GiB byte offset. BF16 is used here only to expose the incorrect index-width propagation through element-to-byte lowering; the coordinate-to-index issue is not BF16-specific. The dynamic-stride case keeps the index in 64 bits:

mul.lo.s64 %rd4, %rd2, %rd3;
shl.b64    %rd5, %rd4, 1;
add.s64    %rd6, %rd1, %rd5;

Start by running the supplied repro.py cases and inspect the lowering path for cute.domain_offset when a layout has a static stride, comparing it with the dynamic-stride path and the emitted PTX. Done means all three cases write out[-1, 0] successfully without CUDA_ERROR_ILLEGAL_ADDRESS and preserve the existing assertions.