[Mlir-commits] [mlir] 3463d98 - [mlir][Linalg] Add a hoistViewAllocOps helper function
Nicolas Vasilache
llvmlistbot at llvm.org
Thu Jun 4 16:03:23 PDT 2020
Author: Nicolas Vasilache
Date: 2020-06-04T18:59:03-04:00
New Revision: 3463d9835b02a2c1b22fabadb4ba1a175fea6dad
URL: https://github.com/llvm/llvm-project/commit/3463d9835b02a2c1b22fabadb4ba1a175fea6dad
DIFF: https://github.com/llvm/llvm-project/commit/3463d9835b02a2c1b22fabadb4ba1a175fea6dad.diff
LOG: [mlir][Linalg] Add a hoistViewAllocOps helper function
This revision adds a helper function to hoist alloc/dealloc pairs and
alloca op out of immediately enclosing scf::ForOp if both conditions are true:
1. all operands are defined outside the loop.
2. all uses are ViewLikeOp or DeallocOp.
This is now considered Linalg-specific and will be generalized on a per-need basis.
Differential Revision: https://reviews.llvm.org/D81152
Added:
mlir/include/mlir/Dialect/Linalg/Transforms/Hoisting.h
mlir/lib/Dialect/Linalg/Transforms/Hoisting.cpp
mlir/test/Dialect/Linalg/hoisting.mlir
mlir/test/lib/Transforms/TestLinalgHoisting.cpp
Modified:
mlir/lib/Dialect/Linalg/Transforms/CMakeLists.txt
mlir/test/lib/Transforms/CMakeLists.txt
mlir/tools/mlir-opt/mlir-opt.cpp
Removed:
################################################################################
diff --git a/mlir/include/mlir/Dialect/Linalg/Transforms/Hoisting.h b/mlir/include/mlir/Dialect/Linalg/Transforms/Hoisting.h
new file mode 100644
index 000000000000..283e68e8a76b
--- /dev/null
+++ b/mlir/include/mlir/Dialect/Linalg/Transforms/Hoisting.h
@@ -0,0 +1,27 @@
+//===- Hoisting.h - Linalg hoisting transformations -------------*- C++ -*-===//
+//
+// Part of the LLVM 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
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_DIALECT_LINALG_TRANSFORMS_HOISTING_H_
+#define MLIR_DIALECT_LINALG_TRANSFORMS_HOISTING_H_
+
+namespace mlir {
+class FuncOp;
+
+namespace linalg {
+
+/// Hoist alloc/dealloc pairs and alloca op out of immediately enclosing
+/// scf::ForOp if both conditions are true:
+/// 1. all operands are defined outside the loop.
+/// 2. all uses are ViewLikeOp or DeallocOp.
+// TODO: generalize on a per-need basis.
+void hoistViewAllocOps(FuncOp func);
+
+} // namespace linalg
+} // namespace mlir
+
+#endif // MLIR_DIALECT_LINALG_TRANSFORMS_HOISTING_H_
diff --git a/mlir/lib/Dialect/Linalg/Transforms/CMakeLists.txt b/mlir/lib/Dialect/Linalg/Transforms/CMakeLists.txt
index c87e3d4f15b6..4f85603ac148 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/CMakeLists.txt
+++ b/mlir/lib/Dialect/Linalg/Transforms/CMakeLists.txt
@@ -1,6 +1,7 @@
add_mlir_dialect_library(MLIRLinalgTransforms
DropUnitDims.cpp
Fusion.cpp
+ Hoisting.cpp
Interchange.cpp
Loops.cpp
Promotion.cpp
diff --git a/mlir/lib/Dialect/Linalg/Transforms/Hoisting.cpp b/mlir/lib/Dialect/Linalg/Transforms/Hoisting.cpp
new file mode 100644
index 000000000000..737e69021bc8
--- /dev/null
+++ b/mlir/lib/Dialect/Linalg/Transforms/Hoisting.cpp
@@ -0,0 +1,77 @@
+//===- Hoisting.cpp - Linalg hoisting transformations ---------------------===//
+//
+// Part of the LLVM 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
+//
+//===----------------------------------------------------------------------===//
+//
+// This file implements functions concerned with hoisting invariant operations
+// in the context of Linalg transformations.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/Linalg/Transforms/Hoisting.h"
+#include "mlir/Dialect/Linalg/IR/LinalgOps.h"
+#include "mlir/Dialect/SCF/SCF.h"
+#include "mlir/Dialect/StandardOps/IR/Ops.h"
+#include "mlir/IR/Function.h"
+#include "llvm/ADT/StringRef.h"
+#include "llvm/Support/Debug.h"
+
+#define DEBUG_TYPE "linalg-hoisting"
+
+#define DBGS() (dbgs() << '[' << DEBUG_TYPE << "] ")
+
+using namespace mlir;
+using namespace mlir::linalg;
+
+using llvm::dbgs;
+
+void mlir::linalg::hoistViewAllocOps(FuncOp func) {
+ bool changed = true;
+ while (changed) {
+ changed = false;
+ func.walk([&changed](Operation *op) {
+ if (!isa<AllocOp>(op) && !isa<AllocaOp>(op) && !isa<DeallocOp>(op))
+ return;
+
+ LLVM_DEBUG(DBGS() << "Candidate for hoisting: " << *op << "\n");
+ auto loop = dyn_cast<scf::ForOp>(op->getParentOp());
+ LLVM_DEBUG(DBGS() << "Parent op: " << *op->getParentOp() << "\n");
+
+ // Only hoist out of immediately enclosing scf::ForOp.
+ if (!loop)
+ return;
+
+ // If any operand is defined inside the loop don't hoist.
+ if (llvm::any_of(op->getOperands(), [&](Value v) {
+ return !loop.isDefinedOutsideOfLoop(v);
+ }))
+ return;
+
+ LLVM_DEBUG(DBGS() << "All operands defined outside \n");
+
+ // If alloc has other uses than ViewLikeOp and DeallocOp don't hoist.
+ Value v;
+ if (op->getNumResults() > 0) {
+ assert(op->getNumResults() == 1 && "Unexpected multi-result alloc");
+ v = op->getResult(0);
+ }
+ if (v && !llvm::all_of(v.getUses(), [&](OpOperand &operand) {
+ return isa<ViewLikeOpInterface>(operand.getOwner()) ||
+ isa<DeallocOp>(operand.getOwner());
+ })) {
+ LLVM_DEBUG(DBGS() << "Found non view-like or dealloc use: bail\n");
+ return;
+ }
+
+ // Move AllocOp before the loop.
+ if (isa<AllocOp>(op) || isa<AllocaOp>(op))
+ loop.moveOutOfLoop({op});
+ else // Move DeallocOp outside of the loop.
+ op->moveAfter(loop);
+ changed = true;
+ });
+ }
+}
diff --git a/mlir/test/Dialect/Linalg/hoisting.mlir b/mlir/test/Dialect/Linalg/hoisting.mlir
new file mode 100644
index 000000000000..db686655d64d
--- /dev/null
+++ b/mlir/test/Dialect/Linalg/hoisting.mlir
@@ -0,0 +1,82 @@
+// RUN: mlir-opt %s -test-linalg-hoisting=test-hoist-view-allocs | FileCheck %s
+
+// CHECK-LABEL: func @hoist(
+// CHECK-SAME: %[[VAL:[a-zA-Z0-9]*]]: index,
+// CHECK-SAME: %[[LB:[a-zA-Z0-9]*]]: index,
+// CHECK-SAME: %[[UB:[a-zA-Z0-9]*]]: index,
+// CHECK-SAME: %[[STEP:[a-zA-Z0-9]*]]: index,
+// CHECK-SAME: %[[CMP:[a-zA-Z0-9]*]]: i1
+func @hoist(%val: index, %lb : index, %ub : index, %step: index, %cmp: i1) {
+// CHECK-DAG: alloca(%[[VAL]]) : memref<?xi8>
+// CHECK-DAG: %[[A0:.*]] = alloc(%[[VAL]]) : memref<?xi8>
+// CHECK: scf.for %[[I:.*]] = %[[LB]] to %[[UB]] step %[[STEP]] {
+// CHECK: alloca(%[[I]]) : memref<?xi8>
+// CHECK: %[[A1:.*]] = alloc(%[[I]]) : memref<?xi8>
+// CHECK: scf.for %[[J:.*]] = %[[LB]] to %[[UB]] step %[[STEP]] {
+// CHECK-DAG: alloca(%[[J]]) : memref<?xi8>
+// CHECK-DAG: %[[A2:.*]] = alloc(%[[J]]) : memref<?xi8>
+// CHECK: scf.for %[[K:.*]] = %[[LB]] to %[[UB]] step %[[STEP]] {
+ scf.for %i = %lb to %ub step %step {
+ scf.for %j = %lb to %ub step %step {
+ scf.for %k = %lb to %ub step %step {
+ // Hoist allocs / deallocs outermost, keep view/subview below k.
+ %sa0 = alloca(%val) : memref<? x i8>
+ %a0 = alloc(%val) : memref<? x i8>
+// CHECK: std.view %[[A0]][%[[LB]]][] : memref<?xi8> to memref<16xf32>
+// CHECK: subview %[[A0]][0] [42] [1] : memref<?xi8> to memref<42xi8>
+ %v0 = view %a0[%lb][] : memref<? x i8> to memref<16 x f32>
+ %sv0 = subview %a0[0][42][1] : memref<? x i8> to memref<42 x i8>
+ dealloc %a0 : memref<? x i8>
+
+ // Hoist below i.
+ %sa1 = alloca(%i) : memref<? x i8>
+ %a1 = alloc(%i) : memref<? x i8>
+ dealloc %a1 : memref<? x i8>
+
+ // Hoist below j.
+ %sa2 = alloca(%j) : memref<? x i8>
+ %a2 = alloc(%j) : memref<? x i8>
+ dealloc %a2 : memref<? x i8>
+
+ // Don't hoist since k innermost.
+// CHECK: alloca(%[[K]]) : memref<?xi8>
+// CHECK: %[[A3:.*]] = alloc(%[[K]]) : memref<?xi8>
+// CHECK: dealloc %[[A3]] : memref<?xi8>
+ %sa3 = alloca(%k) : memref<? x i8>
+ %a3 = alloc(%k) : memref<? x i8>
+ dealloc %a3 : memref<? x i8>
+
+ // No hoisting due to control flow.
+// CHECK: scf.if %[[CMP]] {
+// CHECK: alloca(%[[VAL]]) : memref<?xi8>
+// CHECK: %[[A4:.*]] = alloc(%[[VAL]]) : memref<?xi8>
+// CHECK: dealloc %[[A4]] : memref<?xi8>
+ scf.if %cmp {
+ %sa4 = alloca(%val) : memref<? x i8>
+ %a4 = alloc(%val) : memref<? x i8>
+ dealloc %a4 : memref<? x i8>
+ }
+
+ // No hoisting due to load/store.
+// CHECK: %[[SA5:.*]] = alloca(%[[VAL]]) : memref<?xi8>
+// CHECK: %[[A5:.*]] = alloc(%[[VAL]]) : memref<?xi8>
+// CHECK: load %[[A5]][%[[LB]]] : memref<?xi8>
+// CHECK: store %{{.*}}, %[[SA5]][%[[LB]]] : memref<?xi8>
+// CHECK: dealloc %[[A5]] : memref<?xi8>
+ %sa5 = alloca(%val) : memref<? x i8>
+ %a5 = alloc(%val) : memref<? x i8>
+ %v5 = load %a5[%lb] : memref<? x i8>
+ store %v5, %sa5[%lb] : memref<? x i8>
+ dealloc %a5 : memref<? x i8>
+
+ }
+ }
+ }
+// CHECK: }
+// CHECK: dealloc %[[A2]] : memref<?xi8>
+// CHECK: }
+// CHECK: dealloc %[[A1]] : memref<?xi8>
+// CHECK: }
+// CHECK: dealloc %[[A0]] : memref<?xi8>
+ return
+}
diff --git a/mlir/test/lib/Transforms/CMakeLists.txt b/mlir/test/lib/Transforms/CMakeLists.txt
index 3f2befefe704..044270276c11 100644
--- a/mlir/test/lib/Transforms/CMakeLists.txt
+++ b/mlir/test/lib/Transforms/CMakeLists.txt
@@ -11,6 +11,7 @@ add_mlir_library(MLIRTestTransforms
TestGpuMemoryPromotion.cpp
TestGpuParallelLoopMapping.cpp
TestInlining.cpp
+ TestLinalgHoisting.cpp
TestLinalgTransforms.cpp
TestLiveness.cpp
TestLoopMapping.cpp
diff --git a/mlir/test/lib/Transforms/TestLinalgHoisting.cpp b/mlir/test/lib/Transforms/TestLinalgHoisting.cpp
new file mode 100644
index 000000000000..2fb5c091aed4
--- /dev/null
+++ b/mlir/test/lib/Transforms/TestLinalgHoisting.cpp
@@ -0,0 +1,47 @@
+//===- TestLinalgHoisting.cpp - Test Linalg hoisting functions ------------===//
+//
+// Part of the LLVM 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
+//
+//===----------------------------------------------------------------------===//
+//
+// This file implements logic for testing Linalg hoisting functions.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/Linalg/IR/LinalgOps.h"
+#include "mlir/Dialect/Linalg/Transforms/Hoisting.h"
+#include "mlir/Pass/Pass.h"
+
+using namespace mlir;
+using namespace mlir::linalg;
+
+namespace {
+struct TestLinalgHoisting
+ : public PassWrapper<TestLinalgHoisting, FunctionPass> {
+ TestLinalgHoisting() = default;
+ TestLinalgHoisting(const TestLinalgHoisting &pass) {}
+
+ void runOnFunction() override;
+
+ Option<bool> testHoistViewAllocs{
+ *this, "test-hoist-view-allocs",
+ llvm::cl::desc("Test hoisting alloc used by view"),
+ llvm::cl::init(false)};
+};
+} // end anonymous namespace
+
+void TestLinalgHoisting::runOnFunction() {
+ if (testHoistViewAllocs) {
+ hoistViewAllocOps(getFunction());
+ return;
+ }
+}
+
+namespace mlir {
+void registerTestLinalgHoisting() {
+ PassRegistration<TestLinalgHoisting> testTestLinalgHoistingPass(
+ "test-linalg-hoisting", "Test Linalg hoisting functions.");
+}
+} // namespace mlir
diff --git a/mlir/tools/mlir-opt/mlir-opt.cpp b/mlir/tools/mlir-opt/mlir-opt.cpp
index 165d6ad159ac..2d286e112f99 100644
--- a/mlir/tools/mlir-opt/mlir-opt.cpp
+++ b/mlir/tools/mlir-opt/mlir-opt.cpp
@@ -50,6 +50,7 @@ void registerTestConvertGPUKernelToHsacoPass();
void registerTestDominancePass();
void registerTestFunc();
void registerTestGpuMemoryPromotionPass();
+void registerTestLinalgHoisting();
void registerTestLinalgTransforms();
void registerTestLivenessPass();
void registerTestLoopFusion();
@@ -121,6 +122,7 @@ void registerTestPasses() {
registerTestDominancePass();
registerTestFunc();
registerTestGpuMemoryPromotionPass();
+ registerTestLinalgHoisting();
registerTestLinalgTransforms();
registerTestLivenessPass();
registerTestLoopFusion();
More information about the Mlir-commits
mailing list