[mlir]: Add folding for DimOp with AllocaOp operand

[akuegel]

Summary

This adds AllocaOp to the operand ops that we support for folding.
This folding can be very useful for extent tensor computations.

[bkramer]

mlir/lib/Dialect/StandardOps/IR/Ops.cpp

1528

This should also be safe for any AllocLikeOp

1537

I don't understand why this must be a rank op

[akuegel]

Make the pattern more general.

[akuegel]

mlir/lib/Dialect/StandardOps/IR/Ops.cpp

1537

You are completely right, it doesn't have to be a rank op. This was my motivating example, but no need to restrict it to that.

[bondhugula]

mlir/test/Dialect/Standard/canonicalize.mlir

137–146

This folding already exists and happens upstream on this test case. Is this an (extra) test case for existing functionality?

The fold pattern on DimOp already handles this for any dimensionality memrefs.

155

Please change the name of this function - it gives the impression that this has something to do with the RankOp but in reality it has nothing to. It could as well be the result of any operation or even a block argument. Also, this can be done for any dimensional memref.

157

This is really a folding pattern and should be implemented as part of DimOp::fold - it allows faster progress in the canonicalizer (greedy rewrite driver) and that's the standard approach. All such replacements on other ops are actually implemented on the folding hook - some that were canonicalization patterns earlier were migrated almost a year ago to the folding patterns. The fold methods are also called during the greedy rewrite driver. It just allows more and earlier simplification with this being part of fold.

[akuegel]

Move canonicalization pattern to folding.

[akuegel]

mlir/test/Dialect/Standard/canonicalize.mlir

137–146

You are right, I just hadn't seen that this already works. I started with noticing that dim(alloca) of a 1D memref was not canonicalized, and that was what I initially wanted to fix. I removed the test case again.

157

Thanks, that makes sense :)
I moved it to fold.

[akuegel]

I think I did the requested changes. Can you please take another look? :)

[bondhugula]

A couple more comments @akuegel

mlir/lib/Dialect/StandardOps/IR/Ops.cpp

1501–1505

Can these two ifs be combined into a single one with the dyn_cast_or_null on AllocLikeOp?

mlir/test/Dialect/Standard/canonicalize.mlir

147

Could you add extra lines to test the case when the alloc operand is a block argument? That should just fold as well.

[bondhugula]

mlir/lib/Dialect/StandardOps/IR/Ops.cpp

1447–1448

I still don't understand what the connection with rank 1 is. This folding for DimOp or the propagation of the AllocLikeOp operand can happen for any ranked memref operand of a DimOp.

[bondhugula]

mlir/lib/Dialect/StandardOps/IR/Ops.cpp

1448

if the code is correct <- This would be a bit misleading. I think you are referring to undefined behavior which would happen if the index of a dim op is >= rank of memref/tensor. Basically, you would be exploiting undefined behavior.

if the code is correct -> for the behavior to be defined ?

@mehdi_amini @rriddle Should we exploiting undefined behavior this way?

1450–1455

I think I now see what you were handling and the other thing I was thinking about. I was just thinking (separately from the intention of your patch) that you could get the indexth operand from the AllocLikeOp that's defining this memref and fold the DimOp by replacing its result with that value. A Value is returned as OpFoldResult in such cases. It is orthogonal to what you are doing here.

[akuegel]

Address review comments.

[akuegel]

mlir/lib/Dialect/StandardOps/IR/Ops.cpp

1448

I thought about this a bit more, and even if we allow exploiting undefined behavior (which seems to be not clear), it is not worth it. If the emitted code knows it is dealing with a 1D memref, I think the index operand of DimOp would always be a constant. So I removed this part again.

1450–1455

If it is orthogonal, I guess it could be done in a separate patch.
For now, I only care about that the folding also works for AllocaOp :)

1501–1505

As far as I can tell, the class AllocLikeOp only exists in the Ops.td file, it is not generated as part of the code created with tablegen. If I look at the generated AllocOp class, it derives from ::mlir::Op, not from AllocLikeOp. But I am not too familiar with the MLIR infrastructure yet, so it is possible I overlooked something.
I have checked for instances dealing with both AllocOp and AllocaOp, they seem to always use a template parameter AllocLikeOp. But I guess we cannot use a template for the fold method?

[akuegel]

Please take another look :)
I have removed the controversial part of the change.

[bondhugula]

LGTM.

mlir/lib/Dialect/StandardOps/IR/Ops.cpp

1501–1505

Does it compile when you use AllocLikeOp? Yes, it only appears in the td file. A class doesn't have to inherit from another for cast/dyn_cast etc. to work. (For eg. Operation * cast to any derived OpTy.)

[akuegel]

Rebase.

[akuegel]

Fix test comment.

[akuegel]

Does it compile when you use AllocLikeOp? Yes, it only appears in the td file. A class doesn't have to inherit from another for cast/dyn_cast etc. to work. (For eg. Operation * cast to any derived OpTy.)

Unfortunately it doesn't compile. No C++ class seems to be generated for it (I didn't find it in the Ops.h.inc (now MemRefOps.h.inc) file where e.g. AllocOp and AllocaOp classes are forward declared). Maybe it would be generated in some other file, and I would need to include that other file?

mlir/lib/Dialect/StandardOps/IR/Ops.cpp

static LogicalResult verify(DimOp op) {
   return success();
 }
 
 OpFoldResult DimOp::fold(ArrayRef<Attribute> operands) {
   auto index = operands[1].dyn_cast_or_null<IntegerAttr>();
 
   // All forms of folding require a known index.
   if (!index)
     return {};
 
   auto argTy = memrefOrTensor().getType();
   // Fold if the shape extent along the given index is known.
   if (auto shapedTy = argTy.dyn_cast<ShapedType>()) {
     // Folding for unranked types (UnrankedMemRefType, UnrankedTensorType) is
     // not supported.
     if (!shapedTy.hasRank())
       return {};
     if (!shapedTy.isDynamicDim(index.getInt())) {
       Builder builder(getContext());
       return builder.getIndexAttr(shapedTy.getShape()[index.getInt()]);
     }
   }
 
   Operation *definingOp = memrefOrTensor().getDefiningOp();
   // dim(tensor_load(memref)) -> dim(memref)
if (auto subtensor = dyn_cast_or_null<SubTensorOp>(definingOp)) {
     return subtensor.getDynamicSize(unsignedIndex);
   }
 
   // Fold dim to the size argument for an `AllocOp`, `ViewOp`, or `SubViewOp`.
   auto memrefType = argTy.dyn_cast<MemRefType>();
   if (!memrefType)
     return {};
 
   if (auto alloc = dyn_cast_or_null<AllocOp>(definingOp))
     return *(alloc.getDynamicSizes().begin() +
              memrefType.getDynamicDimIndex(unsignedIndex));
 
   if (auto view = dyn_cast_or_null<ViewOp>(definingOp))
     return *(view.getDynamicSizes().begin() +
              memrefType.getDynamicDimIndex(unsignedIndex));
 
   if (auto subview = dyn_cast_or_null<SubViewOp>(definingOp)) {
     assert(subview.isDynamicSize(unsignedIndex) &&
            "Expected dynamic subview size");
     return subview.getDynamicSize(unsignedIndex);
   }
 
   // dim(memrefcast) -> dim
   if (succeeded(foldMemRefCast(*this)))
     return getResult();
 
   return {};
 }
 
 namespace {
+template <typename AllocLikeOp>
+struct DimOf1DAllocLikeOp : public OpRewritePattern<DimOp> {
+  using OpRewritePattern<DimOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(DimOp dim,
+                                PatternRewriter &rewriter) const override {
+    auto alloc = dim.memrefOrTensor().getDefiningOp<AllocLikeOp>();
+
+    if (!alloc)
+      return failure();
+
+    auto type = alloc.getType().template cast<ShapedType>();
+    if (!type || type.getRank() != 1)
+      return failure();
+
+    if (alloc.dynamicSizes().size() == 1)
+      rewriter.replaceOp(dim, {alloc.dynamicSizes()[0]});
+    else
+      rewriter.replaceOpWithNewOp<ConstantIndexOp>(dim, type.getShape()[0]);
+
+    return success();
+  }
+};
+
 /// Fold dim of a memref reshape operation to a load into the reshape's shape
 /// operand.
 struct DimOfMemRefReshape : public OpRewritePattern<DimOp> {
   using OpRewritePattern<DimOp>::OpRewritePattern;
 
   LogicalResult matchAndRewrite(DimOp dim,
                                 PatternRewriter &rewriter) const override {
     auto reshape = dim.memrefOrTensor().getDefiningOp<MemRefReshapeOp>();
LogicalResult matchAndRewrite(DimOp dimOp,
     rewriter.replaceOpWithNewOp<DimOp>(dimOp, newSource, dimOp.index());
     return success();
   }
 };
 } // end anonymous namespace.
 
 void DimOp::getCanonicalizationPatterns(OwningRewritePatternList &results,
                                         MLIRContext *context) {
-  results.insert<DimOfMemRefReshape, DimOfCastOp<TensorToMemrefOp>,
+  results.insert<DimOf1DAllocLikeOp<AllocOp>, DimOf1DAllocLikeOp<AllocaOp>,
+                 DimOfMemRefReshape, DimOfCastOp<TensorToMemrefOp>,
                  DimOfCastOp<tensor::CastOp>>(context);
 }
 
 // ---------------------------------------------------------------------------
 // DivFOp
 // ---------------------------------------------------------------------------
 
 OpFoldResult DivFOp::fold(ArrayRef<Attribute> operands) {

mlir/test/Dialect/Standard/canonicalize.mlir

func @cmpi_equal_operands(%arg0: i64)
   %8 = cmpi ult, %arg0, %arg0 : i64
   %9 = cmpi ugt, %arg0, %arg0 : i64
   return %0, %1, %2, %3, %4, %5, %6, %7, %8, %9
       : i1, i1, i1, i1, i1, i1, i1, i1, i1, i1
 }
 
 // -----
 
+// Test case: Folding of dim((alloc(%idx) -> rank(%v)
+// CHECK-LABEL: func @dim_of_1d_alloc_with_static_shape
+//  CHECK-NEXT:   %[[C10:.*]] = constant 10 : index
+//  CHECK-NEXT:   return %[[C10]] : index
+func @dim_of_1d_alloc_with_static_shape() -> index {
+  %0 = alloc() : memref<10xindex>
+  %c0 = constant 0 : index
+  %1 = dim %0, %c0 : memref<10xindex>
+  return %1 : index
+}
+
+// -----
+
+// Test case: Folding of dim((alloca(rank(%v)), %idx) -> rank(%v)
+// CHECK-LABEL: func @dim_of_1d_alloca_from_rank(
+//  CHECK-SAME:     %[[MEM:[0-9a-z]+]]: memref<*xf32>
+//  CHECK-NEXT:   %[[RANK:.*]] = rank %[[MEM]] : memref<*xf32>
+//  CHECK-NEXT:   return %[[RANK]] : index
+func @dim_of_1d_alloca_from_rank(%arg0: memref<*xf32>)
+    -> index {
+  %0 = rank %arg0 : memref<*xf32>
+  %1 = alloca(%0) : memref<?xindex>
+  %c0 = constant 0 : index
+  %2 = dim %1, %c0 : memref<?xindex>
+  return %2 : index
+}
+
+// -----
+
 // Test case: Folding of dim(memref_reshape %v %shp, %idx) -> load %shp[%idx]
 // CHECK-LABEL: func @dim_of_memref_reshape(
 //  CHECK-SAME:     %[[MEM:[0-9a-z]+]]: memref<*xf32>,
 //  CHECK-SAME:     %[[SHP:[0-9a-z]+]]: memref<?xindex>
 //  CHECK-NEXT:   %[[IDX:.*]] = constant 3
 //  CHECK-NEXT:   %[[DIM:.*]] = load %[[SHP]][%[[IDX]]]
 //  CHECK-NEXT:   store
 //   CHECK-NOT:   dim