Make legality check in GPU->SPIR-V lowering of FuncOp kernel specific.

Existing check that sets FuncOp to be dynamically legal was just
checking that the types of the argument are SPIR-V compatible. Since
the current conversion from GPU to SPIR-V does not handle lowering
non-kernel functions, change the legality check to verify that the
FuncOp has the gpu.kernel attribute and has void(void) return type.

PiperOrigin-RevId: 280032782

diff --git a/mlir/lib/Conversion/GPUToSPIRV/GPUToSPIRV.cpp b/mlir/lib/Conversion/GPUToSPIRV/GPUToSPIRV.cpp
index ece4269..bf35c21 100644 (file)
--- a/mlir/lib/Conversion/GPUToSPIRV/GPUToSPIRV.cpp
+++ b/mlir/lib/Conversion/GPUToSPIRV/GPUToSPIRV.cpp
@@ -241,8 +241,13 @@ void GPUToSPIRVPass::runOnModule() {
 
   ConversionTarget target(*context);
   target.addLegalDialect<spirv::SPIRVDialect>();
-  target.addDynamicallyLegalOp<FuncOp>([&](FuncOp Op) {
-    return basicTypeConverter.isSignatureLegal(Op.getType());
+  target.addDynamicallyLegalOp<FuncOp>([&](FuncOp op) {
+    // TODO(ravishankarm) : Currently lowering does not support handling
+    // function conversion of non-kernel functions. This is to be added.
+
+    // For kernel functions, verify that the signature is void(void).
+    return gpu::GPUDialect::isKernel(op) && op.getNumResults() == 0 &&
+           op.getNumArguments() == 0;
   });
 
   if (failed(applyFullConversion(spirvModules, target, patterns,