[Mlir-commits] [mlir] [mlir] Lower math dialect later in gpu-lower-to-nvvm-pipeline (PR #78556)

[Guray Ozen]

https://github.com/grypp updated https://github.com/llvm/llvm-project/pull/78556

>From 1974e5e7fdc3e5ceaf9fda8191ffb5f0ab083efc Mon Sep 17 00:00:00 2001
From: Guray Ozen <guray.ozen at gmail.com>
Date: Thu, 18 Jan 2024 11:28:16 +0100
Subject: [PATCH 1/2] [mlir] Lower math dialect later in
 gpu-lower-to-nvvm-pipeline

This PR moves lowering of math dialect later in the pipeline. Because math dialect is lowered correctly by `createConvertGpuOpsToNVVMOps` for GPU target, and it needs to run it first.
---
 .../GPU/Pipelines/GPUToNVVMPipeline.cpp       |  2 +-
 mlir/test/Dialect/GPU/test-nvvm-pipeline.mlir | 29 +++++++++++++++++++
 2 files changed, 30 insertions(+), 1 deletion(-)
 create mode 100644 mlir/test/Dialect/GPU/test-nvvm-pipeline.mlir

diff --git a/mlir/lib/Dialect/GPU/Pipelines/GPUToNVVMPipeline.cpp b/mlir/lib/Dialect/GPU/Pipelines/GPUToNVVMPipeline.cpp
index 0b4739214bf2f1..935f0deaf9c8a6 100644
--- a/mlir/lib/Dialect/GPU/Pipelines/GPUToNVVMPipeline.cpp
+++ b/mlir/lib/Dialect/GPU/Pipelines/GPUToNVVMPipeline.cpp
@@ -51,7 +51,6 @@ void buildCommonPassPipeline(
   pm.addPass(createConvertVectorToSCFPass());
   pm.addPass(createConvertSCFToCFPass());
   pm.addPass(createConvertNVVMToLLVMPass());
-  pm.addPass(createConvertMathToLLVMPass());
   pm.addPass(createConvertFuncToLLVMPass());
   pm.addPass(memref::createExpandStridedMetadataPass());
 
@@ -98,6 +97,7 @@ void buildHostPostPipeline(OpPassManager &pm,
   GpuModuleToBinaryPassOptions gpuModuleToBinaryPassOptions;
   gpuModuleToBinaryPassOptions.compilationTarget = options.cubinFormat;
   pm.addPass(createGpuModuleToBinaryPass(gpuModuleToBinaryPassOptions));
+  pm.addPass(createConvertMathToLLVMPass());
   pm.addPass(createCanonicalizerPass());
   pm.addPass(createCSEPass());
   pm.addPass(createReconcileUnrealizedCastsPass());
diff --git a/mlir/test/Dialect/GPU/test-nvvm-pipeline.mlir b/mlir/test/Dialect/GPU/test-nvvm-pipeline.mlir
new file mode 100644
index 00000000000000..83d79409f3fecf
--- /dev/null
+++ b/mlir/test/Dialect/GPU/test-nvvm-pipeline.mlir
@@ -0,0 +1,29 @@
+// RUN: mlir-opt %s \
+// RUN:  | mlir-opt -gpu-lower-to-nvvm-pipeline="cubin-format=isa" \ 
+// RUN:    -split-input-file | FileCheck %s
+
+// RUN: mlir-opt %s \
+// RUN:  | mlir-opt -gpu-lower-to-nvvm-pipeline="cubin-format=isa" \ 
+// RUN:    -split-input-file -debug-only=serialize-to-isa \
+// RUN:    2>&1 | FileCheck %s --check-prefixes=CHECK-PTX
+
+// This test checks whether the GPU region is compiled correctly to PTX by 
+// pipeline. It doesn't test IR for GPU side, but it can test Host IR and 
+// generated PTX.
+
+// CHECK-LABEL: llvm.func @test_math(%arg0: f32) {
+func.func @test_math(%arg0 : f32) {
+    %c2 = arith.constant 2 : index
+    %c1 = arith.constant 1 : index
+    // CHECK: gpu.launch_func  @test_math_kernel::@test_math_kernel
+    // CHECK: gpu.binary @test_math_kernel  [#gpu.object<#nvvm.target
+    gpu.launch 
+        blocks(%0, %1, %2) in (%3 = %c1, %4 = %c1, %5 = %c1) 
+        threads(%6, %7, %8) in (%9 = %c2, %10 = %c1, %11 = %c1) { 
+        // CHECK-PTX: ex2.approx.ftz.f32 
+        %s1 = math.exp %arg0 : f32
+        gpu.printf "%f" %s1 : f32
+        gpu.terminator
+    }
+    return
+}

>From 6cf3768f33a702f3464106ad47b6c4ec41a773cb Mon Sep 17 00:00:00 2001
From: Guray Ozen <guray.ozen at gmail.com>
Date: Fri, 19 Jan 2024 14:10:10 +0100
Subject: [PATCH 2/2] check the right thing

---
 mlir/test/Dialect/GPU/test-nvvm-pipeline.mlir | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/mlir/test/Dialect/GPU/test-nvvm-pipeline.mlir b/mlir/test/Dialect/GPU/test-nvvm-pipeline.mlir
index 83d79409f3fecf..fde99ab105d686 100644
--- a/mlir/test/Dialect/GPU/test-nvvm-pipeline.mlir
+++ b/mlir/test/Dialect/GPU/test-nvvm-pipeline.mlir
@@ -20,7 +20,7 @@ func.func @test_math(%arg0 : f32) {
     gpu.launch 
         blocks(%0, %1, %2) in (%3 = %c1, %4 = %c1, %5 = %c1) 
         threads(%6, %7, %8) in (%9 = %c2, %10 = %c1, %11 = %c1) { 
-        // CHECK-PTX: ex2.approx.ftz.f32 
+        // CHECK-PTX: __nv_expf 
         %s1 = math.exp %arg0 : f32
         gpu.printf "%f" %s1 : f32
         gpu.terminator