Remove memref runtime, move the ABI to a fixed form and construct descriptors at compile time (#2635)
diff --git a/iree/compiler/Conversion/LinalgToLLVM/BUILD b/iree/compiler/Conversion/LinalgToLLVM/BUILD index 41d9988..fcd9476 100644 --- a/iree/compiler/Conversion/LinalgToLLVM/BUILD +++ b/iree/compiler/Conversion/LinalgToLLVM/BUILD
@@ -22,7 +22,6 @@ name = "LinalgToLLVM", srcs = [ "ConvertToLLVM.cpp", - "HALInterfaceToMemrefArguments.cpp", "Passes.cpp", ], hdrs = [
diff --git a/iree/compiler/Conversion/LinalgToLLVM/CMakeLists.txt b/iree/compiler/Conversion/LinalgToLLVM/CMakeLists.txt index d21dc19..bc31e4e 100644 --- a/iree/compiler/Conversion/LinalgToLLVM/CMakeLists.txt +++ b/iree/compiler/Conversion/LinalgToLLVM/CMakeLists.txt
@@ -21,7 +21,6 @@ "Passes.h" SRCS "ConvertToLLVM.cpp" - "HALInterfaceToMemrefArguments.cpp" "Passes.cpp" DEPS MLIRAffineToStandard
diff --git a/iree/compiler/Conversion/LinalgToLLVM/ConvertToLLVM.cpp b/iree/compiler/Conversion/LinalgToLLVM/ConvertToLLVM.cpp index 193eb3c..8686dca 100644 --- a/iree/compiler/Conversion/LinalgToLLVM/ConvertToLLVM.cpp +++ b/iree/compiler/Conversion/LinalgToLLVM/ConvertToLLVM.cpp
@@ -124,6 +124,180 @@ } }; +/// Returns true if `aOp` has a desciptor (set, binding) pair smaller than +/// `bOp`. Note that this ignores the offset. +bool operator<(IREE::HAL::InterfaceBindingOp aOp, + IREE::HAL::InterfaceBindingOp bOp) { + if (aOp.set().getZExtValue() == bOp.set().getZExtValue()) + return aOp.binding().getZExtValue() < bOp.binding().getZExtValue(); + return aOp.set().getZExtValue() < bOp.set().getZExtValue(); +} + +// Change signature of entry function to func +// entry_func(%packed_buffers_arg_ptr: +// !<llvm.int8**>, %push_constant: !<llvm.int64*>) and lower IREE and HAL ops to +// corresponding LLVMIR ops to construct memref descriptors and load +// push_constant values. +class ConvertFuncWithHALInterface : public ConvertToLLVMPattern { + public: + explicit ConvertFuncWithHALInterface(MLIRContext *context, + LLVMTypeConverter &typeconverter) + : ConvertToLLVMPattern(FuncOp::getOperationName(), context, + typeconverter) {} + + LogicalResult matchAndRewrite( + Operation *op, ArrayRef<Value> operands, + ConversionPatternRewriter &rewriter) const override { + if (SymbolTable::getSymbolVisibility(op) != SymbolTable::Visibility::Public) + return failure(); + auto funcOp = dyn_cast_or_null<FuncOp>(op); + FunctionType fnType = funcOp.getType(); + if (fnType.getNumInputs() != 0) { + return rewriter.notifyMatchFailure( + funcOp, "entry function should not have inputs"); + } + + // Get interface buffers from all the blocks. + SmallVector<IREE::PlaceholderOp, 8> bufferOps; + SmallVector<IREE::HAL::InterfaceLoadConstantOp, 8> loadOps; + for (Block &block : funcOp.getBlocks()) { + for (Operation &op : block) { + if (auto phOp = dyn_cast<IREE::PlaceholderOp>(op)) + bufferOps.push_back(phOp); + if (auto phOp = dyn_cast<IREE::HAL::InterfaceLoadConstantOp>(op)) { + loadOps.push_back(phOp); + } + } + } + + if (bufferOps.empty()) return failure(); + + // A map from buffer ops to their corresponding interface binding ops. + llvm::DenseMap<Operation *, IREE::HAL::InterfaceBindingOp> bufferBindingMap; + for (auto bufferOp : bufferOps) { + auto symbol = SymbolTable::lookupNearestSymbolFrom( + bufferOp, bufferOp.getAttrOfType<SymbolRefAttr>("binding")); + bufferBindingMap[bufferOp] = cast<IREE::HAL::InterfaceBindingOp>(symbol); + } + + // Sort buffers according to their descriptor (set, binding) pair. + llvm::sort(bufferOps, [&bufferBindingMap](IREE::PlaceholderOp aBuffer, + IREE::PlaceholderOp bBuffer) { + return bufferBindingMap[aBuffer] < bufferBindingMap[bBuffer]; + }); + + // A map from buffer ops to their corresponding function argument indices. + llvm::DenseMap<Operation *, unsigned> bufferArgMap; + // A map from binding ops to their corresponding function argument indices. + llvm::DenseMap<Operation *, unsigned> bindingArgMap; + llvm::SmallVector<MemRefType, 4> inputMemRefTypes; + llvm::SmallVector<LLVM::LLVMType, 4> inputStructPtrs; + unsigned argIndex = 0; + for (auto bufferOp : bufferOps) { + auto binding = bufferBindingMap[bufferOp]; + auto it = bindingArgMap.find(binding); + if (it != bindingArgMap.end()) { + bufferArgMap[bufferOp] = it->second; + } else { + bindingArgMap[binding] = argIndex; + bufferArgMap[bufferOp] = argIndex; + ++argIndex; + } + + auto memrefType = bufferOp.getType().dyn_cast_or_null<MemRefType>(); + inputMemRefTypes.push_back(memrefType); + auto elementType = typeConverter.convertType(memrefType.getElementType()) + .dyn_cast<LLVM::LLVMType>(); + if (!elementType) return failure(); + inputStructPtrs.push_back( + elementType.getPointerTo(memrefType.getMemorySpace())); + } + + TypeConverter::SignatureConversion signatureConverter(/*numOrigInputs=*/0); + + // func foo(%packed_buffer_args: !llvm<i8**>, %push_constant: !llvm<i64*>) + auto packedBuffersArgsTy = + LLVM::LLVMType::getInt8PtrTy(typeConverter.getDialect()).getPointerTo(); + auto pushConstantArgTy = + LLVM::LLVMType::getInt64Ty(typeConverter.getDialect()).getPointerTo(); + signatureConverter.addInputs(packedBuffersArgsTy); + signatureConverter.addInputs(pushConstantArgTy); + + // Create the new function's signature. + Location loc = funcOp.getLoc(); + auto newFuncOp = rewriter.create<FuncOp>( + loc, funcOp.getName(), + rewriter.getFunctionType(signatureConverter.getConvertedTypes(), + llvm::None), + ArrayRef<NamedAttribute>()); + + // Move all ops in the old function's region to the new function. + rewriter.inlineRegionBefore(funcOp.getBody(), newFuncOp.getBody(), + newFuncOp.end()); + rewriter.applySignatureConversion(&newFuncOp.getBody(), signatureConverter); + + auto builder = OpBuilder::atBlockBegin(&(newFuncOp.getBlocks().front())); + + // Cast and unpack input packed_buffer_arguments and construct memref + // descriptors. + Value packedBuffersArgsPtr = builder.create<LLVM::BitcastOp>( + loc, + LLVM::LLVMType::getStructTy(typeConverter.getDialect(), inputStructPtrs) + .getPointerTo(), + newFuncOp.getArgument(0)); + Value packedBuffersArgs = + builder.create<LLVM::LoadOp>(loc, packedBuffersArgsPtr); + for (auto bufferOp : bufferOps) { + MemRefType memrefType = bufferOp.getType().dyn_cast_or_null<MemRefType>(); + if (!memrefType) return failure(); + const auto index = bufferArgMap[bufferOp]; + Value bufferPtr = builder.create<LLVM::ExtractValueOp>( + loc, inputStructPtrs[index], packedBuffersArgs, + rewriter.getI64ArrayAttr(index)); + if (memrefType.hasStaticShape()) { + auto desc = MemRefDescriptor::fromStaticShape( + builder, loc, typeConverter, memrefType, bufferPtr); + rewriter.replaceOp(bufferOp, {desc}); + } else { + auto desc = MemRefDescriptor::undef( + builder, loc, typeConverter.convertType(memrefType)); + desc.setAllocatedPtr(builder, loc, bufferPtr); + desc.setAlignedPtr(builder, loc, bufferPtr); + rewriter.replaceOp(bufferOp, {desc}); + } + } + + // Lower hal.interface.load.constant ops into llvm.getelementptr, llvm.load + for (auto loadOp : loadOps) { + Value offset = builder.create<LLVM::ConstantOp>( + loc, LLVM::LLVMType::getInt64Ty(typeConverter.getDialect()), + builder.getI64IntegerAttr(loadOp.offset().getZExtValue())); + Value constPtr = builder.create<LLVM::GEPOp>(loc, pushConstantArgTy, + newFuncOp.getArgument(1), + ArrayRef<Value>({offset})); + Value dimConstant = builder.create<LLVM::LoadOp>(loc, constPtr); + rewriter.replaceOp(loadOp, dimConstant); + } + + rewriter.eraseOp(funcOp); + return success(); + } +}; + +class RemoveInterfaceOpPattern : public ConvertToLLVMPattern { + public: + explicit RemoveInterfaceOpPattern(MLIRContext *context, + LLVMTypeConverter &typeconverter) + : ConvertToLLVMPattern(IREE::HAL::InterfaceOp::getOperationName(), + context, typeconverter) {} + LogicalResult matchAndRewrite( + Operation *op, ArrayRef<Value> operands, + ConversionPatternRewriter &rewriter) const override { + rewriter.eraseOp(op); + return success(); + } +}; + namespace { struct ConvertToLLVMPass : public PassWrapper<ConvertToLLVMPass, OperationPass<ModuleOp>> { @@ -145,11 +319,12 @@ populateVectorToLLVMConversionPatterns(converter, patterns); populateLinalgToLLVMConversionPatterns(converter, patterns, &getContext()); // The following patterns resolves dynamic shapes by substituting tie_shape - // ops with an updated memref descriptors and replacing RankDimOp with actual - // index loaded from memref<?xi32> that holds all dynamic shapes - // push constants. - patterns.insert<ConvertRankedDimPattern, ConvertTieShapePattern, - RemoveMakeRankedShape>(&getContext(), converter); + // ops with an updated memref descriptors and replacing RankDimOp with + // actual index loaded from memref<?xi32> that holds all dynamic shapes push + // constants. + patterns.insert<ConvertFuncWithHALInterface, ConvertRankedDimPattern, + ConvertTieShapePattern, RemoveMakeRankedShape, + RemoveInterfaceOpPattern>(&getContext(), converter); LLVMConversionTarget target(getContext()); target.addLegalOp<ModuleOp, ModuleTerminatorOp>(); if (failed(applyPartialConversion(module, target, patterns))) @@ -162,7 +337,8 @@ static PassRegistration<ConvertToLLVMPass> pass( "iree-codegen-convert-to-llvm", - "Perform final conversion from Linalg/HAL/Shape/Vector/Standard to LLVMIR " + "Perform final conversion from Linalg/HAL/Shape/Vector/Standard to " + "LLVMIR " "dialect", [] { return std::make_unique<ConvertToLLVMPass>(); });
diff --git a/iree/compiler/Conversion/LinalgToLLVM/HALInterfaceToMemrefArguments.cpp b/iree/compiler/Conversion/LinalgToLLVM/HALInterfaceToMemrefArguments.cpp deleted file mode 100644 index ac968e4..0000000 --- a/iree/compiler/Conversion/LinalgToLLVM/HALInterfaceToMemrefArguments.cpp +++ /dev/null
@@ -1,231 +0,0 @@ -// Copyright 2020 Google LLC -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// https://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -#include <memory> - -#include "iree/compiler/Dialect/HAL/IR/HALDialect.h" -#include "iree/compiler/Dialect/HAL/IR/HALOps.h" -#include "iree/compiler/Dialect/IREE/IR/IREEOps.h" -#include "mlir/Dialect/StandardOps/IR/Ops.h" -#include "mlir/IR/Builders.h" -#include "mlir/Pass/Pass.h" -#include "mlir/Pass/PassRegistry.h" -#include "mlir/Transforms/DialectConversion.h" - -namespace mlir { -namespace iree_compiler { -namespace { - -/// Returns true if the given function contains interface related operations -/// that are used by other ops. -bool containsUsedInterfaceOp(FuncOp funcOp) { - for (Block& block : funcOp.getBlocks()) { - for (Operation& op : block) { - if (!op.getUses().empty() && - (isa<IREE::PlaceholderOp>(op) || - isa<IREE::HAL::InterfaceLoadConstantOp>(op))) { - return true; - } - } - } - return false; -} - -/// Returns true if `aOp` has a desciptor (set, binding) pair smaller than -/// `bOp`. Note that this ignores the offset. -bool operator<(IREE::HAL::InterfaceBindingOp aOp, - IREE::HAL::InterfaceBindingOp bOp) { - if (aOp.set().getZExtValue() == bOp.set().getZExtValue()) - return aOp.binding().getZExtValue() < bOp.binding().getZExtValue(); - return aOp.set().getZExtValue() < bOp.set().getZExtValue(); -} - -/// A pattern to process function interface. It replaces interface related ops -/// with function arguments to match LLVM's CodeGen's ABI contract. -/// -/// IREE scheduler passes interface ABI information via hal.interface.* ops to -/// all backends. We create iree.placeholder ops to represent buffers behind -/// those hal.interface.* ops. However the LLVM CodeGen uses function parameters -/// and memref descriptors for ABI. So we need to bridge the gap somewhere. -/// -/// This pass finds all interface buffers used in the function, sort them -/// according to the descriptor (set, binding) pair, and put unique ones as -/// function parameters in order. -/// Note: This should be kept consistent with LLVM's HAL backend. -struct ProcessFuncInterfacePattern : public OpConversionPattern<FuncOp> { - using OpConversionPattern::OpConversionPattern; - LogicalResult matchAndRewrite( - FuncOp funcOp, ArrayRef<Value> Operands, - ConversionPatternRewriter& rewriter) const override { - // Only process entry functions. - if (SymbolTable::getSymbolVisibility(funcOp) != - SymbolTable::Visibility::Public) - return failure(); - - FunctionType fnType = funcOp.getType(); - if (fnType.getNumInputs() != 0) - return rewriter.notifyMatchFailure( - funcOp, "entry function should not have inputs"); - - // Get interface buffers from all the blocks. - SmallVector<IREE::PlaceholderOp, 8> bufferOps; - SmallVector<IREE::HAL::InterfaceLoadConstantOp, 8> loadOps; - for (Block& block : funcOp.getBlocks()) { - for (Operation& op : block) { - if (auto phOp = dyn_cast<IREE::PlaceholderOp>(op)) - bufferOps.push_back(phOp); - if (auto phOp = dyn_cast<IREE::HAL::InterfaceLoadConstantOp>(op)) { - loadOps.push_back(phOp); - } - } - } - - if (bufferOps.empty()) return failure(); - - // A map from buffer ops to their corresponding interface binding ops. - llvm::DenseMap<Operation*, IREE::HAL::InterfaceBindingOp> bufferBindingMap; - for (auto bufferOp : bufferOps) { - auto symbol = SymbolTable::lookupNearestSymbolFrom( - bufferOp, bufferOp.getAttrOfType<SymbolRefAttr>("binding")); - bufferBindingMap[bufferOp] = cast<IREE::HAL::InterfaceBindingOp>(symbol); - } - - // Sort buffers according to their descriptor (set, binding) pair. - llvm::sort(bufferOps, [&bufferBindingMap](IREE::PlaceholderOp aBuffer, - IREE::PlaceholderOp bBuffer) { - return bufferBindingMap[aBuffer] < bufferBindingMap[bBuffer]; - }); - - // Create a function argument for each of the unique binding pointed by the - // buffer ops. - TypeConverter::SignatureConversion signatureConverter(/*numOrigInputs=*/0); - // A map from buffer ops to their corresponding function argument indices. - llvm::DenseMap<Operation*, unsigned> bufferArgMap; - // A map from binding ops to their corresponding function argument indices. - llvm::DenseMap<Operation*, unsigned> bindingArgMap; - unsigned argIndex = 0; - for (auto bufferOp : bufferOps) { - auto binding = bufferBindingMap[bufferOp]; - auto it = bindingArgMap.find(binding); - if (it != bindingArgMap.end()) { - bufferArgMap[bufferOp] = it->second; - } else { - bindingArgMap[binding] = argIndex; - bufferArgMap[bufferOp] = argIndex; - signatureConverter.addInputs(bufferOp.getType()); - ++argIndex; - } - } - Type dynamicDimsBufferType = - MemRefType::get(ShapedType::kDynamicSize, rewriter.getIntegerType(32)); - signatureConverter.addInputs(dynamicDimsBufferType); - - // Create the new function's signature. - Location loc = funcOp.getLoc(); - auto newFuncOp = rewriter.create<FuncOp>( - loc, funcOp.getName(), - rewriter.getFunctionType(signatureConverter.getConvertedTypes(), - llvm::None), - ArrayRef<NamedAttribute>()); - newFuncOp.setAttr("llvm.emit_c_interface", - mlir::UnitAttr::get(funcOp.getContext())); - - // Move all ops in the old function's region to the new function. - rewriter.inlineRegionBefore(funcOp.getBody(), newFuncOp.getBody(), - newFuncOp.end()); - rewriter.applySignatureConversion(&newFuncOp.getBody(), signatureConverter); - - // Replace all buffer ops' uses with the newly created function arguments - // and erase them. - for (auto bufferOp : bufferOps) { - bufferOp.replaceAllUsesWith( - newFuncOp.getArgument(bufferArgMap[bufferOp])); - - rewriter.eraseOp(bufferOp); - } - - // Lower all hal.interface.load.constant ops into std.load - // from the last buffer holding all dynamic dimensions with the proper - // offset. - Type indexType = rewriter.getIndexType(); - auto builder = OpBuilder::atBlockBegin(&(newFuncOp.getBlocks().front())); - auto newLoc = newFuncOp.front().front().getLoc(); - for (auto loadOp : loadOps) { - SmallVector<Value, 1> indices; - Value constantOffset = builder.create<ConstantOp>( - newLoc, indexType, - rewriter.getIntegerAttr(indexType, loadOp.offset().getZExtValue())); - indices.push_back(constantOffset); - Value loadDim = builder.create<LoadOp>( - newLoc, newFuncOp.getArgument(newFuncOp.getNumArguments() - 1), - indices); - Value loadDimIndex = - builder.create<IndexCastOp>(newLoc, loadDim, indexType); - loadOp.replaceAllUsesWith(loadDimIndex); - rewriter.eraseOp(loadOp); - } - rewriter.eraseOp(funcOp); - return success(); - } -}; - -struct RemoveInterfaceOpPattern - : public OpRewritePattern<IREE::HAL::InterfaceOp> { - using OpRewritePattern::OpRewritePattern; - LogicalResult matchAndRewrite(IREE::HAL::InterfaceOp interfaceOp, - PatternRewriter& rewriter) const override { - rewriter.eraseOp(interfaceOp); - return success(); - } -}; - -/// Converting from Linalg to LLVM needs to run on a module and since it -/// applies a full conversion, make a module with jst the impl function. -struct HALInterfaceToMemrefArgumentsPass - : PassWrapper<HALInterfaceToMemrefArgumentsPass, OperationPass<ModuleOp>> { - void runOnOperation() override { - MLIRContext& context = getContext(); - - OwningRewritePatternList patterns; - patterns.insert<ProcessFuncInterfacePattern>(&context); - patterns.insert<RemoveInterfaceOpPattern>(&context); - - ConversionTarget target(context); - // Convert the interface related ops away. - target.addDynamicallyLegalOp<FuncOp>( - [](FuncOp funcOp) { return !containsUsedInterfaceOp(funcOp); }); - target.addIllegalOp<IREE::PlaceholderOp>(); - target.addIllegalDialect<IREE::HAL::HALDialect>(); - // Allow the rest. - target.markUnknownOpDynamicallyLegal([](Operation*) { return true; }); - - if (failed(applyFullConversion(getOperation(), target, patterns))) - return signalPassFailure(); - } -}; - -} // namespace - -std::unique_ptr<OperationPass<ModuleOp>> -createHALInterfaceToMemrefArgumentsPass() { - return std::make_unique<HALInterfaceToMemrefArgumentsPass>(); -} - -static PassRegistration<HALInterfaceToMemrefArgumentsPass> pass( - "iree-codegen-hal-interface-to-memref-arguments-pass", - "Convert a function with HAL bindings interface to memref arguments", - [] { return std::make_unique<HALInterfaceToMemrefArgumentsPass>(); }); - -} // namespace iree_compiler -} // namespace mlir
diff --git a/iree/compiler/Conversion/LinalgToLLVM/Passes.cpp b/iree/compiler/Conversion/LinalgToLLVM/Passes.cpp index e8c6d9c..8c8eb21 100644 --- a/iree/compiler/Conversion/LinalgToLLVM/Passes.cpp +++ b/iree/compiler/Conversion/LinalgToLLVM/Passes.cpp
@@ -35,10 +35,7 @@ passManager.addPass(createCanonicalizerPass()); passManager.addPass(createCSEPass()); - // Convert ExecuableOp entry function to use memref arguments. - passManager.addPass(createHALInterfaceToMemrefArgumentsPass()); - - // (Linalg, STD) -> LLVM + // (HAL, IREE, Linalg, STD) -> LLVM // OpPassManager& llvmPassManager = passManager.nest<ModuleOp>(); passManager.addPass(createConvertToLLVMPass()); passManager.addPass(createCanonicalizerPass());
diff --git a/iree/compiler/Conversion/LinalgToLLVM/Passes.h b/iree/compiler/Conversion/LinalgToLLVM/Passes.h index 5bfb893..fdad0e6 100644 --- a/iree/compiler/Conversion/LinalgToLLVM/Passes.h +++ b/iree/compiler/Conversion/LinalgToLLVM/Passes.h
@@ -20,11 +20,6 @@ namespace mlir { namespace iree_compiler { -/// Converts function signture type from hal interface op annotation to memref -/// argument. -std::unique_ptr<OperationPass<ModuleOp>> -createHALInterfaceToMemrefArgumentsPass(); - /// Pass to perform final conversion to LLVM dialect. std::unique_ptr<OperationPass<ModuleOp>> createConvertToLLVMPass();
diff --git a/iree/compiler/Conversion/init_conversions.h b/iree/compiler/Conversion/init_conversions.h index 7a190e7..259e3d5 100644 --- a/iree/compiler/Conversion/init_conversions.h +++ b/iree/compiler/Conversion/init_conversions.h
@@ -47,7 +47,6 @@ inline void registerLinalgToLLVMPasses() { static bool init_once = []() { // LinalgToLLVM - createHALInterfaceToMemrefArgumentsPass(); return true; }(); (void)init_once;
diff --git a/iree/compiler/Dialect/HAL/Target/LLVM/LLVMAOTTarget.cpp b/iree/compiler/Dialect/HAL/Target/LLVM/LLVMAOTTarget.cpp index 7269089..8cb47b5 100644 --- a/iree/compiler/Dialect/HAL/Target/LLVM/LLVMAOTTarget.cpp +++ b/iree/compiler/Dialect/HAL/Target/LLVM/LLVMAOTTarget.cpp
@@ -65,13 +65,10 @@ auto executableOp = cast<ExecutableOp>(targetOp.getParentOp()); auto entryPointOps = executableOp.getBlock().getOps<ExecutableEntryPointOp>(); - const bool addCInterface = true; + for (auto entryPointOp : entryPointOps) { - std::string funcName = - addCInterface ? "_mlir_ciface_" + std::string(entryPointOp.sym_name()) - : std::string(entryPointOp.sym_name()); - dyLibExecutableDef.entry_points.push_back("invoke_" + funcName); - createLLVMInvocationFunc(funcName, llvmModule.get()); + dyLibExecutableDef.entry_points.push_back( + std::string(entryPointOp.sym_name())); } // LLVMIR opt passes.
diff --git a/iree/compiler/Dialect/HAL/Target/LLVM/LLVMIRPasses.cpp b/iree/compiler/Dialect/HAL/Target/LLVM/LLVMIRPasses.cpp index 24d5877..e91441d 100644 --- a/iree/compiler/Dialect/HAL/Target/LLVM/LLVMIRPasses.cpp +++ b/iree/compiler/Dialect/HAL/Target/LLVM/LLVMIRPasses.cpp
@@ -44,44 +44,6 @@ return machine; } -void createLLVMInvocationFunc(const std::string& name, llvm::Module* module) { - // TODO(ataei): This is written as a stub in LLVM IR. It would be easier to - // have this using MLIR and lower it to LLVM like the dispatch function - // implementation is. - - auto& ctx = module->getContext(); - llvm::IRBuilder<> builder(ctx); - auto var_func = module->getFunction(name); - - auto new_type = llvm::FunctionType::get( - builder.getVoidTy(), builder.getInt8PtrTy()->getPointerTo(), - /*isVarArg=*/false); - - auto new_name = "invoke_" + name; - auto func_cst = module->getOrInsertFunction(new_name, new_type); - llvm::Function* interface_func = - llvm::cast<llvm::Function>(func_cst.getCallee()); - - auto bb = llvm::BasicBlock::Create(ctx); - bb->insertInto(interface_func); - builder.SetInsertPoint(bb); - llvm::Value* argList = interface_func->arg_begin(); - llvm::SmallVector<llvm::Value*, 8> args; - args.reserve(llvm::size(var_func->args())); - for (auto& indexedArg : llvm::enumerate(var_func->args())) { - llvm::Value* arg_index = llvm::Constant::getIntegerValue( - builder.getInt64Ty(), llvm::APInt(64, indexedArg.index())); - llvm::Value* arg_ptr_ptr = builder.CreateGEP(argList, arg_index); - llvm::Value* arg_ptr = builder.CreateLoad(arg_ptr_ptr); - arg_ptr = builder.CreateBitCast( - arg_ptr, indexedArg.value().getType()->getPointerTo()); - llvm::Value* arg = builder.CreateLoad(arg_ptr); - args.push_back(arg); - } - builder.CreateCall(var_func, args); - builder.CreateRetVoid(); -} - LogicalResult runLLVMIRPasses(const LLVMTargetOptions& options, llvm::TargetMachine* machine, llvm::Module* module) {
diff --git a/iree/compiler/Dialect/HAL/Target/LLVM/LLVMIRPasses.h b/iree/compiler/Dialect/HAL/Target/LLVM/LLVMIRPasses.h index 199e36f..37ee1ba 100644 --- a/iree/compiler/Dialect/HAL/Target/LLVM/LLVMIRPasses.h +++ b/iree/compiler/Dialect/HAL/Target/LLVM/LLVMIRPasses.h
@@ -31,9 +31,6 @@ std::unique_ptr<llvm::TargetMachine> createTargetMachine( const LLVMTargetOptions& options); -// Creates an invocation function in a module for the given function name. -void createLLVMInvocationFunc(const std::string& name, llvm::Module* module); - // Creates and runs LLVMIR optimization passes defined in LLVMTargetOptions. LogicalResult runLLVMIRPasses(const LLVMTargetOptions& options, llvm::TargetMachine* machine,
diff --git a/iree/compiler/Dialect/HAL/Target/LLVM/LLVMIRTarget.cpp b/iree/compiler/Dialect/HAL/Target/LLVM/LLVMIRTarget.cpp index 98c0bf4..96bb5ac 100644 --- a/iree/compiler/Dialect/HAL/Target/LLVM/LLVMIRTarget.cpp +++ b/iree/compiler/Dialect/HAL/Target/LLVM/LLVMIRTarget.cpp
@@ -58,13 +58,9 @@ auto executableOp = cast<IREE::HAL::ExecutableOp>(targetOp.getParentOp()); auto entryPointOps = executableOp.getBlock().getOps<IREE::HAL::ExecutableEntryPointOp>(); - const bool addCInterface = true; for (auto entryPointOp : entryPointOps) { - std::string funcName = - addCInterface ? "_mlir_ciface_" + std::string(entryPointOp.sym_name()) - : std::string(entryPointOp.sym_name()); - llvmIrExecutableDef.entry_points.push_back(funcName); - createLLVMInvocationFunc(funcName, llvmModule.get()); + llvmIrExecutableDef.entry_points.push_back( + std::string(entryPointOp.sym_name())); } // LLVMIR opt passes. @@ -74,8 +70,9 @@ options_.targetTriple); return failure(); } - if (failed( - runLLVMIRPasses(options_, targetMachine.get(), llvmModule.get()))) { + LogicalResult translationResult = + runLLVMIRPasses(options_, targetMachine.get(), llvmModule.get()); + if (failed(translationResult)) { return targetOp.emitError( "Can't build LLVMIR opt passes for ExecutableOp module"); }
diff --git a/iree/hal/dylib/BUILD b/iree/hal/dylib/BUILD index 25c08ea..fc3ccb9 100644 --- a/iree/hal/dylib/BUILD +++ b/iree/hal/dylib/BUILD
@@ -60,7 +60,6 @@ srcs = ["dylib_executable.cc"], hdrs = ["dylib_executable.h"], deps = [ - ":memref_runtime", "//iree/base:dynamic_library", "//iree/base:file_io", "//iree/base:status", @@ -89,10 +88,3 @@ "//iree/hal:executable_format", ], ) - -cc_library( - name = "memref_runtime", - hdrs = [ - "memref_runtime.h", - ], -)
diff --git a/iree/hal/dylib/CMakeLists.txt b/iree/hal/dylib/CMakeLists.txt index 7644d92..d720435 100644 --- a/iree/hal/dylib/CMakeLists.txt +++ b/iree/hal/dylib/CMakeLists.txt
@@ -65,7 +65,6 @@ SRCS "dylib_executable.cc" DEPS - ::memref_runtime absl::inlined_vector absl::span flatbuffers @@ -97,11 +96,3 @@ iree::hal::executable_format PUBLIC ) - -iree_cc_library( - NAME - memref_runtime - HDRS - "memref_runtime.h" - PUBLIC -)
diff --git a/iree/hal/dylib/dylib_executable.cc b/iree/hal/dylib/dylib_executable.cc index e06bb19..e58a003 100644 --- a/iree/hal/dylib/dylib_executable.cc +++ b/iree/hal/dylib/dylib_executable.cc
@@ -17,7 +17,6 @@ #include "flatbuffers/flatbuffers.h" #include "iree/base/file_io.h" #include "iree/base/tracing.h" -#include "iree/hal/dylib/memref_runtime.h" #include "iree/schemas/dylib_executable_def_generated.h" namespace iree { @@ -96,15 +95,9 @@ struct DyLibDispatchState : public HostExecutable::DispatchState { DyLibDispatchState() = default; - ~DyLibDispatchState() override { - for (int i = 0; i < descriptors.size(); ++i) { - freeUnrankedDescriptor(descriptors[i]); - } - } - void* entry_function = nullptr; - absl::InlinedVector<UnrankedMemRefType<uint32_t>*, 4> descriptors; absl::InlinedVector<void*, 4> args; + absl::InlinedVector<int64_t, 4> push_constant; }; StatusOr<ref_ptr<HostExecutable::DispatchState>> @@ -127,17 +120,14 @@ MemoryAccessBitfield::kWrite, io_binding.offset, io_binding.length)); auto data = memory.mutable_data(); - auto descriptor = allocUnrankedDescriptor<uint32_t>(data); - dispatch_state->descriptors.push_back(descriptor); - dispatch_state->args.push_back(&descriptor->descriptor); + + dispatch_state->args.push_back(data); } } - - auto push_constants_descriptor = allocUnrankedDescriptor<uint32_t>( - const_cast<uint32_t*>(params.push_constants->values.data()), - {static_cast<int64_t>(params.push_constants->values.size())}); - dispatch_state->descriptors.push_back(push_constants_descriptor); - dispatch_state->args.push_back(&push_constants_descriptor->descriptor); + // TODO(ataei): Consider moving this casting to codegen side ?! + for (int i = 0; i < params.push_constants->values.size(); ++i) { + dispatch_state->push_constant.push_back(params.push_constants->values[i]); + } return std::move(dispatch_state); } @@ -147,8 +137,10 @@ IREE_TRACE_SCOPE0("DyLibExecutable::DispatchTile"); auto* dispatch_state = static_cast<DyLibDispatchState*>(state); - auto entry_function = (void (*)(void**))dispatch_state->entry_function; - entry_function(dispatch_state->args.data()); + auto entry_function = + (void (*)(void**, int64_t*))dispatch_state->entry_function; + entry_function(dispatch_state->args.data(), + dispatch_state->push_constant.data()); return OkStatus(); }
diff --git a/iree/hal/dylib/memref_runtime.h b/iree/hal/dylib/memref_runtime.h deleted file mode 100644 index 50d3987..0000000 --- a/iree/hal/dylib/memref_runtime.h +++ /dev/null
@@ -1,177 +0,0 @@ -// Copyright 2020 Google LLC -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// https://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. -// - -#ifndef IREE_HAL_DYLIB_MEMREF_RUNTIME_H_ -#define IREE_HAL_DYLIB_MEMREF_RUNTIME_H_ - -#include <assert.h> - -#include <cstdint> -#include <vector> - -namespace iree { -namespace hal { -namespace dylib { - -template <int N> -void dropFront(int64_t arr[N], int64_t *res) { - for (unsigned i = 1; i < N; ++i) *(res + i - 1) = arr[i]; -} - -/// StridedMemRef descriptor type with static rank. -template <typename T, int N> -struct StridedMemRefType { - T *basePtr; - T *data; - int64_t offset; - int64_t sizes[N]; - int64_t strides[N]; - // This operator[] is extremely slow and only for sugaring purposes. - StridedMemRefType<T, N - 1> operator[](int64_t idx) { - StridedMemRefType<T, N - 1> res; - res.basePtr = basePtr; - res.data = data; - res.offset = offset + idx * strides[0]; - dropFront<N>(sizes, res.sizes); - dropFront<N>(strides, res.strides); - return res; - } -}; - -/// StridedMemRef descriptor type specialized for rank 1. -template <typename T> -struct StridedMemRefType<T, 1> { - T *basePtr; - T *data; - int64_t offset; - int64_t sizes[1]; - int64_t strides[1]; - T &operator[](int64_t idx) { return *(data + offset + idx * strides[0]); } -}; - -/// StridedMemRef descriptor type specialized for rank 0. -template <typename T> -struct StridedMemRefType<T, 0> { - T *basePtr; - T *data; - int64_t offset; -}; - -// Unranked MemRef -template <typename T> -struct UnrankedMemRefType { - int64_t rank; - void *descriptor; -}; - -// Given a shape with sizes greater than 0 along all dimensions, -// returns the distance, in number of elements, between a slice in a dimension -// and the next slice in the same dimension. -// e.g. shape[3, 4, 5] -> strides[20, 5, 1] -inline std::vector<int64_t> makeStrides(const std::vector<int64_t> &shape) { - std::vector<int64_t> tmp; - if (shape.empty()) return tmp; - tmp.reserve(shape.size()); - int64_t running = 1; - for (auto rit = shape.rbegin(), reit = shape.rend(); rit != reit; ++rit) { - assert(*rit > 0 && - "size must be greater than 0 along all dimensions of shape"); - tmp.push_back(running); - running *= *rit; - } - return std::vector<int64_t>(tmp.rbegin(), tmp.rend()); -} - -// Mallocs a StridedMemRefDescriptor<T, N>* that matches the MLIR ABI. -// This is an implementation detail that is kept in sync with MLIR codegen -// conventions. -template <typename T, int N> -StridedMemRefType<T, N> *makeStridedMemRefDescriptor( - void *ptr, const std::vector<int64_t> &shape) { - StridedMemRefType<T, N> *descriptor = static_cast<StridedMemRefType<T, N> *>( - malloc(sizeof(StridedMemRefType<T, N>))); - descriptor->basePtr = static_cast<T *>(ptr); - descriptor->data = static_cast<T *>(ptr); - descriptor->offset = 0; - std::copy(shape.begin(), shape.end(), descriptor->sizes); - auto strides = makeStrides(shape); - std::copy(strides.begin(), strides.end(), descriptor->strides); - return descriptor; -} - -// Mallocs a StridedMemRefDescriptor<T, 0>* (i.e. a pointer to scalar) that -// matches the MLIR ABI. This is an implementation detail that is kept in sync -// with MLIR codegen conventions. -template <typename T> -StridedMemRefType<T, 0> *makeStridedMemRefDescriptor( - void *ptr, const std::vector<int64_t> &shape) { - StridedMemRefType<T, 0> *descriptor = static_cast<StridedMemRefType<T, 0> *>( - malloc(sizeof(StridedMemRefType<T, 0>))); - descriptor->basePtr = static_cast<T *>(ptr); - descriptor->data = static_cast<T *>(ptr); - descriptor->offset = 0; - return descriptor; -} - -// Mallocs an UnrankedMemRefType<T>* that contains a ranked -// StridedMemRefDescriptor<T, Rank>* and matches the MLIR ABI. This is an -// implementation detail that is kept in sync with MLIR codegen conventions. -template <typename T> -UnrankedMemRefType<T> *allocUnrankedDescriptor( - void *data, const std::vector<int64_t> &shape) { - UnrankedMemRefType<T> *res = static_cast<UnrankedMemRefType<T> *>( - malloc(sizeof(UnrankedMemRefType<T>))); - res->rank = shape.size(); - if (res->rank == 0) - res->descriptor = makeStridedMemRefDescriptor<T>(data, shape); - else if (res->rank == 1) - res->descriptor = makeStridedMemRefDescriptor<T, 1>(data, shape); - else if (res->rank == 2) - res->descriptor = makeStridedMemRefDescriptor<T, 2>(data, shape); - else if (res->rank == 3) - res->descriptor = makeStridedMemRefDescriptor<T, 3>(data, shape); - else if (res->rank == 4) - res->descriptor = makeStridedMemRefDescriptor<T, 4>(data, shape); - else if (res->rank == 5) - res->descriptor = makeStridedMemRefDescriptor<T, 5>(data, shape); - else if (res->rank == 6) - res->descriptor = makeStridedMemRefDescriptor<T, 6>(data, shape); - else - assert(false && "Unsupported 6+D memref descriptor"); - return res; -} - -// Shape and strides aren't used in the generated code (yet). -// TODO(ataei): Delete this version once we can pass shapes. -template <typename T> -UnrankedMemRefType<T> *allocUnrankedDescriptor(void *data) { - UnrankedMemRefType<T> *res = static_cast<UnrankedMemRefType<T> *>( - malloc(sizeof(UnrankedMemRefType<T>))); - res->descriptor = makeStridedMemRefDescriptor<T>(data, {}); - return res; -} - -// Frees an UnrankedMemRefType<T>* -template <typename T> -void freeUnrankedDescriptor(UnrankedMemRefType<T> *desc) { - free(desc->descriptor); - free(desc); -} - -} // namespace dylib -} // namespace hal -} // namespace iree - -#endif // IREE_HAL_DYLIB_MEMREF_RUNTIME_H_
diff --git a/iree/hal/llvmjit/BUILD b/iree/hal/llvmjit/BUILD index 088bb8b..3ebd609 100644 --- a/iree/hal/llvmjit/BUILD +++ b/iree/hal/llvmjit/BUILD
@@ -64,7 +64,6 @@ srcs = ["llvmjit_executable.cc"], hdrs = ["llvmjit_executable.h"], deps = [ - ":memref_runtime", "//iree/base:status", "//iree/base:tracing", "//iree/hal:buffer", @@ -95,10 +94,3 @@ "//iree/hal:executable_format", ], ) - -cc_library( - name = "memref_runtime", - hdrs = [ - "memref_runtime.h", - ], -)
diff --git a/iree/hal/llvmjit/CMakeLists.txt b/iree/hal/llvmjit/CMakeLists.txt index 8418745..ca40941 100644 --- a/iree/hal/llvmjit/CMakeLists.txt +++ b/iree/hal/llvmjit/CMakeLists.txt
@@ -68,7 +68,6 @@ SRCS "llvmjit_executable.cc" DEPS - ::memref_runtime LLVMAsmParser LLVMCore LLVMOrcJIT @@ -102,11 +101,3 @@ iree::hal::executable_format PUBLIC ) - -iree_cc_library( - NAME - memref_runtime - HDRS - "memref_runtime.h" - PUBLIC -)
diff --git a/iree/hal/llvmjit/llvmjit_executable.cc b/iree/hal/llvmjit/llvmjit_executable.cc index 1596b9e..7d26ccd 100644 --- a/iree/hal/llvmjit/llvmjit_executable.cc +++ b/iree/hal/llvmjit/llvmjit_executable.cc
@@ -21,7 +21,6 @@ #include "iree/base/tracing.h" #include "iree/hal/buffer.h" #include "iree/hal/executable.h" -#include "iree/hal/llvmjit/memref_runtime.h" #include "iree/schemas/llvmir_executable_def_generated.h" #include "llvm/ADT/ArrayRef.h" #include "llvm/ADT/StringRef.h" @@ -82,13 +81,11 @@ make_ref<LLVMJITExecutable>(spec, std::move(ll_jit), allow_aliasing_data); for (const auto func_name : *entry_points) { - auto func_symbol = - executable->ll_jit_->lookup("invoke_" + func_name->str()); + auto func_symbol = executable->ll_jit_->lookup(func_name->str()); if (!func_symbol) { return NotFoundErrorBuilder(IREE_LOC) << "Can't JIT compile function : " << func_name; } - // Map function to its invoke_ symbol. executable->symbols_.push_back(func_symbol.get()); } @@ -111,15 +108,10 @@ struct LLVMJITDispatchState : public HostExecutable::DispatchState { LLVMJITDispatchState() = default; - ~LLVMJITDispatchState() override { - for (int i = 0; i < descriptors.size(); ++i) { - freeUnrankedDescriptor(descriptors[i]); - } - } llvm::JITEvaluatedSymbol symbol; - llvm::SmallVector<UnrankedMemRefType<uint32_t>*, 4> descriptors; llvm::SmallVector<void*, 4> args; + llvm::SmallVector<int64_t, 4> push_constant; }; StatusOr<ref_ptr<HostExecutable::DispatchState>> @@ -142,17 +134,13 @@ MemoryAccessBitfield::kWrite, io_binding.offset, io_binding.length)); auto data = memory.mutable_data(); - auto descriptor = allocUnrankedDescriptor<uint32_t>(data); - dispatch_state->descriptors.push_back(descriptor); - dispatch_state->args.push_back(&descriptor->descriptor); + dispatch_state->args.push_back(data); } } - - auto push_constants_descriptor = allocUnrankedDescriptor<uint32_t>( - const_cast<uint32_t*>(params.push_constants->values.data()), - {static_cast<int64_t>(params.push_constants->values.size())}); - dispatch_state->descriptors.push_back(push_constants_descriptor); - dispatch_state->args.push_back(&push_constants_descriptor->descriptor); + // TODO(ataei): Consider moving this casting to codegen side ?! + for (int i = 0; i < params.push_constants->values.size(); ++i) { + dispatch_state->push_constant.push_back(params.push_constants->values[i]); + } return std::move(dispatch_state); } @@ -162,8 +150,9 @@ IREE_TRACE_SCOPE0("LLVMJITExecutable::DispatchTile"); auto* dispatch_state = static_cast<LLVMJITDispatchState*>(state); - auto func_ptr = (void (*)(void**))dispatch_state->symbol.getAddress(); - func_ptr(dispatch_state->args.data()); + auto func_ptr = + (void (*)(void**, int64_t*))dispatch_state->symbol.getAddress(); + func_ptr(dispatch_state->args.data(), dispatch_state->push_constant.data()); return OkStatus(); }
diff --git a/iree/hal/llvmjit/memref_runtime.h b/iree/hal/llvmjit/memref_runtime.h deleted file mode 100644 index 6b94410..0000000 --- a/iree/hal/llvmjit/memref_runtime.h +++ /dev/null
@@ -1,177 +0,0 @@ -// Copyright 2020 Google LLC -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// https://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. -// - -#ifndef IREE_HAL_LLVMJIT_LLVMJIT_MEMREF_RUNTIME_H_ -#define IREE_HAL_LLVMJIT_LLVMJIT_MEMREF_RUNTIME_H_ - -#include <assert.h> - -#include <cstdint> -#include <vector> - -namespace iree { -namespace hal { -namespace llvmjit { - -template <int N> -void dropFront(int64_t arr[N], int64_t *res) { - for (unsigned i = 1; i < N; ++i) *(res + i - 1) = arr[i]; -} - -/// StridedMemRef descriptor type with static rank. -template <typename T, int N> -struct StridedMemRefType { - T *basePtr; - T *data; - int64_t offset; - int64_t sizes[N]; - int64_t strides[N]; - // This operator[] is extremely slow and only for sugaring purposes. - StridedMemRefType<T, N - 1> operator[](int64_t idx) { - StridedMemRefType<T, N - 1> res; - res.basePtr = basePtr; - res.data = data; - res.offset = offset + idx * strides[0]; - dropFront<N>(sizes, res.sizes); - dropFront<N>(strides, res.strides); - return res; - } -}; - -/// StridedMemRef descriptor type specialized for rank 1. -template <typename T> -struct StridedMemRefType<T, 1> { - T *basePtr; - T *data; - int64_t offset; - int64_t sizes[1]; - int64_t strides[1]; - T &operator[](int64_t idx) { return *(data + offset + idx * strides[0]); } -}; - -/// StridedMemRef descriptor type specialized for rank 0. -template <typename T> -struct StridedMemRefType<T, 0> { - T *basePtr; - T *data; - int64_t offset; -}; - -// Unranked MemRef -template <typename T> -struct UnrankedMemRefType { - int64_t rank; - void *descriptor; -}; - -// Given a shape with sizes greater than 0 along all dimensions, -// returns the distance, in number of elements, between a slice in a dimension -// and the next slice in the same dimension. -// e.g. shape[3, 4, 5] -> strides[20, 5, 1] -inline std::vector<int64_t> makeStrides(const std::vector<int64_t> &shape) { - std::vector<int64_t> tmp; - if (shape.empty()) return tmp; - tmp.reserve(shape.size()); - int64_t running = 1; - for (auto rit = shape.rbegin(), reit = shape.rend(); rit != reit; ++rit) { - assert(*rit > 0 && - "size must be greater than 0 along all dimensions of shape"); - tmp.push_back(running); - running *= *rit; - } - return std::vector<int64_t>(tmp.rbegin(), tmp.rend()); -} - -// Mallocs a StridedMemRefDescriptor<T, N>* that matches the MLIR ABI. -// This is an implementation detail that is kept in sync with MLIR codegen -// conventions. -template <typename T, int N> -StridedMemRefType<T, N> *makeStridedMemRefDescriptor( - void *ptr, const std::vector<int64_t> &shape) { - StridedMemRefType<T, N> *descriptor = static_cast<StridedMemRefType<T, N> *>( - malloc(sizeof(StridedMemRefType<T, N>))); - descriptor->basePtr = static_cast<T *>(ptr); - descriptor->data = static_cast<T *>(ptr); - descriptor->offset = 0; - std::copy(shape.begin(), shape.end(), descriptor->sizes); - auto strides = makeStrides(shape); - std::copy(strides.begin(), strides.end(), descriptor->strides); - return descriptor; -} - -// Mallocs a StridedMemRefDescriptor<T, 0>* (i.e. a pointer to scalar) that -// matches the MLIR ABI. This is an implementation detail that is kept in sync -// with MLIR codegen conventions. -template <typename T> -StridedMemRefType<T, 0> *makeStridedMemRefDescriptor( - void *ptr, const std::vector<int64_t> &shape) { - StridedMemRefType<T, 0> *descriptor = static_cast<StridedMemRefType<T, 0> *>( - malloc(sizeof(StridedMemRefType<T, 0>))); - descriptor->basePtr = static_cast<T *>(ptr); - descriptor->data = static_cast<T *>(ptr); - descriptor->offset = 0; - return descriptor; -} - -// Mallocs an UnrankedMemRefType<T>* that contains a ranked -// StridedMemRefDescriptor<T, Rank>* and matches the MLIR ABI. This is an -// implementation detail that is kept in sync with MLIR codegen conventions. -template <typename T> -UnrankedMemRefType<T> *allocUnrankedDescriptor( - void *data, const std::vector<int64_t> &shape) { - UnrankedMemRefType<T> *res = static_cast<UnrankedMemRefType<T> *>( - malloc(sizeof(UnrankedMemRefType<T>))); - res->rank = shape.size(); - if (res->rank == 0) - res->descriptor = makeStridedMemRefDescriptor<T>(data, shape); - else if (res->rank == 1) - res->descriptor = makeStridedMemRefDescriptor<T, 1>(data, shape); - else if (res->rank == 2) - res->descriptor = makeStridedMemRefDescriptor<T, 2>(data, shape); - else if (res->rank == 3) - res->descriptor = makeStridedMemRefDescriptor<T, 3>(data, shape); - else if (res->rank == 4) - res->descriptor = makeStridedMemRefDescriptor<T, 4>(data, shape); - else if (res->rank == 5) - res->descriptor = makeStridedMemRefDescriptor<T, 5>(data, shape); - else if (res->rank == 6) - res->descriptor = makeStridedMemRefDescriptor<T, 6>(data, shape); - else - assert(false && "Unsupported 6+D memref descriptor"); - return res; -} - -// Shape and strides aren't used in the generated code (yet). -// TODO(ataei): Delete this version once we can pass shapes. -template <typename T> -UnrankedMemRefType<T> *allocUnrankedDescriptor(void *data) { - UnrankedMemRefType<T> *res = static_cast<UnrankedMemRefType<T> *>( - malloc(sizeof(UnrankedMemRefType<T>))); - res->descriptor = makeStridedMemRefDescriptor<T>(data, {}); - return res; -} - -// Frees an UnrankedMemRefType<T>* -template <typename T> -void freeUnrankedDescriptor(UnrankedMemRefType<T> *desc) { - free(desc->descriptor); - free(desc); -} - -} // namespace llvmjit -} // namespace hal -} // namespace iree - -#endif // IREE_HAL_LLVMJIT_LLVMJIT_MEMREF_RUNTIME_H_