Skip to content

[BugFix][Metal] Preserve pointer address spaces for byte offsets - #20101

Open
GY-Bai wants to merge 1 commit into
apache:mainfrom
GY-Bai:fix/metal-pointer-address-spaces
Open

[BugFix][Metal] Preserve pointer address spaces for byte offsets#20101
GY-Bai wants to merge 1 commit into
apache:mainfrom
GY-Bai:fix/metal-pointer-address-spaces

Conversation

@GY-Bai

@GY-Bai GY-Bai commented Aug 6, 2026

Copy link
Copy Markdown

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:

error: pointer type must have explicit address space qualifier
  1. Before this change, the generic C codegen path could emit the equivalent of:
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:

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:

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
  1. Changed-file pre-commit checks pass:
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.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants