Pretransformations to use destination passing style, before upstream …

…bufferization. (#7582)

The upstream ComprehensiveBufferization pass prefers a more
destination-passing style of code, i.e. it has a way to know before
hand what the output buffer to use is for in-place
operations. Specifically, the use of init_tensor to conjure a tensor
(even though it might be fully over-written before its values are
used) results in an allocation. It is easier to do a
pre-transformation that removes these init_tensor operations and
replaces them with a flow.dispatch.tensor.load of the output where
the final result goes into.
The current IREE bufferization handles this. This patch reuses the
relevant parts of IREE bufferization to effect this transformation.

iree/compiler/Codegen/Common/BUILD

@@ -32,7 +32,9 @@ gentbl_cc_library(
 cc_library(
     name = "Common",
     srcs = [
+        "BufferizationAnalysis.cpp",
         "CleanupBufferAllocViewPass.cpp",
+        "ConvertToDestinationPassingStylePass.cpp",
         "DemoteF32ToF16.cpp",
         "FlattenMemRefSubspanPass.cpp",
         "FoldAffineMinInDistributedLoops.cpp",
@@ -46,6 +48,9 @@ cc_library(
         "VectorizeConv.cpp",
         "VectorizeMMT4d.cpp",
     ],
+    hdrs = [
+        "BufferizationAnalysis.h",
+    ],
     deps = [
         "//iree/compiler/Codegen:PassHeaders",
         "//iree/compiler/Codegen/Common:FoldTensorExtractOpIncGen",

iree/compiler/Codegen/Common/BufferizationAnalysis.h

@@ -0,0 +1,96 @@
+// Copyright 2021 The IREE Authors
+//
+// Licensed 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
+
+//===- BufferizationAnalysis.h - Pre bufferization analysis ---------------===//
+//
+// Analysis to group together tensors within a dispatch region into an
+// equivalance class such that all members of a set can be mapped to the same
+// memory region.
+//
+//===----------------------------------------------------------------------===//
+#ifndef IREE_COMPILER_CODEGEN_COMMON_BUFFERIZATIONANALYSIS_H_
+#define IREE_COMPILER_CODEGEN_COMMON_BUFFERIZATIONANALYSIS_H_
+
+#include "llvm/ADT/EquivalenceClasses.h"
+#include "llvm/Support/Debug.h"
+#include "mlir/IR/BuiltinOps.h"
+#include "mlir/IR/Value.h"
+
+namespace mlir {
+namespace iree_compiler {
+
+/// Class that tracks the equivalence relationship between tensors. Its a
+/// light-weight wrapper around `llvm::EquivalenceClasses` to account for
+/// `Value` not directly supported as a value type by this class.
+class BufferizationPlan {
+ public:
+  llvm::EquivalenceClasses<void *>::iterator findValue(Value v) const {
+    return mappedTensors.findValue(getPointer(v));
+  }
+
+  llvm::EquivalenceClasses<void *>::iterator end() const {
+    return mappedTensors.end();
+  }
+
+  SmallVector<Value> getTensorsMappedToSameSet(Value v) const {
+    SmallVector<Value> tensors;
+    for (auto it = mappedTensors.findLeader(getPointer(v)),
+              ie = mappedTensors.member_end();
+         it != ie; ++it) {
+      tensors.push_back(getValue(*it));
+    }
+    return tensors;
+  }
+
+  bool isEquivalent(Value v1, Value v2) const {
+    return mappedTensors.isEquivalent(getPointer(v1), getPointer(v2));
+  }
+
+  void insert(Value v) { mappedTensors.insert(getPointer(v)); }
+
+  void unionSets(Value v1, Value v2) {
+    mappedTensors.unionSets(getPointer(v1), getPointer(v2));
+  }
+
+  /// Sets the equivalance class that contains `v` as the set that contains the
+  /// result tensor of the dispatch region (i.e. a tensor that is the `value`
+  /// operand of a flow.dispatch.tensor.store` op). All operations in this
+  /// equivalence class can use the result buffer of the dispatch region to
+  /// compute their values in place.
+  void storeSet(Value v) { storeLeaders.insert(getLeaderValue(v)); }
+
+  /// Queries if the value `v` is in the same equivalence class as the result of
+  /// the dispatch region.
+  bool isInStoreSet(Value v) { return storeLeaders.count(getLeaderValue(v)); }
+
+  void dump();
+
+ private:
+  Value getLeaderValue(Value v1) const {
+    return getValue(mappedTensors.getLeaderValue(getPointer(v1)));
+  }
+
+  void *getPointer(Value v) const { return v.getAsOpaquePointer(); }
+
+  Value getValue(void *v) const { return Value::getFromOpaquePointer(v); }
+
+  llvm::EquivalenceClasses<void *> mappedTensors;
+
+  /// Leaders of the sets that contain the result tensor of the dispatch
+  /// region, i.e. a tensor that is the `value` operand of a
+  /// flow.dispatch.tensor.store` op
+  llvm::DenseSet<Value> storeLeaders;
+};
+
+/// Analysis the `tensor` values in `funcOp` and groups them together into
+/// equivalence classes such that each class contains tensors that can be mapped
+/// to the same buffer.
+LogicalResult createTensorEquivalenceClasses(FuncOp funcOp,
+                                             BufferizationPlan &plan);
+
+}  // namespace iree_compiler
+}  // namespace mlir
+#endif  // IREE_COMPILER_CODEGEN_COMMON_BUFFERIZATIONANALYSIS_H

iree/compiler/Codegen/Common/CMakeLists.txt

@@ -22,8 +22,12 @@ iree_tablegen_library(
 iree_cc_library(
   NAME
     Common
+  HDRS
+    "BufferizationAnalysis.h"
   SRCS
+    "BufferizationAnalysis.cpp"
     "CleanupBufferAllocViewPass.cpp"
+    "ConvertToDestinationPassingStylePass.cpp"
     "DemoteF32ToF16.cpp"
     "FlattenMemRefSubspanPass.cpp"
     "FoldAffineMinInDistributedLoops.cpp"