Skip to content

[Metal] Resolve target attributes from device properties - #71

Open
anerli wants to merge 1 commit into
tile-ai:tilelang_mainfrom
magnitudedev:pr/metal-target-properties
Open

anerli wants to merge 1 commit into
tile-ai:tilelang_mainfrom
magnitudedev:pr/metal-target-properties

Conversation

@anerli

@anerli anerli commented Sep 12, 2026

Copy link
Copy Markdown

Problem

Metal targets are assembled from constants. Target.from_device("metal") hard-codes 32 KB of threadgroup memory because GetAttr(kMaxSharedMemoryPerBlock) is a stub, and nothing reports whether a device supports bfloat16, SIMD-group permute/reduction/matrix operations, or the Metal 4 family. Codegen consumers such as TileLang have had to guess from xcrun/sysctl output. Separately, MetalModuleNode chooses the MSL language version from the device at compile time, so a module can be compiled with a different language version than the one codegen generated for.

Change

  • Add MetalDeviceProperties, queried once per device when MetalWorkspace initializes, and expose it through device_api.metal.get_target_property, mirroring device_api.vulkan.get_target_property.
  • Register the matching attributes on the metal target kind: metal_language_version (major * 10 + minor, default 23), supports_bfloat16, supports_simdgroup_permute, supports_simdgroup_reduction, supports_simdgroup_matrix, supports_metal4. Existing attribute defaults are unchanged.
  • Serve kMaxSharedMemoryPerBlock from maxThreadgroupMemoryLength, and make Target.from_device("metal") read it and the new properties.
  • Carry metal_language_version on the Metal module. The runtime compiles fmt="metal" sources with exactly that version and rejects a module that needs a newer MSL than the device provides. ffi.Module.create.metal gains a trailing int, and the serialized module format gains one field after fmt; the codegen-side fallback module mirrors both.

Tests

  • tests/python/codegen/test_target_codegen_metal.py
    • test_target_from_device_reports_device_properties: detection matches Device attributes and exposes every new property.
    • test_module_compiles_with_target_language_version: the version survives export and load, and a module generated for a newer MSL than the device compiles is rejected.
  • The touched translation units pass clang -fsyntax-only with the project's build flags. The Python tests still need a run against a rebuilt library before this leaves draft.

Follow-up in TileLang

Once this merges and TileLang bumps 3rdparty/tvm, a TileLang PR titled "[Metal] Resolve Metal targets from device properties" (branch pr/metal-target-resolution) will:

  • replace the xcrun/sysctl detection in tilelang/metal/target.py with Target.from_device and fill device facts into user-supplied metal targets,
  • gate the Metal backend capabilities (bfloat16 dtype, SIMD-group matrix instructions, subgroup_exchange) on these attributes,
  • check supports_simdgroup_matrix and supports_bfloat16 in the Metal GEMM lowering instead of assuming them,
  • read supports_metal4 from the attribute instead of the metal4 target key,
  • forward metal_language_version from src/metal/codegen/codegen_metal.cc into the module.

Query each Metal device once for the facts that decide what codegen may
emit: the highest Metal Shading Language version it compiles, bfloat16,
SIMD-group permute/reduction/matrix support, and the Metal 4 family.
Expose them through "device_api.metal.get_target_property", mirroring the
Vulkan device API, and register the matching attributes on the "metal"
target kind so Target.from_device("metal") reports them. Serve the
threadgroup memory limit through GetAttr instead of hard-coding 32 KB in
target detection.

Carry the MSL version on the Metal module so runtime compilation uses the
version codegen generated for, instead of choosing one from the device at
compile time. The module byte format gains one field after "fmt".
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.

1 participant