[Metal] Derive pointer address spaces from TIR types - #2824
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! 🚀 |
|
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 (1)
🚧 Files skipped from review as they are similar to previous changes (1)
📝 WalkthroughWalkthroughMetal codegen now uses typed, storage-scope-aware pointer inference for declarations and byte-offset casts. It supports additional Metal scopes, removes string-based source rewriting, and adds regression tests for generated source and MPS execution. ChangesMetal address-space code generation
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant BindNode
participant MetalCodeGen
participant GeneratedMetal
participant MPS
BindNode->>MetalCodeGen: provide typed pointer binding
MetalCodeGen->>MetalCodeGen: infer storage scope and pointee type
MetalCodeGen->>GeneratedMetal: emit qualified declarations and casts
GeneratedMetal->>MPS: execute generated GEMM kernel
MPS-->>GeneratedMetal: return computed result
Possibly related PRs
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}
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (1)
testing/python/metal/test_metal_address_space.py (1)
63-68: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueMake the assertions independent of the buffer name hints.
The assertions require the identifiers
AsandBsin the generated Metal source. Those identifiers come from TIR var name hints throughAllocVarID, so any renaming pass breaks this test for a reason unrelated to address spaces. Match the qualifier and the pointee type with a regular expression instead.♻️ Proposed name-independent assertions
- assert "threadgroup half* As = (threadgroup half*)" in compiled_source - assert "threadgroup half* Bs = (threadgroup half*)" in compiled_source - assert "*(threadgroup half2*)(As +" in compiled_source - assert "*(threadgroup half2*)(Bs +" in compiled_source - assert "(half*)As" not in compiled_source - assert "(half*)Bs" not in compiled_source + decls = re.findall(r"threadgroup half\* (\w+) = \(threadgroup half\*\)", compiled_source) + assert len(decls) >= 2, compiled_source + for name in decls: + assert f"*(threadgroup half2*)({name} +" in compiled_source + assert f"(half*){name}" not in compiled_sourceAdd
import reat the top of the file.🤖 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 63 - 68, Update the assertions in the Metal address-space test to avoid depending on the generated buffer names As and Bs. Add the re module import and use regular-expression matching to verify the threadgroup qualifier and half/half2 pointee types, while retaining the checks that reject invalid half-pointer casts.
🤖 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 926-947: Update GetAddrSpaceOf and its pointer-expression handling
to unwrap opaque wrappers such as builtin::tvm_access_ptr and
builtin::reinterpret before resolving the address space. Reuse the underlying
buffer/data expression’s storage scope where appropriate, while preserving
existing handling for Var, address_of(BufferLoad), and handle_add_byte_offset;
avoid allowing unsupported wrapped operands to reach the unconditional
GetStorageScopeOf failure.
In `@testing/python/metal/test_metal_address_space.py`:
- Around line 57-59: Update test_shared_alias_address_spaces_are_type_driven to
call lower_to_metal with enable_device_compile=False, since this source-only
test should use the Metal without-compile lowering path. Remove the ineffective
TVM_COMPILE_FORCE_FALLBACK environment setup unless it is needed elsewhere.
---
Nitpick comments:
In `@testing/python/metal/test_metal_address_space.py`:
- Around line 63-68: Update the assertions in the Metal address-space test to
avoid depending on the generated buffer names As and Bs. Add the re module
import and use regular-expression matching to verify the threadgroup qualifier
and half/half2 pointee types, while retaining the checks that reject invalid
half-pointer casts.
🪄 Autofix (Beta)
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: e99d3463-5a81-4302-921c-5ff5b8e2a6fe
📒 Files selected for processing (3)
src/metal/codegen/codegen_metal.ccsrc/metal/codegen/codegen_metal.htesting/python/metal/test_metal_address_space.py
|
Closing this in favor of a smaller TileLang-side implementation aligned with Apache TVM #20101, which has now been merged. I’ll send the updated version in a new PR. |
1 participant
Problem
Metal requires every pointer type to carry an explicit address-space qualifier. On an Apple M2, a regular FP16 GEMM using dynamic shared-memory aliases named
AsandBsfailed during MSL compilation with errors such as:The existing source-rewrite passes handled a few hard-coded names and then scanned lines containing
_shared. Therefore, correctness depended on variable names:A_sharedcould compile while an equivalent alias namedAsdid not. The compiled and source-only Metal build entry points also exercised different rewrite logic.This failure occurs in the Metal 3 simdgroup path and is independent of Metal 4 cooperative tensors. The invalid pointer types are rejected during MSL compilation, before GPU execution, so the issue can affect any Metal device that reaches this lowering path; the M2 is the hardware used to reproduce and validate it.
Root-cause trace
This was reproduced and traced from a clean
mainat commite01c498b6499367e0a1cd64413e008d223a931b6.T.alloc_sharedbuffers are merged byMergeSharedMemoryAllocations::MakeAliasBindings()insrc/transform/merge_shared_memory_allocations.cc.tirx::Bind(buffer_var, handle_add_byte_offset(merged_buffer, offset)). The alias variable still has a typedPointerType, including its element dtype andshared.dynstorage scope.CodeGenC::VisitStmt_(const BindNode*)did not register that type information inhandle_data_type_oralloc_storage_scope_.BuildTileLangMetalattempted to repair the generated MSL, but some depended on names such asA_shared,B_shared, or the_sharedsuffix.The loss therefore happened at the TIR-to-MSL
BindNodecodegen boundary, not in the shared-memory merge pass and not in Metal runtime compilation.Fix
BindNodevisitor that reads the aliasPointerType, registers its pointee dtype and storage scope, and emits a typed, address-space-qualified declaration.global->deviceshared,shared.dyn,threadgroup->threadgrouplocal,local.var, and Metal register scopes ->threadhandle_add_byte_offsetaddress spaces from the pointer operand instead of hard-codingthreadgroup.address_of(BufferLoad)and the base operand ofhandle_add_byte_offset; reject unknown scopes instead of silently defaulting tothread.CodeGenTileLangMetal, while their existing callback behavior remains unchanged.The failing aliases now generate:
Tests
Run on a MacBook Air with Apple M2, macOS 15.6.1, and Metal 3:
The skip condition above reflects TileLang's current capability detector, not Apple's published hardware boundary. Apple's May 2026 Metal Feature Set Tables list the Metal 4 programming model and tensors starting with Apple7 (the M1 generation). Metal 4 capability detection is separate from this Metal 3 codegen fix and is intentionally unchanged in this PR.
Potentially affected hardware
This defect is in TIR-to-MSL source generation rather than GPU execution. Within TileLang's current macOS/arm64 deployment scope, any Apple silicon Mac generation from M1 through M5 can potentially hit it when a kernel uses merged dynamic shared-memory aliases and reaches the affected pointer-cast path. This does not mean every Metal kernel fails: the previous name-based rewrites masked many aliases containing
_shared.Only the M2/Metal 3 configuration has been validated on hardware in this PR. The M1, M3, M4, and M5 impact is inferred from the generation-independent MSL address-space rule and the shared TileLang codegen path, not from direct testing on each chip. iOS and iPadOS targets are not claimed here because TileLang's current availability detector only enables the backend on arm64 macOS.
Summary
BindNodehandling for typed pointer bindings.threadgroup,metal.simdgroup, andmetal.cooperative_tensorscopes.handle_add_byte_offset.C++ style / lint notes
docs/developer_guide/cpp_style.md.