[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

34

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);
});

67

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

200

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.

205

Same as above, unused.

227

ditto

259

Actually the entire class seems unused?

275

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).

281

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?

301

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

308

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)

310

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;

320

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

321

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

[mehdi_amini]

flang/lib/Optimizer/CodeGen/CodeGen.cpp

59

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

34

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

59

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.

200

No needed now so will bring it back later.

205

Removed for now

227

Removed for now

259

Removed for now

275

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

281

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

310

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

275

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

275

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

216

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

216

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 "PassDetail.h"
 #include "flang/Optimizer/Dialect/FIROps.h"
 #include "flang/Optimizer/Dialect/FIRType.h"
 #include "mlir/Conversion/ArithmeticToLLVM/ArithmeticToLLVM.h"
 #include "mlir/Conversion/LLVMCommon/Pattern.h"
 #include "mlir/Conversion/LLVMCommon/TypeConverter.h"
 #include "mlir/Conversion/StandardToLLVM/ConvertStandardToLLVM.h"
 #include "mlir/IR/BuiltinTypes.h"
+#include "mlir/IR/Matchers.h"
 #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 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 = 0;
+};
+
 // Lower `fir.address_of` operation to `llvm.address_of` operation.
 struct AddrOfOpConversion : public FIROpConversion<fir::AddrOfOp> {
   using FIROpConversion::FIROpConversion;
 
   mlir::LogicalResult
   matchAndRewrite(fir::AddrOfOp addr, OpAdaptor adaptor,
                   mlir::ConversionPatternRewriter &rewriter) const override {
     auto ty = convertType(addr.getType());
struct UnreachableOpConversion : public FIROpConversion<fir::UnreachableOp> {
   using FIROpConversion::FIROpConversion;
 
   mlir::LogicalResult
   matchAndRewrite(fir::UnreachableOp unreach, OpAdaptor adaptor,
                   mlir::ConversionPatternRewriter &rewriter) const override {
     rewriter.replaceOpWithNewOp<mlir::LLVM::UnreachableOp>(unreach);
     return success();
   }
 };
 
 struct ZeroOpConversion : public FIROpConversion<fir::ZeroOp> {
   using FIROpConversion::FIROpConversion;
 
   mlir::LogicalResult
   matchAndRewrite(fir::ZeroOp zero, OpAdaptor,
                   mlir::ConversionPatternRewriter &rewriter) const override {
     auto ty = convertType(zero.getType());
     if (ty.isa<mlir::LLVM::LLVMPointerType>()) {
       rewriter.replaceOpWithNewOp<mlir::LLVM::NullOp>(zero, ty);
     } else if (ty.isa<mlir::IntegerType>()) {
       rewriter.replaceOpWithNewOp<mlir::LLVM::ConstantOp>(
           zero, ty, mlir::IntegerAttr::get(zero.getType(), 0));
     } else if (mlir::LLVM::isCompatibleFloatingPointType(ty)) {
       rewriter.replaceOpWithNewOp<mlir::LLVM::ConstantOp>(
           zero, ty, mlir::FloatAttr::get(zero.getType(), 0.0));
     } else {
       // TODO: create ConstantAggregateZero for FIR aggregate/array types.
       return rewriter.notifyMatchFailure(
           zero,
           "conversion of fir.zero with aggregate type not implemented yet");
     }
     return success();
   }
 };
 
+/// InsertOnRange inserts a value into a sequence over a range of offsets.
+struct InsertOnRangeOpConversion
+    : public FIROpAndTypeConversion<fir::InsertOnRangeOp> {
+  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 {
+
+    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();
+    }
+
+    SmallVector<uint64_t> lBounds;
+    SmallVector<uint64_t> uBounds;
+
+    // Extract integer value from the attribute
+    SmallVector<int64_t> coordinates = llvm::to_vector<4>(
+        llvm::map_range(range.coor(), [](Attribute a) -> int64_t {
+          return a.cast<IntegerAttr>().getInt();
+        }));
+
+    // Unzip the upper and lower bound and convert to a row major format.
+    for (auto i = coordinates.rbegin(), e = coordinates.rend(); i != e; ++i) {
+      uBounds.push_back(*i++);
+      lBounds.push_back(*i);
+    }
+
+    auto &subscripts = lBounds;
+    auto loc = range.getLoc();
+    mlir::Value lastOp = adaptor.getOperands()[0];
+    mlir::Value insertVal = adaptor.getOperands()[1];
+
+    auto i64Ty = rewriter.getI64Type();
+    while (subscripts != uBounds) {
+      // Convert uint64_t's to Attribute's.
+      SmallVector<mlir::Attribute> subscriptAttrs;
+      for (const auto &subscript : subscripts)
+        subscriptAttrs.push_back(IntegerAttr::get(i64Ty, subscript));
+      lastOp = rewriter.create<mlir::LLVM::InsertValueOp>(
+          loc, ty, lastOp, insertVal,
+          ArrayAttr::get(range.getContext(), subscriptAttrs));
+
+      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()};
     mlir::OwningRewritePatternList pattern(context);
-    pattern
-        .insert<AddrOfOpConversion, HasValueOpConversion, GlobalOpConversion,
-                UndefOpConversion, UnreachableOpConversion, ZeroOpConversion>(
-            typeConverter);
+    pattern.insert<AddrOfOpConversion, HasValueOpConversion, GlobalOpConversion,
+                   InsertOnRangeOpConversion, UndefOpConversion,
+                   UnreachableOpConversion, ZeroOpConversion>(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

 
 // 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<32xi32> {
+  %c0_i32 = arith.constant 1 : i32
+  %0 = fir.undefined !fir.array<32xi32>
+  %2 = fir.insert_on_range %0, %c0_i32, [5 : index, 31 : index] : (!fir.array<32xi32>, i32) -> !fir.array<32xi32>
+  fir.has_value %2 : !fir.array<32xi32>
+}
+
+// CHECK:          llvm.mlir.global internal @_QEmultiarray() : !llvm.array<32 x i32> {
+// CHECK:            %[[CST:.*]] = llvm.mlir.constant(1 : i32) : i32
+// CHECK:            %{{.*}} = llvm.mlir.undef : !llvm.array<32 x i32>
+// CHECK:            %{{.*}} = llvm.insertvalue %[[CST]], %{{.*}}[5] : !llvm.array<32 x i32>
+// CHECK-COUNT-24:   %{{.*}} = llvm.insertvalue %[[CST]], %{{.*}}[{{.*}}] : !llvm.array<32 x i32>
+// CHECK:            %{{.*}} = llvm.insertvalue %[[CST]], %{{.*}}[31] : !llvm.array<32 x i32>
+// CHECK-NOT:        llvm.insertvalue
+// CHECK:            llvm.return %{{.*}} : !llvm.array<32 x i32>
+// CHECK:          }
+
+// -----
+
 // Test fir.zero_bits operation with LLVM ptr type
 
 func @zero_test_ptr() {
   %z = fir.zero_bits !llvm.ptr<f32>
   return
 }
 
 // CHECK: %{{.*}} = llvm.mlir.null : !llvm.ptr<f32>