3139cc766c
This patch adds a pattern to decompose a `linalg.generic` operations that - has only parallel iterator types - has more than 2 statements (including the yield) into multiple `linalg.generic` operation such that each operation has a single statement and a yield. The pattern added here just splits the matching `linalg.generic` into two `linalg.generic`s, one containing the first statement, and the other containing the remaining. The same pattern can be applied repeatedly on the second op to ultimately fully decompose the generic op. Differential Revision: https://reviews.llvm.org/D129704
55 lines
1.9 KiB
C++
55 lines
1.9 KiB
C++
//===- TestLinalgDecomposeOps.cpp - Test Linalg decomposition ------------===//
|
|
//
|
|
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
|
|
// See https://llvm.org/LICENSE.txt for license information.
|
|
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
|
|
//
|
|
//===----------------------------------------------------------------------===//
|
|
//
|
|
// This file implements a pass for testing decomposition of Linalg ops.
|
|
//
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
#include "mlir/Dialect/Affine/IR/AffineOps.h"
|
|
#include "mlir/Dialect/Linalg/Transforms/Transforms.h"
|
|
#include "mlir/Pass/Pass.h"
|
|
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
|
|
|
|
using namespace mlir;
|
|
|
|
namespace {
|
|
struct TestLinalgDecomposeOps
|
|
: public PassWrapper<TestLinalgDecomposeOps, OperationPass<>> {
|
|
MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestLinalgDecomposeOps)
|
|
|
|
TestLinalgDecomposeOps() = default;
|
|
TestLinalgDecomposeOps(const TestLinalgDecomposeOps &pass)
|
|
: PassWrapper(pass) {}
|
|
void getDependentDialects(DialectRegistry ®istry) const override {
|
|
registry.insert<AffineDialect, linalg::LinalgDialect>();
|
|
}
|
|
StringRef getArgument() const final { return "test-linalg-decompose-ops"; }
|
|
StringRef getDescription() const final {
|
|
return "Test Linalg decomposition patterns";
|
|
}
|
|
|
|
void runOnOperation() override {
|
|
MLIRContext *context = &this->getContext();
|
|
RewritePatternSet decompositionPatterns(context);
|
|
linalg::populateDecomposeLinalgOpsPattern(decompositionPatterns);
|
|
if (failed(applyPatternsAndFoldGreedily(
|
|
getOperation(), std::move(decompositionPatterns)))) {
|
|
return signalPassFailure();
|
|
}
|
|
}
|
|
};
|
|
} // namespace
|
|
|
|
namespace mlir {
|
|
namespace test {
|
|
void registerTestLinalgDecomposeOps() {
|
|
PassRegistration<TestLinalgDecomposeOps>();
|
|
}
|
|
} // namespace test
|
|
} // namespace mlir
|