252 lines · cpp
1//===- ArmGraphOps.cpp - MLIR SPIR-V SPV_ARM_graph operations -------------===//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 defines the SPV_ARM_graph operations in the SPIR-V dialect.10//11//===----------------------------------------------------------------------===//12 13#include "mlir/Dialect/SPIRV/IR/SPIRVOps.h"14 15#include "SPIRVParsingUtils.h"16 17#include "mlir/Dialect/SPIRV/IR/SPIRVDialect.h"18#include "mlir/Dialect/SPIRV/IR/SPIRVEnums.h"19#include "mlir/Dialect/SPIRV/IR/SPIRVTypes.h"20#include "mlir/IR/Builders.h"21#include "mlir/IR/BuiltinTypes.h"22#include "mlir/IR/Operation.h"23#include "mlir/Interfaces/FunctionImplementation.h"24#include "llvm/Support/InterleavedRange.h"25 26using namespace mlir;27using namespace mlir::spirv::AttrNames;28 29//===----------------------------------------------------------------------===//30// spirv.GraphARM31//===----------------------------------------------------------------------===//32 33ParseResult spirv::GraphARMOp::parse(OpAsmParser &parser,34 OperationState &result) {35 Builder &builder = parser.getBuilder();36 37 // Parse the name as a symbol.38 StringAttr nameAttr;39 if (parser.parseSymbolName(nameAttr, SymbolTable::getSymbolAttrName(),40 result.attributes))41 return failure();42 43 // Parse the function signature.44 bool isVariadic = false;45 SmallVector<OpAsmParser::Argument> entryArgs;46 SmallVector<Type> resultTypes;47 SmallVector<DictionaryAttr> resultAttrs;48 if (function_interface_impl::parseFunctionSignatureWithArguments(49 parser, /*allowVariadic=*/false, entryArgs, isVariadic, resultTypes,50 resultAttrs))51 return failure();52 53 SmallVector<Type> argTypes = llvm::map_to_vector(54 entryArgs, [](const OpAsmParser::Argument &arg) { return arg.type; });55 GraphType grType = builder.getGraphType(argTypes, resultTypes);56 result.addAttribute(getFunctionTypeAttrName(result.name),57 TypeAttr::get(grType));58 59 // If additional attributes are present, parse them.60 if (parser.parseOptionalAttrDictWithKeyword(result.attributes))61 return failure();62 63 // Add the attributes to the function arguments.64 assert(resultAttrs.size() == resultTypes.size());65 call_interface_impl::addArgAndResultAttrs(66 builder, result, entryArgs, resultAttrs, getArgAttrsAttrName(result.name),67 getResAttrsAttrName(result.name));68 69 // Parse the optional function body.70 Region *body = result.addRegion();71 OptionalParseResult parseResult =72 parser.parseOptionalRegion(*body, entryArgs);73 return failure(parseResult.has_value() && failed(*parseResult));74}75 76void spirv::GraphARMOp::print(OpAsmPrinter &printer) {77 // Print graph name, signature, and control.78 printer << " ";79 printer.printSymbolName(getSymName());80 GraphType grType = getFunctionType();81 function_interface_impl::printFunctionSignature(82 printer, *this, grType.getInputs(),83 /*isVariadic=*/false, grType.getResults());84 function_interface_impl::printFunctionAttributes(printer, *this,85 {getFunctionTypeAttrName(),86 getArgAttrsAttrName(),87 getResAttrsAttrName()});88 89 // Print the body.90 Region &body = this->getBody();91 if (!body.empty()) {92 printer << ' ';93 printer.printRegion(body, /*printEntryBlockArgs=*/false,94 /*printBlockTerminators=*/true);95 }96}97 98LogicalResult spirv::GraphARMOp::verifyType() {99 if (getFunctionType().getNumResults() < 1)100 return emitOpError("there should be at least one result");101 return success();102}103 104LogicalResult spirv::GraphARMOp::verifyBody() {105 for (auto [index, graphArgType] : llvm::enumerate(getArgumentTypes())) {106 if (!isa<spirv::TensorArmType>(graphArgType)) {107 return emitOpError("type of argument #")108 << index << " must be a TensorArmType, but got " << graphArgType;109 }110 }111 for (auto [index, graphResType] : llvm::enumerate(getResultTypes())) {112 if (!isa<spirv::TensorArmType>(graphResType)) {113 return emitOpError("type of result #")114 << index << " must be a TensorArmType, but got " << graphResType;115 }116 }117 118 if (!isExternal()) {119 Block &entryBlock = front();120 121 unsigned numArguments = this->getNumArguments();122 if (entryBlock.getNumArguments() != numArguments)123 return emitOpError("entry block must have ")124 << numArguments << " arguments to match graph signature";125 126 for (auto [index, grArgType, blockArgType] :127 llvm::enumerate(getArgumentTypes(), entryBlock.getArgumentTypes())) {128 if (blockArgType != grArgType) {129 return emitOpError("type of entry block argument #")130 << index << '(' << blockArgType131 << ") must match the type of the corresponding argument in "132 << "graph signature(" << grArgType << ')';133 }134 }135 }136 137 GraphType grType = getFunctionType();138 auto walkResult = walk([grType](spirv::GraphOutputsARMOp op) -> WalkResult {139 if (grType.getNumResults() != op.getNumOperands())140 return op.emitOpError("is returning ")141 << op.getNumOperands()142 << " value(s) but enclosing spirv.ARM.Graph requires "143 << grType.getNumResults() << " result(s)";144 145 ValueTypeRange<OperandRange> graphOutputOperandTypes =146 op.getValue().getType();147 for (auto [index, type] : llvm::enumerate(graphOutputOperandTypes)) {148 if (type != grType.getResult(index))149 return op.emitError("type of return operand ")150 << index << " (" << type << ") doesn't match graph result type ("151 << grType.getResult(index) << ")";152 }153 return WalkResult::advance();154 });155 156 return failure(walkResult.wasInterrupted());157}158 159void spirv::GraphARMOp::build(OpBuilder &builder, OperationState &state,160 StringRef name, GraphType type,161 ArrayRef<NamedAttribute> attrs, bool entryPoint) {162 state.addAttribute(SymbolTable::getSymbolAttrName(),163 builder.getStringAttr(name));164 state.addAttribute(getFunctionTypeAttrName(state.name), TypeAttr::get(type));165 state.attributes.append(attrs);166 state.addAttribute(getEntryPointAttrName(state.name),167 builder.getBoolAttr(entryPoint));168 state.addRegion();169}170 171ArrayRef<Type> spirv::GraphARMOp::getArgumentTypes() {172 return getFunctionType().getInputs();173}174 175ArrayRef<Type> spirv::GraphARMOp::getResultTypes() {176 return getFunctionType().getResults();177}178 179Region *spirv::GraphARMOp::getCallableRegion() {180 return isExternal() ? nullptr : &getBody();181}182 183//===----------------------------------------------------------------------===//184// spirv.GraphOutputsARM185//===----------------------------------------------------------------------===//186 187LogicalResult spirv::GraphOutputsARMOp::verify() {188 auto graph = cast<GraphARMOp>((*this)->getParentOp());189 190 // The operand number and types must match the graph signature.191 const ArrayRef<Type> &results = graph.getFunctionType().getResults();192 if (getNumOperands() != results.size())193 return emitOpError("has ")194 << getNumOperands() << " operands, but enclosing spirv.ARM.Graph (@"195 << graph.getName() << ") returns " << results.size();196 197 for (auto [index, result] : llvm::enumerate(results))198 if (getOperand(index).getType() != result)199 return emitError() << "type of return operand " << index << " ("200 << getOperand(index).getType()201 << ") doesn't match spirv.ARM.Graph result type ("202 << result << ")"203 << " in graph @" << graph.getName();204 return success();205}206 207//===----------------------------------------------------------------------===//208// spirv.GraphEntryPointARM209//===----------------------------------------------------------------------===//210 211void spirv::GraphEntryPointARMOp::build(OpBuilder &builder,212 OperationState &state,213 spirv::GraphARMOp graph,214 ArrayRef<Attribute> interfaceVars) {215 build(builder, state, SymbolRefAttr::get(graph),216 builder.getArrayAttr(interfaceVars));217}218 219ParseResult spirv::GraphEntryPointARMOp::parse(OpAsmParser &parser,220 OperationState &result) {221 FlatSymbolRefAttr fn;222 if (parser.parseAttribute(fn, Type(), kFnNameAttrName, result.attributes))223 return failure();224 225 SmallVector<Attribute, 4> interfaceVars;226 if (!parser.parseOptionalComma()) {227 // Parse the interface variables.228 if (parser.parseCommaSeparatedList([&]() -> ParseResult {229 // The name of the interface variable attribute is not important.230 FlatSymbolRefAttr var;231 NamedAttrList attrs;232 if (parser.parseAttribute(var, Type(), "var_symbol", attrs))233 return failure();234 interfaceVars.push_back(var);235 return success();236 }))237 return failure();238 }239 result.addAttribute("interface",240 parser.getBuilder().getArrayAttr(interfaceVars));241 return success();242}243 244void spirv::GraphEntryPointARMOp::print(OpAsmPrinter &printer) {245 printer << " ";246 printer.printSymbolName(getFn());247 ArrayRef<Attribute> interfaceVars = getInterface().getValue();248 if (!interfaceVars.empty()) {249 printer << ", " << llvm::interleaved(interfaceVars);250 }251}252