[Mlir-commits] [mlir] 9ccdc2e - [mlir] Fold memref.dim of OffsetSizeAndStrideOpInterface outputs

[Tres Popp]

Author: Tres Popp
Date: 2021-05-25T12:16:10+02:00
New Revision: 9ccdc2e23be18eca0b09f055fd17115c0366166c

URL: https://github.com/llvm/llvm-project/commit/9ccdc2e23be18eca0b09f055fd17115c0366166c
DIFF: https://github.com/llvm/llvm-project/commit/9ccdc2e23be18eca0b09f055fd17115c0366166c.diff

LOG: [mlir] Fold memref.dim of OffsetSizeAndStrideOpInterface outputs

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

Differential Revision: https://reviews.llvm.org/D103076

Added: 
    

Modified: 
    mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp
    mlir/test/Dialect/MemRef/canonicalize.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp b/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp
index 7e08837474a1b..a4ab6c1d0859f 100644
--- a/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp
+++ b/mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp
@@ -18,6 +18,7 @@
 #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;
@@ -679,10 +680,11 @@ OpFoldResult DimOp::fold(ArrayRef<Attribute> operands) {
     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

diff  --git a/mlir/test/Dialect/MemRef/canonicalize.mlir b/mlir/test/Dialect/MemRef/canonicalize.mlir
index ce04f5048d3f8..354be2237ec30 100644
--- a/mlir/test/Dialect/MemRef/canonicalize.mlir
+++ b/mlir/test/Dialect/MemRef/canonicalize.mlir
@@ -192,3 +192,18 @@ func @alias_is_freed(%arg0 : memref<?xf32>) {
   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
+}
+
+