| // Copyright 2019 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 |
| |
| #include "iree/compiler/Dialect/VM/Conversion/ImportUtils.h" |
| |
| #include "iree/compiler/Dialect/VM/Conversion/ConversionDialectInterface.h" |
| #include "iree/compiler/Dialect/VM/IR/VMOps.h" |
| #include "llvm/ADT/StringExtras.h" |
| #include "mlir/Dialect/Arith/IR/Arith.h" |
| #include "mlir/IR/Attributes.h" |
| #include "mlir/IR/Builders.h" |
| #include "mlir/IR/BuiltinOps.h" |
| #include "mlir/IR/Matchers.h" |
| #include "mlir/Parser/Parser.h" |
| |
| namespace mlir::iree_compiler { |
| |
| LogicalResult ImportTable::build(Operation *rootOp, |
| const TypeConverter &typeConverter) { |
| for (auto funcOp : rootOp->getRegion(0).getOps<FunctionOpInterface>()) { |
| if (!funcOp.isExternal()) { |
| continue; // only external functions are imports |
| } |
| |
| ImportTable::Import import; |
| import.name = funcOp.getNameAttr(); |
| import.fallback = funcOp->getAttrOfType<SymbolRefAttr>("vm.fallback"); |
| |
| // Try to use an assigned signature or fall back to converting the input. |
| if (auto importOp = dyn_cast<IREE::VM::ImportOp>(funcOp.getOperation())) { |
| // Import ops have their signature used directly. |
| import.signature = importOp.getFunctionType(); |
| } else if (auto signatureAttr = |
| funcOp->getAttrOfType<TypeAttr>("vm.signature")) { |
| // Directly use the specified signature. |
| import.signature = |
| dyn_cast_if_present<FunctionType>(signatureAttr.getValue()); |
| } |
| if (!import.signature) { |
| // Convert the signature using the type converter. |
| SmallVector<Type> argumentTypes; |
| if (failed(typeConverter.convertTypes(funcOp.getArgumentTypes(), |
| argumentTypes))) { |
| return funcOp.emitError() << "unable to convert import argument types"; |
| } |
| SmallVector<Type> resultTypes; |
| if (failed(typeConverter.convertTypes(funcOp.getResultTypes(), |
| resultTypes))) { |
| return funcOp.emitError() << "unable to convert import result types"; |
| } |
| import.signature = |
| FunctionType::get(rootOp->getContext(), argumentTypes, resultTypes); |
| } |
| |
| symbols[import.name.getValue()] = std::move(import); |
| } |
| |
| return success(); |
| } |
| |
| std::optional<ImportTable::Import> ImportTable::find(StringRef symbolName) { |
| auto it = symbols.find(symbolName); |
| if (it == symbols.end()) |
| return std::nullopt; |
| return it->second; |
| } |
| |
| // TODO(benvanik): replace with iree/compiler/Utils/ModuleUtils.h. |
| // There may be some special insertion order arrangement required based on the |
| // nested vm.module here. |
| |
| LogicalResult appendImportModule(IREE::VM::ModuleOp importModuleOp, |
| ModuleOp targetModuleOp) { |
| SymbolTable symbolTable(targetModuleOp); |
| OpBuilder targetBuilder(targetModuleOp); |
| targetBuilder.setInsertionPoint(&targetModuleOp.getBody()->back()); |
| importModuleOp.walk([&](IREE::VM::ImportOp importOp) { |
| std::string fullName = |
| (importModuleOp.getName() + "." + importOp.getName()).str(); |
| if (auto *existingOp = symbolTable.lookup(fullName)) { |
| existingOp->erase(); |
| } |
| auto clonedOp = cast<IREE::VM::ImportOp>(targetBuilder.clone(*importOp)); |
| mlir::StringAttr fullNameAttr = |
| mlir::StringAttr::get(clonedOp.getContext(), fullName); |
| clonedOp.setName(fullNameAttr); |
| clonedOp.setPrivate(); |
| }); |
| return success(); |
| } |
| |
| LogicalResult appendImportModule(StringRef importModuleSrc, |
| ModuleOp targetModuleOp) { |
| auto importModuleRef = mlir::parseSourceString<mlir::ModuleOp>( |
| importModuleSrc, targetModuleOp.getContext()); |
| if (!importModuleRef) { |
| return targetModuleOp.emitError() |
| << "unable to append import module; import module failed to parse"; |
| } |
| for (auto importModuleOp : importModuleRef->getOps<IREE::VM::ModuleOp>()) { |
| if (failed(appendImportModule(importModuleOp, targetModuleOp))) { |
| importModuleOp.emitError() << "failed to import module"; |
| } |
| } |
| return success(); |
| } |
| |
| Value castToImportType(Value value, Type targetType, OpBuilder &builder) { |
| auto sourceType = value.getType(); |
| if (sourceType == targetType) |
| return value; |
| bool sourceIsInteger = isa<IntegerType>(sourceType); |
| |
| // Allow bitcast between same width float/int types. This is used for |
| // marshalling to "untyped" VM interfaces, which will have an integer type. |
| if (isa<FloatType>(sourceType) && isa<IntegerType>(targetType) && |
| sourceType.getIntOrFloatBitWidth() == |
| targetType.getIntOrFloatBitWidth()) { |
| return mlir::arith::BitcastOp::create(builder, value.getLoc(), targetType, |
| value); |
| } else if (sourceIsInteger && |
| (targetType.isSignedInteger() || targetType.isSignlessInteger())) { |
| if (targetType.getIntOrFloatBitWidth() > |
| sourceType.getIntOrFloatBitWidth()) { |
| return mlir::arith::ExtSIOp::create(builder, value.getLoc(), targetType, |
| value); |
| } else { |
| return mlir::arith::TruncIOp::create(builder, value.getLoc(), targetType, |
| value); |
| } |
| } else if (sourceIsInteger && targetType.isUnsignedInteger()) { |
| if (targetType.getIntOrFloatBitWidth() > |
| sourceType.getIntOrFloatBitWidth()) { |
| return mlir::arith::ExtUIOp::create(builder, value.getLoc(), targetType, |
| value); |
| } else { |
| return mlir::arith::TruncIOp::create(builder, value.getLoc(), targetType, |
| value); |
| } |
| } else { |
| return value; |
| } |
| } |
| |
| Value castFromImportType(Value value, Type targetType, OpBuilder &builder) { |
| // Right now the to-import and from-import types are the same. |
| return castToImportType(value, targetType, builder); |
| } |
| |
| void copyImportAttrs(IREE::VM::ImportOp importOp, Operation *callOp) { |
| if (importOp->hasAttr("nosideeffects")) { |
| callOp->setAttr("nosideeffects", UnitAttr::get(importOp.getContext())); |
| } |
| } |
| |
| namespace detail { |
| |
| size_t getSegmentSpanSize(Type spanType) { |
| if (auto tupleType = dyn_cast<TupleType>(spanType)) { |
| return tupleType.size(); |
| } else { |
| return 1; |
| } |
| } |
| |
| std::optional<SmallVector<Value>> rewriteAttrToOperands(Location loc, |
| Attribute attrValue, |
| Type inputType, |
| OpBuilder &builder) { |
| if (auto intAttr = dyn_cast<IntegerAttr>(attrValue)) { |
| // NOTE: we intentionally go to std.constant ops so that the standard |
| // conversions can do their job. If we want to remove the dependency |
| // from standard ops in the future we could instead go directly to |
| // one of the vm constant ops. |
| auto constValue = mlir::arith::ConstantOp::create( |
| builder, loc, inputType, |
| IntegerAttr::get(inputType, APInt(inputType.getIntOrFloatBitWidth(), |
| intAttr.getValue().getSExtValue()))); |
| return {{constValue}}; |
| } else if (auto floatAttr = dyn_cast<FloatAttr>(attrValue)) { |
| bool lossy = false; |
| APFloat value = floatAttr.getValue(); |
| value.convert(cast<FloatType>(inputType).getFloatSemantics(), |
| llvm::RoundingMode::NearestTiesToEven, &lossy); |
| auto constValue = mlir::arith::ConstantOp::create( |
| builder, loc, inputType, FloatAttr::get(inputType, value)); |
| return {{constValue}}; |
| } else if (auto elementsAttr = dyn_cast<DenseIntElementsAttr>(attrValue)) { |
| SmallVector<Value> elementValues; |
| elementValues.reserve(elementsAttr.getNumElements()); |
| for (auto intAttr : elementsAttr.getValues<Attribute>()) { |
| elementValues.push_back(mlir::arith::ConstantOp::create( |
| builder, loc, elementsAttr.getType().getElementType(), |
| cast<TypedAttr>(intAttr))); |
| } |
| return elementValues; |
| } else if (auto arrayAttr = dyn_cast<ArrayAttr>(attrValue)) { |
| SmallVector<Value> allValues; |
| for (auto elementAttr : arrayAttr) { |
| auto flattenedValues = |
| rewriteAttrToOperands(loc, elementAttr, inputType, builder); |
| if (!flattenedValues) |
| return std::nullopt; |
| allValues.append(flattenedValues->begin(), flattenedValues->end()); |
| } |
| return allValues; |
| } else if (auto strAttr = dyn_cast<StringAttr>(attrValue)) { |
| return {{IREE::VM::RodataInlineOp::create(builder, loc, strAttr)}}; |
| } |
| |
| // This may be a custom dialect type. As we can't trivially access the storage |
| // of these we need to ask the dialect to do it for us. |
| auto *conversionInterface = |
| attrValue.getDialect() |
| .getRegisteredInterface<VMConversionDialectInterface>(); |
| if (conversionInterface) { |
| bool anyFailed = false; |
| SmallVector<Value> allValues; |
| if (auto tupleType = dyn_cast<TupleType>(inputType)) { |
| // Custom dialect type maps into a tuple; we expect 1:1 tuple elements to |
| // attribute storage elements. |
| auto tupleTypes = llvm::to_vector(tupleType.getTypes()); |
| int ordinal = 0; |
| LogicalResult walkStatus = conversionInterface->walkAttributeStorage( |
| attrValue, [&](Attribute elementAttr) { |
| if (anyFailed) |
| return; |
| auto elementType = tupleTypes[ordinal++]; |
| auto flattenedValues = |
| rewriteAttrToOperands(loc, elementAttr, elementType, builder); |
| if (!flattenedValues) { |
| anyFailed = true; |
| return; |
| } |
| allValues.append(flattenedValues->begin(), flattenedValues->end()); |
| }); |
| if (failed(walkStatus)) |
| return std::nullopt; |
| } else { |
| // Custom dialect type maps into zero or more input types (ala arrays). |
| LogicalResult walkStatus = conversionInterface->walkAttributeStorage( |
| attrValue, [&](Attribute elementAttr) { |
| if (anyFailed) |
| return; |
| auto flattenedValues = |
| rewriteAttrToOperands(loc, elementAttr, inputType, builder); |
| if (!flattenedValues) { |
| anyFailed = true; |
| return; |
| } |
| allValues.append(flattenedValues->begin(), flattenedValues->end()); |
| }); |
| if (failed(walkStatus)) |
| return std::nullopt; |
| } |
| if (anyFailed) |
| return std::nullopt; |
| return allValues; |
| } |
| |
| emitError(loc) << "unsupported attribute encoding: " << attrValue; |
| return std::nullopt; |
| } |
| |
| } // namespace detail |
| |
| } // namespace mlir::iree_compiler |
