[mlir][shape] Add dim op

[jpienaar]

Summary

Convenience op that allows for simple expression of common crossing of
value/shape divide.

[yaochengji]

Could you add test in ops.mlir for this new op?

mlir/lib/Dialect/Shape/IR/Shape.cpp

1072

Do we need to consider more complicated situations here, where dim is not a constant-like op, but it could be folded into one?

1113

it should be dim >= st.getRank()

Further more, how do you plan to support negative dim index?

[jpienaar]

Fix bound calculation

mlir/lib/Dialect/Shape/IR/Shape.cpp

1072

Yes but not here. That would be done as part of analysis (probably sparse data flow analysis one, we are going to rebase the TOSA shape inference one to use that framework) where one would do (potentially) expensive propagations and partial evaluations to get dims along with an inference context (well when we combine the analysis we can take advantage of SCCP's analysis state too and feed back into it, which will be nice). Here it is just for the cheap/trivial case.

1113

Negative dim index convention is not used anywhere where dim is used as operand/runtime value. I'm guessing you are thinking of Python's convention with array slicing wrt usage, similar to split_at. It is used there as its a fixed attribute so one can't do index calculation, and so if dim was an attribute here, we'd do same but for operand keeping it consistent with rest of system.

(which does remind me we didn't add sub yet as we didn't need/used the arith one as part of lowered form only).

[jpienaar]

Forgot to add ops.mlir in last update

[yaochengji]

mlir/lib/Dialect/Shape/IR/Shape.cpp

1113

I guess we still need to check whether dim < 0 here. Since neither ShapeSizeType nor IndexType limit the value to be non-negative.

[jpienaar]

mlir/lib/Dialect/Shape/IR/Shape.cpp

1113

Good point, will update soon (and sorry, for some reason I didn't get a notification of your comment addressing now)

[jpienaar]

mlir/lib/Dialect/Shape/IR/Shape.cpp

1113

Tensor::DimOp also doesn't check this, that seems like bug.

[jpienaar]

Verify for non-negative dims

mlir/include/mlir/Dialect/Shape/IR/ShapeOps.td

def Shape_ToExtentTensorOp : Shape_Op<"to_extent_tensor", [
   let arguments = (ins Shape_ShapeOrExtentTensorType:$input);
   let results = (outs IndexTensor:$result);
 
   let assemblyFormat = "$input attr-dict `:` type($input) `->` type($result)";
 
   let hasFolder = 1;
 }
 
+def Shape_DimOp : Shape_Op<"dim",
+    [NoSideEffect, DeclareOpInterfaceMethods<InferTypeOpInterface>]> {
+  let summary = "Gets the specified extent from the shape of a shaped input";
+  let description = [{
+    Gets the extent indexed by `dim` from the shape of the `value` operand. If
+    the dim is error or out-of-bound then it returns an invalid size if the
+    return type carries error information else the behavior is undefined.
+
+    This is a convenience op that performs the equivalent of getting the extent
+    of a shape (e.g., `dim(x, i) == get_extent(shape_of(x), i)`).
+  }];
+  let arguments = (ins AnyShaped:$value,
+                       Shape_SizeOrIndexType:$dim);
+  let results = (outs Shape_SizeOrIndexType:$extent);
+  let assemblyFormat = "$value `,` $dim attr-dict `:` type($value) `,` type($dim) `->` "
+                       "type($extent)";
+
+  let builders = [
+    // Builder that allows passing a constant dimension as a simple integer.
+    OpBuilder<(ins "Value":$value, "int64_t":$dim)>
+  ];
+
+  let extraClassDeclaration = [{
+    /// Get the `dim` value as integer if it is constant.
+    Optional<int64_t> getConstantDim();
+
+    /// Returns when two result types are compatible for this op; method used by
+    /// InferTypeOpInterface
+    static bool isCompatibleReturnTypes(TypeRange l, TypeRange r);
+  }];
+
+  let hasFolder = 1;
+  let hasVerifier = 1;
+}
+
 def Shape_GetExtentOp : Shape_Op<"get_extent",
     [NoSideEffect, DeclareOpInterfaceMethods<InferTypeOpInterface>]> {
   let summary = "Gets the specified extent from a shape or extent tensor";
   let description = [{
     Gets the extent indexed by `dim` from the `shape` operand. If the shape is
     an error then it returns an invalid size.
   }];
   let arguments = (ins Shape_ShapeOrExtentTensorType:$shape,

mlir/lib/Conversion/ShapeToStandard/ShapeToStandard.cpp

auto reduceResult = lb.create<ForOp>(
         b.create<scf::YieldOp>(loc, broadcastable);
       });
 
   rewriter.replaceOp(op, reduceResult.getResults().front());
   return success();
 }
 
 namespace {
+class DimOpConverter : public OpConversionPattern<DimOp> {
+  using OpConversionPattern<DimOp>::OpConversionPattern;
+
+  LogicalResult
+  matchAndRewrite(DimOp op, OpAdaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override;
+};
+} // namespace
+
+LogicalResult
+DimOpConverter::matchAndRewrite(DimOp op, OpAdaptor adaptor,
+                                ConversionPatternRewriter &rewriter) const {
+  // Lower to dim(X, i) to get_extent(shape_of(X), i) and rely on further
+  // lowerings. This can be further optimized if needed to avoid intermediate
+  // steps.
+  auto shapeOf = rewriter.create<shape::ShapeOfOp>(op.getLoc(), op.getValue());
+  rewriter.replaceOpWithNewOp<shape::GetExtentOp>(op, op.getType(), shapeOf,
+                                                  op.getDim());
+  return success();
+}
+
+namespace {
 class GetExtentOpConverter : public OpConversionPattern<GetExtentOp> {
   using OpConversionPattern<GetExtentOp>::OpConversionPattern;
 
   LogicalResult
   matchAndRewrite(GetExtentOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override;
 };
 } // namespace
void mlir::populateShapeToStandardConversionPatterns(
   populateWithGenerated(patterns);
   patterns.add<
       AnyOpConversion,
       BinaryOpConversion<AddOp, arith::AddIOp>,
       BinaryOpConversion<MulOp, arith::MulIOp>,
       BroadcastOpConverter,
       ConstShapeOpConverter,
       ConstSizeOpConversion,
+      DimOpConverter,
       IsBroadcastableOpConverter,
       GetExtentOpConverter,
       RankOpConverter,
       ReduceOpConverter,
       ShapeEqOpConverter,
       ShapeOfOpConversion,
       SplitAtOpConversion,
       ToExtentTensorOpConversion>(patterns.getContext());
   // clang-format on
 }
 
 std::unique_ptr<OperationPass<ModuleOp>>
 mlir::createConvertShapeToStandardPass() {
   return std::make_unique<ConvertShapeToStandardPass>();
 }

mlir/lib/Dialect/Shape/IR/Shape.cpp

 // CstrRequireOp
 //===----------------------------------------------------------------------===//
 
 OpFoldResult CstrRequireOp::fold(ArrayRef<Attribute> operands) {
   return operands[0];
 }
 
 //===----------------------------------------------------------------------===//
+// DimOp
+//===----------------------------------------------------------------------===//
+
+Optional<int64_t> DimOp::getConstantDim() {
+  if (auto constSizeOp = getDim().getDefiningOp<ConstSizeOp>())
+    return constSizeOp.getValue().getLimitedValue();
+  if (auto constantOp = getDim().getDefiningOp<arith::ConstantOp>())
+    return constantOp.getValue().cast<IntegerAttr>().getInt();
+  return llvm::None;
+}
+
+OpFoldResult DimOp::fold(ArrayRef<Attribute> operands) {
+  Type valType = getValue().getType();
+  auto valShapedType = valType.dyn_cast<ShapedType>();
+  if (!valShapedType || !valShapedType.hasRank()) return nullptr;
+  Optional<int64_t> dim = getConstantDim();
+  if (!dim.has_value())
+    return nullptr;
+  if (dim.value() >= valShapedType.getRank())
+    return nullptr;
+  auto extent = valShapedType.getDimSize(*dim);
+  if (ShapedType::isDynamic(extent))
+    return nullptr;
+  return IntegerAttr::get(IndexType::get(getContext()), extent);
+}
+
+LogicalResult mlir::shape::DimOp::inferReturnTypes(
+    MLIRContext *context, Optional<Location> location, ValueRange operands,
+    DictionaryAttr attributes, RegionRange regions,
+    SmallVectorImpl<Type> &inferredReturnTypes) {
+  DimOpAdaptor dimOp(operands);
+  inferredReturnTypes.assign({dimOp.getDim().getType()});
+  return success();
+}
+
+bool mlir::shape::DimOp::isCompatibleReturnTypes(TypeRange l,
+                                                       TypeRange r) {
+  return eachHasOnlyOneOfTypes<SizeType, IndexType>(l, r);
+}
+
+LogicalResult mlir::shape::DimOp::verify() {
+  auto st = getValue().getType().cast<ShapedType>();
+  if (!st.hasRank())
+    return success();
+  if (auto dim = getConstantDim()) {
+    if (dim < 0 || dim >= st.getRank())
+      return emitOpError("index is out of range");
+  }
+  return success();
+}
+
+//===----------------------------------------------------------------------===//
 // DivOp
 //===----------------------------------------------------------------------===//
 
 OpFoldResult DivOp::fold(ArrayRef<Attribute> operands) {
   auto lhs = operands[0].dyn_cast_or_null<IntegerAttr>();
   if (!lhs)
     return nullptr;
   auto rhs = operands[1].dyn_cast_or_null<IntegerAttr>();

mlir/test/Conversion/ShapeToStandard/shape-to-standard.mlir

 func.func @rank(%shape : !shape.shape) {
   // CHECK: shape.rank
   %rank = shape.rank %shape : !shape.shape -> !shape.size
   return
 }
 
 // -----
 
+// Express `shape.dim` as `tensor.dim` when valid.
+// CHECK-LABEL: @dim
+// CHECK-SAME:  (%[[ARG:.*]]: tensor<2x3xf32>, %[[IDX:.*]]: index) -> index
+func.func @dim(%arg : tensor<2x3xf32>, %idx : index) -> index {
+  // CHECK: %[[RESULT:.*]] = tensor.dim %[[ARG]], %[[IDX]] : tensor<2x3xf32>
+  // CHECK: return %[[RESULT]] : index
+  %result = shape.dim %arg, %idx : tensor<2x3xf32>, index -> index
+  return %result : index
+}
+
+// -----
+
 // Express `get_extent` as `tensor.dim` when it relies directly on the outcome of a
 // `shape_of` operation.
 // CHECK-LABEL: @get_extent_shape_of
 // CHECK-SAME:  (%[[ARG:.*]]: tensor<2x3xf32>, %[[IDX:.*]]: index) -> index
 func.func @get_extent_shape_of(%arg : tensor<2x3xf32>, %idx : index) -> index {
   // CHECK: %[[RESULT:.*]] = tensor.dim %[[ARG]], %[[IDX]] : tensor<2x3xf32>
   // CHECK: return %[[RESULT]] : index
   %shape = shape.shape_of %arg : tensor<2x3xf32> -> tensor<?xindex>

mlir/test/Dialect/Shape/ops.mlir

 }
 
 func.func @get_extent_on_extent_tensor(%arg : tensor<?xindex>) -> index {
   %c0 = arith.constant 0 : index
   %result = shape.get_extent %arg, %c0 : tensor<?xindex>, index -> index
   return %result : index
 }
 
+func.func @get_dim(%arg : memref<?x?xindex>) -> index {
+  %c0 = arith.constant 0 : index
+  %result = shape.dim %arg, %c0 : memref<?x?xindex>, index -> index
+  return %result : index
+}
+
 func.func @get_extent_on_mixed_operands(%arg : tensor<?xindex>) -> !shape.size {
   %c0 = shape.const_size 0
   %result = shape.get_extent %arg, %c0 : tensor<?xindex>, !shape.size -> !shape.size
   return %result : !shape.size
 }
 
 func.func @any() {
   %0 = shape.const_shape [1, 2, 3] : !shape.shape