[mlir][tosa] Add shape inference for tosa.while

[rsuderman]

Summary

Tosa.while shape inference requires repeatedly running shape inference across
the body of the loop until the types become static as we do not know the number
of iterations required by the loop body. Once the least specific arguments are
known they are propagated to both regions.

To determine the final end type, the least restrictive types are determined
from all yields.

[jpienaar]

Nice, some questions to help my understanding here

mlir/lib/Dialect/Tosa/IR/TosaOps.cpp

1459

Add comment for failure

1466

So these would require that the operand type has been updated?

1475

Mmm, would

if (auto meet = ...)

resultKnowledge[index] = meet;

result in difficult to read if? (removed need for continue and reduces lifetime scope of variable)

[I needed to check meet here to make sure it was doing what I'd expect for while where it could get less refined every time through]

1484

Do we expect to need this pattern more times? (I see it is repeated 2x, so perhaps a helper method to go from resultKnowledge to shapes could help)

mlir/lib/Dialect/Tosa/Transforms/TosaInferShapes.cpp

37

Both here and below, an array would have worked as well, is the map just easier to populate? (I htink only accessed once in each, so cost of indexing probably the same)

48

Nit: hoist getNumOperands() to document it doesn't change during iteration.

77

Is tosa:: needed? seems like you are using mlir::tosa above.

81

expected ?

82

consistent

82

And this is a property of tosa while?

96

originalTypeMap ? (backup feels like something used in case of failure)

121

OOC, why not yieldOp ?

154

Could you expand on this? I thought the types would be more refined now, not sure why it would need to be reset

171

Ditto

175

It feels like if propagating to regions took a list of shapes, then we could avoid the setting above (but I don't get the resetting part yet, which is either helped by this/can be removed, or shows why it is needed to set)

[rsuderman]

Updated for jpienaar comments.

[rsuderman]

Removed extraneous tosa::

[rsuderman]

Few more tweaks to updates.

[rsuderman]

mlir/lib/Dialect/Tosa/IR/TosaOps.cpp

1466

Kinda? It more uses the initial operand type as the starting place for performing shape inference. The idea with while loop propagation is that the block args for cond/while must be compatible with ever invocation. Therefore we start with the initial types (specified by the operand types) and iterate each cycle.

E.g. if we don't know much about the operand types we can't infer as much about the types in the cond/body blocks.

1484

Added a helper to ValueKnowledge.

mlir/lib/Dialect/Tosa/Transforms/TosaInferShapes.cpp

37

This is true however the DenseMap is already defined and it would be the same number of dense map lookups to determine the operand types. So better to just pass the DenseMap.

82

I added more details but overall, yes. Because we execute the regions multiple times, we need their types compatible with every possible iteration.

121

No particular reason. Changed to YieldOp.

175

Yeah, we could go store a map of shapes instead and propagate rather than mutating the IR every iteration. It felt more succinct to accept mutation rather than work around it. Trying to store a type change map introduced a level of indirection that felt unnecessary.

[rsuderman]

Accidentally reviewed it myself..... woo reflexes

[jpienaar]

Looks good thanks

mlir/lib/Dialect/Tosa/IR/TosaOps.cpp

1466

Sure, I was more referring to operands.getShape which could give a more exact shape here

mlir/lib/Dialect/Tosa/Transforms/TosaInferShapes.cpp

128

Micro nit: should TOSA be used when referring to the opset and tosa.X when referring to op X in opset TOSA? (e.g., is there a capitalization convention to follow here)

145

Perhaps op.emitWarning() ? If this should never happen, then we are only silently failing at the moment.

[rsuderman]

Fixed up for last comments.

mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td

def Tosa_ConcatOp : Tosa_Op<"concat", [
   let summary = "Concatenates tensors along one dimension.";
 
   let description = [{
     Concatenate a variadic amount of tensors along a given axis. No data 
     conversion happens during a concat operation.
   }];
 
   let arguments = (ins
-    Variadic<Tosa_RankedTensor>:$input1,
+    Variadic<Tosa_Tensor>:$input1,
     I64Attr:$axis
   );
 
   let results = (outs
-    Tosa_RankedTensor:$output
+    Tosa_Tensor:$output
   );
 
   let hasCanonicalizer = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // Operator: pad
 //===----------------------------------------------------------------------===//

 //===----------------------------------------------------------------------===//
 // Operator: while_loop
 //===----------------------------------------------------------------------===//
 //===----------------------------------------------------------------------===//
 // Further described in docs/Rationale/RationaleTOSADialect.md .
 //===----------------------------------------------------------------------===//
 def Tosa_WhileOp : Tosa_Op<"while_loop", [
        DeclareOpInterfaceMethods<LoopLikeOpInterface>,
+       DeclareOpInterfaceMethods<InferShapedTypeOpInterface,
+                                 ["inferReturnTypeComponents"]>,
        SingleBlockImplicitTerminator<"YieldOp">,
        RecursiveSideEffects]> {
   let summary = "output = input; While (Cond(output)) {output = Body(output)}";
 
   let description = [{
     Generates and evaluates a Bool condition and either executes a loop body or
     exits to another control point. This action is performed repeatedly after
     updating and re-evaluating the Boolean condition every iteration. This

mlir/include/mlir/Dialect/Tosa/Utils/ShapeUtils.h

 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_DIALECT_TOSA_UTILS_SHAPEUTILS_H
 #define MLIR_DIALECT_TOSA_UTILS_SHAPEUTILS_H
 
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/Types.h"
+#include "mlir/Interfaces/InferTypeOpInterface.h"
 #include "llvm/ADT/Sequence.h"
 #include "llvm/ADT/SmallVector.h"
 
 namespace mlir {
 namespace tosa {
 /// Statically known information for a particular Value.
 ///
 /// This struct currently tracks only information relevant for tensor/array-like
struct ValueKnowledge {
   }
 
   // Return a pessimistic/conservative value state without assuming any knowlege
   // about the IR.
   static ValueKnowledge getPessimisticValueState() {
     return ValueKnowledge(false, {}, Type());
   }
 
+  ShapedTypeComponents getShapedTypeComponents() const {
+    return hasRank ? ShapedTypeComponents(sizes) : ShapedTypeComponents();
+  }
+
   Type getType() const {
     if (hasRank)
       return RankedTensorType::get(llvm::makeArrayRef(sizes), dtype);
     return UnrankedTensorType::get(dtype);
   }
 
   bool operator==(const ValueKnowledge &rhs) const {
     return hasRank == rhs.hasRank && sizes == rhs.sizes && dtype == rhs.dtype;

mlir/lib/Dialect/Tosa/IR/TosaOps.cpp

 #include "mlir/Dialect/Tosa/Utils/QuantUtils.h"
 #include "mlir/Dialect/Tosa/Utils/ShapeUtils.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/Matchers.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/Transforms/FoldUtils.h"
 #include "mlir/Transforms/InliningUtils.h"
 #include "mlir/Transforms/RegionUtils.h"
+#include "llvm/ADT/DenseMap.h"
 
 using namespace mlir;
 using namespace mlir::tosa;
 
 #include "mlir/Dialect/Tosa/IR/TosaOpsDialect.cpp.inc"
 
 //===----------------------------------------------------------------------===//
 // Tosa dialect structs and interface includes.
for (auto it : llvm::enumerate(yieldOp.getOperands())) {
           ValueKnowledge::getKnowledgeFromType(it.value().getType()));
       if (!meet)
         continue;
       resultKnowledge[index] = meet;
     }
   }
 
   for (const ValueKnowledge &result : resultKnowledge) {
-    if (result.hasRank) {
-      inferredReturnShapes.push_back(ShapedTypeComponents(result.sizes));
-    } else {
-      inferredReturnShapes.push_back(ShapedTypeComponents());
-    }
+    inferredReturnShapes.push_back(result.getShapedTypeComponents());
+  }
+
+  return success();
+}
+
+LogicalResult WhileOp::inferReturnTypeComponents(
+    MLIRContext *context, ::llvm::Optional<Location> location,
+    ValueShapeRange operands, DictionaryAttr attributes, RegionRange regions,
+    SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) {
+  llvm::SmallVector<tosa::YieldOp> yieldOps;
+  for (auto &block : *regions[1])
+    if (auto returnOp = dyn_cast<tosa::YieldOp>(block.getTerminator()))
+      yieldOps.push_back(returnOp);
+
+  // TOSA's while must have a tosa.yield as its terminator. If not found this
+  // tosa.while is invalid.
+  if (yieldOps.empty())
+    return failure();
+
+  // Get the initial type information from the operand types.
+  llvm::SmallVector<ValueKnowledge> resultKnowledge;
+  resultKnowledge.reserve(yieldOps.front().getNumOperands());
+  for (auto operand : yieldOps.front().getOperands()) {
+    resultKnowledge.push_back(
+        ValueKnowledge::getKnowledgeFromType(operand.getType()));
+  }
+
+  for (auto yieldOp : yieldOps) {
+    if (resultKnowledge.size() != yieldOp.getNumOperands())
+      return failure();
+
+    for (auto it : llvm::enumerate(yieldOp.getOperands())) {
+      int32_t index = it.index();
+      if (auto meet = ValueKnowledge::meet(
+              resultKnowledge[index],
+              ValueKnowledge::getKnowledgeFromType(it.value().getType()))) {
+        resultKnowledge[index] = meet;
+      };
+    }
+
+  for (const ValueKnowledge &result : resultKnowledge) {
+    inferredReturnShapes.push_back(result.getShapedTypeComponents());
   }
 
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // TOSA Operator Definitions.
 //===----------------------------------------------------------------------===//
 
 #define GET_OP_CLASSES
 #include "mlir/Dialect/Tosa/IR/TosaOps.cpp.inc"

mlir/lib/Dialect/Tosa/Transforms/TosaInferShapes.cpp

 
 using namespace mlir;
 using namespace mlir::tosa;
 
 namespace {
 
 void propagateShapesInRegion(Region &region);
 
-void propagateShapesToTosaIf(Operation &op) {
-  tosa::IfOp ifOp = dyn_cast<tosa::IfOp>(op);
+void propagateShapesToTosaIf(
+    Operation &op, DenseMap<Value, ShapedTypeComponents> &shapesStorage) {
+  IfOp ifOp = dyn_cast<IfOp>(op);
   if (!ifOp)
     return;
 
   for (auto &region : op.getRegions()) {
     Block &frontBlock = region.front();
     if (frontBlock.getNumArguments() + 1 != ifOp.getNumOperands())
       return;
 
+    for (unsigned int i = 1, s = op.getNumOperands(); i < s; i++) {
+      auto inferredTy = shapesStorage[op.getOperand(i)];
+      auto blockArg = frontBlock.getArgument(i - 1);
+      auto oldType = blockArg.getType().cast<ShapedType>();
+
+      if (inferredTy.hasRank()) {
+        Type newType = oldType.clone(inferredTy.getDims());
+        blockArg.setType(newType);
+      }
+    }
+
     for (int i = 0, e = frontBlock.getNumArguments(); i < e; i++) {
       ValueKnowledge operandKnowledge = ValueKnowledge::getKnowledgeFromType(
           ifOp.getOperand(i + 1).getType());
       ValueKnowledge blockKnowledge = ValueKnowledge::getKnowledgeFromType(
           frontBlock.getArgument(i).getType());
       ValueKnowledge joinedKnowledge =
           ValueKnowledge::join(operandKnowledge, blockKnowledge);
       if (!joinedKnowledge)
         continue;
       frontBlock.getArgument(i).setType(joinedKnowledge.getType());
     }
 
     propagateShapesInRegion(region);
   }
+}
+
+void propagateShapesToTosaWhile(
+    Operation &op, DenseMap<Value, ShapedTypeComponents> &shapesStorage) {
+  WhileOp whileOp = dyn_cast<WhileOp>(op);
+  if (!whileOp)
+    return;
+
+  // Determine what the expected argument types are to the cond/body blocks.
+  // The expected arguments should be compatible with ever iteration of the
+  // loop body / condition for tosa.while.
+  llvm::SmallVector<Type> argTypes;
+  for (auto operand : op.getOperands()) {
+    auto operandTy = operand.getType().cast<ShapedType>();
+    auto shapedTypeComponent = shapesStorage[operand];
+    if (shapedTypeComponent.hasRank()) {
+      auto newTy = operandTy.clone(shapedTypeComponent.getDims());
+      argTypes.push_back(newTy);
+    } else {
+      argTypes.push_back(operand.getType());
+    }
+  }
 
+  // Save out the type information so we can restore at the end.
+  llvm::DenseMap<Value, Type> originalTypeMap;
+  for (auto &block : op.getRegion(1)) {
+    for (auto arg : block.getArguments())
+      originalTypeMap[arg] = arg.getType();
+    for (auto &op : block)
+      for (auto result : op.getResults())
+        originalTypeMap[result] = result.getType();
+  }
+
+  bool hasNewTypes = true;
+  while (hasNewTypes) {
+
+    // Set types on the block args.
+    Region &bodyRegion = op.getRegion(1);
+    Block &block = bodyRegion.front();
+    for (int i = 0, s = argTypes.size(); i < s; i++) {
+      block.getArgument(i).setType(argTypes[i]);
+    }
+
+    // Propagate to the end.
+    propagateShapesInRegion(bodyRegion);
+
+    // Find all the tosa yield types and verify there is atleast one.
+    llvm::SmallVector<YieldOp> yieldOps;
+    for (auto &block : bodyRegion)
+      if (auto yieldOp = dyn_cast<YieldOp>(block.getTerminator()))
+        yieldOps.push_back(yieldOp);
+
+    if (yieldOps.empty())
       return;
+
+    // Using the new tosa yield types, infer the new subtypes.
+    llvm::SmallVector<ValueKnowledge> yieldTypeInfo;
+    for (auto ty : argTypes) {
+      yieldTypeInfo.push_back(ValueKnowledge::getKnowledgeFromType(ty));
+    }
+
+    for (auto yieldOp : yieldOps) {
+      for (auto it : llvm::enumerate(yieldOp.getOperands())) {
+        auto newKnowledge =
+            ValueKnowledge::getKnowledgeFromType(it.value().getType());
+        yieldTypeInfo[it.index()] =
+            ValueKnowledge::meet(yieldTypeInfo[it.index()], newKnowledge);
+      }
+    }
+
+    // This should never happen.
+    if (yieldTypeInfo.size() != argTypes.size())
+      return;
+
+    // Determine the new block args and see if any changed.
+    hasNewTypes = false;
+    for (int i = 0, s = yieldTypeInfo.size(); i < s; i++) {
+      Type newType = yieldTypeInfo[i].getType();
+      hasNewTypes |= (newType != argTypes[i]);
+      argTypes[i] = newType;
+    }
+
+    // The types inferred in the block assume the operand types specified for
+    // this iteration. We need to restore the original types to ensure that
+    // future iterations only use the already specified types, not possible
+    // types from previous iterations.
+    for (auto &block : bodyRegion) {
+      for (auto arg : block.getArguments())
+        arg.setType(originalTypeMap[arg]);
+      for (auto &op : block)
+        for (auto result : op.getResults())
+          result.setType(originalTypeMap[result]);
+    }
+  }
+
+  // We now set the block arguments according to the most recent shape
+  // inference results.
+  for (auto &region : op.getRegions()) {
+    for (unsigned int i = 0; i < argTypes.size(); i++) {
+      region.front().getArgument(i).setType(argTypes[i]);
+    }
+
+    propagateShapesInRegion(region);
+  }
 }
 
 void propagateShapesInRegion(Region &region) {
   DenseMap<Value, ShapedTypeComponents> shapesStorage;
   auto setShapes = [&](Value val, Type t) {
     if (auto st = t.dyn_cast<ShapedType>())
       shapesStorage[val] = st;
     else
       shapesStorage[val] = t;
   };
   auto operandShape = [&](Value val) -> ShapeAdaptor {
     // Query the WIP mapping rather than the type if set.
     auto it = shapesStorage.find(val);
     if (it == shapesStorage.end())
       return nullptr;
     return it->second;
   };
 
   for (auto &block : region) {
     for (Operation &op : block) {
-      if (op.getDialect()->getNamespace() !=
-          tosa::TosaDialect::getDialectNamespace())
+      if (op.getDialect()->getNamespace() != TosaDialect::getDialectNamespace())
         continue;
 
-      propagateShapesToTosaIf(op);
+      propagateShapesToTosaIf(op, shapesStorage);
+      propagateShapesToTosaWhile(op, shapesStorage);
 
       InferShapedTypeOpInterface shapeInterface =
           dyn_cast<InferShapedTypeOpInterface>(op);
       if (!shapeInterface)
         continue;
 
       SmallVector<ShapedTypeComponents> returnedShapes;
 
for (Operation &op : block) {
 
           // Check whether this use case is replaceable. We define an op as
           // being replaceable if it is used by a ReturnOp or a TosaOp.
           bool replaceable = true;
           for (auto user : result.getUsers()) {
             if (isa<ReturnOp>(user))
               continue;
             if (user->getDialect()->getNamespace() ==
-                tosa::TosaDialect::getDialectNamespace())
+                TosaDialect::getDialectNamespace())
               continue;
 
             replaceable = false;
           }
 
           // Determine the knowledge based on the output type.
           // TODO: should also query WIP type probably
           Type resultTy = result.getType();

mlir/test/Dialect/Tosa/tosa-infer-shapes.mlir

func @resize_fp_offsetted(%arg0: tensor<1x2x4x1xi32>) {
   %0 = "tosa.resize"(%arg0) {mode = "NEAREST_NEIGHBOR", offset = [0, 0], offset_fp = [2.500000e-01 : f32, 2.500000e-01 : f32], output_size = [-1, -1], shift = 0 : i32, stride = [0, 0], stride_fp = [2.500000e-01 : f32, 5.000000e-01 : f32]} : (tensor<1x2x4x1xi32>) -> tensor<?x?x?x?xi32>
   return
 }
 
 // -----
 
 // CHECK-LABEL: @if_test_simple
 func @if_test_simple(%arg0 : tensor<f32>, %arg1 : tensor<f32>, %arg2 : tensor<i1>) -> () {
+  %a = "tosa.log"(%arg0) : (tensor<f32>) -> tensor<*xf32>
+  %b = "tosa.log"(%arg1) : (tensor<f32>) -> tensor<*xf32>
   // CHECK: (tensor<i1>, tensor<f32>, tensor<f32>) -> tensor<f32>
-  %0 = "tosa.cond_if"(%arg2, %arg0, %arg1) ({
-  ^bb1(%arg3 : tensor<f32>, %arg4 : tensor<f32>):
-    "tosa.yield"(%arg3) : (tensor<f32>) -> ()
+  %0 = "tosa.cond_if"(%arg2, %a, %b) ({
+  ^bb1(%arg3 : tensor<*xf32>, %arg4 : tensor<*xf32>):
+    "tosa.yield"(%arg3) : (tensor<*xf32>) -> ()
   }, {
-  ^bb1(%arg5 : tensor<f32>, %arg6 : tensor<f32>):
-    "tosa.yield"(%arg6) : (tensor<f32>) -> ()
-  }) : (tensor<i1>, tensor<f32>, tensor<f32>) -> (tensor<*xf32>)
+  ^bb1(%arg5 : tensor<*xf32>, %arg6 : tensor<*xf32>):
+    "tosa.yield"(%arg6) : (tensor<*xf32>) -> ()
+  }) : (tensor<i1>, tensor<*xf32>, tensor<*xf32>) -> (tensor<*xf32>)
   return
 }
 
 // -----
 
 // CHECK-LABEL: @if_test_dynamic
 func @if_test_dynamic(%arg0 : tensor<2xf32>, %arg1 : tensor<3xf32>, %arg2 : tensor<i1>) -> () {
   // CHECK: (tensor<i1>, tensor<2xf32>, tensor<3xf32>) -> tensor<?xf32>
^bb1(%arg3 : tensor<*xf32>, %arg4 : tensor<*xf32>):
     "tosa.yield"(%1) : (tensor<*xf32>) -> ()
   }, {
   ^bb1(%arg5 : tensor<*xf32>, %arg6 : tensor<*xf32>):
     %1 = "tosa.sub"(%arg5, %arg6) : (tensor<*xf32>, tensor<*xf32>) -> tensor<*xf32>
     "tosa.yield"(%1) : (tensor<*xf32>) -> ()
   }) : (tensor<i1>, tensor<f32>, tensor<f32>) -> (tensor<*xf32>)
   return
 }
+
+// -----
+
+// CHECK-LABEL: @while_test
+func @while_test(%arg0 : tensor<i32>) -> (tensor<*xi32>) {
+  // CHECK:      "tosa.add" 
+  // CHECK-SAME: (tensor<i32>, tensor<i32>) -> tensor<i32>
+  %0 = "tosa.add"(%arg0, %arg0) : (tensor<i32>, tensor<i32>) -> tensor<*xi32>
+
+  // CHECK:      "tosa.while_loop"
+  %1 = "tosa.while_loop"(%0) ( {
+
+  // CHECK:      ^bb0
+  // CHECK-SAME: tensor<i32>
+  ^bb0(%arg2: tensor<*xi32>):
+    %2 = "tosa.const"() {value = dense<3> : tensor<i32>} : () -> tensor<i32>
+    // CHECK:       "tosa.greater_equal"
+    // CHECK-SAME: (tensor<i32>, tensor<i32>) -> tensor<i1>
+    %3 = "tosa.greater_equal"(%2, %arg2) : (tensor<i32>, tensor<*xi32>) -> tensor<*xi1>
+    // CHECK:      "tosa.yield"
+    // CHECK-SAME: tensor<i1>
+    "tosa.yield"(%3) : (tensor<*xi1>) -> ()
+  },  {
+  // CHECK:      ^bb0
+  // CHECK-SAME: tensor<i32>
+  ^bb0(%arg2: tensor<*xi32>):
+    %2 = "tosa.const"() {value = dense<1> : tensor<i32>} : () -> tensor<i32>
+
+    // CHECK:     "tosa.add"
+    // CHECK-SAME: (tensor<i32>, tensor<i32>) -> tensor<i32>
+    %3 = "tosa.add"(%arg2, %2) : (tensor<*xi32>, tensor<i32>) -> tensor<*xi32>
+
+    // CHECK:      "tosa.yield"
+    // CHECK-SAME: tensor<i32>
+    "tosa.yield"(%3) : (tensor<*xi32>) -> ()
+
+  // CHECK:      (tensor<i32>) -> tensor<i32>
+  }) : (tensor<*xi32>) -> (tensor<*xi32>)
+
+  // CHECK:      tensor.cast
+  return %1 : tensor<*xi32>
+}
+
+// -----
+
+// CHECK-LABEL: @while_test
+func @while_test(%arg0 : tensor<i32>, %arg1 : tensor<1xi32>) -> () {
+  // CHECK:      "tosa.while_loop"
+  %1:2 = "tosa.while_loop"(%arg0, %arg1) ( {
+
+  // CHECK:      ^bb0
+  // CHECK-SAME: tensor<i32>
+  // CHECK-SAME: tensor<?xi32>
+  ^bb0(%arg2: tensor<*xi32>, %arg3: tensor<*xi32>):
+    %2 = "tosa.const"() {value = dense<3> : tensor<i32>} : () -> tensor<i32>
+
+    // CHECK:       "tosa.greater_equal"
+    // CHECK-SAME: (tensor<i32>, tensor<i32>) -> tensor<i1>
+    %3 = "tosa.greater_equal"(%2, %arg2) : (tensor<i32>, tensor<*xi32>) -> tensor<*xi1>
+    "tosa.yield"(%3) : (tensor<*xi1>) -> ()
+  },  {
+
+  // CHECK:      ^bb0
+  // CHECK-SAME: tensor<i32>
+  // CHECK-SAME: tensor<?xi32>
+  ^bb0(%arg2: tensor<*xi32>, %arg3: tensor<*xi32>):
+    %2 = "tosa.const"() {value = dense<1> : tensor<i32>} : () -> tensor<i32>
+
+    // CHECK:     "tosa.add"
+    // CHECK-SAME: (tensor<i32>, tensor<i32>) -> tensor<i32>
+    %3 = "tosa.add"(%arg2, %2) : (tensor<*xi32>, tensor<i32>) -> tensor<*xi32>
+
+    // CHECK:      "tosa.concat"
+    // CHECK-SAME: (tensor<?xi32>, tensor<?xi32>) -> tensor<?xi32>
+    %4 = "tosa.concat"(%arg3, %arg3) { axis = 0 : i64 } : (tensor<*xi32>, tensor<*xi32>) -> (tensor<*xi32>)
+
+    // CHECK:      "tosa.yield"
+    // CHECK-SAME: tensor<i32>
+    // CHECK-SAME: tensor<?xi32>
+    "tosa.yield"(%3, %4) : (tensor<*xi32>, tensor<*xi32>) -> ()
+
+  // CHECK:      (tensor<i32>, tensor<1xi32>) -> (tensor<i32>, tensor<?xi32>)
+  }) : (tensor<i32>, tensor<1xi32>) -> (tensor<*xi32>, tensor<*xi32>)
+  return
+}