[flang-commits] [flang] [flang] Fixed LoopVersioning for array slices. (PR #65703)

[Mats Petersson via flang-commits]

================
@@ -224,58 +280,137 @@ void LoopVersioningPass::runOnOperation() {
     }
   }
 
-  if (argsOfInterest.empty())
+  if (argsOfInterest.empty()) {
+    LLVM_DEBUG(llvm::dbgs()
+               << "No suitable arguments.\n=== End " DEBUG_TYPE " ===\n");
     return;
+  }
 
-  struct OpsWithArgs {
-    mlir::Operation *op;
-    mlir::SmallVector<ArgInfo, 4> argsAndDims;
-  };
-  // Now see if those arguments are used inside any loop.
-  mlir::SmallVector<OpsWithArgs, 4> loopsOfInterest;
+  // A list of all loops in the function in post-order.
+  mlir::SmallVector<fir::DoLoopOp> originalLoops;
+  // Information about the arguments usage by the instructions
+  // immediately nested in a loop.
+  llvm::DenseMap<fir::DoLoopOp, ArgsUsageInLoop> argsInLoops;
 
+  // Traverse the loops in post-order and see
+  // if those arguments are used inside any loop.
   func.walk([&](fir::DoLoopOp loop) {
     mlir::Block &body = *loop.getBody();
-    mlir::SmallVector<ArgInfo, 4> argsInLoop;
+    auto &argsInLoop = argsInLoops[loop];
+    originalLoops.push_back(loop);
     body.walk([&](mlir::Operation *op) {
-      // support either fir.array_coor or fir.coordinate_of
-      if (auto arrayCoor = mlir::dyn_cast<fir::ArrayCoorOp>(op)) {
-        // no support currently for sliced arrays
-        if (arrayCoor.getSlice())
-          return;
-      } else if (!mlir::isa<fir::CoordinateOp>(op)) {
+      // Support either fir.array_coor or fir.coordinate_of.
+      if (!mlir::isa<fir::ArrayCoorOp, fir::CoordinateOp>(op))
         return;
-      }
-
-      // The current operation could be inside another loop than
-      // the one we're currently processing. Skip it, we'll get
-      // to it later.
+      // Process only operations immediately nested in the current loop.
       if (op->getParentOfType<fir::DoLoopOp>() != loop)
         return;
       mlir::Value operand = op->getOperand(0);
       for (auto a : argsOfInterest) {
         if (a.arg == normaliseVal(operand)) {
-          // use the reboxed value, not the block arg when re-creating the loop:
+          // Use the reboxed value, not the block arg when re-creating the loop.
+          // TODO: should we check that the operand dominates the loop?
----------------
Leporacanthicus wrote:

Thanks, I'll do some debugging. The assert basically means "we didn't find anything in the loop to update", which of course means versioning didn't achieve anything. 

We should identify this case and just say "nope, not doing that". 

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