166 lines · cpp
1//===- Linalg.cpp - C Interface for Linalg dialect ------------------------===//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-c/Dialect/Linalg.h"10#include "mlir/CAPI/AffineMap.h"11#include "mlir/CAPI/Registration.h"12#include "mlir/Dialect/Linalg/IR/Linalg.h"13 14using namespace mlir;15using namespace mlir::linalg;16 17/// Apply the special region builder for the builtin named Linalg op.18/// Assert that `op` is a builtin named Linalg op.19void mlirLinalgFillBuiltinNamedOpRegion(MlirOperation mlirOp) {20 Operation *op = unwrap(mlirOp);21 auto linalgOp = cast<LinalgOp>(op);22 auto *dialect = static_cast<LinalgDialect *>(linalgOp->getDialect());23 LinalgDialect::RegionBuilderFunType fun =24 dialect->getRegionBuilder(op->getName().getStringRef());25 26 assert(fun && "Expected a builtin named Linalg op.");27 assert(op->getNumRegions() == 1 && "Expected Linalg op with 1 region");28 assert(op->getRegion(0).getBlocks().empty() &&29 "Expected Linalg op with 0 blocks");30 31 SmallVector<Type, 8> argTypes;32 SmallVector<Location, 8> argLocs;33 for (OpOperand &opOperand : linalgOp->getOpOperands()) {34 argTypes.push_back(getElementTypeOrSelf(opOperand.get().getType()));35 argLocs.push_back(opOperand.get().getLoc());36 }37 38 ImplicitLocOpBuilder b(op->getLoc(), op->getContext());39 Region ®ion = op->getRegion(0);40 Block *body = b.createBlock(®ion, /*insertPt=*/{}, argTypes, argLocs);41 b.setInsertionPointToStart(body);42 fun(b, *body, op->getAttrs(), /*emitError=*/{});43}44 45MLIR_CAPI_EXPORTED bool mlirLinalgIsAContractionOp(MlirOperation op) {46 auto linalgOp = llvm::dyn_cast<mlir::linalg::LinalgOp>(unwrap(op));47 // isaContractionOpInterface handles null linalgOp internally.48 return linalg::isaContractionOpInterface(linalgOp);49}50 51MLIR_CAPI_EXPORTED MlirLinalgContractionDimensions52mlirLinalgInferContractionDimensions(MlirOperation op) {53 MlirLinalgContractionDimensions result{};54 auto linalgOp = dyn_cast<linalg::LinalgOp>(unwrap(op));55 if (!linalgOp)56 return result;57 58 FailureOr<linalg::ContractionDimensions> maybeDims =59 linalg::inferContractionDims(linalgOp);60 if (failed(maybeDims))61 return result;62 63 const linalg::ContractionDimensions &contractionDims = *maybeDims;64 MLIRContext *ctx = linalgOp.getContext();65 66 auto toAttr = [ctx](ArrayRef<unsigned> vals) -> MlirAttribute {67 return wrap(DenseI32ArrayAttr::get(ctx, llvm::to_vector_of<int32_t>(vals)));68 };69 70 result.batch = toAttr(contractionDims.batch);71 result.m = toAttr(contractionDims.m);72 result.n = toAttr(contractionDims.n);73 result.k = toAttr(contractionDims.k);74 75 return result;76}77 78MLIR_CAPI_EXPORTED MlirLinalgContractionDimensions79mlirLinalgInferContractionDimensionsFromMaps(const MlirAffineMap *indexingMaps,80 size_t numMaps) {81 MlirLinalgContractionDimensions result{};82 if (!indexingMaps || numMaps == 0)83 return result;84 85 SmallVector<AffineMap, 3> maps;86 maps.reserve(numMaps);87 for (size_t i = 0; i < numMaps; ++i) {88 maps.push_back(unwrap(indexingMaps[i]));89 }90 91 FailureOr<linalg::ContractionDimensions> maybeDims =92 linalg::inferContractionDims(maps);93 if (failed(maybeDims))94 return result;95 96 MLIRContext *ctx = maps[0].getContext();97 98 auto toAttr = [ctx](ArrayRef<unsigned> vals) -> MlirAttribute {99 return wrap(DenseI32ArrayAttr::get(ctx, llvm::to_vector_of<int32_t>(vals)));100 };101 102 result.batch = toAttr(maybeDims->batch);103 result.m = toAttr(maybeDims->m);104 result.n = toAttr(maybeDims->n);105 result.k = toAttr(maybeDims->k);106 107 return result;108}109 110MLIR_CAPI_EXPORTED bool mlirLinalgIsAConvolutionOp(MlirOperation op) {111 auto linalgOp = llvm::dyn_cast<mlir::linalg::LinalgOp>(unwrap(op));112 if (!linalgOp)113 return false;114 115 return linalg::isaConvolutionOpInterface(linalgOp);116}117 118MLIR_CAPI_EXPORTED MlirLinalgConvolutionDimensions119mlirLinalgInferConvolutionDimensions(MlirOperation op) {120 MlirLinalgConvolutionDimensions result{};121 auto linalgOp = llvm::dyn_cast<mlir::linalg::LinalgOp>(unwrap(op));122 if (!linalgOp)123 return result;124 125 FailureOr<linalg::ConvolutionDimensions> maybeDims =126 linalg::inferConvolutionDims(linalgOp);127 if (failed(maybeDims))128 return result;129 130 const linalg::ConvolutionDimensions &dims = *maybeDims;131 MLIRContext *ctx = linalgOp.getContext();132 133 auto toI32Attr =134 [&ctx](const SmallVector<unsigned, 2> &vals) -> MlirAttribute {135 return wrap(DenseI32ArrayAttr::get(ctx, llvm::to_vector_of<int32_t>(vals)));136 };137 138 auto toI64Attr =139 [&ctx](const SmallVector<int64_t, 2> &vals) -> MlirAttribute {140 return wrap(DenseI64ArrayAttr::get(ctx, vals));141 };142 143 result.batch = toI32Attr(dims.batch);144 result.outputImage = toI32Attr(dims.outputImage);145 result.outputChannel = toI32Attr(dims.outputChannel);146 result.filterLoop = toI32Attr(dims.filterLoop);147 result.inputChannel = toI32Attr(dims.inputChannel);148 result.depth = toI32Attr(dims.depth);149 result.strides = toI64Attr(dims.strides);150 result.dilations = toI64Attr(dims.dilations);151 152 return result;153}154 155MLIR_CAPI_EXPORTED MlirAttribute156mlirLinalgGetIndexingMapsAttribute(MlirOperation op) {157 auto linalgOp = llvm::dyn_cast<mlir::linalg::LinalgOp>(unwrap(op));158 if (!linalgOp)159 return MlirAttribute{nullptr};160 161 ArrayAttr attr = linalgOp.getIndexingMaps();162 return wrap(attr);163}164 165MLIR_DEFINE_CAPI_DIALECT_REGISTRATION(Linalg, linalg, LinalgDialect)166