@@ -125,61 +125,86 @@ class MapInfoFinalizationPass
125125 // TODO: map the addendum segment of the descriptor, similarly to the
126126 // above base address/data pointer member.
127127
128- auto addOperands = [&](mlir::OperandRange &operandsArr,
129- mlir::MutableOperandRange &mutableOpRange,
130- auto directiveOp) {
128+ mlir::omp::MapInfoOp newDescParentMapOp =
129+ builder.create <mlir::omp::MapInfoOp>(
130+ op->getLoc (), op.getResult ().getType (), descriptor,
131+ mlir::TypeAttr::get (fir::unwrapRefType (descriptor.getType ())),
132+ /* varPtrPtr=*/ mlir::Value{},
133+ /* members=*/ mlir::SmallVector<mlir::Value>{baseAddr},
134+ /* members_index=*/
135+ mlir::DenseIntElementsAttr::get (
136+ mlir::VectorType::get (
137+ llvm::ArrayRef<int64_t >({1 , 1 }),
138+ mlir::IntegerType::get (builder.getContext (), 32 )),
139+ llvm::ArrayRef<int32_t >({0 })),
140+ /* bounds=*/ mlir::SmallVector<mlir::Value>{},
141+ builder.getIntegerAttr (builder.getIntegerType (64 , false ),
142+ op.getMapType ().value ()),
143+ op.getMapCaptureTypeAttr (), op.getNameAttr (),
144+ op.getPartialMapAttr ());
145+ op.replaceAllUsesWith (newDescParentMapOp.getResult ());
146+ op->erase ();
147+
148+ auto addOperands = [&](mlir::MutableOperandRange &mutableOpRange,
149+ mlir::Operation *directiveOp,
150+ unsigned blockArgInsertIndex = 0 ) {
151+ if (!llvm::is_contained (mutableOpRange.getAsOperandRange (),
152+ newDescParentMapOp.getResult ()))
153+ return ;
154+
155+ // There doesn't appear to be a simple way to convert MutableOperandRange
156+ // to a vector currently, so we instead use a for_each to populate our
157+ // vector.
131158 llvm::SmallVector<mlir::Value> newMapOps;
132- for ( size_t i = 0 ; i < operandsArr .size (); ++i) {
133- if (operandsArr[i] == op) {
134- // Push new implicit maps generated for the descriptor.
135- newMapOps.push_back (baseAddr );
159+ newMapOps. reserve (mutableOpRange .size ());
160+ llvm::for_each (
161+ mutableOpRange. getAsOperandRange (),
162+ [& newMapOps](mlir::Value oper) { newMapOps .push_back (oper); } );
136163
137- // for TargetOp's which have IsolatedFromAbove we must align the
138- // new additional map operand with an appropriate BlockArgument,
139- // as the printing and later processing currently requires a 1:1
140- // mapping of BlockArgs to MapInfoOp's at the same placement in
141- // each array (BlockArgs and MapOperands).
142- if ( directiveOp) {
143- directiveOp. getRegion (). insertArgument (i, baseAddr .getType (), loc );
144- }
164+ for ( auto mapMember : newDescParentMapOp. getMembers ()) {
165+ if ( llvm::is_contained (mutableOpRange. getAsOperandRange (), mapMember))
166+ continue ;
167+ newMapOps. push_back (mapMember);
168+ if (directiveOp) {
169+ directiveOp-> getRegion ( 0 ). insertArgument (
170+ blockArgInsertIndex, mapMember .getType (), mapMember. getLoc () );
171+ blockArgInsertIndex++;
145172 }
146- newMapOps.push_back (operandsArr[i]);
147173 }
174+
148175 mutableOpRange.assign (newMapOps);
149176 };
177+
178+ auto argIface =
179+ llvm::dyn_cast<mlir::omp::BlockArgOpenMPOpInterface>(target);
180+
150181 if (auto mapClauseOwner =
151182 llvm::dyn_cast<mlir::omp::MapClauseOwningOpInterface>(target)) {
152- mlir::OperandRange mapOperandsArr = mapClauseOwner.getMapVars ();
153183 mlir::MutableOperandRange mapMutableOpRange =
154184 mapClauseOwner.getMapVarsMutable ();
155- mlir::omp::TargetOp targetOp =
156- llvm::dyn_cast<mlir::omp::TargetOp>(target);
157- addOperands (mapOperandsArr, mapMutableOpRange, targetOp);
185+ unsigned blockArgInsertIndex =
186+ argIface
187+ ? argIface.getMapBlockArgsStart () + argIface.numMapBlockArgs ()
188+ : 0 ;
189+ addOperands (
190+ mapMutableOpRange,
191+ llvm::dyn_cast_or_null<mlir::omp::TargetOp>(argIface.getOperation ()),
192+ blockArgInsertIndex);
158193 }
194+
159195 if (auto targetDataOp = llvm::dyn_cast<mlir::omp::TargetDataOp>(target)) {
160- mlir::OperandRange useDevAddrArr = targetDataOp.getUseDeviceAddrVars ();
161196 mlir::MutableOperandRange useDevAddrMutableOpRange =
162197 targetDataOp.getUseDeviceAddrVarsMutable ();
163- addOperands (useDevAddrArr, useDevAddrMutableOpRange, targetDataOp);
164- }
198+ addOperands (useDevAddrMutableOpRange, target,
199+ argIface.getUseDeviceAddrBlockArgsStart () +
200+ argIface.numUseDeviceAddrBlockArgs ());
165201
166- mlir::Value newDescParentMapOp = builder.create <mlir::omp::MapInfoOp>(
167- op->getLoc (), op.getResult ().getType (), descriptor,
168- mlir::TypeAttr::get (fir::unwrapRefType (descriptor.getType ())),
169- /* varPtrPtr=*/ mlir::Value{},
170- /* members=*/ mlir::SmallVector<mlir::Value>{baseAddr},
171- /* members_index=*/
172- mlir::DenseIntElementsAttr::get (
173- mlir::VectorType::get (
174- llvm::ArrayRef<int64_t >({1 , 1 }),
175- mlir::IntegerType::get (builder.getContext (), 32 )),
176- llvm::ArrayRef<int32_t >({0 })),
177- /* bounds=*/ mlir::SmallVector<mlir::Value>{},
178- builder.getIntegerAttr (builder.getIntegerType (64 , false ),
179- op.getMapType ().value ()),
180- op.getMapCaptureTypeAttr (), op.getNameAttr (), op.getPartialMapAttr ());
181- op.replaceAllUsesWith (newDescParentMapOp);
182- op->erase ();
202+ mlir::MutableOperandRange useDevPtrMutableOpRange =
203+ targetDataOp.getUseDevicePtrVarsMutable ();
204+ addOperands (useDevPtrMutableOpRange, target,
205+ argIface.getUseDevicePtrBlockArgsStart () +
206+ argIface.numUseDevicePtrBlockArgs ());
207+ }
183208 }
184209
185210 // We add all mapped record members not directly used in the target region
0 commit comments