[Mlir-commits] [mlir] [GPUToXeVMPipeline] Add unsupported data type emulation on `math`, `arith`, and `vector` ops. (PR #197779)

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir-gpu

Author: Md Abdullah Shahneous Bari (mshahneo)

<details>
<summary>Changes</summary>

Add `math-extend-to-supported-types` and `arith-emulate-unsupported-floats` pass to the pipeline. These passes are used to emulate `math`, `arith` and `vector` floating point operations that use float types which are unspported on a target by inserting `extf/truncf` pairs around all such operations.

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


2 Files Affected:

- (modified) mlir/include/mlir/Dialect/GPU/Pipelines/Passes.h (+26) 
- (modified) mlir/lib/Dialect/GPU/Pipelines/GPUToXeVMPipeline.cpp (+23) 


``````````diff
diff --git a/mlir/include/mlir/Dialect/GPU/Pipelines/Passes.h b/mlir/include/mlir/Dialect/GPU/Pipelines/Passes.h
index 6263ea63cbf22..222ebf7179a5c 100644
--- a/mlir/include/mlir/Dialect/GPU/Pipelines/Passes.h
+++ b/mlir/include/mlir/Dialect/GPU/Pipelines/Passes.h
@@ -161,6 +161,32 @@ struct GPUToXeVMPipelineOptions
       *this, "igc-cmd-options",
       llvm::cl::desc("Additional downstream compiler command line options"),
       llvm::cl::init("")};
+  // Options for MathExtendToSupportedTypes and ArithEmulateUnsupportedFloats,
+  // run on the host/device module before conversion to LLVM to legalize math
+  // and arith ops operating on floating-point types that the XeVM target
+  // cannot handle natively (e.g. bf16).
+  PassOptions::ListOption<std::string> unsupportedSourceTypes{
+      *this, "unsupported-source-types",
+      llvm::cl::desc("Floating-point source types without arithmetic/math "
+                     "support on the target (e.g. bf16)"),
+      llvm::cl::list_init<std::string>(ArrayRef<std::string>{"bf16"})};
+  PassOptions::Option<std::string> supportedTargetTypes{
+      *this, "supported-target-types",
+      llvm::cl::desc(
+          "Floating-point target type used to emulate the unsupported "
+          "source types via extf/truncf pairs"),
+      llvm::cl::init("f32")};
+  // Additional types (beyond the implicit f32/f64) that math ops are allowed
+  // to run on directly. Maps to the `extra-types` option of
+  // math-extend-to-supported-types. Leave empty to extend every non-f32/f64
+  // math op to `supported-target-types`.
+  PassOptions::ListOption<std::string> mathExtendExtraTypes{
+      *this, "math-extend-extra-types",
+      llvm::cl::desc(
+          "Extra floating-point types with math op support on the target, "
+          "in addition to f32 and f64 (maps to math-extend-to-supported-"
+          "types `extra-types`)"),
+      llvm::cl::list_init<std::string>(ArrayRef<std::string>{"f16"})};
 };
 
 //===----------------------------------------------------------------------===//
diff --git a/mlir/lib/Dialect/GPU/Pipelines/GPUToXeVMPipeline.cpp b/mlir/lib/Dialect/GPU/Pipelines/GPUToXeVMPipeline.cpp
index 7600ec39fb3f5..a0322a28fe73e 100644
--- a/mlir/lib/Dialect/GPU/Pipelines/GPUToXeVMPipeline.cpp
+++ b/mlir/lib/Dialect/GPU/Pipelines/GPUToXeVMPipeline.cpp
@@ -20,12 +20,15 @@
 #include "mlir/Conversion/VectorToSCF/VectorToSCF.h"
 #include "mlir/Conversion/XeGPUToXeVM/XeGPUToXeVM.h"
 #include "mlir/Conversion/XeVMToLLVM/XeVMToLLVM.h"
+#include "mlir/Dialect/Arith/Transforms/Passes.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/GPU/IR/GPUDialect.h"
 #include "mlir/Dialect/GPU/Pipelines/Passes.h"
 #include "mlir/Dialect/GPU/Transforms/Passes.h"
 #include "mlir/Dialect/LLVMIR/Transforms/RequestCWrappers.h"
+#include "mlir/Dialect/Math/Transforms/Passes.h"
 #include "mlir/Dialect/MemRef/Transforms/Passes.h"
+#include "mlir/Dialect/Vector/Transforms/Passes.h"
 #include "mlir/Dialect/XeGPU/Transforms/Passes.h"
 #include "mlir/Pass/PassManager.h"
 #include "mlir/Pass/PassOptions.h"
@@ -108,6 +111,26 @@ void buildGPUPassPipeline(OpPassManager &pm,
     pm.addNestedPass<gpu::GPUModuleOp>(
         createConvertGpuOpsToLLVMSPVOps(gpuToLLVMSPVOptions));
   }
+  // Legalize math/arith ops on floating-point types that the XeVM target
+  // cannot handle natively (e.g. bf16) by wrapping them with extf/truncf
+  // around a supported type (defaulting to f32).
+  {
+    math::MathExtendToSupportedTypesOptions mathExtendOptions;
+    mathExtendOptions.extraTypeStrs.assign(options.mathExtendExtraTypes.begin(),
+                                           options.mathExtendExtraTypes.end());
+    mathExtendOptions.targetTypeStr = options.supportedTargetTypes;
+    pm.addNestedPass<gpu::GPUModuleOp>(
+        math::createMathExtendToSupportedTypes(mathExtendOptions));
+  }
+  {
+    arith::ArithEmulateUnsupportedFloatsOptions arithEmulateOptions;
+    arithEmulateOptions.sourceTypeStrs.assign(
+        options.unsupportedSourceTypes.begin(),
+        options.unsupportedSourceTypes.end());
+    arithEmulateOptions.targetTypeStr = options.supportedTargetTypes;
+    pm.addNestedPass<gpu::GPUModuleOp>(
+        arith::createArithEmulateUnsupportedFloats(arithEmulateOptions));
+  }
   pm.addNestedPass<gpu::GPUModuleOp>(createCSEPass());
   pm.addNestedPass<gpu::GPUModuleOp>(createReconcileUnrealizedCastsPass());
 }

``````````

</details>


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