From aa8815e42e646a98663af4cf036dbb913ad047a7 Mon Sep 17 00:00:00 2001 From: Mehdi Amini Date: Mon, 13 Dec 2021 20:41:25 +0000 Subject: [PATCH] Revert "[NFC] Generalize a couple of passes so they can operate on any FunctionLike op." This reverts commit 34696e6542894ac63dbfb899b0181c539c223ef1. A test is crashing on the mlir-nvidia bot. --- mlir/include/mlir/Dialect/Linalg/Passes.h | 4 +- mlir/include/mlir/Dialect/Linalg/Passes.td | 16 +---- .../Dialect/Linalg/Transforms/Detensorize.cpp | 65 +++++++++---------- .../Linalg/Transforms/DropUnitDims.cpp | 15 ++--- .../Linalg/Transforms/ElementwiseToLinalg.cpp | 9 +-- .../Linalg/convert-elementwise-to-linalg.mlir | 2 +- mlir/test/Dialect/Linalg/detensorize_0d.mlir | 2 +- .../Linalg/detensorize_br_operands.mlir | 2 +- mlir/test/Dialect/Linalg/detensorize_if.mlir | 2 +- .../Dialect/Linalg/detensorize_trivial.mlir | 4 +- .../Dialect/Linalg/detensorize_while.mlir | 4 +- .../Linalg/detensorize_while_impure_cf.mlir | 4 +- .../Linalg/detensorize_while_pure_cf.mlir | 2 +- .../Dialect/Linalg/drop-unit-extent-dims.mlir | 2 +- .../Dialect/Linalg/fold-unit-trip-loops.mlir | 2 +- 15 files changed, 56 insertions(+), 79 deletions(-) diff --git a/mlir/include/mlir/Dialect/Linalg/Passes.h b/mlir/include/mlir/Dialect/Linalg/Passes.h index f21252800af8..b4f3f070673e 100644 --- a/mlir/include/mlir/Dialect/Linalg/Passes.h +++ b/mlir/include/mlir/Dialect/Linalg/Passes.h @@ -19,9 +19,9 @@ namespace mlir { -std::unique_ptr createConvertElementwiseToLinalgPass(); +std::unique_ptr> createConvertElementwiseToLinalgPass(); -std::unique_ptr createLinalgFoldUnitExtentDimsPass(); +std::unique_ptr> createLinalgFoldUnitExtentDimsPass(); std::unique_ptr createLinalgElementwiseOpFusionPass(); std::unique_ptr createFoldReshapeOpsByLinearizationPass(); diff --git a/mlir/include/mlir/Dialect/Linalg/Passes.td b/mlir/include/mlir/Dialect/Linalg/Passes.td index ae1e806d5842..cfba6ac546a1 100644 --- a/mlir/include/mlir/Dialect/Linalg/Passes.td +++ b/mlir/include/mlir/Dialect/Linalg/Passes.td @@ -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. }]; } diff --git a/mlir/lib/Dialect/Linalg/Transforms/Detensorize.cpp b/mlir/lib/Dialect/Linalg/Transforms/Detensorize.cpp index 276d4b728017..0da4f282cde1 100644 --- a/mlir/lib/Dialect/Linalg/Transforms/Detensorize.cpp +++ b/mlir/lib/Dialect/Linalg/Transforms/Detensorize.cpp @@ -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 blockArgsToDetensor) - : ConversionPattern(converter, MatchTraitOpTypeTag(), - TypeID::get(), /*benefit=*/1, - ctx), + : ConversionPattern(converter, functionLikeOpName, /*benefit=*/1, ctx), blockArgsToDetensor(blockArgsToDetensor) {} LogicalResult @@ -236,8 +235,7 @@ struct LinalgDetensorize : public LinalgDetensorizeBase { /// 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 &opsToDetensor, DenseSet &blockArgsToDetensor) = 0; @@ -288,18 +286,18 @@ struct LinalgDetensorize : public LinalgDetensorizeBase { /// AND can be detensored. class ControlFlowDetectionModel : public CostModel { public: - void compute(Operation *func, DetensorizeTypeConverter typeConverter, + void compute(FuncOp func, DetensorizeTypeConverter typeConverter, DenseSet &opsToDetensor, DenseSet &blockArgsToDetensor) override { SmallVector 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 { /// Detensorize everything that can detensored. class AggressiveDetensoringModel : public CostModel { public: - void compute(Operation *func, DetensorizeTypeConverter typeConverter, + void compute(FuncOp func, DetensorizeTypeConverter typeConverter, DenseSet &opsToDetensor, DenseSet &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() && - "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 { 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 { target.addDynamicallyLegalOp( [&](GenericOp op) { return !opsToDetensor.count(op); }); - target.markUnknownOpDynamicallyLegal([&](Operation *op) { + target.addDynamicallyLegalOp([&](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()) { - 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 { }); patterns.insert(typeConverter, context); - patterns.insert(context, typeConverter, + patterns.insert(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 { 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(context); - if (failed(applyPatternsAndFoldGreedily(getOperation(), + if (failed(applyPatternsAndFoldGreedily(getFunction(), std::move(canonPatterns)))) signalPassFailure(); } diff --git a/mlir/lib/Dialect/Linalg/Transforms/DropUnitDims.cpp b/mlir/lib/Dialect/Linalg/Transforms/DropUnitDims.cpp index 6ae5fe1c076e..cfd582dd97a2 100644 --- a/mlir/lib/Dialect/Linalg/Transforms/DropUnitDims.cpp +++ b/mlir/lib/Dialect/Linalg/Transforms/DropUnitDims.cpp @@ -558,23 +558,20 @@ namespace { /// Pass that removes unit-extent dims within generic ops. struct LinalgFoldUnitExtentDimsPass : public LinalgFoldUnitExtentDimsBase { - void runOnOperation() override { - auto funcOp = getOperation(); - assert(funcOp->hasTrait() && - "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(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 mlir::createLinalgFoldUnitExtentDimsPass() { +std::unique_ptr> +mlir::createLinalgFoldUnitExtentDimsPass() { return std::make_unique(); } diff --git a/mlir/lib/Dialect/Linalg/Transforms/ElementwiseToLinalg.cpp b/mlir/lib/Dialect/Linalg/Transforms/ElementwiseToLinalg.cpp index 38b4ab641cce..9ce3dc5fb42f 100644 --- a/mlir/lib/Dialect/Linalg/Transforms/ElementwiseToLinalg.cpp +++ b/mlir/lib/Dialect/Linalg/Transforms/ElementwiseToLinalg.cpp @@ -126,12 +126,8 @@ namespace { class ConvertElementwiseToLinalgPass : public ConvertElementwiseToLinalgBase { - void runOnOperation() final { + void runOnFunction() final { auto func = getOperation(); - assert(func->hasTrait() && - "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 mlir::createConvertElementwiseToLinalgPass() { +std::unique_ptr> +mlir::createConvertElementwiseToLinalgPass() { return std::make_unique(); } diff --git a/mlir/test/Dialect/Linalg/convert-elementwise-to-linalg.mlir b/mlir/test/Dialect/Linalg/convert-elementwise-to-linalg.mlir index 394e33df321f..9ac90d154c8e 100644 --- a/mlir/test/Dialect/Linalg/convert-elementwise-to-linalg.mlir +++ b/mlir/test/Dialect/Linalg/convert-elementwise-to-linalg.mlir @@ -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<() -> ()> diff --git a/mlir/test/Dialect/Linalg/detensorize_0d.mlir b/mlir/test/Dialect/Linalg/detensorize_0d.mlir index 2d73c5b97c45..5e8c7580efe6 100644 --- a/mlir/test/Dialect/Linalg/detensorize_0d.mlir +++ b/mlir/test/Dialect/Linalg/detensorize_0d.mlir @@ -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<() -> ()> diff --git a/mlir/test/Dialect/Linalg/detensorize_br_operands.mlir b/mlir/test/Dialect/Linalg/detensorize_br_operands.mlir index 2682a298dd2a..f6ddc54f7145 100644 --- a/mlir/test/Dialect/Linalg/detensorize_br_operands.mlir +++ b/mlir/test/Dialect/Linalg/detensorize_br_operands.mlir @@ -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 attributes {} { diff --git a/mlir/test/Dialect/Linalg/detensorize_if.mlir b/mlir/test/Dialect/Linalg/detensorize_if.mlir index c9e843bc7d69..e780829dfe67 100644 --- a/mlir/test/Dialect/Linalg/detensorize_if.mlir +++ b/mlir/test/Dialect/Linalg/detensorize_if.mlir @@ -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<() -> ()> diff --git a/mlir/test/Dialect/Linalg/detensorize_trivial.mlir b/mlir/test/Dialect/Linalg/detensorize_trivial.mlir index 5862327ebe6c..513509b4b707 100644 --- a/mlir/test/Dialect/Linalg/detensorize_trivial.mlir +++ b/mlir/test/Dialect/Linalg/detensorize_trivial.mlir @@ -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<() -> ()> diff --git a/mlir/test/Dialect/Linalg/detensorize_while.mlir b/mlir/test/Dialect/Linalg/detensorize_while.mlir index 9ece0029737c..54f459d57c48 100644 --- a/mlir/test/Dialect/Linalg/detensorize_while.mlir +++ b/mlir/test/Dialect/Linalg/detensorize_while.mlir @@ -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<() -> ()> diff --git a/mlir/test/Dialect/Linalg/detensorize_while_impure_cf.mlir b/mlir/test/Dialect/Linalg/detensorize_while_impure_cf.mlir index 765692fa2d3d..cf380db3c19f 100644 --- a/mlir/test/Dialect/Linalg/detensorize_while_impure_cf.mlir +++ b/mlir/test/Dialect/Linalg/detensorize_while_impure_cf.mlir @@ -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) -> ()> diff --git a/mlir/test/Dialect/Linalg/detensorize_while_pure_cf.mlir b/mlir/test/Dialect/Linalg/detensorize_while_pure_cf.mlir index 5ac7ab520680..1b720160affe 100644 --- a/mlir/test/Dialect/Linalg/detensorize_while_pure_cf.mlir +++ b/mlir/test/Dialect/Linalg/detensorize_while_pure_cf.mlir @@ -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<() -> ()> diff --git a/mlir/test/Dialect/Linalg/drop-unit-extent-dims.mlir b/mlir/test/Dialect/Linalg/drop-unit-extent-dims.mlir index 79765ca9dac7..634f38455192 100644 --- a/mlir/test/Dialect/Linalg/drop-unit-extent-dims.mlir +++ b/mlir/test/Dialect/Linalg/drop-unit-extent-dims.mlir @@ -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)>, diff --git a/mlir/test/Dialect/Linalg/fold-unit-trip-loops.mlir b/mlir/test/Dialect/Linalg/fold-unit-trip-loops.mlir index 7654e520a855..d0c526e441b6 100644 --- a/mlir/test/Dialect/Linalg/fold-unit-trip-loops.mlir +++ b/mlir/test/Dialect/Linalg/fold-unit-trip-loops.mlir @@ -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)>,