[CuTeDSL] Prevent domain_offset overflow with static strides

[kzos]

Fixes #3604.

What this fixes

cute.domain_offset can calculate the wrong global-memory address even when
the runtime coordinate itself fits in Int32. A static stride does not widen
the coordinate before multiplication, so the element offset can overflow
before it is converted to a 64-bit byte offset for pointer addition.

This change widens dynamic coordinate leaves before crd2idx for global/generic
memory pointers. It fixes the reported static-stride and rebuilt-static-stride
cases without requiring callers to make their layouts dynamic or manually cast
coordinates to Int64. The code change is confined to tensor.py; this PR also
adds 10 CPU IR regression tests.

Reproducer and failure mechanism

The full standalone reproducer in #3604
allocates a BF16 tensor with shape (2**29 + 1, 4) and compact strides (4, 1),
then uses a runtime Int32(2**29) to address the last row. Its essential operation
inside the kernel is:

# tensor: global BF16 tensor, layout (2**29 + 1, 4):(4, 1)
# row: runtime cutlass.Int32, value 2**29
shifted = cute.domain_offset((row, 0), tensor)
shifted[(0, 0)] = cutlass.BFloat16(1.0)

The coordinate is valid and fits in signed 32 bits. The intermediate product
does not:

Address-calculation step	Intended value	Unpatched static-stride path
Runtime row	2**29	2**29
Element offset: row times stride 4	2**31	Overflows signed 32-bit arithmetic to -2**31
Byte offset: element offset times 2 for BF16	+2**32 (+4 GiB)	Signed wide multiplication produces -2**32 (-4 GiB)

The resulting store goes before the allocation rather than to its last row.
Casting the result of crd2idx to Int64 would be too late: the narrow product
has already lost its intended value. BF16 exposes the byte-address consequence;
the coordinate-to-index overflow is not specific to BF16.

The issue's dynamic-stride variant succeeds because its runtime stride is
already i64, causing the coordinate/stride calculation to stay wide. Rebuilding
a layout with the static stride (4, 1) removes that wide operand and restores
the failure. It is the arithmetic width, not the presence of make_layout, that
distinguishes these cases.

Why this change is placed in domain_offset

domain_offset is the point where the coordinate becomes an offset for a
particular tensor iterator. For a global/generic Pointer, the patch recursively
promotes narrow dynamic coordinate leaves to Int64 before calling
crd2idx, so the coordinate/stride product is formed at the required width.

The implementation preserves the surrounding contracts:

• Python integer constants and None leaves stay unchanged, preserving static
  coordinate information and constant folding.
• Integer leaves already at least 64 bits wide are left alone. Signed narrow
  coordinates sign-extend; unsigned narrow coordinates zero-extend.
• get_divisibility and assume carry existing divisibility facts across the
  promotion, preserving the derived pointer-alignment information.
• transform_leaf handles nested coordinates rather than only flat tuples.
• Shared-memory pointers and non-pointer iterators, including identity tensors,
  retain their existing arithmetic. General-purpose crd2idx and CuTe integer
  algebra are not changed globally.

This addresses overflow introduced within this domain_offset calculation. It
cannot recover a coordinate that the caller has already overflowed, and it does
not make out-of-bounds coordinates valid. No performance improvement is claimed;
the effect of wider arithmetic on arbitrary kernels has not been benchmarked.

Check against the executed native code

For the BF16 row/stride case above, the executed baseline kernel emits:

shl.b32      %r2, %r1, 2;
mul.wide.s32 %rd2, %r2, 2;

Here the first instruction computes the element offset in 32 bits; the wide
multiply only widens that already-overflowed result. The fixed kernel instead
combines the stride and element size into a wide multiply of the original row:

mul.wide.s32 %rd2, %r1, 8;

Disassembly of 60 repeated baseline/fixed binaries confirms the corresponding
32-bit shift is removed and the wide multiply changes from factor 2 to 8.
A semantically neutral source variant retains identical native instructions
within each arm and does not fix the baseline failure. Baseline memcheck reports
a 2-byte global store 4 GiB before the nearest allocation, matching the predicted
negative byte offset.

Regression tests included in this PR

Run with an environment importing the changed CuTeDSL source:

python test/python/CuTeDSL/test_domain_offset.py

These tests build and verify MLIR on the CPU; they do not allocate a multi-GiB
tensor or require a GPU. 10/10 pass with the fix, versus 5/10 on the baseline.
They check widening before crd2idx, global/generic pointers, nested coordinates,
divisibility/alignment preservation, unsigned extension, static constants,
already-wide coordinates, dynamic-i64 strides, and unchanged shared-memory and
identity-tensor behavior. These are IR regressions, not a substitute for runtime
validation.

Additional B300 validation

A separate runtime harness exercised the real installed package on B300 SM103,
with source hashes checked against the baseline and this patch. Each case used
a fresh process and private compile/cache directory. Baseline failures were
isolated because an illegal address invalidates the CUDA context.

• Static-stride, rebuilt-static-layout, and neutral-source variants each produced
  10/10 baseline failures and 10/10 fixed passes.
• A 29-entry fixed matrix ran both normally and under Compute Sanitizer. It
  covers BF16/FP32/Int32 elements; row offsets 0, 2, 2**29 - 1, 2**29,
  2**29 + 1, and 2**30; and dynamic-i64, explicit-i64, nested, constant,
  neutral, and valid negative-offset cases. This is a selected matrix, not the
  Cartesian product of every option.
• All 29 fixed memcheck cases reported zero errors. A separate baseline
  memcheck reproduced the invalid global store described above.
• 590 fixed-path graph replays and 40 baseline-control replays passed.
  Each replay poisoned the tensor buffer and checked that only the intended
  element changed. Captured scalar arguments stayed fixed; mutation of those
  arguments was not tested.

Across the campaign: 90 fixed cases passed, 5 baseline controls passed, and
32 expected baseline failures were reproduced (127 total cases). Memcheck
cases and graph replays are included in those runs, not extra independent cases.
The GPU harness/results are additional validation, not tests added to this PR
or a report of upstream CI passing.

Toolchain and existing suite limitations

Validation used public nvidia-cutlass-dsl==4.8.0.dev0 with matching base/core/cu12
libraries, Torch 2.11.0+cu130, cuda-bindings 13.3.1, B300 driver 580.126.09, and
Compute Sanitizer 2025.4.1. This is not the exact dated alpha wheel in the report.
The unpatched wheel's tensor.py is byte-identical to the base source at
147295a3; the fixed installed module is byte-identical to this PR's version.

The broader public CuTeDSL Python suite is 39/43 passing with this change,
versus 34/43 on the baseline. Both retain the same four unrelated failures/errors:

• test_exec_defined_function_raises_user_error
• test_uncaught_dsl_error_does_not_kill_interactive_session
• test_preserves_attribute_order_and_deduplicates
• test_fastmath_flag

No unrelated changes are included. Both final GPU arms used lazy CUDA module
loading; earlier eager-loading attempts timed out in PyTorch CUDA initialization
before CuTe compilation/launch and are not counted as kernel results.

[kzos]

Hi @hwu36, could you help route this domain_offset fix for CuTeDSL review? It widens dynamic coordinates before static-stride multiplication for global/generic pointers, with shared-memory and non-pointer behavior unchanged. The PR includes 10 CPU IR regression tests and records separate baseline/fixed B300, memcheck, and CUDA Graph validation, including the remaining broader-suite limitations. Happy to address review feedback; no performance improvement is claimed. Thanks!

[hwu36]

@brandon-yujie-sun

[kzos]

Hi @brandon-yujie-sun, following up on hwu36's review routing: would you be able to review this CuTeDSL domain_offset overflow fix, or suggest the right reviewer? It widens dynamic coordinates before static-stride multiplication for global/generic pointers while preserving shared-memory and non-pointer behavior. The existing validation and its limits are recorded in the PR. Happy to address feedback. Thanks!