tile_matmul ignores tile_view strides, silently computing wrong results on CUDA

[etaoxing]

(I have a fix)

Bug Description

"""Repro: ``wp.tile_matmul`` silently computes incorrect results for a strided tile operand.

``wp.tile_view`` keeps the parent tile's strides, but ``tile_matmul``'s cuBLASDx (MathDx)
path derives the leading dimension from the operand's shape alone and never consults its
strides, so a sub-tile view of a wider parent is read with the wrong row stride. No error
or warning is raised. The scalar (CPU) fallback is stride-generic and unaffected.

Both kernels below compute ``a[:, :16] @ b``: one via a (16, 16) view of a (16, 32) tile
(row stride still 32), one via a dense (16, 16) tile. On affected builds only the dense
variant matches the NumPy reference; on fixed builds both match.

Requires a CUDA device and a Warp build with MathDx (libmathdx) support.
"""

import numpy as np

import warp as wp

wp.init()

TILE_M = wp.constant(16)
TILE_K = wp.constant(16)
TILE_N = wp.constant(16)
PARENT_K = wp.constant(2 * TILE_K)  # inner extent of the wider parent tile

device = "cuda"


@wp.kernel
def matmul_via_view(a: wp.array2d(dtype=float), b: wp.array2d(dtype=float), c: wp.array2d(dtype=float)):
    # Load the full (TILE_M, PARENT_K) tile, then view its left half: the view has
    # shape (TILE_M, TILE_K) but keeps the parent's row stride of PARENT_K.
    a_full = wp.tile_load(a, shape=(TILE_M, PARENT_K), offset=(0, 0))
    a_sub = wp.tile_view(a_full, offset=(0, 0), shape=(TILE_M, TILE_K))
    b_tile = wp.tile_load(b, shape=(TILE_K, TILE_N), offset=(0, 0))
    acc = wp.tile_zeros(shape=(TILE_M, TILE_N), dtype=float)
    wp.tile_matmul(a_sub, b_tile, acc)
    wp.tile_store(c, acc, offset=(0, 0))


@wp.kernel
def matmul_via_dense(a: wp.array2d(dtype=float), b: wp.array2d(dtype=float), c: wp.array2d(dtype=float)):
    # Identical math: load the left half directly as a dense (TILE_M, TILE_K) tile
    # (tile_load handles the source array's strides).
    a_tile = wp.tile_load(a, shape=(TILE_M, TILE_K), offset=(0, 0))
    b_tile = wp.tile_load(b, shape=(TILE_K, TILE_N), offset=(0, 0))
    acc = wp.tile_zeros(shape=(TILE_M, TILE_N), dtype=float)
    wp.tile_matmul(a_tile, b_tile, acc)
    wp.tile_store(c, acc, offset=(0, 0))


rng = np.random.default_rng(0)
a_np = rng.standard_normal((TILE_M, PARENT_K)).astype(np.float32)
b_np = rng.standard_normal((TILE_K, TILE_N)).astype(np.float32)

a = wp.array(a_np, dtype=wp.float32, device=device)
b = wp.array(b_np, dtype=wp.float32, device=device)
c_view = wp.zeros((TILE_M, TILE_N), dtype=wp.float32, device=device)
c_dense = wp.zeros((TILE_M, TILE_N), dtype=wp.float32, device=device)

wp.launch_tiled(matmul_via_view, dim=(1, 1), inputs=[a, b], outputs=[c_view], block_dim=64, device=device)
wp.launch_tiled(matmul_via_dense, dim=(1, 1), inputs=[a, b], outputs=[c_dense], block_dim=64, device=device)
wp.synchronize_device(device)

reference = a_np[:, :TILE_K] @ b_np
err_dense = np.abs(c_dense.numpy() - reference).max()
err_view = np.abs(c_view.numpy() - reference).max()

print(f"dense operand:        max abs err vs reference = {err_dense:.3e}")
print(f"strided view operand: max abs err vs reference = {err_view:.3e}")
if err_view > 1e-3:
    print("Bug reproduced: tile_matmul ignored the view's strides.")
else:
    print("Strided view handled correctly.")

System Information

Software
  Warp (package):   1.16.0.dev0
  Warp (core):      1.16.0.dev0
  warp-clang:       1.16.0.dev0 (LLVM 18.1.3)
  NumPy:            2.5.0
  Python:           3.12.3 (main, Jun 19 2026, 12:46:00) [GCC 13.3.0]
  Platform:         Linux-6.17.0-35-generic-x86_64-with-glibc2.39

CUDA
  Enabled:          True
  Toolkit:          12.8
  Driver:           13.2
  NVRTC:            12.8
  Forward compat:   False

Libraries
  MathDx:           0.3.1
  cuBQL:            enabled
  NanoVDB:          32.8.0

Build
  Compiler:         GCC 9.5.0
  Debug:            False
  Verify FP:        False
  Fast math:        False
  Sanitizer:        none

CPU
  Model:            arrowlake-s
  ISA features:     61
                      64bit, adx, aes, avx, avx2, avxifma, avxneconvert, avxvnni, ...

Devices
  cpu
    Name:             x86_64
  cuda:0
    Name:             NVIDIA GeForce RTX 5090
    Memory:           31.4 GiB
    Arch:             sm_120
    SMs:              170
    PCI:              00000000:02:00
    Mempool:          enabled
    GPU->CPU mem:     True
    CPU->GPU mem:     False
    CPU/GPU atomics:  False