62 lines · cpp
1//===- DialectExtension.cpp - Linalg transform dialect extension ----------===//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#include "mlir/Dialect/Linalg/TransformOps/DialectExtension.h"10#include "mlir/Dialect/Affine/IR/AffineOps.h"11#include "mlir/Dialect/Arith/IR/Arith.h"12#include "mlir/Dialect/GPU/IR/GPUDialect.h"13#include "mlir/Dialect/Index/IR/IndexDialect.h"14#include "mlir/Dialect/Linalg/IR/Linalg.h"15#include "mlir/Dialect/Linalg/TransformOps/LinalgMatchOps.h"16#include "mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.h"17#include "mlir/Dialect/SCF/IR/SCF.h"18#include "mlir/Dialect/Tensor/IR/Tensor.h"19#include "mlir/Dialect/Transform/IR/TransformDialect.h"20#include "mlir/Dialect/Vector/IR/VectorOps.h"21 22using namespace mlir;23 24namespace {25/// Registers new ops and declares PDL as dependent dialect since the26/// additional ops are using PDL types for operands and results.27class LinalgTransformDialectExtension28 : public transform::TransformDialectExtension<29 LinalgTransformDialectExtension> {30public:31 MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(LinalgTransformDialectExtension)32 33 using Base::Base;34 35 void init() {36 declareDependentDialect<linalg::LinalgDialect>();37 38 declareGeneratedDialect<affine::AffineDialect>();39 declareGeneratedDialect<arith::ArithDialect>();40 declareGeneratedDialect<index::IndexDialect>();41 declareGeneratedDialect<scf::SCFDialect>();42 declareGeneratedDialect<vector::VectorDialect>();43 declareGeneratedDialect<gpu::GPUDialect>();44 declareGeneratedDialect<tensor::TensorDialect>();45 46 registerTransformOps<47#define GET_OP_LIST48#include "mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp.inc"49 >();50 registerTransformOps<51#define GET_OP_LIST52#include "mlir/Dialect/Linalg/TransformOps/LinalgMatchOps.cpp.inc"53 >();54 }55};56} // namespace57 58void mlir::linalg::registerTransformDialectExtension(59 DialectRegistry ®istry) {60 registry.addExtensions<LinalgTransformDialectExtension>();61}62