[mlir:GPU] Replace reference to LLVMFuncOp with FuncOpInterface

[rriddle]

Summary

The GPU dialect currently contains an explicit reference to LLVMFuncOp
during verification to handle the situation where the kernel has already been
converted. This commit changes that reference to instead use FunctionOpInterface,
which has two main benefits:

• It allows for removing an otherwise unnecessary dependency on the LLVM dialect
• It removes hardcoded assumptions about the lowering path and use of the GPU dialect

[mehdi_amini]

LG

mlir/lib/Dialect/GPU/IR/GPUDialect.cpp

230–239

Since you're here, can you:

1. do the null chech with this error message after module.lookupSymbol
2. use dyn_cast instead of dyn_cast_or_null for the two checks here and error if null with a message that is more accurate.

[mehdi_amini]

LG

mlir/lib/Dialect/GPU/IR/GPUDialect.cpp

250

the cast to kernelGPUFunction can be moved down here, it isn't used before.

[mehdi_amini]

LG, thanks!

mlir/lib/Dialect/GPU/CMakeLists.txt

add_mlir_dialect_library(MLIRGPUOps
 
   LINK_LIBS PUBLIC
   MLIRArithmetic
   MLIRDLTI
   MLIRIR
   MLIRMemRef
   MLIRSideEffectInterfaces
   MLIRSupport
-  MLIRLLVMIR
   )
 
 add_mlir_dialect_library(MLIRGPUTransforms
   Transforms/AllReduceLowering.cpp
   Transforms/AsyncRegionRewriter.cpp
   Transforms/KernelOutlining.cpp
   Transforms/MemoryPromotion.cpp
   Transforms/ParallelLoopMapper.cpp

mlir/lib/Dialect/GPU/IR/GPUDialect.cpp

 //===- GPUDialect.cpp - MLIR Dialect for GPU Kernels implementation -------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This file implements the GPU kernel-related dialect and its operations.
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/GPU/GPUDialect.h"
 
 #include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
-#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/IR/Attributes.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/DialectImplementation.h"
 #include "mlir/IR/FunctionImplementation.h"
 #include "mlir/IR/Matchers.h"
auto walkResult = module.walk([&module](LaunchFuncOp launchOp) -> WalkResult {
     auto kernelModule = module.lookupSymbol<GPUModuleOp>(kernelModuleName);
     if (!kernelModule)
       return launchOp.emitOpError()
              << "kernel module '" << kernelModuleName.getValue()
              << "' is undefined";
 
     // Check that `launch_func` refers to a well-formed kernel function.
     Operation *kernelFunc = module.lookupSymbol(launchOp.kernelAttr());
-    auto kernelGPUFunction = dyn_cast_or_null<gpu::GPUFuncOp>(kernelFunc);
-    auto kernelLLVMFunction = dyn_cast_or_null<LLVM::LLVMFuncOp>(kernelFunc);
-    if (!kernelGPUFunction && !kernelLLVMFunction)
+    if (!kernelFunc)
       return launchOp.emitOpError("kernel function '")
              << launchOp.kernel() << "' is undefined";
+    auto kernelConvertedFunction = dyn_cast<FunctionOpInterface>(kernelFunc);
+    if (!kernelConvertedFunction) {
+      InFlightDiagnostic diag = launchOp.emitOpError()
+                                << "referenced kernel '" << launchOp.kernel()
+                                << "' is not a function";
+      diag.attachNote(kernelFunc->getLoc()) << "see the kernel definition here";
+      return diag;
+    }
+
     if (!kernelFunc->getAttrOfType<mlir::UnitAttr>(
             GPUDialect::getKernelFuncAttrName()))
       return launchOp.emitOpError("kernel function is missing the '")
              << GPUDialect::getKernelFuncAttrName() << "' attribute";
 
-    // TODO: if the kernel function has been converted to
-    // the LLVM dialect but the caller hasn't (which happens during the
-    // separate compilation), do not check type correspondence as it would
-    // require the verifier to be aware of the LLVM type conversion.
-    if (kernelLLVMFunction)
+    // TODO: If the kernel isn't a GPU function (which happens during separate
+    // compilation), do not check type correspondence as it would require the
+    // verifier to be aware of the type conversion.
+    auto kernelGPUFunction = dyn_cast<gpu::GPUFuncOp>(kernelFunc);
+    if (!kernelGPUFunction)
       return success();
 
     unsigned actualNumArguments = launchOp.getNumKernelOperands();
     unsigned expectedNumArguments = kernelGPUFunction.getNumArguments();
     if (expectedNumArguments != actualNumArguments)
       return launchOp.emitOpError("got ")
              << actualNumArguments << " kernel operands but expected "
              << expectedNumArguments;
 

mlir/test/Dialect/GPU/invalid.mlir

func @launch_func_undefined_function(%sz : index) {
     gpu.launch_func @kernels::@kernel_1 blocks in (%sz, %sz, %sz) threads in (%sz, %sz, %sz)
     return
   }
 }
 
 // -----
 
 module attributes {gpu.container_module} {
+  gpu.module @kernels {
+    // expected-note@+1 {{see the kernel definition here}}
+    memref.global "private" @kernel_1 : memref<4xi32>
+  }
+
+  func @launch_func_undefined_function(%sz : index) {
+    // expected-error@+1 {{referenced kernel '@kernels::@kernel_1' is not a function}}
+    gpu.launch_func @kernels::@kernel_1 blocks in (%sz, %sz, %sz) threads in (%sz, %sz, %sz)
+    return
+  }
+}
+
+// -----
+
+module attributes {gpu.container_module} {
   module @kernels {
     gpu.func @kernel_1(%arg1 : !llvm.ptr<f32>) kernel {
       gpu.return
     }
   }
 
   func @launch_func_missing_kernel_attr(%sz : index, %arg : !llvm.ptr<f32>) {
     // expected-error@+1 {{kernel module 'kernels' is undefined}}