@@ -300,10 +300,6 @@ static LogicalResult checkImplementationStatus(Operation &op) {
300300 if (privatizer.getDataSharingType () ==
301301 omp::DataSharingClauseType::FirstPrivate)
302302 result = todo (" firstprivate" );
303-
304- if (!privatizer.getDeallocRegion ().empty ())
305- result = op.emitError (" not yet implemented: privatization of "
306- " structures in omp.target operation" );
307303 }
308304 }
309305 checkThreadLimit (op, result);
@@ -3810,6 +3806,43 @@ createDeviceArgumentAccessor(MapInfoData &mapData, llvm::Argument &arg,
38103806 return builder.saveIP ();
38113807}
38123808
3809+ // / Return the llvm::Value * corresponding to the `privateVar` that
3810+ // / is being privatized. It isn't always as simple as looking up
3811+ // / moduleTranslation with privateVar. For instance, in case of
3812+ // / an allocatable, the descriptor for the allocatable is privatized.
3813+ // / This descriptor is mapped using an MapInfoOp. So, this function
3814+ // / will return a pointer to the llvm::Value corresponding to the
3815+ // / block argument for the mapped descriptor.
3816+ static llvm::Value *
3817+ findHostAssociatedValue (Value privateVar, omp::TargetOp targetOp,
3818+ llvm::DenseMap<Value, int > &mappedPrivateVars,
3819+ llvm::IRBuilderBase &builder,
3820+ LLVM::ModuleTranslation &moduleTranslation) {
3821+ if (mappedPrivateVars.contains (privateVar)) {
3822+ int blockArgIndex = mappedPrivateVars[privateVar];
3823+ Value blockArg = targetOp.getRegion ().getArgument (blockArgIndex);
3824+ Type privVarType = privateVar.getType ();
3825+ Type blockArgType = blockArg.getType ();
3826+ assert (isa<LLVM::LLVMPointerType>(blockArgType) &&
3827+ " A block argument corresponding to a mapped var should have "
3828+ " !llvm.ptr type" );
3829+
3830+ if (privVarType == blockArg.getType ())
3831+ return moduleTranslation.lookupValue (blockArg);
3832+
3833+ if (!isa<LLVM::LLVMPointerType>(privVarType)) {
3834+ // This typically happens when the privatized type is lowered from
3835+ // boxchar<KIND> and gets lowered to !llvm.struct<(ptr, i64)>. That is the
3836+ // struct/pair is passed by value. But, mapped values are passed only as
3837+ // pointers, so before we privatize, we must load the pointer.
3838+ llvm::Value *load =
3839+ builder.CreateLoad (moduleTranslation.convertType (privVarType),
3840+ moduleTranslation.lookupValue (blockArg));
3841+ return load;
3842+ }
3843+ }
3844+ return moduleTranslation.lookupValue (privateVar);
3845+ }
38133846static LogicalResult
38143847convertOmpTarget (Operation &opInst, llvm::IRBuilderBase &builder,
38153848 LLVM::ModuleTranslation &moduleTranslation) {
@@ -3821,6 +3854,19 @@ convertOmpTarget(Operation &opInst, llvm::IRBuilderBase &builder,
38213854 bool isTargetDevice = ompBuilder->Config .isTargetDevice ();
38223855 auto parentFn = opInst.getParentOfType <LLVM::LLVMFuncOp>();
38233856 auto &targetRegion = targetOp.getRegion ();
3857+ // Holds the private vars that have been mapped along with
3858+ // the block argument that corresponds to the MapInfoOp
3859+ // corresponding to the private var in question.
3860+ // So, for instance
3861+ //
3862+ // %10 = omp.map.info var_ptr(%6#0 : !fir.ref<!fir.box<!fir.heap<i32>>>, ..)
3863+ // omp.target map_entries(%10 -> %arg0) private(@box.privatizer %6#0-> %arg1)
3864+ //
3865+ // Then, %10 has been created so that the descriptor can be used by the
3866+ // privatizer
3867+ // @box.privatizer on the device side. Here we'd record {%6#0, 0} in the
3868+ // mappedPrivateVars map.
3869+ llvm::DenseMap<Value, int > mappedPrivateVars;
38243870 DataLayout dl = DataLayout (opInst.getParentOfType <ModuleOp>());
38253871 SmallVector<Value> mapVars = targetOp.getMapVars ();
38263872 ArrayRef<BlockArgument> mapBlockArgs =
@@ -3832,6 +3878,55 @@ convertOmpTarget(Operation &opInst, llvm::IRBuilderBase &builder,
38323878 bool isOffloadEntry =
38333879 isTargetDevice || !ompBuilder->Config .TargetTriples .empty ();
38343880
3881+ // For some private variables, the MapsForPrivatizedVariablesPass
3882+ // creates MapInfoOp instances. Go through the private variables and
3883+ // the mapped variables so that during codegeneration we are able
3884+ // to quickly look up the corresponding map variable, if any for each
3885+ // private variable.
3886+ if (!targetOp.getPrivateVars ().empty () && !targetOp.getMapVars ().empty ()) {
3887+ auto argIface = llvm::cast<omp::BlockArgOpenMPOpInterface>(*targetOp);
3888+ OperandRange privateVars = targetOp.getPrivateVars ();
3889+ std::optional<ArrayAttr> privateSyms = targetOp.getPrivateSyms ();
3890+ std::optional<DenseI64ArrayAttr> privateMapIndices =
3891+ targetOp.getPrivateMapsAttr ();
3892+
3893+ for (auto [privVarIdx, privVarSymPair] :
3894+ llvm::enumerate (llvm::zip_equal (privateVars, *privateSyms))) {
3895+ auto privVar = std::get<0 >(privVarSymPair);
3896+ auto privSym = std::get<1 >(privVarSymPair);
3897+
3898+ SymbolRefAttr privatizerName = llvm::cast<SymbolRefAttr>(privSym);
3899+ omp::PrivateClauseOp privatizer =
3900+ findPrivatizer (targetOp, privatizerName);
3901+
3902+ if (!privatizer.needsMap ())
3903+ continue ;
3904+
3905+ mlir::Value mappedValue =
3906+ targetOp.getMappedValueForPrivateVar (privVarIdx);
3907+ assert (mappedValue);
3908+
3909+ // The MapInfoOp defining the map var isn't really needed later.
3910+ // So, we don't store it in any datastructure. Instead, we just
3911+ // do some sanity checks on it right now.
3912+ auto mapInfoOp = mappedValue.getDefiningOp <omp::MapInfoOp>();
3913+ Type varType = mapInfoOp.getVarType ();
3914+
3915+ // Check #1: Check that the type of the private variable matches
3916+ // the type of the variable being mapped.
3917+ if (!isa<LLVM::LLVMPointerType>(privVar.getType ()))
3918+ assert (
3919+ varType == privVar.getType () &&
3920+ " Type of private var doesn't match the type of the mapped value" );
3921+
3922+ // Ok, only 1 sanity check for now.
3923+ // Record the index of the block argument corresponding to this
3924+ // mapvar.
3925+ mappedPrivateVars.insert ({privVar, argIface.getMapBlockArgsStart () +
3926+ (*privateMapIndices)[privVarIdx]});
3927+ }
3928+ }
3929+
38353930 using InsertPointTy = llvm::OpenMPIRBuilder::InsertPointTy;
38363931 auto bodyCB = [&](InsertPointTy allocaIP, InsertPointTy codeGenIP)
38373932 -> llvm::OpenMPIRBuilder::InsertPointOrErrorTy {
@@ -3858,9 +3953,10 @@ convertOmpTarget(Operation &opInst, llvm::IRBuilderBase &builder,
38583953 moduleTranslation.lookupValue (mapInfoOp.getVarPtr ());
38593954 moduleTranslation.mapValue (arg, mapOpValue);
38603955 }
3861-
38623956 // Do privatization after moduleTranslation has already recorded
38633957 // mapped values.
3958+ SmallVector<llvm::Value *> llvmPrivateVars;
3959+ SmallVector<Region *> privateCleanupRegions;
38643960 if (!targetOp.getPrivateVars ().empty ()) {
38653961 builder.restoreIP (allocaIP);
38663962
@@ -3876,11 +3972,13 @@ convertOmpTarget(Operation &opInst, llvm::IRBuilderBase &builder,
38763972 omp::PrivateClauseOp privatizer = findPrivatizer (&opInst, privSym);
38773973 assert (privatizer.getDataSharingType () !=
38783974 omp::DataSharingClauseType::FirstPrivate &&
3879- privatizer.getDeallocRegion ().empty () &&
38803975 " unsupported privatizer" );
3881- moduleTranslation.mapValue (privatizer.getAllocMoldArg (),
3882- moduleTranslation.lookupValue (privVar));
38833976 Region &allocRegion = privatizer.getAllocRegion ();
3977+ BlockArgument allocRegionArg = allocRegion.getArgument (0 );
3978+ moduleTranslation.mapValue (
3979+ allocRegionArg,
3980+ findHostAssociatedValue (privVar, targetOp, mappedPrivateVars,
3981+ builder, moduleTranslation));
38843982 SmallVector<llvm::Value *, 1 > yieldedValues;
38853983 if (failed (inlineConvertOmpRegions (
38863984 allocRegion, " omp.targetop.privatizer" , builder,
@@ -3889,7 +3987,12 @@ convertOmpTarget(Operation &opInst, llvm::IRBuilderBase &builder,
38893987 " failed to inline `alloc` region of `omp.private`" );
38903988 }
38913989 assert (yieldedValues.size () == 1 );
3892- moduleTranslation.mapValue (privBlockArg, yieldedValues.front ());
3990+ llvm::Value *llvmReplacementValue = yieldedValues.front ();
3991+ moduleTranslation.mapValue (privBlockArg, llvmReplacementValue);
3992+ if (!privatizer.getDeallocRegion ().empty ()) {
3993+ llvmPrivateVars.push_back (llvmReplacementValue);
3994+ privateCleanupRegions.push_back (&privatizer.getDeallocRegion ());
3995+ }
38933996 moduleTranslation.forgetMapping (allocRegion);
38943997 builder.restoreIP (builder.saveIP ());
38953998 }
@@ -3901,6 +4004,19 @@ convertOmpTarget(Operation &opInst, llvm::IRBuilderBase &builder,
39014004 return exitBlock.takeError ();
39024005
39034006 builder.SetInsertPoint (*exitBlock);
4007+ if (!llvmPrivateVars.empty ()) {
4008+ assert (llvmPrivateVars.size () == privateCleanupRegions.size () &&
4009+ " Number of private variables needing cleanup not equal to number"
4010+ " of privatizers with dealloc regions" );
4011+ if (failed (inlineOmpRegionCleanup (
4012+ privateCleanupRegions, llvmPrivateVars, moduleTranslation,
4013+ builder, " omp.targetop.private.cleanup" ,
4014+ /* shouldLoadCleanupRegionArg=*/ false ))) {
4015+ return llvm::createStringError (
4016+ " failed to inline `dealloc` region of `omp.private` "
4017+ " op in the target region" );
4018+ }
4019+ }
39044020 return builder.saveIP ();
39054021 };
39064022
0 commit comments