Skip to content
This repository was archived by the owner on Jan 25, 2023. It is now read-only.

Commit 15bc678

Browse files
author
Ivan Butygin
authored
[MLIR] Optimizations (#205)
* refac * refac * extend CanonicalizeReduction * PromoteParallelPass pass * promote some simple loads * Rewrite single-element single-store memref * copypaste naiveParallelLoopFusion from mlir * Extend mem rewrites
1 parent 71c8661 commit 15bc678

11 files changed

Lines changed: 594 additions & 54 deletions

File tree

‎mlir-compiler/mlir-compiler/src/pipelines/plier_to_linalg.cpp‎

Lines changed: 43 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1028,9 +1028,50 @@ void PostLinalgOptPass::runOnOperation()
10281028
auto& context = getContext();
10291029
plier::populate_common_opts_patterns(context, patterns);
10301030

1031+
patterns.insert<
1032+
plier::CanonicalizeReduction
1033+
>(&context);
1034+
1035+
mlir::FrozenRewritePatternList frozenPatterns(std::move(patterns));
1036+
1037+
while (true)
1038+
{
1039+
(void)mlir::applyPatternsAndFoldGreedily(getOperation(), frozenPatterns);
1040+
bool rerun = false;
1041+
for (auto& op : getOperation().getRegion().front())
1042+
{
1043+
if (auto func = mlir::dyn_cast<mlir::FuncOp>(op))
1044+
{
1045+
if (mlir::succeeded(plier::naivelyFuseParallelOps(func.getRegion())))
1046+
{
1047+
rerun = true;
1048+
}
1049+
}
1050+
}
1051+
if (!rerun)
1052+
{
1053+
break;
1054+
}
1055+
}
1056+
1057+
}
1058+
1059+
struct PromoteParallelPass :
1060+
public mlir::PassWrapper<PromoteParallelPass, mlir::OperationPass<mlir::ModuleOp>>
1061+
{
1062+
void runOnOperation() override;
1063+
};
1064+
1065+
void PromoteParallelPass::runOnOperation()
1066+
{
1067+
mlir::OwningRewritePatternList patterns;
1068+
1069+
auto& context = getContext();
1070+
plier::populate_common_opts_patterns(context, patterns);
1071+
10311072
patterns.insert<
10321073
plier::CanonicalizeReduction,
1033-
plier::PromoteToParallel
1074+
plier::PromoteToParallel // TODO
10341075
>(&context);
10351076

10361077
(void)mlir::applyPatternsAndFoldGreedily(getOperation(), std::move(patterns));
@@ -1064,11 +1105,9 @@ void populate_plier_to_linalg_opt_pipeline(mlir::OpPassManager& pm)
10641105
pm.addPass(mlir::createCopyRemovalPass());
10651106

10661107
pm.addPass(std::make_unique<LowerLinalgPass>());
1067-
pm.addPass(mlir::createParallelLoopFusionPass());
10681108
pm.addPass(std::make_unique<PostLinalgOptPass>());
10691109
pm.addPass(mlir::createSymbolDCEPass());
1070-
pm.addPass(mlir::createParallelLoopFusionPass()); // TODO: make this rewrite and add to PostLinalgOptPass
1071-
pm.addPass(std::make_unique<PostLinalgOptPass>());
1110+
pm.addPass(std::make_unique<PromoteParallelPass>());
10721111
}
10731112
}
10741113

‎mlir-compiler/plier/CMakeLists.txt‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -22,10 +22,12 @@ set(SOURCES_LIST
2222
src/rewrites/cse.cpp
2323
src/rewrites/force_inline.cpp
2424
src/rewrites/index_type_propagation.cpp
25-
include/plier/rewrites/if_rewrites.cpp
25+
src/rewrites/if_rewrites.cpp
2626
src/rewrites/loop_rewrites.cpp
27+
src/rewrites/memory_rewrites.cpp
2728
src/rewrites/promote_to_parallel.cpp
2829
src/rewrites/type_conversion.cpp
30+
src/transforms/block_utils.cpp
2931
src/transforms/cast_utils.cpp
3032
src/transforms/const_utils.cpp
3133
src/transforms/func_utils.cpp
@@ -47,8 +49,10 @@ set(HEADERS_LIST
4749
include/plier/rewrites/index_type_propagation.hpp
4850
include/plier/rewrites/if_rewrites.hpp
4951
include/plier/rewrites/loop_rewrites.hpp
52+
include/plier/rewrites/memory_rewrites.hpp
5053
include/plier/rewrites/promote_to_parallel.hpp
5154
include/plier/rewrites/type_conversion.hpp
55+
include/plier/transforms/block_utils.hpp
5256
include/plier/transforms/cast_utils.hpp
5357
include/plier/transforms/const_utils.hpp
5458
include/plier/transforms/func_utils.hpp

‎mlir-compiler/plier/include/plier/rewrites/memory_rewrites.hpp‎

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,28 @@
1+
#pragma once
2+
3+
#include <mlir/IR/PatternMatch.h>
4+
5+
namespace mlir
6+
{
7+
class FuncOp;
8+
class StoreOp;
9+
}
10+
11+
namespace plier
12+
{
13+
struct PromoteLoads : public mlir::OpRewritePattern<mlir::FuncOp>
14+
{
15+
using mlir::OpRewritePattern<mlir::FuncOp>::OpRewritePattern;
16+
17+
mlir::LogicalResult matchAndRewrite(
18+
mlir::FuncOp op, mlir::PatternRewriter &rewriter) const override;
19+
};
20+
21+
struct SingeWriteMemref : public mlir::OpRewritePattern<mlir::StoreOp>
22+
{
23+
using mlir::OpRewritePattern<mlir::StoreOp>::OpRewritePattern;
24+
25+
mlir::LogicalResult matchAndRewrite(
26+
mlir::StoreOp op, mlir::PatternRewriter &rewriter) const override;
27+
};
28+
}

‎mlir-compiler/plier/include/plier/transforms/block_utils.hpp‎

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,20 @@
1+
#pragma once
2+
3+
namespace mlir
4+
{
5+
class Operation;
6+
}
7+
8+
namespace plier
9+
{
10+
11+
enum class OpRelation
12+
{
13+
Before,
14+
After,
15+
In,
16+
Unknown
17+
};
18+
19+
OpRelation relativeTo(mlir::Operation* op, mlir::Operation* relativeTo);
20+
}

‎mlir-compiler/plier/include/plier/transforms/loop_utils.hpp‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@ class Value;
1010
class Location;
1111
class OpBuilder;
1212
class Type;
13+
class Region;
1314
namespace scf
1415
{
1516
class ForOp;
@@ -27,4 +28,6 @@ mlir::LogicalResult lower_while_to_for(plier::GetiterOp getiter, mlir::PatternRe
2728
llvm::function_ref<std::tuple<mlir::Value,mlir::Value,mlir::Value>(mlir::OpBuilder&, mlir::Location)> get_bounds,
2829
llvm::function_ref<mlir::Value(mlir::OpBuilder&, mlir::Location, mlir::Type, mlir::Value)> get_iter_val,
2930
llvm::function_ref<void(mlir::scf::ForOp)> results = nullptr);
31+
32+
mlir::LogicalResult naivelyFuseParallelOps(mlir::Region &region);
3033
}

‎mlir-compiler/plier/src/rewrites/canonicalize_reductions.cpp‎

Lines changed: 78 additions & 48 deletions
Original file line numberDiff line numberDiff line change
@@ -4,28 +4,46 @@
44
#include <mlir/Dialect/StandardOps/IR/Ops.h>
55
#include <mlir/IR/BlockAndValueMapping.h>
66

7+
#include "plier/transforms/block_utils.hpp"
8+
79
namespace
810
{
911
bool checkMemrefType(mlir::Value value)
1012
{
1113
if (auto type = value.getType().dyn_cast<mlir::MemRefType>())
1214
{
13-
auto shape = type.getShape();
14-
return shape.empty() || (1 == shape.size() && 1 == shape[0]);
15+
// auto shape = type.getShape();
16+
// return shape.empty() || (1 == shape.size() && 1 == shape[0]);
17+
return true;
1518
}
1619
return false;
1720
}
1821

19-
bool checkForPotentialAliases(mlir::Value value)
22+
bool isOutsideBlock(mlir::ValueRange values, mlir::Block& block)
2023
{
21-
auto def_op = value.getDefiningOp();
22-
if (nullptr == def_op)
24+
auto blockArgs = block.getArguments();
25+
for (auto val : values)
2326
{
24-
return false;
27+
if (llvm::find(blockArgs, val) != blockArgs.end())
28+
{
29+
return false;
30+
}
31+
auto op = val.getDefiningOp();
32+
if (op && block.findAncestorOpInBlock(*op))
33+
{
34+
return false;
35+
}
2536
}
26-
if (auto effects = mlir::dyn_cast<mlir::MemoryEffectOpInterface>(def_op))
37+
return true;
38+
}
39+
40+
bool checkForPotentialAliases(mlir::Value value, mlir::Operation* parent)
41+
{
42+
assert(parent->getRegions().size() == 1);
43+
assert(llvm::hasNItems(parent->getRegions().front(), 1));
44+
if (auto effects = mlir::dyn_cast_or_null<mlir::MemoryEffectOpInterface>(value.getDefiningOp()))
2745
{
28-
if (!effects.hasEffect<mlir::MemoryEffects::Allocate>())
46+
if (!effects.onlyHasEffect<mlir::MemoryEffects::Allocate>())
2947
{
3048
return false;
3149
}
@@ -34,13 +52,53 @@ bool checkForPotentialAliases(mlir::Value value)
3452
{
3553
return false;
3654
}
55+
56+
mlir::LoadOp load;
57+
mlir::StoreOp store;
58+
auto& parentBlock = parent->getRegions().front().front();
3759
for (auto user : value.getUsers())
3860
{
3961
if (mlir::isa<mlir::ViewLikeOpInterface>(user))
4062
{
4163
// TODO: very conservative
4264
return false;
4365
}
66+
auto relation = plier::relativeTo(user, parent);
67+
if (plier::OpRelation::Unknown == relation)
68+
{
69+
return false;
70+
}
71+
if (plier::OpRelation::In == relation)
72+
{
73+
if (auto effects = mlir::dyn_cast_or_null<mlir::MemoryEffectOpInterface>(user))
74+
{
75+
if (user->getBlock() != &parentBlock)
76+
{
77+
return false;
78+
}
79+
if (effects.hasEffect<mlir::MemoryEffects::Read>())
80+
{
81+
if (load || !mlir::isa<mlir::LoadOp>(user))
82+
{
83+
return false;
84+
}
85+
load = mlir::cast<mlir::LoadOp>(user);
86+
}
87+
if (effects.hasEffect<mlir::MemoryEffects::Write>())
88+
{
89+
if (store || !mlir::isa<mlir::StoreOp>(user))
90+
{
91+
return false;
92+
}
93+
store = mlir::cast<mlir::StoreOp>(user);
94+
}
95+
}
96+
}
97+
}
98+
if (!load || !store || !load->isBeforeInBlock(store) ||
99+
load.indices() != store.indices() || !isOutsideBlock(load.indices(), parentBlock))
100+
{
101+
return false;
44102
}
45103
return true;
46104
}
@@ -59,63 +117,35 @@ bool checkSupportedOps(mlir::Value value, mlir::Operation* parent)
59117

60118
bool checkMemref(mlir::Value value, mlir::Operation* parent)
61119
{
62-
return checkMemrefType(value) && checkForPotentialAliases(value) &&
120+
return checkMemrefType(value) &&
121+
checkForPotentialAliases(value, parent) &&
63122
checkSupportedOps(value, parent);
64123
}
65124

66-
mlir::Value createScalarLoad(
67-
mlir::PatternRewriter &builder, mlir::Location loc, mlir::Value memref)
125+
mlir::Value createScalarLoad(mlir::PatternRewriter &builder, mlir::Location loc, mlir::Value memref, mlir::ValueRange indices)
68126
{
69-
auto shape = memref.getType().cast<mlir::MemRefType>().getShape();
70-
if (shape.empty())
71-
{
72-
return builder.create<mlir::LoadOp>(loc, memref);
73-
}
74-
else if (llvm::all_of(shape, [](auto s) { return s == 1; }))
75-
{
76-
auto index = builder.create<mlir::ConstantIndexOp>(loc, 0);
77-
llvm::SmallVector<mlir::Value> indices(shape.size(), index);
78-
return builder.create<mlir::LoadOp>(loc, memref, indices);
79-
}
80-
else
81-
{
82-
llvm_unreachable("Invalid shape");
83-
}
127+
return builder.create<mlir::LoadOp>(loc, memref, indices);
84128
}
85129

86130
void createScalarStore(
87131
mlir::PatternRewriter &builder, mlir::Location loc, mlir::Value val,
88-
mlir::Value memref)
132+
mlir::Value memref, mlir::ValueRange indices)
89133
{
90-
auto shape = memref.getType().cast<mlir::MemRefType>().getShape();
91-
if (shape.empty())
92-
{
93-
builder.create<mlir::StoreOp>(loc, val, memref);
94-
}
95-
else if (llvm::all_of(shape, [](auto s) { return s == 1; }))
96-
{
97-
auto index = builder.create<mlir::ConstantIndexOp>(loc, 0);
98-
llvm::SmallVector<mlir::Value> indices(shape.size(), index);
99-
builder.create<mlir::StoreOp>(loc, val, memref, indices);
100-
}
101-
else
102-
{
103-
llvm_unreachable("Invalid shape");
104-
}
134+
builder.create<mlir::StoreOp>(loc, val, memref, indices);
105135
}
106136
}
107137

108138
mlir::LogicalResult plier::CanonicalizeReduction::matchAndRewrite(mlir::scf::ForOp op, mlir::PatternRewriter& rewriter) const
109139
{
110-
llvm::SmallVector<mlir::Value> to_process;
140+
llvm::SmallVector<std::pair<mlir::Value, mlir::ValueRange>> to_process;
111141
for (auto& current : op.getLoopBody().front())
112142
{
113143
if (auto load = mlir::dyn_cast<mlir::LoadOp>(current))
114144
{
115145
auto memref = load.memref();
116146
if (checkMemref(memref, op))
117147
{
118-
to_process.emplace_back(memref);
148+
to_process.push_back({memref, load.indices()});
119149
}
120150
}
121151
}
@@ -124,9 +154,9 @@ mlir::LogicalResult plier::CanonicalizeReduction::matchAndRewrite(mlir::scf::For
124154
{
125155
auto loc = op.getLoc();
126156
auto init_args = llvm::to_vector<8>(op.initArgs());
127-
for (auto val : to_process)
157+
for (auto it : to_process)
128158
{
129-
init_args.emplace_back(createScalarLoad(rewriter, loc, val));
159+
init_args.emplace_back(createScalarLoad(rewriter, loc, it.first, it.second));
130160
}
131161
auto prev_args_offset = op.initArgs().size();
132162
auto body = [&](mlir::OpBuilder &builder, mlir::Location loc, mlir::Value iter, mlir::ValueRange iter_vals)
@@ -142,7 +172,7 @@ mlir::LogicalResult plier::CanonicalizeReduction::matchAndRewrite(mlir::scf::For
142172
auto get_iter_index = [&](auto op)->unsigned
143173
{
144174
auto arg = op.memref();
145-
for (auto it : llvm::enumerate(to_process))
175+
for (auto it : llvm::enumerate(llvm::make_first_range(to_process)))
146176
{
147177
if (arg == it.value())
148178
{
@@ -189,7 +219,7 @@ mlir::LogicalResult plier::CanonicalizeReduction::matchAndRewrite(mlir::scf::For
189219
{
190220
auto index = prev_args_offset + it.index();
191221
auto result = results[static_cast<unsigned>(index)];
192-
createScalarStore(rewriter, loc, result, it.value());
222+
createScalarStore(rewriter, loc, result, it.value().first, it.value().second);
193223
}
194224
rewriter.replaceOp(op, results.take_front(prev_args_offset));
195225
return mlir::success();

‎mlir-compiler/plier/src/rewrites/common_opts.cpp‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33
#include "plier/rewrites/force_inline.hpp"
44
#include "plier/rewrites/index_type_propagation.hpp"
55
#include "plier/rewrites/loop_rewrites.hpp"
6+
#include "plier/rewrites/memory_rewrites.hpp"
67
#include "plier/rewrites/cse.hpp"
78
#include "plier/rewrites/if_rewrites.hpp"
89

@@ -26,7 +27,9 @@ void plier::populate_common_opts_patterns(mlir::MLIRContext& context, mlir::Owni
2627
plier::IfOpConstCond,
2728
SimplifySelect,
2829
SimplifySelectEq,
29-
plier::CSERewrite<mlir::FuncOp, /*recusive*/false>
30+
plier::CSERewrite<mlir::FuncOp, /*recusive*/false>,
31+
PromoteLoads,
32+
SingeWriteMemref
3033
>(&context);
3134

3235
plier::populate_index_propagate_patterns(context, patterns);

mlir-compiler/plier/include/plier/rewrites/if_rewrites.cpp renamed to mlir-compiler/plier/src/rewrites/if_rewrites.cpp

File renamed without changes.

0 commit comments

Comments
 (0)