[mlir][MemRef] Fix DimOp folding of OffsetSizeAndStrideInterface.

[nicolasvasilache]

Summary

This addresses the issue reported in

https://llvm.discourse.group/t/rank-reducing-memref-subview-offsetsizeandstrideopinterface-interface-issues/3805

[Hardcode84]

I am not sure this is a good solution. Requiring from OffsetSizeAndStrideOpInterface users to to do such index manipulation is very inconvenient and error-prone. Maybe better to separate isDynamicDim APIs into 2 sets isInputDynamicDim/isOutputDynamicDim which can do index translation under the hood?

Also, locally I workarounded this by never using rank-reducing subview and instead implementing a custom op to do a rank reduction on memrefs (with this approach I can also precisely control, which dimensions I want to collapse, not just all-or-nothing).

[nicolasvasilache]

Add a ShapedType::getRelativeIndexOfDynamicDim helper.

[nicolasvasilache]

Make index actually start at 0.

[Hardcode84]

see comment

otherwise LGTM, thanks

mlir/lib/Dialect/Tensor/IR/TensorOps.cpp

274

Is this part of this patch? It needs tests then

[nicolasvasilache]

Add test.

[nicolasvasilache]

This simple implementation is unfortunately not extensible and needs to be reverted.
The extensible way should be to extend https://reviews.llvm.org/D104321.

[silvas]

mlir/lib/Dialect/Tensor/IR/TensorOps.cpp

277

btw, I don't think folders are allowed to mutate the op like this. This needs to be a canonicalization pattern.

[nicolasvasilache]

mlir/lib/Dialect/Tensor/IR/TensorOps.cpp

277

I think this is fine, see the MLIR doc:

"fold has the restriction that no new operations may be created, and only the root operation may be replaced. It allows for updating an operation in-place, or returning a set of pre-existing values (or attributes) to replace the operation with"

[bondhugula]

mlir/lib/Dialect/Tensor/IR/TensorOps.cpp

277

Folders can mutate the op in place - in fact, the folder is prefered to a canonicalization pattern for such local op mutation because it lets the rewrite driver converge faster (to do more work early) -- since the folder would have run on it before the highest benefit canonicalization pattern is picked later in an iteration of the rewriter driver.

[silvas]

mlir/lib/Dialect/Tensor/IR/TensorOps.cpp

277

Oh, fascinating! TIL!

mlir/include/mlir/IR/BuiltinTypes.h

public:
 
   /// If this has a static shape and the shape is equal to `shape` return true.
   bool hasStaticShape(ArrayRef<int64_t> shape) const;
 
   /// If this is a ranked type, return the number of dimensions with dynamic
   /// size. Otherwise, abort.
   int64_t getNumDynamicDims() const;
 
+  /// If `dim` is a dynamic dim, return its relative index among the dynamic
+  /// dims. Otherwise, abort. The result is guaranteed to be nonnegative.
+  int64_t getRelativeIndexOfDynamicDim(unsigned dim) const;
+
   /// If this is ranked type, return the size of the specified dimension.
   /// Otherwise, abort.
   int64_t getDimSize(unsigned idx) const;
 
   /// Returns true if this dimension has a dynamic size (for ranked types);
   /// aborts for unranked types.
   bool isDynamicDim(unsigned idx) const;
 

mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp

 /// Fold alloc operations with no users or only store and dealloc uses.
 template <typename T>
 struct SimplifyDeadAlloc : public OpRewritePattern<T> {
   using OpRewritePattern<T>::OpRewritePattern;
 
   LogicalResult matchAndRewrite(T alloc,
                                 PatternRewriter &rewriter) const override {
     if (llvm::any_of(alloc->getUsers(), [&](Operation *op) {
           if (auto storeOp = dyn_cast<StoreOp>(op))
             return storeOp.value() == alloc;
           return !isa<DeallocOp>(op);
         }))
       return failure();
 
     for (Operation *user : llvm::make_early_inc_range(alloc->getUsers()))
       rewriter.eraseOp(user);
 
     rewriter.eraseOp(alloc);
     return success();
return *(alloca.getDynamicSizes().begin() +
              memrefType.getDynamicDimIndex(unsignedIndex));
 
   if (auto view = dyn_cast_or_null<ViewOp>(definingOp))
     return *(view.getDynamicSizes().begin() +
              memrefType.getDynamicDimIndex(unsignedIndex));
 
   if (auto sizeInterface =
           dyn_cast_or_null<OffsetSizeAndStrideOpInterface>(definingOp)) {
-    assert(sizeInterface.isDynamicSize(unsignedIndex) &&
-           "Expected dynamic subview size");
-    return sizeInterface.getDynamicSize(unsignedIndex);
+    int64_t nthDynamicIndex =
+        memrefType.getRelativeIndexOfDynamicDim(unsignedIndex);
+    return sizeInterface.sizes()[nthDynamicIndex];
   }
 
   // dim(memrefcast) -> dim
   if (succeeded(foldMemRefCast(*this)))
     return getResult();
 
   return {};
 }

mlir/lib/Dialect/Tensor/IR/TensorOps.cpp

if (auto fromElements = dyn_cast_or_null<tensor::GenerateOp>(definingOp)) {
     auto dynExtents = fromElements.dynamicExtents().begin();
     for (auto dim : resultType.getShape().take_front(index.getInt()))
       if (dim == RankedTensorType::kDynamicSize)
         dynExtents++;
 
     return Value{*dynExtents};
   }
 
+  // dim(insert_slice.result()) -> dim(insert_slice.dest())
+  if (auto insertSliceOp =
+          dyn_cast_or_null<tensor::InsertSliceOp>(definingOp)) {
+    this->sourceMutable().assign(insertSliceOp.dest());
+    return getResult();
+  }
+
   // The size at the given index is now known to be a dynamic size.
   unsigned unsignedIndex = index.getValue().getZExtValue();
 
-  if (auto sliceOp = dyn_cast_or_null<tensor::ExtractSliceOp>(definingOp)) {
-    assert(sliceOp.isDynamicSize(unsignedIndex) &&
-           "Expected dynamic slice size");
-    return sliceOp.getDynamicSize(unsignedIndex);
+  if (auto sizeInterface =
+          dyn_cast_or_null<OffsetSizeAndStrideOpInterface>(definingOp)) {
+    int64_t nthDynamicIndex =
+        tensorType.getRelativeIndexOfDynamicDim(unsignedIndex);
+    return sizeInterface.sizes()[nthDynamicIndex];
   }
 
   // dim(cast) -> dim
   if (succeeded(foldTensorCast(*this)))
     return getResult();
 
   return {};
 }

mlir/lib/IR/BuiltinTypes.cpp

if (auto tensorType = dyn_cast<RankedTensorType>())
     return tensorType.getShape();
   return cast<MemRefType>().getShape();
 }
 
 int64_t ShapedType::getNumDynamicDims() const {
   return llvm::count_if(getShape(), isDynamic);
 }
 
+int64_t ShapedType::getRelativeIndexOfDynamicDim(unsigned dim) const {
+  assert(isDynamicDim(dim) && "expected a dynamic dim");
+  int nthDynamicIndex = -1;
+  for (unsigned idx = 0; idx <= dim; ++idx)
+    if (isDynamicDim(idx))
+      ++nthDynamicIndex;
+  return nthDynamicIndex;
+}
+
 bool ShapedType::hasStaticShape() const {
   return hasRank() && llvm::none_of(getShape(), isDynamic);
 }
 
 bool ShapedType::hasStaticShape(ArrayRef<int64_t> shape) const {
   return hasStaticShape() && getShape() == shape;
 }
 

mlir/test/Dialect/MemRef/canonicalize.mlir

 #map0 = affine_map<(d0)[s0, s1] -> (d0 * s1 + s0)>
 func @alloc_const_fold_with_symbols2() -> memref<?xi32, #map0> {
   %c1 = constant 1 : index
   %0 = memref.alloc(%c1)[%c1, %c1] : memref<?xi32, #map0>
   return %0 : memref<?xi32, #map0>
 }
 
 // -----
+
 // CHECK-LABEL: func @allocator
 // CHECK:   %[[alloc:.+]] = memref.alloc
 // CHECK:   memref.store %[[alloc:.+]], %arg0
 func @allocator(%arg0 : memref<memref<?xi32>>, %arg1 : index)  {
   %0 = memref.alloc(%arg1) : memref<?xi32>
   memref.store %0, %arg0[] : memref<memref<?xi32>>
-  return 
+  return
 }
+
+// -----
+
+#map0 = affine_map<(d0, d1)[s0, s1, s2] -> (d0 * s1 + s0 + d1 * s2)>
+
+// CHECK-LABEL: func @rank_reducing_subview_dim
+//  CHECK-SAME:   %[[IDX_0:[0-9a-zA-Z]*]]: index
+//  CHECK-SAME:   %[[IDX_1:[0-9a-zA-Z]*]]: index
+func @rank_reducing_subview_dim(%arg0 : memref<?x?x?xf32>, %arg1 : index,
+    %arg2 : index) -> index
+{
+  %c0 = constant 0 : index
+  %c1 = constant 1 : index
+  %c4 = constant 4 : index
+  %0 = memref.subview %arg0[%c0, %arg1, %c1] [%c4, 1, %arg2] [%c1, %c1, %c1] : memref<?x?x?xf32> to memref<?x?xf32, #map0>
+  %1 = memref.dim %0, %c1 : memref<?x?xf32, #map0>
+
+  // CHECK-NEXT: return %[[IDX_1]] : index
+  return %1 : index
+}

mlir/test/Dialect/Tensor/canonicalize.mlir

 func @fold_dim_of_tensor.cast(%arg0 : tensor<4x?xf32>) -> (index, index) {
   %c0 = constant 0 : index
   %c1 = constant 1 : index
   %0 = tensor.cast %arg0 : tensor<4x?xf32> to tensor<?x?xf32>
   %1 = tensor.dim %0, %c0 : tensor<?x?xf32>
   %2 = tensor.dim %0, %c1 : tensor<?x?xf32>
   return %1, %2: index, index
 }
+
+// -----
+
+// CHECK-LABEL: func @rank_reducing_extract_slice_dim
+//  CHECK-SAME:   %[[IDX_0:[0-9a-zA-Z]*]]: index
+//  CHECK-SAME:   %[[IDX_1:[0-9a-zA-Z]*]]: index
+func @rank_reducing_extract_slice_dim(%arg0 : tensor<?x?x?xf32>, %arg1 : index,
+    %arg2 : index) -> index
+{
+  %c0 = constant 0 : index
+  %c1 = constant 1 : index
+  %c4 = constant 4 : index
+  %0 = tensor.extract_slice %arg0[%c0, %arg1, %c1] [%c4, 1, %arg2] [%c1, %c1, %c1] : tensor<?x?x?xf32> to tensor<?x?xf32>
+  %1 = tensor.dim %0, %c1 : tensor<?x?xf32>
+
+  // CHECK-NEXT: return %[[IDX_1]] : index
+  return %1 : index
+}
+
+// -----
+
+// CHECK-LABEL: func @rank_reducing_insert_slice_dim
+//  CHECK-SAME:   %[[OUT:[0-9a-zA-Z]*]]: tensor<?x?x?xf32>
+func @rank_reducing_insert_slice_dim(%out : tensor<?x?x?xf32>, %in : tensor<?x?xf32>, %arg1 : index,
+    %arg2 : index) -> index
+{
+  // CHECK-NEXT: %[[C1:.*]] = constant 1 : index
+
+  %c0 = constant 0 : index
+  %c1 = constant 1 : index
+  %c4 = constant 4 : index
+  %0 = tensor.insert_slice %in into %out[%c0, %arg1, %c1] [1, 1, 1] [%c1, %c1, %c1] : tensor<?x?xf32> into tensor<?x?x?xf32>
+
+  // CHECK-NEXT: %[[D1:.*]] = tensor.dim %[[OUT]], %[[C1]] : tensor<?x?x?xf32>
+  %1 = tensor.dim %0, %c1 : tensor<?x?x?xf32>
+
+  // CHECK-NEXT: return %[[D1]] : index
+  return %1 : index
+}