[MLIR][TOSA] Resubmit Tosa to Standard/SCF Lowerings (const, if, while)"

[rsuderman]

Summary

Includes a lowering for tosa.const, tosa.if, and tosa.while to Standard/SCF dialects. TosaToStandard is
used for constant lowerings and TosaToSCF handles the if/while ops.

Resubmission of https://reviews.llvm.org/D97518 with ASAN fixes.

[stellaraccident]

Lgtm, but let's get Mehdi's eyes because he spotted the bug I was missing.

[rsuderman]

Clang-tidy fixes.

[mehdi_amini]

mlir/lib/Conversion/TosaToSCF/TosaToSCF.cpp

54

Aren't the region "single block" regions? If so can we "not loop" here?

56

This should be the terminator of the block, can we just get it directly instead of walking the block?

I'm also concerned about nested cases here: you'll catch YieldOp that are potentially nested in the IR while you should only touch the region you're converting here.

71

This delayed deleting shouldn't be needed, you should be able to erase when replacing.

[rsuderman]

Updated to expect only a single block and handle a single terminator. Also referred delayed deletion.

[mehdi_amini]

Nice, I think the code looks much simpler :)

mlir/include/mlir/Conversion/Passes.h

 #include "mlir/Conversion/SCFToOpenMP/SCFToOpenMP.h"
 #include "mlir/Conversion/SCFToSPIRV/SCFToSPIRVPass.h"
 #include "mlir/Conversion/SCFToStandard/SCFToStandard.h"
 #include "mlir/Conversion/SPIRVToLLVM/SPIRVToLLVMPass.h"
 #include "mlir/Conversion/ShapeToStandard/ShapeToStandard.h"
 #include "mlir/Conversion/StandardToLLVM/ConvertStandardToLLVMPass.h"
 #include "mlir/Conversion/StandardToSPIRV/StandardToSPIRVPass.h"
 #include "mlir/Conversion/TosaToLinalg/TosaToLinalg.h"
+#include "mlir/Conversion/TosaToSCF/TosaToSCF.h"
+#include "mlir/Conversion/TosaToStandard/TosaToStandard.h"
 #include "mlir/Conversion/VectorToLLVM/ConvertVectorToLLVM.h"
 #include "mlir/Conversion/VectorToROCDL/VectorToROCDL.h"
 #include "mlir/Conversion/VectorToSCF/VectorToSCF.h"
 #include "mlir/Conversion/VectorToSPIRV/VectorToSPIRVPass.h"
 
 namespace mlir {
 
 /// Generate the code for registering conversion passes.
 #define GEN_PASS_REGISTRATION
 #include "mlir/Conversion/Passes.h.inc"
 
 } // namespace mlir
 
 #endif // MLIR_CONVERSION_PASSES_H

mlir/include/mlir/Conversion/Passes.td

let description = [{
     Pass that converts TOSA operations to the equivalent operations using the
     tensor operations in LinAlg.
   }];
 
   let constructor = "tosa::createTosaToLinalgOnTensors()";
 }
 
 //===----------------------------------------------------------------------===//
+// TosaToSCF
+//===----------------------------------------------------------------------===//
+
+def TosaToSCF : Pass<"tosa-to-scf"> {
+  let summary = "Lower TOSA to the SCF dialect";
+  let dependentDialects = ["tensor::TensorDialect, scf::SCFDialect"];
+  let description = [{
+    Pass that converts TOSA's control flow operations to the equivalent SCF
+    operations.
+  }];
+
+  let constructor = "tosa::createTosaToSCF()";
+}
+
+//===----------------------------------------------------------------------===//
+// TosaToStandard
+//===----------------------------------------------------------------------===//
+
+def TosaToStandard : Pass<"tosa-to-standard"> {
+  let summary = "Lower TOSA to the Standard dialect";
+  let dependentDialects = ["StandardOpsDialect"];
+  let description = [{
+    Pass that converts TOSA operations to the equivalent operations using the
+    operations in the Standard dialect.
+  }];
+
+  let constructor = "tosa::createTosaToStandard()";
+}
+
+//===----------------------------------------------------------------------===//
 // VectorToSCF
 //===----------------------------------------------------------------------===//
 
 def ConvertVectorToSCF : FunctionPass<"convert-vector-to-scf"> {
   let summary = "Lower the operations from the vector dialect into the SCF "
                 "dialect";
   let constructor = "mlir::createConvertVectorToSCFPass()";
   let dependentDialects = ["AffineDialect", "scf::SCFDialect"];

mlir/include/mlir/Conversion/TosaToSCF/TosaToSCF.h

+//===-- TosaToSCF.h - TOSA to SCF dialect lowerings -------------*- C++ -*-===//
+//
+// 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 declares the passes for the TOSA to SCF Dialect conversion.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_CONVERSION_TOSATOSCF_TOSATOSCF_H
+#define MLIR_CONVERSION_TOSATOSCF_TOSATOSCF_H
+
+#include "mlir/Pass/Pass.h"
+
+namespace mlir {
+namespace tosa {
+
+std::unique_ptr<Pass> createTosaToSCF();
+
+void populateTosaToSCFConversionPatterns(MLIRContext *context,
+                                         OwningRewritePatternList *patterns);
+
+/// Populates passes to convert from TOSA to SCF.
+void addTosaToSCFPasses(OpPassManager &pm);
+
+} // namespace tosa
+} // namespace mlir
+
+#endif // MLIR_CONVERSION_TOSATOSCF_TOSATOSCF_H

mlir/include/mlir/Conversion/TosaToStandard/TosaToStandard.h

+//===-- TosaToStandard.h - TOSA optimization pass declarations --*- C++ -*-===//
+//
+// 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 declares the passes for the TOSA to Standard Dialect conversion.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_CONVERSION_TOSATOSTANDARD_TOSATOSTANDARD_H
+#define MLIR_CONVERSION_TOSATOSTANDARD_TOSATOSTANDARD_H
+
+#include "mlir/Pass/Pass.h"
+
+namespace mlir {
+namespace tosa {
+
+std::unique_ptr<Pass> createTosaToStandard();
+
+void populateTosaToStandardConversionPatterns(
+    MLIRContext *context, OwningRewritePatternList *patterns);
+
+/// Populates passes to convert from TOSA to Standard.
+void addTosaToStandardPasses(OpPassManager &pm);
+
+} // namespace tosa
+} // namespace mlir
+
+#endif // MLIR_CONVERSION_TOSATOSTANDARD_TOSATOSTANDARD_H

mlir/lib/Conversion/CMakeLists.txt

 add_subdirectory(SCFToOpenMP)
 add_subdirectory(SCFToSPIRV)
 add_subdirectory(SCFToStandard)
 add_subdirectory(ShapeToStandard)
 add_subdirectory(SPIRVToLLVM)
 add_subdirectory(StandardToLLVM)
 add_subdirectory(StandardToSPIRV)
 add_subdirectory(TosaToLinalg)
+add_subdirectory(TosaToSCF)
+add_subdirectory(TosaToStandard)
 add_subdirectory(ArmSVEToLLVM)
 add_subdirectory(VectorToROCDL)
 add_subdirectory(VectorToLLVM)
 add_subdirectory(VectorToSCF)
 add_subdirectory(VectorToSPIRV)

mlir/lib/Conversion/PassDetail.h

 namespace scf {
 class SCFDialect;
 } // end namespace scf
 
 namespace spirv {
 class SPIRVDialect;
 } // end namespace spirv
 
+namespace tensor {
+class TensorDialect;
+} // end namespace tensor
+
+namespace tosa {
+class TosaDialect;
+} // end namespace tosa
+
 namespace vector {
 class VectorDialect;
 } // end namespace vector
 
 #define GEN_PASS_CLASSES
 #include "mlir/Conversion/Passes.h.inc"
 
 } // end namespace mlir
 
 #endif // CONVERSION_PASSDETAIL_H_

mlir/lib/Conversion/TosaToSCF/CMakeLists.txt

+add_mlir_conversion_library(MLIRTosaToSCF
+  TosaToSCF.cpp
+  TosaToSCFPass.cpp
+
+  ADDITIONAL_HEADER_DIRS
+  ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/Tosa
+  ${MLIR_MAIN_INCLUDE_DIR}/mlir/IR
+
+  DEPENDS
+  MLIRConversionPassIncGen
+
+  LINK_LIBS PUBLIC
+  MLIRIR
+  MLIRSCF
+  MLIRStandard
+  MLIRPass
+  MLIRTensor
+  MLIRTosa
+  MLIRTosaTransforms
+  MLIRSupport
+  )

mlir/lib/Conversion/TosaToSCF/TosaToSCF.cpp

+//===- TosaToSCF.cpp - Lowering Tosa to SCF Dialect -----------------------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// These rewriters lower from the Tosa to the SCF dialect.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Conversion/TosaToSCF/TosaToSCF.h"
+#include "mlir/Dialect/SCF/SCF.h"
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/Dialect/Tosa/IR/TosaOps.h"
+#include "mlir/IR/BlockAndValueMapping.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+
+using namespace mlir;
+using namespace tosa;
+
+static void inlineIfCase(Region &srcRegion, Region &dstRegion,
+                         OperandRange operands, PatternRewriter &rewriter) {
+  rewriter.cloneRegionBefore(srcRegion, &dstRegion.front());
+  rewriter.eraseBlock(&dstRegion.back());
+
+  Block *headBlock = &dstRegion.front();
+  for (auto it : llvm::zip(headBlock->getArguments(), operands))
+    std::get<0>(it).replaceAllUsesWith(std::get<1>(it));
+
+  auto yield = cast<YieldOp>(headBlock->getTerminator());
+  rewriter.setInsertionPoint(yield);
+  rewriter.create<scf::YieldOp>(yield.getLoc(), yield.inputs());
+  rewriter.eraseOp(yield);
+
+  headBlock->eraseArguments(
+      llvm::to_vector<4>(llvm::seq<unsigned>(0, headBlock->getNumArguments())));
+}
+
+static void inlineWhileCase(Region &srcRegion, Region &dstRegion,
+                            PatternRewriter &rewriter, bool isCond) {
+  rewriter.cloneRegionBefore(srcRegion, &dstRegion.back());
+  rewriter.eraseBlock(&dstRegion.back());
+
+  Block *headBlock = &dstRegion.front();
+
+  auto yield = cast<YieldOp>(headBlock->getTerminator());
+  rewriter.setInsertionPoint(yield);
+  if (isCond) {
+    auto condition =
+        rewriter.create<tensor::ExtractOp>(yield.getLoc(), yield.getOperand(0));
+    rewriter.create<scf::ConditionOp>(yield.getLoc(), condition,
+                                      headBlock->getArguments());
+  } else {
+    rewriter.setInsertionPoint(yield);
+    rewriter.create<scf::YieldOp>(yield.getLoc(), yield.inputs());
+  }
+  rewriter.eraseOp(yield);
+}
+
+namespace {
+
+class IfOpConverter : public OpRewritePattern<tosa::IfOp> {
+public:
+  using OpRewritePattern<tosa::IfOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(tosa::IfOp op,
+                                PatternRewriter &rewriter) const final {
+    auto condition = rewriter.create<tensor::ExtractOp>(op.getLoc(), op.cond());
+    auto newIf = rewriter.create<scf::IfOp>(op.getLoc(), op.getResultTypes(),
+                                            condition, true);
+
+    inlineIfCase(op.then_branch(), newIf.thenRegion(), op.inputs(), rewriter);
+    inlineIfCase(op.else_branch(), newIf.elseRegion(), op.inputs(), rewriter);
+
+    rewriter.replaceOp(op, newIf.getResults());
+    return success();
+  }
+};
+
+class WhileOpConverter : public OpRewritePattern<tosa::WhileOp> {
+public:
+  using OpRewritePattern<tosa::WhileOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(tosa::WhileOp op,
+                                PatternRewriter &rewriter) const final {
+    auto newWhile = rewriter.create<scf::WhileOp>(
+        op.getLoc(), op.getResultTypes(), op.inputs());
+    rewriter.createBlock(&newWhile.before());
+    rewriter.createBlock(&newWhile.after());
+
+    inlineWhileCase(op.cond(), newWhile.before(), rewriter, true);
+    inlineWhileCase(op.body(), newWhile.after(), rewriter, false);
+
+    rewriter.replaceOp(op, newWhile.getResults());
+
+    return success();
+  }
+};
+
+} // namespace
+
+void mlir::tosa::populateTosaToSCFConversionPatterns(
+    MLIRContext *context, OwningRewritePatternList *patterns) {
+  patterns->insert<IfOpConverter>(context);
+  patterns->insert<WhileOpConverter>(context);
+}

mlir/lib/Conversion/TosaToSCF/TosaToSCFPass.cpp

+//===- TosaToSCFPass.cpp - Lowering Tosa to SCF Dialect -------------------===//
+//
+// 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 transformation pass legalizes Tosa operations to the SCF dialect.
+//
+//===----------------------------------------------------------------------===//
+
+#include "../PassDetail.h"
+#include "mlir/Conversion/TosaToSCF/TosaToSCF.h"
+#include "mlir/Dialect/SCF/SCF.h"
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/Dialect/Tosa/IR/TosaOps.h"
+#include "mlir/Dialect/Tosa/Transforms/PassDetail.h"
+#include "mlir/Dialect/Tosa/Transforms/Passes.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/Pass/PassManager.h"
+#include "mlir/Transforms/DialectConversion.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+
+using namespace mlir;
+using namespace tosa;
+
+namespace {
+struct TosaToSCF : public TosaToSCFBase<TosaToSCF> {
+public:
+  void runOnOperation() override {
+    OwningRewritePatternList patterns;
+    ConversionTarget target(getContext());
+    target.addLegalDialect<tensor::TensorDialect, scf::SCFDialect>();
+    target.addIllegalOp<tosa::IfOp, tosa::WhileOp>();
+    target.markUnknownOpDynamicallyLegal([](Operation *) { return true; });
+
+    auto *op = getOperation();
+    mlir::tosa::populateTosaToSCFConversionPatterns(op->getContext(),
+                                                    &patterns);
+    if (failed(applyPartialConversion(op, target, std::move(patterns))))
+      signalPassFailure();
+  }
+};
+} // namespace
+
+std::unique_ptr<Pass> mlir::tosa::createTosaToSCF() {
+  return std::make_unique<TosaToSCF>();
+}
+
+void mlir::tosa::addTosaToSCFPasses(OpPassManager &pm) {
+  pm.addNestedPass<FuncOp>(createTosaToSCF());
+}

mlir/lib/Conversion/TosaToStandard/CMakeLists.txt

+add_mlir_conversion_library(MLIRTosaToStandard
+  TosaToStandard.cpp
+  TosaToStandardPass.cpp
+
+  ADDITIONAL_HEADER_DIRS
+  ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/Tosa
+  ${MLIR_MAIN_INCLUDE_DIR}/mlir/IR
+
+  DEPENDS
+  MLIRConversionPassIncGen
+
+  LINK_LIBS PUBLIC
+  MLIRIR
+  MLIRStandard
+  MLIRPass
+  MLIRTosa
+  MLIRTosaTransforms
+  MLIRSupport
+  )

mlir/lib/Conversion/TosaToStandard/TosaToStandard.cpp

+//===- TosaToStandard.cpp - Lowering Tosa to Standard Dialect -------------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// These rewriters lower from the Tosa to the Standard dialect.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Conversion/TosaToStandard/TosaToStandard.h"
+#include "mlir/Dialect/StandardOps/IR/Ops.h"
+#include "mlir/Dialect/Tosa/IR/TosaOps.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+
+using namespace mlir;
+using namespace tosa;
+
+namespace {
+
+class ConstOpConverter : public OpRewritePattern<tosa::ConstOp> {
+public:
+  using OpRewritePattern<tosa::ConstOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(tosa::ConstOp op,
+                                PatternRewriter &rewriter) const final {
+    rewriter.replaceOpWithNewOp<::ConstantOp>(op, op.value());
+    return success();
+  }
+};
+
+} // namespace
+
+void mlir::tosa::populateTosaToStandardConversionPatterns(
+    MLIRContext *context, OwningRewritePatternList *patterns) {
+  patterns->insert<ConstOpConverter>(context);
+}

mlir/lib/Conversion/TosaToStandard/TosaToStandardPass.cpp

+//===- TosaToStandardPass.cpp - Lowering Tosa to Linalg Dialect -----------===//
+//
+// 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 transformation pass legalizes Tosa operations to the Standard dialect.
+//
+//===----------------------------------------------------------------------===//
+
+#include "../PassDetail.h"
+#include "mlir/Conversion/TosaToStandard/TosaToStandard.h"
+#include "mlir/Dialect/StandardOps/IR/Ops.h"
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/Dialect/Tosa/IR/TosaOps.h"
+#include "mlir/Dialect/Tosa/Transforms/PassDetail.h"
+#include "mlir/Dialect/Tosa/Transforms/Passes.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/Pass/PassManager.h"
+#include "mlir/Transforms/DialectConversion.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+
+using namespace mlir;
+using namespace tosa;
+
+namespace {
+struct TosaToStandard : public TosaToStandardBase<TosaToStandard> {
+public:
+  void runOnOperation() override {
+    OwningRewritePatternList patterns;
+    ConversionTarget target(getContext());
+    target.addIllegalOp<tosa::ConstOp>();
+    target.addLegalOp<ConstantOp>();
+
+    auto *op = getOperation();
+    mlir::tosa::populateTosaToStandardConversionPatterns(op->getContext(),
+                                                         &patterns);
+    if (failed(applyPartialConversion(op, target, std::move(patterns))))
+      signalPassFailure();
+  }
+};
+} // namespace
+
+std::unique_ptr<Pass> mlir::tosa::createTosaToStandard() {
+  return std::make_unique<TosaToStandard>();
+}
+
+void mlir::tosa::addTosaToStandardPasses(OpPassManager &pm) {
+  pm.addNestedPass<FuncOp>(createTosaToStandard());
+}

mlir/test/Conversion/TosaToSCF/tosa-to-scf.mlir

+// RUN: mlir-opt --split-input-file --tosa-to-scf %s -verify-diagnostics -o -| FileCheck %s
+
+// CHECK-LABEL: func @while_test
+// CHECK-SAME: ([[ARG0:%.+]]: tensor<i32>)
+func @while_test(%arg0 : tensor<i32>) -> (tensor<i32>) {
+  // CHECK: [[WHILE:%.+]] = scf.while ([[ARG1:%.+]] = [[ARG0]])
+  %1 = "tosa.while_loop"(%arg0) ( {
+  ^bb0(%arg2: tensor<i32>):
+    // CHECK: "tosa.const"
+    %2 = "tosa.const"() {value = dense<3> : tensor<i32>} : () -> tensor<i32>
+
+    // CHECK: [[COMPARE:%.+]] = "tosa.greater_equal"
+    %3 = "tosa.greater_equal"(%2, %arg2) : (tensor<i32>, tensor<i32>) -> tensor<i1>
+
+    // CHECK: [[EX:%.+]] = tensor.extract [[COMPARE]]
+    // CHECK: scf.condition([[EX]]) [[ARG1]]
+    "tosa.yield"(%3) : (tensor<i1>) -> ()
+  },  {
+  // CHECK: ^bb0([[ARG1:%.+]]: tensor<i32>)
+  ^bb0(%arg2: tensor<i32>):
+    // CHECK: tosa.const
+    %2 = "tosa.const"() {value = dense<1> : tensor<i32>} : () -> tensor<i32>
+
+    // CHECK: [[ADD:%.+]] = "tosa.add"
+    %3 = "tosa.add"(%arg2, %2) : (tensor<i32>, tensor<i32>) -> tensor<i32>
+
+    // CHECK: scf.yield [[ADD]]
+    "tosa.yield"(%3) : (tensor<i32>) -> ()
+  }) : (tensor<i32>) -> (tensor<i32>)
+  return %1 : tensor<i32>
+}
+
+// ----
+
+// CHECK-LABEL: func @if_test
+// CHECK-SAME: ([[ARG0:%.+]]: tensor<f32>, [[ARG1:%.+]]: tensor<f32>, [[ARG2:%.+]]: tensor<i1>)
+func @if_test(%arg0 : tensor<f32>, %arg1 : tensor<f32>, %arg2 : tensor<i1>) -> (tensor<f32>) {
+  // CHECK: [[EX:%.+]] = tensor.extract [[ARG2]]
+  // CHECK: [[IF:%.+]] = scf.if [[EX]] -> (tensor<f32>) {
+  %0 = "tosa.cond_if"(%arg2, %arg0, %arg1) ({
+
+  // CHECK:   scf.yield [[ARG0]]
+  ^bb1(%arg3 : tensor<f32>, %arg4 : tensor<f32>):
+    "tosa.yield"(%arg3) : (tensor<f32>) -> ()
+
+  // CHECK: } else {
+  }, {
+
+  // CHECK:   scf.yield [[ARG1]]
+  ^bb1(%arg5 : tensor<f32>, %arg6 : tensor<f32>):
+    "tosa.yield"(%arg6) : (tensor<f32>) -> ()
+
+  // CHECK: }
+  // CHECK: return [[IF]]
+  }) : (tensor<i1>, tensor<f32>, tensor<f32>) -> (tensor<f32>)
+
+  return %0 : tensor<f32>
+}

mlir/test/Conversion/TosaToStandard/tosa-to-standard.mlir

+// RUN: mlir-opt --split-input-file --tosa-to-standard %s -verify-diagnostics -o -| FileCheck %s
+
+// CHECK-LABEL: func @const_test
+func @const_test() -> (tensor<i32>) {
+  // CHECK: [[C3:%.+]] = constant dense<3> : tensor<i32>
+  %0 = "tosa.const"() {value = dense<3> : tensor<i32>} : () -> tensor<i32>
+
+  // CHECK: return [[C3]]
+  return %0 : tensor<i32>
+}