[flang-commits] [flang] [Flang] Add opt-in affine loop optimization pipeline (PR #191854)

[Tom Eccles via flang-commits]

================
@@ -316,146 +514,150 @@ AffineFunctionAnalysis::getChildIfAnalysis(fir::IfOp op) const {
   if (it == ifAnalysisMap.end()) {
     LLVM_DEBUG(llvm::dbgs() << "AffineFunctionAnalysis: not computed for:\n";
                op.dump(););
-    op.emitError("error in fetching if analysis in AffineFunctionAnalysis\n");
+
     return {};
   }
   return it->getSecond();
 }
 
-/// AffineMap rewriting fir.array_coor operation to affine apply,
-/// %dim = fir.gendim %lowerBound, %upperBound, %stride
-/// %a = fir.array_coor %arr(%dim) %i
-/// returning affineMap = affine_map<(i)[lb, ub, st] -> (i*st - lb)>
-static mlir::AffineMap createArrayIndexAffineMap(unsigned dimensions,
-                                                 MLIRContext *context) {
-  auto index = mlir::getAffineConstantExpr(0, context);
-  auto accuExtent = mlir::getAffineConstantExpr(1, context);
-  for (unsigned i = 0; i < dimensions; ++i) {
-    mlir::AffineExpr idx = mlir::getAffineDimExpr(i, context),
-                     lowerBound = mlir::getAffineSymbolExpr(i * 3, context),
-                     currentExtent =
-                         mlir::getAffineSymbolExpr(i * 3 + 1, context),
-                     stride = mlir::getAffineSymbolExpr(i * 3 + 2, context),
-                     currentPart = (idx * stride - lowerBound) * accuExtent;
-    index = currentPart + index;
-    accuExtent = accuExtent * currentExtent;
-  }
-  return mlir::AffineMap::get(dimensions, dimensions * 3, index);
-}
-
 static std::optional<int64_t> constantIntegerLike(const mlir::Value value) {
   if (auto definition = value.getDefiningOp<mlir::arith::ConstantOp>())
     if (auto stepAttr = mlir::dyn_cast<IntegerAttr>(definition.getValue()))
       return stepAttr.getInt();
   return {};
 }
 
-static mlir::Type coordinateArrayElement(fir::ArrayCoorOp op) {
-  if (auto refType =
-          mlir::dyn_cast_or_null<ReferenceType>(op.getMemref().getType())) {
-    if (auto seqType =
-            mlir::dyn_cast_or_null<SequenceType>(refType.getEleTy())) {
-      return seqType.getEleTy();
-    }
-  }
-  op.emitError(
-      "AffineLoopConversion: array type in coordinate operation not valid\n");
-  return mlir::Type();
-}
-
-static void populateIndexArgs(fir::ArrayCoorOp acoOp, fir::ShapeOp shape,
-                              SmallVectorImpl<mlir::Value> &indexArgs,
-                              mlir::PatternRewriter &rewriter) {
-  auto one = mlir::arith::ConstantOp::create(rewriter, acoOp.getLoc(),
-                                             rewriter.getIndexType(),
-                                             rewriter.getIndexAttr(1));
-  auto extents = shape.getExtents();
-  for (auto i = extents.begin(); i < extents.end(); i++) {
-    indexArgs.push_back(one);
-    indexArgs.push_back(*i);
-    indexArgs.push_back(one);
-  }
-}
+/// Holds the result of creating multi-dimensional affine operations.
+struct MultiDimAffineResult {
+  SmallVector<mlir::Value> indices;
+  fir::ConvertOp arrayConvert;
+};
 
-static void populateIndexArgs(fir::ArrayCoorOp acoOp, fir::ShapeShiftOp shape,
-                              SmallVectorImpl<mlir::Value> &indexArgs,
-                              mlir::PatternRewriter &rewriter) {
-  auto one = mlir::arith::ConstantOp::create(rewriter, acoOp.getLoc(),
-                                             rewriter.getIndexType(),
-                                             rewriter.getIndexAttr(1));
-  auto extents = shape.getPairs();
-  for (auto i = extents.begin(); i < extents.end();) {
-    indexArgs.push_back(*i++);
-    indexArgs.push_back(*i++);
-    indexArgs.push_back(one);
-  }
-}
+/// Creates multi-dimensional affine operations preserving array dimensionality.
+/// Instead of linearizing all indices into a single 1D offset, this extracts
+/// the array shape from the FIR SequenceType, creates a matching multi-dim
+/// MemRefType, and adjusts each per-dimension index from Fortran 1-based to
+/// memref 0-based indexing.
+static MultiDimAffineResult
+createMultiDimAffineOps(mlir::Value arrayRef, mlir::PatternRewriter &rewriter,
+                        fir::DoLoopOp outermost) {
+  auto acoOp = arrayRef.getDefiningOp<ArrayCoorOp>();
+  auto loc = acoOp.getLoc();
+  auto *context = acoOp.getContext();
 
-static void populateIndexArgs(fir::ArrayCoorOp acoOp, fir::SliceOp slice,
-                              SmallVectorImpl<mlir::Value> &indexArgs,
-                              mlir::PatternRewriter &rewriter) {
-  auto extents = slice.getTriples();
-  for (auto i = extents.begin(); i < extents.end();) {
-    indexArgs.push_back(*i++);
-    indexArgs.push_back(*i++);
-    indexArgs.push_back(*i++);
+  fir::SequenceType seqType;
+  if (auto refType =
+          mlir::dyn_cast<fir::ReferenceType>(acoOp.getMemref().getType()))
+    seqType = mlir::dyn_cast<fir::SequenceType>(refType.getEleTy());
+  else if (auto heapType =
+               mlir::dyn_cast<fir::HeapType>(acoOp.getMemref().getType()))
+    seqType = mlir::dyn_cast<fir::SequenceType>(heapType.getEleTy());
+
+  // need change because memref is row major order but fir.array is column major
+  // order]=
+  SmallVector<int64_t> reversedShape(seqType.getShape().rbegin(),
+                                     seqType.getShape().rend());
+
+  auto newType = mlir::MemRefType::get(reversedShape, seqType.getEleTy());
----------------
tblah wrote:

Is `seqType.getEleTy()` always a type that can be used with `mlir::MemRefType`? I'm worried about fir types which aren't mlir integer or floating point types.

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