Skip to content

Metal: legalize the simdgroup matrix and tensor intrinsics - #1046

Closed
maleadt wants to merge 1 commit into
mainfrom
tb/metal-intrinsic-abi
Closed

maleadt wants to merge 1 commit into
mainfrom
tb/metal-intrinsic-abi

Conversation

@maleadt

@maleadt maleadt commented Oct 21, 2026 •

Copy link
Copy Markdown
Member

Metal.jl emits the simdgroup matrix load/store and tensor intrinsics in their newest AIR form, and rewrote them for older AIR versions in its own finish_ir! override. That rewrite depends only on the targeted AIR version, like the barrier and atomic legalization this package already does, and is the kind of thing an LLVM back-end handles for the other GPU targets. So this moves it into lower_air!, next to legalize_barrier_abi!:

  • air.simdgroup_matrix_8x8_{load,store}: before AIR 2.8, the dims/strides/origin vectors become the legacy elements-per-row, origin and transpose operands. The transpose flag has to be an immediate, so one of the strides must be the literal 1 (as Metal.jl emits them); other layouts are an error rather than an assertion.
  • air.init_strided_private_tensor.i32.{global,local} and air.get_extent_private_tensor.i32: before AIR 2.9, the operands AIR 2.9 appended (always null, 0, -1 and 0 in Apple's front-end) are dropped.
  • The element type of the simdgroup matrix pointer operand is derived from the intrinsic name for the downgrader, as is done for atomics.

Calls already in the older form are left alone. Released versions of Metal.jl still rewrite them in finish_ir!, which runs before lower_air!, so they keep working with this version. Older forms aren't upgraded for newer targets.

Tests cover both families at the AIR versions on either side of the change, with constant and run-time operands, and that the legacy forms pass through. Validated on Julia 1.10 and 1.13. Metal.jl's simdgroup, linalg and bfloat16 tests pass with JuliaGPU/Metal.jl#1096, which removes its copy of the rewrites, including execution at macos=v"15" (AIR 2.7) on macOS 27.

Metal.jl emits these in their newest form and rewrote them for older AIR
versions in its own `finish_ir!`. Like the barrier and atomic ABIs, that is
a property of the AIR version being targeted, so do it here during AIR
lowering: before AIR 2.8, the simdgroup matrix load/store calls get the
legacy elements-per-row, origin and transpose operands, and before AIR 2.9
the operands AIR 2.9 appended to the tensor intrinsics are dropped. Calls
already in the older form, as emitted by released versions of Metal.jl,
are left alone. The element type of the simdgroup matrix pointer operand
is now also derived here, for the downgrader.
@maleadt

maleadt commented Oct 21, 2026

Copy link
Copy Markdown
Member Author

Closing: the AIR legalization of intrinsics Metal.jl emits stays next to the emitter in Metal.jl, so the two packages don't have to release in lockstep, and GPUCompiler doesn't need to recognize every form an older Metal.jl may emit.

@maleadt maleadt closed this Oct 21, 2026
@maleadt
maleadt deleted the tb/metal-intrinsic-abi branch October 21, 2026 05:17
@codecov

codecov Bot commented Oct 21, 2026

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 98.68421% with 1 line in your changes missing coverage. Please review.
✅ Project coverage is 91.51%. Comparing base (ba243ce) to head (c60a69d).

Files with missing lines Patch % Lines
src/metal.jl 98.68% 1 Missing ⚠️
Additional details and impacted files
@@            Coverage Diff             @@
##             main    #1046      +/-   ##
==========================================
- Coverage   92.34%   91.51%   -0.83%     
==========================================
  Files          31       31              
  Lines        7341     7413      +72     
==========================================
+ Hits         6779     6784       +5     
- Misses        562      629      +67     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

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