[Flang] Fixes for XArrayCoorOp

[kiranchandramohan]

Summary

The upstreamed code was not incrementing the sliceOffset in multiples
of 3. This issue is fixed by using Offsets and incrementing by 3 during
every iteration.
In the conversion pattern, we were comparing the definingOp of an
operand with an FIR::UndefOp. Use LLVM::UndefOp for conversion.

[clementval]

This increment was upstreamed correctly but it was removed by https://reviews.llvm.org/D125967.

Otherwise LGTM.

[kiranchandramohan]

In D128017#3591201, @clementval wrote:

This increment was upstreamed correctly but it was removed by https://reviews.llvm.org/D125967.

Otherwise LGTM.

Ahh, I see. Now I understand why that code was in the original way. I guess it was to check with the fir::UndefOp.

flang/lib/Optimizer/CodeGen/CodeGen.cpp

2107

I have opened a discussion in https://discourse.llvm.org/t/order-of-conversion-of-operations/63259 for checking whether comparing with LLVM::UndefOp is the correct way. It seems it is not entirely safe to do that comparison.

May be we should do the comparison on the original Operands, and create step with the converted operands. But I am getting a crash with the getDefiningOp call in some cases with the original Operand.

[kiranchandramohan]

Use the originalOperands to check whether it is an fir::UndefOp
instead of using the newOperands and checking whether it is an
LLVM::UndefOp.

[Leporacanthicus]

LGTM.
Tested locally, works for me!

[peixin]

Ahh, I see. Now I understand why that code was in the original way.

+1. Sorry for the previous incorrect remove. I go through the conversion of fir::cg::XArrayCoorOp and have several nit comments. In addition, operands[0] is commonly used, should it be assigned to one variable coorBaseAddr/coorBase similar to the conversion of fir::CoordinateOp so to make it more readable? @kiranchandramohan Can you commit one NFC for these changes if they look OK to you?

flang/lib/Optimizer/CodeGen/CodeGen.cpp

2125

2145

2147

2193

2202

2206

2214

[kiranchandramohan]

In D128017#3593802, @peixin wrote:

Ahh, I see. Now I understand why that code was in the original way.

+1. Sorry for the previous incorrect remove. I go through the conversion of fir::cg::XArrayCoorOp and have several nit comments. In addition, operands[0] is commonly used, should it be assigned to one variable coorBaseAddr/coorBase similar to the conversion of fir::CoordinateOp so to make it more readable? @kiranchandramohan Can you commit one NFC for these changes if they look OK to you?

Thanks @peixin. I believe D129071 addresses some of these. I will make any remaining changes post completion of upstreaming.

flang/lib/Optimizer/CodeGen/CodeGen.cpp

doRewrite(fir::cg::XArrayCoorOp coor, mlir::Type ty, OpAdaptor adaptor,
     auto loc = coor.getLoc();
     mlir::ValueRange operands = adaptor.getOperands();
     unsigned rank = coor.getRank();
     assert(coor.indices().size() == rank);
     assert(coor.shape().empty() || coor.shape().size() == rank);
     assert(coor.shift().empty() || coor.shift().size() == rank);
     assert(coor.slice().empty() || coor.slice().size() == 3 * rank);
     mlir::Type idxTy = lowerTy().indexType();
+    unsigned indexOffset = coor.indicesOffset();
+    unsigned shapeOffset = coor.shapeOffset();
+    unsigned shiftOffset = coor.shiftOffset();
+    unsigned sliceOffset = coor.sliceOffset();
+    auto sliceOps = coor.slice().begin();
     mlir::Value one = genConstantIndex(loc, idxTy, rewriter, 1);
     mlir::Value prevExt = one;
     mlir::Value zero = genConstantIndex(loc, idxTy, rewriter, 0);
     mlir::Value offset = zero;
     const bool isShifted = !coor.shift().empty();
     const bool isSliced = !coor.slice().empty();
     const bool baseIsBoxed = coor.memref().getType().isa<fir::BoxType>();
 
     // For each dimension of the array, generate the offset calculation.
-    for (unsigned i = 0; i < rank; ++i) {
+    for (unsigned i = 0; i < rank; ++i, ++indexOffset, ++shapeOffset,
+                  ++shiftOffset, sliceOffset += 3, sliceOps += 3) {
       mlir::Value index =
-          integerCast(loc, rewriter, idxTy, operands[coor.indicesOffset() + i]);
-      mlir::Value lb = isShifted ? integerCast(loc, rewriter, idxTy,
-                                               operands[coor.shiftOffset() + i])
-                                 : one;
+          integerCast(loc, rewriter, idxTy, operands[indexOffset]);
+      mlir::Value lb =
+          isShifted ? integerCast(loc, rewriter, idxTy, operands[shiftOffset])
       mlir::Value step = one;
       bool normalSlice = isSliced;
       // Compute zero based index in dimension i of the element, applying
       // potential triplets and lower bounds.
       if (isSliced) {
-        mlir::Value ub = operands[coor.sliceOffset() + i + 1];
-        normalSlice = !mlir::isa_and_nonnull<fir::UndefOp>(ub.getDefiningOp());
+        mlir::Value originalUb = *(sliceOps + 1);
+        normalSlice =
+            !mlir::isa_and_nonnull<fir::UndefOp>(originalUb.getDefiningOp());
         if (normalSlice)
-          step = integerCast(loc, rewriter, idxTy,
-                             operands[coor.sliceOffset() + i + 2]);
+          step = integerCast(loc, rewriter, idxTy, operands[sliceOffset + 2]);
       }
       auto idx = rewriter.create<mlir::LLVM::SubOp>(loc, idxTy, index, lb);
       mlir::Value diff =
           rewriter.create<mlir::LLVM::MulOp>(loc, idxTy, idx, step);
       if (normalSlice) {
         mlir::Value sliceLb =
-            integerCast(loc, rewriter, idxTy, operands[coor.sliceOffset() + i]);
+            integerCast(loc, rewriter, idxTy, operands[sliceOffset]);
         auto adj = rewriter.create<mlir::LLVM::SubOp>(loc, idxTy, sliceLb, lb);
         diff = rewriter.create<mlir::LLVM::AddOp>(loc, idxTy, diff, adj);
       }
       // Update the offset given the stride and the zero based index `diff`
       // that was just computed.
       if (baseIsBoxed) {
         // Use stride in bytes from the descriptor.
         mlir::Value stride =
             loadStrideFromBox(loc, adaptor.getOperands()[0], i, rewriter);
         auto sc = rewriter.create<mlir::LLVM::MulOp>(loc, idxTy, diff, stride);
         offset = rewriter.create<mlir::LLVM::AddOp>(loc, idxTy, sc, offset);
       } else {
         // Use stride computed at last iteration.
         auto sc = rewriter.create<mlir::LLVM::MulOp>(loc, idxTy, diff, prevExt);
         offset = rewriter.create<mlir::LLVM::AddOp>(loc, idxTy, sc, offset);
         // Compute next stride assuming contiguity of the base array
         // (in element number).
-        auto nextExt =
-            integerCast(loc, rewriter, idxTy, operands[coor.shapeOffset() + i]);
+        auto nextExt = integerCast(loc, rewriter, idxTy, operands[shapeOffset]);
         prevExt =
             rewriter.create<mlir::LLVM::MulOp>(loc, idxTy, prevExt, nextExt);
       }
     }
 
     // Add computed offset to the base address.
     if (baseIsBoxed) {
       // Working with byte offsets. The base address is read from the fir.box.
       // and need to be casted to i8* to do the pointer arithmetic.
       mlir::Type baseTy =
           getBaseAddrTypeFromBox(adaptor.getOperands()[0].getType());
       mlir::Value base =
           loadBaseAddrFromBox(loc, baseTy, adaptor.getOperands()[0], rewriter);
       mlir::Type voidPtrTy = getVoidPtrType();
       base = rewriter.create<mlir::LLVM::BitcastOp>(loc, voidPtrTy, base);
       llvm::SmallVector<mlir::Value> args{offset};
       auto addr =
           rewriter.create<mlir::LLVM::GEPOp>(loc, voidPtrTy, base, args);
       if (coor.subcomponent().empty()) {
         rewriter.replaceOpWithNewOp<mlir::LLVM::BitcastOp>(coor, baseTy, addr);
         return mlir::success();
if (coor.subcomponent().empty()) {
           assert(coor.lenParams().size() == 1);
           auto bitsInChar = lowerTy().getKindMap().getCharacterBitsize(
               eleTy.cast<fir::CharacterType>().getFKind());
           auto scaling = genConstantIndex(loc, idxTy, rewriter, bitsInChar / 8);
           auto scaledBySize =
               rewriter.create<mlir::LLVM::MulOp>(loc, idxTy, offset, scaling);
           auto length =
               integerCast(loc, rewriter, idxTy,
                           adaptor.getOperands()[coor.lenParamsOffset()]);
           offset = rewriter.create<mlir::LLVM::MulOp>(loc, idxTy, scaledBySize,
                                                       length);
         } else {
           TODO(loc, "compute size of derived type with type parameters");
         }
       }
       // Cast the base address to a pointer to T.
       base = rewriter.create<mlir::LLVM::BitcastOp>(loc, ty,
                                                     adaptor.getOperands()[0]);
     } else {
       // Operand #0 must have a pointer type. For subcomponent slicing, we
       // want to cast away the array type and have a plain struct type.
       mlir::Type ty0 = adaptor.getOperands()[0].getType();
       auto ptrTy = ty0.dyn_cast<mlir::LLVM::LLVMPointerType>();
       assert(ptrTy && "expected pointer type");
       mlir::Type eleTy = ptrTy.getElementType();
       while (auto arrTy = eleTy.dyn_cast<mlir::LLVM::LLVMArrayType>())
         eleTy = arrTy.getElementType();
       auto newTy = mlir::LLVM::LLVMPointerType::get(eleTy);
       base = rewriter.create<mlir::LLVM::BitcastOp>(loc, newTy,
                                                     adaptor.getOperands()[0]);
     }
     llvm::SmallVector<mlir::Value> args = {offset};
     for (auto i = coor.subcomponentOffset(); i != coor.indicesOffset(); ++i)
       args.push_back(operands[i]);
     rewriter.replaceOpWithNewOp<mlir::LLVM::GEPOp>(coor, ty, base, args);
     return mlir::success();
   }
 };

flang/test/Fir/convert-to-llvm.fir

 // CHECK:           %[[VAL_8:.*]] = llvm.sub %[[VAL_5]], %[[VAL_6]]  : i64
 // CHECK:           %[[VAL_9:.*]] = llvm.mul %[[VAL_8]], %[[VAL_4]]  : i64
 // CHECK:           %[[VAL_10:.*]] = llvm.sub %[[VAL_2]], %[[VAL_6]]  : i64
 // CHECK:           %[[VAL_11:.*]] = llvm.add %[[VAL_9]], %[[VAL_10]]  : i64
 // CHECK:           %[[VAL_12:.*]] = llvm.mul %[[VAL_11]], %[[VAL_6]]  : i64
 // CHECK:           %[[VAL_13:.*]] = llvm.add %[[VAL_12]], %[[VAL_7]]  : i64
 // CHECK:           %[[VAL_14:.*]] = llvm.mul %[[VAL_6]], %[[VAL_1]]  : i64
 // CHECK:           %[[VAL_15:.*]] = llvm.sub %[[VAL_5]], %[[VAL_6]]  : i64
-// CHECK:           %[[VAL_16:.*]] = llvm.mul %[[VAL_15]], %[[VAL_2]]  : i64
-// CHECK:           %[[VAL_17:.*]] = llvm.sub %[[VAL_3]], %[[VAL_6]]  : i64
+// CHECK:           %[[VAL_16:.*]] = llvm.mul %[[VAL_15]], %[[VAL_4]]  : i64
+// CHECK:           %[[VAL_17:.*]] = llvm.sub %[[VAL_2]], %[[VAL_6]]  : i64
 // CHECK:           %[[VAL_18:.*]] = llvm.add %[[VAL_16]], %[[VAL_17]]  : i64
 // CHECK:           %[[VAL_19:.*]] = llvm.mul %[[VAL_18]], %[[VAL_14]]  : i64
 // CHECK:           %[[VAL_20:.*]] = llvm.add %[[VAL_19]], %[[VAL_13]]  : i64
 // CHECK:           %[[VAL_21:.*]] = llvm.mul %[[VAL_14]], %[[VAL_1]]  : i64
 // CHECK:           %[[VAL_22:.*]] = llvm.sub %[[VAL_5]], %[[VAL_6]]  : i64
-// CHECK:           %[[VAL_23:.*]] = llvm.mul %[[VAL_22]], %[[VAL_3]]  : i64
-// CHECK:           %[[VAL_24:.*]] = llvm.sub %[[VAL_4]], %[[VAL_6]]  : i64
+// CHECK:           %[[VAL_23:.*]] = llvm.mul %[[VAL_22]], %[[VAL_4]]  : i64
+// CHECK:           %[[VAL_24:.*]] = llvm.sub %[[VAL_2]], %[[VAL_6]]  : i64
 // CHECK:           %[[VAL_25:.*]] = llvm.add %[[VAL_23]], %[[VAL_24]]  : i64
 // CHECK:           %[[VAL_26:.*]] = llvm.mul %[[VAL_25]], %[[VAL_21]]  : i64
 // CHECK:           %[[VAL_27:.*]] = llvm.add %[[VAL_26]], %[[VAL_20]]  : i64
 // CHECK:           %[[VAL_28:.*]] = llvm.mul %[[VAL_21]], %[[VAL_1]]  : i64
 // CHECK:           %[[VAL_29:.*]] = llvm.bitcast %[[VAL_0]] : !llvm.ptr<i32> to !llvm.ptr<i32>
 // CHECK:           %[[VAL_30:.*]] = llvm.getelementptr %[[VAL_29]][%[[VAL_27]]] : (!llvm.ptr<i32>, i64) -> !llvm.ptr<i32>
 // CHECK:           llvm.return
 // CHECK:         }