[Mlir-commits] [mlir] c0ccb69 - [mlir][spirv] Convert func.call to spv.FunctionCall

[llvmlistbot at llvm.org]

Author: xndcn
Date: 2022-03-26T19:21:23+08:00
New Revision: c0ccb692285d1fb411f90d4d68f834719be4d009

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

LOG: [mlir][spirv] Convert func.call to spv.FunctionCall

Differential Revision: https://reviews.llvm.org/D122368

Added: 
    

Modified: 
    mlir/lib/Conversion/FuncToSPIRV/FuncToSPIRV.cpp
    mlir/test/Conversion/FuncToSPIRV/func-ops-to-spirv.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Conversion/FuncToSPIRV/FuncToSPIRV.cpp b/mlir/lib/Conversion/FuncToSPIRV/FuncToSPIRV.cpp
index 0ae25eca00a02..5207f3fd094ee 100644
--- a/mlir/lib/Conversion/FuncToSPIRV/FuncToSPIRV.cpp
+++ b/mlir/lib/Conversion/FuncToSPIRV/FuncToSPIRV.cpp
@@ -57,6 +57,32 @@ class ReturnOpPattern final : public OpConversionPattern<func::ReturnOp> {
   }
 };
 
+/// Converts func.call to spv.FunctionCall.
+class CallOpPattern final : public OpConversionPattern<func::CallOp> {
+public:
+  using OpConversionPattern<func::CallOp>::OpConversionPattern;
+
+  LogicalResult
+  matchAndRewrite(func::CallOp callOp, OpAdaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override {
+    // multiple results func was not converted to spv.func
+    if (callOp.getNumResults() > 1)
+      return failure();
+    if (callOp.getNumResults() == 1) {
+      auto resultType =
+          getTypeConverter()->convertType(callOp.getResult(0).getType());
+      if (!resultType)
+        return failure();
+      rewriter.replaceOpWithNewOp<spirv::FunctionCallOp>(
+          callOp, resultType, adaptor.getOperands(), callOp->getAttrs());
+    } else {
+      rewriter.replaceOpWithNewOp<spirv::FunctionCallOp>(
+          callOp, TypeRange(), adaptor.getOperands(), callOp->getAttrs());
+    }
+    return success();
+  }
+};
+
 } // namespace
 
 //===----------------------------------------------------------------------===//
@@ -67,5 +93,5 @@ void mlir::populateFuncToSPIRVPatterns(SPIRVTypeConverter &typeConverter,
                                        RewritePatternSet &patterns) {
   MLIRContext *context = patterns.getContext();
 
-  patterns.add<ReturnOpPattern>(typeConverter, context);
+  patterns.add<ReturnOpPattern, CallOpPattern>(typeConverter, context);
 }

diff  --git a/mlir/test/Conversion/FuncToSPIRV/func-ops-to-spirv.mlir b/mlir/test/Conversion/FuncToSPIRV/func-ops-to-spirv.mlir
index 43f6583b08e89..17329e16f40c9 100644
--- a/mlir/test/Conversion/FuncToSPIRV/func-ops-to-spirv.mlir
+++ b/mlir/test/Conversion/FuncToSPIRV/func-ops-to-spirv.mlir
@@ -8,6 +8,12 @@ module attributes {
   spv.target_env = #spv.target_env<#spv.vce<v1.0, [], []>, {}>
 } {
 
+// CHECK-LABEL: spv.func @return_none_val
+func @return_none_val() {
+  // CHECK: spv.Return
+  return
+}
+
 // CHECK-LABEL: spv.func @return_one_val
 //  CHECK-SAME: (%[[ARG:.+]]: f32)
 func @return_one_val(%arg0: f32) -> f32 {
@@ -22,6 +28,24 @@ func @return_multi_val(%arg0: f32) -> (f32, f32) {
   return %arg0, %arg0: f32, f32
 }
 
+// CHECK-LABEL: spv.func @return_one_index
+//  CHECK-SAME: (%[[ARG:.+]]: i32)
+func @return_one_index(%arg0: index) -> index {
+  // CHECK: spv.ReturnValue %[[ARG]] : i32
+  return %arg0: index
+}
+
+// CHECK-LABEL: spv.func @call_functions
+//  CHECK-SAME: (%[[ARG:.+]]: i32)
+func @call_functions(%arg0: index) -> index {
+  // CHECK: spv.FunctionCall @return_none_val() : () -> ()
+  call @return_none_val(): () -> ()
+  // CHECK: {{%.*}} = spv.FunctionCall @return_one_index(%[[ARG]]) : (i32) -> i32
+  %0 = call @return_one_index(%arg0): (index) -> index
+  // CHECK: spv.ReturnValue {{%.*}} : i32
+  return %0: index
+}
+
 }
 
 // -----