[llvm-branch-commits] [flang] [mlir] [Flang][mlir][OpenMP] Support affinity clause codegen in Flang (PR #182222)

[Tom Eccles via llvm-branch-commits]

================
@@ -917,6 +917,192 @@ void collectLoopRelatedInfo(
   convertLoopBounds(converter, currentLocation, result, loopVarTypeSize);
 }
 
+mlir::Value genAffinityAddr(Fortran::lower::AbstractConverter &converter,
+                            const omp::Object &object,
+                            Fortran::lower::StatementContext &stmtCtx,
+                            mlir::Location loc) {
+  // Get address from expression if it exists: affinity(a(3)), affinity(a(1:10))
+  if (auto expr = object.ref()) {
+    fir::ExtendedValue exv =
+        converter.genExprAddr(toEvExpr(*expr), stmtCtx, &loc);
+    return fir::getBase(exv);
+  }
+
+  // Fallback to base symbol address: affinity(a)
+  const Fortran::semantics::Symbol *sym = object.sym();
+  assert(sym && "expected symbol in affinity object");
+  mlir::Value addr = converter.getSymbolAddress(*sym);
+
+  if (mlir::isa<fir::BoxType>(addr.getType())) {
+    addr = fir::BoxAddrOp::create(converter.getFirOpBuilder(), loc, addr);
+  }
+  return addr;
+}
+
+static mlir::Value buildNumElemsFromMapBound(fir::FirOpBuilder &builder,
+                                             mlir::Location loc,
+                                             mlir::omp::MapBoundsOp mb) {
+  mlir::Value lb = mb.getLowerBound();
+  mlir::Value ub = mb.getUpperBound();
+  mlir::Value stride = mb.getStride();
+
+  // ((ub - lb) / stride) + 1
+  mlir::Value diff = mlir::arith::SubIOp::create(builder, loc, ub, lb);
+  mlir::Value div = mlir::arith::DivUIOp::create(builder, loc, diff, stride);
+  mlir::Value one =
+      builder.createIntegerConstant(loc, builder.getIndexType(), 1);
+  mlir::Value result = mlir::arith::AddIOp::create(builder, loc, div, one);
+
+  return mlir::arith::IndexCastOp::create(builder, loc, builder.getI64Type(),
+                                          result);
+}
+
+mlir::Value genAffinityLen(fir::FirOpBuilder &builder, mlir::Location loc,
+                           const mlir::DataLayout &dl, mlir::Value addr,
+                           llvm::ArrayRef<mlir::Value> bounds, bool hasRef) {
+  auto isDescriptorLike = [](mlir::Type t) -> bool {
+    t = fir::unwrapPassByRefType(t);
+    return mlir::isa<fir::BoxType, fir::ClassType>(t);
+  };
+
+  auto getElementBytesOrZero = [&](mlir::Type baseTy) -> int64_t {
+    if (isDescriptorLike(baseTy))
+      return 0;
+    mlir::Type eleTy = fir::unwrapPassByRefType(baseTy);
+    eleTy = fir::unwrapSequenceType(eleTy);
+    return static_cast<int64_t>(dl.getTypeSize(eleTy));
+  };
+
+  auto getWholeObjectBytesIfStaticOrZero = [&](mlir::Type addrTy) -> int64_t {
+    if (isDescriptorLike(addrTy))
+      return 0;
+
+    mlir::Type eleTy = fir::unwrapPassByRefType(addrTy);
+
+    // Scalar
+    if (!mlir::isa<fir::SequenceType>(eleTy))
+      return static_cast<int64_t>(dl.getTypeSize(eleTy));
+
+    // Array with static extents
----------------
tblah wrote:

We don't always do a good job of maintaining accurate information for array shapes in the array type itself so sometimes an array that should have a statically known size will still have a type containing unknown dimensions. This could be improved but it is how things are now.

The most reliable way to get information about type sizes is to look at the (hl)fir.declare. These (mostly) nop operations collect all of the extended information for a fortran value which doesn't fit neatly into a single mlir::Value. So if auxilery information are needed (e.g. character length, array shape or shapeshift (for non-default lower bounds)) they will all be available through the (hl)fir.declare that this value comes from. You should be able to get much more of the information you need in a more reliable way using the helpers written for hlfir FortranVariableInterface and hlfir entity.

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