diff --git a/src/numba/openmp/libs/pass/CGIntrinsicsOpenMP.cpp b/src/numba/openmp/libs/pass/CGIntrinsicsOpenMP.cpp index 62934eead366..19aabf828b9d 100644 --- a/src/numba/openmp/libs/pass/CGIntrinsicsOpenMP.cpp +++ b/src/numba/openmp/libs/pass/CGIntrinsicsOpenMP.cpp @@ -2,6 +2,8 @@ #include "DebugOpenMP.h" #include +#include +#include #include #include #include @@ -13,6 +15,7 @@ #include #include #include +#include #include #include #include @@ -532,10 +535,18 @@ CGIntrinsicsOpenMP::CGIntrinsicsOpenMP(Module &M) : OMPBuilder(M), M(M) { /* HasRequiresDynamicAllocators */ false); OMPBuilder.initialize(); - TgtOffloadEntryTy = StructType::create({OMPBuilder.Int8Ptr, - OMPBuilder.Int8Ptr, OMPBuilder.SizeTy, - OMPBuilder.Int32, OMPBuilder.Int32}, - "struct.__tgt_offload_entry"); + OMPBuilder.loadOffloadInfoMetadata(M); + OMPBuilder.OffloadInfoManager.actOnTargetRegionEntriesInfo( + [&](const TargetRegionEntryInfo &EntryInfo, + const OffloadEntriesInfoManager::OffloadEntryInfoTargetRegion + &Entry) { + SmallString<128> Name; + TargetRegionEntryInfo::getTargetRegionEntryFnName( + Name, EntryInfo.ParentName, EntryInfo.DeviceID, EntryInfo.FileID, + EntryInfo.Line, EntryInfo.Count); + OffloadEntryNames[Entry.getOrder()] = std::string(Name); + }); + // OpenMP device runtime expects this global that controls debugging, default // to 0 (no debugging enabled). if (isOpenMPDeviceRuntime()) { @@ -1679,35 +1690,14 @@ void CGIntrinsicsOpenMP::emitOMPTask(DSAValueMapTy &DSAValueMap, Function *Fn, } } -void CGIntrinsicsOpenMP::emitOMPOffloadingEntry(const Twine &DevFuncName, - Value *EntryPtr, - Constant *&OMPOffloadEntry) { - - Constant *DevFuncNameConstant = - ConstantDataArray::getString(M.getContext(), DevFuncName.str()); - auto *GV = new GlobalVariable( - M, DevFuncNameConstant->getType(), - /* isConstant */ true, GlobalValue::InternalLinkage, DevFuncNameConstant, - ".omp_offloading.entry_name", nullptr, GlobalVariable::NotThreadLocal, - /* AddressSpace */ 0); - GV->setUnnamedAddr(GlobalValue::UnnamedAddr::Global); - - Constant *EntryConst = dyn_cast(EntryPtr); - assert(EntryConst && "Expected constant entry pointer"); - OMPOffloadEntry = ConstantStruct::get( - TgtOffloadEntryTy, - ConstantExpr::getPointerBitCastOrAddrSpaceCast(EntryConst, - OMPBuilder.VoidPtr), - ConstantExpr::getPointerBitCastOrAddrSpaceCast(GV, OMPBuilder.Int8Ptr), - ConstantInt::get(OMPBuilder.SizeTy, 0), - ConstantInt::get(OMPBuilder.Int32, 0), - ConstantInt::get(OMPBuilder.Int32, 0)); - auto *OMPOffloadEntryGV = new GlobalVariable( - M, TgtOffloadEntryTy, - /* isConstant */ true, GlobalValue::WeakAnyLinkage, OMPOffloadEntry, - ".omp_offloading.entry." + DevFuncName); - OMPOffloadEntryGV->setSection("omp_offloading_entries"); - OMPOffloadEntryGV->setAlignment(Align(1)); +GlobalVariable * +CGIntrinsicsOpenMP::emitOMPOffloadingEntry(StringRef DevFuncName, + Constant *Addr) { + GlobalVariable *EntryGV = offloading::emitOffloadingEntry( + M, object::OffloadKind::OFK_OpenMP, Addr, DevFuncName, /* Size */ 0, + /* Flags */ 0, /* Data */ 0); + OffloadEntryGVs[DevFuncName] = EntryGV; + return EntryGV; } void CGIntrinsicsOpenMP::emitOMPOffloadingMappings( @@ -2055,150 +2045,80 @@ void CGIntrinsicsOpenMP::emitOMPTaskwait(BasicBlock *BBEntry) { } GlobalVariable * -CGIntrinsicsOpenMP::emitOffloadingGlobals(StringRef DevWrapperFuncName, - ConstantDataArray *ELF) { - GlobalVariable *OMPRegionId = nullptr; - GlobalVariable *OMPOffloadEntries = nullptr; - - OMPRegionId = new GlobalVariable( +CGIntrinsicsOpenMP::emitOffloadingGlobals(StringRef DevWrapperFuncName) { + GlobalVariable *OMPRegionId = new GlobalVariable( M, OMPBuilder.Int8, /* isConstant */ true, GlobalValue::WeakAnyLinkage, ConstantInt::get(OMPBuilder.Int8, 0), DevWrapperFuncName + ".region_id", nullptr, GlobalVariable::NotThreadLocal, /* AddressSpace */ 0); - Constant *OMPOffloadEntry; - CGIntrinsicsOpenMP::emitOMPOffloadingEntry(DevWrapperFuncName, OMPRegionId, - OMPOffloadEntry); - - // TODO: do this at finalization when all entries have been - // found. - // TODO: assumes 1 device image, can we call tgt_register_lib - // multiple times? - auto *ArrayTy = ArrayType::get(TgtOffloadEntryTy, 1); - OMPOffloadEntries = - new GlobalVariable(M, ArrayTy, - /* isConstant */ true, GlobalValue::InternalLinkage, - ConstantArray::get(ArrayTy, {OMPOffloadEntry}), - ".omp_offloading.entries"); - - assert(OMPRegionId && "Expected non-null omp region id global"); - assert(OMPOffloadEntries && - "Expected non-null omp offloading entries constant"); - - auto EmitOffloadingBinaryGlobals = [&]() { - auto *GV = new GlobalVariable(M, ELF->getType(), /* isConstant */ true, - GlobalValue::InternalLinkage, ELF, - ".omp_offloading.device_image"); - GV->setUnnamedAddr(GlobalValue::UnnamedAddr::Global); - - auto &Ctx = M.getContext(); - - StructType *TgtDeviceImageTy = StructType::create( - {OMPBuilder.Int8Ptr, OMPBuilder.Int8Ptr, PointerType::getUnqual(Ctx), - PointerType::getUnqual(Ctx)}, - "struct.__tgt_device_image"); - - StructType *TgtBinDescTy = StructType::create( - {OMPBuilder.Int32, PointerType::getUnqual(Ctx), - PointerType::getUnqual(Ctx), PointerType::getUnqual(Ctx)}, - "struct.__tgt_bin_desc"); - - auto *ArrayTy = ArrayType::get(TgtDeviceImageTy, 1); - auto *Zero = ConstantInt::get(OMPBuilder.SizeTy, 0); - auto *One = ConstantInt::get(OMPBuilder.SizeTy, 1); - auto *Size = ConstantInt::get(OMPBuilder.SizeTy, ELF->getNumElements()); - Constant *ZeroZero[] = {Zero, Zero}; - Constant *ZeroOne[] = {Zero, One}; - Constant *ZeroSize[] = {Zero, Size}; - - auto *ImageB = - ConstantExpr::getGetElementPtr(GV->getValueType(), GV, ZeroZero); - auto *ImageE = - ConstantExpr::getGetElementPtr(GV->getValueType(), GV, ZeroSize); - auto *EntriesB = ConstantExpr::getGetElementPtr( - OMPOffloadEntries->getValueType(), OMPOffloadEntries, ZeroZero); - auto *EntriesE = ConstantExpr::getGetElementPtr( - OMPOffloadEntries->getValueType(), OMPOffloadEntries, ZeroOne); - - auto *DeviceImageEntry = ConstantStruct::get(TgtDeviceImageTy, ImageB, - ImageE, EntriesB, EntriesE); - auto *DeviceImages = - new GlobalVariable(M, ArrayTy, - /* isConstant */ true, GlobalValue::InternalLinkage, - ConstantArray::get(ArrayTy, {DeviceImageEntry}), - ".omp_offloading.device_images"); - - auto *ImagesB = ConstantExpr::getGetElementPtr(DeviceImages->getValueType(), - DeviceImages, ZeroZero); - auto *DescInit = - ConstantStruct::get(TgtBinDescTy, - ConstantInt::get(OMPBuilder.Int32, - /* number of images */ 1), - ImagesB, EntriesB, EntriesE); - auto *BinDesc = - new GlobalVariable(M, DescInit->getType(), - /* isConstant */ true, GlobalValue::InternalLinkage, - DescInit, ".omp_offloading.descriptor"); - - // Add tgt_register_lib in global ctors and tgt_unregister_lib in atexit. - auto CreateUnregFunction = [&]() { - auto *FuncTy = FunctionType::get(OMPBuilder.Void, /*isVarArg*/ false); - auto *Func = Function::Create(FuncTy, GlobalValue::InternalLinkage, - ".omp_offloading.descriptor_unreg", &M); - Func->setSection(".text.startup"); - - // Get __tgt_unregister_lib function declaration. - auto *UnRegFuncTy = - FunctionType::get(OMPBuilder.Void, PointerType::getUnqual(Ctx), - /*isVarArg*/ false); - FunctionCallee UnRegFuncC = - M.getOrInsertFunction("__tgt_unregister_lib", UnRegFuncTy); - - // Construct function body - IRBuilder<> Builder(BasicBlock::Create(M.getContext(), "entry", Func)); - Builder.CreateCall(UnRegFuncC, BinDesc); - Builder.CreateRetVoid(); - - return Func; - }; + emitOMPOffloadingEntry(DevWrapperFuncName, OMPRegionId); - // Create the registration function constructor. - auto *FuncTy = FunctionType::get(OMPBuilder.Void, /*isVarArg*/ false); - auto *Func = Function::Create(FuncTy, GlobalValue::InternalLinkage, - ".omp_offloading.descriptor_reg", &M); - Func->setSection(".text.startup"); - - // Get __tgt_register_lib function declaration. - auto *RegFuncTy = - FunctionType::get(OMPBuilder.Void, PointerType::getUnqual(Ctx), - /*isVarArg*/ false); - FunctionCallee RegFuncC = - M.getOrInsertFunction("__tgt_register_lib", RegFuncTy); - - // Get atexit function declaration. - auto *AtExitTy = - FunctionType::get(OMPBuilder.Int32, PointerType::getUnqual(Ctx), - /*isVarArg=*/false); - FunctionCallee AtExit = M.getOrInsertFunction("atexit", AtExitTy); - - // Construct function body. - IRBuilder<> Builder(BasicBlock::Create(M.getContext(), "entry", Func)); - Builder.CreateCall(RegFuncC, BinDesc); - - Function *UnregFunc = CreateUnregFunction(); - Builder.CreateCall(AtExit, UnregFunc); - - Builder.CreateRetVoid(); - - // Add this function to constructors. - // Set priority to 101 so that __tgt_register_lib is executed after system - // constructors but before user constructors. - appendToGlobalCtors(M, Func, /*Priority*/ 101); - }; + return OMPRegionId; +} - EmitOffloadingBinaryGlobals(); +void CGIntrinsicsOpenMP::emitOffloadImageDescriptors() { + NamedMDNode *ImagesMD = M.getNamedMetadata("pyomp.offload_images"); + if (!ImagesMD) + return; + + SmallVector> Images; + for (MDNode *Image : ImagesMD->operands()) + Images.emplace_back( + mdconst::extract(Image->getOperand(0))->getZExtValue(), + mdconst::extract(Image->getOperand(1))); + ImagesMD->eraseFromParent(); + + for (auto [Order, ImageGV] : Images) { + TargetInfoStruct TargetInfo; + TargetInfo.OffloadEntryIdx = Order; + std::string EntryName = getOffloadEntryName(TargetInfo); + GlobalVariable *EntryGV = OffloadEntryGVs.lookup(EntryName); + if (!EntryGV) + FATAL_ERROR("Missing offload entry for " + EntryName); + + // Each image is compiled from one target region, so its descriptor holds + // that region's entry only. + std::string Suffix = "." + std::to_string(Order); + auto *ArrayTy = ArrayType::get(EntryGV->getValueType(), 1); + auto *Entries = new GlobalVariable( + M, ArrayTy, /* isConstant */ true, GlobalValue::InternalLinkage, + ConstantArray::get(ArrayTy, {EntryGV->getInitializer()}), + ".omp_offloading.entries" + Suffix); + // The wrapper takes the entry bounds as globals. No linker defines them + // under JIT, so pass a placeholder end and replace it with the address + // one past the array. + auto *EntriesEnd = new GlobalVariable( + M, EntryGV->getValueType(), /* isConstant */ true, + GlobalValue::ExternalLinkage, nullptr, + ".omp_offloading.entries_end" + Suffix); + + StringRef Bytes = + cast(ImageGV->getInitializer())->getRawDataValues(); + ArrayRef Image(Bytes.data(), Bytes.size()); + if (Error Err = offloading::wrapOpenMPBinaries(M, {Image}, + {Entries, EntriesEnd}, Suffix)) + FATAL_ERROR("Error wrapping the offload image: " + + toString(std::move(Err))); + + EntriesEnd->replaceAllUsesWith(ConstantExpr::getGetElementPtr( + ArrayTy, Entries, + ArrayRef{ConstantInt::get(OMPBuilder.Int64, 0), + ConstantInt::get(OMPBuilder.Int64, 1)})); + EntriesEnd->eraseFromParent(); + ImageGV->eraseFromParent(); + } +} - return OMPRegionId; +std::string +CGIntrinsicsOpenMP::getOffloadEntryName(const TargetInfoStruct &TargetInfo) { + if (!TargetInfo.OffloadEntryIdx) + FATAL_ERROR("Expected QUAL.OMP.OFFLOAD.ENTRY.IDX on the target region"); + auto It = OffloadEntryNames.find(*TargetInfo.OffloadEntryIdx); + if (It == OffloadEntryNames.end()) + FATAL_ERROR("Missing omp_offload.info record for offload entry " + + std::to_string(*TargetInfo.OffloadEntryIdx)); + return It->second; } void CGIntrinsicsOpenMP::emitOMPTarget(Function *Fn, BasicBlock *EntryBB, @@ -2221,10 +2141,10 @@ void CGIntrinsicsOpenMP::emitOMPTargetHost( DSAValueMapTy &DSAValueMap, StructMapTy &StructMappingInfoMap, TargetInfoStruct &TargetInfo, OMPLoopInfoStruct *OMPLoopInfo) { - Twine DevWrapperFuncName = getDevWrapperFuncPrefix() + TargetInfo.DevFuncName; + std::string EntryName = getOffloadEntryName(TargetInfo); GlobalVariable *OMPRegionId = - emitOffloadingGlobals(DevWrapperFuncName.str(), TargetInfo.ELF); + emitOffloadingGlobals(EntryName); const DebugLoc DL = EntryBB->getTerminator()->getDebugLoc(); OpenMPIRBuilder::LocationDescription Loc( @@ -2347,7 +2267,7 @@ void CGIntrinsicsOpenMP::emitOMPTargetDevice(Function *Fn, BasicBlock *EntryBB, } } - Twine DevWrapperFuncName = getDevWrapperFuncPrefix() + Fn->getName(); + std::string DevWrapperFuncName = getOffloadEntryName(TargetInfo); FunctionType *NumbaWrapperFnTy = FunctionType::get(OMPBuilder.Void, WrapperArgsTypes, /* isVarArg */ false); @@ -2447,9 +2367,7 @@ void CGIntrinsicsOpenMP::emitOMPTargetDevice(Function *Fn, BasicBlock *EntryBB, } else { // Generating an offloading entry is required by the x86_64 plugin. - Constant *OMPOffloadEntry; - emitOMPOffloadingEntry(DevWrapperFuncName, NumbaWrapperFunc, - OMPOffloadEntry); + emitOMPOffloadingEntry(DevWrapperFuncName, NumbaWrapperFunc); } // Add llvm.module.flags for "openmp", "openmp-device" to enable // OpenMPOpt. diff --git a/src/numba/openmp/libs/pass/CGIntrinsicsOpenMP.h b/src/numba/openmp/libs/pass/CGIntrinsicsOpenMP.h index 11b1b8ea8ddc..3ca3458ffeb5 100644 --- a/src/numba/openmp/libs/pass/CGIntrinsicsOpenMP.h +++ b/src/numba/openmp/libs/pass/CGIntrinsicsOpenMP.h @@ -4,6 +4,7 @@ #include "DebugOpenMP.h" #include +#include #include #include #include @@ -223,8 +224,8 @@ struct OMPDistributeInfoStruct { }; struct TargetInfoStruct { - StringRef DevFuncName; - ConstantDataArray *ELF = nullptr; + // Order of the region's omp_offload.info record. + std::optional OffloadEntryIdx; Value *NumTeams = nullptr; Value *ThreadLimit = nullptr; OMPTgtExecModeFlags ExecMode = OMPTgtExecModeFlags::OMP_TGT_EXEC_MODE_GENERIC; @@ -422,9 +423,6 @@ class CGIntrinsicsOpenMP { OpenMPIRBuilder OMPBuilder; Module &M; - StructType *TgtOffloadEntryTy; - - StructType *getTgtOffloadEntryTy() { return TgtOffloadEntryTy; } void emitOMPParallel(DSAValueMapTy &DSAValueMap, ValueToValueMapTy *VMap, const DebugLoc &DL, Function *Fn, BasicBlock *BBEntry, @@ -440,8 +438,8 @@ class CGIntrinsicsOpenMP { BasicBlock *BBEntry, BasicBlock *StartBB, BasicBlock *EndBB, BasicBlock *AfterBB); - void emitOMPOffloadingEntry(const Twine &DevFuncName, Value *EntryPtr, - Constant *&OMPOffloadEntry); + GlobalVariable *emitOMPOffloadingEntry(StringRef DevFuncName, + Constant *Addr); void emitOMPOffloadingMappings(InsertPointTy AllocaIP, DSAValueMapTy &DSAValueMap, @@ -509,10 +507,17 @@ class CGIntrinsicsOpenMP { StructMapTy &StructMappingInfoMap, bool IsDeviceTargetRegion); - GlobalVariable *emitOffloadingGlobals(StringRef DevWrapperFuncName, - ConstantDataArray *ELF); + GlobalVariable *emitOffloadingGlobals(StringRef DevWrapperFuncName); + + // Wrap each device image listed in the module's pyomp.offload_images into a + // binary descriptor registered with libomptarget. + void emitOffloadImageDescriptors(); - Twine getDevWrapperFuncPrefix() { return "__omp_offload_numba_"; } + // Kernel and entry name of a target region, built from the module's + // omp_offload.info record for the region's QUAL.OMP.OFFLOAD.ENTRY.IDX. + std::string getOffloadEntryName(const TargetInfoStruct &TargetInfo); + DenseMap OffloadEntryNames; + StringMap OffloadEntryGVs; OutlinedInfoStruct createOutlinedFunction(DSAValueMapTy &DSAValueMap, ValueToValueMapTy *VMap, diff --git a/src/numba/openmp/libs/pass/CMakeLists.txt b/src/numba/openmp/libs/pass/CMakeLists.txt index fcc977fc4c77..c2ec53aef9bc 100644 --- a/src/numba/openmp/libs/pass/CMakeLists.txt +++ b/src/numba/openmp/libs/pass/CMakeLists.txt @@ -35,7 +35,7 @@ if(DEFINED ENV{USE_CXX11_ABI}) endif() # Use static library components to avoid issues with shared library dependencies. -set(llvm_libs LLVMPasses) +set(llvm_libs LLVMPasses LLVMFrontendOffloading) if(NOT APPLE) target_link_options(IntrinsicsOpenMP PRIVATE "-Wl,--no-undefined") diff --git a/src/numba/openmp/libs/pass/IntrinsicsOpenMP.cpp b/src/numba/openmp/libs/pass/IntrinsicsOpenMP.cpp index d415cbbbf150..8992d8c0d967 100644 --- a/src/numba/openmp/libs/pass/IntrinsicsOpenMP.cpp +++ b/src/numba/openmp/libs/pass/IntrinsicsOpenMP.cpp @@ -389,21 +389,12 @@ struct IntrinsicsOpenMP { assert(O.input_size() == 1 && "Expected single if condition value"); ParRegionInfo.IfCondition = TagInputs[0]; - } else if (Tag.starts_with("QUAL.OMP.TARGET.DEV_FUNC")) { + } else if (Tag.starts_with("QUAL.OMP.OFFLOAD.ENTRY.IDX")) { assert(O.input_size() == 1 && - "Expected a single device function name"); - ConstantDataArray *DevFuncArray = - dyn_cast(TagInputs[0]); - assert(DevFuncArray && - "Expected constant string for the device function"); - TargetInfo.DevFuncName = DevFuncArray->getAsString(); - } else if (Tag.starts_with("QUAL.OMP.TARGET.ELF")) { - assert(O.input_size() == 1 && - "Expected a single elf image string"); - ConstantDataArray *ELF = - dyn_cast(TagInputs[0]); - assert(ELF && "Expected constant string for ELF"); - TargetInfo.ELF = ELF; + "Expected a single offload entry index"); + auto *EntryIdx = dyn_cast(TagInputs[0]); + assert(EntryIdx && "Expected constant offload entry index"); + TargetInfo.OffloadEntryIdx = EntryIdx->getZExtValue(); } else if (Tag.starts_with("QUAL.OMP.DEVICE")) { assert(O.input_size() == 1 && "Expected a single device id value"); @@ -707,6 +698,8 @@ struct IntrinsicsOpenMP { } } + CGIOMP.emitOffloadImageDescriptors(); + DEBUG_ENABLE(dbgs() << "=== Dump Lowered Module\n" << M << "=== End of Dump Lowered Module\n"); diff --git a/src/numba/openmp/omp_ir.py b/src/numba/openmp/omp_ir.py index c09c8288552e..4540353a448d 100644 --- a/src/numba/openmp/omp_ir.py +++ b/src/numba/openmp/omp_ir.py @@ -25,6 +25,7 @@ import sys import os import copy +import zlib import tempfile import subprocess import operator @@ -46,7 +47,6 @@ openmp_tag_list_to_str, list_vars_from_tags, get_tags_of_type, - StringLiteral, openmp_tag, NameSlice, ) @@ -79,6 +79,78 @@ def get_unique(): return ret +def reserve_offload_entry(module, parent_name, loc): + """Reserve the offload entry of a target region in the host module. + + The entry is identified as clang identifies it: device id, + file id, parent function, line, and a count that separates regions on the + same line, plus the region's order in the module. The LLVM pass derives + the kernel and entry names from the omp_offload.info record. + """ + state = getattr(module, "_pyomp_offload_state", None) + if state is None: + state = {"next_order": 0, "line_counts": {}} + module._pyomp_offload_state = state + + line = max(int(getattr(loc, "line", 0) or 0), 0) + key = (parent_name, line) + count = state["line_counts"].get(key, 0) + state["line_counts"][key] = count + 1 + order = state["next_order"] + state["next_order"] = order + 1 + + filename = getattr(loc, "filename", None) or "" + return { + "device_id": 0, + "file_id": zlib.crc32(filename.encode("utf-8")), + "parent_name": parent_name, + "line": line, + "count": count, + "order": order, + } + + +def add_offload_info_metadata(module, entry): + """Add the entry's omp_offload.info record, in the layout that + OpenMPIRBuilder::loadOffloadInfoMetadata reads: kind (0 = target region), + device id, file id, parent name, line, count, order.""" + + def i32(value): + # The pass reads the fields zero-extended, so emit the signed form. + if value >= 2**31: + value -= 2**32 + return lir.Constant(lir.IntType(32), value) + + module.add_named_metadata( + "omp_offload.info", + [ + i32(0), + i32(entry["device_id"]), + i32(entry["file_id"]), + entry["parent_name"], + i32(entry["line"]), + i32(entry["count"]), + i32(entry["order"]), + ], + ) + + +def add_offload_image(module, entry, image): + """Embed a target region's device image in the host module. The pass + registers each image listed in pyomp.offload_images with libomptarget.""" + data = lir.Constant(lir.ArrayType(lir.IntType(8), len(image)), bytearray(image)) + image_gv = lir.GlobalVariable( + module, data.type, name=f".pyomp.offload_image.{entry['order']}" + ) + image_gv.linkage = "private" + image_gv.global_constant = True + image_gv.initializer = data + module.add_named_metadata( + "pyomp.offload_images", + [lir.Constant(lir.IntType(32), entry["order"]), image_gv], + ) + + def openmp_region_alloca(obj, alloca_instr, typ): obj.alloca(alloca_instr, typ) @@ -166,7 +238,7 @@ def _internalize(): continue if func.linkage != ll.Linkage.external: continue - if "__omp_offload_numba" in func.name: + if func.name.startswith("__omp_offloading_"): continue func.linkage = "internal" @@ -1271,6 +1343,16 @@ def add_struct_tags(self, var_table): elif target_num is not None and not self.target_copy: var_table = get_name_var_table(lowerer.func_ir.blocks) + # The device copy made below carries the same entry, so both + # modules describe this region with the same omp_offload.info record. + self.offload_entry = reserve_offload_entry( + mod, lowerer.fndesc.mangled_name, self.loc + ) + add_offload_info_metadata(mod, self.offload_entry) + host_side_target_tags.append( + openmp_tag("QUAL.OMP.OFFLOAD.ENTRY.IDX", self.offload_entry["order"]) + ) + ompx_attrs = list( filter(lambda x: x.name == "QUAL.OMP.OMPX_ATTRIBUTE", self.tags) ) @@ -1390,6 +1472,13 @@ def fixup_openmp_pairs(blocks): # extraneous arguments in the kernel function. if start_region.has_target() == target_num: start_region.tags.append(openmp_tag("OMP.DEVICE")) + start_region.offload_entry = self.offload_entry + start_region.tags.append( + openmp_tag( + "QUAL.OMP.OFFLOAD.ENTRY.IDX", + self.offload_entry["order"], + ) + ) end_region = blocks[end_block].body[ebindex] # assert(start_region.omp_region_var is None) # Make start and end copies point at each other. @@ -1820,19 +1909,16 @@ def prepend_device_to_func_name(outlined_ir): f"Unsupported OpenMP device number {selected_device}, type {device_type}, vendor {device_vendor}, arch {get_device_arch(selected_device)}" ) - host_side_target_tags.append( - openmp_tag( - "QUAL.OMP.TARGET.DEV_FUNC", - StringLiteral(cres.fndesc.mangled_name.encode("utf-8")), - ) - ) - host_side_target_tags.append( - openmp_tag("QUAL.OMP.TARGET.ELF", StringLiteral(target_elf)) - ) + add_offload_image(mod, self.offload_entry, target_elf) if DEBUG_OPENMP >= 1: dprint_func_ir(func_ir, "target after outline compiled func_ir") + # The device copy of a target region records the host's entry in the + # device module, so the pass derives the same kernel name on both sides. + if self.target_copy and get_tags_of_type(self.tags, "OMP.DEVICE"): + add_offload_info_metadata(mod, self.offload_entry) + llvm_token_t = TokenType() fnty = lir.FunctionType(llvm_token_t, []) tags_to_include = self.tags + host_side_target_tags diff --git a/src/numba/openmp/tests/test_openmp.py b/src/numba/openmp/tests/test_openmp.py index 4e0e0425b60f..c53abeacd55c 100644 --- a/src/numba/openmp/tests/test_openmp.py +++ b/src/numba/openmp/tests/test_openmp.py @@ -3633,6 +3633,40 @@ def test_impl(a): r = test_impl(a) np.testing.assert_array_equal(r, np.full(n, 43)) + def target_two_regions(self, device): + target_pragma = f"target device({device}) map(tofrom: a)" + + @njit + def test_impl(a): + with openmp(target_pragma): + for i in range(len(a)): + a[i] += 1 + with openmp(target_pragma): + for i in range(len(a)): + a[i] *= 2 + return a + + n = 10 + a = np.full(n, 42) + r = test_impl(a) + np.testing.assert_array_equal(r, np.full(n, 86)) + + def target_two_signatures(self, device): + target_pragma = f"target device({device}) map(tofrom: a)" + + @njit + def test_impl(a): + with openmp(target_pragma): + for i in range(len(a)): + a[i] += 1 + return a + + n = 10 + r_int = test_impl(np.full(n, 42)) + r_float = test_impl(np.full(n, 42.5)) + np.testing.assert_array_equal(r_int, np.full(n, 43)) + np.testing.assert_array_equal(r_float, np.full(n, 43.5)) + def target_nest_parallel_for(self, device): target_pragma = f"target device({device}) map(tofrom: a, sched)"