[FlattenMemRef] Flatten MemRef ReshapeOp when its source memory reference is flattened in other passes (#8265)

Author: jiahanxie353

lib/Transforms/FlattenMemRefs.cpp

@@ -10,6 +10,7 @@
 //
 //===----------------------------------------------------------------------===//
 
+#include "circt/Support/LLVM.h"
 #include "circt/Transforms/Passes.h"
 #include "mlir/Conversion/LLVMCommon/ConversionTarget.h"
 #include "mlir/Conversion/LLVMCommon/Pattern.h"
@@ -25,6 +26,7 @@
 #include "mlir/Transforms/DialectConversion.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 #include "llvm/Support/FormatVariadic.h"
+#include "llvm/Support/LogicalResult.h"
 #include "llvm/Support/MathExtras.h"
 
 namespace circt {
@@ -238,6 +240,27 @@ struct GetGlobalOpConversion : public OpConversionPattern<memref::GetGlobalOp> {
   }
 };
 
+struct ReshapeOpConversion : public OpConversionPattern<memref::ReshapeOp> {
+  using OpConversionPattern::OpConversionPattern;
+
+  LogicalResult
+  matchAndRewrite(memref::ReshapeOp op, OpAdaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override {
+    Value flattenedSource = rewriter.getRemappedValue(op.getSource());
+    if (!flattenedSource)
+      return failure();
+
+    auto flattenedSrcType = cast<MemRefType>(flattenedSource.getType());
+    if (isUniDimensional(flattenedSrcType) ||
+        !flattenedSrcType.hasStaticShape()) {
+      rewriter.replaceOp(op, flattenedSource);
+      return success();
+    }
+
+    return failure();
+  }
+};
+
 // A generic pattern which will replace an op with a new op of the same type
 // but using the adaptor (type converted) operands.
 template <typename TOp>
@@ -403,7 +426,7 @@ struct FlattenMemRefPass
     RewritePatternSet patterns(ctx);
     SetVector<StringRef> rewrittenCallees;
     patterns.add<LoadOpConversion, StoreOpConversion, AllocOpConversion,
-                 GlobalOpConversion, GetGlobalOpConversion,
+                 GlobalOpConversion, GetGlobalOpConversion, ReshapeOpConversion,
                  OperandConversionPattern<func::ReturnOp>,
                  OperandConversionPattern<memref::DeallocOp>,
                  CondBranchOpConversion,

test/Transforms/flatten_memref.mlir

@@ -271,3 +271,46 @@ module {
   }
 }
 
+// -----
+
+// CHECK:   func.func @main(%[[VAL_0:arg0]]: memref<30xf32>) {
+// CHECK:           %[[VAL_1:.*]] = arith.constant 30 : index
+// CHECK:           %[[VAL_2:.*]] = arith.constant 1 : index
+// CHECK:           %[[VAL_3:.*]] = arith.constant 2 : index
+// CHECK:           %[[VAL_4:.*]] = arith.constant 0 : index
+// CHECK:           %[[VAL_5:.*]] = memref.alloc() : memref<30xf32>
+// CHECK:           %[[VAL_6:.*]] = memref.get_global @const_1_30 : memref<2xi64>
+// CHECK:           scf.for %[[VAL_7:.*]] = %[[VAL_4]] to %[[VAL_3]] step %[[VAL_2]] {
+// CHECK:             scf.for %[[VAL_8:.*]] = %[[VAL_4]] to %[[VAL_1]] step %[[VAL_2]] {
+// CHECK:               %[[VAL_9:.*]] = arith.constant 30 : index
+// CHECK:               %[[VAL_10:.*]] = arith.muli %[[VAL_4]], %[[VAL_9]] : index
+// CHECK:               %[[VAL_11:.*]] = arith.addi %[[VAL_10]], %[[VAL_8]] : index
+// CHECK:               %[[VAL_12:.*]] = memref.load %[[VAL_5]]{{\[}}%[[VAL_11]]] : memref<30xf32>
+// CHECK:               %[[VAL_13:.*]] = arith.constant 0 : index
+// CHECK:               %[[VAL_14:.*]] = arith.shli %[[VAL_4]], %[[VAL_13]] : index
+// CHECK:               %[[VAL_15:.*]] = arith.addi %[[VAL_14]], %[[VAL_7]] : index
+// CHECK:               memref.store %[[VAL_12]], %[[VAL_0]]{{\[}}%[[VAL_15]]] : memref<30xf32>
+// CHECK:             }
+// CHECK:           }
+// CHECK:           return
+// CHECK:         }
+
+module {
+  memref.global "private" constant @const_1_30 : memref<2xi64> = dense<[1, 30]>
+  func.func @main(%arg0: memref<30x1xf32>) {
+    %c30 = arith.constant 30 : index
+    %c1 = arith.constant 1 : index
+    %c2 = arith.constant 2 : index
+    %c0 = arith.constant 0 : index
+    %alloc = memref.alloc() : memref<2x5x3xf32>
+    %0 = memref.get_global @const_1_30 : memref<2xi64>
+    %reshape = memref.reshape %alloc(%0) : (memref<2x5x3xf32>, memref<2xi64>) -> memref<1x30xf32>
+    scf.for %arg2 = %c0 to %c2 step %c1 {
+      scf.for %arg3 = %c0 to %c30 step %c1 {
+        %4 = memref.load %reshape[%c0, %arg3] : memref<1x30xf32>
+        memref.store %4, %arg0[%c0, %arg2] : memref<30x1xf32>
+      }
+    }
+    return
+  }
+}