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

diff --git a/mlir/lib/Conversion/GPUCommon/OpToFuncCallLowering.h b/mlir/lib/Conversion/GPUCommon/OpToFuncCallLowering.h
index 0168e6e..6858569 100644 (file)
--- a/mlir/lib/Conversion/GPUCommon/OpToFuncCallLowering.h
+++ b/mlir/lib/Conversion/GPUCommon/OpToFuncCallLowering.h
@@ -107,7 +107,7 @@ private:
     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 ccf8e47..777b03a 100644 (file)
--- 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>
+  }
+}
+