compiler/src/iree/compiler/GlobalOptimization/GeneralizeLinalgNamedOps.cpp

// Copyright 2023 The IREE Authors
//
// Licensed 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

//===- GeneralizeLinalgOps.cpp - Pass to generalize named LinalgOps -------==//
//
// The pass is to generalize Linalg named operations that are better off being
// represented as `linalg.generic` operations in IREE.
//
//===----------------------------------------------------------------------===//

#include "iree/compiler/Dialect/Flow/Transforms/RegionOpUtils.h"
#include "iree/compiler/GlobalOptimization/Passes.h"
#include "mlir/Dialect/Linalg/IR/Linalg.h"
#include "mlir/Dialect/Linalg/Transforms/Transforms.h"
#include "mlir/Pass/Pass.h"

#define DEBUG_TYPE "iree-global-opt-generalize-linalg-named-ops"

namespace mlir::iree_compiler::GlobalOptimization {

#define GEN_PASS_DEF_GENERALIZELINALGNAMEDOPSPASS
#include "iree/compiler/GlobalOptimization/Passes.h.inc"

namespace {
struct GeneralizeLinalgNamedOpsPass
    : impl::GeneralizeLinalgNamedOpsPassBase<GeneralizeLinalgNamedOpsPass> {
  using Base::Base;
  void runOnOperation() override;
};
} // namespace

void GeneralizeLinalgNamedOpsPass::runOnOperation() {
  mlir::FunctionOpInterface funcOp = getOperation();
  SmallVector<linalg::LinalgOp> namedOpCandidates;
  funcOp.walk([&](linalg::LinalgOp linalgOp) {
    if (!IREE::Flow::isNonNullAndOutsideDispatch(linalgOp) ||
        isa<linalg::GenericOp>(linalgOp)) {
      return;
    }
    if (enableGeneralizeMatmul && linalg::isaContractionOpInterface(linalgOp)) {
      namedOpCandidates.push_back(linalgOp);
      return;
    }
    if (isa_and_nonnull<linalg::BroadcastOp, linalg::CopyOp,
                        linalg::ElementwiseOp, linalg::MapOp, linalg::ReduceOp,
                        linalg::TransposeOp>(linalgOp.getOperation()) ||
        linalg::isaConvolutionOpInterface(linalgOp)) {
      namedOpCandidates.push_back(linalgOp);
    }
  });

  IRRewriter rewriter(&getContext());
  for (auto linalgOp : namedOpCandidates) {
    rewriter.setInsertionPoint(linalgOp);
    FailureOr<linalg::LinalgOp> generalizedOp =
        linalg::generalizeNamedOp(rewriter, linalgOp);
    if (failed(generalizedOp)) {
      linalgOp->emitOpError("failed to generalize operation");
      return signalPassFailure();
    }
  }
}

} // namespace mlir::iree_compiler::GlobalOptimization