diff --git a/mlir/include/mlir/Analysis/DataFlow/LivenessAnalysis.h b/mlir/include/mlir/Analysis/DataFlow/LivenessAnalysis.h new file mode 100644 --- /dev/null +++ b/mlir/include/mlir/Analysis/DataFlow/LivenessAnalysis.h @@ -0,0 +1,73 @@ +//===- LivenessAnalysis.h - Liveness analysis -------------------*- 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 +// +//===----------------------------------------------------------------------===// +// +// This file implements liveness analysis using the sparse backward data-flow +// analysis framework. Theoretically, liveness analysis assigns liveness to each +// (value, program point) pair in the program and it is thus a dense analysis. +// However, since values are immutable in MLIR, a sparse analysis, which will +// assign liveness to each value in the program, suffices here. +// +// Liveness analysis has many applications. It can be used to avoid the +// computation of extraneous operations that have no effect on the memory or the +// final output of a program. It can also be used to optimize register +// allocation. Both of these applications help reduce runtime. +// +//===----------------------------------------------------------------------===// + +#ifndef MLIR_ANALYSIS_DATAFLOW_LIVENESSANALYSIS_H +#define MLIR_ANALYSIS_DATAFLOW_LIVENESSANALYSIS_H + +#include +#include + +namespace mlir::dataflow { + +//===----------------------------------------------------------------------===// +// LivenessAnalysis +//===----------------------------------------------------------------------===// + +/// This lattice represents, for a given value, whether or not it is "live". +/// +/// A value is considered "live" iff it: +/// (a) has memory effects OR +/// (b) is used to compute a value of type (a) OR +/// (c) is returned by a public function OR +/// (d) is used to compute a value of type (c). +struct Liveness : public AbstractSparseLattice { + MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(Liveness) + using AbstractSparseLattice::AbstractSparseLattice; + + void print(raw_ostream &os) const override; + + ChangeResult markLive(); + + ChangeResult meet(const AbstractSparseLattice &other) override; + + // At the beginning of the analysis, everything is marked "not live" and as + // the analysis progresses, values are marked "live" if they are found to be + // live. + bool isLive = false; +}; + +/// An analysis that, by going backwards along the dataflow graph, annotates +/// each value with a boolean storing true iff it is "live". +class LivenessAnalysis : public SparseBackwardDataFlowAnalysis { +public: + using SparseBackwardDataFlowAnalysis::SparseBackwardDataFlowAnalysis; + + void visitOperation(Operation *op, ArrayRef operands, + ArrayRef results) override; + + void visitBranchOperand(OpOperand &operand) override; + + void setToExitState(Liveness *lattice) override; +}; + +} // end namespace mlir::dataflow + +#endif // MLIR_ANALYSIS_DATAFLOW_LIVENESSANALYSIS_H diff --git a/mlir/lib/Analysis/CMakeLists.txt b/mlir/lib/Analysis/CMakeLists.txt --- a/mlir/lib/Analysis/CMakeLists.txt +++ b/mlir/lib/Analysis/CMakeLists.txt @@ -13,6 +13,7 @@ DataFlow/DeadCodeAnalysis.cpp DataFlow/DenseAnalysis.cpp DataFlow/IntegerRangeAnalysis.cpp + DataFlow/LivenessAnalysis.cpp DataFlow/SparseAnalysis.cpp ) @@ -34,6 +35,7 @@ DataFlow/DeadCodeAnalysis.cpp DataFlow/DenseAnalysis.cpp DataFlow/IntegerRangeAnalysis.cpp + DataFlow/LivenessAnalysis.cpp DataFlow/SparseAnalysis.cpp ADDITIONAL_HEADER_DIRS diff --git a/mlir/lib/Analysis/DataFlow/LivenessAnalysis.cpp b/mlir/lib/Analysis/DataFlow/LivenessAnalysis.cpp new file mode 100644 --- /dev/null +++ b/mlir/lib/Analysis/DataFlow/LivenessAnalysis.cpp @@ -0,0 +1,107 @@ +//===- LivenessAnalysis.cpp - Liveness analysis ---------------------------===// +// +// 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 +// +//===----------------------------------------------------------------------===// + +#include + +#include +#include +#include +#include +#include +#include + +using namespace mlir; +using namespace mlir::dataflow; + +void Liveness::print(raw_ostream &os) const { + os << (isLive ? "live" : "not live"); +} + +ChangeResult Liveness::markLive() { + bool wasLive = isLive; + isLive = true; + return wasLive ? ChangeResult::NoChange : ChangeResult::Change; +} + +ChangeResult Liveness::meet(const AbstractSparseLattice &other) { + const auto *otherLiveness = reinterpret_cast(&other); + return otherLiveness->isLive ? markLive() : ChangeResult::NoChange; +} + +//===----------------------------------------------------------------------===// +// LivenessAnalysis +//===----------------------------------------------------------------------===// + +/// For every value, liveness analysis determines whether or not it is "live". +/// +/// A value is considered "live" iff it: +/// (a) has memory effects OR +/// (b) is used to compute a value of type (a) OR +/// (c) is returned by a public function OR +/// (d) is used to compute a value of type (c). + +void LivenessAnalysis::visitOperation(Operation *op, + ArrayRef operands, + ArrayRef results) { + // This marks values of type (a) liveness as "live". + if (!isMemoryEffectFree(op)) { + for (auto *operand : operands) + propagateIfChanged(operand, operand->markLive()); + } + + // This marks values of type (b) and (d) liveness as "live". + bool foundLiveResult = false; + for (const Liveness *r : results) { + if (r->isLive && !foundLiveResult) { + // It is assumed that each operand is used to compute each result of an + // op. Thus, if at least one result is live, each operand is live. + for (Liveness *operand : operands) + meet(operand, *r); + + // TODO(srisrivastava): Enhance this backward flow of liveness. One + // potential enhancement: If this op exists in a block which is in a + // region of a region-based control flow op, then mark the non-forwarded + // operands of that op as "live" because such operands are used to compute + // this live result. + + foundLiveResult = true; + } + addDependency(const_cast(r), op); + } +} + +void LivenessAnalysis::visitBranchOperand(OpOperand &operand) { + // The lattices of the non-forwarded branch operands don't get updated like + // the forwarded branch operands or the non-branch operands. Thus they need + // to be handled separately. This is where we handle them. The liveness flows + // backward (or, in other words, the lattices get updated) in such operands by + // visiting their corresponding branch op (with all its operands). + + // Avoid visiting the same operation through multiple operands. + if (operand.getOperandNumber() > 0) + return; + + Operation *branchOp = operand.getOwner(); + + SmallVector operandsLiveness; + for (const Value operand : branchOp->getOperands()) { + operandsLiveness.push_back(getLatticeElement(operand)); + } + + SmallVector resultsLiveness; + for (const Value result : branchOp->getResults()) { + resultsLiveness.push_back(getLatticeElement(result)); + } + + visitOperation(branchOp, operandsLiveness, resultsLiveness); +} + +void LivenessAnalysis::setToExitState(Liveness *lattice) { + // This marks values of type (c) liveness as "live". + lattice->markLive(); +} diff --git a/mlir/test/Analysis/DataFlow/test-liveness-analysis.mlir b/mlir/test/Analysis/DataFlow/test-liveness-analysis.mlir new file mode 100644 --- /dev/null +++ b/mlir/test/Analysis/DataFlow/test-liveness-analysis.mlir @@ -0,0 +1,195 @@ +// RUN: mlir-opt -split-input-file -test-liveness-analysis %s 2>&1 | FileCheck %s + +// Positive test: Type (a) "has memory effects" +// zero is live because it is stored in memory. +// CHECK-LABEL: test_tag: zero: +// CHECK-NEXT: result #0: live +func.func @test_type_a(%arg0: memref) { + %c0_i32 = arith.constant {tag = "zero"} 0 : i32 + memref.store %c0_i32, %arg0[] : memref + return +} + +// ----- + + +func.func private @private0(%0 : i32) -> i32 { + %1 = arith.addi %0, %0 {tag = "in_private0"} : i32 + func.return %1 : i32 +} + +// Positive test: Type (b) "is used to compute a value of type (a)" +// zero, ten, and one are live because they are used to decide the number of +// times the `for` loop executes, which in turn decides the value stored in +// memory. +// Note that if `visitBranchOperand()` was left empty, they would have been +// incorrectly marked as "not live". +// in_private0 and x are also live because they decide the value stored in +// memory. +// CHECK-LABEL: test_tag: in_private0: +// CHECK-NEXT: operand #0: live +// CHECK-NEXT: operand #1: live +// CHECK-NEXT: result #0: live +// CHECK-LABEL: test_tag: zero: +// CHECK-NEXT: result #0: live +// CHECK-LABEL: test_tag: ten: +// CHECK-NEXT: result #0: live +// CHECK-LABEL: test_tag: one: +// CHECK-NEXT: result #0: live +// CHECK-LABEL: test_tag: x: +// CHECK-NEXT: result #0: live +func.func @test_type_b(%arg0: memref) { + %c0 = arith.constant {tag = "zero"} 0 : index + %c10 = arith.constant {tag = "ten"} 10 : index + %c1 = arith.constant {tag = "one"} 1 : index + %x = arith.constant {tag = "x"} 0 : i32 + %0 = scf.for %arg1 = %c0 to %c10 step %c1 iter_args(%arg2 = %x) -> (i32) { + %1 = arith.addi %x, %x : i32 + %2 = func.call @private0(%1) : (i32) -> i32 + scf.yield %2 : i32 + } + memref.store %0, %arg0[] : memref + return +} + +// ----- + +// Positive test: Type (c) "is returned by a public function" +// zero is live because it is returned by a public function. +// CHECK-LABEL: test_tag: zero: +// CHECK-NEXT: result #0: live +func.func @test_type_c() -> (f32){ + %0 = arith.constant {tag = "zero"} 0.0 : f32 + return %0 : f32 +} + +// ----- + +func.func private @private1(%0 : i32) -> i32 { + %1 = func.call @private2(%0) : (i32) -> i32 + %2 = arith.muli %0, %1 {tag = "in_private1"} : i32 + func.return %2 : i32 +} + +func.func private @private2(%0 : i32) -> i32 { + %cond = arith.index_cast %0 {tag = "in_private2"} : i32 to index + %1 = scf.index_switch %cond -> i32 + case 1 { + %ten = arith.constant 10 : i32 + scf.yield %ten : i32 + } + case 2 { + %twenty = arith.constant 20 : i32 + scf.yield %twenty : i32 + } + default { + %thirty = arith.constant 30 : i32 + scf.yield %thirty : i32 + } + func.return %1 : i32 +} + +// Positive test: Type (d) "is used to compute a value of type (c)" +// in_private1, in_private2, and final are live because they are used to compute +// the value returned by this public function. +// CHECK-LABEL: test_tag: in_private1: +// CHECK-NEXT: operand #0: live +// CHECK-NEXT: operand #1: live +// CHECK-NEXT: result #0: live +// CHECK-LABEL: test_tag: in_private2: +// CHECK-NEXT: operand #0: live +// CHECK-NEXT: result #0: live +// CHECK-LABEL: test_tag: final: +// CHECK-NEXT: operand #0: live +// CHECK-NEXT: operand #1: live +// CHECK-NEXT: result #0: live +func.func @test_type_d(%arg: i32) -> (i32) { + %0 = func.call @private1(%arg) : (i32) -> i32 + %final = arith.muli %0, %arg {tag = "final"} : i32 + return %final : i32 +} + +// ----- + +// Negative test: None of the types (a) through (d) +// zero is not live because it has no effect outside the program: it doesn't +// affect the memory or the program output. +// CHECK-LABEL: test_tag: zero: +// CHECK-NEXT: result #0: not live +// CHECK-LABEL: test_tag: one: +// CHECK-NEXT: result #0: live +func.func @test_negative_simple() -> (f32){ + %0 = arith.constant {tag = "zero"} 0.0 : f32 + %1 = arith.constant {tag = "one"} 1.0 : f32 + return %1 : f32 +} + +// ----- + +// The branch operand is incorrectly marked "not live" because, while live +// results in their branch operation can make branch operands live, this branch +// operation doesn't have any result. To mark such an operand live, the analysis +// needs to be enhanced more (already marked as a TODO). +// CHECK-LABEL: test_tag: br: +// CHECK-NEXT: operand #0: not live +// CHECK-NEXT: operand #1: live +// CHECK-NEXT: operand #2: live +func.func @test_negative_blocks(%arg0: memref, %arg1: memref, %arg2: memref, %arg3: i1) { + %c0_i32 = arith.constant 0 : i32 + cf.cond_br %arg3, ^bb1(%c0_i32 : i32), ^bb2(%c0_i32 : i32) {tag = "br"} +^bb1(%0 : i32): + memref.store %0, %arg0[] : memref + cf.br ^bb3 +^bb2(%1 : i32): + memref.store %1, %arg1[] : memref + cf.br ^bb3 +^bb3: + return +} + +// ----- + +// The branch operand is incorrectly marked "not live" because, while live +// results in their branch operation can make branch operands live, this branch +// operation doesn't have any result. To mark such an operand live, the analysis +// needs to be enhanced more (already marked as a TODO). +// CHECK-LABEL: test_tag: flag: +// CHECK-NEXT: operand #0: not live +func.func @test_negative_switch(%arg0: i32, %arg1: memref, %arg2: memref) { + %c0_i32 = arith.constant 0 : i32 + cf.switch %arg0 : i32, [ + default: ^bb1, + 42: ^bb2 + ] {tag = "flag"} +^bb1: + memref.store %c0_i32, %arg1[] : memref + cf.br ^bb3 +^bb2: + memref.store %c0_i32, %arg2[] : memref + cf.br ^bb3 +^bb3: + return +} + +// ----- + +// The branch operand is incorrectly marked "not live" because, while live +// results in their branch operation can make branch operands live, this branch +// operation doesn't have any result. To mark such an operand live, the analysis +// needs to be enhanced more (already marked as a TODO). +// CHECK-LABEL: test_tag: condition: +// CHECK-NEXT: operand #0: not live +// CHECK-NEXT: operand #1: live +func.func @test_negative_condition(%arg0: memref, %arg1: i32, %arg2: i1) { + %c0_i32 = arith.constant 0 : i32 + %0 = scf.while (%arg3 = %c0_i32) : (i32) -> (i32) { + memref.store %arg3, %arg0[] : memref + scf.condition(%arg2) {tag = "condition"} %arg3 : i32 + } do { + ^bb0(%arg3: i32): + memref.store %arg3, %arg0[] : memref + scf.yield %arg3 : i32 + } + memref.store %0, %arg0[] : memref + return +} diff --git a/mlir/test/lib/Analysis/CMakeLists.txt b/mlir/test/lib/Analysis/CMakeLists.txt --- a/mlir/test/lib/Analysis/CMakeLists.txt +++ b/mlir/test/lib/Analysis/CMakeLists.txt @@ -14,6 +14,7 @@ DataFlow/TestDeadCodeAnalysis.cpp DataFlow/TestDenseDataFlowAnalysis.cpp DataFlow/TestBackwardDataFlowAnalysis.cpp + DataFlow/TestLivenessAnalysis.cpp EXCLUDE_FROM_LIBMLIR diff --git a/mlir/test/lib/Analysis/DataFlow/TestLivenessAnalysis.cpp b/mlir/test/lib/Analysis/DataFlow/TestLivenessAnalysis.cpp new file mode 100644 --- /dev/null +++ b/mlir/test/lib/Analysis/DataFlow/TestLivenessAnalysis.cpp @@ -0,0 +1,79 @@ +//===- TestLivenessAnalysis.cpp - Test liveness analysis ------------------===// +// +// 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 +// +//===----------------------------------------------------------------------===// + +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +using namespace mlir; +using namespace mlir::dataflow; + +namespace { +struct TestLivenessAnalysisPass + : public PassWrapper> { + MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestLivenessAnalysisPass) + + StringRef getArgument() const override { return "test-liveness-analysis"; } + + void runOnOperation() override { + Operation *op = getOperation(); + + SymbolTableCollection symbolTable; + + DataFlowSolver solver; + solver.load(); + solver.load(); + solver.load(symbolTable); + if (failed(solver.initializeAndRun(op))) + return signalPassFailure(); + + raw_ostream &os = llvm::outs(); + op->walk([&](Operation *op) { + auto tag = op->getAttrOfType("tag"); + if (!tag) + return; + os << "test_tag: " << tag.getValue() << ":\n"; + for (auto [index, operand] : llvm::enumerate(op->getOperands())) { + const Liveness *liveness = solver.lookupState(operand); + assert(liveness && "expected a sparse lattice"); + os << " operand #" << index << ": "; + liveness->print(os); + os << "\n"; + } + for (auto [index, operand] : llvm::enumerate(op->getResults())) { + const Liveness *liveness = solver.lookupState(operand); + assert(liveness && "expected a sparse lattice"); + os << " result #" << index << ": "; + liveness->print(os); + os << "\n"; + } + }); + } +}; +} // end anonymous namespace + +namespace mlir { +namespace test { +void registerTestLivenessAnalysisPass() { + PassRegistration(); +} +} // end namespace test +} // end 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 @@ -102,6 +102,7 @@ void registerTestLinalgElementwiseFusion(); void registerTestLinalgGreedyFusion(); void registerTestLinalgTransforms(); +void registerTestLivenessAnalysisPass(); void registerTestLivenessPass(); void registerTestLoopFusion(); void registerTestCFGLoopInfoPass(); @@ -220,6 +221,7 @@ mlir::test::registerTestLinalgElementwiseFusion(); mlir::test::registerTestLinalgGreedyFusion(); mlir::test::registerTestLinalgTransforms(); + mlir::test::registerTestLivenessAnalysisPass(); mlir::test::registerTestLivenessPass(); mlir::test::registerTestLoopFusion(); mlir::test::registerTestCFGLoopInfoPass();