[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]

1485

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

38

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

167

Ditto

171

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.

1485

Added a helper to ValueKnowledge.

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

38

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.

171

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/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.hasRank
+                                       ? ShapedTypeComponents(result.sizes)
+                                       : ShapedTypeComponents());
+  }
+
+  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);
+
+  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();
+      auto meet = ValueKnowledge::meet(
+          resultKnowledge[index],
+          ValueKnowledge::getKnowledgeFromType(it.value().getType()));
+      if (!meet)
+        continue;
+      resultKnowledge[index] = meet;
+    }
+  }
+
+  for (const ValueKnowledge &result : resultKnowledge) {
+    inferredReturnShapes.push_back(result.hasRank
+                                       ? ShapedTypeComponents(result.sizes)
+                                       : ShapedTypeComponents());
   }
 
   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) {
+void propagateShapesToTosaIf(
+    Operation &op, DenseMap<Value, ShapedTypeComponents> &shapesStorage) {
   tosa::IfOp ifOp = dyn_cast<tosa::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; i < op.getNumOperands(); 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) {
+  tosa::WhileOp whileOp = dyn_cast<tosa::WhileOp>(op);
+  if (!whileOp)
+    return;
+
+  // Determine what the expect argument types are to the cond/body blocks.
+  // These need to remain consisten across every iteration.
+  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> backupTypeMap;
+  for (auto &block : op.getRegion(1)) {
+    for (auto arg : block.getArguments())
+      backupTypeMap[arg] = arg.getType();
+    for (auto &op : block)
+      for (auto result : op.getResults())
+        backupTypeMap[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<tosa::YieldOp> yieldOps;
+    for (auto &block : bodyRegion)
+      if (auto returnOp = dyn_cast<tosa::YieldOp>(block.getTerminator()))
+        yieldOps.push_back(returnOp);
+
+    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;
+    }
+
+    // Restore the original types in case we need to do a new pass.
+    for (auto &block : bodyRegion) {
+      for (auto arg : block.getArguments())
+        arg.setType(backupTypeMap[arg]);
+      for (auto &op : block)
+        for (auto result : op.getResults())
+          result.setType(backupTypeMap[result]);
+    }
+  }
+
+  // Now we set the eventual types on the args of each block, then propagate
+  // through the region.
+  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())
         continue;
 
-      propagateShapesToTosaIf(op);
+      propagateShapesToTosaIf(op, shapesStorage);
+      propagateShapesToTosaWhile(op, shapesStorage);
 
       InferShapedTypeOpInterface shapeInterface =
           dyn_cast<InferShapedTypeOpInterface>(op);
       if (!shapeInterface)
         continue;
 
       SmallVector<ShapedTypeComponents> returnedShapes;
 

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
+}