169 lines · cpp
1//===- Dialect.cpp - Implementation of the linalg dialect and types -------===//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 dialect types and dialect.10//11//===----------------------------------------------------------------------===//12 13#include "mlir/Dialect/Affine/IR/AffineOps.h"14#include "mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h"15#include "mlir/Dialect/Linalg/IR/Linalg.h"16#include "mlir/Dialect/Math/IR/Math.h"17#include "mlir/Dialect/MemRef/IR/MemRef.h"18#include "mlir/Dialect/Shard/Interfaces/ShardingInterface.h"19#include "mlir/Dialect/Tensor/IR/Tensor.h"20#include "mlir/IR/DialectImplementation.h"21#include "mlir/Interfaces/SubsetOpInterface.h"22#include "mlir/Interfaces/ValueBoundsOpInterface.h"23#include "mlir/Support/LLVM.h"24#include "mlir/Transforms/InliningUtils.h"25 26#include "llvm/ADT/TypeSwitch.h"27 28using namespace mlir;29using namespace mlir::linalg;30 31//===----------------------------------------------------------------------===//32// LinalgDialect Dialect Interfaces33//===----------------------------------------------------------------------===//34 35namespace {36 37struct LinalgInlinerInterface : public DialectInlinerInterface {38 using DialectInlinerInterface::DialectInlinerInterface;39 40 // We don't have any special restrictions on what can be inlined into41 // destination regions (e.g. while/conditional bodies). Always allow it.42 bool isLegalToInline(Region *dest, Region *src, bool wouldBeCloned,43 IRMapping &valueMapping) const final {44 return true;45 }46 // Operations in Linalg dialect are always legal to inline.47 bool isLegalToInline(Operation *, Region *, bool, IRMapping &) const final {48 return true;49 }50 // Handle the given inlined terminator by replacing it with a new operation51 // as necessary. Required when the region has only one block.52 void handleTerminator(Operation *op, ValueRange valuesToRepl) const final {}53};54 55} // namespace56 57//===----------------------------------------------------------------------===//58// LinalgDialect59//===----------------------------------------------------------------------===//60 61/// Attribute name used to memoize indexing maps for named ops.62constexpr const ::llvm::StringLiteral63 LinalgDialect::kMemoizedIndexingMapsAttrName;64 65/// Trait to check if T provides a `regionBuilder` method.66template <typename T, typename... Args>67using has_region_builder = decltype(T::regionBuilder);68template <typename T>69using detect_has_region_builder = llvm::is_detected<has_region_builder, T>;70 71/// SFINAE helper for single C++ class without a `regionBuilder` method (e.g.72/// an OpInterface).73template <typename OpType, typename = std::enable_if_t<74 !detect_has_region_builder<OpType>::value>>75void addNamedOpBuilderImpl(76 llvm::StringMap<LinalgDialect::RegionBuilderFunType> &map) {77 // Do nothing.78}79 80template <typename OpType,81 typename = std::enable_if_t<detect_has_region_builder<OpType>::value>,82 typename = void>83void addNamedOpBuilderImpl(84 llvm::StringMap<LinalgDialect::RegionBuilderFunType> &map) {85 map.insert(std::make_pair(86 OpType::getOperationName(),87 static_cast<LinalgDialect::RegionBuilderFunType>(OpType::regionBuilder)));88}89 90template <typename... OpTypes>91void addNamedOpBuilders(92 llvm::StringMap<LinalgDialect::RegionBuilderFunType> &map) {93 (addNamedOpBuilderImpl<OpTypes>(map), ...);94}95 96void mlir::linalg::LinalgDialect::initialize() {97 addAttributes<98#define GET_ATTRDEF_LIST99#include "mlir/Dialect/Linalg/IR/LinalgOpsAttrDefs.cpp.inc"100 >();101 addOperations<102#define GET_OP_LIST103#include "mlir/Dialect/Linalg/IR/LinalgOps.cpp.inc"104 >();105 addOperations<106#define GET_OP_LIST107#include "mlir/Dialect/Linalg/IR/LinalgStructuredOps.cpp.inc"108 >();109 addOperations<110#define GET_OP_LIST111#include "mlir/Dialect/Linalg/IR/LinalgRelayoutOps.cpp.inc"112 >();113 114 // Fill the Linalg-specific OpName to RegionBuilder map.115 addNamedOpBuilders<116#define GET_OP_LIST117#include "mlir/Dialect/Linalg/IR/LinalgStructuredOps.cpp.inc"118 >(namedStructuredOpRegionBuilders);119 120 addInterfaces<LinalgInlinerInterface>();121 122 declarePromisedInterface<shard::ShardingInterface, GenericOp>();123 declarePromisedInterfaces<shard::ShardingInterface,124#define GET_OP_LIST125#include "mlir/Dialect/Linalg/IR/LinalgStructuredOps.cpp.inc"126 >();127 declarePromisedInterface<SubsetOpInterface, CopyOp>();128 declarePromisedInterface<SubsetInsertionOpInterface, CopyOp>();129 130 // ValueBoundsOpInterface131 declarePromisedInterface<ValueBoundsOpInterface, IndexOp>();132 133 declarePromisedInterface<PartialReductionOpInterface, linalg::GenericOp>();134 135 // Tiling Interface136 declarePromisedInterface<TilingInterface, linalg::GenericOp>();137 declarePromisedInterfaces<TilingInterface,138#define GET_OP_LIST139#include "mlir/Dialect/Linalg/IR/LinalgStructuredOps.cpp.inc"140 >();141 declarePromisedInterfaces<TilingInterface,142#define GET_OP_LIST143#include "mlir/Dialect/Linalg/IR/LinalgRelayoutOps.cpp.inc"144 >();145 declarePromisedInterfaces<PartialReductionOpInterface,146#define GET_OP_LIST147#include "mlir/Dialect/Linalg/IR/LinalgStructuredOps.cpp.inc"148 >();149 declarePromisedInterfaces<bufferization::BufferizableOpInterface,150#define GET_OP_LIST151#include "mlir/Dialect/Linalg/IR/LinalgStructuredOps.cpp.inc"152 >();153}154 155LogicalResult LinalgDialect::verifyOperationAttribute(Operation *op,156 NamedAttribute attr) {157 if (attr.getName() == LinalgDialect::kMemoizedIndexingMapsAttrName)158 return success();159 return op->emitError() << "attribute '" << attr.getName()160 << "' not supported by the linalg dialect";161}162 163#include "mlir/Dialect/Linalg/IR/LinalgOpsEnums.cpp.inc"164 165#define GET_ATTRDEF_CLASSES166#include "mlir/Dialect/Linalg/IR/LinalgOpsAttrDefs.cpp.inc"167 168#include "mlir/Dialect/Linalg/IR/LinalgOpsDialect.cpp.inc"169