[Mlir-commits] [mlir] 663e9ce - [Func][GPU] Use SymbolUserOpInterface in func::ConstantOp (#107748)

[llvmlistbot at llvm.org]

Author: Artem Kroviakov
Date: 2024-09-09T11:49:16+02:00
New Revision: 663e9cec9c96169aa4e72ab9b6bf08b2d6603093

URL: https://github.com/llvm/llvm-project/commit/663e9cec9c96169aa4e72ab9b6bf08b2d6603093
DIFF: https://github.com/llvm/llvm-project/commit/663e9cec9c96169aa4e72ab9b6bf08b2d6603093.diff

LOG: [Func][GPU] Use SymbolUserOpInterface in func::ConstantOp  (#107748)

This PR enables `func::ConstantOp` creation and usage for device
functions inside GPU modules.
The current main returns error for referencing device functions via
`func::ConstantOp`, because during the `ConstantOp` verification it only
checks symbols in `ModuleOp` symbol table, which, of course, does not
contain device functions that are defined in `GPUModuleOp`. This PR
proposes a more general solution.

Co-authored-by: Artem Kroviakov <artem.kroviakov at tum.de>

Added: 
    mlir/test/Dialect/GPU/indirect-device-func-call.mlir

Modified: 
    mlir/include/mlir/Dialect/Func/IR/FuncOps.td
    mlir/lib/Dialect/Func/IR/FuncOps.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/Dialect/Func/IR/FuncOps.td b/mlir/include/mlir/Dialect/Func/IR/FuncOps.td
index f4b43a59047054..cca0008cbc2633 100644
--- a/mlir/include/mlir/Dialect/Func/IR/FuncOps.td
+++ b/mlir/include/mlir/Dialect/Func/IR/FuncOps.td
@@ -183,6 +183,7 @@ def CallIndirectOp : Func_Op<"call_indirect", [
 
 def ConstantOp : Func_Op<"constant",
     [ConstantLike, Pure,
+     DeclareOpInterfaceMethods<SymbolUserOpInterface>,
      DeclareOpInterfaceMethods<OpAsmOpInterface, ["getAsmResultNames"]>]> {
   let summary = "constant";
   let description = [{
@@ -216,7 +217,6 @@ def ConstantOp : Func_Op<"constant",
   }];
 
   let hasFolder = 1;
-  let hasVerifier = 1;
 }
 
 //===----------------------------------------------------------------------===//

diff  --git a/mlir/lib/Dialect/Func/IR/FuncOps.cpp b/mlir/lib/Dialect/Func/IR/FuncOps.cpp
index c719981769b9e1..a490b4c3c4ab43 100644
--- a/mlir/lib/Dialect/Func/IR/FuncOps.cpp
+++ b/mlir/lib/Dialect/Func/IR/FuncOps.cpp
@@ -123,12 +123,13 @@ LogicalResult CallIndirectOp::canonicalize(CallIndirectOp indirectCall,
 // ConstantOp
 //===----------------------------------------------------------------------===//
 
-LogicalResult ConstantOp::verify() {
+LogicalResult ConstantOp::verifySymbolUses(SymbolTableCollection &symbolTable) {
   StringRef fnName = getValue();
   Type type = getType();
 
   // Try to find the referenced function.
-  auto fn = (*this)->getParentOfType<ModuleOp>().lookupSymbol<FuncOp>(fnName);
+  auto fn = symbolTable.lookupNearestSymbolFrom<FuncOp>(
+      this->getOperation(), StringAttr::get(getContext(), fnName));
   if (!fn)
     return emitOpError() << "reference to undefined function '" << fnName
                          << "'";

diff  --git a/mlir/test/Dialect/GPU/indirect-device-func-call.mlir b/mlir/test/Dialect/GPU/indirect-device-func-call.mlir
new file mode 100644
index 00000000000000..91d7f1cd6c67d9
--- /dev/null
+++ b/mlir/test/Dialect/GPU/indirect-device-func-call.mlir
@@ -0,0 +1,21 @@
+// RUN: mlir-opt -test-gpu-rewrite -convert-func-to-llvm %s | FileCheck %s
+
+gpu.module @kernels {
+    // CHECK-LABEL: @hello
+    // CHECK-SAME: %[[ARG0:.*]]: f32
+    func.func @hello(%arg0 : f32) {
+        %tid_x = gpu.thread_id x
+        %csti8 = arith.constant 2 : i8
+        gpu.printf "Hello from %lld, %d, %f\n" %tid_x, %csti8, %arg0  : index, i8, f32
+        return
+    }
+    // CHECK-LABEL: @hello_indirect
+    gpu.func @hello_indirect() kernel { 
+        %cstf32 = arith.constant 3.0 : f32
+        // CHECK: %[[DEVICE_FUNC_ADDR:.*]] = llvm.mlir.addressof @hello : !llvm.ptr
+        %func_ref = func.constant @hello : (f32) -> ()
+        // CHECK: llvm.call %[[DEVICE_FUNC_ADDR]](%{{.*}}) : !llvm.ptr, (f32) -> ()
+        func.call_indirect %func_ref(%cstf32) : (f32) -> ()
+        gpu.return
+    }
+}