Fold full-size subview of static shapes.

[asaadaldien]

Fix check-label name

[rriddle]

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

3504

Can you add local variables for the ShapedType of source() and getResult()? Would simplify here and above.

[asaadaldien]

Use local {result|source}shapedType variables

[rriddle]

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

3500

Are all of these cases just handling the situation where the input and output type are the same?

[asaadaldien]

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

3500

This can be type equality in the static case.

[rriddle]

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

3500

Can you just write it using type comparison then? Type comparison is O(1), whereas shape comparison is O(N). Both of these if statements are handling static shapes.

[asaadaldien]

Use shaped-type equality

[asaadaldien]

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

3500

Done! but hasStaticShape() itself is O(N)

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

 
 void SubViewOp::getCanonicalizationPatterns(OwningRewritePatternList &results,
                                             MLIRContext *context) {
   results.insert<OpWithOffsetSizesAndStridesConstantArgumentFolder<SubViewOp>,
                  SubViewOpMemRefCastFolder>(context);
 }
 
 OpFoldResult SubViewOp::fold(ArrayRef<Attribute> operands) {
-  if (getResult().getType().cast<ShapedType>().getRank() == 0 &&
-      source().getType().cast<ShapedType>().getRank() == 0)
+  auto resultShapedType = getResult().getType().cast<ShapedType>();
+  auto sourceShapedType = source().getType().cast<ShapedType>();
+
+  if (resultShapedType.getRank() == 0 && sourceShapedType.getRank() == 0)
+    return getViewSource();
+
+  if (sourceShapedType.hasStaticShape() &&
+      resultShapedType.hasStaticShape(sourceShapedType.getShape())) {
     return getViewSource();
+  }
 
   return {};
 }
 
 //===----------------------------------------------------------------------===//
 // SubTensorOp
 //===----------------------------------------------------------------------===//
 
 /// A subtensor result type can be fully inferred from the source type and the
 /// static representation of offsets, sizes and strides. Special sentinels

mlir/test/Dialect/Standard/canonicalize.mlir

func @subview_of_memcast(%arg : memref<4x6x16x32xi8>) ->
   %1 = subview %0[0, 1, 0] [1, 1, 16] [1, 1, 1] :
     memref<?x?x16x32xi8> to
     memref<16x32xi8, affine_map<(d0, d1)[s0] -> (d0 * 32 + d1 + s0)>>
   return %1 : memref<16x32xi8, affine_map<(d0, d1)[s0] -> (d0 * 32 + d1 + s0)>>
 }
 
 // -----
 
+// CHECK-LABEL: func @subview_of_static_full_size
+// CHECK-SAME: %[[ARG0:.+]]: memref<4x6x16x32xi8>
+// CHECK-NOT: subview
+// CHECK: return %[[ARG0]] : memref<4x6x16x32xi8>
+func @subview_of_static_full_size(%arg0 : memref<4x6x16x32xi8>) -> memref<4x6x16x32xi8> {
+  %0 = subview %arg0[0, 0, 0, 0] [4, 6, 16, 32] [1, 1, 1, 1] : memref<4x6x16x32xi8> to memref<4x6x16x32xi8>
+  return %0 : memref<4x6x16x32xi8>
+}
+
+// -----
+
 // CHECK-LABEL: func @trivial_subtensor
 //  CHECK-SAME:   %[[ARG0:.[a-z0-9A-Z_]+]]: tensor<4x6x16x32xi8>
 //   CHECK-NOT:   subtensor
 //       CHECK:   return %[[ARG0]] :  tensor<4x6x16x32xi8>
 func @trivial_subtensor(%arg0 : tensor<4x6x16x32xi8>) -> tensor<4x6x16x32xi8> {
   %0 = subtensor %arg0[0, 0, 0, 0] [4, 6, 16, 32] [1, 1, 1, 1] : tensor<4x6x16x32xi8> to tensor<4x6x16x32xi8>
   return %0 : tensor<4x6x16x32xi8>
 }