mirror of
https://github.com/tile-ai/tilelang.git
synced 2026-10-22 06:50:25 +08:00
Yichen YanandLeiWang1999
26c9c6e9cf
* [Metal] Add Metal GEMM support with simdgroup_matrix MMA Add T.gemm support for Apple Metal using simdgroup_matrix 8x8 operations (simdgroup_load/store/multiply_accumulate). Works on all Apple Silicon (M1-M5) without requiring a TVM fork. Key changes: - codegen_metal.cc/h: Fork TVM Metal codegen to tilelang with simdgroup intrinsic emission and 128-bit vectorized copy - gemm_metal.py: GemmMetal tile operator for sharedxshared GEMM - metal_macro_generator.py: MPSIntrinEmitter for simdgroup MMA macros - metal_fragment_to_simdgroup.py: Pass rewrites local.fragment GEMM accumulators to metal.simdgroup scope before layout inference - LowerSIMDGroupCopy in copy.cc for fragment->device simdgroup_store 24 Metal tests (codegen cross-platform + correctness on device). * Move Metal buffer helpers to Metal backend * fix(metal): remove layout_inference bypass, fix build & clarify docs - Remove duplicate BuildTileLangMetal in rt_mod_metal.cc (fixes linker error) - Remove TargetIsMetal bypass in layout_inference.cc. The bypass is unnecessary because MetalFragmentToSimdgroup already converts all GEMM accumulator fragment buffers to metal.simdgroup scope before LayoutInference runs, so IsFragmentBuffer never matches them. - Clarify docstrings for MetalFragmentToSimdgroup pass explaining why it runs before LayoutInference (simdgroup matrices are opaque and have no explicit thread-level layout). - Restore is_metal_simdgroup check in GemmMetal.lower() so the GEMM lowering correctly handles buffers already converted to simdgroup scope by the pass. * fix: tir->tirx migration build fixes for metal backend - Replace tir::Buffer with Buffer, tir:: with tirx:: in metal backend - Update codegen_metal to use AllocBufferNode instead of AllocateNode - Create runtime/metal/metal_module.h with MetalModuleCreate factory - Fix ICHECK -> TVM_FFI_ICHECK macros - Fix 'from tvm import tir' -> 'from tvm import tirx as tir' in Python - Restore tilelang/backend/__init__.py for metal gemm registration - Add metal backend import to tilelang/__init__.py - Fix metal_fragment_to_simdgroup.py: Block->SBlock, Allocate->AllocBuffer * fix: improve metal_fragment_to_simdgroup buffer remapping - Add AllocBuffer node handling for standalone allocations - Remap buffer references in SBlock reads/writes/match_buffers - Add _remap_buffer_region and _remap_match_buffer helpers * lint * fix: add full-body var subst in MetalFragmentToSimdgroup to reach gemm call_intrin args The per-SBlock substitute() in _rewrite_scope missed Var references buried inside BufferRegion objects nested in call_intrin arguments (e.g., the GEMM op's C buffer). This leaked the old local.fragment-typed C_local var into the subsequent gemm lowering, producing a duplicate variable that MakePackedAPI flagged as undefined. Fix: apply tir.stmt_functor.substitute() to the entire body before the ir_transform pass, so all Var references — including those inside opaque Call args — are remapped to the metal.simdgroup-typed replacements. Also remove the redundant per-SBlock substitute() and the duplicate _remap_buffer function. * fix: add missing <cmath> include for std::isinf/isnan in codegen_metal * [Metal] Fix review findings: dedup warp partition, kernel_only, barrier comment, immutable gemm ops - Extract shared ComputeSquareWarpPartition to src/backend/metal/op/utils.h, used by both gemm.cc (Square policy) and copy.cc (simdgroup store lowering). - Implement kernel_only parameter in MetalKernelAdapter.get_kernel_source. - Add comment documenting that Metal's simdgroup_barrier synchronizes the full threadgroup (no per-simdgroup barrier exists). - Replace module-level mutable global in metal_fragment_to_simdgroup.py with @lru_cache and remove CUDA-only gemm ops from the set. * lint * Refactor MetalFragmentToSimdgroup import path in LowerAndLegalize function * Remove metal_fragment_to_simdgroup.py file and update import in metal/__init__.py --------- Co-authored-by: LeiWang1999 <leiwang1999@outlook.com>