diff --git a/mlir/include/mlir/Dialect/Linalg/Transforms/Hoisting.h b/mlir/include/mlir/Dialect/Linalg/Transforms/Hoisting.h new file mode 100644 --- /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 --- 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 --- /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(op) && !isa(op) && !isa(op)) + return; + + LLVM_DEBUG(DBGS() << "Candidate for hoisting: " << *op << "\n"); + auto loop = dyn_cast(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(operand.getOwner()) || + isa(operand.getOwner()); + })) { + LLVM_DEBUG(DBGS() << "Found non view-like or dealloc use: bail\n"); + return; + } + + // Move AllocOp before the loop. + if (isa(op) || isa(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 --- /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 +// CHECK-DAG: %[[A0:.*]] = alloc(%[[VAL]]) : memref +// CHECK: scf.for %[[I:.*]] = %[[LB]] to %[[UB]] step %[[STEP]] { +// CHECK: alloca(%[[I]]) : memref +// CHECK: %[[A1:.*]] = alloc(%[[I]]) : memref +// CHECK: scf.for %[[J:.*]] = %[[LB]] to %[[UB]] step %[[STEP]] { +// CHECK-DAG: alloca(%[[J]]) : memref +// CHECK-DAG: %[[A2:.*]] = alloc(%[[J]]) : memref +// 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 + %a0 = alloc(%val) : memref +// CHECK: std.view %[[A0]][%[[LB]]][] : memref to memref<16xf32> +// CHECK: subview %[[A0]][0] [42] [1] : memref to memref<42xi8> + %v0 = view %a0[%lb][] : memref to memref<16 x f32> + %sv0 = subview %a0[0][42][1] : memref to memref<42 x i8> + dealloc %a0 : memref + + // Hoist below i. + %sa1 = alloca(%i) : memref + %a1 = alloc(%i) : memref + dealloc %a1 : memref + + // Hoist below j. + %sa2 = alloca(%j) : memref + %a2 = alloc(%j) : memref + dealloc %a2 : memref + + // Don't hoist since k innermost. +// CHECK: alloca(%[[K]]) : memref +// CHECK: %[[A3:.*]] = alloc(%[[K]]) : memref +// CHECK: dealloc %[[A3]] : memref + %sa3 = alloca(%k) : memref + %a3 = alloc(%k) : memref + dealloc %a3 : memref + + // No hoisting due to control flow. +// CHECK: scf.if %[[CMP]] { +// CHECK: alloca(%[[VAL]]) : memref +// CHECK: %[[A4:.*]] = alloc(%[[VAL]]) : memref +// CHECK: dealloc %[[A4]] : memref + scf.if %cmp { + %sa4 = alloca(%val) : memref + %a4 = alloc(%val) : memref + dealloc %a4 : memref + } + + // No hoisting due to load/store. +// CHECK: %[[SA5:.*]] = alloca(%[[VAL]]) : memref +// CHECK: %[[A5:.*]] = alloc(%[[VAL]]) : memref +// CHECK: load %[[A5]][%[[LB]]] : memref +// CHECK: store %{{.*}}, %[[SA5]][%[[LB]]] : memref +// CHECK: dealloc %[[A5]] : memref + %sa5 = alloca(%val) : memref + %a5 = alloc(%val) : memref + %v5 = load %a5[%lb] : memref + store %v5, %sa5[%lb] : memref + dealloc %a5 : memref + + } + } + } +// CHECK: } +// CHECK: dealloc %[[A2]] : memref +// CHECK: } +// CHECK: dealloc %[[A1]] : memref +// CHECK: } +// CHECK: dealloc %[[A0]] : memref + return +} diff --git a/mlir/test/lib/Transforms/CMakeLists.txt b/mlir/test/lib/Transforms/CMakeLists.txt --- a/mlir/test/lib/Transforms/CMakeLists.txt +++ b/mlir/test/lib/Transforms/CMakeLists.txt @@ -11,6 +11,7 @@ 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 --- /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() = default; + TestLinalgHoisting(const TestLinalgHoisting &pass) {} + + void runOnFunction() override; + + Option 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 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 --- a/mlir/tools/mlir-opt/mlir-opt.cpp +++ b/mlir/tools/mlir-opt/mlir-opt.cpp @@ -50,6 +50,7 @@ void registerTestDominancePass(); void registerTestFunc(); void registerTestGpuMemoryPromotionPass(); +void registerTestLinalgHoisting(); void registerTestLinalgTransforms(); void registerTestLivenessPass(); void registerTestLoopFusion(); @@ -121,6 +122,7 @@ registerTestDominancePass(); registerTestFunc(); registerTestGpuMemoryPromotionPass(); + registerTestLinalgHoisting(); registerTestLinalgTransforms(); registerTestLivenessPass(); registerTestLoopFusion();