[llvm-branch-commits] [flang] [Flang][OpenMP][MLIR] Initial derived type member map support (PR #82853)

[Sergio Afonso via llvm-branch-commits]

================
@@ -116,6 +119,216 @@ void gatherFuncAndVarSyms(
     symbolAndClause.emplace_back(clause, *object.id());
 }
 
+mlir::omp::MapInfoOp
+createMapInfoOp(fir::FirOpBuilder &builder, mlir::Location loc,
+                mlir::Value baseAddr, mlir::Value varPtrPtr, std::string name,
+                llvm::ArrayRef<mlir::Value> bounds,
+                llvm::ArrayRef<mlir::Value> members,
+                mlir::DenseIntElementsAttr membersIndex, uint64_t mapType,
+                mlir::omp::VariableCaptureKind mapCaptureType, mlir::Type retTy,
+                bool partialMap) {
+  if (auto boxTy = baseAddr.getType().dyn_cast<fir::BaseBoxType>()) {
+    baseAddr = builder.create<fir::BoxAddrOp>(loc, baseAddr);
+    retTy = baseAddr.getType();
+  }
+
+  mlir::TypeAttr varType = mlir::TypeAttr::get(
+      llvm::cast<mlir::omp::PointerLikeType>(retTy).getElementType());
+
+  mlir::omp::MapInfoOp op = builder.create<mlir::omp::MapInfoOp>(
+      loc, retTy, baseAddr, varType, varPtrPtr, members, membersIndex, bounds,
+      builder.getIntegerAttr(builder.getIntegerType(64, false), mapType),
+      builder.getAttr<mlir::omp::VariableCaptureKindAttr>(mapCaptureType),
+      builder.getStringAttr(name), builder.getBoolAttr(partialMap));
+
+  return op;
+}
+
+static int
+getComponentPlacementInParent(const Fortran::semantics::Symbol *componentSym) {
+  const auto *derived =
+      componentSym->owner()
+          .derivedTypeSpec()
+          ->typeSymbol()
+          .detailsIf<Fortran::semantics::DerivedTypeDetails>();
+  assert(derived &&
+         "expected derived type details when processing component symbol");
+  for (auto [placement, name] : llvm::enumerate(derived->componentNames()))
+    if (name == componentSym->name())
+      return placement;
+  return -1;
+}
+
+static std::optional<Object>
+getComponentObject(std::optional<Object> object,
+                   Fortran::semantics::SemanticsContext &semaCtx) {
+  if (!object)
+    return std::nullopt;
+
+  auto ref = evaluate::ExtractDataRef(*object.value().ref());
+  if (!ref)
+    return std::nullopt;
+
+  if (std::holds_alternative<evaluate::Component>(ref->u))
+    return object;
+
+  auto baseObj = getBaseObject(object.value(), semaCtx);
+  if (!baseObj)
+    return std::nullopt;
+
+  return getComponentObject(baseObj.value(), semaCtx);
+}
+
+static void
+generateMemberPlacementIndices(const Object &object,
+                               llvm::SmallVectorImpl<int> &indices,
+                               Fortran::semantics::SemanticsContext &semaCtx) {
+  auto compObj = getComponentObject(object, semaCtx);
+  while (compObj) {
+    indices.push_back(getComponentPlacementInParent(compObj->id()));
+    compObj =
+        getComponentObject(getBaseObject(compObj.value(), semaCtx), semaCtx);
+  }
+
+  indices = llvm::SmallVector<int>{llvm::reverse(indices)};
+}
+
+void addChildIndexAndMapToParent(
+    const omp::Object &object,
+    std::map<const Fortran::semantics::Symbol *,
+             llvm::SmallVector<OmpMapMemberIndicesData>> &parentMemberIndices,
+    mlir::omp::MapInfoOp &mapOp,
+    Fortran::semantics::SemanticsContext &semaCtx) {
+  std::optional<Fortran::evaluate::DataRef> dataRef =
+      ExtractDataRef(object.designator);
+  assert(dataRef.has_value() &&
+         "DataRef could not be extracted during mapping of derived type "
+         "cannot proceed");
+  const Fortran::semantics::Symbol *parentSym = &dataRef->GetFirstSymbol();
+  assert(parentSym && "Could not find parent symbol during lower of "
+                      "a component member in OpenMP map clause");
+  llvm::SmallVector<int> indices;
+  generateMemberPlacementIndices(object, indices, semaCtx);
+  parentMemberIndices[parentSym].push_back({indices, mapOp});
+}
+
+static void calculateShapeAndFillIndices(
+    llvm::SmallVectorImpl<int64_t> &shape,
+    llvm::SmallVectorImpl<OmpMapMemberIndicesData> &memberPlacementData) {
+  shape.push_back(memberPlacementData.size());
+  size_t largestIndicesSize =
+      std::max_element(memberPlacementData.begin(), memberPlacementData.end(),
+                       [](auto a, auto b) {
+                         return a.memberPlacementIndices.size() <
+                                b.memberPlacementIndices.size();
+                       })
+          ->memberPlacementIndices.size();
+  shape.push_back(largestIndicesSize);
+
+  // DenseElementsAttr expects a rectangular shape for the data, so all
+  // index lists have to be of the same length, this emplaces -1 as filler
----------------
skatrak wrote:

```suggestion
  // index lists have to be of the same length, this emplaces -1 as filler.
```

https://github.com/llvm/llvm-project/pull/82853