diff --git a/mlir/lib/Dialect/StandardOps/IR/Ops.cpp b/mlir/lib/Dialect/StandardOps/IR/Ops.cpp --- a/mlir/lib/Dialect/StandardOps/IR/Ops.cpp +++ b/mlir/lib/Dialect/StandardOps/IR/Ops.cpp @@ -1341,9 +1341,9 @@ // The subview op here is expected to have rank dynamic sizes now. if (auto subview = dyn_cast_or_null(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 diff --git a/mlir/test/Transforms/canonicalize.mlir b/mlir/test/Transforms/canonicalize.mlir --- a/mlir/test/Transforms/canonicalize.mlir +++ b/mlir/test/Transforms/canonicalize.mlir @@ -429,8 +429,13 @@ #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 func @dim_op_fold(%arg0: index, %arg1: index, %arg2: index, %BUF: memref, %M : index, %N : index, %K : index) { // CHECK-SAME: [[M:arg[0-9]+]]: index // CHECK-SAME: [[N:arg[0-9]+]]: index @@ -452,11 +457,20 @@ affine.for %arg5 = %l to %u { "foo"() : () -> () } + %sv2 = subview %0[0, 0][17, %arg4][1, 1] : memref to memref<17x?xf32, #map3> + %l2 = dim %v, 1 : memref + %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: }