[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);
----------------
tblah wrote:

I think you can look for BaseBoxType here

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