14 #include "llvm/ADT/SmallSet.h"
15 #include "llvm/Support/CommandLine.h"
16 #include "llvm/Support/Debug.h"
19 #define GEN_PASS_DEF_TESTSCFPARALLELLOOPCOLLAPSING
20 #include "mlir/Dialect/SCF/Transforms/Passes.h.inc"
23 #define DEBUG_TYPE "parallel-loop-collapsing"
28 struct TestSCFParallelLoopCollapsing
29 :
public impl::TestSCFParallelLoopCollapsingBase<
30 TestSCFParallelLoopCollapsing> {
31 void runOnOperation()
override {
40 if (!clCollapsedIndices0.empty())
41 combinedLoops.push_back(clCollapsedIndices0);
42 if (!clCollapsedIndices1.empty()) {
43 if (clCollapsedIndices0.empty()) {
45 <<
"collapsed-indices-1 specified but not collapsed-indices-0";
49 combinedLoops.push_back(clCollapsedIndices1);
51 if (!clCollapsedIndices2.empty()) {
52 if (clCollapsedIndices1.empty()) {
54 <<
"collapsed-indices-2 specified but not collapsed-indices-1";
58 combinedLoops.push_back(clCollapsedIndices2);
61 if (combinedLoops.empty()) {
62 llvm::errs() <<
"No collapsed-indices were specified. This pass is only "
63 "for testing and does not automatically collapse all "
64 "parallel loops or similar.";
71 llvm::SmallSet<unsigned, 8> flattenedCombinedLoops;
72 unsigned maxCollapsedIndex = 0;
73 for (
auto &loops : combinedLoops) {
74 for (
auto &loop : loops) {
75 flattenedCombinedLoops.insert(loop);
76 maxCollapsedIndex =
std::max(maxCollapsedIndex, loop);
80 if (maxCollapsedIndex != flattenedCombinedLoops.size() - 1 ||
81 !flattenedCombinedLoops.contains(maxCollapsedIndex)) {
83 <<
"collapsed-indices arguments must include all values [0,N).";
91 module->
walk([&](scf::ParallelOp op) {
92 if (flattenedCombinedLoops.size() != op.getNumLoops()) {
93 op.emitOpError(
"has ")
95 <<
" iter args while this limited functionality testing pass was "
96 "configured only for loops with exactly "
97 << flattenedCombinedLoops.size() <<
" iter args.";
107 return std::make_unique<TestSCFParallelLoopCollapsing>();
static Value max(ImplicitLocOpBuilder &builder, Value value, Value bound)
Operation is the basic unit of execution within MLIR.
std::enable_if_t< llvm::function_traits< std::decay_t< FnT > >::num_args==1, RetT > walk(FnT &&callback)
Walk the operation by calling the callback for each nested operation (including this one),...
Include the generated interface declarations.
void collapseParallelLoops(scf::ParallelOp loops, ArrayRef< std::vector< unsigned >> combinedDimensions)
Take the ParallelLoop and for each set of dimension indices, combine them into a single dimension.
std::unique_ptr< Pass > createTestSCFParallelLoopCollapsingPass()
Creates a pass that transforms a single ParallelLoop over N induction variables into another Parallel...