[mlir][tosa] Add tosa shape inference with InferReturnTypeComponent

Added InferReturnTypeComponents for NAry operations, reshape, and reverse.
With the additional tosa-infer-shapes pass, we can infer/propagate shapes
across a set of TOSA operations. Current version does not modify the
FuncOp type by inserting an unrealized conversion cast prior to any new
non-matchin returns.

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

Author: rsuderman

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

@@ -15,6 +15,7 @@
 
 #include "mlir/Dialect/Quant/QuantOps.h"
 #include "mlir/Dialect/Traits.h"
+#include "mlir/Interfaces/InferTypeOpInterface.h"
 #include "mlir/Interfaces/LoopLikeInterface.h"
 #include "mlir/Interfaces/SideEffectInterfaces.h"
 

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

@@ -13,11 +13,13 @@
 #ifndef MLIR_DIALECT_TOSA_TRANSFORMS_PASSES_H
 #define MLIR_DIALECT_TOSA_TRANSFORMS_PASSES_H
 
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/Pass/Pass.h"
 
 namespace mlir {
 namespace tosa {
 
+std::unique_ptr<Pass> createTosaInferShapesPass();
 std::unique_ptr<Pass> createTosaMakeBroadcastablePass();
 std::unique_ptr<Pass> createTosaTestQuantUtilAPIPass();
 

mlir/include/mlir/Dialect/Tosa/Transforms/Passes.h

@@ -15,6 +15,21 @@
 
 include "mlir/Pass/PassBase.td"
 
+def TosaInferShapes : FunctionPass<"tosa-infer-shapes"> {
+  let summary = "Propagate shapes across TOSA operations";
+  let description = [{
+    Pass that uses operand types and propagates shapes to TOSA operations.
+    This includes legalizing rankless and dynamic shapes towards static.
+  }];
+
+  let constructor = "createTosaInferShapesPass()";
+  let dependentDialects = [
+    "StandardOpsDialect",
+    "tensor::TensorDialect",
+    "tosa::TosaDialect",
+  ];
+}
+
 def TosaMakeBroadcastable : FunctionPass<"tosa-make-broadcastable"> {
   let summary = "TOSA rank Reshape to enable Broadcasting";
   let description = [{

mlir/include/mlir/Dialect/Tosa/Transforms/Passes.td

@@ -291,6 +291,148 @@ static void buildPadOpWithQuantInfo(OpBuilder &builder, OperationState &result,
   result.types.push_back(outputType);
 }
 
+//===----------------------------------------------------------------------===//
+// TOSA Operator Return Type Inference.
+//===----------------------------------------------------------------------===//
+
+static void getI64Values(ArrayAttr arrayAttr, SmallVector<int64_t> &values) {
+  for (auto it : arrayAttr) {
+    values.push_back(it.cast<IntegerAttr>().getValue().getSExtValue());
+  }
+}
+
+LogicalResult tosa::ReshapeOp::inferReturnTypeComponents(
+    MLIRContext *context, ::llvm::Optional<Location> location,
+    ValueRange operands, DictionaryAttr attributes, RegionRange regions,
+    SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) {
+  ShapedType type = operands.front().getType().cast<ShapedType>();
+
+  auto newShape = attributes.get("new_shape").cast<ArrayAttr>();
+  llvm::SmallVector<int64_t> newShapeValue;
+  getI64Values(newShape, newShapeValue);
+
+  // We cannot infer from the total number of elements so we must take the
+  // shape attribute as exact.
+  if (!type.hasRank() || !type.hasStaticShape()) {
+    inferredReturnShapes.push_back(ShapedTypeComponents(newShapeValue));
+    return success();
+  }
+
+  // Determine the number of elements covered by the slice of all static
+  // dimensions. This allows us to infer the length of the remaining dynamic
+  // dimension.
+  int64_t numElements = type.getNumElements();
+  int64_t staticMul = 1;
+  for (auto val : newShapeValue) {
+    if (val != -1) {
+      staticMul *= val;
+    }
+  }
+
+  // Determine the length of the dynamic dimension.
+  for (auto &val : newShapeValue) {
+    if (val == -1)
+      val = numElements / staticMul;
+  }
+
+  inferredReturnShapes.push_back(ShapedTypeComponents(newShapeValue));
+  return success();
+}
+
+static LogicalResult resolveBroadcastShape(ValueRange operands,
+                                           SmallVector<int64_t> &outShape) {
+  int64_t outRank = 0;
+  for (auto operand : operands) {
+    auto type = operand.getType().cast<ShapedType>();
+    if (!type.hasRank())
+      return failure();
+    outRank = std::max<int64_t>(outRank, type.getRank());
+  }
+
+  outShape.resize(outRank, 1);
+
+  for (auto operand : operands) {
+    auto type = operand.getType().cast<ShapedType>();
+    auto shape = type.getShape();
+    auto rankDiff = outShape.size() - shape.size();
+
+    for (size_t i = 0; i < shape.size(); i++) {
+      auto dim1 = outShape[i + rankDiff];
+      auto dim2 = shape[i];
+      auto resolvedDim = dim1;
+
+      if (dim1 == 1) {
+        resolvedDim = dim2;
+      } else if (dim2 == 1) {
+        resolvedDim = dim1;
+      } else if (dim1 != dim2) {
+        return failure();
+      }
+      outShape[i + rankDiff] = resolvedDim;
+    }
+  }
+
+  return success();
+}
+
+static LogicalResult NAryInferReturnTypes(
+    ValueRange operands,
+    SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) {
+  llvm::SmallVector<int64_t> outShape;
+  if (resolveBroadcastShape(operands, outShape).failed()) {
+    inferredReturnShapes.push_back(ShapedTypeComponents());
+  } else {
+    inferredReturnShapes.push_back(ShapedTypeComponents(outShape));
+  }
+  return success();
+}
+
+#define NARY_SHAPE_INFER(OP)                                                   \
+  LogicalResult OP::inferReturnTypeComponents(                                 \
+      MLIRContext *context, ::llvm::Optional<Location> location,               \
+      ValueRange operands, DictionaryAttr attributes, RegionRange regions,     \
+      SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) {           \
+    return NAryInferReturnTypes(operands, inferredReturnShapes);               \
+  }
+
+NARY_SHAPE_INFER(tosa::AbsOp)
+NARY_SHAPE_INFER(tosa::AddOp)
+NARY_SHAPE_INFER(tosa::ArithmeticRightShiftOp)
+NARY_SHAPE_INFER(tosa::BitwiseAndOp)
+NARY_SHAPE_INFER(tosa::BitwiseOrOp)
+NARY_SHAPE_INFER(tosa::BitwiseXorOp)
+NARY_SHAPE_INFER(tosa::BitwiseNotOp)
+NARY_SHAPE_INFER(tosa::CeilOp)
+NARY_SHAPE_INFER(tosa::ClampOp)
+NARY_SHAPE_INFER(tosa::ClzOp)
+NARY_SHAPE_INFER(tosa::DivOp)
+NARY_SHAPE_INFER(tosa::EqualOp)
+NARY_SHAPE_INFER(tosa::ExpOp)
+NARY_SHAPE_INFER(tosa::FloorOp)
+NARY_SHAPE_INFER(tosa::GreaterEqualOp)
+NARY_SHAPE_INFER(tosa::GreaterOp)
+NARY_SHAPE_INFER(tosa::LogOp)
+NARY_SHAPE_INFER(tosa::LogicalAndOp)
+NARY_SHAPE_INFER(tosa::LogicalLeftShiftOp)
+NARY_SHAPE_INFER(tosa::LogicalNotOp)
+NARY_SHAPE_INFER(tosa::LogicalOrOp)
+NARY_SHAPE_INFER(tosa::LogicalRightShiftOp)
+NARY_SHAPE_INFER(tosa::LogicalXorOp)
+NARY_SHAPE_INFER(tosa::MaximumOp)
+NARY_SHAPE_INFER(tosa::MinimumOp)
+NARY_SHAPE_INFER(tosa::MulOp)
+NARY_SHAPE_INFER(tosa::NegateOp)
+NARY_SHAPE_INFER(tosa::PowOp)
+NARY_SHAPE_INFER(tosa::ReciprocalOp)
+NARY_SHAPE_INFER(tosa::ReluNOp)
+NARY_SHAPE_INFER(tosa::ReverseOp)
+NARY_SHAPE_INFER(tosa::RsqrtOp)
+NARY_SHAPE_INFER(tosa::SelectOp)
+NARY_SHAPE_INFER(tosa::SubOp)
+NARY_SHAPE_INFER(tosa::TanhOp)
+NARY_SHAPE_INFER(tosa::SigmoidOp)
+#undef PRED_SHAPE_INFER
+
 //===----------------------------------------------------------------------===//
 // TOSA Operator Definitions.
 //===----------------------------------------------------------------------===//

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

@@ -1,4 +1,5 @@
 add_mlir_dialect_library(MLIRTosaTransforms
+  TosaInferShapes.cpp
   TosaMakeBroadcastable.cpp
 
   ADDITIONAL_HEADER_DIRS