[Feature] Support None type as input for `T.ptr` and `T.Tensor` #1114

xwhzz · 2025-10-23T10:58:32Z

For issue #828

Summary by CodeRabbit

New Features
- Support for optional or missing buffers in kernel calls, allowing kernels to accept null inputs where appropriate.
Tests
- New test suite validating JIT GEMM-like kernels with tiled/pipelined execution and optional bias, verifying results against reference matmul.
Chores
- Internal type-annotation update to modern union syntax (no runtime behavior changes).

coderabbitai · 2025-10-23T10:58:41Z

Walkthrough

Adds a new test module exercising TileLang JIT tiled GEMM kernels with optional/nullptr bias handling and PyTorch validation; the Cython wrapper now accepts None kernel arguments and passes null pointers; a type-annotation change updates alloc_var to use PEP 604 union syntax.

Changes

Cohort / File(s)	Change Summary
New nullptr/bias test module `testing/python/jit/test_tilelang_jit_nullptr.py`	Adds `ptr_null_test`, `tensor_null_test`, `run_test`, and `test_nullptr`: JIT-wrapped tiled GEMM kernels using shared-memory tiles and a 3-stage pipelined K-loop, optional bias via nullptr or explicit tensor, orchestration and PyTorch-based validation.
Cython wrapper null pointer support `tilelang/jit/adapter/cython/cython_wrapper.pyx`	Extends kernel argument construction to map `None` entries in `tensor_list` to null pointers (`ctypes.c_void_p(0)`), preserving existing handling of tensors and scalars.
Type-annotation update `tilelang/language/allocate.py`	Changes `alloc_var` signature annotation from `Union[PrimExpr] = None` to PEP 604 style `PrimExpr

Sequence Diagram(s)

sequenceDiagram
  participant Test as Test script
  participant Builder as JIT kernel builder
  participant Cython as Cython wrapper
  participant Runtime as Native kernel
  participant Ref as PyTorch

  Note over Test,Builder: Prepare inputs (A,B,C, Bias or None)
  Test->>Builder: build/compile JIT kernel (ptr_null_test / tensor_null_test)
  Test->>Cython: invoke kernel with args (some args may be None)
  Cython->>Runtime: pass arguments (None -> null pointer)
  Runtime-->>Cython: execute tiled GEMM (shared tiles, pipelined K-loop), apply bias if pointer non-null
  Cython-->>Test: return output C
  Test->>Ref: compute reference (A @ B.T [+ bias])
  Test->>Test: compare outputs (validate)

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~20 minutes

Possibly related PRs

[Language] Support tilelang alloc_var(dtype, init=x) #1092 — Touches alloc_var in tilelang/language/allocate.py; related to signature/initializer changes and may overlap with the PEP 604 annotation edit in this PR.

Poem

🐰 I hopped through tiles and tiny loops,
Nulls tucked like carrots in kernel groups,
Cython said "bring me nothing" — I complied,
Bias present or missing, numbers coincide,
I twitched my nose and checked the tides. 🥕

Pre-merge checks and finishing touches

✅ Passed checks (3 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title Check	✅ Passed	The pull request title "[Feature] Support None type as input for `T.ptr` and `T.Tensor`" directly aligns with the main change in the changeset. The modifications add handling for None entries in the Cython wrapper by converting them to null pointers, implement tests validating this nullptr functionality, and update type annotations. The title is concise, specific, and clearly communicates the primary feature being implemented without ambiguity or misleading information. A teammate reviewing the commit history would immediately understand that this PR introduces None/nullptr support for tensor inputs.
Docstring Coverage	✅ Passed	Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%.

✨ 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 3f40b23 and 5466e73.

📒 Files selected for processing (1)

tilelang/language/allocate.py (2 hunks)

⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (1)

GitHub Check: Quick Lint

🔇 Additional comments (2)

tilelang/language/allocate.py (2)

17-17: LGTM! Good modernization.

The addition of from __future__ import annotations enables PEP 604 union syntax and is necessary for the type annotation update on line 70 to work in Python versions before 3.10.

70-70: LGTM! Clean type annotation modernization.

The update to PEP 604 union syntax (PrimExpr | None) is correct and improves code clarity without changing runtime behavior.

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

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

github-actions · 2025-10-23T10:58:47Z

👋 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! 🚀

LeiWang1999 · 2025-10-23T11:49:57Z

@codex review

chatgpt-codex-connector

💡 Codex Review

Here are some automated review suggestions for this pull request.

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

Open a pull request for review
Mark a draft as ready
Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

chatgpt-codex-connector · 2025-10-23T11:54:17Z

tilelang/jit/adapter/cython/cython_wrapper.pyx

                    if dtype not in dtype_to_ctype:
                        raise ValueError(f"Unsupported tensor dtype: {dtype}")
                    call_args.append(dtype_to_ctype[dtype](tensor))
+            elif tensor is None:
+                call_args.append(ctypes.c_void_p(0))
            else:
                raise ValueError(f"Unsupported tensor type: {type(tensor)}")



Handle optional tensors consistently across adapters

Adding elif tensor is None: call_args.append(ctypes.c_void_p(0)) makes the Cython path accept null inputs, but the other adapters (e.g. the ctypes and nvrtc wrappers) still always access .shape/.stride on every argument when building dynamic symbolics. Passing None for a T.ptr/T.Tensor with those backends will therefore still raise an AttributeError instead of forwarding a null pointer, despite the feature claiming generic support for None. To avoid backend‑dependent failures, the same None handling or a guard in their dynamic dimension logic is needed.

Useful? React with 👍 / 👎.

…-ai#1114) * [Feature] Support None type as input for T.ptr and T.Tensor * lint * lint * lint * lint fix

[Feature] Support None type as input for T.ptr and T.Tensor

ac7fe0b

xwhzz added 4 commits October 23, 2025 11:04

lint

3f40b23

lint

528f53b

lint

f8ccff3

lint fix

5466e73

LeiWang1999 approved these changes Oct 23, 2025

View reviewed changes

chatgpt-codex-connector bot reviewed Oct 23, 2025

View reviewed changes

LeiWang1999 merged commit 50e789d into tile-ai:main Oct 23, 2025
6 checks passed

This was referenced Oct 24, 2025

[Language] Initial version of tilelang frontend v2 #1120

Merged

[BugFix] alloc_var init failed to handle complex expression #1144

Merged

This was referenced Nov 12, 2025

[Fix] Fix a type that make wrong T.macro backtrace #1234

Merged

[Language] Add type stubs for tir op #1239

Merged

coderabbitai bot mentioned this pull request Nov 17, 2025

[FFI] Use tvm ffi as the default execution backend #1259

Merged

This was referenced Nov 21, 2025

[Feat] Add missing support for uint32x2, add unsigned implicit cast in bitwise op, add T.Ref as macro annotation #1302

Closed

[Fix] Remove unused let_bindings_ in CodeGenC to fix #1300 #1305

Merged

[Fix] Fix frame scope error in T.macro #1308

Merged

RubiaCx pushed a commit to RubiaCx/tilelang that referenced this pull request Nov 24, 2025

[Feature] Support None type as input for T.ptr and T.Tensor (tile…

371ec2f

…-ai#1114) * [Feature] Support None type as input for T.ptr and T.Tensor * lint * lint * lint * lint fix

coderabbitai bot mentioned this pull request Nov 27, 2025

[Refactor] Improve assertion handling in CodeGenCHost and ArgBinder #1352

Merged

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

[Feature] Support None type as input for `T.ptr` and `T.Tensor` #1114

[Feature] Support None type as input for `T.ptr` and `T.Tensor` #1114

Uh oh!

xwhzz commented Oct 23, 2025 •

edited by coderabbitai bot

Loading

Uh oh!

coderabbitai bot commented Oct 23, 2025 •

edited

Loading

Uh oh!

github-actions bot commented Oct 23, 2025

Uh oh!

LeiWang1999 commented Oct 23, 2025

Uh oh!

chatgpt-codex-connector bot left a comment

Uh oh!

chatgpt-codex-connector bot Oct 23, 2025

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

[Feature] Support None type as input for T.ptr and T.Tensor #1114

[Feature] Support None type as input for T.ptr and T.Tensor #1114

Uh oh!

Conversation

xwhzz commented Oct 23, 2025 • edited by coderabbitai bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary by CodeRabbit

Uh oh!

coderabbitai bot commented Oct 23, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Walkthrough

Changes

Sequence Diagram(s)

Estimated code review effort

Possibly related PRs

Poem

Pre-merge checks and finishing touches

Uh oh!

github-actions bot commented Oct 23, 2025

Uh oh!

LeiWang1999 commented Oct 23, 2025

Uh oh!

chatgpt-codex-connector bot left a comment

Choose a reason for hiding this comment

💡 Codex Review

Uh oh!

chatgpt-codex-connector bot Oct 23, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

[Feature] Support None type as input for `T.ptr` and `T.Tensor` #1114

[Feature] Support None type as input for `T.ptr` and `T.Tensor` #1114

xwhzz commented Oct 23, 2025 •

edited by coderabbitai bot

Loading

coderabbitai bot commented Oct 23, 2025 •

edited

Loading