[MLIR][MathToXeVM] Remove requirement for ModuleOp op type for MathToXeVM

[ianayl]

This dependency is not needed. This PR removes it and checks that MathToXeVM works and respects pass pipelines specified by the user.

[llvmbot]

@llvm/pr-subscribers-mlir

Author: Ian Li (ianayl)

Changes

This dependency is not needed. This PR removes it and checks that MathToXeVM works and respects pass pipelines specified by the user.

---

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

2 Files Affected:

• (modified) mlir/include/mlir/Conversion/Passes.td (+1-1)
• (modified) mlir/test/Conversion/MathToXeVM/math-to-xevm.mlir (+44)

diff --git a/mlir/include/mlir/Conversion/Passes.td b/mlir/include/mlir/Conversion/Passes.td
index 25e9d34f3e653..20db6c7d3f410 100644
--- a/mlir/include/mlir/Conversion/Passes.td
+++ b/mlir/include/mlir/Conversion/Passes.td
@@ -800,7 +800,7 @@ def ConvertMathToSPIRVPass : Pass<"convert-math-to-spirv"> {
 // MathToXeVM
 //===----------------------------------------------------------------------===//
 
-def ConvertMathToXeVM : Pass<"convert-math-to-xevm", "ModuleOp"> {
+def ConvertMathToXeVM : Pass<"convert-math-to-xevm"> {
   let summary =
       "Convert (fast) math operations to native XeVM/SPIRV equivalents";
   let description = [{
diff --git a/mlir/test/Conversion/MathToXeVM/math-to-xevm.mlir b/mlir/test/Conversion/MathToXeVM/math-to-xevm.mlir
index d76627bb4201c..c61640c2afc4f 100644
--- a/mlir/test/Conversion/MathToXeVM/math-to-xevm.mlir
+++ b/mlir/test/Conversion/MathToXeVM/math-to-xevm.mlir
@@ -3,6 +3,15 @@
 // RUN: mlir-opt %s -convert-math-to-xevm='convert-arith=false' \
 // RUN:   | FileCheck %s -check-prefixes='CHECK,CHECK-NO-ARITH'
 
+// RUN: mlir-opt --pass-pipeline="builtin.module(convert-math-to-xevm)" %s \
+// RUN:   | FileCheck %s -check-prefixes='CHECK-MODULE,CHECK-ENTIRE-MODULE'
+// RUN: mlir-opt --pass-pipeline="builtin.module(gpu.module(convert-math-to-xevm))" %s \
+// RUN:   | FileCheck %s -check-prefixes='CHECK-MODULE,CHECK-ONLY-GPU'
+
+// This test:
+// - check that MathToXeVM converts fastmath math/arith ops properly;
+// - check that MathToXeVM handles nested modules while respecting pass manager.
+
 module @test_module {
   // CHECK-DAG: llvm.func @_Z22__spirv_ocl_native_expDh(f16) -> f16
   // CHECK-DAG: llvm.func @_Z22__spirv_ocl_native_expf(f32) -> f32
@@ -152,4 +161,39 @@ module @test_module {
 
     return
   }
+
+  // Check that MathToXeVM handles nested modules while respecting pass manager:
+
+  // CHECK-ENTIRE-MODULE: llvm.func @_Z22__spirv_ocl_native_expf(f32) -> f32
+  // CHECK-ONLY-GPU-NOT: llvm.func @_Z22__spirv_ocl_native_expf(f32) -> f32
+  
+  // CHECK-MODULE-LABEL: @test_gpu
+  gpu.module @test_gpu {
+    // CHECK-MODULE: llvm.func @_Z22__spirv_ocl_native_expf(f32) -> f32
+    gpu.func @exp_gpu() {
+      %c1_f32 = arith.constant 1. : f32
+
+      // CHECK-MODULE: math.exp
+      %exp_normal_f32 = math.exp %c1_f32 : f32
+
+      // CHECK-MODULE: llvm.call @_Z22__spirv_ocl_native_expf(%{{.*}}) {fastmathFlags = #llvm.fastmath<afn>} : (f32) -> f32
+      %exp_fast_f32 = math.exp %c1_f32 fastmath<afn> : f32
+
+      gpu.return
+    }
+  }
+
+  // CHECK-MODULE-LABEL: @exp_func
+  func.func @exp_func() {
+    %c1_f32 = arith.constant 1. : f32
+
+    // CHECK-MODULE: math.exp
+    %exp_normal_f32 = math.exp %c1_f32 : f32
+
+    // CHECK-ENTIRE-MODULE: llvm.call @_Z22__spirv_ocl_native_expf(%{{.*}}) {fastmathFlags = #llvm.fastmath<afn>} : (f32) -> f32
+    // CHECK-ONLY-GPU: math.exp
+    %exp_fast_f32 = math.exp %c1_f32 fastmath<afn> : f32
+
+    return
+  }
 }