[Mlir-commits] [mlir] [MLIR][GPUToXeVMPipeline] Expand MX scaling ops before XeVM conversion (PR #203632)

[Sang Ik Lee]

https://github.com/silee2 created https://github.com/llvm/llvm-project/pull/203632

arith.scaling_extf/scaling_truncf were never lowered by the gpu-lower-to-xevm pipeline, so micro-scaling (MX) GEMM kernels kept these ops (and their narrow-float operands) live all the way to LLVM translation.

Run arith-expand before the XeVM/LLVM conversions to break scaling_extf/scaling_truncf into extf/truncf + mulf and to expand f8E8M0FNU casts into integer bit manipulation. f4E2M1FN expansion is intentionally left disabled: its casts are lowered by the XeVM conversions (xevm.extf), whereas f8E8M0FNU is not handled there and must be expanded here. The generic f4E2M1FN expansion would otherwise emit i4 vector arithmetic that the XeVM backend cannot legalize.

>From d9a860987c1718d9af54ca187c4d413ceb3cd279 Mon Sep 17 00:00:00 2001
From: "Lee, Sang Ik" <sang.ik.lee at intel.com>
Date: Thu, 11 Jun 2026 16:52:54 +0000
Subject: [PATCH] [MLIR][GPUToXeVMPipeline] Expand MX scaling ops before XeVM
 conversion

arith.scaling_extf/scaling_truncf were never lowered by the gpu-lower-to-xevm pipeline, so micro-scaling (MX) GEMM kernels kept these ops (and their narrow-float operands) live all the way to LLVM translation.

Run arith-expand before the XeVM/LLVM conversions to break scaling_extf/scaling_truncf into extf/truncf + mulf and to expand f8E8M0FNU casts into integer bit manipulation. f4E2M1FN expansion is intentionally left disabled: its casts are lowered by the XeVM conversions (xevm.extf), whereas f8E8M0FNU is not handled there and must be expanded here. The generic f4E2M1FN expansion would otherwise emit i4 vector arithmetic that the XeVM backend cannot legalize.
---
 .../lib/Dialect/GPU/Pipelines/GPUToXeVMPipeline.cpp | 13 +++++++++++++
 1 file changed, 13 insertions(+)

diff --git a/mlir/lib/Dialect/GPU/Pipelines/GPUToXeVMPipeline.cpp b/mlir/lib/Dialect/GPU/Pipelines/GPUToXeVMPipeline.cpp
index fa0ab712fa64d..1cbf3bd633633 100644
--- a/mlir/lib/Dialect/GPU/Pipelines/GPUToXeVMPipeline.cpp
+++ b/mlir/lib/Dialect/GPU/Pipelines/GPUToXeVMPipeline.cpp
@@ -100,6 +100,19 @@ void buildGPUPassPipeline(OpPassManager &pm,
     pm.addNestedPass<gpu::GPUModuleOp>(createCanonicalizerPass());
     pm.addNestedPass<gpu::GPUModuleOp>(createCSEPass());
   }
+  // Break down high-level micro-scaling (MX) ops (arith.scaling_extf and
+  // arith.scaling_truncf) into standard arith ops (extf/truncf + mulf), and
+  // expand extf/truncf on f8E8M0FNU into integer bit manipulation. This runs
+  // before the XeVM/LLVM conversions. The f4E2M1FN expansion patterns are
+  // intentionally left disabled: f4E2M1FN extf/truncf are lowered by the XeVM
+  // conversions (xevm.extf), whereas f8E8M0FNU is not supported there and so
+  // must be expanded here.
+  {
+    arith::ArithExpandOpsPassOptions arithExpandOptions;
+    arithExpandOptions.includeF8E8M0 = true;
+    pm.addNestedPass<gpu::GPUModuleOp>(
+        arith::createArithExpandOpsPass(arithExpandOptions));
+  }
   pm.addNestedPass<gpu::GPUModuleOp>(createConvertMathToXeVM());
   ConvertXeGPUToXeVMPassOptions xegpuToXeVMOptions;
   xegpuToXeVMOptions.use64bitIndex = options.use64bitIndex;