[BugFix][Metal] Preserve pointer address spaces for byte offsets

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);

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(*)

tests/python/codegen/test_target_codegen_metal.py

@@ -384,5 +384,78 @@ def kernel():
     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
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(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()