[Enhancement] Remove constraint requiring last dimension stride to be 1

[LJC00118]

After #798, we can remove this constraint.

Summary by CodeRabbit

• New Features

  • Accept tensors with a non‑unit stride on the last dimension, reducing errors for non‑contiguous layouts and enabling more flexible memory layouts.

• Tests

  • Added CUDA vectorization tests that validate kernel vector-width selection and numerical correctness across varied stride configurations.

[github-actions]

👋 Hi! Thank you for contributing to the TileLang project.

Please remember to run pre-commit run --all-files in the root directory of the project to ensure your changes are properly linted and formatted. This will help ensure your contribution passes the format check.

We appreciate you taking this step! Our team will review your contribution, and we look forward to your awesome work! 🚀

[coderabbitai]

Walkthrough

Removed the runtime contiguity check in StridedTensorProxy.call (last-stride == 1) so it always forwards shape/strides/scope to its parent; added a CUDA-focused vectorization test module that runs kernels on varied strided inputs and asserts emitted vector types.

Changes

Cohort / File(s)	Summary
Strided tensor contiguity check removal tilelang/language/proxy.py	Removed the runtime check that enforced the last stride equals 1 and its associated error path; StridedTensorProxy.__call__ now forwards shape, strides, and scope to the parent unconditionally, permitting non-unit last strides.
Vectorization tests (CUDA) testing/python/language/test_tilelang_language_vectorize.py	Added tests: vectorize_test, run_vectorize, and test_vectorize. They compile/run CUDA kernels with varied strided inputs (using torch.as_strided), validate numerical correctness, and assert emitted vector types (e.g., float4, float2) in generated kernels.

Sequence Diagram(s)

sequenceDiagram
  autonumber
  participant C as Caller
  participant P as StridedTensorProxy
  participant B as ParentProxy

  rect rgba(230,230,255,0.35)
  note over P: Previous flow (before change)
  C->>P: __call__(shape, strides, scope)
  P->>P: Check last stride == 1
  alt last stride != 1
    P-->>C: Raise ValueError
  else
    P->>B: __call__(shape, strides, scope)
    B-->>P: Result
    P-->>C: Result
  end
  end

  rect rgba(230,255,230,0.35)
  note over P: New flow (after change)
  C->>P: __call__(shape, strides, scope)
  P->>B: __call__(shape, strides, scope)
  B-->>P: Result
  P-->>C: Result
  end

Loading

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~20 minutes

Possibly related PRs

• [Language] Introduce StridedTensor to support non contigious torch inputs #722: Introduced StridedTensorProxy with the last-stride==1 contiguity validation; directly related to this PR which removes that check.

Poem

I hop through strides with nimble toes,
No last-step chain to block my rows.
Kernels hum and vectors play,
Noncontig paths now find their way.
A little rabbit cheers, hip-hip-hurray! 🐇

Pre-merge checks and finishing touches

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 0.00% which is insufficient. The required threshold is 80.00%.	You can run @coderabbitai generate docstrings to improve docstring coverage.

✅ Passed checks (2 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title Check	✅ Passed	The PR title "[Enhancement] Remove constraint requiring last dimension stride to be 1" directly and accurately describes the primary change in the pull request. The main modification, as evidenced by the raw summary, is the removal of runtime validation in StridedTensorProxy.call that enforces the last stride to be 1, allowing non-contiguous layouts. The title is specific, concise (71 characters), and uses clear language that immediately conveys the key change without vague or misleading terms. The secondary changes (adding test coverage) are appropriately supportive of the main change and do not detract from the title's accuracy.

✨ Finishing touches

• 📝 Generate docstrings

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Post copyable unit tests in a comment

---

📜 Recent review details

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between 466e53c and 2758b95.

📒 Files selected for processing (1)

• testing/python/language/test_tilelang_language_vectorize.py (1 hunks)

🚧 Files skipped from review as they are similar to previous changes (1)

• testing/python/language/test_tilelang_language_vectorize.py

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[LeiWang1999]

Would be better if we can add a simple test for this pr.

[coderabbitai]

Actionable comments posted: 1

📜 Review details

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between 3c9961e and 754c8b5.

📒 Files selected for processing (1)

• testing/python/language/test_tilelang_language_vectorize.py (1 hunks)

🧰 Additional context used

🧬 Code graph analysis (1)

testing/python/language/test_tilelang_language_vectorize.py (3)

tilelang/language/proxy.py (1)

• StridedTensor (254-255)

tilelang/language/kernel.py (1)

• threads (215-219)

tilelang/jit/__init__.py (1)

• compile (34-90)

🪛 GitHub Actions: CI Test on AMD

testing/python/language/test_tilelang_language_vectorize.py

[error] 13-14: F821 Undefined name float32. Likely a missing or misspelled type name in the StridedTensor dtype annotation.

🪛 GitHub Actions: CI Test on Metal

testing/python/language/test_tilelang_language_vectorize.py

[error] 13-15: format.sh: ruff check failed with F821 Undefined name float32 in test_tilelang_language_vectorize.py (likely missing import or type alias).

🪛 Ruff (0.14.0)

testing/python/language/test_tilelang_language_vectorize.py

13-13: Undefined name float32

(F821)

---

14-14: Undefined name float32

(F821)

🔇 Additional comments (3)

testing/python/language/test_tilelang_language_vectorize.py (3)

1-67: Excellent test design for non-contiguous vectorization.

This test effectively validates the removal of the last-stride constraint by exercising multiple non-contiguous layouts via torch.as_strided, compiling with vectorization enabled, and verifying both numerical correctness and the presence of expected vector types in the emitted CUDA code. The test design directly addresses the maintainer's request for test coverage.

---

27-27: Verify the stride constraint assumption.

The assertion stride_A >= N and stride_B >= N assumes strides must be at least as large as the first dimension. This makes sense given the tensor creation approach (lines 32-35), where base tensors have shape (stride_X, M) and views have stride (1, stride_X). However, confirm whether this constraint is fundamental or just a test simplification.

Consider adding a comment explaining why this constraint is required, or testing the opposite case if the constraint is artificial:

def run_vectorize(N, M, stride_A, stride_B):
    # Constraint: stride_A >= N and stride_B >= N
    # Required because base tensors are shaped (stride_X, M) to accommodate
    # the strided view of (N, M) with stride (1, stride_X)
    assert stride_A >= N and stride_B >= N

---

43-54: Vectorization width detection is correct.

The logic correctly computes the maximum vectorization width by doubling until either stride is no longer divisible, then validates the expected vector type appears in the generated code. The mapping of vectorize_size to CUDA vector types (ulonglong4 for 8-wide, float4 for 4-wide, float2 for 2-wide) is appropriate.

[coderabbitai]

⚠️ Potential issue | 🔴 Critical

🧩 Analysis chain

CI failure: Fix float32 dtype specification.

Multiple CI pipelines (AMD, Metal) are failing with F821 Undefined name 'float32' errors on these lines. The dtype specification syntax may be incorrect for StridedTensor type annotations.

Run the following script to check the correct dtype specification syntax in TileLang's codebase:

Additionally, search the web for TileLang's StridedTensor dtype annotation syntax:

---

🏁 Script executed:

#!/bin/bash
# Search for existing StridedTensor usage patterns to determine correct dtype syntax

# Find other StridedTensor type annotations with dtype specifications
rg -nP 'StridedTensor\[.*\]' --type=py -C3

# Find dtype-related imports or constants in tilelang.language
ast-grep --pattern $'import $$$
float32 = $$$
$$$'

# Check if there's a dtype module or enum
fd -e py -x rg -l 'class.*float32|Float32|FLOAT32' --type=py

Length of output: 791

---

🌐 Web query:

TileLang StridedTensor dtype annotation syntax examples

💡 Result:

TileLang uses T.Tensor-style annotations (dtype as a string or a variable) rather than a separate "StridedTensor" keyword. Examples from the official docs / repo:

• Function argument tensor annotation:
  @T.prim_func
  def main(A: T.Tensor((M, K), dtype), B: T.Tensor((K, N), dtype), C: T.Tensor((M, N), dtype)): ...
  (dtype can be a string literal like "float16" or a parameter variable). [1][2]

• alloc / storage annotations include dtype similarly:
  A_shared = T.alloc_shared((block_M, block_N), dtype)
  C_L0 = T.alloc_L0C((block_M, block_N), accum_dtype) # accum_dtype passed as string/param. [2][3]

• Elementwise example with literal dtype:
  @T.prim_func
  def main(A: T.Tensor((N), "bfloat16"), B: T.Tensor((N), "bfloat16"), C: T.Tensor((N), "bfloat16")): ...
  (or use dtype="bfloat16" in an outer function and refer to dtype). [4]

Sources:
[1] TileLang API (language) — examples of T.Tensor annotations.
[2] TileLang GitHub repo — matmul / GEMM examples using dtype and accum_dtype.
[3] tilelang-ascend adapter examples (alloc_L0/alloc_L1 with dtype/accum_dtype).
[4] TileLang docs — elementwise operator example showing literal dtype.

If you want, I can extract exact copy-paste snippets from the linked pages.

---

🏁 Script executed:

#!/bin/bash
# Inspect StridedTensor and Tensor definitions and dtype handling in ir.py
rg -nP 'class\s+StridedTensor' -C10 tilelang/language/ast/ir.py
rg -nP 'class\s+Tensor' -C10 tilelang/language/ast/ir.py
rg -nP '__class_getitem__' -C10 tilelang/language/ast/ir.py
rg -n 'dtype' -C5 tilelang/language/ast/ir.py

Length of output: 22606

---

🏁 Script executed:

#!/bin/bash
# Locate StridedTensor definition or alias in the repository
rg -n 'StridedTensor' -C3

Length of output: 3211

---

🏁 Script executed:

#!/bin/bash
# Extract StridedTensorProxy and BaseTensorProxy definitions and __call__ methods
rg -nP 'class\s+StridedTensorProxy' -C30 tilelang/language/proxy.py
rg -nP 'class\s+BaseTensorProxy' -C30 tilelang/language/proxy.py
rg -nP 'def\s+__call__' -C10 tilelang/language/proxy.py

Length of output: 7349

---

Fix StridedTensor annotation syntax
Replace subscription with call syntax and name the dtype argument:

-testing/python/language/test_tilelang_language_vectorize.py:13-14
-    A: T.StridedTensor[(N, M), (1, stride_A), "float32"],
-    B: T.StridedTensor[(N, M), (1, stride_B), "float32"],
+    A: T.StridedTensor((N, M), (1, stride_A), dtype="float32"),
+    B: T.StridedTensor((N, M), (1, stride_B), dtype="float32"),

📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change

	A: T.StridedTensor[(N, M), (1, stride_A), "float32"],
	B: T.StridedTensor[(N, M), (1, stride_B), "float32"],
	A: T.StridedTensor((N, M), (1, stride_A), dtype="float32"),
	B: T.StridedTensor((N, M), (1, stride_B), dtype="float32"),

🧰 Tools

🪛 GitHub Actions: CI Test on AMD

[error] 13-14: F821 Undefined name float32. Likely a missing or misspelled type name in the StridedTensor dtype annotation.

🪛 GitHub Actions: CI Test on Metal

[error] 13-15: format.sh: ruff check failed with F821 Undefined name float32 in test_tilelang_language_vectorize.py (likely missing import or type alias).

🪛 Ruff (0.14.0)

13-13: Undefined name float32

(F821)

---

14-14: Undefined name float32

(F821)

🤖 Prompt for AI Agents

In testing/python/language/test_tilelang_language_vectorize.py around lines
13-14, the StridedTensor annotations use subscription syntax and an unnamed
dtype; change them to call syntax and pass the dtype as a named argument, e.g.,
replace T.StridedTensor[(N, M), (1, stride_A), "float32"] and
T.StridedTensor[(N, M), (1, stride_B), "float32"] with T.StridedTensor((N, M),
(1, stride_A), dtype="float32") and T.StridedTensor((N, M), (1, stride_B),
dtype="float32") respectively.

[LeiWang1999]

LGTM.