[flang][OpenMP] Implicitly map allocatable record fields #117867
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -24,10 +24,14 @@ | |
| /// indirectly via a parent object. | ||
| //===----------------------------------------------------------------------===// | ||
|
|
||
| #include "flang/Lower/DirectivesCommon.h" | ||
| #include "flang/Optimizer/Builder/FIRBuilder.h" | ||
| #include "flang/Optimizer/Builder/HLFIRTools.h" | ||
| #include "flang/Optimizer/Dialect/FIRType.h" | ||
| #include "flang/Optimizer/Dialect/Support/KindMapping.h" | ||
| #include "flang/Optimizer/HLFIR/HLFIROps.h" | ||
| #include "flang/Optimizer/OpenMP/Passes.h" | ||
| #include "mlir/Analysis/SliceAnalysis.h" | ||
| #include "mlir/Dialect/Func/IR/FuncOps.h" | ||
| #include "mlir/Dialect/OpenMP/OpenMPDialect.h" | ||
| #include "mlir/IR/BuiltinDialect.h" | ||
|
|
@@ -486,6 +490,160 @@ class MapInfoFinalizationPass | |
| // iterations from previous function scopes. | ||
| localBoxAllocas.clear(); | ||
|
|
||
| // First, walk `omp.map.info` ops to see if any record members should be | ||
| // implicitly mapped. | ||
| func->walk([&](mlir::omp::MapInfoOp op) { | ||
| mlir::Type underlyingType = | ||
| fir::unwrapRefType(op.getVarPtr().getType()); | ||
|
|
||
| // TODO Test with and support more complicated cases; like arrays for | ||
| // records, for example. | ||
| if (!fir::isRecordWithAllocatableMember(underlyingType)) | ||
| return mlir::WalkResult::advance(); | ||
|
|
||
| // TODO For now, only consider `omp.target` ops. Other ops that support | ||
| // `map` clauses will follow later. | ||
| mlir::omp::TargetOp target = | ||
| mlir::dyn_cast_if_present<mlir::omp::TargetOp>( | ||
|
There was a problem hiding this comment. I think this will have to work for anything that has an explicit map clause as well. From my understanding we have to map these components if someone was to write map(tofrom: some_dtype_with_allocas). Worth double checking this with @mjklemm. Hopefully isn't too hard to extend, but if it is, I'd be happy with a first pass that works for just TargetOp, which is arguably the hardest to do it for in any case :-) |
||
| getFirstTargetUser(op)); | ||
|
|
||
| if (!target) | ||
| return mlir::WalkResult::advance(); | ||
|
|
||
| auto mapClauseOwner = | ||
| llvm::dyn_cast<mlir::omp::MapClauseOwningOpInterface>(*target); | ||
|
|
||
| int64_t mapVarIdx = mapClauseOwner.getOperandIndexForMap(op); | ||
| assert(mapVarIdx >= 0 && | ||
| mapVarIdx < | ||
| static_cast<int64_t>(mapClauseOwner.getMapVars().size())); | ||
|
|
||
| auto argIface = | ||
| llvm::dyn_cast<mlir::omp::BlockArgOpenMPOpInterface>(*target); | ||
| // TODO How should `map` block argument that correspond to: `private`, | ||
| // `use_device_addr`, `use_device_ptr`, be handled? | ||
| mlir::BlockArgument opBlockArg = argIface.getMapBlockArgs()[mapVarIdx]; | ||
|
There was a problem hiding this comment. Might need to be a tad wary of how this works with use_device_addr/ptr, as they're both map info holders, and soon with Pranav's and your work I believe private may also be, and I don't think they have the same implicit connotations regular map does. We may also eventually need to be careful of this with declare mappers as well, as they're explicit user defined mappings that shouldn't have the implicit behavior but that's a problem for another day/PR when that work is further along :-) |
||
| llvm::SetVector<mlir::Operation *> mapVarForwardSlice; | ||
| mlir::getForwardSlice(opBlockArg, &mapVarForwardSlice); | ||
|
|
||
| mapVarForwardSlice.remove_if([&](mlir::Operation *sliceOp) { | ||
| // TODO Support coordinate_of ops. | ||
| // | ||
| // TODO Support call ops by recursively examining the forward slice of | ||
| // the corresponding parameter to the field in the called function. | ||
| return !mlir::isa<hlfir::DesignateOp>(sliceOp); | ||
| }); | ||
|
|
||
| auto recordType = mlir::cast<fir::RecordType>(underlyingType); | ||
| llvm::SmallVector<mlir::Value> newMapOpsForFields; | ||
| llvm::SmallVector<int64_t> fieldIndicies; | ||
|
|
||
| for (auto fieldMemTyPair : recordType.getTypeList()) { | ||
| auto &field = fieldMemTyPair.first; | ||
| auto memTy = fieldMemTyPair.second; | ||
|
|
||
| bool shouldMapField = | ||
| llvm::find_if(mapVarForwardSlice, [&](mlir::Operation *sliceOp) { | ||
| if (!fir::isAllocatableType(memTy)) | ||
| return false; | ||
|
|
||
| auto designateOp = mlir::dyn_cast<hlfir::DesignateOp>(sliceOp); | ||
| if (!designateOp) | ||
| return false; | ||
|
|
||
| return designateOp.getComponent() && | ||
| designateOp.getComponent()->strref() == field; | ||
| }) != mapVarForwardSlice.end(); | ||
|
|
||
| // TODO Handle recursive record types. Adapting | ||
| // `createParentSymAndGenIntermediateMaps` to work direclty on MLIR | ||
| // entities might be helpful here. | ||
|
|
||
| if (!shouldMapField) | ||
| continue; | ||
|
|
||
| int64_t fieldIdx = recordType.getFieldIndex(field); | ||
| bool alreadyMapped = [&]() { | ||
| if (op.getMembersIndexAttr()) | ||
| for (auto indexList : op.getMembersIndexAttr()) { | ||
| auto indexListAttr = mlir::cast<mlir::ArrayAttr>(indexList); | ||
| if (indexListAttr.size() == 1 && | ||
| mlir::cast<mlir::IntegerAttr>(indexListAttr[0]).getInt() == | ||
| fieldIdx) | ||
| return true; | ||
| } | ||
|
|
||
| return false; | ||
| }(); | ||
|
|
||
| if (alreadyMapped) | ||
| continue; | ||
|
|
||
| builder.setInsertionPoint(op); | ||
| mlir::Value fieldIdxVal = builder.createIntegerConstant( | ||
| op.getLoc(), mlir::IndexType::get(builder.getContext()), | ||
| fieldIdx); | ||
| auto fieldCoord = builder.create<fir::CoordinateOp>( | ||
|
There was a problem hiding this comment. for the whole recursive generation / gather of indices, you may be able to borrow / tweak (improve on ;-)) the explicit map code that does something similar in the function: " There was a problem hiding this comment. Thanks for the pointer. Also added a |
||
| op.getLoc(), builder.getRefType(memTy), op.getVarPtr(), | ||
| fieldIdxVal); | ||
| Fortran::lower::AddrAndBoundsInfo info = | ||
| Fortran::lower::getDataOperandBaseAddr( | ||
| builder, fieldCoord, /*isOptional=*/false, op.getLoc()); | ||
| llvm::SmallVector<mlir::Value> bounds = | ||
| Fortran::lower::genImplicitBoundsOps<mlir::omp::MapBoundsOp, | ||
| mlir::omp::MapBoundsType>( | ||
| builder, info, | ||
| hlfir::translateToExtendedValue(op.getLoc(), builder, | ||
| hlfir::Entity{fieldCoord}) | ||
| .first, | ||
| /*dataExvIsAssumedSize=*/false, op.getLoc()); | ||
|
|
||
| mlir::omp::MapInfoOp fieldMapOp = | ||
| builder.create<mlir::omp::MapInfoOp>( | ||
| op.getLoc(), fieldCoord.getResult().getType(), | ||
| fieldCoord.getResult(), | ||
| mlir::TypeAttr::get( | ||
| fir::unwrapRefType(fieldCoord.getResult().getType())), | ||
| /*varPtrPtr=*/mlir::Value{}, | ||
| /*members=*/mlir::ValueRange{}, | ||
| /*members_index=*/mlir::ArrayAttr{}, | ||
| /*bounds=*/bounds, op.getMapTypeAttr(), | ||
| builder.getAttr<mlir::omp::VariableCaptureKindAttr>( | ||
|
There was a problem hiding this comment. should probably take the capture kind from the original operation as well, to keep it inline with the map type attribute, may save some possible weirdness :-) But taking the map type from the parent does raise the interesting question of if the implicit mapping should take on the specified map type behavior of it's parent or if it should default to the regular implicit map behaviour of it's type. Not too sure on that one, the spec or one of our local specification gurus may have some insight into that. There was a problem hiding this comment. Maybe @mjklemm can provide some guidance here? There was a problem hiding this comment. Yes, I think that's the right behavior. However, we need to bear in mind that a user might have used |
||
| mlir::omp::VariableCaptureKind::ByRef), | ||
| builder.getStringAttr(op.getNameAttr().strref() + "." + | ||
| field + ".implicit_map"), | ||
| /*partial_map=*/builder.getBoolAttr(false)); | ||
| newMapOpsForFields.emplace_back(fieldMapOp); | ||
| fieldIndicies.emplace_back(fieldIdx); | ||
| } | ||
|
|
||
| if (newMapOpsForFields.empty()) | ||
| return mlir::WalkResult::advance(); | ||
|
|
||
| op.getMembersMutable().append(newMapOpsForFields); | ||
| llvm::SmallVector<llvm::SmallVector<int64_t>> newMemberIndices; | ||
| mlir::ArrayAttr oldMembersIdxAttr = op.getMembersIndexAttr(); | ||
|
|
||
| if (oldMembersIdxAttr) | ||
| for (mlir::Attribute indexList : oldMembersIdxAttr) { | ||
| llvm::SmallVector<int64_t> listVec; | ||
|
|
||
| for (mlir::Attribute index : mlir::cast<mlir::ArrayAttr>(indexList)) | ||
| listVec.push_back(mlir::cast<mlir::IntegerAttr>(index).getInt()); | ||
|
|
||
| newMemberIndices.emplace_back(std::move(listVec)); | ||
| } | ||
|
|
||
| for (int64_t newFieldIdx : fieldIndicies) | ||
| newMemberIndices.emplace_back( | ||
| llvm::SmallVector<int64_t>(1, newFieldIdx)); | ||
|
|
||
| op.setMembersIndexAttr(builder.create2DI64ArrayAttr(newMemberIndices)); | ||
| op.setPartialMap(true); | ||
|
|
||
| return mlir::WalkResult::advance(); | ||
| }); | ||
|
|
||
| func->walk([&](mlir::omp::MapInfoOp op) { | ||
| // TODO: Currently only supports a single user for the MapInfoOp. This | ||
| // is fine for the moment, as the Fortran frontend will generate a | ||
|
|
||
Oops, something went wrong.
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Uh oh!
There was an error while loading. Please reload this page.