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

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: Sang Ik Lee (silee2)

<details>
<summary>Changes</summary>

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.

---
Full diff: https://github.com/llvm/llvm-project/pull/203632.diff


1 Files Affected:

- (modified) mlir/lib/Dialect/GPU/Pipelines/GPUToXeVMPipeline.cpp (+13) 


``````````diff
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;

``````````

</details>


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