65 lines · cpp
1//===- TestLinalgRankReduceContractionOps.cpp -----------------------------===//2//3// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.4// See https://llvm.org/LICENSE.txt for license information.5// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception6//7//===----------------------------------------------------------------------===//8//9// This file implements a pass for testing rank reduing patterns for named10// contraction ops with unit dims.11//12//===----------------------------------------------------------------------===//13 14#include "mlir/Dialect/Affine/IR/AffineOps.h"15#include "mlir/Dialect/Func/IR/FuncOps.h"16#include "mlir/Dialect/Linalg/Transforms/Transforms.h"17#include "mlir/Pass/Pass.h"18#include "mlir/Pass/PassManager.h"19#include "mlir/Transforms/GreedyPatternRewriteDriver.h"20 21using namespace mlir;22 23namespace {24 25struct TestLinalgRankReduceContractionOps26 : public PassWrapper<TestLinalgRankReduceContractionOps,27 OperationPass<func::FuncOp>> {28 MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(29 TestLinalgRankReduceContractionOps)30 31 TestLinalgRankReduceContractionOps() = default;32 TestLinalgRankReduceContractionOps(33 const TestLinalgRankReduceContractionOps &pass) = default;34 void getDependentDialects(DialectRegistry ®istry) const override {35 registry.insert<affine::AffineDialect, linalg::LinalgDialect,36 memref::MemRefDialect, tensor::TensorDialect>();37 }38 StringRef getArgument() const final {39 return "test-linalg-rank-reduce-contraction-ops";40 }41 StringRef getDescription() const final {42 return "Test Linalg rank reduce contraction ops with unit dims";43 }44 45 void runOnOperation() override {46 MLIRContext *context = &this->getContext();47 func::FuncOp funcOp = this->getOperation();48 49 RewritePatternSet patterns(context);50 linalg::populateContractionOpRankReducingPatterns(patterns);51 if (failed(applyPatternsGreedily(funcOp.getBody(), std::move(patterns))))52 return signalPassFailure();53 }54};55 56} // namespace57 58namespace mlir {59namespace test {60void registerTestLinalgRankReduceContractionOps() {61 PassRegistration<TestLinalgRankReduceContractionOps>();62}63} // namespace test64} // namespace mlir65