[mlir] Make ValueShapeRange a new class

[jpienaar]

Summary

Retaining old interface and should be constructable as previous, change would have been NFC except it this doesn't implicitly work with OpAdaptor constructor in C++14. Beyond that no new information/accessors in class - follow up to enable passing in shapes compatible with values during refinement without mutating the IR.

[rriddle]

Can you provide the "why"? I'm not sure what the motivation is here.

mlir/include/mlir/Interfaces/InferTypeOpInterface.h

81–101

Why not just inherit from ValueRange?

[jpienaar]

In D106772#2903289, @rriddle wrote:

Can you provide the "why"? I'm not sure what the motivation is here.

Sure, expanded description. This is basically towards avoiding to mutate IR during refinement. And will also be used to allow querying constant values without materializing them.

mlir/include/mlir/Interfaces/InferTypeOpInterface.h

81–101

ValueRange is final and I need to add members and accessors. I could also make it not be final, but assumed there was a reason & it is widely used so didn't want to impede optimizations.

[rsuderman]

Looks good to me assuming @rriddle is satisfied.

[rriddle]

mlir/include/mlir/Interfaces/InferTypeOpInterface.h

81–101

You won't be able to use RangeBaseT here. It is CRTP (on ValueRange) and there are several static_casts that rely on the derived type.

mlir/include/mlir/Interfaces/InferTypeOpInterface.h

private:
   ShapeStorageT dims;
   bool ranked;
   Type elementType;
   Attribute attr;
 };
 
 /// Range of values and shapes (corresponding effectively to Shapes dialect's
 /// ValueShape type concept).
-using ValueShapeRange = ValueRange;
+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))) {}
+  ValueShapeRange(const std::initializer_list<Value> &values)
+      : ValueShapeRange(ValueRange(values)) {}
+
+  /// 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.
+  ValueRange getValues() const { return ValueRange(begin(), end()); };
+};
 
 namespace detail {
-// Helper function to infer return tensor returns types given element and shape
-// inference function.
+// 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.
 LogicalResult inferReturnTensorTypes(
     function_ref<LogicalResult(
         MLIRContext *, Optional<Location> location, ValueShapeRange operands,
         DictionaryAttr attributes, RegionRange regions,
         SmallVectorImpl<ShapedTypeComponents> &retComponents)>

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

static LogicalResult ReduceInferReturnTypes(
   outputShape[axisVal] = 1;
   inferredReturnShapes.push_back(ShapedTypeComponents(outputShape));
   return success();
 }
 
 #define REDUCE_SHAPE_INFER(OP)                                                 \
   LogicalResult OP::inferReturnTypeComponents(                                 \
       MLIRContext *context, ::llvm::Optional<Location> location,               \
-      ValueRange operands, DictionaryAttr attributes, RegionRange regions,     \
+      ValueShapeRange operands, DictionaryAttr attributes,                     \
+      RegionRange regions,                                                     \
       SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) {           \
     return ReduceInferReturnTypes(operands[0],                                 \
                                   attributes.get("axis").cast<IntegerAttr>(),  \
                                   inferredReturnShapes);                       \
   }
 
 REDUCE_SHAPE_INFER(tosa::ReduceAllOp)
 REDUCE_SHAPE_INFER(tosa::ReduceAnyOp)
if (resolveBroadcastShape(operands, outShape).failed()) {
     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,     \
+      ValueShapeRange 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)
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);
+  Conv2DOp::Adaptor adaptor(operands.getValues());
 
   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>()) {
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);
+  Conv2DOp::Adaptor adaptor(operands.getValues());
 
   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;
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);
+  DepthwiseConv2DOp::Adaptor adaptor(operands.getValues());
 
   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;
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);
+  TransposeConv2DOp::Adaptor adaptor(operands.getValues());
   llvm::SmallVector<int64_t> outputShape;
   getI64Values(attributes.get("out_shape").cast<ArrayAttr>(), outputShape);
 
   int32_t inputWidth = ShapedType::kDynamicSize;
   int32_t inputHeight = ShapedType::kDynamicSize;
   int32_t weightWidth = ShapedType::kDynamicSize;
   int32_t weightHeight = ShapedType::kDynamicSize;
 

mlir/test/lib/Dialect/Test/TestDialect.cpp

LogicalResult OpWithInferTypeInterfaceOp::inferReturnTypes(
   return success();
 }
 
 LogicalResult OpWithShapedTypeInferTypeInterfaceOp::inferReturnTypeComponents(
     MLIRContext *context, Optional<Location> location, ValueShapeRange operands,
     DictionaryAttr attributes, RegionRange regions,
     SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) {
   // Create return type consisting of the last element of the first operand.
-  auto operandType = *operands.getTypes().begin();
+  auto operandType = operands.front().getType();
   auto sval = operandType.dyn_cast<ShapedType>();
   if (!sval) {
     return emitOptionalError(location, "only shaped type operands allowed");
   }
   int64_t dim =
       sval.hasRank() ? sval.getShape().front() : ShapedType::kDynamicSize;
   auto type = IntegerType::get(context, 17);
   inferredReturnShapes.push_back(ShapedTypeComponents({dim}, type));