[mlir] Enable specifying querying function in ValueShapeRange

[jpienaar]

Summary

This enables querying shapes/values as shapes without mutating the IR directly
(e.g., towards enabling doing inference in analysis & application steps,
inferring function shape with constant from callsite, ...). Add a new
ShapeAdaptor that abstracts over whether shape is from Type or
ShapedTypeComponents or DenseIntElementsAttribute. This adds new accessors to
ValueShapeRange to get Shape and value as shape, but doesn't restrict or remove
the previous way of accessing Type via the Value for now, that does mean a less
refined shape could be accidentally queried and will be restricted in follow
up.

Currently restricted Value query to what can be represented as Shape. So only
supports cases where constant subgraph evaluation's output is a shape. I had
considered making it more general, but without TBD extern attribute concept or
some such a user cannot today uniformly avoid overhead (it does happen less
frequently shape inference side, but we've run into OOMs here and so rather
starting more restrictive here).

Update TOSA ops and also (but potentially more as an example and for discussion)
the shape inference pass.

[rsuderman]

Just some minor fixes.

mlir/lib/Interfaces/InferTypeOpInterface.cpp

164

Some leaked debug code

mlir/unittests/Interfaces/InferTypeOpInterfaceTest.cpp

58

Is there any reason this is an ASSERT instead of an EXPECT? Seems kinda strange to have the random intersperse.

[rriddle]

mlir/lib/Interfaces/InferTypeOpInterface.cpp

135–136

This looks sketchy, why is this necessary?

[jpienaar]

Remove leftover debugging

mlir/lib/Interfaces/InferTypeOpInterface.cpp

135–136

Left over debugging, good catch. I was debugging a module and just use this to avoid output from different threads getting munged.

mlir/unittests/Interfaces/InferTypeOpInterfaceTest.cpp

58

Yes, so if this fails then the call below (hasRank()) would assert in debug mode and crash the test, so need to ensure we don't try and "dereference" if null.

mlir/include/mlir/Interfaces/InferTypeOpInterface.h

 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_INTERFACES_INFERTYPEOPINTERFACE_H_
 #define MLIR_INTERFACES_INFERTYPEOPINTERFACE_H_
 
 #include "mlir/IR/Attributes.h"
 #include "mlir/IR/Builders.h"
+#include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/Location.h"
 #include "mlir/IR/OpDefinition.h"
 #include "mlir/Support/LLVM.h"
+#include "llvm/ADT/PointerUnion.h"
 #include "llvm/ADT/SmallVector.h"
 
 namespace mlir {
 
 using ReifiedRankedShapedTypeDims = SmallVector<SmallVector<Value>>;
 
 /// ShapedTypeComponents that represents the components of a ShapedType.
 /// The components consist of
 ///  - A ranked or unranked shape with the dimension specification match those
 ///    of ShapeType's getShape() (e.g., dynamic dimension represented using
 ///    ShapedType::kDynamicSize)
 ///  - A element type, may be unset (nullptr)
 ///  - A attribute, may be unset (nullptr)
 /// Used by ShapedType type inferences.
 class ShapedTypeComponents {
   /// Internal storage type for shape.
   using ShapeStorageT = SmallVector<int64_t, 3>;
 
 public:
   /// Default construction is an unranked shape.
-  ShapedTypeComponents() : ranked(false), elementType(nullptr), attr(nullptr){};
+  ShapedTypeComponents() : elementType(nullptr), attr(nullptr), ranked(false){};
   ShapedTypeComponents(Type elementType)
-      : ranked(false), elementType(elementType), attr(nullptr) {}
+      : elementType(elementType), attr(nullptr), ranked(false) {}
+  ShapedTypeComponents(ShapedType shapedType) : attr(nullptr) {
+    ranked = shapedType.hasRank();
+    elementType = shapedType.getElementType();
+    if (ranked)
+      dims = llvm::to_vector<4>(shapedType.getShape());
+  }
   template <typename Arg, typename = typename std::enable_if_t<
                               std::is_constructible<ShapeStorageT, Arg>::value>>
   ShapedTypeComponents(Arg &&arg, Type elementType = nullptr,
                        Attribute attr = nullptr)
-      : dims(std::forward<Arg>(arg)), ranked(true), elementType(elementType),
-        attr(attr) {}
+      : dims(std::forward<Arg>(arg)), elementType(elementType), attr(attr),
+        ranked(true) {}
   ShapedTypeComponents(ArrayRef<int64_t> vec, Type elementType = nullptr,
                        Attribute attr = nullptr)
-      : dims(vec.begin(), vec.end()), ranked(true), elementType(elementType),
-        attr(attr) {}
+      : dims(vec.begin(), vec.end()), elementType(elementType), attr(attr),
+        ranked(true) {}
 
   /// Return the dimensions of the shape.
   /// Requires: shape is ranked.
   ArrayRef<int64_t> getDims() const {
     assert(ranked && "requires ranked shape");
     return dims;
   }
 
   /// Return whether the shape has a rank.
   bool hasRank() const { return ranked; };
 
   /// Return the element type component.
   Type getElementType() const { return elementType; };
 
   /// Return the raw attribute component.
   Attribute getAttribute() const { return attr; };
 
 private:
+  friend class ShapeAdaptor;
+
   ShapeStorageT dims;
-  bool ranked;
   Type elementType;
   Attribute attr;
+  bool ranked;
+};
+
+/// Adaptor class to abstract the differences between whether value is from
+/// a ShapedType or ShapedTypeComponents or DenseIntElementsAttribute.
+class ShapeAdaptor {
+public:
+  ShapeAdaptor(Type t) {
+    if (auto st = t.dyn_cast<ShapedType>())
+      val = st;
+  }
+  ShapeAdaptor(Attribute t) {
+    if (auto da = t.dyn_cast<DenseIntElementsAttr>())
+      val = da;
+  }
+  ShapeAdaptor(ShapedTypeComponents *components) : val(components) {}
+  ShapeAdaptor(ShapedTypeComponents &components) : val(&components) {}
+
+  /// Returns whether the shape has a rank.
+  bool hasRank() const;
+
+  /// Returns the element type.
+  Type getElementType() const;
+
+  /// Populates the dimensions from shape referenced.
+  /// Requires: shape is ranked.
+  void getDims(SmallVectorImpl<int64_t> &res) const;
+
+  /// Populates the dimensions of the ShapeTypeComponents.
+  /// Requires: shape is ranked.
+  void getDims(ShapedTypeComponents &res) const;
+
+  /// Returns the size of the index'th dimension.
+  /// Requires: shape is ranked.
+  int64_t getDimSize(int index) const;
+
+  /// Returns whether the index'th dimension is dynamic.
+  /// Requires: shape is ranked.
+  bool isDynamicDim(int index) const {
+    return ShapedType::isDynamic(getDimSize(index));
+  }
+
+  /// Returns whether the shape is fully static.
+  bool hasStaticShape() const;
+
+  /// Returns the rank of the shape.
+  /// Requires: shape is ranked.
+  int64_t getRank() const;
+
+  /// Returns the number of elements in the shape.
+  /// Requires: hasStaticShape
+  int64_t getNumElements() const;
+
+  /// Returns whether valid (non-null) shape.
+  operator bool() const { return !val.isNull(); }
+
+  /// Dumps textual repesentation to stderr.
+  void dump() const;
+
+private:
+  // Union storing either ShapedTypeComponents, ShapedType (stored as Type and
+  // casted), or DenseIntElementsAttribute (stored as Atrtribute).
+  PointerUnion<ShapedTypeComponents *, Type, Attribute> val = nullptr;
 };
 
 /// Range of values and shapes (corresponding effectively to Shapes dialect's
 /// ValueShape type concept).
+// Currently this exposes the Value (of operands) and Type of the Value. This is
+// not ideal as then one can accidentally reference an out of date shape. This
+// is done to both enable gradual switch and also as OpAdaptor doesn't currently
+// allow returning anything other than Value.
 class ValueShapeRange : public ValueRange::RangeBaseT {
 public:
-  ValueShapeRange(ValueRange values) : RangeBaseT(values) {}
-  template <typename Arg, typename = typename std::enable_if_t<
-                              std::is_constructible<ValueRange, Arg>::value>>
-  ValueShapeRange(Arg &&arg)
-      : ValueShapeRange(ValueRange(std::forward<Arg>(arg))) {}
+  using ValueShapeMapFn = function_ref<ShapeAdaptor(Value)>;
+
+  ValueShapeRange(ValueRange values, ValueShapeMapFn operandShape = nullptr,
+                  ValueShapeMapFn valueToShape = nullptr)
+      : RangeBaseT(values), operandShape(operandShape),
+        valueToShape(valueToShape) {}
   ValueShapeRange(const std::initializer_list<Value> &values)
       : ValueShapeRange(ValueRange(values)) {}
 
+  ValueShapeRange(const ValueShapeRange &other) : RangeBaseT(other) {
+    operandShape = other.operandShape;
+    valueToShape = other.valueToShape;
+  }
+
+  /// Sets the Value to ShapeAdaptor mapping function and returns this.
+  ValueShapeRange &setValueToShapeMapping(ValueShapeMapFn fn) {
+    valueToShape = fn;
+    return *this;
+  }
+
+  ValueShapeRange &setOperandShapeMapping(ValueShapeMapFn fn) {
+    operandShape = fn;
+    return *this;
+  }
+
+  /// Returns the set Value to ShapeAdaptor mapping function.
+  ValueShapeMapFn getValueToShapeMapping() const { return valueToShape; }
+  ValueShapeMapFn getOperandShapeMapping() const { return operandShape; }
+
+  // Accessors.
+
   /// Returns the types of the values within this range.
   /// Note: This returns only the types of Values in the ValueRange and not a
   /// more refined type.
   using type_iterator = ValueTypeIterator<iterator>;
   using type_range = ValueTypeRange<ValueRange>;
   type_range getTypes() const { return {begin(), end()}; }
   auto getType() const { return getTypes(); }
 
   /// Returns the Values in the ValueRange.
+  /// To query the most up to date shape of a Value, query the shape
+  /// using getShape below rather than using the type of the Value.
   ValueRange getValues() const { return ValueRange(begin(), end()); };
+
+  /// Returns an argument as shape. If the argument is not constant or not a
+  /// shape, then the function returns a nullptr.
+  /// This will first query the valueToShape mapping (if set), before querying
+  /// the ValueRange.
+  ShapeAdaptor getValueAsShape(int index);
+
+  /// Returns the shape of index'th operand.
+  // TODO: Update so that operator[] references these instead to avoid
+  // accidentally refering to less refined shape.
+  ShapeAdaptor getShape(int index) const;
+
+  /// Returns the shape of the given Value.
+  ShapeAdaptor getShape(Value val) const;
+
+private:
+  // Mapping from Value to ShapedTypeComponents corresponding to shape of type
+  // of Value.
+  ValueShapeMapFn operandShape;
+
+  // Mapping from Value to ShapedTypeComponents corresponding to constant Value
+  // if interpreted as shape.
+  ValueShapeMapFn valueToShape;
 };
 
 namespace detail {
 // Helper function to infer return tensor returns types given element and
 // shape inference function.
 //
 // TODO: Consider generating typedefs for trait member functions if this usage
 // becomes more common.

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

for (auto it : arrayAttr) {
     values.push_back(it.cast<FloatAttr>().getValueAsDouble());
   }
 }
 
 LogicalResult tosa::ArgMaxOp::inferReturnTypeComponents(
     MLIRContext *context, ::llvm::Optional<Location> location,
     ValueShapeRange operands, DictionaryAttr attributes, RegionRange regions,
     SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) {
-  ShapedType inputTy = operands[0].getType().cast<ShapedType>();
+  ShapeAdaptor inputShape = operands.getShape(0);
   IntegerAttr axis = attributes.get("axis").cast<IntegerAttr>();
   int32_t axisVal = axis.getValue().getSExtValue();
 
-  if (!inputTy.hasRank()) {
+  if (!inputShape.hasRank()) {
     inferredReturnShapes.push_back(ShapedTypeComponents());
     return success();
   }
 
   SmallVector<int64_t> outShape;
-  outShape.reserve(inputTy.getRank() - 1);
-  for (int i = 0, s = inputTy.getRank(); i < s; i++) {
+  outShape.reserve(inputShape.getRank() - 1);
+  for (int i = 0, s = inputShape.getRank(); i < s; i++) {
     if (i == axisVal)
       continue;
-    outShape.push_back(inputTy.getDimSize(i));
+    outShape.push_back(inputShape.getDimSize(i));
   }
 
   inferredReturnShapes.push_back(ShapedTypeComponents(outShape));
   return success();
 }
 
 LogicalResult tosa::ConcatOp::inferReturnTypeComponents(
     MLIRContext *context, ::llvm::Optional<Location> location,
     ValueShapeRange operands, DictionaryAttr attributes, RegionRange regions,
     SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) {
   // Infer all dimension sizes by reducing based on inputs.
   int32_t axis =
       attributes.get("axis").cast<IntegerAttr>().getValue().getSExtValue();
   llvm::SmallVector<int64_t> outputShape;
   bool hasRankedInput = false;
   for (auto operand : operands) {
-    ShapedType operandTy = operand.getType().cast<ShapedType>();
-    if (!operandTy.hasRank())
+    ShapeAdaptor operandShape = operands.getShape(operand);
+    if (!operandShape.hasRank())
       continue;
 
     // Copy the Operand's rank.
     if (!hasRankedInput)
-      outputShape.resize(operandTy.getRank(), ShapedType::kDynamicSize);
+      outputShape.resize(operandShape.getRank(), ShapedType::kDynamicSize);
 
     // Copy shapes until the dim is non-dynamic.
-    for (int i = 0, s = operandTy.getRank(); i < s; i++) {
-      if (i == axis || operandTy.isDynamicDim(i))
+    for (int i = 0, s = operandShape.getRank(); i < s; i++) {
+      if (i == axis || operandShape.isDynamicDim(i))
         continue;
       if (outputShape[i] == ShapedType::kDynamicSize)
-        outputShape[i] = operandTy.getDimSize(i);
-      if (outputShape[i] != operandTy.getDimSize(i))
+        outputShape[i] = operandShape.getDimSize(i);
+      if (outputShape[i] != operandShape.getDimSize(i))
         return failure();
     }
 
     hasRankedInput = true;
   }
 
   if (!hasRankedInput) {
     inferredReturnShapes.push_back(ShapedTypeComponents());
     return success();
   }
 
   // Determine the dimension size along the concatenation axis.
   int concatDimSize = 0;
   for (auto operand : operands) {
-    ShapedType operandTy = operand.getType().cast<ShapedType>();
+    ShapeAdaptor operandShape = operands.getShape(operand);
 
     // We need to know the length of the concatenation axis of all inputs to
     // determine the dimension size of the output shape.
-    if (!operandTy.hasRank() || operandTy.isDynamicDim(axis)) {
+    if (!operandShape.hasRank() || operandShape.isDynamicDim(axis)) {
       concatDimSize = ShapedType::kDynamicSize;
       break;
     }
 
-    concatDimSize += operandTy.getDimSize(axis);
+    concatDimSize += operandShape.getDimSize(axis);
   }
 
   outputShape[axis] = concatDimSize;
 
   inferredReturnShapes.push_back(ShapedTypeComponents(outputShape));
   return success();
 }
 
 LogicalResult tosa::FullyConnectedOp::inferReturnTypeComponents(
     MLIRContext *context, ::llvm::Optional<Location> location,
     ValueShapeRange operands, DictionaryAttr attributes, RegionRange regions,
     SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) {
-  ShapedType inputTy = operands[0].getType().cast<ShapedType>();
-  ShapedType weightTy = operands[1].getType().cast<ShapedType>();
-  ShapedType biasTy = operands[2].getType().cast<ShapedType>();
+  ShapeAdaptor inputShape = operands.getShape(0);
+  ShapeAdaptor weightShape = operands.getShape(1);
+  ShapeAdaptor biasShape = operands.getShape(2);
 
   // All shapes are dynamic.
   SmallVector<int64_t> outShape;
   outShape.resize(2, ShapedType::kDynamicSize);
 
-  if (inputTy.hasRank()) {
-    outShape[0] = inputTy.getDimSize(0);
+  if (inputShape.hasRank()) {
+    outShape[0] = inputShape.getDimSize(0);
   }
 
-  if (weightTy.hasRank()) {
-    outShape[1] = weightTy.getDimSize(0);
+  if (weightShape.hasRank()) {
+    outShape[1] = weightShape.getDimSize(0);
   }
 
-  if (biasTy.hasRank()) {
-    outShape[1] = outShape[1] == ShapedType::kDynamicSize ? biasTy.getDimSize(0)
-                                                          : outShape[1];
+  if (biasShape.hasRank()) {
+    outShape[1] = outShape[1] == ShapedType::kDynamicSize
+                      ? biasShape.getDimSize(0)
   }
 
   inferredReturnShapes.push_back(ShapedTypeComponents(outShape));
   return success();
 }
 
 LogicalResult tosa::MatMulOp::inferReturnTypeComponents(
     MLIRContext *context, ::llvm::Optional<Location> location,
     ValueShapeRange operands, DictionaryAttr attributes, RegionRange regions,
     SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) {
-  ShapedType lhsTy = operands[0].getType().cast<ShapedType>();
-  ShapedType rhsTy = operands[1].getType().cast<ShapedType>();
+  ShapeAdaptor lhsShape = operands.getShape(0);
+  ShapeAdaptor rhsShape = operands.getShape(1);
 
   // All shapes are dynamic.
   SmallVector<int64_t> outShape;
   outShape.resize(3, ShapedType::kDynamicSize);
 
-  if (lhsTy.hasRank()) {
-    outShape[0] = lhsTy.getDimSize(0);
-    outShape[1] = lhsTy.getDimSize(1);
+  if (lhsShape.hasRank()) {
+    outShape[0] = lhsShape.getDimSize(0);
+    outShape[1] = lhsShape.getDimSize(1);
   }
 
-  if (rhsTy.hasRank()) {
-    outShape[0] = outShape[0] == ShapedType::kDynamicSize ? rhsTy.getDimSize(0)
-                                                          : outShape[0];
-    outShape[2] = rhsTy.getDimSize(2);
+  if (rhsShape.hasRank()) {
+    outShape[0] = outShape[0] == ShapedType::kDynamicSize
+                      ? rhsShape.getDimSize(0)
+    outShape[2] = rhsShape.getDimSize(2);
   }
 
   inferredReturnShapes.push_back(ShapedTypeComponents(outShape));
   return success();
 }
 
 LogicalResult tosa::PadOp::inferReturnTypeComponents(
     MLIRContext *context, ::llvm::Optional<Location> location,
     ValueShapeRange operands, DictionaryAttr attributes, RegionRange regions,
     SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) {
-  ShapedType inputTy = operands[0].getType().cast<ShapedType>();
-  ShapedType paddingTy = operands[1].getType().cast<ShapedType>();
+  ShapeAdaptor inputShape = operands.getShape(0);
+  ShapeAdaptor paddingShape = operands.getShape(1);
   SmallVector<int64_t> outputShape;
 
   // If both inputs have unknown shape, we cannot determine the shape of the
   // output.
-  if (!inputTy.hasRank() && !paddingTy.hasRank()) {
+  if (!inputShape.hasRank() && !paddingShape.hasRank()) {
     inferredReturnShapes.push_back(ShapedTypeComponents());
     return success();
   }
 
   // If the input rank is unknown we can info the output rank using the padding
   // shape's first dim.
-  if (!inputTy.hasRank()) {
-    if (paddingTy.isDynamicDim(0)) {
+  if (!inputShape.hasRank()) {
+    if (paddingShape.isDynamicDim(0)) {
       inferredReturnShapes.push_back(ShapedTypeComponents());
       return success();
     }
 
-    outputShape.resize(paddingTy.getDimSize(0), ShapedType::kDynamicSize);
+    outputShape.resize(paddingShape.getDimSize(0), ShapedType::kDynamicSize);
     inferredReturnShapes.push_back(ShapedTypeComponents(outputShape));
     return success();
   }
 
   DenseIntElementsAttr paddings;
   // If the paddings value is not a constant, all dimensions must be dynamic.
   if (!matchPattern(operands[1], m_Constant(&paddings))) {
-    outputShape.resize(inputTy.getRank(), ShapedType::kDynamicSize);
+    outputShape.resize(inputShape.getRank(), ShapedType::kDynamicSize);
     inferredReturnShapes.push_back(ShapedTypeComponents(outputShape));
     return success();
   }
 
   SmallVector<int64_t> paddingValues;
   for (auto val : paddings) {
     paddingValues.push_back(val.getSExtValue());
   }
 
-  outputShape.reserve(inputTy.getRank());
-  for (int i = 0, s = inputTy.getRank(); i < s; i++) {
-    if (inputTy.isDynamicDim(i)) {
+  outputShape.reserve(inputShape.getRank());
+  for (int i = 0, s = inputShape.getRank(); i < s; i++) {
+    if (inputShape.isDynamicDim(i)) {
       outputShape.push_back(ShapedType::kDynamicSize);
       continue;
     }
 
-    outputShape.push_back(inputTy.getDimSize(i) + paddingValues[i * 2] +
+    outputShape.push_back(inputShape.getDimSize(i) + paddingValues[i * 2] +
                           paddingValues[i * 2 + 1]);
   }
 
   inferredReturnShapes.push_back(ShapedTypeComponents(outputShape));
   return success();
 }
 
 LogicalResult tosa::SliceOp::inferReturnTypeComponents(
     MLIRContext *context, ::llvm::Optional<Location> location,
     ValueShapeRange operands, DictionaryAttr attributes, RegionRange regions,
     SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) {
-  auto sizes = attributes.get("size").cast<ArrayAttr>().getValue();
+  ArrayAttr sizes = SliceOpAdaptor(operands, attributes).size();
   SmallVector<int64_t> outputShape;
   outputShape.reserve(sizes.size());
   for (auto val : sizes) {
     outputShape.push_back(val.cast<IntegerAttr>().getValue().getSExtValue());
   }
 
   inferredReturnShapes.push_back(ShapedTypeComponents(outputShape));
   return success();
 }
 
 LogicalResult tosa::TableOp::inferReturnTypeComponents(
     MLIRContext *context, ::llvm::Optional<Location> location,
     ValueShapeRange operands, DictionaryAttr attributes, RegionRange regions,
     SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) {
-  ShapedType inputTy = operands[0].getType().cast<ShapedType>();
+  ShapeAdaptor inputShape = operands.getShape(0);
 
-  if (!inputTy.hasRank()) {
+  if (!inputShape.hasRank()) {
     inferredReturnShapes.push_back(ShapedTypeComponents());
     return success();
   }
 
-  inferredReturnShapes.push_back(inputTy.getShape());
+  inferredReturnShapes.resize(1);
+  inputShape.getDims(inferredReturnShapes[0]);
   return success();
 }
 
 LogicalResult tosa::TileOp::inferReturnTypeComponents(
     MLIRContext *context, ::llvm::Optional<Location> location,
     ValueShapeRange operands, DictionaryAttr attributes, RegionRange regions,
     SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) {
-  auto multiples = attributes.get("multiples").cast<ArrayAttr>().getValue();
-  ShapedType inputTy = operands[0].getType().cast<ShapedType>();
+  TileOpAdaptor adaptor(operands, attributes);
+  ArrayAttr multiples = adaptor.multiples();
+  ShapeAdaptor inputShape = operands.getShape(0);
   SmallVector<int64_t> outputShape;
-  if (!inputTy.hasRank()) {
+  if (!inputShape.hasRank()) {
     outputShape.resize(multiples.size(), ShapedType::kDynamicSize);
     inferredReturnShapes.push_back(ShapedTypeComponents(outputShape));
     return success();
   }
 
   // We need the multiple values to determine the output shape.
   SmallVector<int64_t> multipleValues;
   multipleValues.reserve(multiples.size());
   for (auto val : multiples) {
     multipleValues.push_back(val.cast<IntegerAttr>().getValue().getSExtValue());
   }
 
   // Any non dynamic dimension can be multiplied to a known size.
   outputShape.reserve(multiples.size());
-  for (int i = 0, s = inputTy.getRank(); i < s; i++) {
-    int dim = inputTy.getDimSize(i);
+  for (int i = 0, s = inputShape.getRank(); i < s; i++) {
+    int dim = inputShape.getDimSize(i);
     if (dim != ShapedType::kDynamicSize)
       dim *= multipleValues[i];
     outputShape.push_back(dim);
   }
 
   inferredReturnShapes.push_back(ShapedTypeComponents(outputShape));
   return success();
 }
 
 LogicalResult tosa::ReshapeOp::inferReturnTypeComponents(
     MLIRContext *context, ::llvm::Optional<Location> location,
     ValueShapeRange operands, DictionaryAttr attributes, RegionRange regions,
     SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) {
-  ShapedType type = operands.front().getType().cast<ShapedType>();
+  ReshapeOpAdaptor adaptor(operands, attributes);
+  ShapeAdaptor inputShape = operands.getShape(0);
 
-  auto newShape = attributes.get("new_shape").cast<ArrayAttr>();
+  ArrayAttr newShape = adaptor.new_shape();
   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()) {
+  if (!inputShape.hasRank() || !inputShape.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 numElements = inputShape.getNumElements();
   int64_t staticMul = 1;
   for (auto val : newShapeValue) {
     if (val != ShapedType::kDynamicSize) {
       staticMul *= val;
     }
   }
 
   // Determine the length of the dynamic dimension.
   for (auto &val : newShapeValue) {
     if (val == ShapedType::kDynamicSize)
       val = numElements / staticMul;
   }
 
   inferredReturnShapes.push_back(ShapedTypeComponents(newShapeValue));
   return success();
 }
 
 LogicalResult tosa::TransposeOp::inferReturnTypeComponents(
     MLIRContext *context, ::llvm::Optional<Location> location,
     ValueShapeRange operands, DictionaryAttr attributes, RegionRange regions,
     SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) {
-  ShapedType inputTy = operands[0].getType().cast<ShapedType>();
-  ShapedType permsTy = operands[1].getType().cast<ShapedType>();
+  ShapeAdaptor inputShape = operands.getShape(0);
+  ShapeAdaptor permsShape = operands.getShape(1);
 
   // If input rank and permutation length is unknown, the output rank is
   // unknown.
-  if (!inputTy.hasRank() && (!permsTy.hasRank() || permsTy.isDynamicDim(0))) {
+  if (!inputShape.hasRank() &&
+      (!permsShape.hasRank() || permsShape.isDynamicDim(0))) {
     inferredReturnShapes.push_back(ShapedTypeComponents());
     return success();
   }
 
   // Without the input dims we cannot determine the output dim sizes but we
   // can determine the output rank.
   SmallVector<int64_t> outputShape;
-  if (!inputTy.hasRank()) {
-    outputShape.resize(permsTy.getDimSize(0), ShapedType::kDynamicSize);
+  if (!inputShape.hasRank()) {
+    outputShape.resize(permsShape.getDimSize(0), ShapedType::kDynamicSize);
     inferredReturnShapes.push_back(ShapedTypeComponents(outputShape));
     return success();
   }
 
   // Rank-0 means no permutations matter.
-  if (inputTy.getRank() == 0) {
+  if (inputShape.getRank() == 0) {
     inferredReturnShapes.push_back(ShapedTypeComponents(outputShape));
     return success();
   }
 
   // Check whether the input dimensions are all the same.
   bool allTheSame = true;
-  for (int i = 1, s = inputTy.getRank(); i < s; i++) {
-    if (inputTy.getDimSize(0) != inputTy.getDimSize(i)) {
+  for (int i = 1, s = inputShape.getRank(); i < s; i++) {
+    if (inputShape.getDimSize(0) != inputShape.getDimSize(i)) {
       allTheSame = false;
       break;
     }
   }
 
   // If all of the input dimensions are the same we don't care about the
   // permutation.
   if (allTheSame) {
-    outputShape.resize(inputTy.getRank(), inputTy.getDimSize(0));
+    outputShape.resize(inputShape.getRank(), inputShape.getDimSize(0));
     inferredReturnShapes.push_back(ShapedTypeComponents(outputShape));
     return success();
   }
 
-  DenseIntElementsAttr perms;
-  outputShape.resize(inputTy.getRank(), ShapedType::kDynamicSize);
+  outputShape.resize(inputShape.getRank(), ShapedType::kDynamicSize);
   // If the permuations are a constant we can directly determine the output
   // shape.
-  if (matchPattern(operands[1], m_Constant(&perms))) {
-    llvm::SmallVector<int64_t> permValues;
-    for (auto val : perms) {
-      permValues.push_back(val.getSExtValue());
-    }
-
-    outputShape.reserve(inputTy.getRank());
-    for (int i = 0, s = inputTy.getRank(); i < s; i++) {
-      outputShape[i] = inputTy.getDimSize(permValues[i]);
+  if (ShapeAdaptor permShape = operands.getValueAsShape(1)) {
+    outputShape.reserve(inputShape.getRank());
+    for (int i = 0, s = inputShape.getRank(); i < s; i++) {
+      outputShape[i] = inputShape.getDimSize(permShape.getDimSize(i));
     }
   }
 
   inferredReturnShapes.push_back(ShapedTypeComponents(outputShape));
   return success();
 }
 
 LogicalResult tosa::GatherOp::inferReturnTypeComponents(
     MLIRContext *context, ::llvm::Optional<Location> location,
     ValueShapeRange operands, DictionaryAttr attributes, RegionRange regions,
     SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) {
   llvm::SmallVector<int64_t> outputShape;
   outputShape.resize(3, ShapedType::kDynamicSize);
 
-  if (auto ty = operands[0].getType().dyn_cast<RankedTensorType>()) {
-    outputShape[0] = ty.getDimSize(0);
-    outputShape[2] = ty.getDimSize(2);
+  ShapeAdaptor valuesShape = operands.getShape(0);
+  if (valuesShape.hasRank()) {
+    outputShape[0] = valuesShape.getDimSize(0);
+    outputShape[2] = valuesShape.getDimSize(2);
   }
 
-  if (auto ty = operands[1].getType().dyn_cast<RankedTensorType>()) {
+  ShapeAdaptor indicesShape = operands.getShape(1);
+  if (indicesShape.hasRank()) {
     if (outputShape[0] == ShapedType::kDynamicSize)
-      outputShape[0] = ty.getDimSize(0);
+      outputShape[0] = indicesShape.getDimSize(0);
     if (outputShape[1] == ShapedType::kDynamicSize)
-      outputShape[1] = ty.getDimSize(1);
+      outputShape[1] = indicesShape.getDimSize(1);
   }
 
   inferredReturnShapes.push_back(ShapedTypeComponents(outputShape));
   return success();
 }
 
 LogicalResult tosa::ResizeOp::inferReturnTypeComponents(
     MLIRContext *context, ::llvm::Optional<Location> location,
     ValueShapeRange operands, DictionaryAttr attributes, RegionRange regions,
     SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) {
+  ResizeOpAdaptor adaptor(operands, attributes);
   llvm::SmallVector<int64_t, 4> outputShape;
   outputShape.resize(4, ShapedType::kDynamicSize);
 
   int32_t inHeight = ShapedType::kDynamicSize;
   int32_t inWidth = ShapedType::kDynamicSize;
 
-  if (auto ty = operands[0].getType().dyn_cast<RankedTensorType>()) {
-    outputShape[0] = ty.getDimSize(0);
-    outputShape[3] = ty.getDimSize(3);
+  ShapeAdaptor inputShape = operands.getShape(adaptor.input());
+  if (inputShape.hasRank()) {
+    outputShape[0] = inputShape.getDimSize(0);
+    outputShape[3] = inputShape.getDimSize(3);
 
-    inHeight = ty.getDimSize(1);
-    inWidth = ty.getDimSize(2);
+    inHeight = inputShape.getDimSize(1);
+    inWidth = inputShape.getDimSize(2);
   }
 
-  int32_t shift =
-      attributes.get("shift").cast<IntegerAttr>().getValue().getSExtValue();
+  int32_t shift = adaptor.shift().getValue().getSExtValue();
   llvm::SmallVector<int64_t> newShape;
-  getI64Values(attributes.get("output_size").cast<ArrayAttr>(), newShape);
+  getI64Values(adaptor.output_size(), newShape);
   outputShape[1] = newShape[0];
   outputShape[2] = newShape[1];
 
   llvm::SmallVector<int64_t> strideInt;
   llvm::SmallVector<int64_t> offsetInt;
   llvm::SmallVector<double> strideFp;
   llvm::SmallVector<double> offsetFp;
-  getI64Values(attributes.get("offset").cast<ArrayAttr>(), offsetInt);
-  getF64Values(attributes.get("offset_fp").cast<ArrayAttr>(), offsetFp);
-  getI64Values(attributes.get("stride").cast<ArrayAttr>(), strideInt);
-  getF64Values(attributes.get("stride_fp").cast<ArrayAttr>(), strideFp);
+  getI64Values(adaptor.offset(), offsetInt);
+  getF64Values(adaptor.offset_fp(), offsetFp);
+  getI64Values(adaptor.stride(), strideInt);
+  getF64Values(adaptor.stride_fp(), strideFp);
 
   // If we have a 0 zero in integers we know that the resize indexing needs to
   // be performed in floating point. Use the floating point varient to compute
   // the resize shape.
   bool fpMode = strideInt[0] == 0;
 
   // We can compute the output shape if attribute specifies unknown dimensions
   // based on the offset and stride. If we perfectly line up to the last index

 
 LogicalResult tosa::ScatterOp::inferReturnTypeComponents(
     MLIRContext *context, ::llvm::Optional<Location> location,
     ValueShapeRange operands, DictionaryAttr attributes, RegionRange regions,
     SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) {
   llvm::SmallVector<int64_t> outputShape;
   outputShape.resize(3, ShapedType::kDynamicSize);
 
-  if (auto ty = operands[0].getType().dyn_cast<RankedTensorType>()) {
-    outputShape[0] = ty.getDimSize(0);
-    outputShape[1] = ty.getDimSize(1);
-    outputShape[2] = ty.getDimSize(2);
+  ShapeAdaptor valuesInShape = operands.getShape(0);
+  if (valuesInShape.hasRank()) {
+    outputShape[0] = valuesInShape.getDimSize(0);
+    outputShape[1] = valuesInShape.getDimSize(1);
+    outputShape[2] = valuesInShape.getDimSize(2);
   }
 
-  if (auto ty = operands[1].getType().dyn_cast<RankedTensorType>()) {
+  ShapeAdaptor indicesShape = operands.getShape(1);
+  if (indicesShape.hasRank()) {
     if (outputShape[0] == ShapedType::kDynamicSize)
-      outputShape[0] = ty.getDimSize(0);
+      outputShape[0] = indicesShape.getDimSize(0);
   }
 
-  if (auto ty = operands[2].getType().dyn_cast<RankedTensorType>()) {
+  ShapeAdaptor inputShape = operands.getShape(2);
+  if (inputShape.hasRank()) {
     if (outputShape[0] == ShapedType::kDynamicSize)
-      outputShape[0] = ty.getDimSize(0);
+      outputShape[0] = inputShape.getDimSize(0);
     if (outputShape[2] == ShapedType::kDynamicSize)
-      outputShape[2] = ty.getDimSize(2);
+      outputShape[2] = inputShape.getDimSize(2);
   }
 
   inferredReturnShapes.push_back(ShapedTypeComponents(outputShape));
   return success();
 }
 
 static LogicalResult ReduceInferReturnTypes(
-    Value operand, IntegerAttr axis,
+    ShapeAdaptor operandShape, IntegerAttr axis,
     SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) {
-  auto operandTy = operand.getType().cast<ShapedType>();
-  if (!operandTy.hasRank()) {
+  if (!operandShape.hasRank()) {
     inferredReturnShapes.push_back(ShapedTypeComponents());
     return success();
   }
 
-  int64_t axisVal = axis.getValue().getSExtValue();
   SmallVector<int64_t> outputShape;
-  outputShape.reserve(operandTy.getRank());
-  for (auto dim : operandTy.getShape()) {
-    outputShape.push_back(dim);
-  }
-
+  operandShape.getDims(outputShape);
+  int64_t axisVal = axis.getValue().getSExtValue();
   outputShape[axisVal] = 1;
   inferredReturnShapes.push_back(ShapedTypeComponents(outputShape));
   return success();
 }
 
 #define REDUCE_SHAPE_INFER(OP)                                                 \
   LogicalResult OP::inferReturnTypeComponents(                                 \
       MLIRContext *context, ::llvm::Optional<Location> location,               \
       ValueShapeRange operands, DictionaryAttr attributes,                     \
       RegionRange regions,                                                     \
       SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) {           \
-    return ReduceInferReturnTypes(operands[0],                                 \
+    return ReduceInferReturnTypes(operands.getShape(0),                        \
                                   attributes.get("axis").cast<IntegerAttr>(),  \
                                   inferredReturnShapes);                       \
   }
 
 REDUCE_SHAPE_INFER(tosa::ReduceAllOp)
 REDUCE_SHAPE_INFER(tosa::ReduceAnyOp)
 REDUCE_SHAPE_INFER(tosa::ReduceMaxOp)
 REDUCE_SHAPE_INFER(tosa::ReduceMinOp)
 REDUCE_SHAPE_INFER(tosa::ReduceProdOp)
 REDUCE_SHAPE_INFER(tosa::ReduceSumOp)
 #undef REDUCE_SHAPE_INFER
 
-static LogicalResult resolveBroadcastShape(ValueRange operands,
+static LogicalResult resolveBroadcastShape(const ValueShapeRange &operands,
                                            SmallVector<int64_t> &outShape) {
   int64_t outRank = 0;
-  for (auto operand : operands) {
-    auto type = operand.getType().cast<ShapedType>();
-    if (!type.hasRank())
+  for (int i = 0, e = operands.size(); i != e; ++i) {
+    auto shape = operands.getShape(i);
+    if (!shape.hasRank()) {
       return failure();
-    outRank = std::max<int64_t>(outRank, type.getRank());
+    }
+    outRank = std::max<int64_t>(outRank, shape.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 (int i = 0, e = operands.size(); i != e; ++i) {
+    auto shape = operands.getShape(i);
+    auto rankDiff = outShape.size() - shape.getRank();
 
-    for (size_t i = 0; i < shape.size(); i++) {
+    for (size_t i = 0, e = shape.getRank(); i < e; ++i) {
       auto dim1 = outShape[i + rankDiff];
-      auto dim2 = shape[i];
+      auto dim2 = shape.getDimSize(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,
+    const ValueShapeRange &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();

 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
 
 static LogicalResult poolingInferReturnTypes(
-    ValueRange operands, DictionaryAttr attributes,
+    const ValueShapeRange &operands, DictionaryAttr attributes,
     SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) {
-  RankedTensorType inputTy = operands[0].getType().dyn_cast<RankedTensorType>();
+  ShapeAdaptor inputShape = operands.getShape(0);
   llvm::SmallVector<int64_t> outputShape;
   outputShape.resize(4, -1);
 
   // We only know the rank if the input type is unranked.
-  if (!inputTy) {
+  if (!inputShape) {
     inferredReturnShapes.push_back(ShapedTypeComponents(outputShape));
     return success();
   }
 
   // Batch and number of channels are identical for pooling layer.
-  outputShape[0] = inputTy.getDimSize(0);
-  outputShape[3] = inputTy.getDimSize(3);
+  outputShape[0] = inputShape.getDimSize(0);
+  outputShape[3] = inputShape.getDimSize(3);
 
-  int32_t height = inputTy.getDimSize(1);
-  int32_t width = inputTy.getDimSize(2);
+  int32_t height = inputShape.getDimSize(1);
+  int32_t width = inputShape.getDimSize(2);
 
   llvm::SmallVector<int64_t> kernel;
   llvm::SmallVector<int64_t> stride;
   llvm::SmallVector<int64_t> pad;
 
   getI64Values(attributes.get("kernel").cast<ArrayAttr>(), kernel);
   getI64Values(attributes.get("stride").cast<ArrayAttr>(), stride);
   getI64Values(attributes.get("pad").cast<ArrayAttr>(), pad);
static LogicalResult poolingInferReturnTypes(
   return success();
 }
 
 LogicalResult Conv2DOp::inferReturnTypeComponents(
     MLIRContext *context, ::llvm::Optional<Location> location,
     ValueShapeRange operands, DictionaryAttr attributes, RegionRange regions,
     SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) {
   llvm::SmallVector<int64_t> outputShape(4, ShapedType::kDynamicSize);
-  Conv2DOp::Adaptor adaptor(operands.getValues());
+  Conv2DOp::Adaptor adaptor(operands.getValues(), attributes);
 
   int32_t inputWidth = ShapedType::kDynamicSize;
   int32_t inputHeight = ShapedType::kDynamicSize;
   int32_t weightWidth = ShapedType::kDynamicSize;
   int32_t weightHeight = ShapedType::kDynamicSize;
 
   // Input shape describes input width/height and batch.
-  if (auto inputTy = adaptor.input().getType().dyn_cast<RankedTensorType>()) {
-    outputShape[0] = inputTy.getDimSize(0);
-    inputHeight = inputTy.getDimSize(1);
-    inputWidth = inputTy.getDimSize(2);
+
+  ShapeAdaptor inputShape = operands.getShape(adaptor.input());
+  if (inputShape.hasRank()) {
+    outputShape[0] = inputShape.getDimSize(0);
+    inputHeight = inputShape.getDimSize(1);
+    inputWidth = inputShape.getDimSize(2);
   }
 
   // Weight shapes describes the filter width/height and the output channels.
-  if (auto weightTy = adaptor.weight().getType().dyn_cast<RankedTensorType>()) {
-    outputShape[3] = weightTy.getDimSize(0);
-    weightHeight = weightTy.getDimSize(1);
-    weightWidth = weightTy.getDimSize(2);
+  ShapeAdaptor weightShape = operands.getShape(adaptor.weight());
+  if (weightShape.hasRank()) {
+    outputShape[3] = weightShape.getDimSize(0);
+    weightHeight = weightShape.getDimSize(1);
+    weightWidth = weightShape.getDimSize(2);
   }
 
   // Bias shape can describe the output channels.
-  if (auto biasTy = adaptor.bias().getType().dyn_cast<RankedTensorType>()) {
+  ShapeAdaptor biasShape = operands.getShape(adaptor.bias());
+  if (biasShape.hasRank()) {
     outputShape[3] = ShapedType::isDynamic(outputShape[3])
-                         ? biasTy.getDimSize(0)
+                         ? biasShape.getDimSize(0)
                          : outputShape[3];
   }
 
   llvm::SmallVector<int64_t> dilation;
   llvm::SmallVector<int64_t> padding;
   llvm::SmallVector<int64_t> stride;
 
-  getI64Values(attributes.get("dilation").cast<ArrayAttr>(), dilation);
-  getI64Values(attributes.get("pad").cast<ArrayAttr>(), padding);
-  getI64Values(attributes.get("stride").cast<ArrayAttr>(), stride);
+  getI64Values(adaptor.dilation(), dilation);
+  getI64Values(adaptor.pad(), padding);
+  getI64Values(adaptor.stride(), stride);
 
   if (!ShapedType::isDynamic(inputHeight) &&
       !ShapedType::isDynamic(weightHeight)) {
     int32_t inputSize = inputHeight + padding[0] + padding[1];
     int32_t filterSize = (weightHeight - 1) * dilation[0] + 1;
     int32_t unstridedResult = inputSize - filterSize + 1;
     outputShape[1] = (unstridedResult - 1) / stride[0] + 1;
   }
LogicalResult Conv2DOp::inferReturnTypeComponents(
   return success();
 }
 
 LogicalResult Conv3DOp::inferReturnTypeComponents(
     MLIRContext *context, ::llvm::Optional<Location> location,
     ValueShapeRange operands, DictionaryAttr attributes, RegionRange regions,
     SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) {
   llvm::SmallVector<int64_t> outputShape(5, ShapedType::kDynamicSize);
-  Conv2DOp::Adaptor adaptor(operands.getValues());
+  Conv2DOp::Adaptor adaptor(operands.getValues(), attributes);
 
   int32_t inputWidth = ShapedType::kDynamicSize;
   int32_t inputHeight = ShapedType::kDynamicSize;
   int32_t inputDepth = ShapedType::kDynamicSize;
 
   int32_t weightWidth = ShapedType::kDynamicSize;
   int32_t weightHeight = ShapedType::kDynamicSize;
   int32_t weightDepth = ShapedType::kDynamicSize;
 
   // Input shape describes input width/height and batch.
-  if (auto inputTy = adaptor.input().getType().dyn_cast<RankedTensorType>()) {
-    outputShape[0] = inputTy.getDimSize(0);
-    inputHeight = inputTy.getDimSize(1);
-    inputWidth = inputTy.getDimSize(2);
-    inputDepth = inputTy.getDimSize(3);
+  ShapeAdaptor inputShape = operands.getShape(adaptor.input());
+  if (inputShape.hasRank()) {
+    outputShape[0] = inputShape.getDimSize(0);
+    inputHeight = inputShape.getDimSize(1);
+    inputWidth = inputShape.getDimSize(2);
+    inputDepth = inputShape.getDimSize(3);
   }
 
   // Weight shapes describes the filter width/height and the output channels.
-  if (auto weightTy = adaptor.weight().getType().dyn_cast<RankedTensorType>()) {
-    outputShape[4] = weightTy.getDimSize(0);
-    weightHeight = weightTy.getDimSize(1);
-    weightWidth = weightTy.getDimSize(2);
-    weightDepth = weightTy.getDimSize(3);
+  ShapeAdaptor weightShape = operands.getShape(adaptor.weight());
+  if (weightShape.hasRank()) {
+    outputShape[4] = weightShape.getDimSize(0);
+    weightHeight = weightShape.getDimSize(1);
+    weightWidth = weightShape.getDimSize(2);
+    weightDepth = weightShape.getDimSize(3);
   }
 
   // Bias shape can describe the output channels.
-  if (auto biasTy = adaptor.bias().getType().dyn_cast<RankedTensorType>()) {
+  ShapeAdaptor biasShape = operands.getShape(adaptor.bias());
+  if (biasShape.hasRank()) {
     outputShape[4] =
-        (outputShape[4] == -1) ? biasTy.getDimSize(0) : outputShape[4];
+        (outputShape[4] == -1) ? biasShape.getDimSize(0) : outputShape[4];
   }
 
   llvm::SmallVector<int64_t> dilation;
   llvm::SmallVector<int64_t> padding;
   llvm::SmallVector<int64_t> stride;
 
-  getI64Values(attributes.get("dilation").cast<ArrayAttr>(), dilation);
-  getI64Values(attributes.get("pad").cast<ArrayAttr>(), padding);
-  getI64Values(attributes.get("stride").cast<ArrayAttr>(), stride);
+  getI64Values(adaptor.dilation(), dilation);
+  getI64Values(adaptor.pad(), padding);
+  getI64Values(adaptor.stride(), stride);
 
   if (!ShapedType::isDynamic(inputHeight) &&
       !ShapedType::isDynamic(weightHeight)) {
     int32_t inputSize = inputHeight + padding[0] + padding[1];
     int32_t filterSize = (weightHeight - 1) * dilation[0] + 1;
     int32_t unstridedResult = inputSize - filterSize + 1;
     outputShape[1] = (unstridedResult - 1) / stride[0] + 1;
   }
LogicalResult MaxPool2dOp::inferReturnTypeComponents(
   return poolingInferReturnTypes(operands, attributes, inferredReturnShapes);
 }
 
 LogicalResult DepthwiseConv2DOp::inferReturnTypeComponents(
     MLIRContext *context, ::llvm::Optional<Location> location,
     ValueShapeRange operands, DictionaryAttr attributes, RegionRange regions,
     SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) {
   llvm::SmallVector<int64_t> outputShape(4, ShapedType::kDynamicSize);
-  DepthwiseConv2DOp::Adaptor adaptor(operands.getValues());
+  DepthwiseConv2DOp::Adaptor adaptor(operands.getValues(), attributes);
 
   int32_t inputWidth = ShapedType::kDynamicSize;
   int32_t inputHeight = ShapedType::kDynamicSize;
   int32_t inputChannels = ShapedType::kDynamicSize;
 
   int32_t weightWidth = ShapedType::kDynamicSize;
   int32_t weightHeight = ShapedType::kDynamicSize;
   int32_t depthChannels = ShapedType::kDynamicSize;
 
   // Input shape describes input width/height and batch.
-  if (auto inputTy = adaptor.input().getType().dyn_cast<RankedTensorType>()) {
-    outputShape[0] = inputTy.getDimSize(0);
-    inputHeight = inputTy.getDimSize(1);
-    inputWidth = inputTy.getDimSize(2);
-    inputChannels = inputTy.getDimSize(3);
+  ShapeAdaptor inputShape = operands.getShape(adaptor.input());
+  if (inputShape.hasRank()) {
+    outputShape[0] = inputShape.getDimSize(0);
+    inputHeight = inputShape.getDimSize(1);
+    inputWidth = inputShape.getDimSize(2);
+    inputChannels = inputShape.getDimSize(3);
   }
 
   // Weight shapes describes the filter width/height and the output channels.
-  if (auto weightTy = adaptor.weight().getType().dyn_cast<RankedTensorType>()) {
-    weightHeight = weightTy.getDimSize(0);
-    weightWidth = weightTy.getDimSize(1);
+  ShapeAdaptor weightShape = operands.getShape(adaptor.weight());
+  if (weightShape.hasRank()) {
+    weightHeight = weightShape.getDimSize(0);
+    weightWidth = weightShape.getDimSize(1);
     inputChannels = ShapedType::isDynamic(inputChannels)
-                        ? weightTy.getDimSize(2)
+                        ? weightShape.getDimSize(2)
                         : inputChannels;
-    depthChannels = weightTy.getDimSize(3);
+    depthChannels = weightShape.getDimSize(3);
   }
 
   // If both inputChannels and depthChannels are available we can determine
   // the output channels.
   if (!ShapedType::isDynamic(inputChannels) &&
       !ShapedType::isDynamic(depthChannels)) {
     outputShape[3] = inputChannels * depthChannels;
   }
 
   // Bias shape can describe the output channels.
-  if (auto biasTy = adaptor.bias().getType().dyn_cast<RankedTensorType>()) {
+  ShapeAdaptor biasShape = operands.getShape(adaptor.bias());
+  if (biasShape.hasRank()) {
     outputShape[3] = ShapedType::isDynamic(outputShape[3])
-                         ? biasTy.getDimSize(0)
+                         ? biasShape.getDimSize(0)
                          : outputShape[3];
   }
 
   llvm::SmallVector<int64_t> dilation;
   llvm::SmallVector<int64_t> padding;
   llvm::SmallVector<int64_t> stride;
 
-  getI64Values(attributes.get("dilation").cast<ArrayAttr>(), dilation);
-  getI64Values(attributes.get("pad").cast<ArrayAttr>(), padding);
-  getI64Values(attributes.get("stride").cast<ArrayAttr>(), stride);
+  getI64Values(adaptor.dilation(), dilation);
+  getI64Values(adaptor.pad(), padding);
+  getI64Values(adaptor.stride(), stride);
 
   if (!ShapedType::isDynamic(inputHeight) &&
       !ShapedType::isDynamic(weightHeight)) {
     int32_t inputSize = inputHeight + padding[0] + padding[1];
     int32_t filterSize = (weightHeight - 1) * dilation[0] + 1;
     int32_t unstridedResult = inputSize - filterSize + 1;
     outputShape[1] = (unstridedResult - 1) / stride[0] + 1;
   }
LogicalResult DepthwiseConv2DOp::inferReturnTypeComponents(
   inferredReturnShapes.push_back(ShapedTypeComponents(outputShape));
   return success();
 }
 
 LogicalResult TransposeConv2DOp::inferReturnTypeComponents(
     MLIRContext *context, ::llvm::Optional<Location> location,
     ValueShapeRange operands, DictionaryAttr attributes, RegionRange regions,
     SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) {
-  TransposeConv2DOp::Adaptor adaptor(operands.getValues());
+  TransposeConv2DOp::Adaptor adaptor(operands.getValues(), attributes);
   llvm::SmallVector<int64_t> outputShape;
-  getI64Values(attributes.get("out_shape").cast<ArrayAttr>(), outputShape);
+  getI64Values(adaptor.out_shape(), outputShape);
 
   int32_t inputWidth = ShapedType::kDynamicSize;
   int32_t inputHeight = ShapedType::kDynamicSize;
   int32_t weightWidth = ShapedType::kDynamicSize;
   int32_t weightHeight = ShapedType::kDynamicSize;
 
   // Input shape describes input width/height and batch.
-  if (auto inputTy = adaptor.input().getType().dyn_cast<RankedTensorType>()) {
+  ShapeAdaptor inputShape = operands.getShape(adaptor.input());
+  if (inputShape.hasRank()) {
     outputShape[0] = ShapedType::isDynamic(outputShape[0])
-                         ? inputTy.getDimSize(0)
+                         ? inputShape.getDimSize(0)
                          : outputShape[0];
-    inputHeight = inputTy.getDimSize(1);
-    inputWidth = inputTy.getDimSize(2);
+    inputHeight = inputShape.getDimSize(1);
+    inputWidth = inputShape.getDimSize(2);
   }
 
   // Weight shapes describes the filter width/height and the output channels.
-  if (auto weightTy = adaptor.filter().getType().dyn_cast<RankedTensorType>()) {
+  ShapeAdaptor weightShape = operands.getShape(adaptor.input());
+  if (weightShape.hasRank()) {
     outputShape[3] = ShapedType::isDynamic(outputShape[3])
-                         ? weightTy.getDimSize(0)
+                         ? weightShape.getDimSize(0)
                          : outputShape[3];
-    weightHeight = weightTy.getDimSize(1);
-    weightWidth = weightTy.getDimSize(2);
+    weightHeight = weightShape.getDimSize(1);
+    weightWidth = weightShape.getDimSize(2);
   }
 
   // Bias shape can describe the output channels.
-  if (auto biasTy = adaptor.bias().getType().dyn_cast<RankedTensorType>()) {
+  ShapeAdaptor biasShape = operands.getShape(adaptor.input());
+  if (biasShape.hasRank()) {
     outputShape[3] = ShapedType::isDynamic(outputShape[3])
-                         ? biasTy.getDimSize(0)
+                         ? biasShape.getDimSize(0)
                          : outputShape[3];
   }
 
   llvm::SmallVector<int64_t> dilation;
   llvm::SmallVector<int64_t> padding;
   llvm::SmallVector<int64_t> stride;
 
-  getI64Values(attributes.get("dilation").cast<ArrayAttr>(), dilation);
-  getI64Values(attributes.get("out_pad").cast<ArrayAttr>(), padding);
-  getI64Values(attributes.get("stride").cast<ArrayAttr>(), stride);
+  getI64Values(adaptor.dilation(), dilation);
+  getI64Values(adaptor.out_pad(), padding);
+  getI64Values(adaptor.stride(), stride);
 
   if (!ShapedType::isDynamic(inputHeight) &&
       !ShapedType::isDynamic(weightHeight)) {
     int32_t dilated = (weightHeight - 1) * dilation[0] + 1;
     int32_t calculateSize =
         (inputHeight - 1) * stride[0] - padding[0] + dilated;
     outputShape[1] = outputShape[1] == -1 ? calculateSize : outputShape[1];
   }
for (auto it : llvm::enumerate(yieldOp.getOperands())) {
           resultKnowledge[index],
           ValueKnowledge::getKnowledgeFromType(it.value().getType()));
       if (!meet)
         continue;
       resultKnowledge[index] = meet;
     }
   }
 
-  for (auto result : resultKnowledge) {
+  for (const ValueKnowledge &result : resultKnowledge) {
     if (result.hasRank) {
       inferredReturnShapes.push_back(ShapedTypeComponents(result.sizes));
     } else {
       inferredReturnShapes.push_back(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

 #include "mlir/Dialect/Tosa/Utils/ShapeUtils.h"
 #include "mlir/IR/BlockAndValueMapping.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/Matchers.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Transforms/DialectConversion.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+#include "llvm/Support/FormatVariadic.h"
 
 using namespace mlir;
 using namespace mlir::tosa;
 
 namespace {
 
 void propagateShapesInRegion(Region &region);
 
for (auto &region : op.getRegions()) {
 
     propagateShapesInRegion(region);
   }
 
   return;
 }
 
 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);
 
       InferShapedTypeOpInterface shapeInterface =
           dyn_cast<InferShapedTypeOpInterface>(op);
       if (!shapeInterface)
         continue;
 
       SmallVector<ShapedTypeComponents> returnedShapes;
+
+      ValueShapeRange range(op.getOperands(), operandShape);
       if (shapeInterface
-              .inferReturnTypeComponents(
-                  op.getContext(), op.getLoc(), op.getOperands(),
-                  op.getAttrDictionary(), op.getRegions(), returnedShapes)
+              .inferReturnTypeComponents(op.getContext(), op.getLoc(), range,
+                                         op.getAttrDictionary(),
+                                         op.getRegions(), returnedShapes)
               .succeeded()) {
         for (auto it : llvm::zip(op.getResults(), returnedShapes)) {
           Value result = std::get<0>(it);
           ShapedTypeComponents predictedShape = std::get<1>(it);
 
           // 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())
               continue;
 
             replaceable = false;
           }
 
           // Determine the knowledge based on the output type.
+          // TODO: should also query WIP type probably
           Type resultTy = result.getType();
           auto currentKnowledge =
               ValueKnowledge::getKnowledgeFromType(resultTy);
 
           // Compute the knowledge based on the inferred type.
           auto inferredKnowledge = ValueKnowledge::getPessimisticValueState();
           inferredKnowledge.dtype =
               resultTy.cast<ShapedType>().getElementType();
           inferredKnowledge.hasRank = predictedShape.hasRank();
           if (predictedShape.hasRank()) {
             for (auto dim : predictedShape.getDims()) {
               inferredKnowledge.sizes.push_back(dim);
             }
           }
 
           if (!replaceable)
             continue;
 
           // Compute the new type based on the joined version.
           auto newKnowledge =
               ValueKnowledge::join(currentKnowledge, inferredKnowledge);
           if (!newKnowledge)
             continue;
-          result.setType(newKnowledge.getType());
+          setShapes(result, newKnowledge.getType());
         }
       }
     }
   }
+
+  // Actually update types with updated shape knowledge.
+  for (auto it : shapesStorage) {
+    auto result = it.second;
+    if (result.hasRank()) {
+      Type t = it.first.getType().cast<ShapedType>().clone(result.getDims());
+      it.first.setType(t);
+    }
+  }
 }
 
 /// Pass that performs shape propagation across TOSA operations. This includes
 /// migrating to within the regions of if/while operations.
 struct TosaInferShapes : public TosaInferShapesBase<TosaInferShapes> {
 public:
   void runOnFunction() override {
     FuncOp func = getOperation();

mlir/lib/Interfaces/InferTypeOpInterface.cpp

 //===- InferTypeOpInterface.cpp - Infer Type Interfaces ---------*- 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 contains the definitions of the infer op interfaces defined in
 // `InferTypeOpInterface.td`.
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Interfaces/InferTypeOpInterface.h"
 #include "mlir/IR/BuiltinTypes.h"
+#include "mlir/IR/Matchers.h"
+#include "llvm/Support/FormatVariadic.h"
 
 using namespace mlir;
 
 namespace mlir {
 #include "mlir/Interfaces/InferTypeOpInterface.cpp.inc"
 } // namespace mlir
 
+bool ShapeAdaptor::hasRank() const {
+  if (val.isNull())
+    return false;
+  if (auto t = val.dyn_cast<Type>())
+    return t.cast<ShapedType>().hasRank();
+  if (val.is<Attribute>())
+    return true;
+  return val.get<ShapedTypeComponents *>()->hasRank();
+}
+
+Type ShapeAdaptor::getElementType() const {
+  if (val.isNull())
+    return nullptr;
+  if (auto t = val.dyn_cast<Type>())
+    return t.cast<ShapedType>().getElementType();
+  if (val.is<Attribute>())
+    return nullptr;
+  return val.get<ShapedTypeComponents *>()->getElementType();
+}
+
+void ShapeAdaptor::getDims(SmallVectorImpl<int64_t> &res) const {
+  assert(hasRank());
+  if (auto t = val.dyn_cast<Type>()) {
+    ArrayRef<int64_t> vals = t.cast<ShapedType>().getShape();
+    res.assign(vals.begin(), vals.end());
+  } else if (auto attr = val.dyn_cast<Attribute>()) {
+    auto dattr = attr.cast<DenseIntElementsAttr>();
+    res.clear();
+    res.reserve(dattr.size());
+    for (auto it : dattr.getIntValues())
+      res.push_back(it.getSExtValue());
+  } else {
+    auto vals = val.get<ShapedTypeComponents *>()->getDims();
+    res.assign(vals.begin(), vals.end());
+  }
+}
+
+void ShapeAdaptor::getDims(ShapedTypeComponents &res) const {
+  assert(hasRank());
+  res.ranked = true;
+  getDims(res.dims);
+}
+
+int64_t ShapeAdaptor::getDimSize(int index) const {
+  assert(hasRank());
+  if (auto t = val.dyn_cast<Type>())
+    return t.cast<ShapedType>().getDimSize(index);
+  if (auto attr = val.dyn_cast<Attribute>())
+    return attr.cast<DenseIntElementsAttr>()
+        .getValue<APInt>({static_cast<uint64_t>(index)})
+        .getSExtValue();
+  auto *stc = val.get<ShapedTypeComponents *>();
+  return stc->getDims()[index];
+}
+
+int64_t ShapeAdaptor::getRank() const {
+  assert(hasRank());
+  if (auto t = val.dyn_cast<Type>())
+    return t.cast<ShapedType>().getRank();
+  if (auto attr = val.dyn_cast<Attribute>())
+    return attr.cast<DenseIntElementsAttr>().size();
+  return val.get<ShapedTypeComponents *>()->getDims().size();
+}
+
+bool ShapeAdaptor::hasStaticShape() const {
+  if (!hasRank())
+    return false;
+
+  if (auto t = val.dyn_cast<Type>())
+    return t.cast<ShapedType>().hasStaticShape();
+  if (auto attr = val.dyn_cast<Attribute>()) {
+    auto dattr = attr.cast<DenseIntElementsAttr>();
+    for (auto index : dattr.getIntValues())
+      if (ShapedType::isDynamic(index.getSExtValue()))
+        return false;
+    return true;
+  }
+  auto *stc = val.get<ShapedTypeComponents *>();
+  for (int64_t dim : stc->getDims())
+    if (ShapedType::isDynamic(dim))
+      return false;
+  return true;
+}
+
+int64_t ShapeAdaptor::getNumElements() const {
+  assert(hasStaticShape() && "cannot get element count of dynamic shaped type");
+
+  if (auto t = val.dyn_cast<Type>())
+    return t.cast<ShapedType>().getNumElements();
+
+  if (auto attr = val.dyn_cast<Attribute>()) {
+    auto dattr = attr.cast<DenseIntElementsAttr>();
+    int64_t num = 1;
+    for (auto index : dattr.getIntValues()) {
+      num *= index.getZExtValue();
+      assert(num >= 0 && "integer overflow in element count computation");
+    }
+    return num;
+  }
+
+  auto *stc = val.get<ShapedTypeComponents *>();
+  int64_t num = 1;
+  for (int64_t dim : stc->getDims()) {
+    num *= dim;
+    assert(num >= 0 && "integer overflow in element count computation");
+  }
+  return num;
+}
+
+void ShapeAdaptor::dump() const {
+  if (!hasRank()) {
+    llvm::errs() << "<<unranked>>\n";
+    return;
+  }
+
+  SmallVector<int64_t> dims;
+  getDims(dims);
+  auto mapped = llvm::map_range(dims, [](int64_t dim) -> std::string {
+    if (ShapedType::isDynamic(dim))
+      return "?";
+    return llvm::formatv("{0}", dim).str();
+  });
+  llvm::errs() << "rank = " << getRank() << " dims = [";
+  llvm::interleave(mapped, llvm::errs(), "x");
+  llvm::errs() << "]\n";
+}
+
+ShapeAdaptor ValueShapeRange::getValueAsShape(int index) {
+  Value val = operator[](index);
+  if (valueToShape)
+    if (ShapeAdaptor ret = valueToShape(val))
+      return ret;
+
+  DenseIntElementsAttr attr;
+  if (!matchPattern(val, m_Constant(&attr)))
+    return nullptr;
+  if (attr.getType().getRank() != 1)
+    return nullptr;
+  return attr;
+}
+
+ShapeAdaptor ValueShapeRange::getShape(Value val) const {
+  if (operandShape)
+    if (ShapeAdaptor ret = operandShape(val))
+      return ret;
+  return val.getType();
+}
+
+ShapeAdaptor ValueShapeRange::getShape(int index) const {
+  if (index < 0 || static_cast<size_t>(index) >= size())
+    return nullptr;
+  return getShape(operator[](index));
+}
+
 LogicalResult mlir::detail::inferReturnTensorTypes(
     function_ref<LogicalResult(
         MLIRContext *, Optional<Location> location, ValueShapeRange operands,
         DictionaryAttr attributes, RegionRange regions,
         SmallVectorImpl<ShapedTypeComponents> &retComponents)>
         componentTypeFn,
     MLIRContext *context, Optional<Location> location, ValueRange operands,
     DictionaryAttr attributes, RegionRange regions,

mlir/unittests/Interfaces/CMakeLists.txt

 add_mlir_unittest(MLIRInterfacesTests
   DataLayoutInterfacesTest.cpp
+  InferTypeOpInterfaceTest.cpp
 )
 
 target_link_libraries(MLIRInterfacesTests
   PRIVATE
   MLIRDataLayoutInterfaces
   MLIRDLTI
+  MLIRInferTypeOpInterface
   MLIRParser
+  MLIRStandard
 )

mlir/unittests/Interfaces/InferTypeOpInterfaceTest.cpp

+//===- InferTypeOpInterfaceTest.cpp - Unit Test for type interface --------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Interfaces/InferTypeOpInterface.h"
+#include "mlir/Dialect/StandardOps/IR/Ops.h"
+#include "mlir/IR/Builders.h"
+#include "mlir/IR/BuiltinOps.h"
+#include "mlir/IR/Dialect.h"
+#include "mlir/IR/DialectImplementation.h"
+#include "mlir/IR/ImplicitLocOpBuilder.h"
+#include "mlir/IR/OpDefinition.h"
+#include "mlir/IR/OpImplementation.h"
+#include "mlir/Parser.h"
+
+#include <gtest/gtest.h>
+
+using namespace mlir;
+
+class ValueShapeRangeTest : public testing::Test {
+protected:
+  void SetUp() override {
+    const char *ir = R"MLIR(
+      func @map(%arg : tensor<1xi64>) {
+        %0 = constant dense<[10]> : tensor<1xi64>
+        %1 = addi %arg, %0 : tensor<1xi64>
+        return
+      }
+    )MLIR";
+
+    registry.insert<StandardOpsDialect>();
+    ctx.appendDialectRegistry(registry);
+    module = parseSourceString(ir, &ctx);
+    mapFn = cast<FuncOp>(module->front());
+  }
+
+  // Create ValueShapeRange on the addi operation.
+  ValueShapeRange addiRange() {
+    auto &fnBody = mapFn.body();
+    return std::next(fnBody.front().begin())->getOperands();
+  }
+
+  DialectRegistry registry;
+  MLIRContext ctx;
+  OwningModuleRef module;
+  FuncOp mapFn;
+};
+
+TEST_F(ValueShapeRangeTest, ShapesFromValues) {
+  ValueShapeRange range = addiRange();
+
+  EXPECT_FALSE(range.getValueAsShape(0));
+  ASSERT_TRUE(range.getValueAsShape(1));
+  EXPECT_TRUE(range.getValueAsShape(1).hasRank());
+  EXPECT_EQ(range.getValueAsShape(1).getRank(), 1);
+  EXPECT_EQ(range.getValueAsShape(1).getDimSize(0), 10);
+  EXPECT_EQ(range.getShape(1).getRank(), 1);
+  EXPECT_EQ(range.getShape(1).getDimSize(0), 1);
+}
+
+TEST_F(ValueShapeRangeTest, MapValuesToShapes) {
+  ValueShapeRange range = addiRange();
+  ShapedTypeComponents fixed(SmallVector<int64_t>{30});
+  auto mapping = [&](Value val) -> ShapeAdaptor {
+    if (val == mapFn.getArgument(0))
+      return &fixed;
+    return nullptr;
+  };
+  range.setValueToShapeMapping(mapping);
+
+  ASSERT_TRUE(range.getValueAsShape(0));
+  EXPECT_TRUE(range.getValueAsShape(0).hasRank());
+  EXPECT_EQ(range.getValueAsShape(0).getRank(), 1);
+  EXPECT_EQ(range.getValueAsShape(0).getDimSize(0), 30);
+  ASSERT_TRUE(range.getValueAsShape(1));
+  EXPECT_TRUE(range.getValueAsShape(1).hasRank());
+  EXPECT_EQ(range.getValueAsShape(1).getRank(), 1);
+  EXPECT_EQ(range.getValueAsShape(1).getDimSize(0), 10);
+}
+
+TEST_F(ValueShapeRangeTest, SettingShapes) {
+  ShapedTypeComponents shape(SmallVector<int64_t>{10, 20});
+  ValueShapeRange range = addiRange();
+  auto mapping = [&](Value val) -> ShapeAdaptor {
+    if (val == mapFn.getArgument(0))
+      return &shape;
+    return nullptr;
+  };
+  range.setOperandShapeMapping(mapping);
+
+  ASSERT_TRUE(range.getShape(0));
+  EXPECT_EQ(range.getShape(0).getRank(), 2);
+  EXPECT_EQ(range.getShape(0).getDimSize(0), 10);
+  EXPECT_EQ(range.getShape(0).getDimSize(1), 20);
+  ASSERT_TRUE(range.getShape(1));
+  EXPECT_EQ(range.getShape(1).getRank(), 1);
+  EXPECT_EQ(range.getShape(1).getDimSize(0), 1);
+  EXPECT_FALSE(range.getShape(2));
+}