[Transforms] lower memref.copy (#507)

* [Transforms] lower memref.copy

* address review issues

Author: Dewei-Wang-sh

include/imex/Transforms/Passes.h

@@ -25,6 +25,7 @@ std::unique_ptr<mlir::Pass> createInsertGPUAllocsPass();
 std::unique_ptr<mlir::Pass> createSetSPIRVCapabilitiesPass();
 std::unique_ptr<mlir::Pass> createSetSPIRVAbiAttributePass();
 std::unique_ptr<mlir::Pass> createAddOuterParallelLoopPass();
+std::unique_ptr<mlir::Pass> createLowerMemRefCopyPass();
 
 #define GEN_PASS_DECL
 #include "imex/Transforms/Passes.h.inc"

include/imex/Transforms/Passes.td

@@ -75,4 +75,19 @@ def AddOuterParallelLoop : Pass<"imex-add-outer-parallel-loop", "::mlir::func::F
     ];
 }
 
+def LowerMemRefCopy : Pass<"imex-lower-memref-copy", "::mlir::func::FuncOp"> {
+  let summary = "lower memref.copy to linalg.generic";
+  let description = [{
+    This Pass transforms memref.copy to linalg.generic with identity index map and
+    parallel iterator. If satisfied, this pass also does memref.copy canonicalization.
+
+    This pass is supposed to work after bufferization and before linalg-lowering.
+  }];
+  let constructor = "imex::createLowerMemRefCopyPass()";
+  let dependentDialects = [
+    "::mlir::linalg::LinalgDialect",
+    "::mlir::memref::MemRefDialect"
+    ];
+}
+
 #endif // _IMEX_TRANSFORMS_PASSES_TD_INCLUDED_

lib/Transforms/CMakeLists.txt

@@ -4,6 +4,7 @@ add_mlir_library(IMEXTransforms
   SetSPIRVCapabilities.cpp
   SetSPIRVAbiAttribute.cpp
   AddOuterParallelLoop.cpp
+  LowerMemRefCopy.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${PROJECT_SOURCE_DIR}/imex/Transforms

lib/Transforms/LowerMemRefCopy.cpp

@@ -0,0 +1,89 @@
+//===- LowerMemRefCopy.cpp - lower memref.copy pass --------*- C++ -*-===//
+//
+// Copyright 2022 Intel Corporation
+// Part of the IMEX Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+///
+/// \file
+/// This pass lowers memref copyOp to linalg generic operations and enables
+/// simple memref copyOp canonicalization
+///
+//===----------------------------------------------------------------------===//
+
+#include "imex/Transforms/Passes.h"
+
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/Dialect/Linalg/IR/Linalg.h"
+#include "mlir/Dialect/Linalg/Utils/Utils.h"
+#include "mlir/Dialect/MemRef/IR/MemRef.h"
+#include "mlir/IR/Dominance.h"
+#include "mlir/Pass/Pass.h"
+#include "mlir/Support/LogicalResult.h"
+
+namespace imex {
+#define GEN_PASS_DEF_LOWERMEMREFCOPY
+#include "imex/Transforms/Passes.h.inc"
+} // namespace imex
+
+using namespace mlir;
+using namespace imex;
+
+namespace {
+struct LowerMemRefCopy
+    : public imex::impl::LowerMemRefCopyBase<LowerMemRefCopy> {
+  void runOnOperation() override {
+    auto &domInfo = getAnalysis<DominanceInfo>();
+    auto func = getOperation();
+    // walk through memref.copy ops in the funcion body
+    WalkResult result =
+        func.walk<WalkOrder::PreOrder>([&](memref::CopyOp op) -> WalkResult {
+          if (op->getParentOp() != func)
+            return WalkResult::skip();
+          auto src = op.getSource();
+          auto dst = op.getTarget();
+          // supposed to work on same memref type
+          auto srcType = src.getType().cast<MemRefType>();
+          auto dstType = dst.getType().cast<MemRefType>();
+          if (srcType != dstType)
+            return WalkResult::skip();
+          // supposed to work on memref.alloc
+          auto srcOp = src.getDefiningOp<memref::AllocOp>();
+          auto dstOp = dst.getDefiningOp<memref::AllocOp>();
+          if (!srcOp || !dstOp)
+            return WalkResult::skip();
+          // check use of src after this copyOp, being conservative
+          // FIXME: handle dealloc of src and dst
+          bool hasSubsequentUse = false;
+          for (auto user : src.getUsers()) {
+            if (isa<memref::DeallocOp>(user)) {
+              continue;
+            }
+            if (domInfo.properlyDominates(op, user)) {
+              hasSubsequentUse = true;
+              break;
+            }
+          }
+
+          // replace copy with linalg.generic
+          if (hasSubsequentUse) {
+            OpBuilder builder(op);
+            linalg::makeMemRefCopyOp(builder, op.getLoc(), src, dst);
+          } else {
+            // coalesce buffer
+            dst.replaceAllUsesWith(src);
+          }
+          op.erase();
+          return WalkResult::advance();
+        });
+  }
+};
+} // namespace
+
+namespace imex {
+std::unique_ptr<mlir::Pass> createLowerMemRefCopyPass() {
+  return std::make_unique<LowerMemRefCopy>();
+}
+} // namespace imex

lib/Transforms/PassDetail.h

@@ -30,6 +30,10 @@ class SPIRVDialect;
 namespace scf {
 class SCFDialect;
 }
+
+namespace linalg {
+class LinalgDialect;
+}
 } // end namespace mlir
 
 #include "mlir/Dialect/Affine/IR/AffineOps.h"

test/Transforms/lower-memref-copy.mlir

@@ -0,0 +1,39 @@
+// RUN: imex-opt -imex-lower-memref-copy -allow-unregistered-dialect %s | FileCheck %s
+#map = affine_map<(d0, d1) -> (d0, d1)>
+module {
+  func.func @copy_with_later_use(%arg0: memref<10x20xf32>) -> memref<10x20xf32> {
+    %cst = arith.constant 0.000000e+00 : f32
+    %alloc = memref.alloc() {alignment = 128 : i64} : memref<10x20xf32>
+    %alloc_0 = memref.alloc() {alignment = 128 : i64} : memref<10x20xf32>
+    linalg.fill ins(%cst : f32) outs(%alloc_0 : memref<10x20xf32>)
+    %alloc_1 = memref.alloc() {alignment = 128 : i64} : memref<10x20xf32>
+    memref.copy %alloc_0, %alloc_1 : memref<10x20xf32> to memref<10x20xf32>
+    "some_use" (%alloc_0) {} : (memref<10x20xf32>) -> ()
+    // CHECK-LABEL: func @copy_with_later_use
+    // CHECK:       linalg.generic
+    // CHECK:       linalg.generic
+    linalg.generic {indexing_maps = [#map, #map], iterator_types = ["parallel", "parallel"]} ins(%arg0 : memref<10x20xf32>) outs(%alloc_1 : memref<10x20xf32>) attrs =  {iterator_ranges = [10, 20]} {
+    ^bb0(%in: f32, %out: f32):
+      %0 = arith.addf %out, %in : f32
+      linalg.yield %0 : f32
+    }
+    return %alloc_1 : memref<10x20xf32>
+  }
+  func.func @copy_without_later_use(%arg0: memref<10x20xf32>) -> memref<10x20xf32> {
+    %cst = arith.constant 0.000000e+00 : f32
+    %alloc = memref.alloc() {alignment = 128 : i64} : memref<10x20xf32>
+    %alloc_0 = memref.alloc() {alignment = 128 : i64} : memref<10x20xf32>
+    linalg.fill ins(%cst : f32) outs(%alloc_0 : memref<10x20xf32>)
+    %alloc_1 = memref.alloc() {alignment = 128 : i64} : memref<10x20xf32>
+    memref.copy %alloc_0, %alloc_1 : memref<10x20xf32> to memref<10x20xf32>
+    // CHECK-LABEL: func @copy_without_later_use
+    // CHECK:       linalg.generic
+    // CHECK-NOT:   linalg.generic
+    linalg.generic {indexing_maps = [#map, #map], iterator_types = ["parallel", "parallel"]} ins(%arg0 : memref<10x20xf32>) outs(%alloc_1 : memref<10x20xf32>) attrs =  {iterator_ranges = [10, 20]} {
+    ^bb0(%in: f32, %out: f32):
+      %0 = arith.addf %out, %in : f32
+      linalg.yield %0 : f32
+    }
+    return %alloc_1 : memref<10x20xf32>
+  }
+}