diff --git a/include/dataflow-scheduler/Analysis/SliceAnalysis.h b/include/dataflow-scheduler/Analysis/SliceAnalysis.h new file mode 100644 index 0000000..0d9cd9d --- /dev/null +++ b/include/dataflow-scheduler/Analysis/SliceAnalysis.h @@ -0,0 +1,208 @@ +//===-- SliceAnalysis.h -----------------------------------------*- c++ -*-===// +// +// Part of the Dataflow Scheduler project. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// +//===----------------------------------------------------------------------===// +// +// SSA slice analyses +// +// A slice through an SSA program contains all values that are reachable (either +// by going forward or backward) via def-use chains. In MLIR, there is also +// (structured) control flow, i.e., there are operations which may pass values +// along edges defined by other means. The slice analyses provide cached queries +// for the built-in MLIR interfaces that define these relationships. +// +//===----------------------------------------------------------------------===// + +#ifndef DATAFLOW_SCHEDULER_ANALYSIS_SLICEANALYSIS_H_ +#define DATAFLOW_SCHEDULER_ANALYSIS_SLICEANALYSIS_H_ + +#include +#include +#include +#include +#include +#include +#include + +#include + +namespace scheduler { + +/// Implements a backward dataflow SSA slice analysis. +/// +/// This analysis will (cache and) return all immediate (control flow) SSA value +/// predecessors for a given input. +class BackwardSliceAnalysis { + using set_type = llvm::SmallPtrSet; + + public: + /// Indicates the set of predecessors of an SSA value. + struct Predecessors { + /// Initializes an open set (lower bound) of @p values . + [[nodiscard]] static auto lowerBound(mlir::ValueRange values = {}) + -> Predecessors { + return Predecessors(false, values); + } + /// Initializes a closed set (exhaustive) of @p values . + [[nodiscard]] static auto exhaustive(mlir::ValueRange values = {}) + -> Predecessors { + return Predecessors(true, values); + } + + /// Initializes an empty lower bound. + /*implicit*/ Predecessors() = default; + + /// Updates this set to include @p values that are @p is_exhaustive . + void unite(mlir::ValueRange values, bool is_exhaustive = true) { + is_exhaustive_ &= is_exhaustive; + values_.insert_range(values); + } + /// Updates this set to include @p rhs . + void unite(const Predecessors& rhs) { + is_exhaustive_ &= rhs.is_exhaustive_; + values_.insert_range(rhs.values_); + } + + /// Determines whether the set of predecessor values is known to be closed. + /// + /// If `true`, then there are no predecessors besides those enumerated by + /// this container. If `false`, then there were types of control flow that + /// could not be followed, possibly due to unregistered operations. + [[nodiscard]] auto isExhaustive() const -> bool { return is_exhaustive_; } + + /// Determines whether there are no known and unknown predecessors. + [[nodiscard]] auto isKnownEmpty() const -> bool { + return isExhaustive() && values_.empty(); + } + + /// Gets the known predecessor values. + [[nodiscard]] auto getValues() const -> const set_type& { return values_; } + + //===------------------------------------------------------------------===// + // Container interface + //===------------------------------------------------------------------===// + + using value_type = set_type::value_type; + using size_type = set_type::size_type; + using iterator = set_type::const_iterator; + + [[nodiscard]] auto empty() const -> bool { return values_.empty(); } + [[nodiscard]] auto size() const -> size_type { return values_.size(); } + + [[nodiscard]] auto begin() const -> iterator { return values_.begin(); } + [[nodiscard]] auto end() const -> iterator { return values_.end(); } + + private: + friend class BackwardSliceAnalysis; + + explicit Predecessors(bool is_exhaustive, mlir::ValueRange values) + : is_exhaustive_(is_exhaustive), values_(llvm::from_range, values) {} + + bool is_exhaustive_ = false; + set_type values_; + }; + + using key_type = mlir::Value; + using mapped_type = Predecessors; + using map_type = llvm::DenseMap; + + // Allow construction as an MLIR analysis. + explicit BackwardSliceAnalysis(mlir::Operation* /*op*/ = nullptr); + + /// Gets the immediate @p predecessors of @p value . + /// + /// @param value Value to query the predecessors of. + /// @param [in,out] is_exhaustive Whether the result is exhaustive. + /// @param [out] predecessors Set of predecessors. + void getPredecessors(mlir::Value value, bool& is_exhaustive, + llvm::SmallPtrSetImpl& predecessors); + /// Gets the immediate @p predecessors of @p value . + void getPredecessors(mlir::Value value, Predecessors& predecessors) { + getPredecessors(value, predecessors.is_exhaustive_, predecessors.values_); + } + /// Gets the immediate predecessors of @p value . + [[nodiscard]] auto getPredecessors(mlir::Value value) -> Predecessors { + Predecessors result; + getPredecessors(value, result); + return result; + } + + /// Gets the immediate control flow predecessors of @p value . + [[nodiscard]] auto getControlFlowPredecessors(mlir::Value value) + -> const Predecessors&; + + private: + map_type control_flow_; +}; + +/// Base class for implementing a forward slice analysis. +/// +/// This analysis will determine whether SSA values are reachable from forward +/// dataflow starting with an initial set of values. +class ForwardSlice { + public: + /// Result of a slice membership check. + enum class Result : char { + /// Value is not in the slice. + NoContain = 0, + /// Value might be in the slice (lower bound). + MayContain = 0b01, + /// Value must be in the slice (upper bound). + MustContain = 0b11, + }; + + using key_type = mlir::Value; + using mapped_type = Result; + using map_type = llvm::DenseMap; + + /// Initializes a ForwardSlice using @p backward and the initial @p values . + explicit ForwardSlice(BackwardSliceAnalysis& backward, + mlir::ValueRange values); + + /// Inserts additional @p values into the slice. + /// + /// @retval false @p values were already contained. + /// @retval true New values were added, and the cache was invalidated. + auto insert(mlir::ValueRange values) -> bool; + + /// Determines whether @p value is in the slice. + auto contains(mlir::Value value) -> Result; + + private: + BackwardSliceAnalysis& backward_; + map_type cache_; +}; + +/// Implements a ForwardSlice based on the loop variables of an operation. +/// +/// Loop variables are induction variables and inter-iteration dependencies +/// carried by region arguments, as advertised by the mlir::LoopLikeOpInterface. +/// If the operation does not implement this interface, the slice is empty. +class LoopSliceAnalysis : public ForwardSlice { + public: + // Allow construction as an MLIR analysis. + explicit LoopSliceAnalysis(mlir::Operation* op, + mlir::AnalysisManager& analyses); +}; + +/// Backport of llvm-project/pull/188758. +[[nodiscard]] +auto getControlFlowPredecessors(mlir::Value value) + -> std::optional>; + +} // namespace scheduler + +#endif // DATAFLOW_SCHEDULER_ANALYSIS_SLICEANALYSIS_H_ diff --git a/lib/Analysis/CMakeLists.txt b/lib/Analysis/CMakeLists.txt index bae2b72..02aafcd 100644 --- a/lib/Analysis/CMakeLists.txt +++ b/lib/Analysis/CMakeLists.txt @@ -6,6 +6,7 @@ add_dataflow_scheduler_library(DataflowSchedulerAnalysis MemoryTrackerAnalysis.cpp OperationTree.cpp PipelineTree.cpp + SliceAnalysis.cpp Utils.cpp WriteSetScan.cpp diff --git a/lib/Analysis/SliceAnalysis.cpp b/lib/Analysis/SliceAnalysis.cpp new file mode 100644 index 0000000..9649cd8 --- /dev/null +++ b/lib/Analysis/SliceAnalysis.cpp @@ -0,0 +1,268 @@ +//===-- SliceAnalysis.cpp ---------------------------------------*- c++ -*-===// +// +// Part of the Dataflow Scheduler project. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// +//===----------------------------------------------------------------------===// + +#include "dataflow-scheduler/Analysis/SliceAnalysis.h" + +#include +#include +#include +#include +#include +#include +#include +#include + +using namespace scheduler; + +namespace { + +void visitControlFlow(mlir::SelectLikeOpInterface select, + BackwardSliceAnalysis::map_type& result) { + if (select->getNumResults() != 1) { + return; + } + + result.emplace_or_assign( + select->getResult(0), + BackwardSliceAnalysis::Predecessors::exhaustive( + {select.getTrueValue(), select.getFalseValue()})); +} + +void visitControlFlow(mlir::RegionBranchOpInterface branch, + BackwardSliceAnalysis::map_type& result) { + for (auto branch_point : branch.getAllRegionBranchPoints()) { + llvm::SmallVector successors; + branch.getSuccessorRegions(branch_point, successors); + for (auto& successor : successors) { + const auto entry_values = successor.getSuccessorInputs(); + const auto exit_values = + branch.getSuccessorOperands(branch_point, successor); + for (auto [entry, exit] : llvm::zip_equal(entry_values, exit_values)) { + auto& cached = + result + .try_emplace(entry, + BackwardSliceAnalysis::Predecessors::exhaustive()) + .first->getSecond(); + cached.unite(exit); + } + } + } +} + +void visitControlFlow(mlir::BranchOpInterface branch, + BackwardSliceAnalysis::map_type& result) { + for (auto& successor : branch->getBlockOperands()) { + const auto operands = + branch.getSuccessorOperands(successor.getOperandNumber()); + + for (auto argument : successor.get()->getArguments()) { + auto& cached = + result + .try_emplace(argument, + BackwardSliceAnalysis::Predecessors::exhaustive()) + .first->getSecond(); + if (const auto forwarded = operands[argument.getArgNumber()]; forwarded) { + cached.unite(forwarded); + } + } + } +} + +} // namespace + +//===----------------------------------------------------------------------===// +// BackwardSliceAnalysis +//===----------------------------------------------------------------------===// + +void BackwardSliceAnalysis::getPredecessors( + mlir::Value value, bool& is_exhaustive, + SmallPtrSetImpl& predecessors) { + if (const auto result = llvm::dyn_cast(value); result) { + is_exhaustive &= result.getOwner()->isRegistered(); + predecessors.insert_range(result.getOwner()->getOperands()); + } + + const auto& control_flow = getControlFlowPredecessors(value); + is_exhaustive &= control_flow.is_exhaustive_; + predecessors.insert_range(control_flow.values_); +} + +auto BackwardSliceAnalysis::getControlFlowPredecessors(mlir::Value value) + -> const Predecessors& { + { + // Lookup in cache, emplacing inexact result if none exists. + auto [it, inserted] = control_flow_.try_emplace(value); + if (!inserted) { + return it->second; + } + } + + if (const auto argument = llvm::dyn_cast(value); + argument) { + if (!argument.getOwner()->hasNoPredecessors()) { + // Argument is fully determined by predecessors' branch operations. + auto is_exact = argument.getOwner()->getParentOp()->isRegistered(); + for (auto* const pred : argument.getOwner()->getPredecessors()) { + if (auto iface = llvm::dyn_cast( + pred->getTerminator())) { + visitControlFlow(iface, control_flow_); + } else { + // We don't understand that one. + is_exact = false; + } + } + + auto& result = control_flow_[value]; + result.is_exhaustive_ = is_exact; + return result; + } + + if (auto iface = llvm::dyn_cast( + argument.getOwner()->getParentOp()); + iface) { + // We do not perform any inter-procedural analyses. + return control_flow_[value] = Predecessors::exhaustive(); + } + + if (auto iface = llvm::dyn_cast( + argument.getOwner()->getParentOp()); + iface) { + // Argument is determined by region branch semantics. + visitControlFlow(iface, control_flow_); + } + } else { + const auto result = llvm::cast(value); + + if (auto iface = + llvm::dyn_cast(result.getOwner()); + iface) { + // Result is determined by select semantics. + visitControlFlow(iface, control_flow_); + } else if (auto iface = llvm::dyn_cast( + result.getOwner()); + iface) { + // Result is determined by region branch semantics. + visitControlFlow(iface, control_flow_); + } + } + + return control_flow_[value]; +} + +//===----------------------------------------------------------------------===// +// ForwardSliceAnalysis +//===----------------------------------------------------------------------===// + +ForwardSlice::ForwardSlice(BackwardSliceAnalysis& backward, + mlir::ValueRange values) + : backward_(backward) { + for (auto value : values) { + cache_[value] = Result::MustContain; + } +} + +auto ForwardSlice::insert(mlir::ValueRange values) -> bool { + if (llvm::all_of(values, [&](mlir::Value value) -> bool { + return cache_[value] == Result::MustContain; + })) { + return false; + } + + // This invalidates the cache apart from MustContain. + map_type temp; + using std::swap; + swap(temp, cache_); + + for (auto value : values) { + cache_[value] = Result::MustContain; + } + for (auto [key, value] : temp) { + if (value == Result::MustContain) { + cache_[key] = Result::MustContain; + } + } + + return true; +} + +auto ForwardSlice::contains(mlir::Value value) -> Result { + if (cache_.empty()) { + // Short-circuit on the known empty slice. + return Result::NoContain; + } + + // Lookup cached result or initialize with MayContain. + auto [it, invalid] = cache_.try_emplace(value, Result::MayContain); + if (!invalid) { + return it->second; + } + + // Find all predecessors of the value. + bool is_exhaustive; + llvm::SmallPtrSet predecessors; + backward_.getPredecessors(value, is_exhaustive, predecessors); + + // If the set of predecessors is exhaustive, we may assume that the value is + // independent for now. If it is visited in the recursive search (which can + // only happen within a graph region, or if we follow block arguments), then + // it being part of its own cycle should not be an obstacle to independence. + auto result = it->second = + is_exhaustive ? Result::NoContain : Result::MayContain; + + for (auto predecessor : predecessors) { + switch (contains(predecessor)) { + case Result::MustContain: + return cache_[value] = Result::MustContain; + case Result::MayContain: + result = cache_[value] = Result::MayContain; + continue; + case Result::NoContain: + continue; + } + } + + return result; +} + +//===----------------------------------------------------------------------===// +// LoopSliceAnalysis +//===----------------------------------------------------------------------===// + +namespace { + +auto getLoopVariables(mlir::Operation* op) -> llvm::SmallVector { + llvm::SmallVector result; + + auto iface = llvm::dyn_cast(op); + if (iface == nullptr) { + return result; + } + + if (const auto ivs = iface.getLoopInductionVars(); ivs) { + llvm::append_range(result, ivs.value()); + } + llvm::append_range(result, iface.getRegionIterArgs()); + return result; +} + +} // namespace + +LoopSliceAnalysis::LoopSliceAnalysis(mlir::Operation* op, + mlir::AnalysisManager& analyses) + : ForwardSlice(analyses.getAnalysis(), + getLoopVariables(op)) {}