[MLIR][Shape] Add conversion for missing ops to standard

[herhut]

Summary

This adds conversions for const_size and to_extent_tensor. Also, cast-like operations are now folded away if the source and target types are the same.

[herhut]

Remove a test

[herhut]

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

543

I don't know how to test this case, as it is illegal. I just added it for completeness if this every materializes during some conversion.

[silvas]

mlir/lib/Conversion/ShapeToStandard/ShapeToStandard.cpp

133

nit: notifyMatchFailure

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

543

I think this would do it:

func @f(%arg0: index) {

%0 = size_to_index(%arg0)
return s

}

foldCastOp is for non-constant but where we otherwise know the op is structurally identity.

[frgossen]

mlir/lib/Conversion/ShapeToStandard/ShapeToStandard.cpp

67

Should we make this pattern fail in case the result type is a shape.size?

[herhut]

mlir/lib/Conversion/ShapeToStandard/ShapeToStandard.cpp

67

But even if the result is a shape.size, we would lower it to a index value when going to standard. It cannot be an error.

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

543

If I write a test case where %arg0 is of type index, then I cannot even parse it because passing and index to this operation is invalid. I guess I should remove it again. I also have no good use for it.

This op is the structural identity of the two types are the same. Would you prefer to spell this out here?

[herhut]

Comments

[silvas]

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

543

Sounds like it can be a TODO for when size_to_index accepts index inputs.

[jpienaar]

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

443

Ouch, thanks :)

[herhut]

Rebase

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

def Shape_SizeToIndexOp : Shape_Op<"size_to_index", [NoSideEffect]> {
   let summary = "Casts between index types of the shape and standard dialect";
   let description = [{
     Converts a `shape.size` to a standard index. This operation and its
     inverse, `index_to_size`, facilitate index conversion between the standard
     and the shape dialect. The behavior is undefined for unknown and invalid
     arguments.
   }];
 
-  let arguments = (outs Shape_SizeOrIndexType:$arg);
+  let arguments = (ins Shape_SizeOrIndexType:$arg);
   let results = (outs Index:$result);
 
   let assemblyFormat = "$arg attr-dict `:` type($arg)";
 
   let hasFolder = 1;
   let hasCanonicalizer = 1;
 }
 

mlir/lib/Conversion/ShapeToStandard/ShapeToStandard.cpp

matchAndRewrite(SrcOpTy op, ArrayRef<Value> operands,
     typename SrcOpTy::Adaptor adaptor(operands);
     rewriter.replaceOpWithNewOp<DstOpTy>(op, adaptor.lhs(), adaptor.rhs());
     return success();
   }
 };
 } // namespace
 
 namespace {
+class ConstSizeOpConversion : public OpConversionPattern<ConstSizeOp> {
+public:
+  using OpConversionPattern<ConstSizeOp>::OpConversionPattern;
+
+  LogicalResult
+  matchAndRewrite(ConstSizeOp op, ArrayRef<Value> operands,
+                  ConversionPatternRewriter &rewriter) const override {
+
+    rewriter.replaceOpWithNewOp<ConstantIndexOp>(op, op.value().getSExtValue());
+    return success();
+  }
+};
+} // namespace
+
+namespace {
 class ShapeOfOpConversion : public OpConversionPattern<ShapeOfOp> {
 public:
   using OpConversionPattern<ShapeOfOp>::OpConversionPattern;
 
   LogicalResult
   matchAndRewrite(ShapeOfOp op, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const override;
 };
public:
 
   LogicalResult
   matchAndRewrite(ConstShapeOp op, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const override;
 };
 } // namespace
 
 LogicalResult ConstShapeOpConverter::matchAndRewrite(
     ConstShapeOp op, ArrayRef<Value> operands,
     ConversionPatternRewriter &rewriter) const {
 
   // For now, this lowering supports only extent tensors, not `shape.shape`
   // types.
   if (op.getType().isa<ShapeType>())
     return failure();
 
   auto loc = op.getLoc();
   SmallVector<Value, 4> extentOperands;
   for (auto extent : op.shape()) {
     extentOperands.push_back(
         rewriter.create<ConstantIndexOp>(loc, extent.getLimitedValue()));
   }
   Value tensor = rewriter.create<TensorFromElementsOp>(loc, extentOperands);
   Type indexTy = rewriter.getIndexType();
   Type resultTy = RankedTensorType::get({ShapedType::kDynamicSize}, indexTy);
   rewriter.replaceOpWithNewOp<TensorCastOp>(op, tensor, resultTy);
   return success();
 }
 
 namespace {
+class ToExtentTensorOpConversion
+    : public OpConversionPattern<ToExtentTensorOp> {
+public:
+  using OpConversionPattern<ToExtentTensorOp>::OpConversionPattern;
+
+  LogicalResult
+  matchAndRewrite(ToExtentTensorOp op, ArrayRef<Value> operands,
+                  ConversionPatternRewriter &rewriter) const override {
+    ToExtentTensorOpAdaptor adaptor(operands);
+
+    if (!adaptor.input().getType().isa<RankedTensorType>())
+      return rewriter.notifyMatchFailure(op, "input needs to be a tensor");
+
+    rewriter.replaceOpWithNewOp<TensorCastOp>(op, adaptor.input(),
+                                              op.getType());
+    return success();
+  }
+};
+} // namespace
+
+namespace {
 class GetExtentOpConverter : public OpConversionPattern<GetExtentOp> {
   using OpConversionPattern<GetExtentOp>::OpConversionPattern;
 
   LogicalResult
   matchAndRewrite(GetExtentOp op, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const override;
 };
 } // namespace

 void mlir::populateShapeToStandardConversionPatterns(
     OwningRewritePatternList &patterns, MLIRContext *ctx) {
   // clang-format off
   patterns.insert<
       AnyOpConversion,
       BinaryOpConversion<AddOp, AddIOp>,
       ConstShapeOpConverter,
       BinaryOpConversion<MulOp, MulIOp>,
+      ConstSizeOpConversion,
       GetExtentOpConverter,
       RankOpConverter,
-      ShapeOfOpConversion>(ctx);
+      ShapeOfOpConversion,
+      ToExtentTensorOpConversion>(ctx);
   // clang-format on
 }
 
 std::unique_ptr<OperationPass<ModuleOp>>
 mlir::createConvertShapeToStandardPass() {
   return std::make_unique<ConvertShapeToStandardPass>();
 }

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

 // IndexToSizeOp
 //===----------------------------------------------------------------------===//
 
 OpFoldResult IndexToSizeOp::fold(ArrayRef<Attribute> operands) {
   // Constant values of both types, `shape.size` and `index`, are represented as
   // `IntegerAttr`s which makes constant folding simple.
   if (Attribute arg = operands[0])
     return arg;
   return {};
 }
 
 void IndexToSizeOp::getCanonicalizationPatterns(
     OwningRewritePatternList &patterns, MLIRContext *context) {
   patterns.insert<SizeToIndexToSizeCanonicalization>(context);
 }
 
 //===----------------------------------------------------------------------===//

 // SizeToIndexOp
 //===----------------------------------------------------------------------===//
 
 OpFoldResult SizeToIndexOp::fold(ArrayRef<Attribute> operands) {
   // Constant values of both types, `shape.size` and `index`, are represented as
   // `IntegerAttr`s which makes constant folding simple.
   if (Attribute arg = operands[0])
     return arg;
-  return {};
+  return impl::foldCastOp(*this);
 }
 
 void SizeToIndexOp::getCanonicalizationPatterns(
     OwningRewritePatternList &patterns, MLIRContext *context) {
   patterns.insert<IndexToSizeToIndexCanonicalization>(context);
 }
 
 //===----------------------------------------------------------------------===//

 }
 
 //===----------------------------------------------------------------------===//
 // ToExtentTensorOp
 //===----------------------------------------------------------------------===//
 
 OpFoldResult ToExtentTensorOp::fold(ArrayRef<Attribute> operands) {
   if (!operands[0])
-    return nullptr;
+    return impl::foldCastOp(*this);
   Builder builder(getContext());
   auto shape = llvm::to_vector<6>(
       operands[0].cast<DenseIntElementsAttr>().getValues<int64_t>());
   auto type = RankedTensorType::get({static_cast<int64_t>(shape.size())},
                                     builder.getIndexType());
   return DenseIntElementsAttr::get(type, shape);
 }
 

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

 // CHECK-LABEL: @any_of_one
 // CHECK-SAME:  (%[[A:.*]]: tensor<?xindex>) -> tensor<?xindex>
 func @any_of_one(%a : tensor<?xindex>) -> tensor<?xindex> {
   // CHECK: return %[[A]] : tensor<?xindex>
   %result = "shape.any"(%a) : (tensor<?xindex>) -> tensor<?xindex>
   return %result : tensor<?xindex>
 }
 
+// -----
+
+// Lower 'const_size` to `std.constant`
+// CHECK-LABEL: @const_size
+func @const_size() -> index {
+  // CHECK: %[[RES:.*]] = constant 42 : index
+  %size = shape.const_size 42
+  %result = shape.size_to_index %size : !shape.size
+  // CHECK: return %[[RES]]
+  return %result : index
+}
+
+// -----
+
+// Lower `to_extent_tensor` to `std.tensor_cast`
+// Fold to_extent_tensor when already on tensor.
+// CHECK-LABEL: @to_extent_tensor
+// CHECK-SAME: (%[[ARG:.*]]: tensor<?xindex>
+func @to_extent_tensor(%arg: tensor<?xindex>) -> tensor<3xindex> {
+  // CHECK-NOT: to_extent_tensor
+  // CHECK: %[[RES:.*]] = tensor_cast %[[ARG]] : tensor<?xindex> to tensor<3xindex
+  %casted = shape.to_extent_tensor %arg : tensor<?xindex> -> tensor<3xindex>
+  // CHECK: return %[[RES]]
+  return %casted : tensor<3xindex>
+}

mlir/test/Dialect/Shape/canonicalize.mlir

func @fold_mul_mixed() -> !shape.size {
   // CHECK: %[[RESULT:.*]] = shape.const_size 6
   // CHECK: return %[[RESULT]] : !shape.size
   %c2 = shape.const_size 2
   %c3 = constant 3 : index
   %result = shape.mul %c2, %c3 : !shape.size, index -> !shape.size
   return %result : !shape.size
 }
 
+// -----
+
+// Fold index_cast when already on index.
+// CHECK-LABEL: @fold_index_cast_on_index
+func @fold_index_cast_on_index(%arg: index) -> index {
+  // CHECK-NOT: size_to_index
+  %casted = shape.size_to_index %arg : index
+  return %casted : index
+}
+
+// -----
+
+// Fold to_extent_tensor when already on tensor.
+// CHECK-LABEL: @fold_to_extent_tensor_on_tensor
+func @fold_to_extent_tensor_on_tensor(%arg: tensor<?xindex>) -> tensor<?xindex> {
+  // CHECK-NOT: to_extent_tensor
+  %casted = shape.to_extent_tensor %arg : tensor<?xindex> -> tensor<?xindex>
+  return %casted : tensor<?xindex>
+}