[SimplifyMemRefAccess] start of this pass; [Analysis] update getLoadStoresMap method in utils

include/Analysis/Utils.h

@@ -88,9 +88,9 @@ public:
 //===----------------------------------------------------------------------===//
 
 // For storing all affine memory access operations (including AffineLoadOp and
-// AffineStoreOp) indexed by the array (ArrayOp).
+// AffineStoreOp) indexed by the corresponding memref.
 using LoadStores = SmallVector<Operation *, 16>;
-using LoadStoresMap = DenseMap<Operation *, LoadStores>;
+using LoadStoresMap = DenseMap<Value, LoadStores>;
 
 // Check if the lhsOp and rhsOp is at the same scheduling level. In this check,
 // AffineIfOp is transparent.
@@ -101,7 +101,9 @@ Optional<std::pair<Operation *, Operation *>> checkSameLevel(Operation *lhsOp,
 // level with dstOp's any parent loop.
 Operation *getSameLevelDstOp(Operation *srcOp, Operation *dstOp);
 
-/// Get the definition ArrayOp given any memory access operation.
+/// Get the definition ArrayOp given any memref or memory access operation.
+hlscpp::ArrayOp getArrayOp(Value memref);
+
 hlscpp::ArrayOp getArrayOp(Operation *op);
 
 /// Collect all load and store operations in the block.

include/Transforms/Passes.h

@@ -31,7 +31,10 @@ std::unique_ptr<Pass> createLegalizeDataflowPass();
 
 /// Bufferization passes.
 std::unique_ptr<Pass> createHLSKernelBufferizePass();
+
+/// MemRef Optimization Passes.
 std::unique_ptr<Pass> createStoreOpForwardPass();
+std::unique_ptr<Pass> createSimplifyMemRefAccessPass();
 
 void registerTransformsPasses();
 

include/Transforms/Passes.td

@@ -136,6 +136,10 @@ def HLSKernelBufferize : Pass<"hlskernel-bufferize", "FuncOp"> {
   let constructor = "mlir::scalehls::createHLSKernelBufferizePass()";
 }
 
+//===----------------------------------------------------------------------===//
+// MemRef Optimization Passes
+//===----------------------------------------------------------------------===//
+
 def StoreOpForward : Pass<"store-op-forward", "FuncOp"> {
   let summary = "Forward store to load, including conditional stores";
   let description = [{
@@ -146,4 +150,14 @@ def StoreOpForward : Pass<"store-op-forward", "FuncOp"> {
   let constructor = "mlir::scalehls::createStoreOpForwardPass()";
 }
 
+def SimplifyMemRefAccess : Pass<"simplify-memref-access", "FuncOp"> {
+  let summary = "Forward store to load, including conditional stores";
+  let description = [{
+    This simplify-memref-access pass will eliminate the known redundant load and
+    store operations for simplifying the memref access.
+  }];
+
+  let constructor = "mlir::scalehls::createSimplifyMemRefAccessPass()";
+}
+
 #endif // SCALEHLS_TRANSFORMS_PASSES_TD

lib/Analysis/QoREstimation.cpp

@@ -329,7 +329,7 @@ unsigned HLSCppEstimator::getResMinII(LoadStoresMap &map) {
   unsigned II = 1;
 
   for (auto &pair : map) {
-    auto arrayOp = pair.first;
+    auto arrayOp = getArrayOp(pair.first);
     // Partition number should at least be 1.
     auto partitionNum =
         max((unsigned)1, getUIntAttrValue(arrayOp, "partition_num"));

lib/Analysis/Utils.cpp

@@ -85,9 +85,11 @@ Operation *scalehls::getSameLevelDstOp(Operation *srcOp, Operation *dstOp) {
   return nullptr;
 }
 
-/// Get the definition ArrayOp given any memory access operation.
-hlscpp::ArrayOp scalehls::getArrayOp(Operation *op) {
-  auto defOp = MemRefAccess(op).memref.getDefiningOp();
+/// Get the definition ArrayOp given any memref or memory access operation.
+hlscpp::ArrayOp scalehls::getArrayOp(Value memref) {
+  assert(memref.getType().isa<MemRefType>() && "isn't a MemRef type value");
+
+  auto defOp = memref.getDefiningOp();
   assert(defOp && "MemRef is block argument");
 
   auto arrayOp = dyn_cast<hlscpp::ArrayOp>(defOp);
@@ -96,11 +98,15 @@ hlscpp::ArrayOp scalehls::getArrayOp(Operation *op) {
   return arrayOp;
 }
 
+hlscpp::ArrayOp scalehls::getArrayOp(Operation *op) {
+  return getArrayOp(MemRefAccess(op).memref);
+}
+
 /// Collect all load and store operations in the block.
 void scalehls::getLoadStoresMap(Block &block, LoadStoresMap &map) {
   for (auto &op : block) {
     if (isa<AffineReadOpInterface, AffineWriteOpInterface>(op))
-      map[getArrayOp(&op)].push_back(&op);
+      map[MemRefAccess(&op).memref].push_back(&op);
     else if (op.getNumRegions()) {
       for (auto &region : op.getRegions())
         for (auto &block : region)

lib/Transforms/ArrayPartition.cpp

@@ -34,7 +34,7 @@ static mlir::AffineForOp getPipelineLoop(mlir::AffineForOp root) {
 template <typename OpType>
 static void applyArrayPartition(LoadStoresMap &map, OpBuilder &builder) {
   for (auto pair : map) {
-    auto arrayOp = cast<ArrayOp>(pair.first);
+    auto arrayOp = getArrayOp(pair.first);
     auto arrayShape = arrayOp.getShapedType().getShape();
     auto arrayAccesses = pair.second;
 
@@ -120,14 +120,12 @@ void ArrayPartition::runOnOperation() {
       // Collect memory access information.
       LoadStoresMap loadMap;
       outermost.walk([&](mlir::AffineLoadOp loadOp) {
-        auto arrayOp = loadOp.getMemRef().getDefiningOp();
-        loadMap[arrayOp].push_back(loadOp);
+        loadMap[loadOp.getMemRef()].push_back(loadOp);
       });
 
       LoadStoresMap storeMap;
       outermost.walk([&](mlir::AffineStoreOp storeOp) {
-        auto arrayOp = storeOp.getMemRef().getDefiningOp();
-        storeMap[arrayOp].push_back(storeOp);
+        storeMap[storeOp.getMemRef()].push_back(storeOp);
       });
 
       // Apply array partition pragma.

lib/Transforms/SimplifyMemRefAccess.cpp

@@ -0,0 +1,31 @@
+//===------------------------------------------------------------*- C++ -*-===//
+//
+//===----------------------------------------------------------------------===//
+
+#include "Analysis/Utils.h"
+#include "Transforms/Passes.h"
+#include "mlir/Analysis/AffineAnalysis.h"
+#include "mlir/Analysis/Utils.h"
+#include "mlir/Dialect/Affine/IR/AffineOps.h"
+#include "mlir/Dialect/StandardOps/IR/Ops.h"
+#include "mlir/IR/Dominance.h"
+#include "mlir/IR/IntegerSet.h"
+#include "llvm/ADT/SmallPtrSet.h"
+#include <algorithm>
+
+using namespace mlir;
+using namespace scalehls;
+
+namespace {
+struct SimplifyMemRefAccess
+    : public SimplifyMemRefAccessBase<SimplifyMemRefAccess> {
+  void runOnOperation() override;
+};
+
+} // end anonymous namespace
+
+void SimplifyMemRefAccess::runOnOperation() {}
+
+std::unique_ptr<Pass> scalehls::createSimplifyMemRefAccessPass() {
+  return std::make_unique<SimplifyMemRefAccess>();
+}

test/Transforms/test_simplify_memref_access.mlir

@@ -0,0 +1,6 @@
+// RUN: scalehls-opt -simplify-memref-access %s | FileCheck %s
+
+// CHECK-LABEL: func @test_for
+func @test_for() {
+  return
+}