[flang][LoopVersioning] support reboxed operands

[tblah]

Summary

Since https://reviews.llvm.org/D158119, many boxes lowered via HLFIR are
reboxed with better lower bounds information after they are declared.

For the loop versioning pass to support FIR lowered via HLFIR, it needs
to dereference fir.rebox operations to figure out that the variable was
a function argument.

I decided to modify the existing dereferencing of fir.declare so that
the declared/reboxed value is used in the versioned loop instead of the
function argument. This makes it easier for the improved lower bounds
information to be accessed. In doing this, I changed ArgInfo to store
ArgInfo::arg by value instead of by pointer because mlir::Value has
value-type semantics.

[vzakhari]

LGTM. Thank you!

flang/lib/Optimizer/Transforms/LoopVersioning.cpp

 static mlir::Value unwrapFirDeclare(mlir::Value val) {
   // fir.declare is for source code variables. We don't have declares of
   // declares
   if (fir::DeclareOp declare = val.getDefiningOp<fir::DeclareOp>())
     return declare.getMemref();
   return val;
 }
 
+/// if a value comes from a fir.rebox, follow the rebox to the original source,
+/// of the value, otherwise return the value
+static mlir::Value unwrapReboxOp(mlir::Value val) {
+  // don't support reboxes of reboxes
+  if (fir::ReboxOp rebox = val.getDefiningOp<fir::ReboxOp>())
+    val = rebox.getBox();
+  return val;
+}
+
+/// normalize a value (removing fir.declare and fir.rebox) so that we can
+/// more conveniently spot values which came from function arguments
+static mlir::Value normaliseVal(mlir::Value val) {
+  return unwrapFirDeclare(unwrapReboxOp(val));
+}
+
 void LoopVersioningPass::runOnOperation() {
   LLVM_DEBUG(llvm::dbgs() << "=== Begin " DEBUG_TYPE " ===\n");
   mlir::func::FuncOp func = getOperation();
 
   /// @c ArgInfo
   /// A structure to hold an argument, the size of the argument and dimension
   /// information.
   struct ArgInfo {
-    mlir::Value *arg;
+    mlir::Value arg;
     size_t size;
     unsigned rank;
     fir::BoxDimsOp dims[CFI_MAX_RANK];
   };
 
   // First look for arguments with assumed shape = unknown extent in the lowest
   // dimension.
   LLVM_DEBUG(llvm::dbgs() << "Func-name:" << func.getSymName() << "\n");
if (auto seqTy = getAsSequenceType(&arg)) {
         size_t typeSize = 0;
         mlir::Type elementType = fir::unwrapSeqOrBoxedSeqType(arg.getType());
         if (elementType.isa<mlir::FloatType>() ||
             elementType.isa<mlir::IntegerType>())
           typeSize = elementType.getIntOrFloatBitWidth() / 8;
         else if (auto cty = elementType.dyn_cast<fir::ComplexType>())
           typeSize = 2 * cty.getEleType(kindMap).getIntOrFloatBitWidth() / 8;
         if (typeSize)
-          argsOfInterest.push_back({&arg, typeSize, rank, {}});
+          argsOfInterest.push_back({arg, typeSize, rank, {}});
         else
           LLVM_DEBUG(llvm::dbgs() << "Type not supported\n");
       }
     }
   }
 
   if (argsOfInterest.empty())
     return;
func.walk([&](fir::DoLoopOp loop) {
     body.walk([&](fir::CoordinateOp op) {
       // The current operation could be inside another loop than
       // the one we're currently processing. Skip it, we'll get
       // to it later.
       if (op->getParentOfType<fir::DoLoopOp>() != loop)
         return;
       mlir::Value operand = op->getOperand(0);
       for (auto a : argsOfInterest) {
-        if (*a.arg == unwrapFirDeclare(operand)) {
+        if (a.arg == normaliseVal(operand)) {
+          // use the reboxed value, not the block arg when re-creating the loop:
+          a.arg = operand;
           // Only add if it's not already in the list.
           if (std::find_if(argsInLoop.begin(), argsInLoop.end(), [&](auto it) {
                 return it.arg == a.arg;
               }) == argsInLoop.end()) {
 
             argsInLoop.push_back(a);
             break;
           }
for (auto op : loopsOfInterest) {
     for (auto &arg : op.argsAndDims) {
       // Fetch all the dimensions of the array, except the last dimension.
       // Always fetch the first dimension, however, so set ndims = 1 if
       // we have one dim
       unsigned ndims = arg.rank;
       for (unsigned i = 0; i < ndims; i++) {
         mlir::Value dimIdx = builder.createIntegerConstant(loc, idxTy, i);
         arg.dims[i] = builder.create<fir::BoxDimsOp>(loc, idxTy, idxTy, idxTy,
-                                                     *arg.arg, dimIdx);
+                                                     arg.arg, dimIdx);
       }
       // We only care about lowest order dimension, here.
       mlir::Value elemSize =
           builder.createIntegerConstant(loc, idxTy, arg.size);
       mlir::Value cmp = builder.create<mlir::arith::CmpIOp>(
           loc, mlir::arith::CmpIPredicate::eq, arg.dims[0].getResult(2),
           elemSize);
       if (!allCompares) {
for (auto op : loopsOfInterest) {
     builder.setInsertionPointToStart(&ifOp.getThenRegion().front());
 
     LLVM_DEBUG(llvm::dbgs() << "Creating cloned loop\n");
     mlir::Operation *clonedLoop = op.op->clone();
     bool changed = false;
     for (auto &arg : op.argsAndDims) {
       fir::SequenceType::Shape newShape;
       newShape.push_back(fir::SequenceType::getUnknownExtent());
-      auto elementType = fir::unwrapSeqOrBoxedSeqType(arg.arg->getType());
+      auto elementType = fir::unwrapSeqOrBoxedSeqType(arg.arg.getType());
       mlir::Type arrTy = fir::SequenceType::get(newShape, elementType);
       mlir::Type boxArrTy = fir::BoxType::get(arrTy);
       mlir::Type refArrTy = builder.getRefType(arrTy);
-      auto carg = builder.create<fir::ConvertOp>(loc, boxArrTy, *arg.arg);
+      auto carg = builder.create<fir::ConvertOp>(loc, boxArrTy, arg.arg);
       auto caddr = builder.create<fir::BoxAddrOp>(loc, refArrTy, carg);
       auto insPt = builder.saveInsertionPoint();
       // Use caddr instead of arg.
       clonedLoop->walk([&](fir::CoordinateOp coop) {
         // Reduce the multi-dimensioned index to a single index.
         // This is required becase fir arrays do not support multiple dimensions
         // with unknown dimensions at compile time.
         // We then calculate the multidimensional array like this:
         // arr(x, y, z) bedcomes arr(z * stride(2) + y * stride(1) + x)
         // where stride is the distance between elements in the dimensions
         // 0, 1 and 2 or x, y and z.
-        if (unwrapFirDeclare(coop->getOperand(0)) == *arg.arg &&
-            coop->getOperands().size() >= 2) {
+        if (coop->getOperand(0) == arg.arg && coop->getOperands().size() >= 2) {
           builder.setInsertionPoint(coop);
           mlir::Value totalIndex;
           for (unsigned i = arg.rank - 1; i > 0; i--) {
             // Operand(1) = array; Operand(2) = index1; Operand(3) = index2
             mlir::Value curIndex =
                 builder.createConvert(loc, idxTy, coop->getOperand(i + 1));
             // Multiply by the stride of this array. Later we'll divide by the
             // element size.

flang/test/Transforms/loop-versioning.fir

 // RUN: fir-opt --loop-versioning %s | FileCheck %s
 
 
 //  subroutine sum1d(a, n)
 //    real*8 :: a(:)
 //    integer :: n
 //    real*8 :: sum
 //    integer :: i
 //    sum = 0
 //    do i=1,n
 //       sum = sum + a(i)
 //    end do
 //  end subroutine sum1d
 module {
   func.func @sum1d(%arg0: !fir.box<!fir.array<?xf64>> {fir.bindc_name = "a"}, %arg1: !fir.ref<i32> {fir.bindc_name = "n"}) {
     %decl = fir.declare %arg0 {uniq_name = "a"} : (!fir.box<!fir.array<?xf64>>) -> !fir.box<!fir.array<?xf64>>
+    %rebox = fir.rebox %decl : (!fir.box<!fir.array<?xf64>>) -> !fir.box<!fir.array<?xf64>>
     %0 = fir.alloca i32 {bindc_name = "i", uniq_name = "_QMmoduleFsum1dEi"}
     %1 = fir.alloca f64 {bindc_name = "sum", uniq_name = "_QMmoduleFsum1dEsum"}
     %cst = arith.constant 0.000000e+00 : f64
     fir.store %cst to %1 : !fir.ref<f64>
     %c1_i32 = arith.constant 1 : i32
     %2 = fir.convert %c1_i32 : (i32) -> index
     %3 = fir.load %arg1 : !fir.ref<i32>
     %4 = fir.convert %3 : (i32) -> index
     %c1 = arith.constant 1 : index
     %5 = fir.convert %2 : (index) -> i32
     %6:2 = fir.do_loop %arg2 = %2 to %4 step %c1 iter_args(%arg3 = %5) -> (index, i32) {
       fir.store %arg3 to %0 : !fir.ref<i32>
       %7 = fir.load %1 : !fir.ref<f64>
       %8 = fir.load %0 : !fir.ref<i32>
       %9 = fir.convert %8 : (i32) -> i64
       %c1_i64 = arith.constant 1 : i64
       %10 = arith.subi %9, %c1_i64 : i64
-      %11 = fir.coordinate_of %decl, %10 : (!fir.box<!fir.array<?xf64>>, i64) -> !fir.ref<f64>
+      %11 = fir.coordinate_of %rebox, %10 : (!fir.box<!fir.array<?xf64>>, i64) -> !fir.ref<f64>
       %12 = fir.load %11 : !fir.ref<f64>
       %13 = arith.addf %7, %12 fastmath<contract> : f64
       fir.store %13 to %1 : !fir.ref<f64>
       %14 = arith.addi %arg2, %c1 : index
       %15 = fir.convert %c1 : (index) -> i32
       %16 = fir.load %0 : !fir.ref<i32>
       %17 = arith.addi %16, %15 : i32
       fir.result %14, %17 : index, i32
     }
     fir.store %6#1 to %0 : !fir.ref<i32>
     return
   }
 
 // Note this only checks the expected transformation, not the entire generated code:
 // CHECK-LABEL: func.func @sum1d(
 // CHECK-SAME:                  %[[ARG0:.*]]: !fir.box<!fir.array<?xf64>> {{.*}})
 // CHECK: %[[DECL:.*]] = fir.declare %arg0 {uniq_name = "a"} : (!fir.box<!fir.array<?xf64>>) -> !fir.box<!fir.array<?xf64>>
+// CHECK: %[[REBOX:.*]] = fir.rebox %[[DECL]]
 // CHECK: %[[ZERO:.*]] = arith.constant 0 : index
-// CHECK: %[[DIMS:.*]]:3 = fir.box_dims %[[ARG0]], %[[ZERO]] : {{.*}}
+// CHECK: %[[DIMS:.*]]:3 = fir.box_dims %[[REBOX]], %[[ZERO]] : {{.*}}
 // CHECK: %[[SIZE:.*]] = arith.constant 8 : index
 // CHECK: %[[CMP:.*]] = arith.cmpi eq, %[[DIMS]]#2, %[[SIZE]]
 // CHECK: %[[IF_RES:.*]]:2 = fir.if %[[CMP]] -> {{.*}}
-// CHECK: %[[NEWARR:.*]] = fir.convert %[[ARG0]]
+// CHECK: %[[NEWARR:.*]] = fir.convert %[[REBOX]]
 // CHECK: %[[BOXADDR:.*]] = fir.box_addr %[[NEWARR]] : {{.*}} -> !fir.ref<!fir.array<?xf64>>
 // CHECK: %[[LOOP_RES:.*]]:2 = fir.do_loop {{.*}}
 // CHECK: %[[COORD:.*]] = fir.coordinate_of %[[BOXADDR]], %{{.*}} : (!fir.ref<!fir.array<?xf64>>, index) -> !fir.ref<f64>
 // CHECK: %{{.*}} = fir.load %[[COORD]] : !fir.ref<f64>
 // CHECK: fir.result %{{.*}}, %{{.*}}
 // CHECK: }
 // CHECK  fir.result %[[LOOP_RES]]#0, %[[LOOP_RES]]#1
 // CHECK: } else {
 // CHECK: %[[LOOP_RES2:.*]]:2 = fir.do_loop {{.*}}
-// CHECK: %[[COORD2:.*]] = fir.coordinate_of %[[DECL]], %{{.*}} : (!fir.box<!fir.array<?xf64>>, i64) -> !fir.ref<f64>
+// CHECK: %[[COORD2:.*]] = fir.coordinate_of %[[REBOX]], %{{.*}} : (!fir.box<!fir.array<?xf64>>, i64) -> !fir.ref<f64>
 // CHECK: %{{.*}}= fir.load %[[COORD2]] : !fir.ref<f64>
 // CHECK: fir.result %{{.*}}, %{{.*}}
 // CHECK: }
 // CHECK  fir.result %[[LOOP_RES2]]#0, %[[LOOP_RES2]]#1
 // CHECK: }
 // CHECK: fir.store %[[IF_RES]]#1 to %{{.*}}
 // CHECK: return