[Mlir-commits] [mlir] [mlir][GPUToNVVM] enable fallback to generic LLVM lowering for math dialect in convert-gpu-to-nvvm pass (PR #165728)

[Yang Bai]

https://github.com/yangtetris updated https://github.com/llvm/llvm-project/pull/165728

>From 79e28d83be55211c9b1a218cb77bc8563415e27c Mon Sep 17 00:00:00 2001
From: Yang Bai <yangb at nvidia.com>
Date: Thu, 30 Oct 2025 07:23:55 -0700
Subject: [PATCH 1/2] enable fallback to generic LLVM lowering for math dialect
 in convert-gpu-to-nvvm pass

---
 .../GPUToNVVM/LowerGpuOpsToNVVMOps.cpp        |  6 ++--
 .../GPUToNVVM/gpu-to-generic-llvm.mlir        | 29 +++++++++++++++++++
 2 files changed, 31 insertions(+), 4 deletions(-)
 create mode 100644 mlir/test/Conversion/GPUToNVVM/gpu-to-generic-llvm.mlir

diff --git a/mlir/lib/Conversion/GPUToNVVM/LowerGpuOpsToNVVMOps.cpp b/mlir/lib/Conversion/GPUToNVVM/LowerGpuOpsToNVVMOps.cpp
index d64c4d64cad84..70c97b3566662 100644
--- a/mlir/lib/Conversion/GPUToNVVM/LowerGpuOpsToNVVMOps.cpp
+++ b/mlir/lib/Conversion/GPUToNVVM/LowerGpuOpsToNVVMOps.cpp
@@ -383,16 +383,14 @@ struct LowerGpuOpsToNVVMOpsPass final
     LLVMConversionTarget target(getContext());
 
     // Set higher benefit, so patterns will run before generic LLVM lowering.
+    // Make sure the benefit here is higher than ArithToLLVMDialectInterface and
+    // MathToLLVMDialectInterface.
     populateGpuToNVVMConversionPatterns(converter, llvmPatterns,
                                         /*benefit=*/10);
 
     llvm::SmallDenseSet<StringRef> allowedDialectsSet(allowedDialects.begin(),
                                                       allowedDialects.end());
     for (Dialect *dialect : getContext().getLoadedDialects()) {
-      // Skip math patterns as nvvm needs custom math lowering.
-      if (isa<math::MathDialect>(dialect))
-        continue;
-
       bool allowed = allowedDialectsSet.contains(dialect->getNamespace());
       // Empty `allowedDialectsSet` means all dialects are allowed.
       if (!allowedDialectsSet.empty() && !allowed)
diff --git a/mlir/test/Conversion/GPUToNVVM/gpu-to-generic-llvm.mlir b/mlir/test/Conversion/GPUToNVVM/gpu-to-generic-llvm.mlir
new file mode 100644
index 0000000000000..5be7938aae8ef
--- /dev/null
+++ b/mlir/test/Conversion/GPUToNVVM/gpu-to-generic-llvm.mlir
@@ -0,0 +1,29 @@
+// RUN: mlir-opt %s -convert-gpu-to-nvvm -split-input-file | FileCheck %s
+
+/// Math/arith ops that are not supported by libdevice
+/// should be converted by generic LLVM lowering patterns.
+
+gpu.module @generic_llvm_test_module_0 {
+  // CHECK-LABEL: @arith_add
+  func.func @arith_add(%left: i64, %right: i64) -> i64 {
+    // CHECK: llvm.add {{.*}}, {{.*}} : i64
+    %result = arith.addi %left, %right : i64
+    return %result : i64
+  }
+}
+
+gpu.module @generic_llvm_test_module_1 {
+  // CHECK-LABEL: @math_abs_non_i32
+  func.func @math_abs_non_i32(%arg_i64: i64, %arg_i16: i16, %arg_i8: i8, %arg_i1: i1) 
+      -> (i64, i16, i8, i1) {
+    // CHECK: "llvm.intr.abs"{{.*}} : (i64) -> i64
+    %abs_i64 = math.absi %arg_i64 : i64
+    // CHECK: "llvm.intr.abs"{{.*}} : (i16) -> i16
+    %abs_i16 = math.absi %arg_i16 : i16
+    // CHECK: "llvm.intr.abs"{{.*}} : (i8) -> i8
+    %abs_i8 = math.absi %arg_i8 : i8
+    // CHECK: "llvm.intr.abs"{{.*}} : (i1) -> i1
+    %abs_i1 = math.absi %arg_i1 : i1
+    return %abs_i64, %abs_i16, %abs_i8, %abs_i1 : i64, i16, i8, i1
+  }
+}

>From 9089dba88a20f071123ce7f80aca7dad4fbacd21 Mon Sep 17 00:00:00 2001
From: Yang Bai <yangb at nvidia.com>
Date: Thu, 30 Oct 2025 09:41:58 -0700
Subject: [PATCH 2/2] remove unrelated test

---
 mlir/test/Conversion/GPUToNVVM/gpu-to-generic-llvm.mlir | 9 ---------
 1 file changed, 9 deletions(-)

diff --git a/mlir/test/Conversion/GPUToNVVM/gpu-to-generic-llvm.mlir b/mlir/test/Conversion/GPUToNVVM/gpu-to-generic-llvm.mlir
index 5be7938aae8ef..4594025bdfc90 100644
--- a/mlir/test/Conversion/GPUToNVVM/gpu-to-generic-llvm.mlir
+++ b/mlir/test/Conversion/GPUToNVVM/gpu-to-generic-llvm.mlir
@@ -4,15 +4,6 @@
 /// should be converted by generic LLVM lowering patterns.
 
 gpu.module @generic_llvm_test_module_0 {
-  // CHECK-LABEL: @arith_add
-  func.func @arith_add(%left: i64, %right: i64) -> i64 {
-    // CHECK: llvm.add {{.*}}, {{.*}} : i64
-    %result = arith.addi %left, %right : i64
-    return %result : i64
-  }
-}
-
-gpu.module @generic_llvm_test_module_1 {
   // CHECK-LABEL: @math_abs_non_i32
   func.func @math_abs_non_i32(%arg_i64: i64, %arg_i16: i16, %arg_i8: i8, %arg_i1: i1) 
       -> (i64, i16, i8, i1) {