This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new 28f4bc47c9 [BugFix][Metal] Preserve pointer address spaces for byte
offsets (#20101)
28f4bc47c9 is described below
commit 28f4bc47c9152cd4542deca4d95ba0b1a46706ec
Author: Gengyuan Bai <[email protected]>
AuthorDate: Sun Aug 9 04:13:02 2026 -0400
[BugFix][Metal] Preserve pointer address spaces for byte offsets (#20101)
## Problem
1. A Metal kernel that derives typed and untyped aliases from shared
memory can reach the MSL compiler with pointer declarations and
byte-offset casts that have no address-space qualifier.
2. Metal rejects those generic pointers before kernel execution with a
diagnostic such as:
```text
error: pointer type must have explicit address space qualifier
```
3. Before this change, the generic C codegen path could emit the
equivalent of:
```metal
half* typed_alias = (half*)((half*)(((char*)shared) + 4));
void* void_alias = (void*)((char*)shared + 8);
```
The source allocation is `threadgroup`, but the alias declarations,
result casts, and intermediate `char*` casts do not carry that address
space.
## Root cause
1. The TIR result types of `tirx.ptr_byte_offset` and
`tirx.handle_add_byte_offset` retain the pointer storage scope.
2. `CodeGenMetal` previously inherited the generic `CodeGenC` handling
for both the pointer-valued `Bind` nodes and these byte-offset
operations. That path emits ordinary C pointer casts and does not
express Metal address spaces.
3. The loss occurs at the TIR-to-MSL codegen boundary, before Metal
runtime compilation or GPU execution.
## What changed
1. Handle pointer-valued `Bind` nodes in `CodeGenMetal` and emit the
storage scope recorded in their `PointerType`.
2. Lower typed and untyped byte-offset operations with the same address
space on both the result pointer and the intermediate `char*` cast.
3. Keep the existing generic C behavior for pointers without an explicit
storage scope.
4. The corrected source now has the equivalent form:
```metal
threadgroup half* typed_alias =
(threadgroup half*)((threadgroup half*)(((threadgroup char*)shared) +
4));
threadgroup void* void_alias =
(threadgroup void*)((threadgroup void*)(((threadgroup char*)shared) +
8));
```
## Validation
1. Added a source-generation regression covering both typed and untyped
aliases derived from a `float16` shared allocation.
2. Added an end-to-end Metal test using the standard compilation
pipeline. The generated source contains the expected `threadgroup`
allocation, typed alias, void alias, and byte-offset casts.
3. Ran the end-to-end test on a MacBook Air (`Mac14,2`) with an Apple M2
GPU, macOS 15.6.1 (`24G90`), Metal 3, Command Line Tools, and the macOS
15.5 SDK. TVM compiled MSL 2.3 through the Metal runtime without the
offline `metal` or `metallib` tools. The kernel returned `[10.0, 11.0,
12.0]`, matching the expected output.
4. The focused tests pass:
```text
python -m pytest -q \
tests/python/codegen/test_target_codegen_metal.py::test_codegen_pointer_byte_offsets_preserve_storage_scope
\
tests/python/codegen/test_target_codegen_metal.py::test_pointer_byte_offsets_execute_in_threadgroup_memory
# 2 passed
```
5. Changed-file pre-commit checks pass:
```text
pre-commit run --files \
src/backend/metal/codegen/codegen_metal.cc \
src/backend/metal/codegen/codegen_metal.h \
tests/python/codegen/test_target_codegen_metal.py
```
## Scope
1. This change only affects Metal pointer code generation when the TIR
pointer type carries an explicit storage scope.
2. The issue is independent of the pointee dtype: the source-generation
regression uses `float16`, while the end-to-end Metal test exercises the
same path with `float32`.
3. It does not change compilation pipelines, Metal runtime
implementation, public APIs, or serialized formats.
4. Hardware validation is limited to the Apple M2 configuration above.
---
src/backend/metal/codegen/codegen_metal.cc | 46 ++++++++++++++
src/backend/metal/codegen/codegen_metal.h | 1 +
tests/python/codegen/test_target_codegen_metal.py | 73 +++++++++++++++++++++++
3 files changed, 120 insertions(+)
diff --git a/src/backend/metal/codegen/codegen_metal.cc
b/src/backend/metal/codegen/codegen_metal.cc
index 78a0d7495b..361a9ee3d2 100644
--- a/src/backend/metal/codegen/codegen_metal.cc
+++ b/src/backend/metal/codegen/codegen_metal.cc
@@ -327,6 +327,31 @@ void CodeGenMetal::PrintStorageScope(const std::string&
scope, std::ostream& os)
}
}
+void CodeGenMetal::VisitStmt_(const BindNode* op) {
+ const auto* pointer_type = op->var->ty.as<PointerTypeNode>();
+ if (pointer_type == nullptr || pointer_type->storage_scope.empty()) {
+ return CodeGenC::VisitStmt_(op);
+ }
+
+ const std::string& storage_scope = pointer_type->storage_scope;
+ alloc_storage_scope_[op->var.get()] = storage_scope;
+ RegisterHandleTypeFromPointer(op->var, &op->value);
+ std::string value = PrintExpr(op->value);
+ if (print_ssa_form_) {
+ TVM_FFI_ICHECK(!var_idmap_.count(op->var.get()));
+ var_idmap_[op->var.get()] = value;
+ return;
+ }
+
+ PrintIndent();
+ PrintStorageScope(storage_scope, stream);
+ PrintType(pointer_type->element_type, stream);
+ stream << "* " << AllocVarID(op->var.get()) << " = (";
+ PrintStorageScope(storage_scope, stream);
+ PrintType(pointer_type->element_type, stream);
+ stream << "*)" << value << ";\n";
+}
+
void CodeGenMetal::VisitStmt_(const AllocBufferNode* op) {
TVM_FFI_ICHECK(op->buffer.defined());
std::string vid = AllocVarID(op->buffer.get());
@@ -445,6 +470,27 @@ void CodeGenMetal::VisitExpr_(const CallNode* op,
std::ostream& os) { // NOLINT
<< PrintExpr(a) << "[" << PrintExpr(op->args[3]) << "], " //
<< PrintExpr(b) << "[" << PrintExpr(op->args[5]) << "], " //
<< PrintExpr(c) << "[" << PrintExpr(op->args[7]) << "])";
+ } else if (op->op.same_as(builtin::ptr_byte_offset()) ||
+ op->op.same_as(builtin::handle_add_byte_offset())) {
+ bool is_typed_offset = op->op.same_as(builtin::ptr_byte_offset());
+ TVM_FFI_ICHECK_EQ(op->args.size(), is_typed_offset ? 3U : 2U);
+ const auto* pointer_type = op->ty.as<PointerTypeNode>();
+ TVM_FFI_ICHECK(pointer_type)
+ << "Metal pointer byte offsets must have a pointer result type, but
got " << op->ty;
+ if (pointer_type->storage_scope.empty()) {
+ return CodeGenC::VisitExpr_(op, os);
+ }
+
+ os << "((";
+ PrintStorageScope(pointer_type->storage_scope, os);
+ PrintType(pointer_type->element_type, os);
+ os << "*)(((";
+ PrintStorageScope(pointer_type->storage_scope, os);
+ os << "char*)";
+ PrintExpr(op->args[0], os);
+ os << ") + ";
+ PrintExpr(op->args[1], os);
+ os << "))";
} else if (op->op.same_as(builtin::reinterpret())) {
if (!op->ty.as<PrimTypeNode>() || !op->args[0]->ty.as<PrimTypeNode>()) {
return CodeGenC::VisitExpr_(op, os);
diff --git a/src/backend/metal/codegen/codegen_metal.h
b/src/backend/metal/codegen/codegen_metal.h
index b54852ffa1..2ad84578ac 100644
--- a/src/backend/metal/codegen/codegen_metal.h
+++ b/src/backend/metal/codegen/codegen_metal.h
@@ -52,6 +52,7 @@ class CodeGenMetal final : public CodeGenC {
void PrintVecElemStore(const std::string& vec, const PrimType& t, int i,
const std::string& value) final;
// overload visitor
+ void VisitStmt_(const BindNode* op) final; //
NOLINT(*)
void VisitStmt_(const AllocBufferNode* op) final; //
NOLINT(*)
void VisitExpr_(const SelectNode* op, std::ostream& os) final; //
NOLINT(*)
void VisitExpr_(const BroadcastNode* op, std::ostream& os) final; //
NOLINT(*)
diff --git a/tests/python/codegen/test_target_codegen_metal.py
b/tests/python/codegen/test_target_codegen_metal.py
index a1fb713656..f0d9998b4c 100644
--- a/tests/python/codegen/test_target_codegen_metal.py
+++ b/tests/python/codegen/test_target_codegen_metal.py
@@ -384,5 +384,78 @@ def test_codegen_simdgroup_buffer_data():
assert "simdgroup_multiply_accumulate(" in source
+def test_codegen_pointer_byte_offsets_preserve_storage_scope():
+ """Pointer byte offsets should preserve the source Metal address space."""
+
+ @I.ir_module(s_tir=True)
+ class Module:
+ @T.prim_func(s_tir=True)
+ def kernel():
+ T.func_attr(
+ {
+ "calling_conv": 2,
+ "global_symbol": "kernel",
+ "tirx.kernel_launch_params": [],
+ }
+ )
+ shared = T.alloc_buffer((16,), "float16", scope="shared")
+ typed_alias = T.ptr_byte_offset(shared.data, 4, "float16")
+ typed_buffer = T.decl_buffer((14,), "float16", data=typed_alias,
scope="shared")
+ void_alias = T.handle_add_byte_offset(shared.data, 8)
+ void_buffer = T.decl_buffer((12,), "float16", data=void_alias,
scope="shared")
+ typed_buffer[0] = T.float16(1)
+ void_buffer[0] = T.float16(2)
+
+ metal_codegen = tvm.get_global_func("target.build.metal")
+ module = metal_codegen(Module, tvm.target.Target("metal"))
+ source = module.inspect_source()
+
+ assert "threadgroup half* typed_alias" in source
+ assert "threadgroup void* void_alias" in source
+ assert source.count("threadgroup char*") == 2
+
+
[email protected]
[email protected](not env.has_metal(), reason="need metal")
+def test_pointer_byte_offsets_execute_in_threadgroup_memory():
+ """Pointer byte offsets should execute in Metal threadgroup memory."""
+
+ @I.ir_module(s_tir=True)
+ class Module:
+ @T.prim_func(s_tir=True)
+ def main(A: T.Buffer((16,), "float32"), B: T.Buffer((16,), "float32")):
+ for bx in T.thread_binding(1, thread="blockIdx.x"):
+ for tx in T.thread_binding(1, thread="threadIdx.x"):
+ shared = T.alloc_buffer((16,), "float32", scope="shared")
+ typed_alias = T.ptr_byte_offset(shared.data, 4, "float32")
+ typed_buffer = T.decl_buffer((15,), "float32",
data=typed_alias, scope="shared")
+ void_alias = T.handle_add_byte_offset(shared.data, 8)
+ void_buffer = T.decl_buffer((14,), "float32",
data=void_alias, scope="shared")
+ shared[0] = A[0]
+ typed_buffer[0] = A[1]
+ void_buffer[0] = A[2]
+ B[0] = shared[0]
+ B[1] = typed_buffer[0]
+ B[2] = void_buffer[0]
+
+ executable = tvm.compile(Module, target="metal")
+ source = executable.mod.imports[0].inspect_source()
+
+ assert "threadgroup float shared[16]" in source
+ assert "threadgroup float* typed_alias" in source
+ assert "threadgroup void* void_alias" in source
+ assert source.count("threadgroup char*") == 2
+
+ def run_and_check():
+ dev = tvm.metal(0)
+ host_input = np.arange(16, dtype="float32") + 10
+ input_tensor = tvm.runtime.tensor(host_input, dev)
+ output_tensor = tvm.runtime.tensor(np.zeros(16, dtype="float32"), dev)
+ executable(input_tensor, output_tensor)
+ tvm.testing.assert_allclose(output_tensor.numpy()[:3], host_input[:3])
+
+ tvm.testing.run_with_gpu_lock(run_and_check)
+
+
if __name__ == "__main__":
tvm.testing.main()
