From 8fa5e551c2f535a55311a606237827be264d45eb Mon Sep 17 00:00:00 2001 From: Greg Stoner Date: Thu, 7 May 2026 22:55:33 -0500 Subject: [PATCH 1/2] =?UTF-8?q?Phase=208.4.4=20=E2=80=94=20fp16=20/=20bf16?= =?UTF-8?q?=20matmul=20on=20apple=5Fgpu=20(mirror=20of=20CPU=20BNNS=20bf16?= =?UTF-8?q?)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Extends the apple_gpu matmul runtime path with fp16 and bf16 dtype variants. Mirrors the Phase 8.2 BNNS bf16 follow-up on the CPU side: - fp16: native MPSDataTypeFloat16 (Apple Silicon GPUs run fp16 at higher throughput than fp32 on most ops). - bf16: fp32-conversion path inside the runtime shim because MPS does NOT natively support bf16 matrix descriptors as of macOS 14. Same pattern as the CPU bf16 cblas_sgemm fallback. Scope is intentionally narrow — only the matmul kernel gets dtype variants this phase. The other custom MSL kernels (rope, softmax, gelu, flash_attn, matmul_softmax_fusion) remain f32-only; their dtype variants are 8.4.4.x followups. MLIR / runtime - Two new C symbols in apple_gpu_runtime.mm: * tessera_apple_gpu_mps_matmul_f16 — native MPSDataTypeFloat16. ABI is uint16_t* for fp16 bit-pattern transmission (no _Float16 dep). * tessera_apple_gpu_mps_matmul_bf16 — fp32 conversion path. Decodes bf16 bit-pattern via shift, runs MPSDataTypeFloat32 matmul, encodes back with round-to-nearest-even. - apple_gpu_runtime_stub.cpp gets matching reference fallbacks (fp32-via-conversion) for non-Darwin builds. - MatmulToAppleGPU.cpp picks the runtime symbol by input element type: f32 / f16 / bf16. Same i64×3 + i32×3 ABI shape across all three — the element type is encoded in the symbol name only. Python - driver.py: _apple_gpu_matmul_dtype_suffix extracts the dtype from the Graph IR operand types (tensor<*xf16>, tensor<*xbf16>) and routes the backend artifact's runtime symbol selection accordingly. - schedule_ir.py: _base_attrs now surfaces dtype on every schedule op by parsing the IROp's operand_types. The attr propagates through Schedule -> Tile -> Target IR layers so target_ir's mps_matmul emission carries the right dtype attr. - runtime.py: _apple_gpu_dispatch_matmul detects input array dtype at launch time (call-site dtypes are runtime-only since the @jit function signatures are type-polymorphic) and routes to the matching ctypes wrapper. fp16 and bf16 paths both use uint16_t* ABI via numpy's .view(np.uint16); ml_dtypes.bfloat16 is byte-compatible. - New ctypes wrappers _apple_gpu_mps_matmul_f16 / _bf16. Loader gate now requires both new symbols (forces rebuild after Phase 8.4.4). Tests - New lit fixture apple_gpu_matmul_dtypes.mlir — three positive cases (f32, f16, bf16 matmul lower to the right runtime symbol with the shared i64×3 + i32×3 ABI) and one negative case (mixed-dtype operands fall back to the artifact path). - 4 new unit tests in test_apple_backend_roadmap.py: * f32 default artifact contract (compile-time symbol selection) * fp16 end-to-end matches fp32-converted reference at fp16 tolerance * bf16 end-to-end matches fp32-converted reference at bf16 tolerance (gated on ml_dtypes presence, mirrors the CPU bf16 soft-dep) * fp16 + bf16 ABI shim correctness (direct ctypes invocation against a freshly-compiled shim) Verified on Apple Silicon (LLVM/MLIR 21, Metal active): 1994 unit tests passing (1991 + 3 net new fp16/bf16 tests); 12/12 Phase 8 lit fixtures passing against the in-tree tessera-opt. fp16 matmul matches fp32-converted reference at rtol=5e-2 (MPS does fp16 internal accumulation; minor drift from the per-element reference is expected). bf16 matches at rtol=2e-2. Co-Authored-By: Claude Opus 4.7 --- python/tessera/compiler/driver.py | 32 ++- python/tessera/compiler/schedule_ir.py | 26 ++ python/tessera/compiler/target_ir.py | 14 +- python/tessera/runtime.py | 141 ++++++++--- .../Apple/Lowering/MatmulToAppleGPU.cpp | 45 +++- .../runtime/apple_gpu_runtime.mm | 229 ++++++++++++++++++ .../runtime/apple_gpu_runtime_stub.cpp | 110 +++++++++ tests/tessera-ir/.lit_test_times.txt | 23 +- .../Output/apple_cpu_lowering.mlir.script | 2 +- .../Output/apple_cpu_runtime.mlir.script | 2 +- .../apple_dialect_roundtrip.mlir.script | 2 +- .../Output/apple_gpu_flash_attn.mlir.script | 2 +- .../Output/apple_gpu_lowering.mlir.script | 2 +- .../apple_gpu_matmul_dtypes.mlir.script | 1 + ...pple_gpu_matmul_softmax_fusion.mlir.script | 2 +- .../phase8/Output/apple_gpu_msl.mlir.script | 2 +- .../Output/apple_gpu_runtime.mlir.script | 2 +- .../Output/apple_gpu_softmax_gelu.mlir.script | 2 +- .../Output/target_ir_contracts.mlir.script | 2 +- .../Output/tmem_tcgen05_contract.mlir.script | 2 +- .../phase8/apple_gpu_matmul_dtypes.mlir | 49 ++++ tests/unit/test_apple_backend_roadmap.py | 173 +++++++++++++ 22 files changed, 800 insertions(+), 65 deletions(-) create mode 100644 tests/tessera-ir/phase8/Output/apple_gpu_matmul_dtypes.mlir.script create mode 100644 tests/tessera-ir/phase8/apple_gpu_matmul_dtypes.mlir diff --git a/python/tessera/compiler/driver.py b/python/tessera/compiler/driver.py index 89f67e6c7..28d49241b 100644 --- a/python/tessera/compiler/driver.py +++ b/python/tessera/compiler/driver.py @@ -418,6 +418,32 @@ def _is_apple_gpu_mps_executable(cpu_plan: CPUPlan | None) -> bool: return _apple_gpu_chain_kind(cpu_plan) is not None +def _apple_gpu_matmul_dtype_suffix(cpu_plan: CPUPlan | None) -> str: + """Phase 8.4.4 — extract the matmul element type from a CPUPlan and + map it to the runtime symbol's dtype suffix. + + The graph IR operand types are strings like "tensor<*xf32>", "tensor<*xf16>", + or "tensor<*xbf16>". The matching runtime symbol is one of: + - tessera_apple_gpu_mps_matmul_f32 (Phase 8.3, native MPSDataTypeFloat32) + - tessera_apple_gpu_mps_matmul_f16 (Phase 8.4.4, native MPSDataTypeFloat16) + - tessera_apple_gpu_mps_matmul_bf16 (Phase 8.4.4, fp32-conversion path) + + Defaults to f32 when the operand type can't be parsed — matches the pre- + Phase 8.4.4 behavior so existing single-op f32 programs are unchanged. + """ + + if cpu_plan is None or not cpu_plan.ops: + return "f32" + op = cpu_plan.ops[0] + operand_types = list(getattr(op, "operand_types", ()) or ()) + for t in operand_types: + if "bf16" in t: + return "bf16" + if "f16" in t and "bf16" not in t: + return "f16" + return "f32" + + def _apple_gpu_chain_kind(cpu_plan: CPUPlan | None) -> str | None: """Phase 8.4.3: classify multi-op apple_gpu plans against the recognized fusion patterns. Returns the chain kind ("matmul_softmax" today) or None @@ -492,7 +518,11 @@ def _backend_artifact_for(target_kind: str, cpu_plan: CPUPlan | None) -> Lowerin else: only_op = cpu_plan.ops[0].op_name if only_op in _APPLE_GPU_MPS_OPS: - symbol = "tessera_apple_gpu_mps_matmul_f32" + # Phase 8.4.4 — pick the matmul symbol by element type. The + # actual element type comes from the Graph IR operand types + # ("tensor<*xf16>" etc); fall back to f32 if the parse fails. + dtype_suffix = _apple_gpu_matmul_dtype_suffix(cpu_plan) + symbol = f"tessera_apple_gpu_mps_matmul_{dtype_suffix}" framework = "MetalPerformanceShaders" abi = "MPSMatrixMultiplication" elif only_op == "tessera.rope": diff --git a/python/tessera/compiler/schedule_ir.py b/python/tessera/compiler/schedule_ir.py index 1739c2b49..e81ab6690 100644 --- a/python/tessera/compiler/schedule_ir.py +++ b/python/tessera/compiler/schedule_ir.py @@ -384,9 +384,35 @@ def _base_attrs(op: IROp, ordinal: int) -> dict[str, Any]: } if op.result is not None: attrs["result"] = op.result + # Phase 8.4.4 — surface the operand element type so downstream Tile and + # Target IR layers can pick dtype-specific runtime symbols (e.g. mps_matmul + # f32/f16/bf16). The Graph IR encodes types as strings like "tensor<*xf16>"; + # we extract the trailing element-type token. Defaults to f32 when the + # operand types aren't parseable, preserving the pre-Phase 8.4.4 contract. + dtype = _resolve_element_dtype(op) + if dtype: + attrs["dtype"] = dtype return attrs +def _resolve_element_dtype(op: IROp) -> str | None: + operand_types = list(getattr(op, "operand_types", ()) or ()) + if not operand_types: + return None + # Pick the dtype from the first operand. The IROp invariant for compute + # ops is uniform element type across operands (e.g. matmul A and B match). + t = operand_types[0] + if "bf16" in t: + return "bf16" + if "f16" in t and "bf16" not in t: + return "f16" + if "f32" in t: + return "f32" + if "f64" in t: + return "f64" + return None + + def _shape_key(ops: list[IROp]) -> str: parts = [] for op in ops: diff --git a/python/tessera/compiler/target_ir.py b/python/tessera/compiler/target_ir.py index f0dd52a94..7cdc1c6e6 100644 --- a/python/tessera/compiler/target_ir.py +++ b/python/tessera/compiler/target_ir.py @@ -798,15 +798,23 @@ def _lower_apple_gpu_op(op: TileOp, *, mps_runtime: bool = False) -> list[Target if op.op_name.startswith("tessera.queue.") or op.op_name in {"tile.async_copy", "tile.wait_async"}: return [] # Phase 8.3 MPS runtime path: a single-matmul module is lowered to - # mps_matmul + mps_dispatch with execution_mode="metal_runtime". The - # AppleGPUToMPS pass and the apple_gpu_runtime.mm shim consume this op. + # mps_matmul + mps_dispatch with execution_mode="metal_runtime". Phase + # 8.4.4 — the dtype attr now reflects the element type the runtime will + # dispatch to (f32, f16, or bf16). The MatmulToAppleGPU lowering pass + # picks the matching runtime symbol; the IR-level dtype attr is the + # introspection mirror. if mps_runtime and op.op_name == "tile.mma" and source in {"tessera.matmul", "tessera.gemm"}: + # Resolve dtype from the tile op's attrs if present (Phase 8.4.4), + # else default to f32 (preserves Phase 8.3 contract). + tile_dtype = str(op.attrs.get("dtype", "f32")) + if tile_dtype not in {"f32", "f16", "bf16"}: + tile_dtype = "f32" return [ TargetOp("tessera_apple.gpu.mps_matmul", { **base, "framework": "MetalPerformanceShaders", "abi": "MPSMatrixMultiplication", - "dtype": "f32", + "dtype": tile_dtype, }), TargetOp("tessera_apple.gpu.mps_dispatch", { "ordinal": base["ordinal"], diff --git a/python/tessera/runtime.py b/python/tessera/runtime.py index eee1f5f0c..275f3a96d 100644 --- a/python/tessera/runtime.py +++ b/python/tessera/runtime.py @@ -1965,10 +1965,12 @@ def _apple_gpu_metadata_is_matmul_softmax_chain(ops: list[dict]) -> bool: def _apple_gpu_dispatch_matmul(op_name: str, operands: list[Any], np: Any) -> Any: - """Phase 8.3: dispatch a single rank-2 f32 matmul through the apple_gpu - runtime shim (MPSMatrixMultiplication when Metal is available, portable - reference fallback otherwise). Inputs outside the supported envelope fall - back to numpy.matmul — same shape as the apple_cpu dispatcher. + """Phase 8.3 + 8.4.4: dispatch a single rank-2 matmul through the + apple_gpu runtime shim. Picks the runtime symbol by element type: + - f32: native MPSDataTypeFloat32 (Phase 8.3) + - f16: native MPSDataTypeFloat16 (Phase 8.4.4) + - bf16: fp32-conversion path inside the shim (Phase 8.4.4) + Other dtypes fall back to numpy.matmul. """ if len(operands) != 2: @@ -1976,32 +1978,69 @@ def _apple_gpu_dispatch_matmul(op_name: str, operands: list[Any], np: Any) -> An a = np.asarray(operands[0]) b = np.asarray(operands[1]) - rank2_fast_path = ( - a.dtype == np.float32 - and b.dtype == np.float32 - and a.ndim == 2 - and b.ndim == 2 - ) - if not rank2_fast_path: + if a.ndim != 2 or b.ndim != 2 or a.shape[1] != b.shape[0]: + return np.matmul(a, b) + if a.dtype != b.dtype: return np.matmul(a, b) - if a.shape[1] != b.shape[0]: - raise ValueError(f"matmul shape mismatch: {a.shape} x {b.shape}") - if not a.flags.c_contiguous: - a = np.ascontiguousarray(a, dtype=np.float32) - if not b.flags.c_contiguous: - b = np.ascontiguousarray(b, dtype=np.float32) - out = np.zeros((a.shape[0], b.shape[1]), dtype=np.float32) - gemm = _apple_gpu_mps_matmul_f32() - gemm( - a.ctypes.data_as(ctypes.POINTER(ctypes.c_float)), - b.ctypes.data_as(ctypes.POINTER(ctypes.c_float)), - out.ctypes.data_as(ctypes.POINTER(ctypes.c_float)), - ctypes.c_int32(a.shape[0]), - ctypes.c_int32(b.shape[1]), - ctypes.c_int32(a.shape[1]), - ) - return out + if a.dtype == np.float32: + if not a.flags.c_contiguous: + a = np.ascontiguousarray(a, dtype=np.float32) + if not b.flags.c_contiguous: + b = np.ascontiguousarray(b, dtype=np.float32) + out = np.zeros((a.shape[0], b.shape[1]), dtype=np.float32) + gemm = _apple_gpu_mps_matmul_f32() + gemm( + a.ctypes.data_as(ctypes.POINTER(ctypes.c_float)), + b.ctypes.data_as(ctypes.POINTER(ctypes.c_float)), + out.ctypes.data_as(ctypes.POINTER(ctypes.c_float)), + ctypes.c_int32(a.shape[0]), + ctypes.c_int32(b.shape[1]), + ctypes.c_int32(a.shape[1]), + ) + return out + + if a.dtype == np.float16: + if not a.flags.c_contiguous: + a = np.ascontiguousarray(a, dtype=np.float16) + if not b.flags.c_contiguous: + b = np.ascontiguousarray(b, dtype=np.float16) + out = np.zeros((a.shape[0], b.shape[1]), dtype=np.float16) + gemm_f16 = _apple_gpu_mps_matmul_f16() + if gemm_f16 is not None: + gemm_f16( + a.view(np.uint16).ctypes.data_as(ctypes.POINTER(ctypes.c_uint16)), + b.view(np.uint16).ctypes.data_as(ctypes.POINTER(ctypes.c_uint16)), + out.view(np.uint16).ctypes.data_as(ctypes.POINTER(ctypes.c_uint16)), + ctypes.c_int32(a.shape[0]), + ctypes.c_int32(b.shape[1]), + ctypes.c_int32(a.shape[1]), + ) + return out + # Older runtime build without the f16 symbol — convert to f32 and back. + return (a.astype(np.float32) @ b.astype(np.float32)).astype(np.float16) + + bf16_dtype = _bfloat16_dtype() + if bf16_dtype is not None and a.dtype == bf16_dtype: + if not a.flags.c_contiguous: + a = np.ascontiguousarray(a, dtype=bf16_dtype) + if not b.flags.c_contiguous: + b = np.ascontiguousarray(b, dtype=bf16_dtype) + out = np.zeros((a.shape[0], b.shape[1]), dtype=bf16_dtype) + gemm_bf16 = _apple_gpu_mps_matmul_bf16() + if gemm_bf16 is not None: + gemm_bf16( + a.view(np.uint16).ctypes.data_as(ctypes.POINTER(ctypes.c_uint16)), + b.view(np.uint16).ctypes.data_as(ctypes.POINTER(ctypes.c_uint16)), + out.view(np.uint16).ctypes.data_as(ctypes.POINTER(ctypes.c_uint16)), + ctypes.c_int32(a.shape[0]), + ctypes.c_int32(b.shape[1]), + ctypes.c_int32(a.shape[1]), + ) + return out + return (a.astype(np.float32) @ b.astype(np.float32)).astype(bf16_dtype) + + return np.matmul(a, b) def _apple_gpu_mps_matmul_f32() -> Any: @@ -2019,6 +2058,48 @@ def _apple_gpu_mps_matmul_f32() -> Any: return sym +def _apple_gpu_mps_matmul_f16() -> Any: + """Phase 8.4.4 — fp16 matmul via MPSDataTypeFloat16. Inputs/outputs are + bit-pattern uint16_t* (numpy float16 layout). Symbol may be absent on + older runtime builds; the dispatcher falls through to fp32 conversion.""" + + runtime = _load_apple_gpu_runtime() + sym = getattr(runtime, "tessera_apple_gpu_mps_matmul_f16", None) + if sym is None: + return None + sym.argtypes = [ + ctypes.POINTER(ctypes.c_uint16), + ctypes.POINTER(ctypes.c_uint16), + ctypes.POINTER(ctypes.c_uint16), + ctypes.c_int32, + ctypes.c_int32, + ctypes.c_int32, + ] + sym.restype = None + return sym + + +def _apple_gpu_mps_matmul_bf16() -> Any: + """Phase 8.4.4 — bf16 matmul. The runtime shim does fp32 conversion + inside since MPS doesn't natively support bf16 matrix descriptors as of + macOS 14. ml_dtypes.bfloat16 dtype is byte-compatible with the C ABI.""" + + runtime = _load_apple_gpu_runtime() + sym = getattr(runtime, "tessera_apple_gpu_mps_matmul_bf16", None) + if sym is None: + return None + sym.argtypes = [ + ctypes.POINTER(ctypes.c_uint16), + ctypes.POINTER(ctypes.c_uint16), + ctypes.POINTER(ctypes.c_uint16), + ctypes.c_int32, + ctypes.c_int32, + ctypes.c_int32, + ] + sym.restype = None + return sym + + def _apple_gpu_dispatch_rope(op_name: str, operands: list[Any], np: Any) -> Any: """Phase 8.4: dispatch a single rank-2 f32 rope through the apple_gpu runtime shim's custom MSL kernel. Inputs outside the supported envelope @@ -2318,6 +2399,10 @@ def _load_apple_gpu_runtime() -> ctypes.CDLL: getattr(lib, "tessera_apple_gpu_softmax_f32") getattr(lib, "tessera_apple_gpu_gelu_f32") getattr(lib, "tessera_apple_gpu_matmul_softmax_f32") + # Phase 8.4.4 — require fp16/bf16 matmul symbols too. Older + # builds lack them; falling through forces a rebuild. + getattr(lib, "tessera_apple_gpu_mps_matmul_f16") + getattr(lib, "tessera_apple_gpu_mps_matmul_bf16") _apple_gpu_runtime = lib return _apple_gpu_runtime except (OSError, AttributeError): diff --git a/src/compiler/codegen/Tessera_Apple_Backend/lib/Target/Apple/Lowering/MatmulToAppleGPU.cpp b/src/compiler/codegen/Tessera_Apple_Backend/lib/Target/Apple/Lowering/MatmulToAppleGPU.cpp index 8eda0752c..6b60d0bed 100644 --- a/src/compiler/codegen/Tessera_Apple_Backend/lib/Target/Apple/Lowering/MatmulToAppleGPU.cpp +++ b/src/compiler/codegen/Tessera_Apple_Backend/lib/Target/Apple/Lowering/MatmulToAppleGPU.cpp @@ -37,6 +37,10 @@ namespace { constexpr llvm::StringLiteral kGemmF32Symbol = "tessera_apple_gpu_mps_matmul_f32"; +constexpr llvm::StringLiteral kGemmF16Symbol = + "tessera_apple_gpu_mps_matmul_f16"; +constexpr llvm::StringLiteral kGemmBF16Symbol = + "tessera_apple_gpu_mps_matmul_bf16"; static func::FuncOp ensureExternalDecl(ModuleOp mod, StringRef name, FunctionType fnTy) { @@ -74,9 +78,26 @@ struct LowerMatmulToAppleGPU : public RewritePattern { Type lhsElem = lhsTy.getElementType(); Type rhsElem = rhsTy.getElementType(); - if (!lhsElem.isF32() || !rhsElem.isF32()) + if (lhsElem != rhsElem) return rewriter.notifyMatchFailure( - op, "AppleGPU MPS path is f32-only in Phase 8.3"); + op, "AppleGPU MPS path requires matching matmul element types"); + + // Phase 8.4.4 — pick the runtime symbol based on the input element type. + // f32 routes to MPSDataTypeFloat32 (Phase 8.3); f16 routes to native + // MPSDataTypeFloat16; bf16 routes to a fp32-conversion path inside the + // shim because MPS doesn't natively accept bf16 matrix descriptors as of + // macOS 14. + StringRef symbol; + if (lhsElem.isF32()) { + symbol = kGemmF32Symbol; + } else if (lhsElem.isF16()) { + symbol = kGemmF16Symbol; + } else if (lhsElem.isBF16()) { + symbol = kGemmBF16Symbol; + } else { + return rewriter.notifyMatchFailure( + op, "AppleGPU MPS path supports f32, f16, and bf16 in Phase 8.4.4"); + } if (lhsTy.isDynamicDim(0) || lhsTy.isDynamicDim(1) || rhsTy.isDynamicDim(0) || rhsTy.isDynamicDim(1)) @@ -95,11 +116,10 @@ struct LowerMatmulToAppleGPU : public RewritePattern { Type i64Ty = rewriter.getI64Type(); Type i32Ty = rewriter.getI32Type(); - Type f32Ty = rewriter.getF32Type(); - auto lhsMemTy = MemRefType::get({M, K}, f32Ty); - auto rhsMemTy = MemRefType::get({K, N}, f32Ty); - auto outMemTy = MemRefType::get({M, N}, f32Ty); + auto lhsMemTy = MemRefType::get({M, K}, lhsElem); + auto rhsMemTy = MemRefType::get({K, N}, rhsElem); + auto outMemTy = MemRefType::get({M, N}, lhsElem); Value aPtr = extractPtr(rewriter, loc, lhs, lhsMemTy); Value bPtr = extractPtr(rewriter, loc, rhs, rhsMemTy); @@ -115,15 +135,18 @@ struct LowerMatmulToAppleGPU : public RewritePattern { Value Nv = rewriter.create(loc, N, 32); Value Kv = rewriter.create(loc, K, 32); + // The runtime ABI is the same shape for all three dtypes — three i64 + // pointers + three i32 dim sizes. The element type is encoded in the + // symbol name, not the signature. FunctionType gemmFnTy = FunctionType::get( ctx, {i64Ty, i64Ty, i64Ty, i32Ty, i32Ty, i32Ty}, {}); - ensureExternalDecl(mod, kGemmF32Symbol, gemmFnTy); + ensureExternalDecl(mod, symbol, gemmFnTy); rewriter.create( - loc, kGemmF32Symbol, TypeRange{}, + loc, symbol, TypeRange{}, ValueRange{aPtr, bPtr, cPtr, Mv, Nv, Kv}); - auto outTensorTy = RankedTensorType::get({M, N}, f32Ty); + auto outTensorTy = RankedTensorType::get({M, N}, lhsElem); Value result = rewriter.create(loc, outTensorTy, cAlloc); rewriter.replaceOp(op, result); @@ -140,8 +163,8 @@ struct LowerMatmulToAppleGPUPass return "tessera-matmul-to-apple_gpu"; } StringRef getDescription() const override { - return "Lower tessera.matmul (rank-2, f32) to Apple GPU runtime calls " - "(MPSMatrixMultiplication)"; + return "Lower tessera.matmul (rank-2, f32/f16/bf16) to Apple GPU runtime " + "calls (MPSMatrixMultiplication)"; } void getDependentDialects(DialectRegistry ®istry) const override { diff --git a/src/compiler/codegen/Tessera_Apple_Backend/runtime/apple_gpu_runtime.mm b/src/compiler/codegen/Tessera_Apple_Backend/runtime/apple_gpu_runtime.mm index cd6266826..e50305bca 100644 --- a/src/compiler/codegen/Tessera_Apple_Backend/runtime/apple_gpu_runtime.mm +++ b/src/compiler/codegen/Tessera_Apple_Backend/runtime/apple_gpu_runtime.mm @@ -185,6 +185,235 @@ bool dispatch_mps_gemm_f32(MetalDeviceContext &ctx, const float* A, reference_gemm_f32(A, B, C, M, N, K); } +//===---------------------------------------------------------------------===// +// Phase 8.4.4 — fp16 / bf16 matmul (mirrors Phase 8.2 BNNS bf16 pattern) +// +// fp16: native MPSDataTypeFloat16 path. Apple Silicon GPUs run fp16 +// natively at 2x throughput vs fp32 on most ops. +// bf16: MPS does NOT directly support bf16 matmul as of macOS 14, so this +// path uses fp32 conversion at the boundary — load with bit-shift +// (bf16 -> fp32), run MPSDataTypeFloat32, convert back. Same shape +// as the BNNS bf16 fallback in apple_cpu_runtime.cpp. +// +// At the C ABI boundary fp16/bf16 inputs are passed as uint16_t* (the bit +// pattern). This keeps the ABI portable across compilers regardless of +// _Float16 / __bf16 availability. +//===---------------------------------------------------------------------===// + +namespace { + +// fp16 <-> float helpers (bit-pattern; no _Float16 dependency). +inline float half_to_float_gpu(uint16_t h) { + uint32_t sign = (uint32_t(h) & 0x8000u) << 16; + uint32_t exp = (uint32_t(h) & 0x7C00u) >> 10; + uint32_t frac = uint32_t(h) & 0x03FFu; + uint32_t f; + if (exp == 0) { + if (frac == 0) { + f = sign; + } else { + while ((frac & 0x0400u) == 0) { frac <<= 1; exp -= 1; } + exp += 1; + frac &= ~0x0400u; + f = sign | ((exp + 112) << 23) | (frac << 13); + } + } else if (exp == 0x1F) { + f = sign | 0x7F800000u | (frac << 13); + } else { + f = sign | ((exp + 112) << 23) | (frac << 13); + } + float out; + std::memcpy(&out, &f, sizeof(out)); + return out; +} + +inline uint16_t float_to_half_gpu(float v) { + uint32_t f; + std::memcpy(&f, &v, sizeof(f)); + uint32_t sign = (f >> 16) & 0x8000u; + int32_t exp = int32_t((f >> 23) & 0xFFu) - 127 + 15; + uint32_t frac = f & 0x007FFFFFu; + if (exp <= 0) { + if (exp < -10) return uint16_t(sign); + frac = (frac | 0x00800000u) >> (1 - exp); + if (frac & 0x00001000u) frac += 0x00002000u; + return uint16_t(sign | (frac >> 13)); + } + if (exp >= 0x1F) { + if (((f >> 23) & 0xFFu) == 0xFFu) { + return uint16_t(sign | 0x7C00u | (frac ? (frac >> 13) | 0x200u : 0)); + } + return uint16_t(sign | 0x7C00u); + } + if (frac & 0x00001000u) { + frac += 0x00002000u; + if (frac & 0x00800000u) { + frac = 0; + exp += 1; + if (exp >= 0x1F) return uint16_t(sign | 0x7C00u); + } + } + return uint16_t(sign | (uint32_t(exp) << 10) | (frac >> 13)); +} + +// bf16 <-> float helpers (bit-shift; round-to-nearest-even on store). +inline float bfloat16_to_float_gpu(uint16_t b) { + uint32_t f = static_cast(b) << 16; + float out; + std::memcpy(&out, &f, sizeof(out)); + return out; +} + +inline uint16_t float_to_bfloat16_gpu(float v) { + uint32_t f; + std::memcpy(&f, &v, sizeof(f)); + if ((f & 0x7FC00000u) == 0x7F800000u && (f & 0x007FFFFFu) != 0) { + return static_cast((f >> 16) | 0x40u); + } + uint32_t lsb = (f >> 16) & 1u; + uint32_t rounded = f + 0x7FFFu + lsb; + return static_cast(rounded >> 16); +} + +bool dispatch_mps_gemm_f16(MetalDeviceContext &ctx, const uint16_t* A, + const uint16_t* B, uint16_t* C, + int32_t M, int32_t N, int32_t K) { + // Same shape as dispatch_mps_gemm_f32 — only the MPS data type and the + // per-element byte count change. Apple GPUs run fp16 natively at higher + // throughput than fp32 so this is a real perf win, not just convenience. + @autoreleasepool { + NSUInteger byteCountA = sizeof(uint16_t) * static_cast(M) * + static_cast(K); + NSUInteger byteCountB = sizeof(uint16_t) * static_cast(K) * + static_cast(N); + NSUInteger byteCountC = sizeof(uint16_t) * static_cast(M) * + static_cast(N); + + id bufA = [ctx.device newBufferWithBytes:A + length:byteCountA + options:MTLResourceStorageModeShared]; + id bufB = [ctx.device newBufferWithBytes:B + length:byteCountB + options:MTLResourceStorageModeShared]; + id bufC = [ctx.device newBufferWithLength:byteCountC + options:MTLResourceStorageModeShared]; + if (!bufA || !bufB || !bufC) return false; + + NSUInteger rowBytesA = sizeof(uint16_t) * static_cast(K); + NSUInteger rowBytesB = sizeof(uint16_t) * static_cast(N); + NSUInteger rowBytesC = sizeof(uint16_t) * static_cast(N); + + MPSMatrixDescriptor *descA = + [MPSMatrixDescriptor matrixDescriptorWithRows:static_cast(M) + columns:static_cast(K) + rowBytes:rowBytesA + dataType:MPSDataTypeFloat16]; + MPSMatrixDescriptor *descB = + [MPSMatrixDescriptor matrixDescriptorWithRows:static_cast(K) + columns:static_cast(N) + rowBytes:rowBytesB + dataType:MPSDataTypeFloat16]; + MPSMatrixDescriptor *descC = + [MPSMatrixDescriptor matrixDescriptorWithRows:static_cast(M) + columns:static_cast(N) + rowBytes:rowBytesC + dataType:MPSDataTypeFloat16]; + + MPSMatrix *matA = [[MPSMatrix alloc] initWithBuffer:bufA descriptor:descA]; + MPSMatrix *matB = [[MPSMatrix alloc] initWithBuffer:bufB descriptor:descB]; + MPSMatrix *matC = [[MPSMatrix alloc] initWithBuffer:bufC descriptor:descC]; + + MPSMatrixMultiplication *kernel = [[MPSMatrixMultiplication alloc] + initWithDevice:ctx.device + transposeLeft:NO + transposeRight:NO + resultRows:static_cast(M) + resultColumns:static_cast(N) + interiorColumns:static_cast(K) + alpha:1.0 + beta:0.0]; + + id cb = [ctx.queue commandBuffer]; + [kernel encodeToCommandBuffer:cb + leftMatrix:matA + rightMatrix:matB + resultMatrix:matC]; + [cb commit]; + [cb waitUntilCompleted]; + + if (cb.status != MTLCommandBufferStatusCompleted) return false; + std::memcpy(C, [bufC contents], byteCountC); + return true; + } +} + +inline void reference_gemm_f16_via_fp32(const uint16_t* A, const uint16_t* B, + uint16_t* C, int32_t M, int32_t N, + int32_t K) { + // Convert each operand to fp32, run the existing reference kernel, convert + // back. Same numerical contract as the BNNS-fallback path in + // apple_cpu_runtime.cpp. + std::vector Af(static_cast(M) * K); + std::vector Bf(static_cast(K) * N); + std::vector Cf(static_cast(M) * N, 0.0f); + for (std::size_t i = 0; i < Af.size(); ++i) Af[i] = half_to_float_gpu(A[i]); + for (std::size_t i = 0; i < Bf.size(); ++i) Bf[i] = half_to_float_gpu(B[i]); + reference_gemm_f32(Af.data(), Bf.data(), Cf.data(), M, N, K); + for (std::size_t i = 0; i < Cf.size(); ++i) C[i] = float_to_half_gpu(Cf[i]); +} + +inline void reference_gemm_bf16_via_fp32(const uint16_t* A, const uint16_t* B, + uint16_t* C, int32_t M, int32_t N, + int32_t K) { + std::vector Af(static_cast(M) * K); + std::vector Bf(static_cast(K) * N); + std::vector Cf(static_cast(M) * N, 0.0f); + for (std::size_t i = 0; i < Af.size(); ++i) Af[i] = bfloat16_to_float_gpu(A[i]); + for (std::size_t i = 0; i < Bf.size(); ++i) Bf[i] = bfloat16_to_float_gpu(B[i]); + reference_gemm_f32(Af.data(), Bf.data(), Cf.data(), M, N, K); + for (std::size_t i = 0; i < Cf.size(); ++i) C[i] = float_to_bfloat16_gpu(Cf[i]); +} + +bool dispatch_mps_gemm_bf16_via_fp32(MetalDeviceContext &ctx, + const uint16_t* A, const uint16_t* B, + uint16_t* C, int32_t M, int32_t N, + int32_t K) { + // bf16 -> fp32 -> MPSDataTypeFloat32 -> fp32 -> bf16. MPS does not + // natively accept bf16 as of macOS 14; this path keeps the bf16 ABI + // honest while still spending the heavy compute on the GPU. + std::vector Af(static_cast(M) * K); + std::vector Bf(static_cast(K) * N); + std::vector Cf(static_cast(M) * N, 0.0f); + for (std::size_t i = 0; i < Af.size(); ++i) Af[i] = bfloat16_to_float_gpu(A[i]); + for (std::size_t i = 0; i < Bf.size(); ++i) Bf[i] = bfloat16_to_float_gpu(B[i]); + if (!dispatch_mps_gemm_f32(ctx, Af.data(), Bf.data(), Cf.data(), M, N, K)) + return false; + for (std::size_t i = 0; i < Cf.size(); ++i) C[i] = float_to_bfloat16_gpu(Cf[i]); + return true; +} + +} // namespace + +extern "C" void tessera_apple_gpu_mps_matmul_f16(const uint16_t* A, + const uint16_t* B, + uint16_t* C, + int32_t M, int32_t N, + int32_t K) { + MetalDeviceContext &ctx = deviceContext(); + if (ctx.ok && dispatch_mps_gemm_f16(ctx, A, B, C, M, N, K)) return; + reference_gemm_f16_via_fp32(A, B, C, M, N, K); +} + +extern "C" void tessera_apple_gpu_mps_matmul_bf16(const uint16_t* A, + const uint16_t* B, + uint16_t* C, + int32_t M, int32_t N, + int32_t K) { + MetalDeviceContext &ctx = deviceContext(); + if (ctx.ok && dispatch_mps_gemm_bf16_via_fp32(ctx, A, B, C, M, N, K)) return; + reference_gemm_bf16_via_fp32(A, B, C, M, N, K); +} + //===---------------------------------------------------------------------===// // Phase 8.4 — Custom MSL kernel infrastructure // diff --git a/src/compiler/codegen/Tessera_Apple_Backend/runtime/apple_gpu_runtime_stub.cpp b/src/compiler/codegen/Tessera_Apple_Backend/runtime/apple_gpu_runtime_stub.cpp index 02c07cbbe..2953c0ada 100644 --- a/src/compiler/codegen/Tessera_Apple_Backend/runtime/apple_gpu_runtime_stub.cpp +++ b/src/compiler/codegen/Tessera_Apple_Backend/runtime/apple_gpu_runtime_stub.cpp @@ -13,6 +13,7 @@ #include #include #include +#include #if !defined(__APPLE__) @@ -60,6 +61,115 @@ extern "C" void tessera_apple_gpu_mps_matmul_f32(const float* A, reference_gemm_f32(A, B, C, M, N, K); } +namespace { + +// Phase 8.4.4 — fp16/bf16 bit-pattern conversion helpers for the non-Darwin +// stub. Same shape as the conversion helpers in apple_gpu_runtime.mm; kept +// inline here so the stub TU is self-contained. + +inline float half_to_float_stub(uint16_t h) { + uint32_t sign = (uint32_t(h) & 0x8000u) << 16; + uint32_t exp = (uint32_t(h) & 0x7C00u) >> 10; + uint32_t frac = uint32_t(h) & 0x03FFu; + uint32_t f; + if (exp == 0) { + if (frac == 0) { + f = sign; + } else { + while ((frac & 0x0400u) == 0) { frac <<= 1; exp -= 1; } + exp += 1; + frac &= ~0x0400u; + f = sign | ((exp + 112) << 23) | (frac << 13); + } + } else if (exp == 0x1F) { + f = sign | 0x7F800000u | (frac << 13); + } else { + f = sign | ((exp + 112) << 23) | (frac << 13); + } + float out; + std::memcpy(&out, &f, sizeof(out)); + return out; +} + +inline uint16_t float_to_half_stub(float v) { + uint32_t f; + std::memcpy(&f, &v, sizeof(f)); + uint32_t sign = (f >> 16) & 0x8000u; + int32_t exp = int32_t((f >> 23) & 0xFFu) - 127 + 15; + uint32_t frac = f & 0x007FFFFFu; + if (exp <= 0) { + if (exp < -10) return uint16_t(sign); + frac = (frac | 0x00800000u) >> (1 - exp); + if (frac & 0x00001000u) frac += 0x00002000u; + return uint16_t(sign | (frac >> 13)); + } + if (exp >= 0x1F) { + if (((f >> 23) & 0xFFu) == 0xFFu) { + return uint16_t(sign | 0x7C00u | (frac ? (frac >> 13) | 0x200u : 0)); + } + return uint16_t(sign | 0x7C00u); + } + if (frac & 0x00001000u) { + frac += 0x00002000u; + if (frac & 0x00800000u) { + frac = 0; + exp += 1; + if (exp >= 0x1F) return uint16_t(sign | 0x7C00u); + } + } + return uint16_t(sign | (uint32_t(exp) << 10) | (frac >> 13)); +} + +inline float bfloat16_to_float_stub(uint16_t b) { + uint32_t f = static_cast(b) << 16; + float out; + std::memcpy(&out, &f, sizeof(out)); + return out; +} + +inline uint16_t float_to_bfloat16_stub(float v) { + uint32_t f; + std::memcpy(&f, &v, sizeof(f)); + if ((f & 0x7FC00000u) == 0x7F800000u && (f & 0x007FFFFFu) != 0) { + return static_cast((f >> 16) | 0x40u); + } + uint32_t lsb = (f >> 16) & 1u; + uint32_t rounded = f + 0x7FFFu + lsb; + return static_cast(rounded >> 16); +} + +} // namespace + +extern "C" void tessera_apple_gpu_mps_matmul_f16(const uint16_t* A, + const uint16_t* B, + uint16_t* C, + int32_t M, int32_t N, + int32_t K) { + // Convert each operand to fp32, run the reference fp32 GEMM, convert back. + // Same numerical contract as the BNNS-fallback path on CPU. + std::vector Af(static_cast(M) * K); + std::vector Bf(static_cast(K) * N); + std::vector Cf(static_cast(M) * N, 0.0f); + for (std::size_t i = 0; i < Af.size(); ++i) Af[i] = half_to_float_stub(A[i]); + for (std::size_t i = 0; i < Bf.size(); ++i) Bf[i] = half_to_float_stub(B[i]); + reference_gemm_f32(Af.data(), Bf.data(), Cf.data(), M, N, K); + for (std::size_t i = 0; i < Cf.size(); ++i) C[i] = float_to_half_stub(Cf[i]); +} + +extern "C" void tessera_apple_gpu_mps_matmul_bf16(const uint16_t* A, + const uint16_t* B, + uint16_t* C, + int32_t M, int32_t N, + int32_t K) { + std::vector Af(static_cast(M) * K); + std::vector Bf(static_cast(K) * N); + std::vector Cf(static_cast(M) * N, 0.0f); + for (std::size_t i = 0; i < Af.size(); ++i) Af[i] = bfloat16_to_float_stub(A[i]); + for (std::size_t i = 0; i < Bf.size(); ++i) Bf[i] = bfloat16_to_float_stub(B[i]); + reference_gemm_f32(Af.data(), Bf.data(), Cf.data(), M, N, K); + for (std::size_t i = 0; i < Cf.size(); ++i) C[i] = float_to_bfloat16_stub(Cf[i]); +} + extern "C" void tessera_apple_gpu_rope_f32(const float* X, const float* Theta, float* Out, int32_t M, int32_t K) { reference_rope_f32(X, Theta, Out, M, K); diff --git a/tests/tessera-ir/.lit_test_times.txt b/tests/tessera-ir/.lit_test_times.txt index f2326a875..c5bd069de 100644 --- a/tests/tessera-ir/.lit_test_times.txt +++ b/tests/tessera-ir/.lit_test_times.txt @@ -1,8 +1,8 @@ -7.576761e-01 phase8/apple_cpu_lowering.mlir -7.509289e-01 phase8/apple_dialect_roundtrip.mlir -7.781518e-01 phase8/apple_gpu_lowering.mlir -7.821021e-01 phase8/apple_cpu_runtime.mlir -1.081014e-02 phase8/target_ir_contracts.mlir +8.501682e-01 phase8/apple_cpu_lowering.mlir +8.418021e-01 phase8/apple_dialect_roundtrip.mlir +8.678038e-01 phase8/apple_gpu_lowering.mlir +8.748050e-01 phase8/apple_cpu_runtime.mlir +1.045299e-02 phase8/target_ir_contracts.mlir 4.739904e-02 phase2/distribution_lowering.mlir 4.920816e-02 phase2/effect_annotation.mlir 4.286575e-02 phase2/full_pipeline.mlir @@ -29,9 +29,10 @@ 4.314494e-02 phase7/neighbors_stencil_lower.mlir 4.376101e-02 phase7/shardy_export.mlir 3.875470e-02 pipelines/cleanup_pipeline.mlir -7.705491e-01 phase8/apple_gpu_runtime.mlir -1.354003e-02 phase8/tmem_tcgen05_contract.mlir -7.526319e-01 phase8/apple_gpu_msl.mlir -7.745039e-01 phase8/apple_gpu_flash_attn.mlir -8.078930e-01 phase8/apple_gpu_softmax_gelu.mlir -7.541583e-01 phase8/apple_gpu_matmul_softmax_fusion.mlir +8.595231e-01 phase8/apple_gpu_runtime.mlir +1.159406e-02 phase8/tmem_tcgen05_contract.mlir +8.450270e-01 phase8/apple_gpu_msl.mlir +8.636880e-01 phase8/apple_gpu_flash_attn.mlir +8.775210e-01 phase8/apple_gpu_softmax_gelu.mlir +8.536379e-01 phase8/apple_gpu_matmul_softmax_fusion.mlir +8.237970e-01 phase8/apple_gpu_matmul_dtypes.mlir diff --git a/tests/tessera-ir/phase8/Output/apple_cpu_lowering.mlir.script b/tests/tessera-ir/phase8/Output/apple_cpu_lowering.mlir.script index aa8b056a3..8ffc9539c 100644 --- a/tests/tessera-ir/phase8/Output/apple_cpu_lowering.mlir.script +++ b/tests/tessera-ir/phase8/Output/apple_cpu_lowering.mlir.script @@ -1 +1 @@ -set -o pipefail;set -x;{ { set +x; } 2>/dev/null && echo 'RUN: at line 1': 'tessera-opt -allow-unregistered-dialect /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-3-msl-fusion/tests/tessera-ir/phase8/apple_cpu_lowering.mlir -tessera-lower-to-apple_cpu --allow-unregistered-dialect | FileCheck /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-3-msl-fusion/tests/tessera-ir/phase8/apple_cpu_lowering.mlir' >&2 && { set -x; } 2>/dev/null && { tessera-opt -allow-unregistered-dialect /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-3-msl-fusion/tests/tessera-ir/phase8/apple_cpu_lowering.mlir -tessera-lower-to-apple_cpu --allow-unregistered-dialect | FileCheck /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-3-msl-fusion/tests/tessera-ir/phase8/apple_cpu_lowering.mlir; }; } +set -o pipefail;set -x;{ { set +x; } 2>/dev/null && echo 'RUN: at line 1': 'tessera-opt -allow-unregistered-dialect /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-4-fp16-bf16/tests/tessera-ir/phase8/apple_cpu_lowering.mlir -tessera-lower-to-apple_cpu --allow-unregistered-dialect | FileCheck /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-4-fp16-bf16/tests/tessera-ir/phase8/apple_cpu_lowering.mlir' >&2 && { set -x; } 2>/dev/null && { tessera-opt -allow-unregistered-dialect /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-4-fp16-bf16/tests/tessera-ir/phase8/apple_cpu_lowering.mlir -tessera-lower-to-apple_cpu --allow-unregistered-dialect | FileCheck /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-4-fp16-bf16/tests/tessera-ir/phase8/apple_cpu_lowering.mlir; }; } diff --git a/tests/tessera-ir/phase8/Output/apple_cpu_runtime.mlir.script b/tests/tessera-ir/phase8/Output/apple_cpu_runtime.mlir.script index eea81a58b..dea9258e6 100644 --- a/tests/tessera-ir/phase8/Output/apple_cpu_runtime.mlir.script +++ b/tests/tessera-ir/phase8/Output/apple_cpu_runtime.mlir.script @@ -1 +1 @@ -set -o pipefail;set -x;{ { set +x; } 2>/dev/null && echo 'RUN: at line 1': 'tessera-opt -allow-unregistered-dialect /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-3-msl-fusion/tests/tessera-ir/phase8/apple_cpu_runtime.mlir --pass-pipeline='"'"'builtin.module(tessera-lower-to-apple_cpu-runtime)'"'"' --allow-unregistered-dialect | FileCheck /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-3-msl-fusion/tests/tessera-ir/phase8/apple_cpu_runtime.mlir' >&2 && { set -x; } 2>/dev/null && { tessera-opt -allow-unregistered-dialect /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-3-msl-fusion/tests/tessera-ir/phase8/apple_cpu_runtime.mlir --pass-pipeline='builtin.module(tessera-lower-to-apple_cpu-runtime)' --allow-unregistered-dialect | FileCheck /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-3-msl-fusion/tests/tessera-ir/phase8/apple_cpu_runtime.mlir; }; } +set -o pipefail;set -x;{ { set +x; } 2>/dev/null && echo 'RUN: at line 1': 'tessera-opt -allow-unregistered-dialect /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-4-fp16-bf16/tests/tessera-ir/phase8/apple_cpu_runtime.mlir --pass-pipeline='"'"'builtin.module(tessera-lower-to-apple_cpu-runtime)'"'"' --allow-unregistered-dialect | FileCheck /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-4-fp16-bf16/tests/tessera-ir/phase8/apple_cpu_runtime.mlir' >&2 && { set -x; } 2>/dev/null && { tessera-opt -allow-unregistered-dialect /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-4-fp16-bf16/tests/tessera-ir/phase8/apple_cpu_runtime.mlir --pass-pipeline='builtin.module(tessera-lower-to-apple_cpu-runtime)' --allow-unregistered-dialect | FileCheck /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-4-fp16-bf16/tests/tessera-ir/phase8/apple_cpu_runtime.mlir; }; } diff --git a/tests/tessera-ir/phase8/Output/apple_dialect_roundtrip.mlir.script b/tests/tessera-ir/phase8/Output/apple_dialect_roundtrip.mlir.script index 30b8a82a3..0225a1d31 100644 --- a/tests/tessera-ir/phase8/Output/apple_dialect_roundtrip.mlir.script +++ b/tests/tessera-ir/phase8/Output/apple_dialect_roundtrip.mlir.script @@ -1 +1 @@ -set -o pipefail;set -x;{ { set +x; } 2>/dev/null && echo 'RUN: at line 1': 'tessera-opt -allow-unregistered-dialect /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-3-msl-fusion/tests/tessera-ir/phase8/apple_dialect_roundtrip.mlir --allow-unregistered-dialect | FileCheck /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-3-msl-fusion/tests/tessera-ir/phase8/apple_dialect_roundtrip.mlir' >&2 && { set -x; } 2>/dev/null && { tessera-opt -allow-unregistered-dialect /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-3-msl-fusion/tests/tessera-ir/phase8/apple_dialect_roundtrip.mlir --allow-unregistered-dialect | FileCheck /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-3-msl-fusion/tests/tessera-ir/phase8/apple_dialect_roundtrip.mlir; }; } +set -o pipefail;set -x;{ { set +x; } 2>/dev/null && echo 'RUN: at line 1': 'tessera-opt -allow-unregistered-dialect /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-4-fp16-bf16/tests/tessera-ir/phase8/apple_dialect_roundtrip.mlir --allow-unregistered-dialect | FileCheck /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-4-fp16-bf16/tests/tessera-ir/phase8/apple_dialect_roundtrip.mlir' >&2 && { set -x; } 2>/dev/null && { tessera-opt -allow-unregistered-dialect /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-4-fp16-bf16/tests/tessera-ir/phase8/apple_dialect_roundtrip.mlir --allow-unregistered-dialect | FileCheck /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-4-fp16-bf16/tests/tessera-ir/phase8/apple_dialect_roundtrip.mlir; }; } diff --git a/tests/tessera-ir/phase8/Output/apple_gpu_flash_attn.mlir.script b/tests/tessera-ir/phase8/Output/apple_gpu_flash_attn.mlir.script index f6ef661c1..f00b9bdcb 100644 --- a/tests/tessera-ir/phase8/Output/apple_gpu_flash_attn.mlir.script +++ b/tests/tessera-ir/phase8/Output/apple_gpu_flash_attn.mlir.script @@ -1 +1 @@ -set -o pipefail;set -x;{ { set +x; } 2>/dev/null && echo 'RUN: at line 1': 'tessera-opt -allow-unregistered-dialect /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-3-msl-fusion/tests/tessera-ir/phase8/apple_gpu_flash_attn.mlir --pass-pipeline='"'"'builtin.module(tessera-lower-to-apple_gpu-runtime)'"'"' --allow-unregistered-dialect | FileCheck /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-3-msl-fusion/tests/tessera-ir/phase8/apple_gpu_flash_attn.mlir' >&2 && { set -x; } 2>/dev/null && { tessera-opt -allow-unregistered-dialect /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-3-msl-fusion/tests/tessera-ir/phase8/apple_gpu_flash_attn.mlir --pass-pipeline='builtin.module(tessera-lower-to-apple_gpu-runtime)' --allow-unregistered-dialect | FileCheck /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-3-msl-fusion/tests/tessera-ir/phase8/apple_gpu_flash_attn.mlir; }; } +set -o pipefail;set -x;{ { set +x; } 2>/dev/null && echo 'RUN: at line 1': 'tessera-opt -allow-unregistered-dialect /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-4-fp16-bf16/tests/tessera-ir/phase8/apple_gpu_flash_attn.mlir --pass-pipeline='"'"'builtin.module(tessera-lower-to-apple_gpu-runtime)'"'"' --allow-unregistered-dialect | FileCheck /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-4-fp16-bf16/tests/tessera-ir/phase8/apple_gpu_flash_attn.mlir' >&2 && { set -x; } 2>/dev/null && { tessera-opt -allow-unregistered-dialect /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-4-fp16-bf16/tests/tessera-ir/phase8/apple_gpu_flash_attn.mlir --pass-pipeline='builtin.module(tessera-lower-to-apple_gpu-runtime)' --allow-unregistered-dialect | FileCheck /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-4-fp16-bf16/tests/tessera-ir/phase8/apple_gpu_flash_attn.mlir; }; } diff --git a/tests/tessera-ir/phase8/Output/apple_gpu_lowering.mlir.script b/tests/tessera-ir/phase8/Output/apple_gpu_lowering.mlir.script index 05384828b..6b0ec2fe7 100644 --- a/tests/tessera-ir/phase8/Output/apple_gpu_lowering.mlir.script +++ b/tests/tessera-ir/phase8/Output/apple_gpu_lowering.mlir.script @@ -1 +1 @@ -set -o pipefail;set -x;{ { set +x; } 2>/dev/null && echo 'RUN: at line 1': 'tessera-opt -allow-unregistered-dialect /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-3-msl-fusion/tests/tessera-ir/phase8/apple_gpu_lowering.mlir -tessera-lower-to-apple_gpu --allow-unregistered-dialect | FileCheck /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-3-msl-fusion/tests/tessera-ir/phase8/apple_gpu_lowering.mlir' >&2 && { set -x; } 2>/dev/null && { tessera-opt -allow-unregistered-dialect /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-3-msl-fusion/tests/tessera-ir/phase8/apple_gpu_lowering.mlir -tessera-lower-to-apple_gpu --allow-unregistered-dialect | FileCheck /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-3-msl-fusion/tests/tessera-ir/phase8/apple_gpu_lowering.mlir; }; } +set -o pipefail;set -x;{ { set +x; } 2>/dev/null && echo 'RUN: at line 1': 'tessera-opt -allow-unregistered-dialect /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-4-fp16-bf16/tests/tessera-ir/phase8/apple_gpu_lowering.mlir -tessera-lower-to-apple_gpu --allow-unregistered-dialect | FileCheck /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-4-fp16-bf16/tests/tessera-ir/phase8/apple_gpu_lowering.mlir' >&2 && { set -x; } 2>/dev/null && { tessera-opt -allow-unregistered-dialect /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-4-fp16-bf16/tests/tessera-ir/phase8/apple_gpu_lowering.mlir -tessera-lower-to-apple_gpu --allow-unregistered-dialect | FileCheck /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-4-fp16-bf16/tests/tessera-ir/phase8/apple_gpu_lowering.mlir; }; } diff --git a/tests/tessera-ir/phase8/Output/apple_gpu_matmul_dtypes.mlir.script b/tests/tessera-ir/phase8/Output/apple_gpu_matmul_dtypes.mlir.script new file mode 100644 index 000000000..a4f54a63e --- /dev/null +++ b/tests/tessera-ir/phase8/Output/apple_gpu_matmul_dtypes.mlir.script @@ -0,0 +1 @@ +set -o pipefail;set -x;{ { set +x; } 2>/dev/null && echo 'RUN: at line 1': 'tessera-opt -allow-unregistered-dialect /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-4-fp16-bf16/tests/tessera-ir/phase8/apple_gpu_matmul_dtypes.mlir --pass-pipeline='"'"'builtin.module(tessera-lower-to-apple_gpu-runtime)'"'"' --allow-unregistered-dialect | FileCheck /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-4-fp16-bf16/tests/tessera-ir/phase8/apple_gpu_matmul_dtypes.mlir' >&2 && { set -x; } 2>/dev/null && { tessera-opt -allow-unregistered-dialect /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-4-fp16-bf16/tests/tessera-ir/phase8/apple_gpu_matmul_dtypes.mlir --pass-pipeline='builtin.module(tessera-lower-to-apple_gpu-runtime)' --allow-unregistered-dialect | FileCheck /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-4-fp16-bf16/tests/tessera-ir/phase8/apple_gpu_matmul_dtypes.mlir; }; } diff --git a/tests/tessera-ir/phase8/Output/apple_gpu_matmul_softmax_fusion.mlir.script b/tests/tessera-ir/phase8/Output/apple_gpu_matmul_softmax_fusion.mlir.script index ef10cec87..efa2bbde9 100644 --- a/tests/tessera-ir/phase8/Output/apple_gpu_matmul_softmax_fusion.mlir.script +++ b/tests/tessera-ir/phase8/Output/apple_gpu_matmul_softmax_fusion.mlir.script @@ -1 +1 @@ -set -o pipefail;set -x;{ { set +x; } 2>/dev/null && echo 'RUN: at line 1': 'tessera-opt -allow-unregistered-dialect /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-3-msl-fusion/tests/tessera-ir/phase8/apple_gpu_matmul_softmax_fusion.mlir --pass-pipeline='"'"'builtin.module(tessera-lower-to-apple_gpu-runtime)'"'"' --allow-unregistered-dialect | FileCheck /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-3-msl-fusion/tests/tessera-ir/phase8/apple_gpu_matmul_softmax_fusion.mlir' >&2 && { set -x; } 2>/dev/null && { tessera-opt -allow-unregistered-dialect /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-3-msl-fusion/tests/tessera-ir/phase8/apple_gpu_matmul_softmax_fusion.mlir --pass-pipeline='builtin.module(tessera-lower-to-apple_gpu-runtime)' --allow-unregistered-dialect | FileCheck /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-3-msl-fusion/tests/tessera-ir/phase8/apple_gpu_matmul_softmax_fusion.mlir; }; } +set -o pipefail;set -x;{ { set +x; } 2>/dev/null && echo 'RUN: at line 1': 'tessera-opt -allow-unregistered-dialect /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-4-fp16-bf16/tests/tessera-ir/phase8/apple_gpu_matmul_softmax_fusion.mlir --pass-pipeline='"'"'builtin.module(tessera-lower-to-apple_gpu-runtime)'"'"' --allow-unregistered-dialect | FileCheck /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-4-fp16-bf16/tests/tessera-ir/phase8/apple_gpu_matmul_softmax_fusion.mlir' >&2 && { set -x; } 2>/dev/null && { tessera-opt -allow-unregistered-dialect /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-4-fp16-bf16/tests/tessera-ir/phase8/apple_gpu_matmul_softmax_fusion.mlir --pass-pipeline='builtin.module(tessera-lower-to-apple_gpu-runtime)' --allow-unregistered-dialect | FileCheck /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-4-fp16-bf16/tests/tessera-ir/phase8/apple_gpu_matmul_softmax_fusion.mlir; }; } diff --git a/tests/tessera-ir/phase8/Output/apple_gpu_msl.mlir.script b/tests/tessera-ir/phase8/Output/apple_gpu_msl.mlir.script index 2994cfcb3..a54a4e6b8 100644 --- a/tests/tessera-ir/phase8/Output/apple_gpu_msl.mlir.script +++ b/tests/tessera-ir/phase8/Output/apple_gpu_msl.mlir.script @@ -1 +1 @@ -set -o pipefail;set -x;{ { set +x; } 2>/dev/null && echo 'RUN: at line 1': 'tessera-opt -allow-unregistered-dialect /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-3-msl-fusion/tests/tessera-ir/phase8/apple_gpu_msl.mlir --allow-unregistered-dialect | FileCheck /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-3-msl-fusion/tests/tessera-ir/phase8/apple_gpu_msl.mlir' >&2 && { set -x; } 2>/dev/null && { tessera-opt -allow-unregistered-dialect /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-3-msl-fusion/tests/tessera-ir/phase8/apple_gpu_msl.mlir --allow-unregistered-dialect | FileCheck /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-3-msl-fusion/tests/tessera-ir/phase8/apple_gpu_msl.mlir; }; } +set -o pipefail;set -x;{ { set +x; } 2>/dev/null && echo 'RUN: at line 1': 'tessera-opt -allow-unregistered-dialect /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-4-fp16-bf16/tests/tessera-ir/phase8/apple_gpu_msl.mlir --allow-unregistered-dialect | FileCheck /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-4-fp16-bf16/tests/tessera-ir/phase8/apple_gpu_msl.mlir' >&2 && { set -x; } 2>/dev/null && { tessera-opt -allow-unregistered-dialect /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-4-fp16-bf16/tests/tessera-ir/phase8/apple_gpu_msl.mlir --allow-unregistered-dialect | FileCheck /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-4-fp16-bf16/tests/tessera-ir/phase8/apple_gpu_msl.mlir; }; } diff --git a/tests/tessera-ir/phase8/Output/apple_gpu_runtime.mlir.script b/tests/tessera-ir/phase8/Output/apple_gpu_runtime.mlir.script index d767be70f..a3cb3dc52 100644 --- a/tests/tessera-ir/phase8/Output/apple_gpu_runtime.mlir.script +++ b/tests/tessera-ir/phase8/Output/apple_gpu_runtime.mlir.script @@ -1 +1 @@ -set -o pipefail;set -x;{ { set +x; } 2>/dev/null && echo 'RUN: at line 1': 'tessera-opt -allow-unregistered-dialect /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-3-msl-fusion/tests/tessera-ir/phase8/apple_gpu_runtime.mlir --pass-pipeline='"'"'builtin.module(tessera-lower-to-apple_gpu-runtime)'"'"' --allow-unregistered-dialect | FileCheck /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-3-msl-fusion/tests/tessera-ir/phase8/apple_gpu_runtime.mlir' >&2 && { set -x; } 2>/dev/null && { tessera-opt -allow-unregistered-dialect /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-3-msl-fusion/tests/tessera-ir/phase8/apple_gpu_runtime.mlir --pass-pipeline='builtin.module(tessera-lower-to-apple_gpu-runtime)' --allow-unregistered-dialect | FileCheck /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-3-msl-fusion/tests/tessera-ir/phase8/apple_gpu_runtime.mlir; }; } +set -o pipefail;set -x;{ { set +x; } 2>/dev/null && echo 'RUN: at line 1': 'tessera-opt -allow-unregistered-dialect /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-4-fp16-bf16/tests/tessera-ir/phase8/apple_gpu_runtime.mlir --pass-pipeline='"'"'builtin.module(tessera-lower-to-apple_gpu-runtime)'"'"' --allow-unregistered-dialect | FileCheck /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-4-fp16-bf16/tests/tessera-ir/phase8/apple_gpu_runtime.mlir' >&2 && { set -x; } 2>/dev/null && { tessera-opt -allow-unregistered-dialect /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-4-fp16-bf16/tests/tessera-ir/phase8/apple_gpu_runtime.mlir --pass-pipeline='builtin.module(tessera-lower-to-apple_gpu-runtime)' --allow-unregistered-dialect | FileCheck /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-4-fp16-bf16/tests/tessera-ir/phase8/apple_gpu_runtime.mlir; }; } diff --git a/tests/tessera-ir/phase8/Output/apple_gpu_softmax_gelu.mlir.script b/tests/tessera-ir/phase8/Output/apple_gpu_softmax_gelu.mlir.script index 9c223f81c..a57989598 100644 --- a/tests/tessera-ir/phase8/Output/apple_gpu_softmax_gelu.mlir.script +++ b/tests/tessera-ir/phase8/Output/apple_gpu_softmax_gelu.mlir.script @@ -1 +1 @@ -set -o pipefail;set -x;{ { set +x; } 2>/dev/null && echo 'RUN: at line 1': 'tessera-opt -allow-unregistered-dialect /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-3-msl-fusion/tests/tessera-ir/phase8/apple_gpu_softmax_gelu.mlir --pass-pipeline='"'"'builtin.module(tessera-lower-to-apple_gpu-runtime)'"'"' --allow-unregistered-dialect | FileCheck /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-3-msl-fusion/tests/tessera-ir/phase8/apple_gpu_softmax_gelu.mlir' >&2 && { set -x; } 2>/dev/null && { tessera-opt -allow-unregistered-dialect /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-3-msl-fusion/tests/tessera-ir/phase8/apple_gpu_softmax_gelu.mlir --pass-pipeline='builtin.module(tessera-lower-to-apple_gpu-runtime)' --allow-unregistered-dialect | FileCheck /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-3-msl-fusion/tests/tessera-ir/phase8/apple_gpu_softmax_gelu.mlir; }; } +set -o pipefail;set -x;{ { set +x; } 2>/dev/null && echo 'RUN: at line 1': 'tessera-opt -allow-unregistered-dialect /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-4-fp16-bf16/tests/tessera-ir/phase8/apple_gpu_softmax_gelu.mlir --pass-pipeline='"'"'builtin.module(tessera-lower-to-apple_gpu-runtime)'"'"' --allow-unregistered-dialect | FileCheck /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-4-fp16-bf16/tests/tessera-ir/phase8/apple_gpu_softmax_gelu.mlir' >&2 && { set -x; } 2>/dev/null && { tessera-opt -allow-unregistered-dialect /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-4-fp16-bf16/tests/tessera-ir/phase8/apple_gpu_softmax_gelu.mlir --pass-pipeline='builtin.module(tessera-lower-to-apple_gpu-runtime)' --allow-unregistered-dialect | FileCheck /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-4-fp16-bf16/tests/tessera-ir/phase8/apple_gpu_softmax_gelu.mlir; }; } diff --git a/tests/tessera-ir/phase8/Output/target_ir_contracts.mlir.script b/tests/tessera-ir/phase8/Output/target_ir_contracts.mlir.script index 62de16adf..3ff8bef57 100644 --- a/tests/tessera-ir/phase8/Output/target_ir_contracts.mlir.script +++ b/tests/tessera-ir/phase8/Output/target_ir_contracts.mlir.script @@ -1 +1 @@ -set -o pipefail;set -x;{ { set +x; } 2>/dev/null && echo 'RUN: at line 1': 'FileCheck /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-3-msl-fusion/tests/tessera-ir/phase8/target_ir_contracts.mlir --check-prefixes=ROCM,METALIUM,APPLE-CPU,APPLE-GPU < /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-3-msl-fusion/tests/tessera-ir/phase8/target_ir_contracts.mlir' >&2 && { set -x; } 2>/dev/null && { FileCheck /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-3-msl-fusion/tests/tessera-ir/phase8/target_ir_contracts.mlir --check-prefixes=ROCM,METALIUM,APPLE-CPU,APPLE-GPU < /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-3-msl-fusion/tests/tessera-ir/phase8/target_ir_contracts.mlir; }; } +set -o pipefail;set -x;{ { set +x; } 2>/dev/null && echo 'RUN: at line 1': 'FileCheck /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-4-fp16-bf16/tests/tessera-ir/phase8/target_ir_contracts.mlir --check-prefixes=ROCM,METALIUM,APPLE-CPU,APPLE-GPU < /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-4-fp16-bf16/tests/tessera-ir/phase8/target_ir_contracts.mlir' >&2 && { set -x; } 2>/dev/null && { FileCheck /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-4-fp16-bf16/tests/tessera-ir/phase8/target_ir_contracts.mlir --check-prefixes=ROCM,METALIUM,APPLE-CPU,APPLE-GPU < /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-4-fp16-bf16/tests/tessera-ir/phase8/target_ir_contracts.mlir; }; } diff --git a/tests/tessera-ir/phase8/Output/tmem_tcgen05_contract.mlir.script b/tests/tessera-ir/phase8/Output/tmem_tcgen05_contract.mlir.script index a68f75eb6..c18936445 100644 --- a/tests/tessera-ir/phase8/Output/tmem_tcgen05_contract.mlir.script +++ b/tests/tessera-ir/phase8/Output/tmem_tcgen05_contract.mlir.script @@ -1 +1 @@ -set -o pipefail;set -x;{ { set +x; } 2>/dev/null && echo 'RUN: at line 1': 'FileCheck /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-3-msl-fusion/tests/tessera-ir/phase8/tmem_tcgen05_contract.mlir < /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-3-msl-fusion/tests/tessera-ir/phase8/tmem_tcgen05_contract.mlir' >&2 && { set -x; } 2>/dev/null && { FileCheck /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-3-msl-fusion/tests/tessera-ir/phase8/tmem_tcgen05_contract.mlir < /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-3-msl-fusion/tests/tessera-ir/phase8/tmem_tcgen05_contract.mlir; }; } +set -o pipefail;set -x;{ { set +x; } 2>/dev/null && echo 'RUN: at line 1': 'FileCheck /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-4-fp16-bf16/tests/tessera-ir/phase8/tmem_tcgen05_contract.mlir < /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-4-fp16-bf16/tests/tessera-ir/phase8/tmem_tcgen05_contract.mlir' >&2 && { set -x; } 2>/dev/null && { FileCheck /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-4-fp16-bf16/tests/tessera-ir/phase8/tmem_tcgen05_contract.mlir < /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-4-fp16-bf16/tests/tessera-ir/phase8/tmem_tcgen05_contract.mlir; }; } diff --git a/tests/tessera-ir/phase8/apple_gpu_matmul_dtypes.mlir b/tests/tessera-ir/phase8/apple_gpu_matmul_dtypes.mlir new file mode 100644 index 000000000..c435a98e7 --- /dev/null +++ b/tests/tessera-ir/phase8/apple_gpu_matmul_dtypes.mlir @@ -0,0 +1,49 @@ +// RUN: tessera-opt %s --pass-pipeline='builtin.module(tessera-lower-to-apple_gpu-runtime)' --allow-unregistered-dialect | FileCheck %s + +// Phase 8.4.4 — Apple GPU fp16 + bf16 matmul. Verifies that the runtime +// pipeline picks the right runtime symbol based on the matmul's input +// element type: +// f32 -> tessera_apple_gpu_mps_matmul_f32 (Phase 8.3, native MPS) +// f16 -> tessera_apple_gpu_mps_matmul_f16 (Phase 8.4.4, native MPS) +// bf16 -> tessera_apple_gpu_mps_matmul_bf16 (Phase 8.4.4, fp32 conversion) +// +// Each runtime symbol shares the same i64×3 + i32×3 ABI shape — the element +// type is encoded in the symbol name, not the signature. + +// CHECK-DAG: func.func private @tessera_apple_gpu_mps_matmul_f32(i64, i64, i64, i32, i32, i32) +// CHECK-DAG: func.func private @tessera_apple_gpu_mps_matmul_f16(i64, i64, i64, i32, i32, i32) +// CHECK-DAG: func.func private @tessera_apple_gpu_mps_matmul_bf16(i64, i64, i64, i32, i32, i32) + +func.func @gemm_f32(%A: tensor<8x16xf32>, %B: tensor<16x32xf32>) -> tensor<8x32xf32> { + // CHECK-LABEL: func.func @gemm_f32 + // CHECK: call @tessera_apple_gpu_mps_matmul_f32 + // CHECK-NOT: tessera.matmul + %C = "tessera.matmul"(%A, %B) : (tensor<8x16xf32>, tensor<16x32xf32>) -> tensor<8x32xf32> + return %C : tensor<8x32xf32> +} + +func.func @gemm_f16(%A: tensor<8x16xf16>, %B: tensor<16x32xf16>) -> tensor<8x32xf16> { + // CHECK-LABEL: func.func @gemm_f16 + // CHECK: call @tessera_apple_gpu_mps_matmul_f16 + // CHECK-NOT: tessera.matmul + %C = "tessera.matmul"(%A, %B) : (tensor<8x16xf16>, tensor<16x32xf16>) -> tensor<8x32xf16> + return %C : tensor<8x32xf16> +} + +func.func @gemm_bf16(%A: tensor<8x16xbf16>, %B: tensor<16x32xbf16>) -> tensor<8x32xbf16> { + // CHECK-LABEL: func.func @gemm_bf16 + // CHECK: call @tessera_apple_gpu_mps_matmul_bf16 + // CHECK-NOT: tessera.matmul + %C = "tessera.matmul"(%A, %B) : (tensor<8x16xbf16>, tensor<16x32xbf16>) -> tensor<8x32xbf16> + return %C : tensor<8x32xbf16> +} + +// Negative case: mismatched element types fall back to the artifact path. + +// CHECK-LABEL: func.func @gemm_mixed_dtypes +// CHECK: tessera.matmul + +func.func @gemm_mixed_dtypes(%A: tensor<8x16xf32>, %B: tensor<16x32xf16>) -> tensor<8x32xf32> { + %C = "tessera.matmul"(%A, %B) : (tensor<8x16xf32>, tensor<16x32xf16>) -> tensor<8x32xf32> + return %C : tensor<8x32xf32> +} diff --git a/tests/unit/test_apple_backend_roadmap.py b/tests/unit/test_apple_backend_roadmap.py index 1ebb160ff..ae4faa023 100644 --- a/tests/unit/test_apple_backend_roadmap.py +++ b/tests/unit/test_apple_backend_roadmap.py @@ -1269,6 +1269,179 @@ def test_apple_gpu_matmul_softmax_fusion_runtime_shim_correctness(tmp_path): np.testing.assert_allclose(O, ref, rtol=1e-4, atol=1e-5) +# ───────────────────────────────────────────────────────────────────────────── +# Phase 8.4.4: fp16 / bf16 matmul on apple_gpu (mirrors BNNS bf16 from CPU). +# +# Single rank-2 matmul programs now flip to metal_runtime regardless of dtype: +# f32 -> native MPSDataTypeFloat32 +# f16 -> native MPSDataTypeFloat16 +# bf16 -> fp32-conversion path (MPS doesn't support bf16 matrix +# descriptors as of macOS 14) +# Mixed-dtype operands fall back to the artifact-only path. +# ───────────────────────────────────────────────────────────────────────────── + + +def _bfloat16_or_skip(): + pytest.importorskip("ml_dtypes") + import ml_dtypes + return ml_dtypes.bfloat16 + + +def test_apple_gpu_matmul_f32_artifact_reports_metal_runtime(): + """Phase 8.4.4 — the compile-time artifact stays type-polymorphic + (defaults to f32 in the static Graph IR) because call-site dtypes are + only known when @jit functions are invoked. The runtime dispatcher + selects the matching MPS symbol by inspecting input array dtypes. + This test pins the compile-time contract; the runtime dtype dispatch + is exercised by the executes_through_mps tests below. + """ + + @ts.jit(target="apple_gpu") + def mm(A, B): + return ts.ops.matmul(A, B) + + artifact = mm.runtime_artifact() + assert artifact.metadata["execution_mode"] == "metal_runtime" + assert artifact.metadata["compiler_path"] == "apple_gpu_mps" + backend_text = mm.compile_bundle.artifact("backend").text + # The default f32 symbol is named in the artifact since the Graph IR + # operand types are f32 absent explicit type hints. Runtime dtype + # dispatch overrides this at launch time when inputs are f16/bf16. + assert "tessera_apple_gpu_mps_matmul_f32" in backend_text + + +def test_apple_gpu_matmul_f16_executes_through_mps(): + """End-to-end: @jit(target='apple_gpu') fp16 matmul. The runtime + dispatcher detects fp16 inputs and routes to MPSDataTypeFloat16. MPS + does fp16 internal accumulation, which can drift slightly from the + fp32-converted reference; assert at fp16 tolerance.""" + + @ts.jit(target="apple_gpu") + def mm(A, B): + return ts.ops.matmul(A, B) + + rng = np.random.RandomState(53) + for M, K, N in ((4, 8, 8), (8, 16, 32)): + A = rng.randn(M, K).astype(np.float16) + B = rng.randn(K, N).astype(np.float16) + out = mm(A, B) + assert out.dtype == np.float16 + assert out.shape == (M, N) + # Reference: convert to fp32, matmul, convert back. MPS does fp16 + # internal accumulation; modest rel tolerance covers the drift. + ref = (A.astype(np.float32) @ B.astype(np.float32)).astype(np.float16) + np.testing.assert_allclose( + out.astype(np.float32), ref.astype(np.float32), + rtol=5e-2, atol=5e-2, + ) + + +def test_apple_gpu_matmul_bf16_executes_through_fp32_conversion_path(): + """End-to-end: @jit(target='apple_gpu') bf16 matmul matches an + fp32-converted reference at bf16 tolerance. The runtime shim does the + fp32 conversion internally so the host sees a bf16 in/out ABI.""" + + bf16 = _bfloat16_or_skip() + + @ts.jit(target="apple_gpu") + def mm(A, B): + return ts.ops.matmul(A, B) + + rng = np.random.RandomState(59) + for M, K, N in ((4, 8, 8), (8, 16, 32)): + A = rng.randn(M, K).astype(bf16) + B = rng.randn(K, N).astype(bf16) + out = mm(A, B) + assert out.dtype == bf16 + assert out.shape == (M, N) + ref = (A.astype(np.float32) @ B.astype(np.float32)).astype(bf16) + # bf16 has ~7-bit mantissa; tile-order rounding can drift by ~2% on + # K=16. Same tolerance pattern as the apple_cpu BNNS bf16 test. + np.testing.assert_allclose( + out.astype(np.float32), ref.astype(np.float32), + rtol=2e-2, atol=2e-2, + ) + + +def test_apple_gpu_matmul_runtime_shim_exposes_f16_and_bf16_symbols(tmp_path): + """Compile the apple_gpu runtime shim from source and verify the C ABI + of the new fp16 + bf16 matmul symbols. On Darwin this exercises the + Metal/MPS path; on Linux/CI the portable reference fallback.""" + + cxx = shutil.which("c++") or shutil.which("clang++") or shutil.which("g++") + if cxx is None: + pytest.skip("C++ compiler is not available") + + backend = ROOT / "src/compiler/codegen/Tessera_Apple_Backend/runtime" + if sys.platform == "darwin": + source = backend / "apple_gpu_runtime.mm" + lib = tmp_path / "libtessera_apple_gpu_runtime.dylib" + cmd = [cxx, "-std=c++17", "-shared", "-fPIC", "-fobjc-arc", + "-x", "objective-c++", str(source), "-o", str(lib), + "-framework", "Foundation", + "-framework", "Metal", + "-framework", "MetalPerformanceShaders"] + else: + source = backend / "apple_gpu_runtime_stub.cpp" + lib = tmp_path / "libtessera_apple_gpu_runtime.so" + cmd = [cxx, "-std=c++17", "-shared", "-fPIC", str(source), "-o", str(lib)] + subprocess.run(cmd, check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True) + + runtime = ctypes.CDLL(str(lib)) + + # fp16 ABI test + gemm_f16 = runtime.tessera_apple_gpu_mps_matmul_f16 + gemm_f16.argtypes = [ + ctypes.POINTER(ctypes.c_uint16), + ctypes.POINTER(ctypes.c_uint16), + ctypes.POINTER(ctypes.c_uint16), + ctypes.c_int32, ctypes.c_int32, ctypes.c_int32, + ] + gemm_f16.restype = None + + A = np.array([[1, 2, 3], [4, 5, 6]], dtype=np.float16) + B = np.array([[7, 8], [9, 10], [11, 12]], dtype=np.float16) + C = np.zeros((2, 2), dtype=np.float16) + gemm_f16( + A.ctypes.data_as(ctypes.POINTER(ctypes.c_uint16)), + B.ctypes.data_as(ctypes.POINTER(ctypes.c_uint16)), + C.ctypes.data_as(ctypes.POINTER(ctypes.c_uint16)), + 2, 2, 3, + ) + np.testing.assert_array_equal( + C, (A.astype(np.float32) @ B.astype(np.float32)).astype(np.float16) + ) + + # bf16 ABI test (only when ml_dtypes is available) + try: + import ml_dtypes + bf16 = ml_dtypes.bfloat16 + except Exception: + return + + gemm_bf16 = runtime.tessera_apple_gpu_mps_matmul_bf16 + gemm_bf16.argtypes = [ + ctypes.POINTER(ctypes.c_uint16), + ctypes.POINTER(ctypes.c_uint16), + ctypes.POINTER(ctypes.c_uint16), + ctypes.c_int32, ctypes.c_int32, ctypes.c_int32, + ] + gemm_bf16.restype = None + + A_bf = np.array([[1, 2, 3], [4, 5, 6]], dtype=bf16) + B_bf = np.array([[7, 8], [9, 10], [11, 12]], dtype=bf16) + C_bf = np.zeros((2, 2), dtype=bf16) + gemm_bf16( + A_bf.view(np.uint16).ctypes.data_as(ctypes.POINTER(ctypes.c_uint16)), + B_bf.view(np.uint16).ctypes.data_as(ctypes.POINTER(ctypes.c_uint16)), + C_bf.view(np.uint16).ctypes.data_as(ctypes.POINTER(ctypes.c_uint16)), + 2, 2, 3, + ) + np.testing.assert_array_equal( + C_bf, (A_bf.astype(np.float32) @ B_bf.astype(np.float32)).astype(bf16) + ) + + def test_apple_cpu_bf16_disabled_when_ml_dtypes_missing(monkeypatch): """When ml_dtypes isn't installed the bf16 dtype probe returns None and the runtime falls through to numpy. Verified by stubbing the import to From 0639fe473c40759695018b1b781f00aac5bad936 Mon Sep 17 00:00:00 2001 From: Greg Stoner Date: Sat, 9 May 2026 06:13:34 -0500 Subject: [PATCH 2/2] =?UTF-8?q?Phase=208.4.4.1=20=E2=80=94=20fp16=20/=20bf?= =?UTF-8?q?16=20for=20simple=20MSL=20kernels=20(rope,=20softmax,=20gelu)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Extends each of the simple custom MSL kernels — rope, softmax, gelu — with fp16 + bf16 dtype variants. Mirrors the Phase 8.4.4 matmul pattern: - fp16: native MSL `half` kernel with `float` internal compute for accuracy. Apple Silicon GPUs run `half` at higher throughput than `float` for the elementwise math involved. - bf16: fp32-conversion path inside the runtime shim. MSL has no stable `bfloat` type pre-Metal-3.1, and the cost of decode + re-encode is negligible relative to the actual GPU compute. Six new C-ABI symbols (3 kernels × 2 dtypes), all with `uint16_t*` boundary types. `numpy.view(np.uint16)` and `ml_dtypes.bfloat16` are both byte-compatible. MLIR / runtime - Three pairs of new C symbols in apple_gpu_runtime.mm: * tessera_apple_gpu_rope_f16/bf16 * tessera_apple_gpu_softmax_f16/bf16 * tessera_apple_gpu_gelu_f16/bf16 - New native MSL kernels rope_f16, softmax_f16, gelu_f16 use `half` I/O with `float` internal compute (cos/sin, exp/sum, tanh). - bf16 paths convert each operand to fp32 at the boundary, run the existing f32 MSL kernel, encode back with round-to-nearest-even. - apple_gpu_runtime_stub.cpp gets matching reference fallbacks (all fp32-conversion) for non-Darwin builds. - RopeToAppleGPU / SoftmaxToAppleGPU / GeluToAppleGPU passes pick the runtime symbol by input element type. Same i64 + i32 ABI shape across all three dtypes per kernel; the element type is encoded in the symbol name only. Python - target_ir.py: three new fp16 MSL source constants (_APPLE_GPU_{ROPE,SOFTMAX,GELU}_MSL_SOURCE_F16) + sha256 cache_keys. bf16 reuses the f32 source (the runtime does the conversion); the IR-level marker just flips entry_point + cache_key + dtype attr. New helper _apple_gpu_kernel_msl_for_dtype maps (kernel, dtype) pairs to (msl_source, entry_point, cache_key, dtype_attr) tuples so the rope/softmax/gelu emission blocks share the dtype dispatch logic. - runtime.py: each of _apple_gpu_dispatch_{rope,softmax,gelu} now detects input array dtype at launch time and routes to the matching ctypes wrapper. fp16 + bf16 use uint16_t* ABI via numpy.view. Six new wrappers _apple_gpu_{rope,softmax,gelu}_{f16,bf16}. Loader gate now requires all six new symbols (forces rebuild after Phase 8.4.4.1). Tests - New lit fixture apple_gpu_msl_dtypes.mlir — verifies dtype-aware symbol selection for softmax + gelu with f32/f16/bf16 input tensors. rope is omitted because tessera.rope is not a registered dialect op (it's covered by Python tests instead). - Seven new unit tests in test_apple_backend_roadmap.py: * rope/softmax/gelu fp16 end-to-end (3 tests, native MSL path) * rope/softmax/gelu bf16 end-to-end (3 tests, fp32-conversion path, gated on ml_dtypes presence) * runtime shim ABI exposure for all 6 new symbols (1 test) Test bug fix in 8.4.4.1 contract tests - The bf16 input fixtures had `(rng.randn(...).astype(bf16)) * 0.5` — multiplying a bf16 array by a Python float promotes back to fp32 because numpy's bf16 (via ml_dtypes) doesn't special-case Python scalar mixing the way native fp16 does. Fixed by applying the multiplication in fp32 BEFORE the .astype(bf16) cast. Verified on Apple Silicon (LLVM/MLIR 21, Metal active): 2001 unit tests passing (1994 + 7 net new fp16/bf16 tests); 13/13 Phase 8 lit fixtures passing against the in-tree tessera-opt. fp16 paths match fp32-converted reference at rtol=5e-3; bf16 at rtol=2e-2 across rope/softmax/gelu. Co-Authored-By: Claude Opus 4.7 --- python/tessera/compiler/target_ir.py | 175 +++++++- python/tessera/runtime.py | 338 +++++++++++++--- .../Target/Apple/Lowering/GeluToAppleGPU.cpp | 28 +- .../Target/Apple/Lowering/RopeToAppleGPU.cpp | 36 +- .../Apple/Lowering/SoftmaxToAppleGPU.cpp | 30 +- .../runtime/apple_gpu_runtime.mm | 376 ++++++++++++++++++ .../runtime/apple_gpu_runtime_stub.cpp | 66 +++ tests/tessera-ir/.lit_test_times.txt | 25 +- .../Output/apple_cpu_lowering.mlir.script | 2 +- .../Output/apple_cpu_runtime.mlir.script | 2 +- .../apple_dialect_roundtrip.mlir.script | 2 +- .../Output/apple_gpu_flash_attn.mlir.script | 2 +- .../Output/apple_gpu_lowering.mlir.script | 2 +- .../apple_gpu_matmul_dtypes.mlir.script | 2 +- ...pple_gpu_matmul_softmax_fusion.mlir.script | 2 +- .../phase8/Output/apple_gpu_msl.mlir.script | 2 +- .../Output/apple_gpu_msl_dtypes.mlir.script | 1 + .../Output/apple_gpu_runtime.mlir.script | 2 +- .../Output/apple_gpu_softmax_gelu.mlir.script | 2 +- .../Output/target_ir_contracts.mlir.script | 2 +- .../Output/tmem_tcgen05_contract.mlir.script | 2 +- .../phase8/apple_gpu_msl_dtypes.mlir | 64 +++ tests/unit/test_apple_backend_roadmap.py | 178 +++++++++ 23 files changed, 1211 insertions(+), 130 deletions(-) create mode 100644 tests/tessera-ir/phase8/Output/apple_gpu_msl_dtypes.mlir.script create mode 100644 tests/tessera-ir/phase8/apple_gpu_msl_dtypes.mlir diff --git a/python/tessera/compiler/target_ir.py b/python/tessera/compiler/target_ir.py index 7cdc1c6e6..9de11d0c0 100644 --- a/python/tessera/compiler/target_ir.py +++ b/python/tessera/compiler/target_ir.py @@ -174,6 +174,135 @@ def _sha256_short(text: str) -> str: _APPLE_GPU_GELU_MSL_CACHE_KEY = _sha256_short(_APPLE_GPU_GELU_MSL_SOURCE) +# ───────────────────────────────────────────────────────────────────────────── +# Phase 8.4.4.1 — fp16 / bf16 MSL source constants for the simple kernels. +# fp16: native MSL `half` kernels with `float` internal compute for accuracy. +# bf16: emit a `bf16` cache_key marker; the runtime shim dispatches the +# fp32-conversion path internally (no native MSL bf16 source). +# ───────────────────────────────────────────────────────────────────────────── + +_APPLE_GPU_ROPE_MSL_SOURCE_F16 = ( + "#include \n" + "using namespace metal;\n" + "kernel void rope_f16(\n" + " device const half* x [[buffer(0)]],\n" + " device const half* theta [[buffer(1)]],\n" + " device half* out [[buffer(2)]],\n" + " constant int& M [[buffer(3)]],\n" + " constant int& K [[buffer(4)]],\n" + " uint2 gid [[thread_position_in_grid]])\n" + "{\n" + " if (gid.x >= (uint)(K / 2) || gid.y >= (uint)M) return;\n" + " int row = (int)gid.y;\n" + " int pair = (int)gid.x;\n" + " int idx_even = row * K + pair * 2;\n" + " int idx_odd = idx_even + 1;\n" + " float xe = float(x[idx_even]);\n" + " float xo = float(x[idx_odd]);\n" + " float c = cos(float(theta[idx_even]));\n" + " float s = sin(float(theta[idx_even]));\n" + " out[idx_even] = half(xe * c - xo * s);\n" + " out[idx_odd] = half(xe * s + xo * c);\n" + "}\n" +) +_APPLE_GPU_ROPE_MSL_CACHE_KEY_F16 = _sha256_short(_APPLE_GPU_ROPE_MSL_SOURCE_F16) + +_APPLE_GPU_SOFTMAX_MSL_SOURCE_F16 = ( + "#include \n" + "using namespace metal;\n" + "kernel void softmax_f16(\n" + " device const half* x [[buffer(0)]],\n" + " device half* out [[buffer(1)]],\n" + " constant int& M [[buffer(2)]],\n" + " constant int& K [[buffer(3)]],\n" + " uint gid [[thread_position_in_grid]])\n" + "{\n" + " if (gid >= (uint)M) return;\n" + " int row = (int)gid;\n" + " int row_off = row * K;\n" + " float row_max = -INFINITY;\n" + " for (int j = 0; j < K; ++j) row_max = max(row_max, float(x[row_off + j]));\n" + " float denom = 0.0f;\n" + " for (int j = 0; j < K; ++j) {\n" + " float e = exp(float(x[row_off + j]) - row_max);\n" + " out[row_off + j] = half(e);\n" + " denom += e;\n" + " }\n" + " float inv = 1.0f / denom;\n" + " for (int j = 0; j < K; ++j) out[row_off + j] = half(float(out[row_off + j]) * inv);\n" + "}\n" +) +_APPLE_GPU_SOFTMAX_MSL_CACHE_KEY_F16 = _sha256_short(_APPLE_GPU_SOFTMAX_MSL_SOURCE_F16) + +_APPLE_GPU_GELU_MSL_SOURCE_F16 = ( + "#include \n" + "using namespace metal;\n" + "kernel void gelu_f16(\n" + " device const half* x [[buffer(0)]],\n" + " device half* out [[buffer(1)]],\n" + " constant int& N [[buffer(2)]],\n" + " uint gid [[thread_position_in_grid]])\n" + "{\n" + " if (gid >= (uint)N) return;\n" + " float v = float(x[gid]);\n" + " float t = 0.7978845608028654f * (v + 0.044715f * v * v * v);\n" + " out[gid] = half(0.5f * v * (1.0f + tanh(t)));\n" + "}\n" +) +_APPLE_GPU_GELU_MSL_CACHE_KEY_F16 = _sha256_short(_APPLE_GPU_GELU_MSL_SOURCE_F16) + + +# bf16 doesn't get native MSL kernels in Phase 8.4.4.1 — the runtime shim +# does fp32 conversion at the boundary then dispatches the existing f32 +# kernel. The IR-level marker reuses the f32 source text but flips the +# entry_point + cache_key + dtype attr so downstream tooling can tell the +# difference. Same shape as Phase 8.4.4 bf16 matmul (no native MPS bf16). +def _apple_gpu_kernel_msl_for_dtype( + kernel: str, dtype: str +) -> tuple[str, str, str, str]: + """Return (msl_source, entry_point, cache_key, dtype) for the given + (kernel, dtype) pair. dtype is one of {"f32", "f16", "bf16"}.""" + + if kernel == "rope": + if dtype == "f16": + return (_APPLE_GPU_ROPE_MSL_SOURCE_F16, "rope_f16", + _APPLE_GPU_ROPE_MSL_CACHE_KEY_F16, "f16") + if dtype == "bf16": + return (_APPLE_GPU_ROPE_MSL_SOURCE, "rope_bf16", + _APPLE_GPU_ROPE_MSL_CACHE_KEY, "bf16") + return (_APPLE_GPU_ROPE_MSL_SOURCE, "rope_f32", + _APPLE_GPU_ROPE_MSL_CACHE_KEY, "f32") + if kernel == "softmax": + if dtype == "f16": + return (_APPLE_GPU_SOFTMAX_MSL_SOURCE_F16, "softmax_f16", + _APPLE_GPU_SOFTMAX_MSL_CACHE_KEY_F16, "f16") + if dtype == "bf16": + return (_APPLE_GPU_SOFTMAX_MSL_SOURCE, "softmax_bf16", + _APPLE_GPU_SOFTMAX_MSL_CACHE_KEY, "bf16") + return (_APPLE_GPU_SOFTMAX_MSL_SOURCE, "softmax_f32", + _APPLE_GPU_SOFTMAX_MSL_CACHE_KEY, "f32") + if kernel == "gelu": + if dtype == "f16": + return (_APPLE_GPU_GELU_MSL_SOURCE_F16, "gelu_f16", + _APPLE_GPU_GELU_MSL_CACHE_KEY_F16, "f16") + if dtype == "bf16": + return (_APPLE_GPU_GELU_MSL_SOURCE, "gelu_bf16", + _APPLE_GPU_GELU_MSL_CACHE_KEY, "bf16") + return (_APPLE_GPU_GELU_MSL_SOURCE, "gelu_f32", + _APPLE_GPU_GELU_MSL_CACHE_KEY, "f32") + raise ValueError(f"unknown apple_gpu kernel: {kernel!r}") + + +def _apple_gpu_dtype_from_op(op) -> str: + """Pick a runtime-supported dtype from a tile op's attrs. Defaults to f32 + when the attr is absent or not in the supported envelope.""" + + raw = str(op.attrs.get("dtype", "f32")) + if raw in {"f32", "f16", "bf16"}: + return raw + return "f32" + + # Phase 8.4.3 — embedded MSL source for the fused matmul -> softmax(axis=-1) # kernel. One thread per output row: computes the row of A@B into a stack # array, then row-wise softmax in place. Cap N <= 256 to keep the stack @@ -823,21 +952,27 @@ def _lower_apple_gpu_op(op: TileOp, *, mps_runtime: bool = False) -> list[Target "execution_mode": "metal_runtime", }), ] - # Phase 8.4 custom MSL path: a single-rope module is lowered to + # Phase 8.4 + 8.4.4.1 custom MSL path: a single-rope module is lowered to # msl_kernel + mps_dispatch carrying the MSL source as a StringAttr. The # RopeToAppleGPU pass and the apple_gpu_runtime.mm shim consume this op. + # dtype attr (Phase 8.4.4.1) picks between f32 / f16 / bf16 source + + # entry point. if mps_runtime and ( op.op_name in {"tile.rotary_pair", "tile.rope"} or source == "tessera.rope" ): + dtype = _apple_gpu_dtype_from_op(op) + msl_source, entry_point, cache_key, dtype_attr = ( + _apple_gpu_kernel_msl_for_dtype("rope", dtype) + ) return [ TargetOp("tessera_apple.gpu.msl_kernel", { **base, "framework": "Metal", - "dtype": "f32", - "entry_point": "rope_f32", - "msl_source": _APPLE_GPU_ROPE_MSL_SOURCE, - "cache_key": _APPLE_GPU_ROPE_MSL_CACHE_KEY, + "dtype": dtype_attr, + "entry_point": entry_point, + "msl_source": msl_source, + "cache_key": cache_key, "grid": "tokens_pairs", "threadgroup": "32x?", }), @@ -871,16 +1006,21 @@ def _lower_apple_gpu_op(op: TileOp, *, mps_runtime: bool = False) -> list[Target "execution_mode": "metal_runtime", }), ] - # Phase 8.4.2 custom MSL path: single-softmax (axis=-1) module. + # Phase 8.4.2 + 8.4.4.1 custom MSL path: single-softmax (axis=-1) module + # with dtype-aware MSL source selection. if mps_runtime and source in {"tessera.softmax", "tessera.softmax_safe"}: + dtype = _apple_gpu_dtype_from_op(op) + msl_source, entry_point, cache_key, dtype_attr = ( + _apple_gpu_kernel_msl_for_dtype("softmax", dtype) + ) return [ TargetOp("tessera_apple.gpu.msl_kernel", { **base, "framework": "Metal", - "dtype": "f32", - "entry_point": "softmax_f32", - "msl_source": _APPLE_GPU_SOFTMAX_MSL_SOURCE, - "cache_key": _APPLE_GPU_SOFTMAX_MSL_CACHE_KEY, + "dtype": dtype_attr, + "entry_point": entry_point, + "msl_source": msl_source, + "cache_key": cache_key, "grid": "rows", "threadgroup": "?x1x1", }), @@ -891,16 +1031,21 @@ def _lower_apple_gpu_op(op: TileOp, *, mps_runtime: bool = False) -> list[Target "execution_mode": "metal_runtime", }), ] - # Phase 8.4.2 custom MSL path: single-gelu (elementwise) module. + # Phase 8.4.2 + 8.4.4.1 custom MSL path: single-gelu (elementwise) module + # with dtype-aware MSL source selection. if mps_runtime and source == "tessera.gelu": + dtype = _apple_gpu_dtype_from_op(op) + msl_source, entry_point, cache_key, dtype_attr = ( + _apple_gpu_kernel_msl_for_dtype("gelu", dtype) + ) return [ TargetOp("tessera_apple.gpu.msl_kernel", { **base, "framework": "Metal", - "dtype": "f32", - "entry_point": "gelu_f32", - "msl_source": _APPLE_GPU_GELU_MSL_SOURCE, - "cache_key": _APPLE_GPU_GELU_MSL_CACHE_KEY, + "dtype": dtype_attr, + "entry_point": entry_point, + "msl_source": msl_source, + "cache_key": cache_key, "grid": "elements", "threadgroup": "?x1x1", }), diff --git a/python/tessera/runtime.py b/python/tessera/runtime.py index 275f3a96d..8c073869b 100644 --- a/python/tessera/runtime.py +++ b/python/tessera/runtime.py @@ -2101,9 +2101,12 @@ def _apple_gpu_mps_matmul_bf16() -> Any: def _apple_gpu_dispatch_rope(op_name: str, operands: list[Any], np: Any) -> Any: - """Phase 8.4: dispatch a single rank-2 f32 rope through the apple_gpu - runtime shim's custom MSL kernel. Inputs outside the supported envelope - fall back to the numpy reference path used by the default `cpu` target. + """Phase 8.4 + 8.4.4.1: dispatch a single rank-2 rope through the apple_gpu + runtime shim's custom MSL kernel. Picks the runtime symbol by element type: + - f32: native MSL kernel (Phase 8.4) + - f16: native MSL `half` kernel (Phase 8.4.4.1) + - bf16: fp32-conversion path inside the shim (Phase 8.4.4.1) + Inputs outside the supported envelope fall back to the numpy reference. """ if len(operands) != 2: @@ -2111,31 +2114,69 @@ def _apple_gpu_dispatch_rope(op_name: str, operands: list[Any], np: Any) -> Any: x = np.asarray(operands[0]) theta = np.asarray(operands[1]) - rank2_fast_path = ( - x.dtype == np.float32 - and theta.dtype == np.float32 - and x.ndim == 2 - and theta.ndim == 2 - and x.shape == theta.shape - and x.shape[1] % 2 == 0 - ) - if not rank2_fast_path: + if ( + x.ndim != 2 or theta.ndim != 2 + or x.shape != theta.shape + or x.shape[1] % 2 != 0 + or x.dtype != theta.dtype + ): return _runtime_rope(np, x, theta) - if not x.flags.c_contiguous: - x = np.ascontiguousarray(x, dtype=np.float32) - if not theta.flags.c_contiguous: - theta = np.ascontiguousarray(theta, dtype=np.float32) - - out = np.zeros(x.shape, dtype=np.float32) - rope = _apple_gpu_rope_f32() - rope( - x.ctypes.data_as(ctypes.POINTER(ctypes.c_float)), - theta.ctypes.data_as(ctypes.POINTER(ctypes.c_float)), - out.ctypes.data_as(ctypes.POINTER(ctypes.c_float)), - ctypes.c_int32(x.shape[0]), - ctypes.c_int32(x.shape[1]), - ) - return out + + if x.dtype == np.float32: + if not x.flags.c_contiguous: + x = np.ascontiguousarray(x, dtype=np.float32) + if not theta.flags.c_contiguous: + theta = np.ascontiguousarray(theta, dtype=np.float32) + out = np.zeros(x.shape, dtype=np.float32) + rope = _apple_gpu_rope_f32() + rope( + x.ctypes.data_as(ctypes.POINTER(ctypes.c_float)), + theta.ctypes.data_as(ctypes.POINTER(ctypes.c_float)), + out.ctypes.data_as(ctypes.POINTER(ctypes.c_float)), + ctypes.c_int32(x.shape[0]), + ctypes.c_int32(x.shape[1]), + ) + return out + + if x.dtype == np.float16: + if not x.flags.c_contiguous: + x = np.ascontiguousarray(x, dtype=np.float16) + if not theta.flags.c_contiguous: + theta = np.ascontiguousarray(theta, dtype=np.float16) + out = np.zeros(x.shape, dtype=np.float16) + rope_f16 = _apple_gpu_rope_f16() + if rope_f16 is not None: + rope_f16( + x.view(np.uint16).ctypes.data_as(ctypes.POINTER(ctypes.c_uint16)), + theta.view(np.uint16).ctypes.data_as(ctypes.POINTER(ctypes.c_uint16)), + out.view(np.uint16).ctypes.data_as(ctypes.POINTER(ctypes.c_uint16)), + ctypes.c_int32(x.shape[0]), + ctypes.c_int32(x.shape[1]), + ) + return out + # Older runtime build without f16 — fall back via fp32. + return _runtime_rope(np, x.astype(np.float32), theta.astype(np.float32)).astype(np.float16) + + bf16_dtype = _bfloat16_dtype() + if bf16_dtype is not None and x.dtype == bf16_dtype: + if not x.flags.c_contiguous: + x = np.ascontiguousarray(x, dtype=bf16_dtype) + if not theta.flags.c_contiguous: + theta = np.ascontiguousarray(theta, dtype=bf16_dtype) + out = np.zeros(x.shape, dtype=bf16_dtype) + rope_bf16 = _apple_gpu_rope_bf16() + if rope_bf16 is not None: + rope_bf16( + x.view(np.uint16).ctypes.data_as(ctypes.POINTER(ctypes.c_uint16)), + theta.view(np.uint16).ctypes.data_as(ctypes.POINTER(ctypes.c_uint16)), + out.view(np.uint16).ctypes.data_as(ctypes.POINTER(ctypes.c_uint16)), + ctypes.c_int32(x.shape[0]), + ctypes.c_int32(x.shape[1]), + ) + return out + return _runtime_rope(np, x.astype(np.float32), theta.astype(np.float32)).astype(bf16_dtype) + + return _runtime_rope(np, x, theta) def _apple_gpu_rope_f32() -> Any: @@ -2152,6 +2193,38 @@ def _apple_gpu_rope_f32() -> Any: return sym +def _apple_gpu_rope_f16() -> Any: + runtime = _load_apple_gpu_runtime() + sym = getattr(runtime, "tessera_apple_gpu_rope_f16", None) + if sym is None: + return None + sym.argtypes = [ + ctypes.POINTER(ctypes.c_uint16), + ctypes.POINTER(ctypes.c_uint16), + ctypes.POINTER(ctypes.c_uint16), + ctypes.c_int32, + ctypes.c_int32, + ] + sym.restype = None + return sym + + +def _apple_gpu_rope_bf16() -> Any: + runtime = _load_apple_gpu_runtime() + sym = getattr(runtime, "tessera_apple_gpu_rope_bf16", None) + if sym is None: + return None + sym.argtypes = [ + ctypes.POINTER(ctypes.c_uint16), + ctypes.POINTER(ctypes.c_uint16), + ctypes.POINTER(ctypes.c_uint16), + ctypes.c_int32, + ctypes.c_int32, + ] + sym.restype = None + return sym + + def _apple_gpu_dispatch_flash_attn(op_name: str, operands: list[Any], kwargs: Mapping[str, Any], np: Any) -> Any: """Phase 8.4.1: dispatch a single rank-3 f32 flash-attention forward @@ -2228,34 +2301,71 @@ def _apple_gpu_flash_attn_f32() -> Any: def _apple_gpu_dispatch_softmax(op_name: str, operands: list[Any], kwargs: Mapping[str, Any], np: Any) -> Any: - """Phase 8.4.2: dispatch a single rank-2 f32 softmax (axis=-1) through - the apple_gpu runtime shim's custom MSL kernel. Inputs outside the - supported envelope (rank, dtype, axis) fall back to the numpy reference. + """Phase 8.4.2 + 8.4.4.1: dispatch a single rank-2 softmax (axis=-1) + through the apple_gpu runtime shim's custom MSL kernel. Picks symbol by + element type (f32, f16, bf16). Inputs outside the supported envelope + fall back to the numpy reference. """ if len(operands) < 1: raise ValueError(f"{op_name!r} requires one operand") x = np.asarray(operands[0]) axis = int(kwargs.get("axis", -1)) - rank2_fast_path = ( - x.dtype == np.float32 and x.ndim == 2 and (axis == -1 or axis == 1) - ) - if not rank2_fast_path: - # Numpy reference — same op as the default `cpu` target. + if x.ndim != 2 or (axis != -1 and axis != 1): e = np.exp(x - np.max(x, axis=axis, keepdims=True)) return e / np.sum(e, axis=axis, keepdims=True) - if not x.flags.c_contiguous: - x = np.ascontiguousarray(x, dtype=np.float32) - M, K = x.shape - out = np.zeros((M, K), dtype=np.float32) - softmax = _apple_gpu_softmax_f32() - softmax( - x.ctypes.data_as(ctypes.POINTER(ctypes.c_float)), - out.ctypes.data_as(ctypes.POINTER(ctypes.c_float)), - ctypes.c_int32(M), - ctypes.c_int32(K), - ) - return out + + if x.dtype == np.float32: + if not x.flags.c_contiguous: + x = np.ascontiguousarray(x, dtype=np.float32) + M, K = x.shape + out = np.zeros((M, K), dtype=np.float32) + softmax = _apple_gpu_softmax_f32() + softmax( + x.ctypes.data_as(ctypes.POINTER(ctypes.c_float)), + out.ctypes.data_as(ctypes.POINTER(ctypes.c_float)), + ctypes.c_int32(M), + ctypes.c_int32(K), + ) + return out + + if x.dtype == np.float16: + if not x.flags.c_contiguous: + x = np.ascontiguousarray(x, dtype=np.float16) + M, K = x.shape + out = np.zeros((M, K), dtype=np.float16) + softmax_f16 = _apple_gpu_softmax_f16() + if softmax_f16 is not None: + softmax_f16( + x.view(np.uint16).ctypes.data_as(ctypes.POINTER(ctypes.c_uint16)), + out.view(np.uint16).ctypes.data_as(ctypes.POINTER(ctypes.c_uint16)), + ctypes.c_int32(M), + ctypes.c_int32(K), + ) + return out + e = np.exp(x.astype(np.float32) - np.max(x.astype(np.float32), axis=-1, keepdims=True)) + return (e / np.sum(e, axis=-1, keepdims=True)).astype(np.float16) + + bf16_dtype = _bfloat16_dtype() + if bf16_dtype is not None and x.dtype == bf16_dtype: + if not x.flags.c_contiguous: + x = np.ascontiguousarray(x, dtype=bf16_dtype) + M, K = x.shape + out = np.zeros((M, K), dtype=bf16_dtype) + softmax_bf16 = _apple_gpu_softmax_bf16() + if softmax_bf16 is not None: + softmax_bf16( + x.view(np.uint16).ctypes.data_as(ctypes.POINTER(ctypes.c_uint16)), + out.view(np.uint16).ctypes.data_as(ctypes.POINTER(ctypes.c_uint16)), + ctypes.c_int32(M), + ctypes.c_int32(K), + ) + return out + e = np.exp(x.astype(np.float32) - np.max(x.astype(np.float32), axis=-1, keepdims=True)) + return (e / np.sum(e, axis=-1, keepdims=True)).astype(bf16_dtype) + + e = np.exp(x - np.max(x, axis=axis, keepdims=True)) + return e / np.sum(e, axis=axis, keepdims=True) def _apple_gpu_softmax_f32() -> Any: @@ -2271,29 +2381,94 @@ def _apple_gpu_softmax_f32() -> Any: return sym +def _apple_gpu_softmax_f16() -> Any: + runtime = _load_apple_gpu_runtime() + sym = getattr(runtime, "tessera_apple_gpu_softmax_f16", None) + if sym is None: + return None + sym.argtypes = [ + ctypes.POINTER(ctypes.c_uint16), + ctypes.POINTER(ctypes.c_uint16), + ctypes.c_int32, + ctypes.c_int32, + ] + sym.restype = None + return sym + + +def _apple_gpu_softmax_bf16() -> Any: + runtime = _load_apple_gpu_runtime() + sym = getattr(runtime, "tessera_apple_gpu_softmax_bf16", None) + if sym is None: + return None + sym.argtypes = [ + ctypes.POINTER(ctypes.c_uint16), + ctypes.POINTER(ctypes.c_uint16), + ctypes.c_int32, + ctypes.c_int32, + ] + sym.restype = None + return sym + + def _apple_gpu_dispatch_gelu(op_name: str, operands: list[Any], np: Any) -> Any: - """Phase 8.4.2: dispatch a single rank-2 f32 gelu through the apple_gpu - runtime shim's custom MSL kernel. Tanh-approximation matching the numpy - reference. Inputs outside the supported envelope fall back to numpy. - """ + """Phase 8.4.2 + 8.4.4.1: dispatch a single rank-2 gelu through the + apple_gpu runtime shim's custom MSL kernel. Picks symbol by element type + (f32, f16, bf16). Tanh-approximation matching the numpy reference.""" if len(operands) < 1: raise ValueError(f"{op_name!r} requires one operand") x = np.asarray(operands[0]) - rank2_fast_path = x.dtype == np.float32 and x.ndim == 2 - if not rank2_fast_path: + if x.ndim != 2: return 0.5 * x * (1.0 + np.tanh(np.sqrt(2.0 / np.pi) * (x + 0.044715 * x**3))) - if not x.flags.c_contiguous: - x = np.ascontiguousarray(x, dtype=np.float32) - M, K = x.shape - out = np.zeros((M, K), dtype=np.float32) - gelu = _apple_gpu_gelu_f32() - gelu( - x.ctypes.data_as(ctypes.POINTER(ctypes.c_float)), - out.ctypes.data_as(ctypes.POINTER(ctypes.c_float)), - ctypes.c_int32(M * K), - ) - return out + + if x.dtype == np.float32: + if not x.flags.c_contiguous: + x = np.ascontiguousarray(x, dtype=np.float32) + M, K = x.shape + out = np.zeros((M, K), dtype=np.float32) + gelu = _apple_gpu_gelu_f32() + gelu( + x.ctypes.data_as(ctypes.POINTER(ctypes.c_float)), + out.ctypes.data_as(ctypes.POINTER(ctypes.c_float)), + ctypes.c_int32(M * K), + ) + return out + + if x.dtype == np.float16: + if not x.flags.c_contiguous: + x = np.ascontiguousarray(x, dtype=np.float16) + M, K = x.shape + out = np.zeros((M, K), dtype=np.float16) + gelu_f16 = _apple_gpu_gelu_f16() + if gelu_f16 is not None: + gelu_f16( + x.view(np.uint16).ctypes.data_as(ctypes.POINTER(ctypes.c_uint16)), + out.view(np.uint16).ctypes.data_as(ctypes.POINTER(ctypes.c_uint16)), + ctypes.c_int32(M * K), + ) + return out + x32 = x.astype(np.float32) + return (0.5 * x32 * (1.0 + np.tanh(np.sqrt(2.0 / np.pi) * (x32 + 0.044715 * x32**3)))).astype(np.float16) + + bf16_dtype = _bfloat16_dtype() + if bf16_dtype is not None and x.dtype == bf16_dtype: + if not x.flags.c_contiguous: + x = np.ascontiguousarray(x, dtype=bf16_dtype) + M, K = x.shape + out = np.zeros((M, K), dtype=bf16_dtype) + gelu_bf16 = _apple_gpu_gelu_bf16() + if gelu_bf16 is not None: + gelu_bf16( + x.view(np.uint16).ctypes.data_as(ctypes.POINTER(ctypes.c_uint16)), + out.view(np.uint16).ctypes.data_as(ctypes.POINTER(ctypes.c_uint16)), + ctypes.c_int32(M * K), + ) + return out + x32 = x.astype(np.float32) + return (0.5 * x32 * (1.0 + np.tanh(np.sqrt(2.0 / np.pi) * (x32 + 0.044715 * x32**3)))).astype(bf16_dtype) + + return 0.5 * x * (1.0 + np.tanh(np.sqrt(2.0 / np.pi) * (x + 0.044715 * x**3))) def _apple_gpu_gelu_f32() -> Any: @@ -2308,6 +2483,34 @@ def _apple_gpu_gelu_f32() -> Any: return sym +def _apple_gpu_gelu_f16() -> Any: + runtime = _load_apple_gpu_runtime() + sym = getattr(runtime, "tessera_apple_gpu_gelu_f16", None) + if sym is None: + return None + sym.argtypes = [ + ctypes.POINTER(ctypes.c_uint16), + ctypes.POINTER(ctypes.c_uint16), + ctypes.c_int32, + ] + sym.restype = None + return sym + + +def _apple_gpu_gelu_bf16() -> Any: + runtime = _load_apple_gpu_runtime() + sym = getattr(runtime, "tessera_apple_gpu_gelu_bf16", None) + if sym is None: + return None + sym.argtypes = [ + ctypes.POINTER(ctypes.c_uint16), + ctypes.POINTER(ctypes.c_uint16), + ctypes.c_int32, + ] + sym.restype = None + return sym + + def _apple_gpu_dispatch_matmul_softmax(operands: list[Any], np: Any) -> Any: """Phase 8.4.3 — dispatch a fused matmul -> softmax(axis=-1) chain through the apple_gpu runtime shim's purpose-built MSL kernel. Inputs @@ -2403,6 +2606,13 @@ def _load_apple_gpu_runtime() -> ctypes.CDLL: # builds lack them; falling through forces a rebuild. getattr(lib, "tessera_apple_gpu_mps_matmul_f16") getattr(lib, "tessera_apple_gpu_mps_matmul_bf16") + # Phase 8.4.4.1 — fp16/bf16 for the simple MSL kernels. + getattr(lib, "tessera_apple_gpu_rope_f16") + getattr(lib, "tessera_apple_gpu_rope_bf16") + getattr(lib, "tessera_apple_gpu_softmax_f16") + getattr(lib, "tessera_apple_gpu_softmax_bf16") + getattr(lib, "tessera_apple_gpu_gelu_f16") + getattr(lib, "tessera_apple_gpu_gelu_bf16") _apple_gpu_runtime = lib return _apple_gpu_runtime except (OSError, AttributeError): diff --git a/src/compiler/codegen/Tessera_Apple_Backend/lib/Target/Apple/Lowering/GeluToAppleGPU.cpp b/src/compiler/codegen/Tessera_Apple_Backend/lib/Target/Apple/Lowering/GeluToAppleGPU.cpp index 67b61bb06..11da73717 100644 --- a/src/compiler/codegen/Tessera_Apple_Backend/lib/Target/Apple/Lowering/GeluToAppleGPU.cpp +++ b/src/compiler/codegen/Tessera_Apple_Backend/lib/Target/Apple/Lowering/GeluToAppleGPU.cpp @@ -35,6 +35,8 @@ namespace apple { namespace { constexpr llvm::StringLiteral kGeluF32Symbol = "tessera_apple_gpu_gelu_f32"; +constexpr llvm::StringLiteral kGeluF16Symbol = "tessera_apple_gpu_gelu_f16"; +constexpr llvm::StringLiteral kGeluBF16Symbol = "tessera_apple_gpu_gelu_bf16"; static func::FuncOp ensureExternalDecl(ModuleOp mod, StringRef name, FunctionType fnTy) { @@ -67,9 +69,18 @@ struct LowerGeluToAppleGPU : public RewritePattern { if (!xTy || xTy.getRank() != 2) return rewriter.notifyMatchFailure( op, "AppleGPU gelu MSL path is rank-2 only in Phase 8.4.2"); - if (!xTy.getElementType().isF32()) + Type xElem = xTy.getElementType(); + StringRef symbol; + if (xElem.isF32()) { + symbol = kGeluF32Symbol; + } else if (xElem.isF16()) { + symbol = kGeluF16Symbol; + } else if (xElem.isBF16()) { + symbol = kGeluBF16Symbol; + } else { return rewriter.notifyMatchFailure( - op, "AppleGPU gelu MSL path is f32-only in Phase 8.4.2"); + op, "AppleGPU gelu MSL path supports f32, f16, and bf16 in Phase 8.4.4.1"); + } if (xTy.isDynamicDim(0) || xTy.isDynamicDim(1)) return rewriter.notifyMatchFailure( op, "AppleGPU gelu MSL path requires static shapes"); @@ -84,9 +95,8 @@ struct LowerGeluToAppleGPU : public RewritePattern { Type i64Ty = rewriter.getI64Type(); Type i32Ty = rewriter.getI32Type(); - Type f32Ty = rewriter.getF32Type(); - auto memTy = MemRefType::get({M, K}, f32Ty); + auto memTy = MemRefType::get({M, K}, xElem); Value xPtr = extractPtr(rewriter, loc, x, memTy); auto outAlloc = rewriter.create(loc, memTy); Value outPtr; @@ -100,12 +110,12 @@ struct LowerGeluToAppleGPU : public RewritePattern { FunctionType fnTy = FunctionType::get(ctx, {i64Ty, i64Ty, i32Ty}, {}); - ensureExternalDecl(mod, kGeluF32Symbol, fnTy); + ensureExternalDecl(mod, symbol, fnTy); rewriter.create( - loc, kGeluF32Symbol, TypeRange{}, ValueRange{xPtr, outPtr, Nv}); + loc, symbol, TypeRange{}, ValueRange{xPtr, outPtr, Nv}); - auto outTensorTy = RankedTensorType::get({M, K}, f32Ty); + auto outTensorTy = RankedTensorType::get({M, K}, xElem); Value result = rewriter.create(loc, outTensorTy, outAlloc); rewriter.replaceOp(op, result); @@ -121,8 +131,8 @@ struct LowerGeluToAppleGPUPass return "tessera-gelu-to-apple_gpu"; } StringRef getDescription() const override { - return "Lower tessera.gelu (rank-2, f32) to Apple GPU runtime calls " - "(custom MSL kernel)"; + return "Lower tessera.gelu (rank-2, f32/f16/bf16) to Apple GPU runtime " + "calls (custom MSL kernel)"; } void getDependentDialects(DialectRegistry ®istry) const override { diff --git a/src/compiler/codegen/Tessera_Apple_Backend/lib/Target/Apple/Lowering/RopeToAppleGPU.cpp b/src/compiler/codegen/Tessera_Apple_Backend/lib/Target/Apple/Lowering/RopeToAppleGPU.cpp index e13a23afa..fa1f1af3f 100644 --- a/src/compiler/codegen/Tessera_Apple_Backend/lib/Target/Apple/Lowering/RopeToAppleGPU.cpp +++ b/src/compiler/codegen/Tessera_Apple_Backend/lib/Target/Apple/Lowering/RopeToAppleGPU.cpp @@ -39,6 +39,8 @@ namespace apple { namespace { constexpr llvm::StringLiteral kRopeF32Symbol = "tessera_apple_gpu_rope_f32"; +constexpr llvm::StringLiteral kRopeF16Symbol = "tessera_apple_gpu_rope_f16"; +constexpr llvm::StringLiteral kRopeBF16Symbol = "tessera_apple_gpu_rope_bf16"; static func::FuncOp ensureExternalDecl(ModuleOp mod, StringRef name, FunctionType fnTy) { @@ -74,9 +76,26 @@ struct LowerRopeToAppleGPU : public RewritePattern { if (!xTy || !thetaTy || xTy.getRank() != 2 || thetaTy.getRank() != 2) return failure(); - if (!xTy.getElementType().isF32() || !thetaTy.getElementType().isF32()) + Type xElem = xTy.getElementType(); + Type thetaElem = thetaTy.getElementType(); + if (xElem != thetaElem) return rewriter.notifyMatchFailure( - op, "AppleGPU rope MSL path is f32-only in Phase 8.4"); + op, "AppleGPU rope MSL path requires matching x/theta dtypes"); + + // Phase 8.4.4.1 — pick the runtime symbol by dtype. Same i64×3 + i32×2 + // ABI shape across all three; the element type is encoded in the + // symbol name, not the signature. + StringRef symbol; + if (xElem.isF32()) { + symbol = kRopeF32Symbol; + } else if (xElem.isF16()) { + symbol = kRopeF16Symbol; + } else if (xElem.isBF16()) { + symbol = kRopeBF16Symbol; + } else { + return rewriter.notifyMatchFailure( + op, "AppleGPU rope MSL path supports f32, f16, and bf16 in Phase 8.4.4.1"); + } if (xTy.isDynamicDim(0) || xTy.isDynamicDim(1) || thetaTy.isDynamicDim(0) || thetaTy.isDynamicDim(1)) @@ -99,9 +118,8 @@ struct LowerRopeToAppleGPU : public RewritePattern { Type i64Ty = rewriter.getI64Type(); Type i32Ty = rewriter.getI32Type(); - Type f32Ty = rewriter.getF32Type(); - auto memTy = MemRefType::get({M, K}, f32Ty); + auto memTy = MemRefType::get({M, K}, xElem); Value xPtr = extractPtr(rewriter, loc, x, memTy); Value thetaPtr = extractPtr(rewriter, loc, theta, memTy); auto outAlloc = rewriter.create(loc, memTy); @@ -117,13 +135,13 @@ struct LowerRopeToAppleGPU : public RewritePattern { FunctionType ropeFnTy = FunctionType::get(ctx, {i64Ty, i64Ty, i64Ty, i32Ty, i32Ty}, {}); - ensureExternalDecl(mod, kRopeF32Symbol, ropeFnTy); + ensureExternalDecl(mod, symbol, ropeFnTy); rewriter.create( - loc, kRopeF32Symbol, TypeRange{}, + loc, symbol, TypeRange{}, ValueRange{xPtr, thetaPtr, outPtr, Mv, Kv}); - auto outTensorTy = RankedTensorType::get({M, K}, f32Ty); + auto outTensorTy = RankedTensorType::get({M, K}, xElem); Value result = rewriter.create(loc, outTensorTy, outAlloc); rewriter.replaceOp(op, result); @@ -140,8 +158,8 @@ struct LowerRopeToAppleGPUPass return "tessera-rope-to-apple_gpu"; } StringRef getDescription() const override { - return "Lower tessera.rope (rank-2, f32) to Apple GPU runtime calls " - "(custom MSL kernel)"; + return "Lower tessera.rope (rank-2, f32/f16/bf16) to Apple GPU runtime " + "calls (custom MSL kernel)"; } void getDependentDialects(DialectRegistry ®istry) const override { diff --git a/src/compiler/codegen/Tessera_Apple_Backend/lib/Target/Apple/Lowering/SoftmaxToAppleGPU.cpp b/src/compiler/codegen/Tessera_Apple_Backend/lib/Target/Apple/Lowering/SoftmaxToAppleGPU.cpp index 0051ce805..065cd72d1 100644 --- a/src/compiler/codegen/Tessera_Apple_Backend/lib/Target/Apple/Lowering/SoftmaxToAppleGPU.cpp +++ b/src/compiler/codegen/Tessera_Apple_Backend/lib/Target/Apple/Lowering/SoftmaxToAppleGPU.cpp @@ -37,6 +37,10 @@ namespace { constexpr llvm::StringLiteral kSoftmaxF32Symbol = "tessera_apple_gpu_softmax_f32"; +constexpr llvm::StringLiteral kSoftmaxF16Symbol = + "tessera_apple_gpu_softmax_f16"; +constexpr llvm::StringLiteral kSoftmaxBF16Symbol = + "tessera_apple_gpu_softmax_bf16"; static func::FuncOp ensureExternalDecl(ModuleOp mod, StringRef name, FunctionType fnTy) { @@ -69,9 +73,18 @@ struct LowerSoftmaxToAppleGPU : public RewritePattern { if (!xTy || xTy.getRank() != 2) return rewriter.notifyMatchFailure( op, "AppleGPU softmax MSL path is rank-2 only in Phase 8.4.2"); - if (!xTy.getElementType().isF32()) + Type xElem = xTy.getElementType(); + StringRef symbol; + if (xElem.isF32()) { + symbol = kSoftmaxF32Symbol; + } else if (xElem.isF16()) { + symbol = kSoftmaxF16Symbol; + } else if (xElem.isBF16()) { + symbol = kSoftmaxBF16Symbol; + } else { return rewriter.notifyMatchFailure( - op, "AppleGPU softmax MSL path is f32-only in Phase 8.4.2"); + op, "AppleGPU softmax MSL path supports f32, f16, and bf16 in Phase 8.4.4.1"); + } if (xTy.isDynamicDim(0) || xTy.isDynamicDim(1)) return rewriter.notifyMatchFailure( op, "AppleGPU softmax MSL path requires static shapes"); @@ -94,9 +107,8 @@ struct LowerSoftmaxToAppleGPU : public RewritePattern { Type i64Ty = rewriter.getI64Type(); Type i32Ty = rewriter.getI32Type(); - Type f32Ty = rewriter.getF32Type(); - auto memTy = MemRefType::get({M, K}, f32Ty); + auto memTy = MemRefType::get({M, K}, xElem); Value xPtr = extractPtr(rewriter, loc, x, memTy); auto outAlloc = rewriter.create(loc, memTy); Value outPtr; @@ -111,13 +123,13 @@ struct LowerSoftmaxToAppleGPU : public RewritePattern { FunctionType fnTy = FunctionType::get(ctx, {i64Ty, i64Ty, i32Ty, i32Ty}, {}); - ensureExternalDecl(mod, kSoftmaxF32Symbol, fnTy); + ensureExternalDecl(mod, symbol, fnTy); rewriter.create( - loc, kSoftmaxF32Symbol, TypeRange{}, + loc, symbol, TypeRange{}, ValueRange{xPtr, outPtr, Mv, Kv}); - auto outTensorTy = RankedTensorType::get({M, K}, f32Ty); + auto outTensorTy = RankedTensorType::get({M, K}, xElem); Value result = rewriter.create(loc, outTensorTy, outAlloc); rewriter.replaceOp(op, result); @@ -134,8 +146,8 @@ struct LowerSoftmaxToAppleGPUPass return "tessera-softmax-to-apple_gpu"; } StringRef getDescription() const override { - return "Lower tessera.softmax (rank-2, f32, axis=-1) to Apple GPU " - "runtime calls (custom MSL kernel)"; + return "Lower tessera.softmax (rank-2, f32/f16/bf16, axis=-1) to Apple " + "GPU runtime calls (custom MSL kernel)"; } void getDependentDialects(DialectRegistry ®istry) const override { diff --git a/src/compiler/codegen/Tessera_Apple_Backend/runtime/apple_gpu_runtime.mm b/src/compiler/codegen/Tessera_Apple_Backend/runtime/apple_gpu_runtime.mm index e50305bca..5d8e58934 100644 --- a/src/compiler/codegen/Tessera_Apple_Backend/runtime/apple_gpu_runtime.mm +++ b/src/compiler/codegen/Tessera_Apple_Backend/runtime/apple_gpu_runtime.mm @@ -569,6 +569,148 @@ inline void reference_rope_f32(const float* X, const float* Theta, float* Out, reference_rope_f32(X, Theta, Out, M, K); } +//===---------------------------------------------------------------------===// +// Phase 8.4.4.1 — fp16 + bf16 rope variants. +// +// fp16: native MSL `half` kernel. Compute is in `float` for accuracy +// (cos/sin); load/store as `half`. Apple Silicon GPUs run this +// at higher throughput than the fp32 variant. +// bf16: same fp32-conversion pattern as Phase 8.4.4 bf16 matmul. Decode +// bf16 bit-pattern, run the fp32 reference, encode back. +//===---------------------------------------------------------------------===// + +namespace { + +bool dispatch_rope_msl_f16(MetalDeviceContext &ctx, const uint16_t* X, + const uint16_t* Theta, uint16_t* Out, + int32_t M, int32_t K) { + static NSString *const kRopeSourceF16 = @R"MSL( +#include +using namespace metal; + +kernel void rope_f16( + device const half* x [[buffer(0)]], + device const half* theta [[buffer(1)]], + device half* out [[buffer(2)]], + constant int& M [[buffer(3)]], + constant int& K [[buffer(4)]], + uint2 gid [[thread_position_in_grid]]) +{ + if (gid.x >= (uint)(K / 2) || gid.y >= (uint)M) return; + int row = (int)gid.y; + int pair = (int)gid.x; + int idx_even = row * K + pair * 2; + int idx_odd = idx_even + 1; + float xe = float(x[idx_even]); + float xo = float(x[idx_odd]); + float c = cos(float(theta[idx_even])); + float s = sin(float(theta[idx_even])); + out[idx_even] = half(xe * c - xo * s); + out[idx_odd] = half(xe * s + xo * c); +} +)MSL"; + + @autoreleasepool { + id pso = + compile_msl_kernel(ctx, kRopeSourceF16, @"rope_f16"); + if (!pso) return false; + + NSUInteger byteCount = sizeof(uint16_t) * static_cast(M) * + static_cast(K); + id bufX = [ctx.device newBufferWithBytes:X + length:byteCount + options:MTLResourceStorageModeShared]; + id bufT = [ctx.device newBufferWithBytes:Theta + length:byteCount + options:MTLResourceStorageModeShared]; + id bufO = [ctx.device newBufferWithLength:byteCount + options:MTLResourceStorageModeShared]; + if (!bufX || !bufT || !bufO) return false; + + id cb = [ctx.queue commandBuffer]; + id enc = [cb computeCommandEncoder]; + [enc setComputePipelineState:pso]; + [enc setBuffer:bufX offset:0 atIndex:0]; + [enc setBuffer:bufT offset:0 atIndex:1]; + [enc setBuffer:bufO offset:0 atIndex:2]; + [enc setBytes:&M length:sizeof(int32_t) atIndex:3]; + [enc setBytes:&K length:sizeof(int32_t) atIndex:4]; + + NSUInteger half_k = static_cast(K / 2); + MTLSize grid = MTLSizeMake(half_k, static_cast(M), 1); + NSUInteger tg_x = std::min(half_k, 32); + NSUInteger tg_y = std::min(static_cast(M), + pso.maxTotalThreadsPerThreadgroup / + std::max(tg_x, 1)); + if (tg_y == 0) tg_y = 1; + MTLSize tg = MTLSizeMake(tg_x, tg_y, 1); + [enc dispatchThreads:grid threadsPerThreadgroup:tg]; + [enc endEncoding]; + [cb commit]; + [cb waitUntilCompleted]; + + if (cb.status != MTLCommandBufferStatusCompleted) return false; + std::memcpy(Out, [bufO contents], byteCount); + return true; + } +} + +inline void reference_rope_f16_via_fp32(const uint16_t* X, const uint16_t* Theta, + uint16_t* Out, int32_t M, int32_t K) { + std::vector Xf(static_cast(M) * K); + std::vector Tf(static_cast(M) * K); + std::vector Of(static_cast(M) * K); + for (std::size_t i = 0; i < Xf.size(); ++i) Xf[i] = half_to_float_gpu(X[i]); + for (std::size_t i = 0; i < Tf.size(); ++i) Tf[i] = half_to_float_gpu(Theta[i]); + reference_rope_f32(Xf.data(), Tf.data(), Of.data(), M, K); + for (std::size_t i = 0; i < Of.size(); ++i) Out[i] = float_to_half_gpu(Of[i]); +} + +inline void reference_rope_bf16_via_fp32(const uint16_t* X, const uint16_t* Theta, + uint16_t* Out, int32_t M, int32_t K) { + std::vector Xf(static_cast(M) * K); + std::vector Tf(static_cast(M) * K); + std::vector Of(static_cast(M) * K); + for (std::size_t i = 0; i < Xf.size(); ++i) Xf[i] = bfloat16_to_float_gpu(X[i]); + for (std::size_t i = 0; i < Tf.size(); ++i) Tf[i] = bfloat16_to_float_gpu(Theta[i]); + reference_rope_f32(Xf.data(), Tf.data(), Of.data(), M, K); + for (std::size_t i = 0; i < Of.size(); ++i) Out[i] = float_to_bfloat16_gpu(Of[i]); +} + +bool dispatch_rope_bf16_via_fp32(MetalDeviceContext &ctx, const uint16_t* X, + const uint16_t* Theta, uint16_t* Out, + int32_t M, int32_t K) { + std::vector Xf(static_cast(M) * K); + std::vector Tf(static_cast(M) * K); + std::vector Of(static_cast(M) * K); + for (std::size_t i = 0; i < Xf.size(); ++i) Xf[i] = bfloat16_to_float_gpu(X[i]); + for (std::size_t i = 0; i < Tf.size(); ++i) Tf[i] = bfloat16_to_float_gpu(Theta[i]); + if (!dispatch_rope_msl(ctx, Xf.data(), Tf.data(), Of.data(), M, K)) + return false; + for (std::size_t i = 0; i < Of.size(); ++i) Out[i] = float_to_bfloat16_gpu(Of[i]); + return true; +} + +} // namespace + +extern "C" void tessera_apple_gpu_rope_f16(const uint16_t* X, + const uint16_t* Theta, + uint16_t* Out, + int32_t M, int32_t K) { + MetalDeviceContext &ctx = deviceContext(); + if (ctx.ok && dispatch_rope_msl_f16(ctx, X, Theta, Out, M, K)) return; + reference_rope_f16_via_fp32(X, Theta, Out, M, K); +} + +extern "C" void tessera_apple_gpu_rope_bf16(const uint16_t* X, + const uint16_t* Theta, + uint16_t* Out, + int32_t M, int32_t K) { + MetalDeviceContext &ctx = deviceContext(); + if (ctx.ok && dispatch_rope_bf16_via_fp32(ctx, X, Theta, Out, M, K)) return; + reference_rope_bf16_via_fp32(X, Theta, Out, M, K); +} + extern "C" int32_t tessera_apple_gpu_runtime_msl_cache_size(void) { MetalDeviceContext &ctx = deviceContext(); if (!ctx.ok) return -1; @@ -918,6 +1060,133 @@ inline void reference_softmax_f32(const float* X, float* Out, int32_t M, reference_softmax_f32(X, Out, M, K); } +//===---------------------------------------------------------------------===// +// Phase 8.4.4.1 — fp16 + bf16 softmax variants. +// fp16: native MSL `half` kernel; per-row reduction in `float` for numerical +// stability (small per-row range matters for softmax denom accuracy). +// bf16: fp32-conversion path. +//===---------------------------------------------------------------------===// + +namespace { + +bool dispatch_softmax_msl_f16(MetalDeviceContext &ctx, const uint16_t* X, + uint16_t* Out, int32_t M, int32_t K) { + static NSString *const kSoftmaxSourceF16 = @R"MSL( +#include +using namespace metal; + +kernel void softmax_f16( + device const half* x [[buffer(0)]], + device half* out [[buffer(1)]], + constant int& M [[buffer(2)]], + constant int& K [[buffer(3)]], + uint gid [[thread_position_in_grid]]) +{ + if (gid >= (uint)M) return; + int row = (int)gid; + int row_off = row * K; + + float row_max = -INFINITY; + for (int j = 0; j < K; ++j) { + row_max = max(row_max, float(x[row_off + j])); + } + float denom = 0.0f; + // Pass 2: store exp values in `out` as fp16 — slight precision loss vs + // f32 reference but matches the MSL native fp16 throughput contract. + for (int j = 0; j < K; ++j) { + float e = exp(float(x[row_off + j]) - row_max); + out[row_off + j] = half(e); + denom += e; + } + float inv = 1.0f / denom; + for (int j = 0; j < K; ++j) { + out[row_off + j] = half(float(out[row_off + j]) * inv); + } +} +)MSL"; + + @autoreleasepool { + id pso = + compile_msl_kernel(ctx, kSoftmaxSourceF16, @"softmax_f16"); + if (!pso) return false; + + NSUInteger byteCount = sizeof(uint16_t) * static_cast(M) * + static_cast(K); + id bufX = [ctx.device newBufferWithBytes:X + length:byteCount + options:MTLResourceStorageModeShared]; + id bufO = [ctx.device newBufferWithLength:byteCount + options:MTLResourceStorageModeShared]; + if (!bufX || !bufO) return false; + + id cb = [ctx.queue commandBuffer]; + id enc = [cb computeCommandEncoder]; + [enc setComputePipelineState:pso]; + [enc setBuffer:bufX offset:0 atIndex:0]; + [enc setBuffer:bufO offset:0 atIndex:1]; + [enc setBytes:&M length:sizeof(int32_t) atIndex:2]; + [enc setBytes:&K length:sizeof(int32_t) atIndex:3]; + + MTLSize grid = MTLSizeMake(static_cast(M), 1, 1); + NSUInteger tg_x = std::min(static_cast(M), + pso.maxTotalThreadsPerThreadgroup); + if (tg_x == 0) tg_x = 1; + MTLSize tg = MTLSizeMake(tg_x, 1, 1); + [enc dispatchThreads:grid threadsPerThreadgroup:tg]; + [enc endEncoding]; + [cb commit]; + [cb waitUntilCompleted]; + + if (cb.status != MTLCommandBufferStatusCompleted) return false; + std::memcpy(Out, [bufO contents], byteCount); + return true; + } +} + +inline void reference_softmax_f16_via_fp32(const uint16_t* X, uint16_t* Out, + int32_t M, int32_t K) { + std::vector Xf(static_cast(M) * K); + std::vector Of(static_cast(M) * K); + for (std::size_t i = 0; i < Xf.size(); ++i) Xf[i] = half_to_float_gpu(X[i]); + reference_softmax_f32(Xf.data(), Of.data(), M, K); + for (std::size_t i = 0; i < Of.size(); ++i) Out[i] = float_to_half_gpu(Of[i]); +} + +inline void reference_softmax_bf16_via_fp32(const uint16_t* X, uint16_t* Out, + int32_t M, int32_t K) { + std::vector Xf(static_cast(M) * K); + std::vector Of(static_cast(M) * K); + for (std::size_t i = 0; i < Xf.size(); ++i) Xf[i] = bfloat16_to_float_gpu(X[i]); + reference_softmax_f32(Xf.data(), Of.data(), M, K); + for (std::size_t i = 0; i < Of.size(); ++i) Out[i] = float_to_bfloat16_gpu(Of[i]); +} + +bool dispatch_softmax_bf16_via_fp32(MetalDeviceContext &ctx, const uint16_t* X, + uint16_t* Out, int32_t M, int32_t K) { + std::vector Xf(static_cast(M) * K); + std::vector Of(static_cast(M) * K); + for (std::size_t i = 0; i < Xf.size(); ++i) Xf[i] = bfloat16_to_float_gpu(X[i]); + if (!dispatch_softmax_msl(ctx, Xf.data(), Of.data(), M, K)) return false; + for (std::size_t i = 0; i < Of.size(); ++i) Out[i] = float_to_bfloat16_gpu(Of[i]); + return true; +} + +} // namespace + +extern "C" void tessera_apple_gpu_softmax_f16(const uint16_t* X, uint16_t* Out, + int32_t M, int32_t K) { + MetalDeviceContext &ctx = deviceContext(); + if (ctx.ok && dispatch_softmax_msl_f16(ctx, X, Out, M, K)) return; + reference_softmax_f16_via_fp32(X, Out, M, K); +} + +extern "C" void tessera_apple_gpu_softmax_bf16(const uint16_t* X, uint16_t* Out, + int32_t M, int32_t K) { + MetalDeviceContext &ctx = deviceContext(); + if (ctx.ok && dispatch_softmax_bf16_via_fp32(ctx, X, Out, M, K)) return; + reference_softmax_bf16_via_fp32(X, Out, M, K); +} + //===---------------------------------------------------------------------===// // Phase 8.4.2 — GeLU (elementwise, f32) // @@ -1005,6 +1274,113 @@ inline void reference_gelu_f32(const float* X, float* Out, int32_t N) { reference_gelu_f32(X, Out, N); } +//===---------------------------------------------------------------------===// +// Phase 8.4.4.1 — fp16 + bf16 gelu variants. +// fp16: native MSL `half` kernel; tanh/cube in `float` for accuracy. +// bf16: fp32-conversion path. +//===---------------------------------------------------------------------===// + +namespace { + +bool dispatch_gelu_msl_f16(MetalDeviceContext &ctx, const uint16_t* X, + uint16_t* Out, int32_t N) { + static NSString *const kGeluSourceF16 = @R"MSL( +#include +using namespace metal; + +kernel void gelu_f16( + device const half* x [[buffer(0)]], + device half* out [[buffer(1)]], + constant int& N [[buffer(2)]], + uint gid [[thread_position_in_grid]]) +{ + if (gid >= (uint)N) return; + float v = float(x[gid]); + float t = 0.7978845608028654f * (v + 0.044715f * v * v * v); + out[gid] = half(0.5f * v * (1.0f + tanh(t))); +} +)MSL"; + + @autoreleasepool { + id pso = + compile_msl_kernel(ctx, kGeluSourceF16, @"gelu_f16"); + if (!pso) return false; + + NSUInteger byteCount = sizeof(uint16_t) * static_cast(N); + id bufX = [ctx.device newBufferWithBytes:X + length:byteCount + options:MTLResourceStorageModeShared]; + id bufO = [ctx.device newBufferWithLength:byteCount + options:MTLResourceStorageModeShared]; + if (!bufX || !bufO) return false; + + id cb = [ctx.queue commandBuffer]; + id enc = [cb computeCommandEncoder]; + [enc setComputePipelineState:pso]; + [enc setBuffer:bufX offset:0 atIndex:0]; + [enc setBuffer:bufO offset:0 atIndex:1]; + [enc setBytes:&N length:sizeof(int32_t) atIndex:2]; + + MTLSize grid = MTLSizeMake(static_cast(N), 1, 1); + NSUInteger tg_x = std::min(static_cast(N), + pso.maxTotalThreadsPerThreadgroup); + if (tg_x == 0) tg_x = 1; + MTLSize tg = MTLSizeMake(tg_x, 1, 1); + [enc dispatchThreads:grid threadsPerThreadgroup:tg]; + [enc endEncoding]; + [cb commit]; + [cb waitUntilCompleted]; + + if (cb.status != MTLCommandBufferStatusCompleted) return false; + std::memcpy(Out, [bufO contents], byteCount); + return true; + } +} + +inline void reference_gelu_f16_via_fp32(const uint16_t* X, uint16_t* Out, + int32_t N) { + std::vector Xf(static_cast(N)); + std::vector Of(static_cast(N)); + for (int32_t i = 0; i < N; ++i) Xf[i] = half_to_float_gpu(X[i]); + reference_gelu_f32(Xf.data(), Of.data(), N); + for (int32_t i = 0; i < N; ++i) Out[i] = float_to_half_gpu(Of[i]); +} + +inline void reference_gelu_bf16_via_fp32(const uint16_t* X, uint16_t* Out, + int32_t N) { + std::vector Xf(static_cast(N)); + std::vector Of(static_cast(N)); + for (int32_t i = 0; i < N; ++i) Xf[i] = bfloat16_to_float_gpu(X[i]); + reference_gelu_f32(Xf.data(), Of.data(), N); + for (int32_t i = 0; i < N; ++i) Out[i] = float_to_bfloat16_gpu(Of[i]); +} + +bool dispatch_gelu_bf16_via_fp32(MetalDeviceContext &ctx, const uint16_t* X, + uint16_t* Out, int32_t N) { + std::vector Xf(static_cast(N)); + std::vector Of(static_cast(N)); + for (int32_t i = 0; i < N; ++i) Xf[i] = bfloat16_to_float_gpu(X[i]); + if (!dispatch_gelu_msl(ctx, Xf.data(), Of.data(), N)) return false; + for (int32_t i = 0; i < N; ++i) Out[i] = float_to_bfloat16_gpu(Of[i]); + return true; +} + +} // namespace + +extern "C" void tessera_apple_gpu_gelu_f16(const uint16_t* X, uint16_t* Out, + int32_t N) { + MetalDeviceContext &ctx = deviceContext(); + if (ctx.ok && dispatch_gelu_msl_f16(ctx, X, Out, N)) return; + reference_gelu_f16_via_fp32(X, Out, N); +} + +extern "C" void tessera_apple_gpu_gelu_bf16(const uint16_t* X, uint16_t* Out, + int32_t N) { + MetalDeviceContext &ctx = deviceContext(); + if (ctx.ok && dispatch_gelu_bf16_via_fp32(ctx, X, Out, N)) return; + reference_gelu_bf16_via_fp32(X, Out, N); +} + //===---------------------------------------------------------------------===// // Phase 8.4.3 — Fused matmul → softmax (rank-2, f32, axis=-1) // diff --git a/src/compiler/codegen/Tessera_Apple_Backend/runtime/apple_gpu_runtime_stub.cpp b/src/compiler/codegen/Tessera_Apple_Backend/runtime/apple_gpu_runtime_stub.cpp index 2953c0ada..4ec3f504b 100644 --- a/src/compiler/codegen/Tessera_Apple_Backend/runtime/apple_gpu_runtime_stub.cpp +++ b/src/compiler/codegen/Tessera_Apple_Backend/runtime/apple_gpu_runtime_stub.cpp @@ -269,6 +269,72 @@ extern "C" void tessera_apple_gpu_gelu_f32(const float* X, float* Out, reference_gelu_f32(X, Out, N); } +// Phase 8.4.4.1 — fp16 / bf16 stubs for rope / softmax / gelu. All routes +// fp32-convert at the boundary, run the existing f32 reference, convert +// back. Same shape as the matmul fp16/bf16 stub from Phase 8.4.4. + +extern "C" void tessera_apple_gpu_rope_f16(const uint16_t* X, + const uint16_t* Theta, + uint16_t* Out, + int32_t M, int32_t K) { + std::vector Xf(static_cast(M) * K); + std::vector Tf(static_cast(M) * K); + std::vector Of(static_cast(M) * K); + for (std::size_t i = 0; i < Xf.size(); ++i) Xf[i] = half_to_float_stub(X[i]); + for (std::size_t i = 0; i < Tf.size(); ++i) Tf[i] = half_to_float_stub(Theta[i]); + reference_rope_f32(Xf.data(), Tf.data(), Of.data(), M, K); + for (std::size_t i = 0; i < Of.size(); ++i) Out[i] = float_to_half_stub(Of[i]); +} + +extern "C" void tessera_apple_gpu_rope_bf16(const uint16_t* X, + const uint16_t* Theta, + uint16_t* Out, + int32_t M, int32_t K) { + std::vector Xf(static_cast(M) * K); + std::vector Tf(static_cast(M) * K); + std::vector Of(static_cast(M) * K); + for (std::size_t i = 0; i < Xf.size(); ++i) Xf[i] = bfloat16_to_float_stub(X[i]); + for (std::size_t i = 0; i < Tf.size(); ++i) Tf[i] = bfloat16_to_float_stub(Theta[i]); + reference_rope_f32(Xf.data(), Tf.data(), Of.data(), M, K); + for (std::size_t i = 0; i < Of.size(); ++i) Out[i] = float_to_bfloat16_stub(Of[i]); +} + +extern "C" void tessera_apple_gpu_softmax_f16(const uint16_t* X, uint16_t* Out, + int32_t M, int32_t K) { + std::vector Xf(static_cast(M) * K); + std::vector Of(static_cast(M) * K); + for (std::size_t i = 0; i < Xf.size(); ++i) Xf[i] = half_to_float_stub(X[i]); + reference_softmax_f32(Xf.data(), Of.data(), M, K); + for (std::size_t i = 0; i < Of.size(); ++i) Out[i] = float_to_half_stub(Of[i]); +} + +extern "C" void tessera_apple_gpu_softmax_bf16(const uint16_t* X, uint16_t* Out, + int32_t M, int32_t K) { + std::vector Xf(static_cast(M) * K); + std::vector Of(static_cast(M) * K); + for (std::size_t i = 0; i < Xf.size(); ++i) Xf[i] = bfloat16_to_float_stub(X[i]); + reference_softmax_f32(Xf.data(), Of.data(), M, K); + for (std::size_t i = 0; i < Of.size(); ++i) Out[i] = float_to_bfloat16_stub(Of[i]); +} + +extern "C" void tessera_apple_gpu_gelu_f16(const uint16_t* X, uint16_t* Out, + int32_t N) { + std::vector Xf(static_cast(N)); + std::vector Of(static_cast(N)); + for (int32_t i = 0; i < N; ++i) Xf[i] = half_to_float_stub(X[i]); + reference_gelu_f32(Xf.data(), Of.data(), N); + for (int32_t i = 0; i < N; ++i) Out[i] = float_to_half_stub(Of[i]); +} + +extern "C" void tessera_apple_gpu_gelu_bf16(const uint16_t* X, uint16_t* Out, + int32_t N) { + std::vector Xf(static_cast(N)); + std::vector Of(static_cast(N)); + for (int32_t i = 0; i < N; ++i) Xf[i] = bfloat16_to_float_stub(X[i]); + reference_gelu_f32(Xf.data(), Of.data(), N); + for (int32_t i = 0; i < N; ++i) Out[i] = float_to_bfloat16_stub(Of[i]); +} + namespace { inline void reference_matmul_softmax_f32(const float* A, const float* B, diff --git a/tests/tessera-ir/.lit_test_times.txt b/tests/tessera-ir/.lit_test_times.txt index c5bd069de..7f4caa430 100644 --- a/tests/tessera-ir/.lit_test_times.txt +++ b/tests/tessera-ir/.lit_test_times.txt @@ -1,8 +1,8 @@ -8.501682e-01 phase8/apple_cpu_lowering.mlir -8.418021e-01 phase8/apple_dialect_roundtrip.mlir -8.678038e-01 phase8/apple_gpu_lowering.mlir -8.748050e-01 phase8/apple_cpu_runtime.mlir -1.045299e-02 phase8/target_ir_contracts.mlir +4.056001e-02 phase8/apple_cpu_lowering.mlir +3.585601e-02 phase8/apple_dialect_roundtrip.mlir +4.633713e-02 phase8/apple_gpu_lowering.mlir +4.500675e-02 phase8/apple_cpu_runtime.mlir +1.523685e-02 phase8/target_ir_contracts.mlir 4.739904e-02 phase2/distribution_lowering.mlir 4.920816e-02 phase2/effect_annotation.mlir 4.286575e-02 phase2/full_pipeline.mlir @@ -29,10 +29,11 @@ 4.314494e-02 phase7/neighbors_stencil_lower.mlir 4.376101e-02 phase7/shardy_export.mlir 3.875470e-02 pipelines/cleanup_pipeline.mlir -8.595231e-01 phase8/apple_gpu_runtime.mlir -1.159406e-02 phase8/tmem_tcgen05_contract.mlir -8.450270e-01 phase8/apple_gpu_msl.mlir -8.636880e-01 phase8/apple_gpu_flash_attn.mlir -8.775210e-01 phase8/apple_gpu_softmax_gelu.mlir -8.536379e-01 phase8/apple_gpu_matmul_softmax_fusion.mlir -8.237970e-01 phase8/apple_gpu_matmul_dtypes.mlir +4.584289e-02 phase8/apple_gpu_runtime.mlir +1.471019e-02 phase8/tmem_tcgen05_contract.mlir +4.270697e-02 phase8/apple_gpu_msl.mlir +4.811597e-02 phase8/apple_gpu_flash_attn.mlir +5.533004e-02 phase8/apple_gpu_softmax_gelu.mlir +4.885697e-02 phase8/apple_gpu_matmul_softmax_fusion.mlir +3.492808e-02 phase8/apple_gpu_matmul_dtypes.mlir +5.615902e-02 phase8/apple_gpu_msl_dtypes.mlir diff --git a/tests/tessera-ir/phase8/Output/apple_cpu_lowering.mlir.script b/tests/tessera-ir/phase8/Output/apple_cpu_lowering.mlir.script index 8ffc9539c..e16777431 100644 --- a/tests/tessera-ir/phase8/Output/apple_cpu_lowering.mlir.script +++ b/tests/tessera-ir/phase8/Output/apple_cpu_lowering.mlir.script @@ -1 +1 @@ -set -o pipefail;set -x;{ { set +x; } 2>/dev/null && echo 'RUN: at line 1': 'tessera-opt -allow-unregistered-dialect /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-4-fp16-bf16/tests/tessera-ir/phase8/apple_cpu_lowering.mlir -tessera-lower-to-apple_cpu --allow-unregistered-dialect | FileCheck /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-4-fp16-bf16/tests/tessera-ir/phase8/apple_cpu_lowering.mlir' >&2 && { set -x; } 2>/dev/null && { tessera-opt -allow-unregistered-dialect /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-4-fp16-bf16/tests/tessera-ir/phase8/apple_cpu_lowering.mlir -tessera-lower-to-apple_cpu --allow-unregistered-dialect | FileCheck /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-4-fp16-bf16/tests/tessera-ir/phase8/apple_cpu_lowering.mlir; }; } +set -o pipefail;set -x;{ { set +x; } 2>/dev/null && echo 'RUN: at line 1': 'tessera-opt -allow-unregistered-dialect /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-4-1-msl-dtypes/tests/tessera-ir/phase8/apple_cpu_lowering.mlir -tessera-lower-to-apple_cpu --allow-unregistered-dialect | FileCheck /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-4-1-msl-dtypes/tests/tessera-ir/phase8/apple_cpu_lowering.mlir' >&2 && { set -x; } 2>/dev/null && { tessera-opt -allow-unregistered-dialect /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-4-1-msl-dtypes/tests/tessera-ir/phase8/apple_cpu_lowering.mlir -tessera-lower-to-apple_cpu --allow-unregistered-dialect | FileCheck /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-4-1-msl-dtypes/tests/tessera-ir/phase8/apple_cpu_lowering.mlir; }; } diff --git a/tests/tessera-ir/phase8/Output/apple_cpu_runtime.mlir.script b/tests/tessera-ir/phase8/Output/apple_cpu_runtime.mlir.script index dea9258e6..56549a88c 100644 --- a/tests/tessera-ir/phase8/Output/apple_cpu_runtime.mlir.script +++ b/tests/tessera-ir/phase8/Output/apple_cpu_runtime.mlir.script @@ -1 +1 @@ -set -o pipefail;set -x;{ { set +x; } 2>/dev/null && echo 'RUN: at line 1': 'tessera-opt -allow-unregistered-dialect /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-4-fp16-bf16/tests/tessera-ir/phase8/apple_cpu_runtime.mlir --pass-pipeline='"'"'builtin.module(tessera-lower-to-apple_cpu-runtime)'"'"' --allow-unregistered-dialect | FileCheck /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-4-fp16-bf16/tests/tessera-ir/phase8/apple_cpu_runtime.mlir' >&2 && { set -x; } 2>/dev/null && { tessera-opt -allow-unregistered-dialect /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-4-fp16-bf16/tests/tessera-ir/phase8/apple_cpu_runtime.mlir --pass-pipeline='builtin.module(tessera-lower-to-apple_cpu-runtime)' --allow-unregistered-dialect | FileCheck /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-4-fp16-bf16/tests/tessera-ir/phase8/apple_cpu_runtime.mlir; }; } +set -o pipefail;set -x;{ { set +x; } 2>/dev/null && echo 'RUN: at line 1': 'tessera-opt -allow-unregistered-dialect /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-4-1-msl-dtypes/tests/tessera-ir/phase8/apple_cpu_runtime.mlir --pass-pipeline='"'"'builtin.module(tessera-lower-to-apple_cpu-runtime)'"'"' --allow-unregistered-dialect | FileCheck /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-4-1-msl-dtypes/tests/tessera-ir/phase8/apple_cpu_runtime.mlir' >&2 && { set -x; } 2>/dev/null && { tessera-opt -allow-unregistered-dialect /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-4-1-msl-dtypes/tests/tessera-ir/phase8/apple_cpu_runtime.mlir --pass-pipeline='builtin.module(tessera-lower-to-apple_cpu-runtime)' --allow-unregistered-dialect | FileCheck /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-4-1-msl-dtypes/tests/tessera-ir/phase8/apple_cpu_runtime.mlir; }; } diff --git a/tests/tessera-ir/phase8/Output/apple_dialect_roundtrip.mlir.script b/tests/tessera-ir/phase8/Output/apple_dialect_roundtrip.mlir.script index 0225a1d31..834461c03 100644 --- a/tests/tessera-ir/phase8/Output/apple_dialect_roundtrip.mlir.script +++ b/tests/tessera-ir/phase8/Output/apple_dialect_roundtrip.mlir.script @@ -1 +1 @@ -set -o pipefail;set -x;{ { set +x; } 2>/dev/null && echo 'RUN: at line 1': 'tessera-opt -allow-unregistered-dialect /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-4-fp16-bf16/tests/tessera-ir/phase8/apple_dialect_roundtrip.mlir --allow-unregistered-dialect | FileCheck /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-4-fp16-bf16/tests/tessera-ir/phase8/apple_dialect_roundtrip.mlir' >&2 && { set -x; } 2>/dev/null && { tessera-opt -allow-unregistered-dialect /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-4-fp16-bf16/tests/tessera-ir/phase8/apple_dialect_roundtrip.mlir --allow-unregistered-dialect | FileCheck /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-4-fp16-bf16/tests/tessera-ir/phase8/apple_dialect_roundtrip.mlir; }; } +set -o pipefail;set -x;{ { set +x; } 2>/dev/null && echo 'RUN: at line 1': 'tessera-opt -allow-unregistered-dialect /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-4-1-msl-dtypes/tests/tessera-ir/phase8/apple_dialect_roundtrip.mlir --allow-unregistered-dialect | FileCheck /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-4-1-msl-dtypes/tests/tessera-ir/phase8/apple_dialect_roundtrip.mlir' >&2 && { set -x; } 2>/dev/null && { tessera-opt -allow-unregistered-dialect /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-4-1-msl-dtypes/tests/tessera-ir/phase8/apple_dialect_roundtrip.mlir --allow-unregistered-dialect | FileCheck /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-4-1-msl-dtypes/tests/tessera-ir/phase8/apple_dialect_roundtrip.mlir; }; } diff --git a/tests/tessera-ir/phase8/Output/apple_gpu_flash_attn.mlir.script b/tests/tessera-ir/phase8/Output/apple_gpu_flash_attn.mlir.script index f00b9bdcb..2fa668890 100644 --- a/tests/tessera-ir/phase8/Output/apple_gpu_flash_attn.mlir.script +++ b/tests/tessera-ir/phase8/Output/apple_gpu_flash_attn.mlir.script @@ -1 +1 @@ -set -o pipefail;set -x;{ { set +x; } 2>/dev/null && echo 'RUN: at line 1': 'tessera-opt -allow-unregistered-dialect /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-4-fp16-bf16/tests/tessera-ir/phase8/apple_gpu_flash_attn.mlir --pass-pipeline='"'"'builtin.module(tessera-lower-to-apple_gpu-runtime)'"'"' --allow-unregistered-dialect | FileCheck /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-4-fp16-bf16/tests/tessera-ir/phase8/apple_gpu_flash_attn.mlir' >&2 && { set -x; } 2>/dev/null && { tessera-opt -allow-unregistered-dialect /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-4-fp16-bf16/tests/tessera-ir/phase8/apple_gpu_flash_attn.mlir --pass-pipeline='builtin.module(tessera-lower-to-apple_gpu-runtime)' --allow-unregistered-dialect | FileCheck /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-4-fp16-bf16/tests/tessera-ir/phase8/apple_gpu_flash_attn.mlir; }; } +set -o pipefail;set -x;{ { set +x; } 2>/dev/null && echo 'RUN: at line 1': 'tessera-opt -allow-unregistered-dialect /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-4-1-msl-dtypes/tests/tessera-ir/phase8/apple_gpu_flash_attn.mlir --pass-pipeline='"'"'builtin.module(tessera-lower-to-apple_gpu-runtime)'"'"' --allow-unregistered-dialect | FileCheck /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-4-1-msl-dtypes/tests/tessera-ir/phase8/apple_gpu_flash_attn.mlir' >&2 && { set -x; } 2>/dev/null && { tessera-opt -allow-unregistered-dialect /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-4-1-msl-dtypes/tests/tessera-ir/phase8/apple_gpu_flash_attn.mlir --pass-pipeline='builtin.module(tessera-lower-to-apple_gpu-runtime)' --allow-unregistered-dialect | FileCheck /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-4-1-msl-dtypes/tests/tessera-ir/phase8/apple_gpu_flash_attn.mlir; }; } diff --git a/tests/tessera-ir/phase8/Output/apple_gpu_lowering.mlir.script b/tests/tessera-ir/phase8/Output/apple_gpu_lowering.mlir.script index 6b0ec2fe7..e4eb82769 100644 --- a/tests/tessera-ir/phase8/Output/apple_gpu_lowering.mlir.script +++ b/tests/tessera-ir/phase8/Output/apple_gpu_lowering.mlir.script @@ -1 +1 @@ -set -o pipefail;set -x;{ { set +x; } 2>/dev/null && echo 'RUN: at line 1': 'tessera-opt -allow-unregistered-dialect /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-4-fp16-bf16/tests/tessera-ir/phase8/apple_gpu_lowering.mlir -tessera-lower-to-apple_gpu --allow-unregistered-dialect | FileCheck /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-4-fp16-bf16/tests/tessera-ir/phase8/apple_gpu_lowering.mlir' >&2 && { set -x; } 2>/dev/null && { tessera-opt -allow-unregistered-dialect /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-4-fp16-bf16/tests/tessera-ir/phase8/apple_gpu_lowering.mlir -tessera-lower-to-apple_gpu --allow-unregistered-dialect | FileCheck /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-4-fp16-bf16/tests/tessera-ir/phase8/apple_gpu_lowering.mlir; }; } +set -o pipefail;set -x;{ { set +x; } 2>/dev/null && echo 'RUN: at line 1': 'tessera-opt -allow-unregistered-dialect /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-4-1-msl-dtypes/tests/tessera-ir/phase8/apple_gpu_lowering.mlir -tessera-lower-to-apple_gpu --allow-unregistered-dialect | FileCheck /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-4-1-msl-dtypes/tests/tessera-ir/phase8/apple_gpu_lowering.mlir' >&2 && { set -x; } 2>/dev/null && { tessera-opt -allow-unregistered-dialect /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-4-1-msl-dtypes/tests/tessera-ir/phase8/apple_gpu_lowering.mlir -tessera-lower-to-apple_gpu --allow-unregistered-dialect | FileCheck /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-4-1-msl-dtypes/tests/tessera-ir/phase8/apple_gpu_lowering.mlir; }; } diff --git a/tests/tessera-ir/phase8/Output/apple_gpu_matmul_dtypes.mlir.script b/tests/tessera-ir/phase8/Output/apple_gpu_matmul_dtypes.mlir.script index a4f54a63e..7b9a7080b 100644 --- a/tests/tessera-ir/phase8/Output/apple_gpu_matmul_dtypes.mlir.script +++ b/tests/tessera-ir/phase8/Output/apple_gpu_matmul_dtypes.mlir.script @@ -1 +1 @@ -set -o pipefail;set -x;{ { set +x; } 2>/dev/null && echo 'RUN: at line 1': 'tessera-opt -allow-unregistered-dialect /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-4-fp16-bf16/tests/tessera-ir/phase8/apple_gpu_matmul_dtypes.mlir --pass-pipeline='"'"'builtin.module(tessera-lower-to-apple_gpu-runtime)'"'"' --allow-unregistered-dialect | FileCheck /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-4-fp16-bf16/tests/tessera-ir/phase8/apple_gpu_matmul_dtypes.mlir' >&2 && { set -x; } 2>/dev/null && { tessera-opt -allow-unregistered-dialect /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-4-fp16-bf16/tests/tessera-ir/phase8/apple_gpu_matmul_dtypes.mlir --pass-pipeline='builtin.module(tessera-lower-to-apple_gpu-runtime)' --allow-unregistered-dialect | FileCheck /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-4-fp16-bf16/tests/tessera-ir/phase8/apple_gpu_matmul_dtypes.mlir; }; } +set -o pipefail;set -x;{ { set +x; } 2>/dev/null && echo 'RUN: at line 1': 'tessera-opt -allow-unregistered-dialect /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-4-1-msl-dtypes/tests/tessera-ir/phase8/apple_gpu_matmul_dtypes.mlir --pass-pipeline='"'"'builtin.module(tessera-lower-to-apple_gpu-runtime)'"'"' --allow-unregistered-dialect | FileCheck /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-4-1-msl-dtypes/tests/tessera-ir/phase8/apple_gpu_matmul_dtypes.mlir' >&2 && { set -x; } 2>/dev/null && { tessera-opt -allow-unregistered-dialect /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-4-1-msl-dtypes/tests/tessera-ir/phase8/apple_gpu_matmul_dtypes.mlir --pass-pipeline='builtin.module(tessera-lower-to-apple_gpu-runtime)' --allow-unregistered-dialect | FileCheck /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-4-1-msl-dtypes/tests/tessera-ir/phase8/apple_gpu_matmul_dtypes.mlir; }; } diff --git a/tests/tessera-ir/phase8/Output/apple_gpu_matmul_softmax_fusion.mlir.script b/tests/tessera-ir/phase8/Output/apple_gpu_matmul_softmax_fusion.mlir.script index efa2bbde9..5c14f6c2e 100644 --- a/tests/tessera-ir/phase8/Output/apple_gpu_matmul_softmax_fusion.mlir.script +++ b/tests/tessera-ir/phase8/Output/apple_gpu_matmul_softmax_fusion.mlir.script @@ -1 +1 @@ -set -o pipefail;set -x;{ { set +x; } 2>/dev/null && echo 'RUN: at line 1': 'tessera-opt -allow-unregistered-dialect /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-4-fp16-bf16/tests/tessera-ir/phase8/apple_gpu_matmul_softmax_fusion.mlir --pass-pipeline='"'"'builtin.module(tessera-lower-to-apple_gpu-runtime)'"'"' --allow-unregistered-dialect | FileCheck /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-4-fp16-bf16/tests/tessera-ir/phase8/apple_gpu_matmul_softmax_fusion.mlir' >&2 && { set -x; } 2>/dev/null && { tessera-opt -allow-unregistered-dialect /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-4-fp16-bf16/tests/tessera-ir/phase8/apple_gpu_matmul_softmax_fusion.mlir --pass-pipeline='builtin.module(tessera-lower-to-apple_gpu-runtime)' --allow-unregistered-dialect | FileCheck /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-4-fp16-bf16/tests/tessera-ir/phase8/apple_gpu_matmul_softmax_fusion.mlir; }; } +set -o pipefail;set -x;{ { set +x; } 2>/dev/null && echo 'RUN: at line 1': 'tessera-opt -allow-unregistered-dialect /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-4-1-msl-dtypes/tests/tessera-ir/phase8/apple_gpu_matmul_softmax_fusion.mlir --pass-pipeline='"'"'builtin.module(tessera-lower-to-apple_gpu-runtime)'"'"' --allow-unregistered-dialect | FileCheck /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-4-1-msl-dtypes/tests/tessera-ir/phase8/apple_gpu_matmul_softmax_fusion.mlir' >&2 && { set -x; } 2>/dev/null && { tessera-opt -allow-unregistered-dialect /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-4-1-msl-dtypes/tests/tessera-ir/phase8/apple_gpu_matmul_softmax_fusion.mlir --pass-pipeline='builtin.module(tessera-lower-to-apple_gpu-runtime)' --allow-unregistered-dialect | FileCheck /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-4-1-msl-dtypes/tests/tessera-ir/phase8/apple_gpu_matmul_softmax_fusion.mlir; }; } diff --git a/tests/tessera-ir/phase8/Output/apple_gpu_msl.mlir.script b/tests/tessera-ir/phase8/Output/apple_gpu_msl.mlir.script index a54a4e6b8..7de1ac77a 100644 --- a/tests/tessera-ir/phase8/Output/apple_gpu_msl.mlir.script +++ b/tests/tessera-ir/phase8/Output/apple_gpu_msl.mlir.script @@ -1 +1 @@ -set -o pipefail;set -x;{ { set +x; } 2>/dev/null && echo 'RUN: at line 1': 'tessera-opt -allow-unregistered-dialect /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-4-fp16-bf16/tests/tessera-ir/phase8/apple_gpu_msl.mlir --allow-unregistered-dialect | FileCheck /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-4-fp16-bf16/tests/tessera-ir/phase8/apple_gpu_msl.mlir' >&2 && { set -x; } 2>/dev/null && { tessera-opt -allow-unregistered-dialect /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-4-fp16-bf16/tests/tessera-ir/phase8/apple_gpu_msl.mlir --allow-unregistered-dialect | FileCheck /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-4-fp16-bf16/tests/tessera-ir/phase8/apple_gpu_msl.mlir; }; } +set -o pipefail;set -x;{ { set +x; } 2>/dev/null && echo 'RUN: at line 1': 'tessera-opt -allow-unregistered-dialect /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-4-1-msl-dtypes/tests/tessera-ir/phase8/apple_gpu_msl.mlir --allow-unregistered-dialect | FileCheck /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-4-1-msl-dtypes/tests/tessera-ir/phase8/apple_gpu_msl.mlir' >&2 && { set -x; } 2>/dev/null && { tessera-opt -allow-unregistered-dialect /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-4-1-msl-dtypes/tests/tessera-ir/phase8/apple_gpu_msl.mlir --allow-unregistered-dialect | FileCheck /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-4-1-msl-dtypes/tests/tessera-ir/phase8/apple_gpu_msl.mlir; }; } diff --git a/tests/tessera-ir/phase8/Output/apple_gpu_msl_dtypes.mlir.script b/tests/tessera-ir/phase8/Output/apple_gpu_msl_dtypes.mlir.script new file mode 100644 index 000000000..7eb30fda3 --- /dev/null +++ b/tests/tessera-ir/phase8/Output/apple_gpu_msl_dtypes.mlir.script @@ -0,0 +1 @@ +set -o pipefail;set -x;{ { set +x; } 2>/dev/null && echo 'RUN: at line 1': 'tessera-opt -allow-unregistered-dialect /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-4-1-msl-dtypes/tests/tessera-ir/phase8/apple_gpu_msl_dtypes.mlir --pass-pipeline='"'"'builtin.module(tessera-lower-to-apple_gpu-runtime)'"'"' --allow-unregistered-dialect | FileCheck /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-4-1-msl-dtypes/tests/tessera-ir/phase8/apple_gpu_msl_dtypes.mlir' >&2 && { set -x; } 2>/dev/null && { tessera-opt -allow-unregistered-dialect /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-4-1-msl-dtypes/tests/tessera-ir/phase8/apple_gpu_msl_dtypes.mlir --pass-pipeline='builtin.module(tessera-lower-to-apple_gpu-runtime)' --allow-unregistered-dialect | FileCheck /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-4-1-msl-dtypes/tests/tessera-ir/phase8/apple_gpu_msl_dtypes.mlir; }; } diff --git a/tests/tessera-ir/phase8/Output/apple_gpu_runtime.mlir.script b/tests/tessera-ir/phase8/Output/apple_gpu_runtime.mlir.script index a3cb3dc52..bc89a1bdf 100644 --- a/tests/tessera-ir/phase8/Output/apple_gpu_runtime.mlir.script +++ b/tests/tessera-ir/phase8/Output/apple_gpu_runtime.mlir.script @@ -1 +1 @@ -set -o pipefail;set -x;{ { set +x; } 2>/dev/null && echo 'RUN: at line 1': 'tessera-opt -allow-unregistered-dialect /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-4-fp16-bf16/tests/tessera-ir/phase8/apple_gpu_runtime.mlir --pass-pipeline='"'"'builtin.module(tessera-lower-to-apple_gpu-runtime)'"'"' --allow-unregistered-dialect | FileCheck /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-4-fp16-bf16/tests/tessera-ir/phase8/apple_gpu_runtime.mlir' >&2 && { set -x; } 2>/dev/null && { tessera-opt -allow-unregistered-dialect /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-4-fp16-bf16/tests/tessera-ir/phase8/apple_gpu_runtime.mlir --pass-pipeline='builtin.module(tessera-lower-to-apple_gpu-runtime)' --allow-unregistered-dialect | FileCheck /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-4-fp16-bf16/tests/tessera-ir/phase8/apple_gpu_runtime.mlir; }; } +set -o pipefail;set -x;{ { set +x; } 2>/dev/null && echo 'RUN: at line 1': 'tessera-opt -allow-unregistered-dialect /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-4-1-msl-dtypes/tests/tessera-ir/phase8/apple_gpu_runtime.mlir --pass-pipeline='"'"'builtin.module(tessera-lower-to-apple_gpu-runtime)'"'"' --allow-unregistered-dialect | FileCheck /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-4-1-msl-dtypes/tests/tessera-ir/phase8/apple_gpu_runtime.mlir' >&2 && { set -x; } 2>/dev/null && { tessera-opt -allow-unregistered-dialect /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-4-1-msl-dtypes/tests/tessera-ir/phase8/apple_gpu_runtime.mlir --pass-pipeline='builtin.module(tessera-lower-to-apple_gpu-runtime)' --allow-unregistered-dialect | FileCheck /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-4-1-msl-dtypes/tests/tessera-ir/phase8/apple_gpu_runtime.mlir; }; } diff --git a/tests/tessera-ir/phase8/Output/apple_gpu_softmax_gelu.mlir.script b/tests/tessera-ir/phase8/Output/apple_gpu_softmax_gelu.mlir.script index a57989598..34a86dcb9 100644 --- a/tests/tessera-ir/phase8/Output/apple_gpu_softmax_gelu.mlir.script +++ b/tests/tessera-ir/phase8/Output/apple_gpu_softmax_gelu.mlir.script @@ -1 +1 @@ -set -o pipefail;set -x;{ { set +x; } 2>/dev/null && echo 'RUN: at line 1': 'tessera-opt -allow-unregistered-dialect /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-4-fp16-bf16/tests/tessera-ir/phase8/apple_gpu_softmax_gelu.mlir --pass-pipeline='"'"'builtin.module(tessera-lower-to-apple_gpu-runtime)'"'"' --allow-unregistered-dialect | FileCheck /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-4-fp16-bf16/tests/tessera-ir/phase8/apple_gpu_softmax_gelu.mlir' >&2 && { set -x; } 2>/dev/null && { tessera-opt -allow-unregistered-dialect /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-4-fp16-bf16/tests/tessera-ir/phase8/apple_gpu_softmax_gelu.mlir --pass-pipeline='builtin.module(tessera-lower-to-apple_gpu-runtime)' --allow-unregistered-dialect | FileCheck /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-4-fp16-bf16/tests/tessera-ir/phase8/apple_gpu_softmax_gelu.mlir; }; } +set -o pipefail;set -x;{ { set +x; } 2>/dev/null && echo 'RUN: at line 1': 'tessera-opt -allow-unregistered-dialect /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-4-1-msl-dtypes/tests/tessera-ir/phase8/apple_gpu_softmax_gelu.mlir --pass-pipeline='"'"'builtin.module(tessera-lower-to-apple_gpu-runtime)'"'"' --allow-unregistered-dialect | FileCheck /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-4-1-msl-dtypes/tests/tessera-ir/phase8/apple_gpu_softmax_gelu.mlir' >&2 && { set -x; } 2>/dev/null && { tessera-opt -allow-unregistered-dialect /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-4-1-msl-dtypes/tests/tessera-ir/phase8/apple_gpu_softmax_gelu.mlir --pass-pipeline='builtin.module(tessera-lower-to-apple_gpu-runtime)' --allow-unregistered-dialect | FileCheck /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-4-1-msl-dtypes/tests/tessera-ir/phase8/apple_gpu_softmax_gelu.mlir; }; } diff --git a/tests/tessera-ir/phase8/Output/target_ir_contracts.mlir.script b/tests/tessera-ir/phase8/Output/target_ir_contracts.mlir.script index 3ff8bef57..94c477959 100644 --- a/tests/tessera-ir/phase8/Output/target_ir_contracts.mlir.script +++ b/tests/tessera-ir/phase8/Output/target_ir_contracts.mlir.script @@ -1 +1 @@ -set -o pipefail;set -x;{ { set +x; } 2>/dev/null && echo 'RUN: at line 1': 'FileCheck /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-4-fp16-bf16/tests/tessera-ir/phase8/target_ir_contracts.mlir --check-prefixes=ROCM,METALIUM,APPLE-CPU,APPLE-GPU < /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-4-fp16-bf16/tests/tessera-ir/phase8/target_ir_contracts.mlir' >&2 && { set -x; } 2>/dev/null && { FileCheck /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-4-fp16-bf16/tests/tessera-ir/phase8/target_ir_contracts.mlir --check-prefixes=ROCM,METALIUM,APPLE-CPU,APPLE-GPU < /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-4-fp16-bf16/tests/tessera-ir/phase8/target_ir_contracts.mlir; }; } +set -o pipefail;set -x;{ { set +x; } 2>/dev/null && echo 'RUN: at line 1': 'FileCheck /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-4-1-msl-dtypes/tests/tessera-ir/phase8/target_ir_contracts.mlir --check-prefixes=ROCM,METALIUM,APPLE-CPU,APPLE-GPU < /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-4-1-msl-dtypes/tests/tessera-ir/phase8/target_ir_contracts.mlir' >&2 && { set -x; } 2>/dev/null && { FileCheck /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-4-1-msl-dtypes/tests/tessera-ir/phase8/target_ir_contracts.mlir --check-prefixes=ROCM,METALIUM,APPLE-CPU,APPLE-GPU < /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-4-1-msl-dtypes/tests/tessera-ir/phase8/target_ir_contracts.mlir; }; } diff --git a/tests/tessera-ir/phase8/Output/tmem_tcgen05_contract.mlir.script b/tests/tessera-ir/phase8/Output/tmem_tcgen05_contract.mlir.script index c18936445..1856b75dd 100644 --- a/tests/tessera-ir/phase8/Output/tmem_tcgen05_contract.mlir.script +++ b/tests/tessera-ir/phase8/Output/tmem_tcgen05_contract.mlir.script @@ -1 +1 @@ -set -o pipefail;set -x;{ { set +x; } 2>/dev/null && echo 'RUN: at line 1': 'FileCheck /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-4-fp16-bf16/tests/tessera-ir/phase8/tmem_tcgen05_contract.mlir < /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-4-fp16-bf16/tests/tessera-ir/phase8/tmem_tcgen05_contract.mlir' >&2 && { set -x; } 2>/dev/null && { FileCheck /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-4-fp16-bf16/tests/tessera-ir/phase8/tmem_tcgen05_contract.mlir < /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-4-fp16-bf16/tests/tessera-ir/phase8/tmem_tcgen05_contract.mlir; }; } +set -o pipefail;set -x;{ { set +x; } 2>/dev/null && echo 'RUN: at line 1': 'FileCheck /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-4-1-msl-dtypes/tests/tessera-ir/phase8/tmem_tcgen05_contract.mlir < /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-4-1-msl-dtypes/tests/tessera-ir/phase8/tmem_tcgen05_contract.mlir' >&2 && { set -x; } 2>/dev/null && { FileCheck /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-4-1-msl-dtypes/tests/tessera-ir/phase8/tmem_tcgen05_contract.mlir < /Users/gregorystoner/dev_project/tessera/.claude/worktrees/phase-8-4-4-1-msl-dtypes/tests/tessera-ir/phase8/tmem_tcgen05_contract.mlir; }; } diff --git a/tests/tessera-ir/phase8/apple_gpu_msl_dtypes.mlir b/tests/tessera-ir/phase8/apple_gpu_msl_dtypes.mlir new file mode 100644 index 000000000..defa3d20e --- /dev/null +++ b/tests/tessera-ir/phase8/apple_gpu_msl_dtypes.mlir @@ -0,0 +1,64 @@ +// RUN: tessera-opt %s --pass-pipeline='builtin.module(tessera-lower-to-apple_gpu-runtime)' --allow-unregistered-dialect | FileCheck %s + +// Phase 8.4.4.1 — Apple GPU custom MSL kernels with fp16 / bf16 dtype +// variants for softmax and gelu. Verifies that the runtime pipeline picks +// the right runtime symbol per (kernel, dtype) pair. All three dtypes +// share the same i64 + i32 ABI shape per kernel — element type is encoded +// in the symbol name only. +// +// rope dtype dispatch is exercised by Python unit tests instead — the +// tessera.rope op is not registered in the Tessera dialect, which prevents +// pass-level lit testing of it. + +// Runtime declarations: one per (kernel, dtype) pair. Order between +// declarations is implementation-defined, so use CHECK-DAG. +// CHECK-DAG: func.func private @tessera_apple_gpu_softmax_f32(i64, i64, i32, i32) +// CHECK-DAG: func.func private @tessera_apple_gpu_softmax_f16(i64, i64, i32, i32) +// CHECK-DAG: func.func private @tessera_apple_gpu_softmax_bf16(i64, i64, i32, i32) +// CHECK-DAG: func.func private @tessera_apple_gpu_gelu_f32(i64, i64, i32) +// CHECK-DAG: func.func private @tessera_apple_gpu_gelu_f16(i64, i64, i32) +// CHECK-DAG: func.func private @tessera_apple_gpu_gelu_bf16(i64, i64, i32) + +func.func @softmax_f16(%X: tensor<8x16xf16>) -> tensor<8x16xf16> { + // CHECK-LABEL: func.func @softmax_f16 + // CHECK: call @tessera_apple_gpu_softmax_f16 + %Out = "tessera.softmax"(%X) : (tensor<8x16xf16>) -> tensor<8x16xf16> + return %Out : tensor<8x16xf16> +} + +func.func @softmax_bf16(%X: tensor<8x16xbf16>) -> tensor<8x16xbf16> { + // CHECK-LABEL: func.func @softmax_bf16 + // CHECK: call @tessera_apple_gpu_softmax_bf16 + %Out = "tessera.softmax"(%X) : (tensor<8x16xbf16>) -> tensor<8x16xbf16> + return %Out : tensor<8x16xbf16> +} + +func.func @gelu_f16(%X: tensor<8x16xf16>) -> tensor<8x16xf16> { + // CHECK-LABEL: func.func @gelu_f16 + // CHECK: call @tessera_apple_gpu_gelu_f16 + %Out = "tessera.gelu"(%X) : (tensor<8x16xf16>) -> tensor<8x16xf16> + return %Out : tensor<8x16xf16> +} + +func.func @gelu_bf16(%X: tensor<8x16xbf16>) -> tensor<8x16xbf16> { + // CHECK-LABEL: func.func @gelu_bf16 + // CHECK: call @tessera_apple_gpu_gelu_bf16 + %Out = "tessera.gelu"(%X) : (tensor<8x16xbf16>) -> tensor<8x16xbf16> + return %Out : tensor<8x16xbf16> +} + +// f32 cases — these paths shouldn't regress. + +func.func @softmax_f32(%X: tensor<8x16xf32>) -> tensor<8x16xf32> { + // CHECK-LABEL: func.func @softmax_f32 + // CHECK: call @tessera_apple_gpu_softmax_f32 + %Out = "tessera.softmax"(%X) : (tensor<8x16xf32>) -> tensor<8x16xf32> + return %Out : tensor<8x16xf32> +} + +func.func @gelu_f32(%X: tensor<8x16xf32>) -> tensor<8x16xf32> { + // CHECK-LABEL: func.func @gelu_f32 + // CHECK: call @tessera_apple_gpu_gelu_f32 + %Out = "tessera.gelu"(%X) : (tensor<8x16xf32>) -> tensor<8x16xf32> + return %Out : tensor<8x16xf32> +} diff --git a/tests/unit/test_apple_backend_roadmap.py b/tests/unit/test_apple_backend_roadmap.py index ae4faa023..056d77d73 100644 --- a/tests/unit/test_apple_backend_roadmap.py +++ b/tests/unit/test_apple_backend_roadmap.py @@ -1442,6 +1442,184 @@ def test_apple_gpu_matmul_runtime_shim_exposes_f16_and_bf16_symbols(tmp_path): ) +# ───────────────────────────────────────────────────────────────────────────── +# Phase 8.4.4.1: fp16 / bf16 for the simple MSL kernels (rope, softmax, gelu). +# +# fp16 path uses native MSL `half` kernels with `float` internal compute. +# bf16 path uses fp32-conversion at the runtime boundary (no native MSL bf16). +# Same pattern as Phase 8.4.4 matmul. The Python dispatcher detects input +# array dtype at runtime and routes to the matching ctypes wrapper. +# ───────────────────────────────────────────────────────────────────────────── + + +def test_apple_gpu_rope_f16_executes_through_native_msl(): + @ts.jit(target="apple_gpu") + def rope(X, Theta): + return ts.ops.rope(X, Theta) + + rng = np.random.RandomState(67) + M, K = 8, 16 + X = rng.randn(M, K).astype(np.float16) * 0.5 + Theta = rng.uniform(-np.pi, np.pi, size=(M, K)).astype(np.float16) + out = rope(X, Theta) + assert out.dtype == np.float16 + assert out.shape == (M, K) + Xf = X.astype(np.float32) + Tf = Theta.astype(np.float32) + even = Xf[:, 0::2] + odd = Xf[:, 1::2] + theta_even = Tf[:, 0::2] + expected = np.empty_like(Xf) + expected[:, 0::2] = even * np.cos(theta_even) - odd * np.sin(theta_even) + expected[:, 1::2] = even * np.sin(theta_even) + odd * np.cos(theta_even) + np.testing.assert_allclose( + out.astype(np.float32), expected, rtol=5e-3, atol=5e-3, + ) + + +def test_apple_gpu_rope_bf16_executes_through_fp32_conversion_path(): + pytest.importorskip("ml_dtypes") + import ml_dtypes + bf16 = ml_dtypes.bfloat16 + + @ts.jit(target="apple_gpu") + def rope(X, Theta): + return ts.ops.rope(X, Theta) + + rng = np.random.RandomState(71) + M, K = 8, 16 + # Multiplication BEFORE astype — `bf16_arr * python_float` would promote + # back to float32 because numpy treats Python scalars as float64 and + # downcast goes through fp32. Apply scaling in fp32, then cast. + X = (rng.randn(M, K) * 0.5).astype(bf16) + Theta = rng.uniform(-np.pi, np.pi, size=(M, K)).astype(bf16) + out = rope(X, Theta) + assert out.dtype == bf16 + assert out.shape == (M, K) + Xf = X.astype(np.float32) + Tf = Theta.astype(np.float32) + even = Xf[:, 0::2] + odd = Xf[:, 1::2] + theta_even = Tf[:, 0::2] + expected = np.empty_like(Xf) + expected[:, 0::2] = even * np.cos(theta_even) - odd * np.sin(theta_even) + expected[:, 1::2] = even * np.sin(theta_even) + odd * np.cos(theta_even) + np.testing.assert_allclose( + out.astype(np.float32), expected, rtol=2e-2, atol=2e-2, + ) + + +def test_apple_gpu_softmax_f16_executes_through_native_msl(): + @ts.jit(target="apple_gpu") + def sm(X): + return ts.ops.softmax(X) + + rng = np.random.RandomState(73) + for shape in ((4, 8), (8, 32)): + X = rng.randn(*shape).astype(np.float16) + out = sm(X) + assert out.dtype == np.float16 + assert out.shape == shape + ref_e = np.exp(X.astype(np.float32) - np.max(X.astype(np.float32), axis=-1, keepdims=True)) + ref = ref_e / np.sum(ref_e, axis=-1, keepdims=True) + np.testing.assert_allclose( + out.astype(np.float32), ref, rtol=5e-3, atol=5e-3, + ) + + +def test_apple_gpu_softmax_bf16_executes_through_fp32_conversion_path(): + pytest.importorskip("ml_dtypes") + import ml_dtypes + bf16 = ml_dtypes.bfloat16 + + @ts.jit(target="apple_gpu") + def sm(X): + return ts.ops.softmax(X) + + rng = np.random.RandomState(79) + X = rng.randn(8, 16).astype(bf16) + out = sm(X) + assert out.dtype == bf16 + ref_e = np.exp(X.astype(np.float32) - np.max(X.astype(np.float32), axis=-1, keepdims=True)) + ref = ref_e / np.sum(ref_e, axis=-1, keepdims=True) + np.testing.assert_allclose( + out.astype(np.float32), ref, rtol=2e-2, atol=2e-2, + ) + + +def test_apple_gpu_gelu_f16_executes_through_native_msl(): + @ts.jit(target="apple_gpu") + def gelu(X): + return ts.ops.gelu(X) + + rng = np.random.RandomState(83) + X = rng.randn(8, 16).astype(np.float16) * 1.5 + out = gelu(X) + assert out.dtype == np.float16 + Xf = X.astype(np.float32) + ref = 0.5 * Xf * (1.0 + np.tanh(np.sqrt(2.0 / np.pi) * (Xf + 0.044715 * Xf**3))) + np.testing.assert_allclose( + out.astype(np.float32), ref, rtol=5e-3, atol=5e-3, + ) + + +def test_apple_gpu_gelu_bf16_executes_through_fp32_conversion_path(): + pytest.importorskip("ml_dtypes") + import ml_dtypes + bf16 = ml_dtypes.bfloat16 + + @ts.jit(target="apple_gpu") + def gelu(X): + return ts.ops.gelu(X) + + rng = np.random.RandomState(89) + X = (rng.randn(8, 16) * 1.5).astype(bf16) + out = gelu(X) + assert out.dtype == bf16 + Xf = X.astype(np.float32) + ref = 0.5 * Xf * (1.0 + np.tanh(np.sqrt(2.0 / np.pi) * (Xf + 0.044715 * Xf**3))) + np.testing.assert_allclose( + out.astype(np.float32), ref, rtol=2e-2, atol=2e-2, + ) + + +def test_apple_gpu_msl_dtype_runtime_shim_exposes_all_symbols(tmp_path): + """Compile the apple_gpu runtime shim from source and verify all 6 new + fp16/bf16 symbols (rope_{f16,bf16}, softmax_{f16,bf16}, gelu_{f16,bf16}) + are exported.""" + + cxx = shutil.which("c++") or shutil.which("clang++") or shutil.which("g++") + if cxx is None: + pytest.skip("C++ compiler is not available") + + backend = ROOT / "src/compiler/codegen/Tessera_Apple_Backend/runtime" + if sys.platform == "darwin": + source = backend / "apple_gpu_runtime.mm" + lib = tmp_path / "libtessera_apple_gpu_runtime.dylib" + cmd = [cxx, "-std=c++17", "-shared", "-fPIC", "-fobjc-arc", + "-x", "objective-c++", str(source), "-o", str(lib), + "-framework", "Foundation", + "-framework", "Metal", + "-framework", "MetalPerformanceShaders"] + else: + source = backend / "apple_gpu_runtime_stub.cpp" + lib = tmp_path / "libtessera_apple_gpu_runtime.so" + cmd = [cxx, "-std=c++17", "-shared", "-fPIC", str(source), "-o", str(lib)] + subprocess.run(cmd, check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True) + + runtime = ctypes.CDLL(str(lib)) + for name in ( + "tessera_apple_gpu_rope_f16", + "tessera_apple_gpu_rope_bf16", + "tessera_apple_gpu_softmax_f16", + "tessera_apple_gpu_softmax_bf16", + "tessera_apple_gpu_gelu_f16", + "tessera_apple_gpu_gelu_bf16", + ): + sym = getattr(runtime, name, None) + assert sym is not None, f"missing C ABI symbol: {name}" + + def test_apple_cpu_bf16_disabled_when_ml_dtypes_missing(monkeypatch): """When ml_dtypes isn't installed the bf16 dtype probe returns None and the runtime falls through to numpy. Verified by stubbing the import to