[mlir] Fix incorrect indexing of subview in DimOp folding.

[nicolasvasilache]

Summary

DimOp folding is using bare accesses to underlying SubViewOp operands.
This is generally incorrect and is fixed in this revision.

[mravishankar]

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

1339–1341

Does this has to be an assert? Is it possible that the subview op has been canonicalized, but the dim op hasnt in which case the subview op will have a static size. You can just return {} if it is not dynamic for now.

[nicolasvasilache]

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

1339–1341

I think so: if the size was static we would have been done at the beginning gettin the information from the memref type.
If not, then we expect to have a dynamic index.
Otherwise the subview itself would be inconsistent and we'd definitely want to fail hard there.

I note however I also need to update the comment.

[nicolasvasilache]

Cleanups.

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

static LogicalResult verify(DimOp op) {
   }
 
   return success();
 }
 
 OpFoldResult DimOp::fold(ArrayRef<Attribute> operands) {
   // Constant fold dim when the size along the index referred to is a constant.
   auto opType = memrefOrTensor().getType();
-  int64_t dimSize = ShapedType::kDynamicSize;
-  if (auto tensorType = opType.dyn_cast<RankedTensorType>())
-    dimSize = tensorType.getShape()[getIndex()];
-  else if (auto memrefType = opType.dyn_cast<MemRefType>())
-    dimSize = memrefType.getShape()[getIndex()];
-
-  if (!ShapedType::isDynamic(dimSize))
-    return IntegerAttr::get(IndexType::get(getContext()), dimSize);
+  if (auto shapedType = opType.dyn_cast<ShapedType>())
+    if (!shapedType.isDynamicDim(getIndex()))
+      return IntegerAttr::get(IndexType::get(getContext()),
+                              shapedType.getShape()[getIndex()]);
 
   // Fold dim to the size argument for an AllocOp/ViewOp/SubViewOp.
   auto memrefType = opType.dyn_cast<MemRefType>();
   if (!memrefType)
     return {};
 
-  // The size at getIndex() is now a dynamic size of a memref.
+  // The size at getIndex() is now known to be a dynamic size of a memref.
   auto memref = memrefOrTensor().getDefiningOp();
   if (auto alloc = dyn_cast_or_null<AllocOp>(memref))
     return *(alloc.getDynamicSizes().begin() +
              memrefType.getDynamicDimIndex(getIndex()));
 
   if (auto view = dyn_cast_or_null<ViewOp>(memref))
     return *(view.getDynamicSizes().begin() +
              memrefType.getDynamicDimIndex(getIndex()));
 
-  // The subview op here is expected to have rank dynamic sizes now.
   if (auto subview = dyn_cast_or_null<SubViewOp>(memref)) {
-    auto sizes = subview.sizes();
-    if (!sizes.empty())
-      return *(sizes.begin() + getIndex());
+    assert(subview.isDynamicSize(getIndex()) &&
+           "Expected dynamic subview size");
+    return subview.getDynamicSize(getIndex());
   }
 
   /// dim(memrefcast) -> dim
   if (succeeded(foldMemRefCast(*this)))
     return getResult();
 
   return {};
 }

mlir/test/Transforms/canonicalize.mlir

affine.for %i = 0 to %L {
     }
   }
 
   return %c, %d : memref<? x ? x i32>, memref<? x ? x f32>
 }
 
 #map1 = affine_map<(d0, d1)[s0, s1, s2] -> (d0 * s1 + s0 + d1 * s2)>
 #map2 = affine_map<(d0, d1, d2)[s0, s1, s2] -> (d0 * s2 + d1 * s1 + d2 + s0)>
+#map3 = affine_map<(d0, d1)[s0, s1] -> (d0 * s1 + s0 + d1)>
 
-// CHECK-LABEL: func @dim_op_fold(%arg0: index, %arg1: index, %arg2: index,
+// CHECK-LABEL: func @dim_op_fold(
+// CHECK-SAME: %[[ARG0:[a-z0-9]*]]: index
+// CHECK-SAME: %[[ARG1:[a-z0-9]*]]: index
+// CHECK-SAME: %[[ARG2:[a-z0-9]*]]: index
+// CHECK-SAME: %[[BUF:[a-z0-9]*]]: memref<?xi8>
 func @dim_op_fold(%arg0: index, %arg1: index, %arg2: index, %BUF: memref<?xi8>, %M : index, %N : index, %K : index) {
 // CHECK-SAME: [[M:arg[0-9]+]]: index
 // CHECK-SAME: [[N:arg[0-9]+]]: index
 // CHECK-SAME: [[K:arg[0-9]+]]: index
   %c0 = constant 0 : index
   %c1 = constant 1 : index
   %0 = alloc(%arg0, %arg1) : memref<?x?xf32>
   %1 = alloc(%arg1, %arg2) : memref<?x8x?xf32>
   %2 = dim %1, 2 : memref<?x8x?xf32>
   affine.for %arg3 = 0 to %2 {
     %3 = alloc(%arg0) : memref<?xi8>
     %ub = dim %3, 0 : memref<?xi8>
     affine.for %arg4 = 0 to %ub {
       %s = dim %0, 0 : memref<?x?xf32>
       %v = std.view %3[%c0][%arg4, %s] : memref<?xi8> to memref<?x?xf32>
       %sv = subview %0[%c0, %c0][%s,%arg4][%c1,%c1] : memref<?x?xf32> to memref<?x?xf32, #map1>
       %l = dim %v, 1 : memref<?x?xf32>
       %u = dim %sv, 0 : memref<?x?xf32, #map1>
       affine.for %arg5 = %l to %u {
         "foo"() : () -> ()
       }
+      %sv2 = subview %0[0, 0][17, %arg4][1, 1] : memref<?x?xf32> to memref<17x?xf32, #map3>
+      %l2 = dim %v, 1 : memref<?x?xf32>
+      %u2 = dim %sv2, 1 : memref<17x?xf32, #map3>
+      scf.for %arg5 = %l2 to %u2 step %c1 {
+        "foo"() : () -> ()
+      }
     }
   }
-  // CHECK-NEXT: affine.for %arg7 = 0 to %arg2 {
-  // CHECK-NEXT:   affine.for %arg8 = 0 to %arg0 {
-  // CHECK-NEXT:     affine.for %arg9 = %arg0 to %arg0 {
+  //      CHECK: affine.for %[[I:.*]] = 0 to %[[ARG2]] {
+  // CHECK-NEXT:   affine.for %[[J:.*]] = 0 to %[[ARG0]] {
+  // CHECK-NEXT:     affine.for %[[K:.*]] = %[[ARG0]] to %[[ARG0]] {
+  // CHECK-NEXT:       "foo"() : () -> ()
+  // CHECK-NEXT:     }
+  // CHECK-NEXT:     scf.for %[[KK:.*]] = %[[ARG0]] to %[[J]] step %{{.*}} {
   // CHECK-NEXT:       "foo"() : () -> ()
   // CHECK-NEXT:     }
   // CHECK-NEXT:   }
   // CHECK-NEXT: }
 
   %A = view %BUF[%c0][%M, %K] : memref<?xi8> to memref<?x?xf32>
   %B = view %BUF[%c0][%K, %N] : memref<?xi8> to memref<?x?xf32>
   %C = view %BUF[%c0][%M, %N] : memref<?xi8> to memref<?x?xf32>