maxjhandsome
/
llvm-project
forked from OSchip/llvm-project
1
0
Fork 0
Code Issues Pull Requests Packages Releases Wiki Activity

Revert "[NFC] Generalize a couple of passes so they can operate on any FunctionLike op." 

This reverts commit 34696e6542.

A test is crashing on the mlir-nvidia bot.
This commit is contained in:
Mehdi Amini 2021-12-13 20:41:25 +00:00
parent 6c3bc57f25
commit aa8815e42e
15 changed files with 56 additions and 79 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
mlir/include/mlir/Dialect/Linalg/Passes.h
View File

@ -19,9 +19,9 @@
namespace mlir {
std::unique_ptr<Pass> createConvertElementwiseToLinalgPass();
std::unique_ptr<OperationPass<FuncOp>> createConvertElementwiseToLinalgPass();
std::unique_ptr<Pass> createLinalgFoldUnitExtentDimsPass();
std::unique_ptr<OperationPass<FuncOp>> createLinalgFoldUnitExtentDimsPass();
std::unique_ptr<Pass> createLinalgElementwiseOpFusionPass();
std::unique_ptr<Pass> createFoldReshapeOpsByLinearizationPass();

16
mlir/include/mlir/Dialect/Linalg/Passes.td
View File

@ -11,15 +11,12 @@
include "mlir/Pass/PassBase.td"
def ConvertElementwiseToLinalg : Pass<"convert-elementwise-to-linalg", ""> {
def ConvertElementwiseToLinalg : FunctionPass<"convert-elementwise-to-linalg"> {
let summary = "Convert ElementwiseMappable ops to linalg";
let description = [{
Convert ops with the `ElementwiseMappable` trait to linalg parallel loops.
This pass only converts ops that operate on ranked tensors.
This can be run on any op with the FunctionLike trait and must not be
run on others.
}];
let constructor = "mlir::createConvertElementwiseToLinalgPass()";
let dependentDialects = ["linalg::LinalgDialect", "memref::MemRefDialect"];
@ -57,12 +54,8 @@ def LinalgComprehensiveModuleBufferize :
let constructor = "mlir::createLinalgComprehensiveModuleBufferizePass()";
}
def LinalgFoldUnitExtentDims : Pass<"linalg-fold-unit-extent-dims", ""> {
def LinalgFoldUnitExtentDims : FunctionPass<"linalg-fold-unit-extent-dims"> {
let summary = "Remove unit-extent dimension in Linalg ops on tensors";
let description = [{
This can be run on any op with the FunctionLike trait and must not be
run on others.
}];
let constructor = "mlir::createLinalgFoldUnitExtentDimsPass()";
let options = [
Option<"foldOneTripLoopsOnly", "fold-one-trip-loops-only", "bool",
@ -204,7 +197,7 @@ def LinalgGeneralization : FunctionPass<"linalg-generalize-named-ops"> {
let dependentDialects = ["linalg::LinalgDialect"];
}
def LinalgDetensorize : Pass<"linalg-detensorize", ""> {
def LinalgDetensorize : FunctionPass<"linalg-detensorize"> {
let summary = "Detensorize linalg ops";
let constructor = "mlir::createLinalgDetensorizePass()";
let dependentDialects = [];
@ -229,9 +222,6 @@ def LinalgDetensorize : Pass<"linalg-detensorize", ""> {
In addition to detensoring individual ops, this pass detensors internal
control flow inside a function. All blocks except for the entry block are
detensored by converting their arguments whenever possible.
This can be run on any op with the FunctionLike trait and must not be
run on others.
}];
}

65
mlir/lib/Dialect/Linalg/Transforms/Detensorize.cpp
View File

@ -93,11 +93,10 @@ public:
/// A conversion pattern for detensoring internal (non-entry) blocks within a
/// function.
struct FunctionNonEntryBlockConversion : public ConversionPattern {
FunctionNonEntryBlockConversion(MLIRContext *ctx, TypeConverter &converter,
FunctionNonEntryBlockConversion(StringRef functionLikeOpName,
MLIRContext *ctx, TypeConverter &converter,
DenseSet<BlockArgument> blockArgsToDetensor)
: ConversionPattern(converter, MatchTraitOpTypeTag(),
TypeID::get<OpTrait::FunctionLike>(), /*benefit=*/1,
ctx),
: ConversionPattern(converter, functionLikeOpName, /*benefit=*/1, ctx),
blockArgsToDetensor(blockArgsToDetensor) {}
LogicalResult
@ -236,8 +235,7 @@ struct LinalgDetensorize : public LinalgDetensorizeBase<LinalgDetensorize> {
/// detensored, then:
/// - opsToDetensor should be = {linalg.generic{add}}.
/// - blockArgsToDetensor should be = {bb1 -> {0}, bb2 -> {0}}.
virtual void compute(Operation *func,
DetensorizeTypeConverter typeConverter,
virtual void compute(FuncOp func, DetensorizeTypeConverter typeConverter,
DenseSet<Operation *> &opsToDetensor,
DenseSet<BlockArgument> &blockArgsToDetensor) = 0;
@ -288,18 +286,18 @@ struct LinalgDetensorize : public LinalgDetensorizeBase<LinalgDetensorize> {
/// AND can be detensored.
class ControlFlowDetectionModel : public CostModel {
public:
void compute(Operation *func, DetensorizeTypeConverter typeConverter,
void compute(FuncOp func, DetensorizeTypeConverter typeConverter,
DenseSet<Operation *> &opsToDetensor,
DenseSet<BlockArgument> &blockArgsToDetensor) override {
SmallVector<Value> workList;
func->walk([&](CondBranchOp condBr) {
func.walk([&](CondBranchOp condBr) {
for (auto operand : condBr.getOperands()) {
workList.push_back(operand);
}
});
func->walk([&](BranchOp br) {
func.walk([&](BranchOp br) {
for (auto operand : br.getOperands()) {
workList.push_back(operand);
}
@ -493,24 +491,21 @@ struct LinalgDetensorize : public LinalgDetensorizeBase<LinalgDetensorize> {
/// Detensorize everything that can detensored.
class AggressiveDetensoringModel : public CostModel {
public:
void compute(Operation *func, DetensorizeTypeConverter typeConverter,
void compute(FuncOp func, DetensorizeTypeConverter typeConverter,
DenseSet<Operation *> &opsToDetensor,
DenseSet<BlockArgument> &blockArgsToDetensor) override {
func->walk([&](GenericOp genericOp) {
func.walk([&](GenericOp genericOp) {
if (shouldBeDetensored(genericOp, typeConverter))
opsToDetensor.insert(genericOp);
});
for (Block &block :
llvm::drop_begin(function_like_impl::getFunctionBody(func), 1))
for (Block &block : llvm::drop_begin(func.getBody(), 1))
for (BlockArgument blockArgument : block.getArguments())
blockArgsToDetensor.insert(blockArgument);
}
};
void runOnOperation() override {
assert(getOperation()->hasTrait<OpTrait::FunctionLike>() &&
"DetensorizePass can only be run on FunctionLike operations");
void runOnFunction() override {
MLIRContext *context = &getContext();
DetensorizeTypeConverter typeConverter;
RewritePatternSet patterns(context);
@ -521,12 +516,12 @@ struct LinalgDetensorize : public LinalgDetensorizeBase<LinalgDetensorize> {
if (aggressiveMode.getValue()) {
AggressiveDetensoringModel costModel;
costModel.compute(getOperation(), typeConverter, opsToDetensor,
costModel.compute(getFunction(), typeConverter, opsToDetensor,
blockArgsToDetensor);
} else {
ControlFlowDetectionModel costModel;
costModel.compute(getOperation(), typeConverter, opsToDetensor,
costModel.compute(getFunction(), typeConverter, opsToDetensor,
blockArgsToDetensor);
}
@ -536,26 +531,24 @@ struct LinalgDetensorize : public LinalgDetensorizeBase<LinalgDetensorize> {
target.addDynamicallyLegalOp<GenericOp>(
[&](GenericOp op) { return !opsToDetensor.count(op); });
target.markUnknownOpDynamicallyLegal([&](Operation *op) {
target.addDynamicallyLegalOp<FuncOp>([&](FuncOp op) {
// A function is legal if all of its non-entry blocks are legal. We
// don't legalize the entry block (i.e. the function's signature)
// since detensoring can't happen along external calling convention
// boundaries, which we conservatively approximate as all function
// signatures.
if (op->hasTrait<OpTrait::FunctionLike>()) {
auto &body = function_like_impl::getFunctionBody(op);
return llvm::all_of(llvm::drop_begin(body, 1), [&](Block &block) {
if (llvm::any_of(
blockArgsToDetensor, [&](BlockArgument blockArgument) {
return blockArgument.getOwner() == &block &&
!typeConverter.isLegal(blockArgument.getType());
})) {
return false;
}
return true;
});
}
return llvm::all_of(llvm::drop_begin(op.getBody(), 1), [&](Block &block) {
if (llvm::any_of(blockArgsToDetensor, [&](BlockArgument blockArgument) {
return blockArgument.getOwner() == &block &&
!typeConverter.isLegal(blockArgument.getType());
})) {
return false;
}
return true;
});
});
target.markUnknownOpDynamicallyLegal([&](Operation *op) {
if (isNotBranchOpInterfaceOrReturnLikeOp(op) ||
isLegalForReturnOpTypeConversionPattern(op, typeConverter,
/*returnOpAlwaysLegal*/ true))
@ -577,7 +570,8 @@ struct LinalgDetensorize : public LinalgDetensorizeBase<LinalgDetensorize> {
});
patterns.insert<DetensorizeGenericOp>(typeConverter, context);
patterns.insert<FunctionNonEntryBlockConversion>(context, typeConverter,
patterns.insert<FunctionNonEntryBlockConversion>(FuncOp::getOperationName(),
context, typeConverter,
blockArgsToDetensor);
// Since non-entry block arguments get detensorized, we also need to
// update the control flow inside the function to reflect the correct
@ -591,13 +585,12 @@ struct LinalgDetensorize : public LinalgDetensorizeBase<LinalgDetensorize> {
populateBranchOpInterfaceTypeConversionPattern(patterns, typeConverter,
shouldConvertBranchOperand);
if (failed(
applyFullConversion(getOperation(), target, std::move(patterns))))
if (failed(applyFullConversion(getFunction(), target, std::move(patterns))))
signalPassFailure();
RewritePatternSet canonPatterns(context);
canonPatterns.add<ExtractFromReshapeFromElements>(context);
if (failed(applyPatternsAndFoldGreedily(getOperation(),
if (failed(applyPatternsAndFoldGreedily(getFunction(),
std::move(canonPatterns))))
signalPassFailure();
}

15
mlir/lib/Dialect/Linalg/Transforms/DropUnitDims.cpp
View File

@ -558,23 +558,20 @@ namespace {
/// Pass that removes unit-extent dims within generic ops.
struct LinalgFoldUnitExtentDimsPass
: public LinalgFoldUnitExtentDimsBase<LinalgFoldUnitExtentDimsPass> {
void runOnOperation() override {
auto funcOp = getOperation();
assert(funcOp->hasTrait<OpTrait::FunctionLike>() &&
"LinalgFoldUnitExtentDimsPass can only be run on FunctionLike "
"operations");
MLIRContext *context = funcOp->getContext();
void runOnFunction() override {
FuncOp funcOp = getFunction();
MLIRContext *context = funcOp.getContext();
RewritePatternSet patterns(context);
if (foldOneTripLoopsOnly)
patterns.add<FoldUnitDimLoops>(context);
else
populateFoldUnitExtentDimsPatterns(patterns);
(void)applyPatternsAndFoldGreedily(
function_like_impl::getFunctionBody(funcOp), std::move(patterns));
(void)applyPatternsAndFoldGreedily(funcOp.getBody(), std::move(patterns));
}
};
} // namespace
std::unique_ptr<Pass> mlir::createLinalgFoldUnitExtentDimsPass() {
std::unique_ptr<OperationPass<FuncOp>>
mlir::createLinalgFoldUnitExtentDimsPass() {
return std::make_unique<LinalgFoldUnitExtentDimsPass>();
}

9
mlir/lib/Dialect/Linalg/Transforms/ElementwiseToLinalg.cpp
View File

@ -126,12 +126,8 @@ namespace {
class ConvertElementwiseToLinalgPass
: public ConvertElementwiseToLinalgBase<ConvertElementwiseToLinalgPass> {
void runOnOperation() final {
void runOnFunction() final {
auto func = getOperation();
assert(func->hasTrait<OpTrait::FunctionLike>() &&
"ConvertElementwiseToLinalgPass can only be run on FunctionLike "
"operations");
auto *context = &getContext();
ConversionTarget target(*context);
RewritePatternSet patterns(context);
@ -147,6 +143,7 @@ class ConvertElementwiseToLinalgPass
};
} // namespace
std::unique_ptr<Pass> mlir::createConvertElementwiseToLinalgPass() {
std::unique_ptr<OperationPass<FuncOp>>
mlir::createConvertElementwiseToLinalgPass() {
return std::make_unique<ConvertElementwiseToLinalgPass>();
}

2
mlir/test/Dialect/Linalg/convert-elementwise-to-linalg.mlir
View File

@ -1,4 +1,4 @@
// RUN: mlir-opt -pass-pipeline="builtin.func(convert-elementwise-to-linalg)" -split-input-file %s | FileCheck %s
// RUN: mlir-opt -convert-elementwise-to-linalg -split-input-file %s | FileCheck %s
// In-depth checking of the linalg.generic op for a very trivial case.
// CHECK: #[[$MAP:.*]] = affine_map<() -> ()>

2
mlir/test/Dialect/Linalg/detensorize_0d.mlir
View File

@ -1,4 +1,4 @@
// RUN: mlir-opt %s -allow-unregistered-dialect -pass-pipeline="builtin.func(linalg-detensorize{aggressive-mode})" | FileCheck %s
// RUN: mlir-opt %s -allow-unregistered-dialect -linalg-detensorize=aggressive-mode | FileCheck %s
#map = affine_map<() -> ()>

2
mlir/test/Dialect/Linalg/detensorize_br_operands.mlir
View File

@ -1,4 +1,4 @@
// RUN: mlir-opt %s -split-input-file -allow-unregistered-dialect -pass-pipeline="builtin.func(linalg-detensorize)" | FileCheck %s
// RUN: mlir-opt %s -split-input-file -allow-unregistered-dialect -linalg-detensorize | FileCheck %s
// TODO: Detensoring breaks if %arg0 or %arg1 are passed directly as tensors. Fix that.
func @if_true_test(%arg0: i1, %arg1: i32) -> tensor<i32> attributes {} {

2
mlir/test/Dialect/Linalg/detensorize_if.mlir
View File

@ -1,4 +1,4 @@
// RUN: mlir-opt %s -split-input-file -allow-unregistered-dialect -pass-pipeline="builtin.func(linalg-detensorize)" | FileCheck %s
// RUN: mlir-opt %s -split-input-file -allow-unregistered-dialect -linalg-detensorize | FileCheck %s
#map0 = affine_map<() -> ()>

4
mlir/test/Dialect/Linalg/detensorize_trivial.mlir
View File

@ -1,5 +1,5 @@
// RUN: mlir-opt %s -pass-pipeline="builtin.func(linalg-detensorize{aggressive-mode})" | FileCheck %s -check-prefix=DET-ALL
// RUN: mlir-opt %s -pass-pipeline="builtin.func(linalg-detensorize)" | FileCheck %s -check-prefix=DET-CF
// RUN: mlir-opt %s -linalg-detensorize=aggressive-mode | FileCheck %s -check-prefix=DET-ALL
// RUN: mlir-opt %s -linalg-detensorize | FileCheck %s -check-prefix=DET-CF
#map0 = affine_map<() -> ()>

4
mlir/test/Dialect/Linalg/detensorize_while.mlir
View File

@ -1,5 +1,5 @@
// RUN: mlir-opt %s -pass-pipeline="builtin.func(linalg-detensorize{aggressive-mode})" | FileCheck %s -check-prefix=DET-ALL
// RUN: mlir-opt %s -pass-pipeline="builtin.func(linalg-detensorize)" | FileCheck %s -check-prefix=DET-CF
// RUN: mlir-opt %s -linalg-detensorize=aggressive-mode | FileCheck %s -check-prefix=DET-ALL
// RUN: mlir-opt %s -linalg-detensorize | FileCheck %s -check-prefix=DET-CF
#map0 = affine_map<() -> ()>

4
mlir/test/Dialect/Linalg/detensorize_while_impure_cf.mlir
View File

@ -1,5 +1,5 @@
// RUN: mlir-opt %s -pass-pipeline="builtin.func(linalg-detensorize{aggressive-mode})" | FileCheck %s -check-prefix=DET-ALL
// RUN: mlir-opt %s -pass-pipeline="builtin.func(linalg-detensorize)" | FileCheck %s -check-prefix=DET-CF
// RUN: mlir-opt %s -linalg-detensorize=aggressive-mode | FileCheck %s -check-prefix=DET-ALL
// RUN: mlir-opt %s -linalg-detensorize | FileCheck %s -check-prefix=DET-CF
#map0 = affine_map<() -> ()>
#map1 = affine_map<(i) -> ()>

2
mlir/test/Dialect/Linalg/detensorize_while_pure_cf.mlir
View File

@ -1,4 +1,4 @@
// RUN: mlir-opt %s -allow-unregistered-dialect -pass-pipeline="builtin.func(linalg-detensorize)" | FileCheck %s
// RUN: mlir-opt %s -allow-unregistered-dialect -linalg-detensorize | FileCheck %s
#map0 = affine_map<() -> ()>

2
mlir/test/Dialect/Linalg/drop-unit-extent-dims.mlir
View File

@ -1,4 +1,4 @@
// RUN: mlir-opt %s -split-input-file -pass-pipeline="builtin.func(linalg-fold-unit-extent-dims)" | FileCheck %s
// RUN: mlir-opt %s -split-input-file -linalg-fold-unit-extent-dims | FileCheck %s
#accesses = [
affine_map<(i, j, k, l, m) -> (i, k, m)>,

2
mlir/test/Dialect/Linalg/fold-unit-trip-loops.mlir
View File

@ -1,4 +1,4 @@
// RUN: mlir-opt %s -split-input-file -pass-pipeline="builtin.func(linalg-fold-unit-extent-dims{fold-one-trip-loops-only})" | FileCheck %s
// RUN: mlir-opt %s -split-input-file -linalg-fold-unit-extent-dims="fold-one-trip-loops-only" | FileCheck %s
#accesses = [
affine_map<(i, j, k, l, m) -> (i, k, m)>,