[fir] Add fir.insert_on_range conversion

[clementval]

Summary

Convert fir.insert_on_range operation to corresponding
llvm.insertvalue operations.

This patch is part of the upstreaming effort from fir-dev branch.

[mehdi_amini]

Looks good in general, just some minor comments inline.

flang/lib/Optimizer/CodeGen/CodeGen.cpp

33

Nit: seems like it could be using llvm::all_of:

return llvm::all_of(operands, [] (Value opnd) { 
  auto *defop = opnd.getDefiningOp();
  return defop && isa<mlir::LLVM::ConstantOp, mlir::arith::ConstantOp>(defop);
});

Should this use isConstantLike? https://github.com/llvm/llvm-project/blob/main/mlir/include/mlir/IR/Matchers.h#L53-L56

return llvm::all_of(operands, [] (Value opnd) { 
  auto *defop = opnd.getDefiningOp();
  return defop && isConstantLike(defop);
});

80

If you make it a pure virtual = 0 it'll be a compile time failure.
(Could also use CRTP here)

210

The usual pattern would look like: https://github.com/llvm/llvm-project/blob/main/mlir/lib/Dialect/Linalg/Transforms/ElementwiseOpFusion.cpp#L1182-L1204

But also this function isn't used here, so I rather remove this until it's needed.

215

Same as above, unused.

237

ditto

269

Actually the entire class seems unused?

285

It is worth documenting that if the subscript is full of zeros then it means we reached the end of the array.
This "wrap around" could also be signaled by a returned LogicalResult (or a bool maybe).

291

May be worth a message && "..." in the assert on the broken invariant.

Also the InsertOnRangeOp op is defined with only two operands, how we can get more?

311

Nit: std:: is spurious here (you don't use it line 295)
(same below)

318

Please fuse this loop with the previous one and eliminate the temporary storage lowerBound / upperBound.
Unless I'm mistaken this is just an extra indirection that doesn't seem useful here.
(You can iterate on the reverse iterators)

320

auto really doesn't help here.
Which reminds me: did you follow up on https://reviews.llvm.org/D111288#inline-1072217 ?

Also this seems like another spurious copy here?
If this is just for the sake of the naming, a lighter way to do this may be auto &subscripts = lBounds;

330

rewriter.getI64Type isn't necessarily "cheap", it is quite common to get it into a variable outside of the enclosing loop and reuse it.

331

Do you need this? (In general SmallVector implicitly convert to ArrayRef)

[mehdi_amini]

flang/lib/Optimizer/CodeGen/CodeGen.cpp

72

Seem like this line is the entire justification for this class? Is this enough for it be pulling its weight?

Previously added pattern (GlobalOpConversion) started with auto tyAttr = convertType(global.getType()); ; should it inherit from this class as well or do I miss something?

[clementval]

Address review comments

flang/lib/Optimizer/CodeGen/CodeGen.cpp

33

We could use the proposed code. Just removed it now since it is not used currently after removing the assert.

72

Yeah GlobalOpConversion might inherit from this. I'll explore that and open a separate revision if needed. I would prefer to differ any refactoring (keeping or removing this class) after we have the full picture on this pass.

210

No needed now so will bring it back later.

215

Removed for now

237

Removed for now

269

Removed for now

285

When subscripts is equal to ubounds we reach the end of the array and end of the insertion.

291

Removed it since this assert does not really makes sense anymore. Looks like it was useful previously when we had more operands.

320

Somehow I missed your last comment on D111288. I'll check with the others what we will do. It seems a bit overkill to go through a RFC for this.

[mehdi_amini]

flang/lib/Optimizer/CodeGen/CodeGen.cpp

285

Right, but I think you're describing a behavior of the caller, while I'm referring to the contract of this function alone.
If we don't expect to wrap-around, then we should just add an assert (and document it).
(I shouldn't have to analyze call-site(s) to figure out the behavior of a function in isolation).

[clementval]

flang/lib/Optimizer/CodeGen/CodeGen.cpp

285

Right my bad. That's true the function can have some more description without looking how it is used. I'll update that.

[clementval]

Add comment about the wrap around in incrementSubscripts.

[clementval]

Replace auto for subscripts ref

[mehdi_amini]

flang/lib/Optimizer/CodeGen/CodeGen.cpp

226

How do we guarantee that we have a static array type? This will crash with an IR with dynamic shape.

Right now the InsertOnRangeOp is defined as taking a fir_SequenceType which isn't restricted to static shapes.

[clementval]

flang/lib/Optimizer/CodeGen/CodeGen.cpp

226

This is actually smth that was missing in the verifier. D113092 adds the check directly in the verifier.

[clementval]

Thanks for the review @mehdi_amini. I'll wait until the verifier is updated to land this.

[clementval]

Rebase

[clementval]

Fix rebase

flang/lib/Optimizer/CodeGen/CodeGen.cpp

 #include "mlir/Pass/Pass.h"
 #include "llvm/ADT/ArrayRef.h"
 
 #define DEBUG_TYPE "flang-codegen"
 
 // fir::LLVMTypeConverter for converting to LLVM IR dialect types.
 #include "TypeConverter.h"
 
+namespace fir {
+/// return true if all `Value`s in `operands` are `ConstantOp`s
+bool allConstants(mlir::ValueRange operands) {
+  for (auto opnd : operands) {
+    if (auto *defop = opnd.getDefiningOp())
+      if (isa<mlir::LLVM::ConstantOp>(defop) ||
+          isa<mlir::arith::ConstantOp>(defop))
+        continue;
+    return false;
+  }
+  return true;
+}
+} // namespace fir
+
 namespace {
 /// FIR conversion pattern template
 template <typename FromOp>
 class FIROpConversion : public mlir::ConvertOpToLLVMPattern<FromOp> {
 public:
   explicit FIROpConversion(fir::LLVMTypeConverter &lowering)
       : mlir::ConvertOpToLLVMPattern<FromOp>(lowering) {}
 
 protected:
   mlir::Type convertType(mlir::Type ty) const {
     return lowerTy().convertType(ty);
   }
 
   fir::LLVMTypeConverter &lowerTy() const {
     return *static_cast<fir::LLVMTypeConverter *>(this->getTypeConverter());
   }
 };
+
+/// FIR conversion pattern template
+template <typename FromOp>
+class FIROpAndTypeConversion : public FIROpConversion<FromOp> {
+public:
+  using FIROpConversion<FromOp>::FIROpConversion;
+  using OpAdaptor = typename FromOp::Adaptor;
+
+  mlir::LogicalResult
+  matchAndRewrite(FromOp op, OpAdaptor adaptor,
+                  mlir::ConversionPatternRewriter &rewriter) const final {
+    mlir::Type ty = this->convertType(op.getType());
+    return doRewrite(op, ty, adaptor, rewriter);
+  }
+
+  virtual mlir::LogicalResult
+  doRewrite(FromOp addr, mlir::Type ty, OpAdaptor adaptor,
+            mlir::ConversionPatternRewriter &rewriter) const {
+    llvm_unreachable("derived class must override");
+  }
+};
 } // namespace
 
 namespace {
 struct AddrOfOpConversion : public FIROpConversion<fir::AddrOfOp> {
   using FIROpConversion::FIROpConversion;
 
   mlir::LogicalResult
   matchAndRewrite(fir::AddrOfOp addr, OpAdaptor adaptor,
struct UndefOpConversion : public FIROpConversion<fir::UndefOp> {
   mlir::LogicalResult
   matchAndRewrite(fir::UndefOp undef, OpAdaptor,
                   mlir::ConversionPatternRewriter &rewriter) const override {
     rewriter.replaceOpWithNewOp<mlir::LLVM::UndefOp>(
         undef, convertType(undef.getType()));
     return success();
   }
 };
+
+// Code shared between insert_value and extract_value Ops.
+struct ValueOpCommon {
+  static mlir::Attribute getValue(mlir::Value value) {
+    auto *defOp = value.getDefiningOp();
+    if (auto v = dyn_cast<mlir::LLVM::ConstantOp>(defOp))
+      return v.value();
+    if (auto v = dyn_cast<mlir::arith::ConstantOp>(defOp))
+      return v.value();
+    llvm_unreachable("must be a constant op");
+    return {};
+  }
+
+  // Translate the arguments pertaining to any multidimensional array to
+  // row-major order for LLVM-IR.
+  static void toRowMajor(SmallVectorImpl<mlir::Attribute> &attrs,
+                         mlir::Type ty) {
+    assert(ty && "type is null");
+    const auto end = attrs.size();
+    for (std::remove_const_t<decltype(end)> i = 0; i < end; ++i) {
+      if (auto seq = ty.dyn_cast<mlir::LLVM::LLVMArrayType>()) {
+        const auto dim = getDimension(seq);
+        if (dim > 1) {
+          auto ub = std::min(i + dim, end);
+          std::reverse(attrs.begin() + i, attrs.begin() + ub);
+          i += dim - 1;
+        }
+        ty = getArrayElementType(seq);
+      } else if (auto st = ty.dyn_cast<mlir::LLVM::LLVMStructType>()) {
+        ty = st.getBody()[attrs[i].cast<mlir::IntegerAttr>().getInt()];
+      } else {
+        llvm_unreachable("index into invalid type");
+      }
+    }
+  }
+
+  static llvm::SmallVector<mlir::Attribute>
+  collectIndices(mlir::ConversionPatternRewriter &rewriter,
+                 mlir::ArrayAttr arrAttr) {
+    llvm::SmallVector<mlir::Attribute> attrs;
+    for (auto i = arrAttr.begin(), e = arrAttr.end(); i != e; ++i) {
+      if (i->isa<mlir::IntegerAttr>()) {
+        attrs.push_back(*i);
+      } else {
+        auto fieldName = i->cast<mlir::StringAttr>().getValue();
+        ++i;
+        auto ty = i->cast<mlir::TypeAttr>().getValue();
+        auto index = ty.cast<fir::RecordType>().getFieldIndex(fieldName);
+        attrs.push_back(mlir::IntegerAttr::get(rewriter.getI32Type(), index));
+      }
+    }
+    return attrs;
+  }
+
+private:
+  static unsigned getDimension(mlir::LLVM::LLVMArrayType ty) {
+    unsigned result = 1;
+    for (auto eleTy = ty.getElementType().dyn_cast<mlir::LLVM::LLVMArrayType>();
+         eleTy;
+         eleTy = eleTy.getElementType().dyn_cast<mlir::LLVM::LLVMArrayType>())
+      ++result;
+    return result;
+  }
+
+  static mlir::Type getArrayElementType(mlir::LLVM::LLVMArrayType ty) {
+    auto eleTy = ty.getElementType();
+    while (auto arrTy = eleTy.dyn_cast<mlir::LLVM::LLVMArrayType>())
+      eleTy = arrTy.getElementType();
+    return eleTy;
+  }
+};
+
+/// InsertOnRange inserts a value into a sequence over a range of offsets.
+struct InsertOnRangeOpConversion
+    : public FIROpAndTypeConversion<fir::InsertOnRangeOp>,
+      public ValueOpCommon {
+  using FIROpAndTypeConversion::FIROpAndTypeConversion;
+
+  // Increments an array of subscripts in a row major fasion.
+  void incrementSubscripts(const SmallVector<uint64_t> &dims,
+                           SmallVector<uint64_t> &subscripts) const {
+    for (size_t i = dims.size(); i > 0; --i) {
+      if (++subscripts[i - 1] < dims[i - 1]) {
+        return;
+      }
+      subscripts[i - 1] = 0;
+    }
+  }
+
+  mlir::LogicalResult
+  doRewrite(fir::InsertOnRangeOp range, mlir::Type ty, OpAdaptor adaptor,
+            mlir::ConversionPatternRewriter &rewriter) const override {
+    assert(fir::allConstants(adaptor.getOperands().drop_front(2)));
+
+    llvm::SmallVector<mlir::Attribute> lowerBound;
+    llvm::SmallVector<mlir::Attribute> upperBound;
+    llvm::SmallVector<uint64_t> dims;
+    auto type = adaptor.getOperands()[0].getType();
+
+    // Iteratively extract the array dimensions from the type.
+    while (auto t = type.dyn_cast<mlir::LLVM::LLVMArrayType>()) {
+      dims.push_back(t.getNumElements());
+      type = t.getElementType();
+    }
+
+    // Unzip the upper and lower bound subscripts.
+    for (auto i = range.coor().begin(), e = range.coor().end(); i != e; ++i) {
+      lowerBound.push_back(*i++);
+      upperBound.push_back(*i);
+    }
+
+    SmallVector<std::uint64_t> lBounds;
+    SmallVector<std::uint64_t> uBounds;
+
+    // Extract the integer value from the attribute bounds and convert to row
+    // major format.
+    for (std::size_t i = lowerBound.size(); i > 0; --i) {
+      lBounds.push_back(lowerBound[i - 1].cast<IntegerAttr>().getInt());
+      uBounds.push_back(upperBound[i - 1].cast<IntegerAttr>().getInt());
+    }
+
+    auto subscripts(lBounds);
+    auto loc = range.getLoc();
+    mlir::Value lastOp = adaptor.getOperands()[0];
+    mlir::Value insertVal = adaptor.getOperands()[1];
+
+    while (subscripts != uBounds) {
+      // Convert uint64_t's to Attribute's.
+      SmallVector<mlir::Attribute> subscriptAttrs;
+      for (const auto &subscript : subscripts)
+        subscriptAttrs.push_back(
+            IntegerAttr::get(rewriter.getI64Type(), subscript));
+      mlir::ArrayRef<mlir::Attribute> arrayRef(subscriptAttrs);
+      lastOp = rewriter.create<mlir::LLVM::InsertValueOp>(
+          loc, ty, lastOp, insertVal,
+          ArrayAttr::get(range.getContext(), arrayRef));
+
+      incrementSubscripts(dims, subscripts);
+    }
+
+    // Convert uint64_t's to Attribute's.
+    SmallVector<mlir::Attribute> subscriptAttrs;
+    for (const auto &subscript : subscripts)
+      subscriptAttrs.push_back(
+          IntegerAttr::get(rewriter.getI64Type(), subscript));
+    mlir::ArrayRef<mlir::Attribute> arrayRef(subscriptAttrs);
+
+    rewriter.replaceOpWithNewOp<mlir::LLVM::InsertValueOp>(
+        range, ty, lastOp, insertVal,
+        ArrayAttr::get(range.getContext(), arrayRef));
+
+    return success();
+  }
+};
+
 } // namespace
 
 namespace {
 /// Convert FIR dialect to LLVM dialect
 ///
 /// This pass lowers all FIR dialect operations to LLVM IR dialect. An
 /// MLIR pass is used to lower residual Std dialect to LLVM IR dialect.
 ///
 /// This pass is not complete yet. We are upstreaming it in small patches.
 class FIRToLLVMLowering : public fir::FIRToLLVMLoweringBase<FIRToLLVMLowering> {
 public:
   mlir::ModuleOp getModule() { return getOperation(); }
 
   void runOnOperation() override final {
     auto *context = getModule().getContext();
     fir::LLVMTypeConverter typeConverter{getModule()};
     auto loc = mlir::UnknownLoc::get(context);
     mlir::OwningRewritePatternList pattern(context);
     pattern.insert<AddrOfOpConversion, HasValueOpConversion, GlobalOpConversion,
-                   UndefOpConversion>(typeConverter);
+                   InsertOnRangeOpConversion, UndefOpConversion>(typeConverter);
     mlir::populateStdToLLVMConversionPatterns(typeConverter, pattern);
     mlir::arith::populateArithmeticToLLVMConversionPatterns(typeConverter,
                                                             pattern);
     mlir::ConversionTarget target{*context};
     target.addLegalDialect<mlir::LLVM::LLVMDialect>();
 
     // required NOPs for applying a full conversion
     target.addLegalOp<mlir::ModuleOp>();

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

fir.global internal @_QEmultiarray : !fir.array<32x32xi32> {
   %2 = fir.insert_on_range %0, %c0_i32, [0 : index, 31 : index, 0 : index, 31 : index] : (!fir.array<32x32xi32>, i32) -> !fir.array<32x32xi32>
   fir.has_value %2 : !fir.array<32x32xi32>
 }
 
 // CHECK: llvm.mlir.global internal @_QEmultiarray() : !llvm.array<32 x array<32 x i32>> {
 // CHECK:   %[[CST:.*]] = llvm.mlir.constant(dense<1> : vector<32x32xi32>) : !llvm.array<32 x array<32 x i32>>
 // CHECK:   llvm.return %[[CST]] : !llvm.array<32 x array<32 x i32>>
 // CHECK: }
+
+// -----
+
+// Test global with insert_on_range operation not covering the full array
+// in initializer region.
+
+fir.global internal @_QEmultiarray : !fir.array<32x32xi32> {
+  %c0_i32 = arith.constant 1 : i32
+  %0 = fir.undefined !fir.array<32x32xi32>
+  %2 = fir.insert_on_range %0, %c0_i32, [2 : index, 4 : index, 30 : index, 31 : index] : (!fir.array<32x32xi32>, i32) -> !fir.array<32x32xi32>
+  fir.has_value %2 : !fir.array<32x32xi32>
+}
+
+// CHECK:          llvm.mlir.global internal @_QEmultiarray() : !llvm.array<32 x array<32 x i32>> {
+// CHECK:            %[[CST:.*]] = llvm.mlir.constant(1 : i32) : i32
+// CHECK:            %{{.*}} = llvm.mlir.undef : !llvm.array<32 x array<32 x i32>>
+// CHECK-COUNT-35:   %{{.*}} = llvm.insertvalue %[[CST]], %{{.*}}[{{.*}}, {{.*}}] : !llvm.array<32 x array<32 x i32>>
+// CHECK:            llvm.return %{{.*}} : !llvm.array<32 x array<32 x i32>>
+// CHECK:          }