[BugFix] alloc_var init failed to handle complex expression

[kurisu6912]

When init value of T.alloc_var is not const, tilelang generates wrong code. This pr fix this by fallback to assignment.

@T.prim_func
def test_var_assign(A: T.Tensor((2,), 'int32')):
    with T.Kernel(1) as _:
        a = T.alloc_var('int32', init=1)
        b = T.alloc_var('int32', init=a) # b gets value of a
        a = 2
        d = T.alloc_var('int32', init=a) # c gets new value of a
        A[0] = b
        A[1] = d

Summary by CodeRabbit

• New Features

  • Enhanced variable allocation with improved initialization handling and more flexible parameter patterns.
  • Added support for both primitive and non-primitive initializers in variable allocation.

• Tests

  • Introduced test coverage for variable initialization and assignment in TileLang language integration.

[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

Adds a new test for TileLang variable allocation and updates alloc_var with typing overloads and branching init logic: primitive inits set tl.local_var_init via block attributes; non-primitive inits use T.buffer_store. Also adds related imports and docstring examples.

Changes

Cohort / File(s)	Summary
Variable Initialization Test testing/python/language/test_tilelang_language_var_init.py	New test validating alloc_var usage in TileLang: allocation, aliasing, re-assignment, kernel generation, execution, and assertion of outputs.
Allocate Module Enhancement tilelang/language/allocate.py	Added typing.overload declarations for alloc_var, new imports (Buffer, IntImm, FloatImm), expanded docstring examples, and conditional init handling: primitive inits (int/float/IntImm/FloatImm) set tl.local_var_init via block attributes; non-primitive inits initialize via T.buffer_store.

Sequence Diagram

sequenceDiagram
    participant User
    participant alloc_var
    participant BlockAttr as block_attr\n(tl.local_var_init)
    participant BufferStore as T.buffer_store
    participant Buffer as Buffer (returned)

    User->>alloc_var: call alloc_var(dtype, init, scope)
    alt init is primitive (int/float/IntImm/FloatImm)
        alloc_var->>BlockAttr: emit tl.local_var_init attribute
        BlockAttr->>Buffer: allocation with init metadata
    else init is non-primitive
        alloc_var->>BufferStore: emit T.buffer_store to write init
        BufferStore->>Buffer: allocation with stored init
    end
    Buffer->>User: Buffer handle

Loading

Estimated code review effort

🎯 2 (Simple) | ⏱️ ~12 minutes

• Focus review on alloc_var branching between block attribute vs buffer_store.
• Check type signatures/overloads and their runtime behavior.
• Validate test covers aliasing and both init paths (primitive vs non-primitive).

Possibly related PRs

• [Language] Support tilelang alloc_var(dtype, init=x)  #1092: Overlapping changes to alloc_var init handling and kLocalVarInit propagation; likely the primary related PR.

Poem

🐰 I nibble bytes and hop on code,

Variables born on a tidy road.
Primitives get tags, deep inits get stored,
Kernels hum softly — allocations adored.
✨

Pre-merge checks and finishing touches

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 37.50% 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 "[BugFix] alloc_var init failed to handle complex expression" directly addresses the main change in the changeset. The core modification is in tilelang/language/allocate.py, where the alloc_var function has been enhanced to handle non-primitive initializers by using T.buffer_store instead of block_attr. The new test file validates this fix. The title is specific and clearly identifies both the affected function and the nature of the problem being fixed, making it easy for reviewers to understand the primary change without requiring full context of the implementation details.

✨ 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 d3c37e3 and 85119da.

📒 Files selected for processing (1)

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

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

• testing/python/language/test_tilelang_language_var_init.py

⏰ 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

---

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.}

[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 60567ba and d3c37e3.

📒 Files selected for processing (2)

• testing/python/language/test_tilelang_language_var_init.py (1 hunks)
• tilelang/language/allocate.py (4 hunks)

🧰 Additional context used

🧬 Code graph analysis (2)

tilelang/language/allocate.py (3)

src/op/builtin.h (1)

• tvm (13-517)

tilelang/language/ast/ir.py (3)

• block_attr (430-438)
• init (361-369)
• buffer_store (1263-1300)

testing/python/language/test_tilelang_language_alloc.py (1)

• alloc_var (4-23)

testing/python/language/test_tilelang_language_var_init.py (1)

tilelang/language/allocate.py (3)

• alloc_var (74-75)
• alloc_var (79-83)
• alloc_var (86-141)

⏰ 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 (6)

testing/python/language/test_tilelang_language_var_init.py (2)

1-10: LGTM! Test structure is well-formed.

The test setup correctly uses tilelang.jit and T.prim_func to create a kernel that exercises variable initialization behavior.

---

17-27: LGTM! Assertions validate the aliasing semantics.

The test correctly verifies that b captures the initial value of a (1) and d captures the reassigned value (2), confirming that variable initialization semantics work as expected for both primitive and complex expressions.

tilelang/language/allocate.py (4)

18-18: LGTM! Imports support the new type overloads and init logic.

The overload import enables proper type hints, while Buffer, IntImm, and FloatImm are used for return type annotations and runtime type checking of initializers.

Also applies to: 23-24

---

73-83: LGTM! Overload signatures correctly model the two main usage patterns.

The overloads provide precise type information for type checkers:

• First overload: positional init parameter
• Second overload: keyword-only init parameter

This matches the flexible argument parsing in the implementation (lines 113-130).

---

101-106: LGTM! Documentation examples are clear and comprehensive.

The examples cover all supported usage patterns and align with the overload signatures, making the API easy to understand.

---

137-140: LGTM! Conditional init logic correctly handles complex expressions.

This is the core fix:

• Primitives (int, float, IntImm, FloatImm) use the optimized block_attr path with tl.local_var_init
• Complex expressions (e.g., Buffer references, arithmetic expressions) use T.buffer_store to properly evaluate and store the result

This enables initializing variables with expressions like T.alloc_var('int32', init=some_buffer), which previously failed. The attribute name matches the constant defined in src/op/builtin.h.