93 lines · cpp
1//===- Generalization.cpp - linalg named ops to generic ops --------------===//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 the Linalg generalization pass. It converts named10// Linalg ops to linalg.generic ops.11//12//===----------------------------------------------------------------------===//13 14#include "mlir/Dialect/Linalg/Passes.h"15 16#include "mlir/Dialect/Linalg/IR/Linalg.h"17#include "mlir/Dialect/Linalg/Transforms/Transforms.h"18#include "mlir/IR/AffineMap.h"19#include "mlir/IR/Builders.h"20#include "mlir/IR/PatternMatch.h"21#include "mlir/Transforms/GreedyPatternRewriteDriver.h"22 23namespace mlir {24#define GEN_PASS_DEF_LINALGGENERALIZENAMEDOPSPASS25#include "mlir/Dialect/Linalg/Passes.h.inc"26} // namespace mlir27 28#define DEBUG_TYPE "linalg-generalization"29 30using namespace mlir;31using namespace mlir::linalg;32 33static LogicalResult generalizeNamedOpPrecondition(LinalgOp linalgOp) {34 // Bailout if `linalgOp` is already a generic.35 if (isa<GenericOp>(linalgOp))36 return failure();37 // Check if the operation has exactly one region.38 if (linalgOp->getNumRegions() != 1) {39 assert(linalgOp->getNumRegions() == 0 && "op with multiple regions");40 // TOD: Otherwise it needs to be built explicitly from the region builder.41 return failure();42 }43 return success();44}45 46FailureOr<GenericOp> mlir::linalg::generalizeNamedOp(RewriterBase &rewriter,47 LinalgOp linalgOp) {48 if (failed(generalizeNamedOpPrecondition(linalgOp)))49 return rewriter.notifyMatchFailure(linalgOp, "preconditions not met");50 51 SmallVector<Value> inputs = linalgOp.getDpsInputs();52 ValueRange outputs = linalgOp.getDpsInits();53 SmallVector<AffineMap> indexingMaps = linalgOp.getIndexingMapsArray();54 SmallVector<utils::IteratorType> iterators = linalgOp.getIteratorTypesArray();55 SmallVector<Type> resultTypes = linalgOp.hasPureTensorSemantics()56 ? TypeRange(ValueRange(outputs))57 : TypeRange{};58 59 // All named ops have a region attached that can be inlined.60 assert(linalgOp->getNumRegions() == 1 &&61 "expect named op to have one region attached");62 GenericOp genericOp =63 GenericOp::create(rewriter, linalgOp.getLoc(), resultTypes, inputs,64 outputs, indexingMaps, iterators);65 rewriter.inlineRegionBefore(linalgOp->getRegion(0), genericOp.getRegion(),66 genericOp.getRegion().begin());67 rewriter.replaceOp(linalgOp, genericOp->getResults());68 return genericOp;69}70 71namespace {72 73struct LinalgGeneralizeNamedOpsPass74 : public impl::LinalgGeneralizeNamedOpsPassBase<75 LinalgGeneralizeNamedOpsPass> {76 using impl::LinalgGeneralizeNamedOpsPassBase<77 LinalgGeneralizeNamedOpsPass>::LinalgGeneralizeNamedOpsPassBase;78 void runOnOperation() override;79};80 81} // namespace82 83void LinalgGeneralizeNamedOpsPass::runOnOperation() {84 RewritePatternSet patterns(&getContext());85 populateLinalgNamedOpsGeneralizationPatterns(patterns);86 (void)applyPatternsGreedily(getOperation(), std::move(patterns));87}88 89void mlir::linalg::populateLinalgNamedOpsGeneralizationPatterns(90 RewritePatternSet &patterns) {91 patterns.add<LinalgGeneralizationPattern>(patterns.getContext());92}93