From 6331ab4737d953266ed606a7879bacbc4a8bdc48 Mon Sep 17 00:00:00 2001 From: Giorgis Georgakoudis Date: Tue, 29 Sep 2026 13:48:36 -0700 Subject: [PATCH 1/3] refactor(offload): identify target regions by omp_offload.info Replace QUAL.OMP.TARGET.DEV_FUNC with clang's region identity: an omp_offload.info record, linked from DIR.OMP.TARGET by QUAL.OMP.OFFLOAD.ENTRY.IDX. The pass derives the kernel and entry names from the record in both modules, so host and device names cannot diverge. Assisted-by: Claude Code (Claude Opus 5.5) --- .../openmp/libs/pass/CGIntrinsicsOpenMP.cpp | 29 ++++++- .../openmp/libs/pass/CGIntrinsicsOpenMP.h | 8 +- .../openmp/libs/pass/IntrinsicsOpenMP.cpp | 12 ++- src/numba/openmp/omp_ir.py | 87 +++++++++++++++++-- 4 files changed, 117 insertions(+), 19 deletions(-) diff --git a/src/numba/openmp/libs/pass/CGIntrinsicsOpenMP.cpp b/src/numba/openmp/libs/pass/CGIntrinsicsOpenMP.cpp index 62934eead366..b1f9a3b537d1 100644 --- a/src/numba/openmp/libs/pass/CGIntrinsicsOpenMP.cpp +++ b/src/numba/openmp/libs/pass/CGIntrinsicsOpenMP.cpp @@ -532,6 +532,18 @@ CGIntrinsicsOpenMP::CGIntrinsicsOpenMP(Module &M) : OMPBuilder(M), M(M) { /* HasRequiresDynamicAllocators */ false); OMPBuilder.initialize(); + 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); + }); + TgtOffloadEntryTy = StructType::create({OMPBuilder.Int8Ptr, OMPBuilder.Int8Ptr, OMPBuilder.SizeTy, OMPBuilder.Int32, OMPBuilder.Int32}, @@ -2201,6 +2213,17 @@ CGIntrinsicsOpenMP::emitOffloadingGlobals(StringRef DevWrapperFuncName, 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, BasicBlock *StartBB, BasicBlock *EndBB, DSAValueMapTy &DSAValueMap, @@ -2221,10 +2244,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, TargetInfo.ELF); const DebugLoc DL = EntryBB->getTerminator()->getDebugLoc(); OpenMPIRBuilder::LocationDescription Loc( @@ -2347,7 +2370,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); diff --git a/src/numba/openmp/libs/pass/CGIntrinsicsOpenMP.h b/src/numba/openmp/libs/pass/CGIntrinsicsOpenMP.h index 11b1b8ea8ddc..c4c5fe997e4e 100644 --- a/src/numba/openmp/libs/pass/CGIntrinsicsOpenMP.h +++ b/src/numba/openmp/libs/pass/CGIntrinsicsOpenMP.h @@ -223,7 +223,8 @@ struct OMPDistributeInfoStruct { }; struct TargetInfoStruct { - StringRef DevFuncName; + // Order of the region's omp_offload.info record. + std::optional OffloadEntryIdx; ConstantDataArray *ELF = nullptr; Value *NumTeams = nullptr; Value *ThreadLimit = nullptr; @@ -512,7 +513,10 @@ class CGIntrinsicsOpenMP { GlobalVariable *emitOffloadingGlobals(StringRef DevWrapperFuncName, ConstantDataArray *ELF); - 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; OutlinedInfoStruct createOutlinedFunction(DSAValueMapTy &DSAValueMap, ValueToValueMapTy *VMap, diff --git a/src/numba/openmp/libs/pass/IntrinsicsOpenMP.cpp b/src/numba/openmp/libs/pass/IntrinsicsOpenMP.cpp index d415cbbbf150..608362786e75 100644 --- a/src/numba/openmp/libs/pass/IntrinsicsOpenMP.cpp +++ b/src/numba/openmp/libs/pass/IntrinsicsOpenMP.cpp @@ -389,14 +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(); + "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.TARGET.ELF")) { assert(O.input_size() == 1 && "Expected a single elf image string"); diff --git a/src/numba/openmp/omp_ir.py b/src/numba/openmp/omp_ir.py index c09c8288552e..e514d13325b0 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 @@ -79,6 +80,62 @@ 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 openmp_region_alloca(obj, alloca_instr, typ): obj.alloca(alloca_instr, typ) @@ -166,7 +223,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 +1328,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 +1457,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,12 +1894,6 @@ 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)) ) @@ -1833,6 +1901,11 @@ def prepend_device_to_func_name(outlined_ir): 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 From bc8f9c2a640747d18d06fe8ada85454acf675aad Mon Sep 17 00:00:00 2001 From: Giorgis Georgakoudis Date: Tue, 29 Sep 2026 13:55:59 -0700 Subject: [PATCH 2/3] refactor(offload): embed device images with wrapOpenMPBinaries Drop QUAL.OMP.TARGET.ELF and the hand-written descriptor and registration code. Images travel as host-module globals listed in !pyomp.offload_images and get one descriptor each after lowering. No linker defines the entry bounds under JIT, so the end bound is a placeholder replaced with the array end. Add tests for two target regions in one function and for one function compiled for two signatures. Assisted-by: Claude Code (Claude Opus 5.5) --- .../openmp/libs/pass/CGIntrinsicsOpenMP.cpp | 194 ++++++------------ .../openmp/libs/pass/CGIntrinsicsOpenMP.h | 8 +- src/numba/openmp/libs/pass/CMakeLists.txt | 2 +- .../openmp/libs/pass/IntrinsicsOpenMP.cpp | 9 +- src/numba/openmp/omp_ir.py | 21 +- src/numba/openmp/tests/test_openmp.py | 34 +++ 6 files changed, 118 insertions(+), 150 deletions(-) diff --git a/src/numba/openmp/libs/pass/CGIntrinsicsOpenMP.cpp b/src/numba/openmp/libs/pass/CGIntrinsicsOpenMP.cpp index b1f9a3b537d1..1435268f8de4 100644 --- a/src/numba/openmp/libs/pass/CGIntrinsicsOpenMP.cpp +++ b/src/numba/openmp/libs/pass/CGIntrinsicsOpenMP.cpp @@ -2,6 +2,7 @@ #include "DebugOpenMP.h" #include +#include #include #include #include @@ -13,6 +14,7 @@ #include #include #include +#include #include #include #include @@ -2067,12 +2069,8 @@ 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, @@ -2082,137 +2080,63 @@ CGIntrinsicsOpenMP::emitOffloadingGlobals(StringRef DevWrapperFuncName, 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; - }; - - // 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); - }; - - EmitOffloadingBinaryGlobals(); - return OMPRegionId; } +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 = + M.getNamedGlobal(".omp_offloading.entry." + 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(); + } +} + std::string CGIntrinsicsOpenMP::getOffloadEntryName(const TargetInfoStruct &TargetInfo) { if (!TargetInfo.OffloadEntryIdx) @@ -2247,7 +2171,7 @@ void CGIntrinsicsOpenMP::emitOMPTargetHost( std::string EntryName = getOffloadEntryName(TargetInfo); GlobalVariable *OMPRegionId = - emitOffloadingGlobals(EntryName, TargetInfo.ELF); + emitOffloadingGlobals(EntryName); const DebugLoc DL = EntryBB->getTerminator()->getDebugLoc(); OpenMPIRBuilder::LocationDescription Loc( diff --git a/src/numba/openmp/libs/pass/CGIntrinsicsOpenMP.h b/src/numba/openmp/libs/pass/CGIntrinsicsOpenMP.h index c4c5fe997e4e..50bc646ce6ef 100644 --- a/src/numba/openmp/libs/pass/CGIntrinsicsOpenMP.h +++ b/src/numba/openmp/libs/pass/CGIntrinsicsOpenMP.h @@ -225,7 +225,6 @@ struct OMPDistributeInfoStruct { struct TargetInfoStruct { // Order of the region's omp_offload.info record. std::optional OffloadEntryIdx; - ConstantDataArray *ELF = nullptr; Value *NumTeams = nullptr; Value *ThreadLimit = nullptr; OMPTgtExecModeFlags ExecMode = OMPTgtExecModeFlags::OMP_TGT_EXEC_MODE_GENERIC; @@ -510,8 +509,11 @@ 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(); // 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. 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 608362786e75..8992d8c0d967 100644 --- a/src/numba/openmp/libs/pass/IntrinsicsOpenMP.cpp +++ b/src/numba/openmp/libs/pass/IntrinsicsOpenMP.cpp @@ -395,13 +395,6 @@ struct IntrinsicsOpenMP { auto *EntryIdx = dyn_cast(TagInputs[0]); assert(EntryIdx && "Expected constant offload entry index"); TargetInfo.OffloadEntryIdx = EntryIdx->getZExtValue(); - } 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; } else if (Tag.starts_with("QUAL.OMP.DEVICE")) { assert(O.input_size() == 1 && "Expected a single device id value"); @@ -705,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 e514d13325b0..4540353a448d 100644 --- a/src/numba/openmp/omp_ir.py +++ b/src/numba/openmp/omp_ir.py @@ -47,7 +47,6 @@ openmp_tag_list_to_str, list_vars_from_tags, get_tags_of_type, - StringLiteral, openmp_tag, NameSlice, ) @@ -136,6 +135,22 @@ def i32(value): ) +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) @@ -1894,9 +1909,7 @@ 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.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") 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)" From 2d63e3968260e5d414ea57b8d40ad8b3433a1c7e Mon Sep 17 00:00:00 2001 From: Giorgis Georgakoudis Date: Tue, 29 Sep 2026 14:39:37 -0700 Subject: [PATCH 3/3] refactor(offload): emit entries in llvm_offload_entries layout Use offloading::emitOffloadingEntry instead of the hand-built 5-field entry. The runtime still upgrades the legacy layout, so this only matches what clang emits. Assisted-by: Claude Code (Claude Opus 5.5) --- .../openmp/libs/pass/CGIntrinsicsOpenMP.cpp | 53 +++++-------------- .../openmp/libs/pass/CGIntrinsicsOpenMP.h | 9 ++-- 2 files changed, 16 insertions(+), 46 deletions(-) diff --git a/src/numba/openmp/libs/pass/CGIntrinsicsOpenMP.cpp b/src/numba/openmp/libs/pass/CGIntrinsicsOpenMP.cpp index 1435268f8de4..19aabf828b9d 100644 --- a/src/numba/openmp/libs/pass/CGIntrinsicsOpenMP.cpp +++ b/src/numba/openmp/libs/pass/CGIntrinsicsOpenMP.cpp @@ -3,6 +3,7 @@ #include #include +#include #include #include #include @@ -546,10 +547,6 @@ CGIntrinsicsOpenMP::CGIntrinsicsOpenMP(Module &M) : OMPBuilder(M), M(M) { OffloadEntryNames[Entry.getOrder()] = std::string(Name); }); - TgtOffloadEntryTy = StructType::create({OMPBuilder.Int8Ptr, - OMPBuilder.Int8Ptr, OMPBuilder.SizeTy, - OMPBuilder.Int32, OMPBuilder.Int32}, - "struct.__tgt_offload_entry"); // OpenMP device runtime expects this global that controls debugging, default // to 0 (no debugging enabled). if (isOpenMPDeviceRuntime()) { @@ -1693,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( @@ -2076,9 +2052,7 @@ CGIntrinsicsOpenMP::emitOffloadingGlobals(StringRef DevWrapperFuncName) { nullptr, GlobalVariable::NotThreadLocal, /* AddressSpace */ 0); - Constant *OMPOffloadEntry; - CGIntrinsicsOpenMP::emitOMPOffloadingEntry(DevWrapperFuncName, OMPRegionId, - OMPOffloadEntry); + emitOMPOffloadingEntry(DevWrapperFuncName, OMPRegionId); return OMPRegionId; } @@ -2099,8 +2073,7 @@ void CGIntrinsicsOpenMP::emitOffloadImageDescriptors() { TargetInfoStruct TargetInfo; TargetInfo.OffloadEntryIdx = Order; std::string EntryName = getOffloadEntryName(TargetInfo); - GlobalVariable *EntryGV = - M.getNamedGlobal(".omp_offloading.entry." + EntryName); + GlobalVariable *EntryGV = OffloadEntryGVs.lookup(EntryName); if (!EntryGV) FATAL_ERROR("Missing offload entry for " + EntryName); @@ -2394,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 50bc646ce6ef..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 @@ -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, @@ -519,6 +517,7 @@ class CGIntrinsicsOpenMP { // 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,