[mlir] Fold memref.dim of OffsetSizeAndStrideOpInterface outputs

[tpopp]

Summary

This previously handled memref::SubviewOp, but this can be extended to
all ops implementing the interface.

[frgossen]

mlir/test/Dialect/MemRef/canonicalize.mlir

199

nit: not used

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

 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/IR/AffineMap.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/Matchers.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/IR/TypeUtilities.h"
 #include "mlir/Interfaces/InferTypeOpInterface.h"
+#include "mlir/Interfaces/ViewLikeInterface.h"
 #include "llvm/ADT/STLExtras.h"
 
 using namespace mlir;
 using namespace mlir::memref;
 
 /// Materialize a single constant operation from a given attribute value with
 /// the desired resultant type.
 Operation *MemRefDialect::materializeConstant(OpBuilder &builder,
OpFoldResult DimOp::fold(ArrayRef<Attribute> operands) {
   if (auto alloca = dyn_cast_or_null<AllocaOp>(definingOp))
     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 subview = dyn_cast_or_null<SubViewOp>(definingOp)) {
-    assert(subview.isDynamicSize(unsignedIndex) &&
+  if (auto sizeInterface =
+          dyn_cast_or_null<OffsetSizeAndStrideOpInterface>(definingOp)) {
+    assert(sizeInterface.isDynamicSize(unsignedIndex) &&
            "Expected dynamic subview size");
-    return subview.getDynamicSize(unsignedIndex);
+    return sizeInterface.getDynamicSize(unsignedIndex);
   }
 
   // dim(memrefcast) -> dim
   if (succeeded(foldMemRefCast(*this)))
     return getResult();
 
   return {};
 }

mlir/test/Dialect/MemRef/canonicalize.mlir

func @alias_is_freed(%arg0 : memref<?xf32>) {
   // CHECK: memref.dealloc
   %0 = memref.cast %arg0 : memref<?xf32> to memref<32xf32>
   %1 = memref.clone %0 : memref<32xf32> to memref<32xf32>
   memref.dealloc %arg0 : memref<?xf32>
   "use"(%1) : (memref<32xf32>) -> ()
   memref.dealloc %1 : memref<32xf32>
   return
 }
+
+// -----
+
+// CHECK-LABEL: func @dim_of_sized_view
+//  CHECK-SAME:   %{{[a-z0-9A-Z_]+}}: memref<?xi8>
+//  CHECK-SAME:   %[[SIZE:.[a-z0-9A-Z_]+]]: index
+//       CHECK:   return %[[SIZE]] : index
+func @dim_of_sized_view(%arg : memref<?xi8>, %size: index) -> index {
+  %c0 = constant 0 : index
+  %0 = memref.reinterpret_cast %arg to offset: [0], sizes: [%size], strides: [0] : memref<?xi8> to memref<?xi8>
+  %1 = memref.dim %0, %c0 : memref<?xi8>
+  return %1 : index
+}
+
+