[ODS] Use Adaptor Trait for Shaped Type Inference

[amandatang]

Summary

Author inferReturnTypeComponents methods with the Op Adaptor by using the InferShapedTypeOpAdaptor.

[amandatang]

clang format

[jpienaar]

mlir/include/mlir/Interfaces/InferTypeOpInterface.td

229

Why not just rely on overloading?

[amandatang]

Use function overloading

[amandatang]

rebase

[jpienaar]

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

410

You are querying inputShape's rank without checking if it was null here. cast is preferred in that case.

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

1456

MLIR style is to elide trivial braces.

[amandatang]

Address comments

[amandatang]

Clang format

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

 //===----------------------------------------------------------------------===//
 // TOSA Spec Section 2.2
 // Operator Class: Tensor Data Engine Operators.
 //===----------------------------------------------------------------------===//
 
 //===----------------------------------------------------------------------===//
 // Operator: argmax
 //===----------------------------------------------------------------------===//
-def Tosa_ArgMaxOp : Tosa_Op<"argmax", [
-    DeclareOpInterfaceMethods<InferShapedTypeOpInterface,
-                              ["inferReturnTypeComponents"]>,
-    Pure]> {
+def Tosa_ArgMaxOp : Tosa_Op<"argmax", [InferShapedTypeOpAdaptor, Pure]> {
   let summary = "Perform argmax on the input.";
 
   let description = [{
     This returns the index with the largest value across the given axis of the
     input tensor.
   }];
 
   let arguments = (ins

 // Accumulator types.
 //===----------------------------------------------------------------------===//
 
 def Tosa_AccType : AnyTypeOf<[I<32>, SI<32>, F16, F32]>;
 
 //===----------------------------------------------------------------------===//
 // Operator: avg_pool2d
 //===----------------------------------------------------------------------===//
-def Tosa_AvgPool2dOp : Tosa_Op<"avg_pool2d", [
-    DeclareOpInterfaceMethods<InferShapedTypeOpInterface,
-                              ["inferReturnTypeComponents"]>,
-    Pure]> {
+def Tosa_AvgPool2dOp : Tosa_Op<"avg_pool2d", [InferShapedTypeOpAdaptor, Pure]> {
   let summary = "Performs max pooling on the input.";
 
   let description = [{
     This performs an average pooling over the given input tensor. A sliding
     window of size given by <kernel size> is passed over the input tensor, with
     the mean value being placed in the output tensor.
   }];
 
def Tosa_AvgPool2dOp : Tosa_Op<"avg_pool2d", [InferShapedTypeOpAdaptor, Pure]> {
 
   let builders = [Tosa_AvgPool2dOpQuantInfoBuilder];
   let hasVerifier = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // Operator: conv2d
 //===----------------------------------------------------------------------===//
-def Tosa_Conv2DOp : Tosa_Op<"conv2d", [
-    DeclareOpInterfaceMethods<InferShapedTypeOpInterface,
-                              ["inferReturnTypeComponents"]>,
-    Pure]> {
+def Tosa_Conv2DOp : Tosa_Op<"conv2d", [InferShapedTypeOpAdaptor, Pure]> {
   let summary = "2D Convolution Operator";
 
   let description = [{
     Performs a 2D convolution over the given tensor input, using the weight
     tensor.
   }];
 
   let arguments = (ins
def Tosa_Conv2DOp : Tosa_Op<"conv2d", [InferShapedTypeOpAdaptor, Pure]> {
 
   let builders = [Tosa_ConvOpQuantInfoBuilder];
   let hasVerifier = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // Operator: conv3d
 //===----------------------------------------------------------------------===//
-def Tosa_Conv3DOp : Tosa_Op<"conv3d", [
-    DeclareOpInterfaceMethods<InferShapedTypeOpInterface,
-                              ["inferReturnTypeComponents"]>,
-    Pure]> {
+def Tosa_Conv3DOp : Tosa_Op<"conv3d", [InferShapedTypeOpAdaptor, Pure]> {
   let summary = "3D Convolution operator";
 
   let description = [{
     Performs a 3D convolution over the given input tensor.
   }];
 
   let arguments = (ins
     Tosa_Tensor5D:$input,
def Tosa_Conv3DOp : Tosa_Op<"conv3d", [InferShapedTypeOpAdaptor, Pure]> {
 
   let builders = [Tosa_ConvOpQuantInfoBuilder];
   let hasVerifier = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // Operator: depthwise_conv2d
 //===----------------------------------------------------------------------===//
-def Tosa_DepthwiseConv2DOp : Tosa_Op<"depthwise_conv2d", [
-    DeclareOpInterfaceMethods<InferShapedTypeOpInterface,
-                              ["inferReturnTypeComponents"]>,
-    Pure]> {
+def Tosa_DepthwiseConv2DOp : Tosa_Op<"depthwise_conv2d",
+    [InferShapedTypeOpAdaptor, Pure]> {
   let summary = "Depthwise 2D Convolution operator";
 
   let description = [{
     Performs 2D convolutions separately over each channel of the given tensor
     input, using the weight tensor.
   }];
 
   let arguments = (ins
def Tosa_DepthwiseConv2DOp : Tosa_Op<"depthwise_conv2d",
 
   let builders = [Tosa_ConvOpQuantInfoBuilder];
   let hasVerifier = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // Operator: fft2d
 //===----------------------------------------------------------------------===//
-def Tosa_FFT2dOp : Tosa_Op<"fft2d", [
-    DeclareOpInterfaceMethods<InferShapedTypeOpInterface,
-                              ["inferReturnTypeComponents"]>,
-    Pure]> {
+def Tosa_FFT2dOp : Tosa_Op<"fft2d", [InferShapedTypeOpAdaptor, Pure]> {
   let summary = "Performs FFT2D operation on the input.";
 
   let description = [{
     Performs a batched complex 2D Fast Fourier Transform over the input. The
     complex input values are constructed from the corresponding values in the
     input_real and input_imag tensors. The resulting values in the output are
     split into the output_real and output_imag tensors. No normalization is
     applied on either the forward or inverse versions of the operation.
let results = (outs
     Tosa_Tensor3D:$output_imag
   );
 }
 
 //===----------------------------------------------------------------------===//
 // Operator: fully_connected
 //===----------------------------------------------------------------------===//
 def Tosa_FullyConnectedOp : Tosa_Op<"fully_connected", [
-    DeclareOpInterfaceMethods<InferShapedTypeOpInterface,
-                              ["inferReturnTypeComponents"]>,
-    Pure]> {
+    InferShapedTypeOpAdaptor, Pure]> {
   let summary = "Fully Connected operator";
 
   let description = [{
     Performs a fully connected network.
   }];
 
   let arguments = (ins
     Tosa_Tensor2D:$input,
     2DTensorOf<[Tosa_Weight]>:$weight,
     Tosa_Tensor1D:$bias,
     OptionalAttr<Tosa_ConvOpQuantizationAttr>:$quantization_info
   );
 
   let results = (outs
     Tosa_Tensor2D:$output
   );
 
   let builders = [Tosa_FCOpQuantInfoBuilder];
   let hasVerifier = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // Operator: matmul
 //===----------------------------------------------------------------------===//
-def Tosa_MatMulOp : Tosa_Op<"matmul", [
-    DeclareOpInterfaceMethods<InferShapedTypeOpInterface,
-                              ["inferReturnTypeComponents"]>,
-    Pure]> {
+def Tosa_MatMulOp : Tosa_Op<"matmul", [InferShapedTypeOpAdaptor, Pure]> {
   let summary = "Matrix multiplication with bias";
 
   let description = [{
     Performs a two dimensional matrix multiplication. This allows both inputs to
     be activations, rather than reserving weights as an attribute in the
     FULLY_CONNECTED operator.
   }];
 
   let arguments = (ins
     Tosa_Tensor3D:$a,
     Tosa_Tensor3D:$b,
     OptionalAttr<Tosa_MatMulOpQuantizationAttr>:$quantization_info
   );
 
   let results = (outs
     Tosa_Tensor3D:$c
   );
 
   let builders = [Tosa_MatMulOpQuantInfoBuilder];
 }
 
 //===----------------------------------------------------------------------===//
 // Operator: max_pool2d
 //===----------------------------------------------------------------------===//
-def Tosa_MaxPool2dOp : Tosa_Op<"max_pool2d", [
-    DeclareOpInterfaceMethods<InferShapedTypeOpInterface,
-                              ["inferReturnTypeComponents"]>,
-    Pure]> {
+def Tosa_MaxPool2dOp : Tosa_Op<"max_pool2d", [InferShapedTypeOpAdaptor, Pure]> {
   let summary = "Performs max pooling on the input.";
 
   let description = [{
     This performs a max pooling over the given input tensor. A sliding window of
     size given by <kernel size> is passed over the input tensor, with the
     maximum value being placed in the
     output tensor.
   }];
def Tosa_MaxPool2dOp : Tosa_Op<"max_pool2d", [InferShapedTypeOpAdaptor, Pure]> {
   );
 
   let hasCanonicalizer = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // Operator: rfft2d
 //===----------------------------------------------------------------------===//
-def Tosa_RFFT2dOp : Tosa_Op<"rfft2d", [
-    DeclareOpInterfaceMethods<InferShapedTypeOpInterface,
-                              ["inferReturnTypeComponents"]>,
-    Pure]> {
+def Tosa_RFFT2dOp : Tosa_Op<"rfft2d", [InferShapedTypeOpAdaptor, Pure]> {
   let summary = "Performs RFFT2D operation on the input.";
 
   let description = [{
     Performs a batched 2D real-valued Fast Fourier Transform over the input where
     the input tensor consists of real values producing complex valued output. The
     complex output values will be split into the output_real and output_imag
     tensor arguments. RFFT2D takes advantage of Hermitian symmetry to only
     calculate the first half of the final output axis. Imaginary values with
     locations (0,0), (0,W/2), (H/2,0) and (H/2,W/2) are zero.
   }];
 
   let arguments = (ins
     Tosa_Tensor3D:$input
   );
 
   let results = (outs
     Tosa_Tensor3D:$output_real,
     Tosa_Tensor3D:$output_imag
   );
 }
 
 //===----------------------------------------------------------------------===//
 // Operator: transpose_conv2d
 //===----------------------------------------------------------------------===//
-def Tosa_TransposeConv2DOp : Tosa_Op<"transpose_conv2d", [
-    DeclareOpInterfaceMethods<InferShapedTypeOpInterface,
-                              ["inferReturnTypeComponents"]>,
-    Pure]> {
+def Tosa_TransposeConv2DOp : Tosa_Op<"transpose_conv2d", 
+    [InferShapedTypeOpAdaptor, Pure]> {
   let summary = "Transpose 2D Convolution operator.";
 
   let description = [{
     Performs a 2D transposed convolution over the given tensor input, using the
     weights tensor.
   }];
 
   let arguments = (ins
def Tosa_SubOp : Tosa_ElemWiseBinaryOp<"sub"> {
   );
 
   let hasFolder = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // Operator: table
 //===----------------------------------------------------------------------===//
-def Tosa_TableOp : Tosa_Op<"table", [
-    DeclareOpInterfaceMethods<InferShapedTypeOpInterface,
-                              ["inferReturnTypeComponents"]>,
-    Pure]> {
+def Tosa_TableOp : Tosa_Op<"table", [InferShapedTypeOpAdaptor, Pure]> {
   let summary = "Table lookup op";
 
   let description = [{
     Interpolated table lookup operation. Input values are scaled to create a
     fixed-point 9.7 value.    The high 9 bits are used to index into the table.
     The fractional bits are used to interpolate based on the looked up value and
     the index+1 value in the table. The TABLE operator then returns a 16.7
     interpolated value. Note that there must be 513 values to handle the full

 // TOSA Spec Section 2.8
 // Operator Class: Reduction Ops.
 //===----------------------------------------------------------------------===//
 
 //===----------------------------------------------------------------------===//
 // Operator: reduce_all
 //===----------------------------------------------------------------------===//
 def Tosa_ReduceAllOp : Tosa_Op<"reduce_all", [
-    InferTensorType, Pure]> {
+    InferTensorTypeAdaptor, Pure]> {
   let summary = "Reduce All operator";
 
   let description = [{
     Reduce a tensor along the given axis with a logical AND operation
   }];
 
   let arguments = (ins
     Tosa_Tensor:$input,
let extraClassDeclaration = [{
     static bool isCompatibleReturnTypes(TypeRange l, TypeRange r);
   }];
 }
 
 //===----------------------------------------------------------------------===//
 // Operator: reduce_any
 //===----------------------------------------------------------------------===//
 def Tosa_ReduceAnyOp : Tosa_Op<"reduce_any", [
-    InferTensorType, Pure]> {
+    InferTensorTypeAdaptor, Pure]> {
   let summary = "Reduce Any operator";
 
   let description = [{
     Reduce a tensor along the given axis with a logical OR operation
   }];
 
   let arguments = (ins
     Tosa_Tensor:$input,
let extraClassDeclaration = [{
     static bool isCompatibleReturnTypes(TypeRange l, TypeRange r);
   }];
 }
 
 //===----------------------------------------------------------------------===//
 // Operator: reduce_max
 //===----------------------------------------------------------------------===//
 def Tosa_ReduceMaxOp : Tosa_Op<"reduce_max", [
-    InferTensorType, Pure]> {
+    InferTensorTypeAdaptor, Pure]> {
   let summary = "Reduce Max operator";
 
   let description = [{
     Reduce a tensor along the given axis with a maximum operation
   }];
 
   let arguments = (ins
     Tosa_Tensor:$input,
let extraClassDeclaration = [{
     static bool isCompatibleReturnTypes(TypeRange l, TypeRange r);
   }];
 }
 
 //===----------------------------------------------------------------------===//
 // Operator: reduce_min
 //===----------------------------------------------------------------------===//
 def Tosa_ReduceMinOp : Tosa_Op<"reduce_min", [
-    InferTensorType, Pure]> {
+    InferTensorTypeAdaptor, Pure]> {
   let summary = "Reduce Min operator";
 
   let description = [{
     Reduce a tensor along the given axis with a minimum operation
   }];
 
   let arguments = (ins
     Tosa_Tensor:$input,
let extraClassDeclaration = [{
     static bool isCompatibleReturnTypes(TypeRange l, TypeRange r);
   }];
 }
 
 //===----------------------------------------------------------------------===//
 // Operator: reduce_prod
 //===----------------------------------------------------------------------===//
 def Tosa_ReduceProdOp : Tosa_Op<"reduce_prod", [
-    InferTensorType, Pure]> {
+    InferTensorTypeAdaptor, Pure]> {
   let summary = "Reduce Prod operator";
 
   let description = [{
     Reduce a tensor along the given axis by computing the product of the axis.
   }];
 
   let arguments = (ins
     Tosa_Tensor:$input,
let extraClassDeclaration = [{
     static bool isCompatibleReturnTypes(TypeRange l, TypeRange r);
   }];
 }
 
 //===----------------------------------------------------------------------===//
 // Operator: reduce_sum
 //===----------------------------------------------------------------------===//
 def Tosa_ReduceSumOp : Tosa_Op<"reduce_sum", [
-    InferTensorType, Pure]> {
+    InferTensorTypeAdaptor, Pure]> {
   let summary = "Reduce Sum operator";
 
   let description = [{
     Reduce a tensor along the given axis by computing the sum of the axis.
   }];
 
   let arguments = (ins
     Tosa_Tensor:$input,

 // TOSA Spec Section 2.9
 // Operator Class: Data Layout / Memory Reinterpretation.
 //===----------------------------------------------------------------------===//
 
 //===----------------------------------------------------------------------===//
 // Operator: concat
 //===----------------------------------------------------------------------===//
 def Tosa_ConcatOp : Tosa_Op<"concat", [
-    InferTensorType, Pure]> {
+    InferTensorTypeAdaptor, Pure]> {
   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
let extraClassDeclaration = [{
     /// Method used by InferTypeOpInterface.
     static bool isCompatibleReturnTypes(TypeRange l, TypeRange r);
   }];
 }
 
 //===----------------------------------------------------------------------===//
 // Operator: pad
 //===----------------------------------------------------------------------===//
-def Tosa_PadOp : Tosa_Op<"pad", [
-    DeclareOpInterfaceMethods<InferShapedTypeOpInterface,
-                              ["inferReturnTypeComponents"]>,
-    Pure]> {
+def Tosa_PadOp : Tosa_Op<"pad", [InferShapedTypeOpAdaptor, Pure]> {
   let summary = "Pads a tensor with value specified.";
 
   let description = [{
     The `tosa.pad` operation pads a tensor along borders of each dimension with
     `pad_const` (defaults to zero), given a padding configuration `padding`
     specifying low and high values along the dimensions.
 
     Example:
def Tosa_PadOp : Tosa_Op<"pad", [InferShapedTypeOpAdaptor, Pure]> {
   let hasCanonicalizer = 1;
   let hasFolder = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // Operator: reshape
 //===----------------------------------------------------------------------===//
 def Tosa_ReshapeOp: Tosa_Op<"reshape", [
-  InferTensorType, Pure]> {
+  InferTensorTypeAdaptor, Pure]> {
   let summary = "Reshape operator";
 
   let description = [{
     Returns a tensor with the same type/values as the input, with a new shape
     specified by the shape argument. Reshape may operate on tensors of any rank.
     No data conversion happens during a reshape operation.
   }];
 
def Tosa_ReverseOp: Tosa_Op<"reverse", [
   );
 
   let hasFolder = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // Operator: slice
 //===----------------------------------------------------------------------===//
-def Tosa_SliceOp: Tosa_Op<"slice", [
-      DeclareOpInterfaceMethods<InferShapedTypeOpInterface,
-                              ["inferReturnTypeComponents"]>, Pure]> {
+def Tosa_SliceOp: Tosa_Op<"slice", [InferShapedTypeOpAdaptor, Pure]> {
   let summary = "Slice operator";
 
   let description = [{
     Extracts a slice of the input1 on the given axis, beginning at the
     start coordinates, and extending for size elements in each direction.  No
     data conversion happens during a slice operation.
   }];
 
def Tosa_SliceOp: Tosa_Op<"slice", [InferShapedTypeOpAdaptor, Pure]> {
 
   let hasCanonicalizer = 1;
   let hasFolder = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // Operator: tile
 //===----------------------------------------------------------------------===//
-def Tosa_TileOp: Tosa_Op<"tile", [
-      DeclareOpInterfaceMethods<InferShapedTypeOpInterface,
-                              ["inferReturnTypeComponents"]>,
-      Pure]> {
+def Tosa_TileOp: Tosa_Op<"tile", [InferShapedTypeOpAdaptor, Pure]> {
   let summary = "Tile operator";
 
   let description = [{
     Replicates input 0 multiplies times along each dimension.
   }];
 
   let arguments = (ins
     Tosa_Tensor:$input1,
     DenseI64ArrayAttr:$multiples);
 
   let results = (outs
     Tosa_Tensor:$output
   );
 
   let hasFolder = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // Operator: transpose
 //===----------------------------------------------------------------------===//
-def Tosa_TransposeOp : Tosa_Op<"transpose", [
-      DeclareOpInterfaceMethods<InferShapedTypeOpInterface,
-                              ["inferReturnTypeComponents"]>,
-      Pure]> {
+def Tosa_TransposeOp : Tosa_Op<"transpose", [InferShapedTypeOpAdaptor, Pure]> {
   let summary = "Transpose operator";
 
   let description = [{
     Permutes the dimensions based on perm.
   }];
 
   let arguments = (ins
     Tosa_Tensor:$input1,

 //===----------------------------------------------------------------------===//
 // TOSA Spec Section 2.10
 // Operator Class: Scatter/gather Operations.
 //===----------------------------------------------------------------------===//
 
 //===----------------------------------------------------------------------===//
 // Operator: gather
 //===----------------------------------------------------------------------===//
-def Tosa_GatherOp : Tosa_Op<"gather", [
-      DeclareOpInterfaceMethods<InferShapedTypeOpInterface,
-                              ["inferReturnTypeComponents"]>,
-      Pure]> {
+def Tosa_GatherOp : Tosa_Op<"gather", [InferShapedTypeOpAdaptor, Pure]> {
   let summary = "Gather operation,";
 
   let description = [{
     Generate a tensor for which each element in the output is a slice of the
     values tensor based on the value of indices.
   }];
 
   let arguments = (ins
     Tosa_Tensor3D:$values,
     2DTensorOf<[Tosa_Int32]>:$indices
   );
 
   let results = (outs
     Tosa_Tensor3D:$output
   );
 }
 
 //===----------------------------------------------------------------------===//
 // Operator: scatter
 //===----------------------------------------------------------------------===//
-def Tosa_ScatterOp : Tosa_Op<"scatter", [
-      DeclareOpInterfaceMethods<InferShapedTypeOpInterface,
-                              ["inferReturnTypeComponents"]>,
-      Pure]> {
+def Tosa_ScatterOp : Tosa_Op<"scatter", [InferShapedTypeOpAdaptor, Pure]> {
   let summary = "Scatter operation,";
 
   let description = [{
     The values_out tensor is set to the values_in tensor with data modified as follows:
     data from the input tensor is inserted at the positions specified by the indices tensor.
   }];
 
   let arguments = (ins

 //===----------------------------------------------------------------------===//
 // TOSA Spec Section 2.11
 // Operator Class: Image Frontend Functions.
 //===----------------------------------------------------------------------===//
 
 //===----------------------------------------------------------------------===//
 // Operator: resize
 //===----------------------------------------------------------------------===//
-def Tosa_ResizeOp : Tosa_Op<"resize", [
-      DeclareOpInterfaceMethods<InferShapedTypeOpInterface,
-                              ["inferReturnTypeComponents"]>,
-      Pure]> {
+def Tosa_ResizeOp : Tosa_Op<"resize", [InferShapedTypeOpAdaptor, Pure]> {
 
   let summary = "Resize operation, supports various resize/upsample modes";
 
   let description = [{
     Resizes a tensor. Resize is only allowed in the H and W dimensions. In
     expected use, The height dimension is scaled by factor (scale_y_n/scale_y_d).
     And the width dimension is scaled by factor (scale_x_n/scale_x_d). Thus the
     output dimensions can be derived from the input dimensions by inverting the

 //===----------------------------------------------------------------------===//
 
 //===----------------------------------------------------------------------===//
 // Operator: cond_if
 //===----------------------------------------------------------------------===//
 //===----------------------------------------------------------------------===//
 // Further described in docs/Rationale/RationaleTOSADialect.md .
 //===----------------------------------------------------------------------===//
-def Tosa_IfOp : Tosa_Op<"cond_if", [
-      DeclareOpInterfaceMethods<InferShapedTypeOpInterface,
-                                ["inferReturnTypeComponents"]>,
+def Tosa_IfOp : Tosa_Op<"cond_if",
+       [InferShapedTypeOpAdaptor,
        SingleBlockImplicitTerminator<"YieldOp">,
        RecursiveMemoryEffects]> {
   let summary = "Conditional if operator";
 
   let description = [{
     Evaluates a Boolean condition and then takes one of two distinct execution
     paths. This implements the semantic If-then-else structure.
   }];

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

mlir/include/mlir/Interfaces/InferTypeOpInterface.h

 
 namespace mlir {
 namespace OpTrait {
 
 template <typename ConcreteType>
 class InferTypeOpInterfaceAdaptor
     : public TraitBase<ConcreteType, InferTypeOpInterfaceAdaptor> {};
 
+template <typename ConcreteType>
+class InferShapedTypeOpAdaptor
+    : public TraitBase<ConcreteType, InferShapedTypeOpAdaptor> {};
+
 /// Tensor type inference trait that constructs a tensor from the inferred
 /// shape and elemental types.
 /// Requires: Op implements InferShapedTypeOpInterface and InferTypeOpInterface.
 ///   Less strict is possible (e.g., implements inferReturnTypeComponents and
 ///   these always populates all element types and shapes or fails, but this\
 ///   trait is currently only used where the interfaces are, so keep it
 ///   restricted for now).
 template <typename ConcreteType>
 class InferTensorType : public TraitBase<ConcreteType, InferTensorType> {};
 
 } // namespace OpTrait
 } // namespace mlir
 
 #endif // MLIR_INTERFACES_INFERTYPEOPINTERFACE_H_

mlir/include/mlir/Interfaces/InferTypeOpInterface.td

NativeOpTrait<
           $cppClass::Adaptor adaptor(operands, attributes, properties, regions);
           return $cppClass::inferReturnTypesAdaptor(context,
             location, adaptor, inferredReturnTypes);
         }
       }]
     >
   ]>;
 
+// Convenient trait to define a wrapper to inferReturnTypeComponents that passes
+// in the Op Adaptor directly
+class InferShapedTypeOpAdaptorBase<list<string> overridenMethods = []> : TraitList<
+  [
+    // Op implements infer type op interface.
+    DeclareOpInterfaceMethods<InferShapedTypeOpInterface, overridenMethods>,
+    NativeOpTrait<
+      /*name=*/"InferShapedTypeOpAdaptor",
+      /*traits=*/[],
+      /*extraOpDeclaration=*/[{
+        static ::mlir::LogicalResult
+        inferReturnTypeComponents(::mlir::MLIRContext *context,
+                                std::optional<::mlir::Location> location,
+                                Adaptor adaptor,
+                                ::llvm::SmallVectorImpl<::mlir::ShapedTypeComponents> &inferredReturnShapes);
+      }],
+      /*extraOpDefinition=*/[{
+        ::mlir::LogicalResult
+        $cppClass::inferReturnTypeComponents(::mlir::MLIRContext *context,
+                          std::optional<::mlir::Location> location,
+                          ::mlir::ValueShapeRange operands, ::mlir::DictionaryAttr attributes,
+                          ::mlir::OpaqueProperties properties, ::mlir::RegionRange regions,
+                          ::llvm::SmallVectorImpl<::mlir::ShapedTypeComponents> &inferredReturnShapes) {
+          $cppClass::Adaptor adaptor(operands, attributes, properties, regions);
+          return $cppClass::inferReturnTypeComponents(context,
+            location, adaptor, inferredReturnShapes);
+        }
+      }]
+    >
+  ]>;
+
+def InferShapedTypeOpAdaptor : InferShapedTypeOpAdaptorBase<[
+  "inferReturnTypeComponents"]>;
+def InferShapedTypeOpAdaptorWithReify : InferShapedTypeOpAdaptorBase<[
+  "inferReturnTypeComponents", "reifyReturnTypeShapes"]>;
+
 // Convenience class grouping together type and shaped type op interfaces for
 // ops that have tensor return types.
 class InferTensorTypeBase<list<string> overridenMethods = []> : TraitList<
   [
     // Op implements infer type op interface.
     DeclareOpInterfaceMethods<InferTypeOpInterface>,
     // The op will have methods implementing the ShapedType type inference
     // interface.
     DeclareOpInterfaceMethods<InferShapedTypeOpInterface, overridenMethods>,
     // The op produces tensors and will use the ShapedType type infer interface
     // along with knowledge that it is producing Tensors to infer the type.
     NativeOpTrait<
       /*name=*/"InferTensorType",
       /*traits=*/[],
       /*extraOpDeclaration=*/[{}],
       /*extraOpDefinition=*/[{
         LogicalResult
-        $cppClass::inferReturnTypes(MLIRContext *context,
-                          std::optional<Location> location,
-                          ValueRange operands, DictionaryAttr attributes,
-                          OpaqueProperties properties, RegionRange regions,
-                          SmallVectorImpl<Type> &inferredReturnTypes) {
+        $cppClass::inferReturnTypes(::mlir::MLIRContext *context,
+                          std::optional<::mlir::Location> location,
+                          ::mlir::ValueRange operands, ::mlir::DictionaryAttr attributes,
+                          ::mlir::OpaqueProperties properties, ::mlir::RegionRange regions,
+                          ::llvm::SmallVectorImpl<::mlir::Type> &inferredReturnTypes) {
           SmallVector<ShapedTypeComponents, 2> retComponents;
           if (failed($cppClass::inferReturnTypeComponents(context, location,
                                     operands, attributes, properties, regions,
                                     retComponents)))
             return failure();
           return ::mlir::detail::inferReturnTensorTypes(retComponents,
                                     inferredReturnTypes);
         }
       }]
     >
   ]>;
 
 def InferTensorType : InferTensorTypeBase<["inferReturnTypeComponents"]>;
 def InferTensorTypeWithReify: InferTensorTypeBase<[
     "inferReturnTypeComponents", "reifyReturnTypeShapes"]>;
 
+// Convenience class grouping together type and shaped type op interfaces for
+// ops that have tensor return types.
+class InferTensorTypeAdaptorBase<list<string> overridenMethods = []> : TraitList<
+  [
+    // Op implements infer type op interface.
+    DeclareOpInterfaceMethods<InferTypeOpInterface>,
+    // The op will have methods implementing the ShapedType type inference
+    // interface.
+    InferShapedTypeOpAdaptorBase<overridenMethods>,
+    // The op produces tensors and will use the ShapedType type infer interface
+    // along with knowledge that it is producing Tensors to infer the type.
+    NativeOpTrait<
+      /*name=*/"InferTensorType",
+      /*traits=*/[],
+      /*extraOpDeclaration=*/[{}],
+      /*extraOpDefinition=*/[{
+        LogicalResult
+        $cppClass::inferReturnTypes(::mlir::MLIRContext *context,
+                          std::optional<::mlir::Location> location,
+                          ::mlir::ValueRange operands, ::mlir::DictionaryAttr attributes,
+                          ::mlir::OpaqueProperties properties, ::mlir::RegionRange regions,
+                          ::llvm::SmallVectorImpl<::mlir::Type> &inferredReturnTypes) {
+          SmallVector<ShapedTypeComponents, 2> retComponents;
+          if (failed($cppClass::inferReturnTypeComponents(context, location,
+                                    operands, attributes, properties, regions,
+                                    retComponents)))
+            return failure();
+          return ::mlir::detail::inferReturnTensorTypes(retComponents,
+                                    inferredReturnTypes);
+        }
+      }]
+    >
+  ]>;
+
+def InferTensorTypeAdaptor : InferTensorTypeAdaptorBase<["inferReturnTypeComponents"]>;
+def InferTensorTypeAdaptorWithReify: InferTensorTypeAdaptorBase<[
+    "inferReturnTypeComponents", "reifyReturnTypeShapes"]>;
+
 def ReifyRankedShapedTypeOpInterface :
     OpInterface<"ReifyRankedShapedTypeOpInterface"> {
   let description = [{
     Interface to compute the shape of the result of an operation when
     the result is a ranked shape type, i.e. `RankedTensorType` or
     `MemRefType`.
   }];
   let cppNamespace = "::mlir";

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

 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/Dialect/Tosa/Utils/QuantUtils.h"
 #include "mlir/Dialect/Tosa/Utils/ShapeUtils.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/DialectImplementation.h"
 #include "mlir/IR/Matchers.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/IR/TypeUtilities.h"
+#include "mlir/Interfaces/InferTypeOpInterface.h"
 #include "mlir/Transforms/InliningUtils.h"
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/TypeSwitch.h"
 
 using namespace mlir;
 using namespace mlir::tosa;
 
 #include "mlir/Dialect/Tosa/IR/TosaOpsDialect.cpp.inc"
for (int i = 0, e = operands.size(); i != e; ++i) {
     }
   }
 
   return success();
 }
 
 LogicalResult tosa::ArgMaxOp::inferReturnTypeComponents(
     MLIRContext *context, ::std::optional<Location> location,
-    ValueShapeRange operands, DictionaryAttr attributes,
-    OpaqueProperties properties, RegionRange regions,
+    ArgMaxOp::Adaptor adaptor,
     SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) {
-  ShapeAdaptor inputShape = operands.getShape(0);
-  auto *prop = properties.as<Properties *>();
-  IntegerAttr axis = prop->axis;
+  ShapeAdaptor inputShape = dyn_cast<ShapedType>(adaptor.getInput().getType());
+  IntegerAttr axis = adaptor.getProperties().axis;
   int32_t axisVal = axis.getValue().getSExtValue();
 
   if (!inputShape.hasRank()) {
     inferredReturnShapes.push_back(ShapedTypeComponents());
     return success();
   }
 
   SmallVector<int64_t> outShape;
   outShape.reserve(inputShape.getRank() - 1);
   for (int i = 0, s = inputShape.getRank(); i < s; i++) {
     if (i == axisVal)
       continue;
     outShape.push_back(inputShape.getDimSize(i));
   }
 
   inferredReturnShapes.push_back(ShapedTypeComponents(outShape));
   return success();
 }
 
 LogicalResult tosa::RFFT2dOp::inferReturnTypeComponents(
     MLIRContext *context, ::std::optional<Location> location,
-    ValueShapeRange operands, DictionaryAttr attributes,
-    OpaqueProperties properties, RegionRange regions,
+    RFFT2dOp::Adaptor adaptor,
     SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) {
-  ShapeAdaptor inputShape = operands.getShape(0);
+  ShapeAdaptor inputShape = dyn_cast<ShapedType>(adaptor.getInput().getType());
 
   if (!inputShape.hasRank())
     return failure();
 
   llvm::SmallVector<int64_t> outputShape;
   outputShape.resize(3, ShapedType::kDynamic);
   outputShape[0] = inputShape.getDimSize(0);
   outputShape[1] = inputShape.getDimSize(1);
   int64_t inWidth = inputShape.getDimSize(2);
 
   // Note that we can support this calculation symbolically
   // in the future e.g. [x, y, z] -> [x, y, z / 2 - 1]
   if (inWidth != ShapedType::kDynamic)
     outputShape[2] = inWidth / 2 + 1;
 
   inferredReturnShapes.push_back(ShapedTypeComponents(outputShape));
   inferredReturnShapes.push_back(ShapedTypeComponents(outputShape));
 
   return success();
 }
 
 LogicalResult tosa::FFT2dOp::inferReturnTypeComponents(
     MLIRContext *context, ::std::optional<Location> location,
-    ValueShapeRange operands, DictionaryAttr attributes,
-    OpaqueProperties properties, RegionRange regions,
+    FFT2dOp::Adaptor adaptor,
     SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) {
-  inferredReturnShapes.push_back(ShapedTypeComponents(operands.getShape(0)));
-  inferredReturnShapes.push_back(ShapedTypeComponents(operands.getShape(1)));
+  inferredReturnShapes.push_back(ShapedTypeComponents(
+      dyn_cast<ShapedType>(adaptor.getInputReal().getType())));
+  inferredReturnShapes.push_back(ShapedTypeComponents(
+      dyn_cast<ShapedType>(adaptor.getInputImag().getType())));
   return success();
 }
 
 LogicalResult tosa::ConcatOp::inferReturnTypeComponents(
     MLIRContext *context, ::std::optional<Location> location,
-    ValueShapeRange operands, DictionaryAttr attributes,
-    OpaqueProperties properties, RegionRange regions,
+    ConcatOp::Adaptor adaptor,
     SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) {
   // Infer all dimension sizes by reducing based on inputs.
-  auto *prop = properties.as<Properties *>();
-  int32_t axis = prop->axis.getValue().getSExtValue();
+  const Properties &prop = adaptor.getProperties();
+  int32_t axis = prop.axis.getValue().getSExtValue();
   llvm::SmallVector<int64_t> outputShape;
   bool hasRankedInput = false;
-  for (auto operand : operands) {
-    ShapeAdaptor operandShape = operands.getShape(operand);
+  for (auto operand : adaptor.getOperands()) {
+    ShapeAdaptor operandShape = dyn_cast<ShapedType>(operand.getType());
     if (!operandShape.hasRank())
       continue;
 
     // Copy the Operand's rank.
     if (!hasRankedInput)
       outputShape.resize(operandShape.getRank(), ShapedType::kDynamic);
 
     // Copy shapes until the dim is non-dynamic.
     for (int i = 0, s = operandShape.getRank(); i < s; i++) {
       if (i == axis || operandShape.isDynamicDim(i))
         continue;
       if (outputShape[i] == ShapedType::kDynamic)
         outputShape[i] = operandShape.getDimSize(i);
       if (outputShape[i] != operandShape.getDimSize(i))
         return emitOptionalError(location,
                                  "Cannot concat tensors with different sizes"
                                  " on the non-axis dimension ",
                                  i);
     }
 
     hasRankedInput = true;
   }
   Type inputType =
-      llvm::cast<TensorType>(operands.getType()[0]).getElementType();
+      llvm::cast<TensorType>(adaptor.getInput1().getType()[0]).getElementType();
   if (!hasRankedInput) {
     inferredReturnShapes.push_back(ShapedTypeComponents(inputType));
     return success();
   }
 
   // Determine the dimension size along the concatenation axis.
   int64_t concatDimSize = 0;
-  for (auto operand : operands) {
-    ShapeAdaptor operandShape = operands.getShape(operand);
+  for (auto operand : adaptor.getOperands()) {
+    ShapeAdaptor operandShape = dyn_cast<ShapedType>(operand.getType());
 
     // We need to know the length of the concatenation axis of all inputs to
     // determine the dimension size of the output shape.
     if (!operandShape.hasRank() || operandShape.isDynamicDim(axis)) {
       concatDimSize = ShapedType::kDynamic;
       break;
     }
 

 bool tosa::EqualOp::isCompatibleReturnTypes(TypeRange l, TypeRange r) {
   if (l.size() != r.size() || l.size() != 1)
     return false;
   return succeeded(verifyCompatibleShape(l[0], r[0]));
 }
 
 LogicalResult tosa::FullyConnectedOp::inferReturnTypeComponents(
     MLIRContext *context, ::std::optional<Location> location,
-    ValueShapeRange operands, DictionaryAttr attributes,
-    OpaqueProperties properties, RegionRange regions,
+    FullyConnectedOp::Adaptor adaptor,
     SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) {
-  ShapeAdaptor inputShape = operands.getShape(0);
-  ShapeAdaptor weightShape = operands.getShape(1);
-  ShapeAdaptor biasShape = operands.getShape(2);
+  ShapeAdaptor inputShape = dyn_cast<ShapedType>(adaptor.getInput().getType());
+  ShapeAdaptor weightShape =
+      dyn_cast<ShapedType>(adaptor.getWeight().getType());
+  ShapeAdaptor biasShape = dyn_cast<ShapedType>(adaptor.getBias().getType());
 
   // All shapes are dynamic.
   SmallVector<int64_t> outShape;
   outShape.resize(2, ShapedType::kDynamic);
 
   if (inputShape.hasRank()) {
     outShape[0] = inputShape.getDimSize(0);
   }
LogicalResult tosa::FullyConnectedOp::inferReturnTypeComponents(
   inferredReturnShapes.push_back(ShapedTypeComponents(outShape));
   return success();
 }
 
 LogicalResult FullyConnectedOp::verify() { return verifyConvOp(*this); }
 
 LogicalResult tosa::MatMulOp::inferReturnTypeComponents(
     MLIRContext *context, ::std::optional<Location> location,
-    ValueShapeRange operands, DictionaryAttr attributes,
-    OpaqueProperties properties, RegionRange regions,
+    MatMulOp::Adaptor adaptor,
     SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) {
-  ShapeAdaptor lhsShape = operands.getShape(0);
-  ShapeAdaptor rhsShape = operands.getShape(1);
+  ShapeAdaptor lhsShape = dyn_cast<ShapedType>(adaptor.getA().getType());
+  ShapeAdaptor rhsShape = dyn_cast<ShapedType>(adaptor.getB().getType());
 
   // All shapes are dynamic.
   SmallVector<int64_t> outShape;
   outShape.resize(3, ShapedType::kDynamic);
 
   if (lhsShape.hasRank()) {
     outShape[0] = lhsShape.getDimSize(0);
     outShape[1] = lhsShape.getDimSize(1);
   }
 
   if (rhsShape.hasRank()) {
     outShape[0] = outShape[0] == ShapedType::kDynamic ? rhsShape.getDimSize(0)
                                                       : outShape[0];
     outShape[2] = rhsShape.getDimSize(2);
   }
 
   inferredReturnShapes.push_back(ShapedTypeComponents(outShape));
   return success();
 }
 
 LogicalResult tosa::PadOp::inferReturnTypeComponents(
     MLIRContext *context, ::std::optional<Location> location,
-    ValueShapeRange operands, DictionaryAttr attributes,
-    OpaqueProperties properties, RegionRange regions,
+    PadOp::Adaptor adaptor,
     SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) {
-  ShapeAdaptor inputShape = operands.getShape(0);
-  ShapeAdaptor paddingShape = operands.getShape(1);
+  ShapeAdaptor inputShape = dyn_cast<ShapedType>(adaptor.getInput1().getType());
+  ShapeAdaptor paddingShape =
+      dyn_cast<ShapedType>(adaptor.getPadding().getType());
   SmallVector<int64_t> outputShape;
 
   // If both inputs have unknown shape, we cannot determine the shape of the
   // output.
   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 (!inputShape.hasRank()) {
     if (paddingShape.isDynamicDim(0)) {
       inferredReturnShapes.push_back(ShapedTypeComponents());
       return success();
     }
 
     outputShape.resize(paddingShape.getDimSize(0), ShapedType::kDynamic);
     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))) {
+  if (!matchPattern(adaptor.getPadding(), m_Constant(&paddings))) {
     outputShape.resize(inputShape.getRank(), ShapedType::kDynamic);
     inferredReturnShapes.push_back(ShapedTypeComponents(outputShape));
     return success();
   }
 
   SmallVector<int64_t> paddingValues;
   for (auto val : paddings) {
     paddingValues.push_back(val.getSExtValue());

 static SmallVector<int64_t> convertToMlirShape(ArrayRef<int64_t> shape) {
   return to_vector(llvm::map_range(shape, [](int64_t dim) {
     return dim == -1 ? ShapedType::kDynamic : dim;
   }));
 }
 
 LogicalResult tosa::SliceOp::inferReturnTypeComponents(
     MLIRContext *context, ::std::optional<Location> location,
-    ValueShapeRange operands, DictionaryAttr attributes,
-    OpaqueProperties properties, RegionRange regions,
+    SliceOp::Adaptor adaptor,
     SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) {
-  inferredReturnShapes.push_back(ShapedTypeComponents(
-      convertToMlirShape(SliceOpAdaptor(operands, attributes,
-                                        *properties.as<Properties *>(), regions)
-                             .getSize())));
+  inferredReturnShapes.push_back(
+      ShapedTypeComponents(convertToMlirShape(adaptor.getSize())));
   return success();
 }
 
 LogicalResult tosa::TableOp::inferReturnTypeComponents(
     MLIRContext *context, ::std::optional<Location> location,
-    ValueShapeRange operands, DictionaryAttr attributes,
-    OpaqueProperties properties, RegionRange regions,
+    TableOp::Adaptor adaptor,
     SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) {
-  ShapeAdaptor inputShape = operands.getShape(0);
+  ShapeAdaptor inputShape = dyn_cast<ShapedType>(adaptor.getInput().getType());
 
   if (!inputShape.hasRank()) {
     inferredReturnShapes.push_back(ShapedTypeComponents());
     return success();
   }
 
   inferredReturnShapes.resize(1);
   inputShape.getDims(inferredReturnShapes[0]);
   return success();
 }
 
 LogicalResult tosa::TileOp::inferReturnTypeComponents(
     MLIRContext *context, ::std::optional<Location> location,
-    ValueShapeRange operands, DictionaryAttr attributes,
-    OpaqueProperties properties, RegionRange regions,
+    TileOp::Adaptor adaptor,
     SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) {
-  TileOpAdaptor adaptor(operands, attributes, *properties.as<Properties *>(),
-                        regions);
   ArrayRef<int64_t> multiples = adaptor.getMultiples();
-  ShapeAdaptor inputShape = operands.getShape(0);
+  ShapeAdaptor inputShape = dyn_cast<ShapedType>(adaptor.getInput1().getType());
   SmallVector<int64_t> outputShape;
   if (!inputShape.hasRank()) {
     outputShape.resize(multiples.size(), ShapedType::kDynamic);
     inferredReturnShapes.push_back(ShapedTypeComponents(outputShape));
     return success();
   }
 
   // Any non dynamic dimension can be multiplied to a known size.

 bool tosa::ReshapeOp::isCompatibleReturnTypes(TypeRange l, TypeRange r) {
   if (l.size() != r.size() || l.size() != 1)
     return false;
   return getElementTypeOrSelf(l[0]) == getElementTypeOrSelf(r[0]);
 }
 
 LogicalResult tosa::ReshapeOp::inferReturnTypeComponents(
     MLIRContext *context, ::std::optional<Location> location,
-    ValueShapeRange operands, DictionaryAttr attributes,
-    OpaqueProperties properties, RegionRange regions,
+    ReshapeOp::Adaptor adaptor,
     SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) {
-  ReshapeOpAdaptor adaptor(operands, attributes, *properties.as<Properties *>(),
-                           regions);
-  ShapeAdaptor inputShape = operands.getShape(0);
-  Type inputType = getElementTypeOrSelf(operands.getType()[0]);
+  ShapeAdaptor inputShape = dyn_cast<ShapedType>(adaptor.getInput1().getType());
+  Type inputType = getElementTypeOrSelf(adaptor.getInput1().getType());
   llvm::SmallVector<int64_t> newShapeValue =
       convertToMlirShape(adaptor.getNewShape());
 
   // We cannot infer from the total number of elements so we must take the
   // shape attribute as exact.
   if (!inputShape.hasRank() || !inputShape.hasStaticShape()) {
     inferredReturnShapes.push_back(
         ShapedTypeComponents(newShapeValue, inputType));
perms = llvm::to_vector(
       llvm::map_range(permsAttr.getValues<APInt>(),
                       [](const APInt &val) { return val.getSExtValue(); }));
 
   return success();
 }
 
 LogicalResult tosa::TransposeOp::inferReturnTypeComponents(
     MLIRContext *context, ::std::optional<Location> location,
-    ValueShapeRange operands, DictionaryAttr attributes,
-    OpaqueProperties properties, RegionRange regions,
+    TransposeOp::Adaptor adaptor,
     SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) {
-  ShapeAdaptor inputShape = operands.getShape(0);
-  ShapeAdaptor permsShape = operands.getShape(1);
+  ShapeAdaptor inputShape = dyn_cast<ShapedType>(adaptor.getInput1().getType());
+  ShapeAdaptor permsShape = dyn_cast<ShapedType>(adaptor.getPerms().getType());
 
   // If input rank and permutation length is unknown, the output rank is
   // unknown.
   if (!inputShape.hasRank() || !permsShape.hasRank() ||
       permsShape.isDynamicDim(0)) {
     inferredReturnShapes.push_back(ShapedTypeComponents());
     return success();
   }
if (allTheSame) {
     outputShape.resize(inputShape.getRank(), inputShape.getDimSize(0));
     inferredReturnShapes.push_back(ShapedTypeComponents(outputShape));
     return success();
   }
 
   outputShape.resize(inputShape.getRank(), ShapedType::kDynamic);
   // If the permuations are a constant we can directly determine the output
   // shape.
-  if (ShapeAdaptor permShape = operands.getValueAsShape(1)) {
+  DenseIntElementsAttr attr;
+  if (matchPattern(adaptor.getPerms(), m_Constant(&attr)) &&
+      attr.getType().getRank() == 1) {
+    ShapeAdaptor permShape = attr;
     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, ::std::optional<Location> location,
-    ValueShapeRange operands, DictionaryAttr attributes,
-    OpaqueProperties properties, RegionRange regions,
+    GatherOp::Adaptor adaptor,
     SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) {
   llvm::SmallVector<int64_t> outputShape;
   outputShape.resize(3, ShapedType::kDynamic);
 
-  ShapeAdaptor valuesShape = operands.getShape(0);
+  ShapeAdaptor valuesShape =
+      dyn_cast<ShapedType>(adaptor.getValues().getType());
   if (valuesShape.hasRank()) {
     outputShape[0] = valuesShape.getDimSize(0);
     outputShape[2] = valuesShape.getDimSize(2);
   }
 
-  ShapeAdaptor indicesShape = operands.getShape(1);
+  ShapeAdaptor indicesShape =
+      dyn_cast<ShapedType>(adaptor.getIndices().getType());
   if (indicesShape.hasRank()) {
     if (outputShape[0] == ShapedType::kDynamic)
       outputShape[0] = indicesShape.getDimSize(0);
     if (outputShape[1] == ShapedType::kDynamic)
       outputShape[1] = indicesShape.getDimSize(1);
   }
 
   inferredReturnShapes.push_back(ShapedTypeComponents(outputShape));
   return success();
 }
 
 LogicalResult tosa::ResizeOp::inferReturnTypeComponents(
     MLIRContext *context, ::std::optional<Location> location,
-    ValueShapeRange operands, DictionaryAttr attributes,
-    OpaqueProperties properties, RegionRange regions,
+    ResizeOp::Adaptor adaptor,
     SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) {
-  ResizeOpAdaptor adaptor(operands, attributes, *properties.as<Properties *>(),
-                          regions);
   llvm::SmallVector<int64_t, 4> outputShape;
   outputShape.resize(4, ShapedType::kDynamic);
 
-  ShapeAdaptor inputShape = operands.getShape(adaptor.getInput());
+  ShapeAdaptor inputShape = dyn_cast<ShapedType>(adaptor.getInput().getType());
   if (!inputShape.hasRank())
     return failure();
 
   outputShape[0] = inputShape.getDimSize(0);
   outputShape[3] = inputShape.getDimSize(3);
   int64_t inputHeight = inputShape.getDimSize(1);
   int64_t inputWidth = inputShape.getDimSize(2);
 
outputShape[2] =
       1;
 
   inferredReturnShapes.push_back(ShapedTypeComponents(outputShape));
   return success();
 }
 
 LogicalResult tosa::ScatterOp::inferReturnTypeComponents(
     MLIRContext *context, ::std::optional<Location> location,
-    ValueShapeRange operands, DictionaryAttr attributes,
-    OpaqueProperties properties, RegionRange regions,
+    ScatterOp::Adaptor adaptor,
     SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) {
   llvm::SmallVector<int64_t> outputShape;
   outputShape.resize(3, ShapedType::kDynamic);
 
-  ShapeAdaptor valuesInShape = operands.getShape(0);
+  ShapeAdaptor valuesInShape =
+      dyn_cast<ShapedType>(adaptor.getValuesIn().getType());
   if (valuesInShape.hasRank()) {
     outputShape[0] = valuesInShape.getDimSize(0);
     outputShape[1] = valuesInShape.getDimSize(1);
     outputShape[2] = valuesInShape.getDimSize(2);
   }
 
-  ShapeAdaptor indicesShape = operands.getShape(1);
+  ShapeAdaptor indicesShape =
+      dyn_cast<ShapedType>(adaptor.getIndices().getType());
   if (indicesShape.hasRank()) {
     if (outputShape[0] == ShapedType::kDynamic)
       outputShape[0] = indicesShape.getDimSize(0);
   }
 
-  ShapeAdaptor inputShape = operands.getShape(2);
+  ShapeAdaptor inputShape = dyn_cast<ShapedType>(adaptor.getInput().getType());
   if (inputShape.hasRank()) {
     if (outputShape[0] == ShapedType::kDynamic)
       outputShape[0] = inputShape.getDimSize(0);
     if (outputShape[2] == ShapedType::kDynamic)
       outputShape[2] = inputShape.getDimSize(2);
   }
 
   inferredReturnShapes.push_back(ShapedTypeComponents(outputShape));
bool OP::isCompatibleReturnTypes(TypeRange l, TypeRange r) {                 \
     if (getElementTypeOrSelf(l[0]) != getElementTypeOrSelf(r[0]))              \
       return false;                                                            \
     return succeeded(verifyCompatibleShape(l[0], r[0]));                       \
   }
 
 #define REDUCE_SHAPE_INFER(OP)                                                 \
   LogicalResult OP::inferReturnTypeComponents(                                 \
       MLIRContext *context, ::std::optional<Location> location,                \
-      ValueShapeRange operands, DictionaryAttr attributes,                     \
-      OpaqueProperties properties, RegionRange regions,                        \
+      OP::Adaptor adaptor,                                                     \
       SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) {           \
     Type inputType =                                                           \
-        llvm::cast<TensorType>(operands.getType()[0]).getElementType();        \
-    return ReduceInferReturnTypes(operands.getShape(0), inputType,             \
-                                  properties.as<Properties *>()->axis,         \
+        llvm::cast<TensorType>(adaptor.getInput().getType()).getElementType(); \
+    ShapeAdaptor inputShape =                                                  \
+        dyn_cast<ShapedType>(adaptor.getInput().getType());                    \
+    const Properties &prop = adaptor.getProperties();                          \
+    return ReduceInferReturnTypes(inputShape, inputType, prop.axis,            \
                                   inferredReturnShapes);                       \
   }                                                                            \
   COMPATIBLE_RETURN_TYPES(OP)
 
 REDUCE_SHAPE_INFER(tosa::ReduceAllOp)
 REDUCE_SHAPE_INFER(tosa::ReduceAnyOp)
 REDUCE_SHAPE_INFER(tosa::ReduceMaxOp)
 REDUCE_SHAPE_INFER(tosa::ReduceMinOp)

 NARY_SHAPE_INFER(tosa::SelectOp)
 NARY_SHAPE_INFER(tosa::SubOp)
 NARY_SHAPE_INFER(tosa::TanhOp)
 NARY_SHAPE_INFER(tosa::ErfOp)
 NARY_SHAPE_INFER(tosa::SigmoidOp)
 #undef PRED_SHAPE_INFER
 
 static LogicalResult poolingInferReturnTypes(
-    const ValueShapeRange &operands, DictionaryAttr attributes,
-    ArrayRef<int64_t> kernel, ArrayRef<int64_t> stride, ArrayRef<int64_t> pad,
+    ShapeAdaptor inputShape, ArrayRef<int64_t> kernel, ArrayRef<int64_t> stride,
+    ArrayRef<int64_t> pad,
     SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) {
-  ShapeAdaptor inputShape = operands.getShape(0);
   llvm::SmallVector<int64_t> outputShape;
   outputShape.resize(4, ShapedType::kDynamic);
 
   // We only know the rank if the input type is unranked.
   if (!inputShape) {
     inferredReturnShapes.push_back(ShapedTypeComponents(outputShape));
     return success();
   }
static LogicalResult poolingInferReturnTypes(
   }
 
   inferredReturnShapes.push_back(ShapedTypeComponents(outputShape));
   return success();
 }
 
 LogicalResult Conv2DOp::inferReturnTypeComponents(
     MLIRContext *context, ::std::optional<Location> location,
-    ValueShapeRange operands, DictionaryAttr attributes,
-    OpaqueProperties properties, RegionRange regions,
+    Conv2DOp::Adaptor adaptor,
     SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) {
   llvm::SmallVector<int64_t> outputShape(4, ShapedType::kDynamic);
-  Conv2DOp::Adaptor adaptor(operands, attributes,
-                            *properties.as<Properties *>(), regions);
 
   int64_t inputWidth = ShapedType::kDynamic;
   int64_t inputHeight = ShapedType::kDynamic;
   int64_t weightWidth = ShapedType::kDynamic;
   int64_t weightHeight = ShapedType::kDynamic;
 
   // Input shape describes input width/height and batch.
 
-  ShapeAdaptor inputShape = operands.getShape(adaptor.getInput());
+  ShapeAdaptor inputShape = dyn_cast<ShapedType>(adaptor.getInput().getType());
   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.
-  ShapeAdaptor weightShape = operands.getShape(adaptor.getWeight());
+  ShapeAdaptor weightShape =
+      dyn_cast<ShapedType>(adaptor.getWeight().getType());
   if (weightShape.hasRank()) {
     outputShape[3] = weightShape.getDimSize(0);
     weightHeight = weightShape.getDimSize(1);
     weightWidth = weightShape.getDimSize(2);
   }
 
   // Bias shape can describe the output channels.
-  ShapeAdaptor biasShape = operands.getShape(adaptor.getBias());
+  ShapeAdaptor biasShape = dyn_cast<ShapedType>(adaptor.getBias().getType());
   if (biasShape.hasRank()) {
     outputShape[3] = ShapedType::isDynamic(outputShape[3])
                          ? biasShape.getDimSize(0)
                          : outputShape[3];
   }
 
   llvm::ArrayRef<int64_t> dilation = adaptor.getDilation();
   llvm::ArrayRef<int64_t> stride = adaptor.getStride();
LogicalResult Conv2DOp::inferReturnTypeComponents(
   inferredReturnShapes.push_back(ShapedTypeComponents(outputShape));
   return success();
 }
 
 LogicalResult Conv2DOp::verify() { return verifyConvOp(*this); }
 
 LogicalResult Conv3DOp::inferReturnTypeComponents(
     MLIRContext *context, ::std::optional<Location> location,
-    ValueShapeRange operands, DictionaryAttr attributes,
-    OpaqueProperties properties, RegionRange regions,
+    Conv3DOp::Adaptor adaptor,
     SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) {
   llvm::SmallVector<int64_t> outputShape(5, ShapedType::kDynamic);
-  Conv3DOp::Adaptor adaptor(operands, attributes,
-                            *properties.as<Properties *>(), regions);
 
   int64_t inputWidth = ShapedType::kDynamic;
   int64_t inputHeight = ShapedType::kDynamic;
   int64_t inputDepth = ShapedType::kDynamic;
 
   int64_t weightWidth = ShapedType::kDynamic;
   int64_t weightHeight = ShapedType::kDynamic;
   int64_t weightDepth = ShapedType::kDynamic;
 
   // Input shape describes input width/height and batch.
-  ShapeAdaptor inputShape = operands.getShape(adaptor.getInput());
+  ShapeAdaptor inputShape = dyn_cast<ShapedType>(adaptor.getInput().getType());
   if (inputShape.hasRank()) {
     outputShape[0] = inputShape.getDimSize(0);
     inputDepth = inputShape.getDimSize(1);
     inputHeight = inputShape.getDimSize(2);
     inputWidth = inputShape.getDimSize(3);
   }
 
   // Weight shapes describes the filter width/height and the output channels.
-  ShapeAdaptor weightShape = operands.getShape(adaptor.getWeight());
+  ShapeAdaptor weightShape =
+      dyn_cast<ShapedType>(adaptor.getWeight().getType());
   if (weightShape.hasRank()) {
     outputShape[4] = weightShape.getDimSize(0);
     weightDepth = weightShape.getDimSize(1);
     weightHeight = weightShape.getDimSize(2);
     weightWidth = weightShape.getDimSize(3);
   }
 
   // Bias shape can describe the output channels.
-  ShapeAdaptor biasShape = operands.getShape(adaptor.getBias());
+  ShapeAdaptor biasShape = dyn_cast<ShapedType>(adaptor.getBias().getType());
   if (biasShape.hasRank() && ShapedType::isDynamic(outputShape[4])) {
     outputShape[4] = biasShape.getDimSize(0);
   }
 
   llvm::ArrayRef<int64_t> dilation = adaptor.getDilation();
   llvm::ArrayRef<int64_t> stride = adaptor.getStride();
   llvm::ArrayRef<int64_t> pad = adaptor.getPad();
 
LogicalResult Conv3DOp::inferReturnTypeComponents(
   inferredReturnShapes.push_back(ShapedTypeComponents(outputShape));
   return success();
 }
 
 LogicalResult Conv3DOp::verify() { return verifyConvOp(*this); }
 
 LogicalResult AvgPool2dOp::inferReturnTypeComponents(
     MLIRContext *context, ::std::optional<Location> location,
-    ValueShapeRange operands, DictionaryAttr attributes,
-    OpaqueProperties properties, RegionRange regions,
+    AvgPool2dOp::Adaptor adaptor,
     SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) {
-  Properties &prop = *properties.as<Properties *>();
-  return poolingInferReturnTypes(operands, attributes, prop.kernel, prop.stride,
-                                 prop.pad, inferredReturnShapes);
+  ShapeAdaptor inputShape = dyn_cast<ShapedType>(adaptor.getInput().getType());
+  const Properties &prop = adaptor.getProperties();
+  return poolingInferReturnTypes(inputShape, prop.kernel, prop.stride, prop.pad,
+                                 inferredReturnShapes);
 }
 
 LogicalResult MaxPool2dOp::inferReturnTypeComponents(
     MLIRContext *context, ::std::optional<Location> location,
-    ValueShapeRange operands, DictionaryAttr attributes,
-    OpaqueProperties properties, RegionRange regions,
+    MaxPool2dOp::Adaptor adaptor,
     SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) {
-  Properties &prop = *properties.as<Properties *>();
-  return poolingInferReturnTypes(operands, attributes, prop.kernel, prop.stride,
-                                 prop.pad, inferredReturnShapes);
+  ShapeAdaptor inputShape = dyn_cast<ShapedType>(adaptor.getInput().getType());
+  const Properties &prop = adaptor.getProperties();
+  return poolingInferReturnTypes(inputShape, prop.kernel, prop.stride, prop.pad,
+                                 inferredReturnShapes);
 }
 
 LogicalResult DepthwiseConv2DOp::inferReturnTypeComponents(
     MLIRContext *context, ::std::optional<Location> location,
-    ValueShapeRange operands, DictionaryAttr attributes,
-    OpaqueProperties properties, RegionRange regions,
+    DepthwiseConv2DOp::Adaptor adaptor,
     SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) {
   llvm::SmallVector<int64_t> outputShape(4, ShapedType::kDynamic);
-  DepthwiseConv2DOp::Adaptor adaptor(operands, attributes,
-                                     *properties.as<Properties *>(), regions);
 
   int64_t inputWidth = ShapedType::kDynamic;
   int64_t inputHeight = ShapedType::kDynamic;
   int64_t inputChannels = ShapedType::kDynamic;
 
   int64_t weightWidth = ShapedType::kDynamic;
   int64_t weightHeight = ShapedType::kDynamic;
   int64_t depthChannels = ShapedType::kDynamic;
 
   // Input shape describes input width/height and batch.
-  ShapeAdaptor inputShape = operands.getShape(adaptor.getInput());
+  ShapeAdaptor inputShape = dyn_cast<ShapedType>(adaptor.getInput().getType());
   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.
-  ShapeAdaptor weightShape = operands.getShape(adaptor.getWeight());
+  ShapeAdaptor weightShape =
+      dyn_cast<ShapedType>(adaptor.getWeight().getType());
   if (weightShape.hasRank()) {
     weightHeight = weightShape.getDimSize(0);
     weightWidth = weightShape.getDimSize(1);
     inputChannels = ShapedType::isDynamic(inputChannels)
                         ? weightShape.getDimSize(2)
                         : inputChannels;
     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.
-  ShapeAdaptor biasShape = operands.getShape(adaptor.getBias());
+  ShapeAdaptor biasShape = dyn_cast<ShapedType>(adaptor.getBias().getType());
   if (biasShape.hasRank()) {
     outputShape[3] = ShapedType::isDynamic(outputShape[3])
                          ? biasShape.getDimSize(0)
                          : outputShape[3];
   }
 
   llvm::ArrayRef<int64_t> dilation = adaptor.getDilation();
   llvm::ArrayRef<int64_t> padding = adaptor.getPad();
LogicalResult DepthwiseConv2DOp::inferReturnTypeComponents(
   inferredReturnShapes.push_back(ShapedTypeComponents(outputShape));
   return success();
 }
 
 LogicalResult DepthwiseConv2DOp::verify() { return verifyConvOp(*this); }
 
 LogicalResult TransposeConv2DOp::inferReturnTypeComponents(
     MLIRContext *context, ::std::optional<Location> location,
-    ValueShapeRange operands, DictionaryAttr attributes,
-    OpaqueProperties properties, RegionRange regions,
+    TransposeConv2DOp::Adaptor adaptor,
     SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) {
-  TransposeConv2DOp::Adaptor adaptor(operands, attributes,
-                                     *properties.as<Properties *>(), regions);
   // outputShape is mutable.
   llvm::SmallVector<int64_t> outputShape =
       convertToMlirShape(adaptor.getOutShape());
 
   int64_t inputWidth = ShapedType::kDynamic;
   int64_t inputHeight = ShapedType::kDynamic;
   int64_t weightWidth = ShapedType::kDynamic;
   int64_t weightHeight = ShapedType::kDynamic;
 
   // Input shape describes input width/height and batch.
-  ShapeAdaptor inputShape = operands.getShape(adaptor.getInput());
+  ShapeAdaptor inputShape = dyn_cast<ShapedType>(adaptor.getInput().getType());
   if (inputShape.hasRank()) {
     outputShape[0] = ShapedType::isDynamic(outputShape[0])
                          ? inputShape.getDimSize(0)
                          : outputShape[0];
     inputHeight = inputShape.getDimSize(1);
     inputWidth = inputShape.getDimSize(2);
   }
 
   // Weight shapes describes the filter width/height and the output channels.
-  ShapeAdaptor weightShape = operands.getShape(adaptor.getFilter());
+  ShapeAdaptor weightShape =
+      dyn_cast<ShapedType>(adaptor.getFilter().getType());
   if (weightShape.hasRank()) {
     outputShape[3] = ShapedType::isDynamic(outputShape[3])
                          ? weightShape.getDimSize(0)
                          : outputShape[3];
     weightHeight = weightShape.getDimSize(1);
     weightWidth = weightShape.getDimSize(2);
   }
 
   // Bias shape can describe the output channels.
-  ShapeAdaptor biasShape = operands.getShape(adaptor.getInput());
+  ShapeAdaptor biasShape = dyn_cast<ShapedType>(adaptor.getInput().getType());
   if (biasShape.hasRank()) {
     outputShape[3] = ShapedType::isDynamic(outputShape[3])
                          ? biasShape.getDimSize(0)
                          : outputShape[3];
   }
 
   llvm::ArrayRef<int64_t> padding = adaptor.getOutPad();
   llvm::ArrayRef<int64_t> stride = adaptor.getStride();
LogicalResult TransposeConv2DOp::inferReturnTypeComponents(
   }
 
   inferredReturnShapes.push_back(ShapedTypeComponents(outputShape));
   return success();
 }
 
 LogicalResult IfOp::inferReturnTypeComponents(
     MLIRContext *context, ::std::optional<Location> location,
-    ValueShapeRange operands, DictionaryAttr attributes,
-    OpaqueProperties properties, RegionRange regions,
+    IfOp::Adaptor adaptor,
     SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) {
   llvm::SmallVector<tosa::YieldOp> yieldOps;
-  for (Region *region : regions) {
+  for (Region *region : adaptor.getRegions()) {
     for (auto &block : *region)
       if (auto returnOp = dyn_cast<tosa::YieldOp>(block.getTerminator()))
         yieldOps.push_back(returnOp);
   }
 
   if (yieldOps.empty())
     return failure();
 
for (const ValueKnowledge &result : resultKnowledge) {
     inferredReturnShapes.push_back(result.getShapedTypeComponents());
   }
 
   return success();
 }
 
 LogicalResult WhileOp::inferReturnTypeComponents(
     MLIRContext *context, ::std::optional<Location> location,
-    ValueShapeRange operands, DictionaryAttr attributes,
-    OpaqueProperties properties, RegionRange regions,
+    WhileOp::Adaptor adaptor,
     SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) {
   llvm::SmallVector<tosa::YieldOp> yieldOps;
-  for (auto &block : *regions[1])
+  for (auto &block : adaptor.getBody())
     if (auto returnOp = dyn_cast<tosa::YieldOp>(block.getTerminator()))
       yieldOps.push_back(returnOp);
 
   // TOSA's while must have a tosa.yield as its terminator. If not found this
   // tosa.while is invalid.
   if (yieldOps.empty())
     return failure();
 

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

 LogicalResult OpWithShapedTypeInferTypeInterfaceOp::reifyReturnTypeShapes(
     OpBuilder &builder, ValueRange operands,
     llvm::SmallVectorImpl<Value> &shapes) {
   shapes = SmallVector<Value, 1>{
       builder.createOrFold<tensor::DimOp>(getLoc(), operands.front(), 0)};
   return success();
 }
 
+LogicalResult
+OpWithShapedTypeInferTypeAdaptorInterfaceOp::inferReturnTypeComponents(
+    MLIRContext *context, std::optional<Location> location,
+    OpWithShapedTypeInferTypeAdaptorInterfaceOp::Adaptor adaptor,
+    SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) {
+  // Create return type consisting of the last element of the first operand.
+  auto operandType = adaptor.getOperand1().getType();
+  auto sval = dyn_cast<ShapedType>(operandType);
+  if (!sval) {
+    return emitOptionalError(location, "only shaped type operands allowed");
+  }
+  int64_t dim = sval.hasRank() ? sval.getShape().front() : ShapedType::kDynamic;
+  auto type = IntegerType::get(context, 17);
+
+  Attribute encoding;
+  if (auto rankedTy = dyn_cast<RankedTensorType>(sval))
+    encoding = rankedTy.getEncoding();
+  inferredReturnShapes.push_back(ShapedTypeComponents({dim}, type, encoding));
+  return success();
+}
+
+LogicalResult
+OpWithShapedTypeInferTypeAdaptorInterfaceOp::reifyReturnTypeShapes(
+    OpBuilder &builder, ValueRange operands,
+    llvm::SmallVectorImpl<Value> &shapes) {
+  shapes = SmallVector<Value, 1>{
+      builder.createOrFold<tensor::DimOp>(getLoc(), operands.front(), 0)};
+  return success();
+}
+
 LogicalResult OpWithResultShapeInterfaceOp::reifyReturnTypeShapes(
     OpBuilder &builder, ValueRange operands,
     llvm::SmallVectorImpl<Value> &shapes) {
   Location loc = getLoc();
   shapes.reserve(operands.size());
   for (Value operand : llvm::reverse(operands)) {
     auto rank = cast<RankedTensorType>(operand.getType()).getRank();
     auto currShape = llvm::to_vector<4>(

mlir/test/lib/Dialect/Test/TestOps.td

 }
 
 def OpWithShapedTypeInferTypeInterfaceOp : TEST_Op<"op_with_shaped_type_infer_type_if",
       [InferTensorTypeWithReify]> {
   let arguments = (ins AnyTensor, AnyTensor);
   let results = (outs AnyTensor);
 }
 
+def OpWithShapedTypeInferTypeAdaptorInterfaceOp : 
+      TEST_Op<"op_with_shaped_type_infer_type_adaptor_if",
+              [InferTensorTypeAdaptorWithReify]> {
+  let arguments = (ins AnyTensor:$operand1, AnyTensor:$operand2);
+  let results = (outs AnyTensor:$result);
+}
+
 def OpWithResultShapeInterfaceOp : TEST_Op<"op_with_result_shape_interface",
       [DeclareOpInterfaceMethods<InferShapedTypeOpInterface,
           ["reifyReturnTypeShapes"]>]> {
   let arguments = (ins AnyRankedTensor:$operand1, AnyRankedTensor:$operand2);
   let results = (outs AnyRankedTensor:$result1, AnyRankedTensor:$result2);
 }
 
 def OpWithResultShapePerDimInterfaceOp :