[Mlir-commits] [mlir] 4e01532 - Check for FunctionOpInterface when looking up a parent function in GPU lowering

[Mehdi Amini]

Author: Mehdi Amini
Date: 2023-01-16T16:40:11Z
New Revision: 4e015324501dfcb57decd6c0639627c9edfd8c5a

URL: https://github.com/llvm/llvm-project/commit/4e015324501dfcb57decd6c0639627c9edfd8c5a
DIFF: https://github.com/llvm/llvm-project/commit/4e015324501dfcb57decd6c0639627c9edfd8c5a.diff

LOG: Check for FunctionOpInterface when looking up a parent function in GPU lowering

This makes it more robust when expanding code in other function than
func.func, like spv.func for example.

Fixes #60072

Added: 
    

Modified: 
    mlir/lib/Conversion/GPUCommon/OpToFuncCallLowering.h
    mlir/test/Conversion/GPUToROCDL/gpu-to-rocdl.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Conversion/GPUCommon/OpToFuncCallLowering.h b/mlir/lib/Conversion/GPUCommon/OpToFuncCallLowering.h
index 0168e6e016449..6858569862535 100644
--- a/mlir/lib/Conversion/GPUCommon/OpToFuncCallLowering.h
+++ b/mlir/lib/Conversion/GPUCommon/OpToFuncCallLowering.h
@@ -107,7 +107,7 @@ struct OpToFuncCallLowering : public ConvertOpToLLVMPattern<SourceOp> {
     if (funcOp)
       return cast<LLVMFuncOp>(*funcOp);
 
-    mlir::OpBuilder b(op->getParentOfType<LLVMFuncOp>());
+    mlir::OpBuilder b(op->getParentOfType<FunctionOpInterface>());
     return b.create<LLVMFuncOp>(op->getLoc(), funcName, funcType);
   }
 

diff  --git a/mlir/test/Conversion/GPUToROCDL/gpu-to-rocdl.mlir b/mlir/test/Conversion/GPUToROCDL/gpu-to-rocdl.mlir
index ccf8e47f8da0b..777b03a0cf19f 100644
--- a/mlir/test/Conversion/GPUToROCDL/gpu-to-rocdl.mlir
+++ b/mlir/test/Conversion/GPUToROCDL/gpu-to-rocdl.mlir
@@ -460,3 +460,15 @@ gpu.module @test_module {
     gpu.return
   }
 }
+
+// ----
+
+gpu.module @module {
+// CHECK-LABEL: @spirv_exp
+// CHECK: llvm.call @__ocml_exp_f32
+  spirv.func @spirv_exp(%arg0: vector<4xf32>) -> vector<4xf32> "None" {
+    %0 = math.exp %arg0 : vector<4xf32>
+    spirv.ReturnValue %0 : vector<4xf32>
+  }
+}
+