[BugFix][Metal] Preserve pointer address spaces for byte offsets - #2925
Merged
Conversation
|
👋 Hi! Thank you for contributing to the TileLang project. Please remember to run We appreciate you taking this step! Our team will review your contribution, and we look forward to your awesome work! 🚀 |
Contributor
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Pro Plus Run ID: 📒 Files selected for processing (2)
🚧 Files skipped from review as they are similar to previous changes (2)
📝 WalkthroughWalkthroughThe Metal backend centralizes storage-scope mapping, validates typed pointer bindings, infers pointer metadata from supported expressions, and emits inferred address spaces for declarations and offsets. Tests cover source generation and MPS execution for shared-memory GEMM kernels. ChangesMetal address-space-aware pointer code generation
Estimated code review effort: 4 (Complex) | ~45 minutes Possibly related PRs
Sequence Diagram(s)sequenceDiagram
participant MetalTest
participant MetalCodegen
participant MPSBackend
MetalTest->>MetalCodegen: lower shared-memory GEMM
MetalCodegen->>MetalTest: return address-space-qualified source
MetalTest->>MPSBackend: execute through TVM FFI or Torch
MPSBackend->>MetalTest: return numerical result
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
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. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 4
🧹 Nitpick comments (3)
testing/python/metal/test_metal_address_space.py (2)
46-56: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueAssert that the kernel source exists instead of substituting an empty string.
artifact.kernel_source or ""converts a missing source into an empty string. The equality assertion on Line 64 then compares""with""and passes. The later alias assertion still catches the case, but the failure message loses the cause.♻️ Proposed change
- return artifact.kernel_source or "" + source = artifact.kernel_source + assert source, f"empty Metal kernel source (enable_device_compile={enable_device_compile})" + return source🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@testing/python/metal/test_metal_address_space.py` around lines 46 - 56, Update lower_to_metal so it asserts artifact.kernel_source is present before returning it, rather than converting a missing value to an empty string. Preserve the existing return type and ensure the assertion clearly identifies missing kernel source; leave the downstream equality and alias assertions unchanged.
84-92: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueSeed the random inputs.
torch.randnuses the global RNG state. If this test fails, the inputs cannot be reproduced. Calltorch.manual_seedat the start of the test.♻️ Proposed change
+ torch.manual_seed(0) a = torch.randn(M, K, dtype=torch.float16, device="mps")🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@testing/python/metal/test_metal_address_space.py` around lines 84 - 92, Seed the test’s random inputs by calling torch.manual_seed at the beginning of the test before the torch.randn calls that create a and b, using the project’s established deterministic seeding convention if available.src/metal/codegen/codegen_metal.cc (1)
484-496: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueReport the enclosing function when a pointer type check fails.
The three checks print only
op->var. A user who hits this error gets a bare variable name with no kernel name and no statement context. Add the value expression to the message so the failing bind is identifiable.♻️ Proposed message improvement
const auto *pointer_type = op->var->type_annotation.as<PointerTypeNode>(); TVM_FFI_ICHECK(pointer_type) - << "Metal handle binding requires a typed pointer: " << op->var; + << "Metal handle binding requires a typed pointer: " << op->var + << " = " << op->value; const auto *element_type = pointer_type->element_type.as<PrimTypeNode>(); TVM_FFI_ICHECK(element_type) - << "Metal handle binding requires a primitive pointee type: " << op->var; + << "Metal handle binding requires a primitive pointee type: " << op->var + << " = " << op->value;🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/metal/codegen/codegen_metal.cc` around lines 484 - 496, Update the three validation messages in the enclosing handle-binding function to include the value expression alongside op->var, so failures identify the kernel statement context. Preserve the existing checks and message distinctions for typed pointers, primitive pointee types, and explicit storage scopes.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@src/metal/codegen/codegen_metal.cc`:
- Around line 946-953: Update GetStorageScopeOf to recursively unwrap CastNode
expressions before inferring the pointer storage scope, while preserving
existing handling for other expression types. In GetAddrSpaceOf, include the
failing kernel context in the check failure message alongside the pointer
expression, and retain the existing storage-scope mapping for successfully
inferred scopes.
- Around line 988-990: Update the builtin::reinterpret() branch in
GetAddrSpaceOf and GetPointeeTypeOf to derive the Metal address space and
pointee type from the reinterpret call’s result dtype rather than recursively
using call->args[0]. Preserve the existing single-argument validation and ensure
cooperative_tensor_load/store receive the target dtype through these helpers.
In `@testing/python/metal/test_metal_address_space.py`:
- Around line 70-72: Update the per-alias assertions in the test loop to stop
requiring a specific vector type or offset syntax, and replace the ineffective
negative alias-cast check. Validate that every half-pointer cast in
compiled_source includes the expected address-space qualifier, using the
existing aliases loop and source inspection while remaining independent of
vectorization width and indexing form.
- Around line 59-62: The
test_both_metal_build_paths_use_type_driven_shared_aliases test must explicitly
apply the existing Metal build-availability gate before calling lower_to_metal.
Add the repository’s established skip mechanism so CUDA-only or non-Metal builds
skip the test, while preserving the current fallback environment setup and both
build-path invocations.
---
Nitpick comments:
In `@src/metal/codegen/codegen_metal.cc`:
- Around line 484-496: Update the three validation messages in the enclosing
handle-binding function to include the value expression alongside op->var, so
failures identify the kernel statement context. Preserve the existing checks and
message distinctions for typed pointers, primitive pointee types, and explicit
storage scopes.
In `@testing/python/metal/test_metal_address_space.py`:
- Around line 46-56: Update lower_to_metal so it asserts artifact.kernel_source
is present before returning it, rather than converting a missing value to an
empty string. Preserve the existing return type and ensure the assertion clearly
identifies missing kernel source; leave the downstream equality and alias
assertions unchanged.
- Around line 84-92: Seed the test’s random inputs by calling torch.manual_seed
at the beginning of the test before the torch.randn calls that create a and b,
using the project’s established deterministic seeding convention if available.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro Plus
Run ID: 7e925cfc-672c-43fa-b7bb-f91431455571
📒 Files selected for processing (3)
src/metal/codegen/codegen_metal.ccsrc/metal/codegen/codegen_metal.htesting/python/metal/test_metal_address_space.py
LeiWang1999
approved these changes
Aug 14, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Problem
Metal requires pointer declarations and pointer casts to carry an explicit address-space qualifier.
Shared-memory aliases created through
BindNodeandhandle_add_byte_offsetcould reach the MSL compiler as unqualified pointers. The current build path then tried to repair the generated source with five string-rewrite passes, including rules tied to names such asA_shared,B_shared, and_shared.Equivalent kernels could therefore compile or fail depending on buffer names. A failing kernel is rejected before GPU execution with
pointer type must have explicit address space qualifier.Root cause
The shared-memory merge pass preserves the pointee type and storage scope on the alias variable.
CodeGenTileLangMetalpreviously inherited generic C handling for pointer-valuedBindNodestatements, so that type information was not used when emitting the alias declaration.handle_add_byte_offsetwas hard-coded tothreadgroupinstead of deriving its address space from the source pointer.Change
Emit pointer-valued
BindNodedeclarations from theirPointerType, including pointee dtype and storage scope.Map TIR storage scopes to Metal
device,threadgroup, andthreadaddress spaces in one place.Trace pointer storage scope and pointee type through variables,
address_of,handle_add_byte_offset,tvm_access_ptr, andreinterpret.Derive byte-offset casts from the source pointer and reject unresolved storage scopes.
Remove the five generated-source string rewrites. Both TileLang Metal build paths now consume the same type-correct codegen output.
Upstream
The corresponding fix was accepted and merged as Apache TVM #20101.
This PR applies the same type-driven correction to TileLang own Metal codegen. It does not modify
3rdparty/tvm.Validation
Hardware validation used a MacBook Air
Mac14,2with Apple M2, macOS 15.6.124G90, Metal 3, Command Line Tools, and macOS SDK 15.5.A full local
tilelangbuild completed successfully from the currentmainbase.The focused regression passed through source-only codegen, compiled codegen, TVM FFI execution, and Torch execution on the M2 GPU:
The three skips are the existing Metal 4 cooperative-tensor tests. They are not supported by the current M2 capability detector.
Repository checks passed:
Scope
The fix is independent of pointee dtype. The regression executes FP16 GEMM aliases, while address-space derivation is based on pointer storage scope.
The change is limited to TileLang Metal code generation and its regression coverage.
It does not change compilation pipelines, runtime dispatch, public APIs, serialized formats, or the
3rdparty/tvmrevision.Summary
Validation
C++ style / lint notes
CodeGenTileLangMetaldeclaration.docs/developer_guide/cpp_style.md.