[fir] Add substr information to fircg.ext_embox and fircg.ext_rebox operations

[clementval]

Summary

This patch adds the substring information to the fircg.ext_embox and
fircg.ext_rebox operations.

Substring is used for CHARACTER types.

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

Co-authored-by: Eric Schweitz <[e-mail]>

[mehdi_amini]

Looks fine,

flang/include/flang/Optimizer/CodeGen/CGOps.td

let description = [{
     arguments.
 
        - memref: The memory reference being emboxed.
        - shape: A vector that is the runtime shape of the underlying array.
        - shift: A vector that is the runtime origin of the first element.
          The default is a vector of the value 1.
        - slice: A vector of triples that describe an array slice.
        - subcomponent: A vector of indices for subobject slicing.
+       - substring: A substring operator (offset, length) for CHARACTER.
        - LEN type parameters: A vector of runtime LEN type parameters that
          describe an correspond to the elemental derived type.
 
     The memref and shape arguments are mandatory. The rest are optional.
   }];
 
   let arguments = (ins
     AnyReferenceLike:$memref,
     Variadic<AnyIntegerType>:$shape,
     Variadic<AnyIntegerType>:$shift,
     Variadic<AnyIntegerType>:$slice,
     Variadic<AnyCoordinateType>:$subcomponent,
+    Variadic<AnyIntegerType>:$substr,
     Variadic<AnyIntegerType>:$lenParams
   );
   let results = (outs fir_BoxType);
 
   let assemblyFormat = [{
     $memref (`(`$shape^`)`)? (`origin` $shift^)? (`[`$slice^`]`)?
-      (`path` $subcomponent^)? (`typeparams` $lenParams^)? attr-dict
-      `:` functional-type(operands, results)
+      (`path` $subcomponent^)? (`substr` $substr^)? (`typeparams` $lenParams^)?
+      attr-dict `:` functional-type(operands, results)
   }];
 
   let extraClassDeclaration = [{
     // The rank of the entity being emboxed
     unsigned getRank() { return shape().size(); }
 
     // The rank of the result. A slice op can reduce the rank.
     unsigned getOutRank();
 
     // The shape operands are mandatory and always start at 1.
     unsigned shapeOffset() { return 1; }
     unsigned shiftOffset() { return shapeOffset() + shape().size(); }
     unsigned sliceOffset() { return shiftOffset() + shift().size(); }
     unsigned subcomponentOffset() { return sliceOffset() + slice().size(); }
-    unsigned lenParamOffset() {
+    unsigned substrOffset() {
        return subcomponentOffset() + subcomponent().size();
     }
+    unsigned lenParamOffset() { return substrOffset() + substr().size(); }
   }];
 }
 
 // Extended rebox operation.
 def fircg_XReboxOp : fircg_Op<"ext_rebox", [AttrSizedOperandSegments]> {
   let summary = "for internal conversion only";
 
   let description = [{
     Prior to lowering to LLVM IR dialect, a non-scalar non-trivial rebox op will
     be converted to an extended rebox. This op will have the following sets of
     arguments.
 
        - box: The box being reboxed.
        - shape: A vector that is the new runtime shape for the array
        - shift: A vector that is the new runtime origin of the first element.
          The default is a vector of the value 1.
        - slice: A vector of triples that describe an array slice.
        - subcomponent: A vector of indices for subobject slicing.
+       - substring: A substring operator (offset, length) for CHARACTER.
 
     The box argument is mandatory, the other arguments are optional.
     There must not both be a shape and slice/subcomponent arguments
   }];
 
   let arguments = (ins
     fir_BoxType:$box,
     Variadic<AnyIntegerType>:$shape,
     Variadic<AnyIntegerType>:$shift,
     Variadic<AnyIntegerType>:$slice,
-    Variadic<AnyCoordinateType>:$subcomponent
+    Variadic<AnyCoordinateType>:$subcomponent,
+    Variadic<AnyIntegerType>:$substr
   );
   let results = (outs fir_BoxType);
 
   let assemblyFormat = [{
     $box (`(`$shape^`)`)? (`origin` $shift^)? (`[`$slice^`]`)?
-      (`path` $subcomponent^) ? attr-dict
-      `:` functional-type(operands, results)
+      (`path` $subcomponent^)? (`substr` $substr^)? attr-dict `:`
+      functional-type(operands, results)
   }];
 
   let extraClassDeclaration = [{
     // The rank of the entity being reboxed
     unsigned getRank();
     // The rank of the result box
     unsigned getOutRank();
   }];

flang/lib/Optimizer/CodeGen/PreCGRewrite.cpp

mlir::LogicalResult rewriteStaticShape(EmboxOp embox,
     auto idxTy = rewriter.getIndexType();
     for (auto ext : seqTy.getShape()) {
       auto iAttr = rewriter.getIndexAttr(ext);
       auto extVal = rewriter.create<mlir::arith::ConstantOp>(loc, idxTy, iAttr);
       shapeOpers.push_back(extVal);
     }
     auto xbox = rewriter.create<cg::XEmboxOp>(
         loc, embox.getType(), embox.memref(), shapeOpers, llvm::None,
-        llvm::None, llvm::None, embox.typeparams());
+        llvm::None, llvm::None, llvm::None, embox.typeparams());
     LLVM_DEBUG(llvm::dbgs() << "rewriting " << embox << " to " << xbox << '\n');
     rewriter.replaceOp(embox, xbox.getOperation()->getResults());
     return mlir::success();
   }
 
   mlir::LogicalResult rewriteDynamicShape(EmboxOp embox,
                                           mlir::PatternRewriter &rewriter,
                                           mlir::Value shapeVal) const {
     auto loc = embox.getLoc();
     auto shapeOp = dyn_cast<ShapeOp>(shapeVal.getDefiningOp());
     llvm::SmallVector<mlir::Value> shapeOpers;
     llvm::SmallVector<mlir::Value> shiftOpers;
     if (shapeOp) {
       populateShape(shapeOpers, shapeOp);
     } else {
       auto shiftOp = dyn_cast<ShapeShiftOp>(shapeVal.getDefiningOp());
       assert(shiftOp && "shape is neither fir.shape nor fir.shape_shift");
       populateShapeAndShift(shapeOpers, shiftOpers, shiftOp);
     }
     llvm::SmallVector<mlir::Value> sliceOpers;
     llvm::SmallVector<mlir::Value> subcompOpers;
+    llvm::SmallVector<mlir::Value> substrOpers;
     if (auto s = embox.getSlice())
       if (auto sliceOp = dyn_cast_or_null<SliceOp>(s.getDefiningOp())) {
-        sliceOpers.append(sliceOp.triples().begin(), sliceOp.triples().end());
-        subcompOpers.append(sliceOp.fields().begin(), sliceOp.fields().end());
+        sliceOpers.assign(sliceOp.triples().begin(), sliceOp.triples().end());
+        subcompOpers.assign(sliceOp.fields().begin(), sliceOp.fields().end());
+        substrOpers.assign(sliceOp.substr().begin(), sliceOp.substr().end());
       }
     auto xbox = rewriter.create<cg::XEmboxOp>(
         loc, embox.getType(), embox.memref(), shapeOpers, shiftOpers,
-        sliceOpers, subcompOpers, embox.typeparams());
+        sliceOpers, subcompOpers, substrOpers, embox.typeparams());
     LLVM_DEBUG(llvm::dbgs() << "rewriting " << embox << " to " << xbox << '\n');
     rewriter.replaceOp(embox, xbox.getOperation()->getResults());
     return mlir::success();
   }
 };
 
 /// Convert fir.rebox to the extended form where necessary.
 ///
if (auto shapeVal = rebox.shape()) {
         populateShapeAndShift(shapeOpers, shiftOpers, shiftOp);
       else if (auto shiftOp = dyn_cast<ShiftOp>(shapeVal.getDefiningOp()))
         populateShift(shiftOpers, shiftOp);
       else
         return mlir::failure();
     }
     llvm::SmallVector<mlir::Value> sliceOpers;
     llvm::SmallVector<mlir::Value> subcompOpers;
+    llvm::SmallVector<mlir::Value> substrOpers;
     if (auto s = rebox.slice())
       if (auto sliceOp = dyn_cast_or_null<SliceOp>(s.getDefiningOp())) {
         sliceOpers.append(sliceOp.triples().begin(), sliceOp.triples().end());
         subcompOpers.append(sliceOp.fields().begin(), sliceOp.fields().end());
+        substrOpers.append(sliceOp.substr().begin(), sliceOp.substr().end());
       }
 
     auto xRebox = rewriter.create<cg::XReboxOp>(
         loc, rebox.getType(), rebox.box(), shapeOpers, shiftOpers, sliceOpers,
-        subcompOpers);
+        subcompOpers, substrOpers);
     LLVM_DEBUG(llvm::dbgs()
                << "rewriting " << rebox << " to " << xRebox << '\n');
     rewriter.replaceOp(rebox, xRebox.getOperation()->getResults());
     return mlir::success();
   }
 };
 
 /// Convert all fir.array_coor to the extended form.
if (auto shapeVal = arrCoor.shape()) {
         return mlir::failure();
     }
     llvm::SmallVector<mlir::Value> sliceOpers;
     llvm::SmallVector<mlir::Value> subcompOpers;
     if (auto s = arrCoor.slice())
       if (auto sliceOp = dyn_cast_or_null<SliceOp>(s.getDefiningOp())) {
         sliceOpers.append(sliceOp.triples().begin(), sliceOp.triples().end());
         subcompOpers.append(sliceOp.fields().begin(), sliceOp.fields().end());
+        assert(sliceOp.substr().empty() &&
+               "Don't allow substring operations on array_coor. This "
+               "restriction may be lifted in the future.");
       }
     auto xArrCoor = rewriter.create<cg::XArrayCoorOp>(
         loc, arrCoor.getType(), arrCoor.memref(), shapeOpers, shiftOpers,
         sliceOpers, subcompOpers, arrCoor.indices(), arrCoor.typeparams());
     LLVM_DEBUG(llvm::dbgs()
                << "rewriting " << arrCoor << " to " << xArrCoor << '\n');
     rewriter.replaceOp(arrCoor, xArrCoor.getOperation()->getResults());
     return mlir::success();

flang/test/Fir/cg-ops.fir

-// RUN: fir-opt --pass-pipeline="builtin.func(cg-rewrite),fir.global(cg-rewrite),cse" %s | FileCheck %s
+// RUN: fir-opt --split-input-file --pass-pipeline="cg-rewrite,cse" %s | FileCheck %s
 
 // CHECK-LABEL: func @codegen(
 // CHECK-SAME: %[[arg:.*]]: !fir
 func @codegen(%addr : !fir.ref<!fir.array<?xi32>>) {
   // CHECK: %[[zero:.*]] = arith.constant 0 : index
   %0 = arith.constant 0 : index
   %1 = fir.shape_shift %0, %0 : (index, index) -> !fir.shapeshift<1>
   %2 = fir.slice %0, %0, %0 : (index, index, index) -> !fir.slice<1>
   // CHECK: %[[box:.*]] = fircg.ext_embox %[[arg]](%[[zero]]) origin %[[zero]][%[[zero]], %[[zero]], %[[zero]]] : (!fir.ref<!fir.array<?xi32>>, index, index, index, index, index) -> !fir.box<!fir.array<?xi32>>
   %3 = fir.embox %addr (%1) [%2] : (!fir.ref<!fir.array<?xi32>>, !fir.shapeshift<1>, !fir.slice<1>) -> !fir.box<!fir.array<?xi32>>
   // CHECK: fircg.ext_array_coor %[[arg]](%[[zero]]) origin %[[zero]][%[[zero]], %[[zero]], %[[zero]]]<%[[zero]]> : (!fir.ref<!fir.array<?xi32>>, index, index, index, index, index, index) -> !fir.ref<i32>
   %4 = fir.array_coor %addr (%1) [%2] %0 : (!fir.ref<!fir.array<?xi32>>, !fir.shapeshift<1>, !fir.slice<1>, index) -> !fir.ref<i32>
   // CHECK: fircg.ext_rebox %[[box]](%[[zero]]) origin %[[zero]] : (!fir.box<!fir.array<?xi32>>, index, index) -> !fir.box<!fir.array<?xi32>>
   %5 = fir.rebox %3(%1) : (!fir.box<!fir.array<?xi32>>, !fir.shapeshift<1>) -> !fir.box<!fir.array<?xi32>>
   return
 }
 
+// -----
+
 // CHECK-LABEL: fir.global @box_global
 fir.global @box_global : !fir.box<!fir.array<?xi32>> {
   // CHECK: %[[arr:.*]] = fir.zero_bits !fir.ref
   %arr = fir.zero_bits !fir.ref<!fir.array<?xi32>>
   // CHECK: %[[zero:.*]] = arith.constant 0 : index
   %0 = arith.constant 0 : index
   %1 = fir.shape_shift %0, %0 : (index, index) -> !fir.shapeshift<1>
   %2 = fir.slice %0, %0, %0 : (index, index, index) -> !fir.slice<1>
   // CHECK: fircg.ext_embox %[[arr]](%[[zero]]) origin %[[zero]][%[[zero]], %[[zero]], %[[zero]]] : (!fir.ref<!fir.array<?xi32>>, index, index, index, index, index) -> !fir.box<!fir.array<?xi32>>
   %3 = fir.embox %arr (%1) [%2] : (!fir.ref<!fir.array<?xi32>>, !fir.shapeshift<1>, !fir.slice<1>) -> !fir.box<!fir.array<?xi32>>
   fir.has_value %3 : !fir.box<!fir.array<?xi32>>
 }
+
+// -----
+
+// fir.embox with slice with substr
+
+// CHECK-LABEL: func @codegen(
+// CHECK-SAME: %[[arg:.*]]: !fir
+func @codegen(%addr : !fir.ref<!fir.array<?xi32>>) {
+  // CHECK: %[[zero:.*]] = arith.constant 0 : index
+  %0 = arith.constant 0 : index
+  %1 = fir.shape_shift %0, %0 : (index, index) -> !fir.shapeshift<1>
+  %2 = fir.slice %0, %0, %0 substr %0, %0: (index, index, index, index, index) -> !fir.slice<1>
+  // CHECK: %[[box:.*]] = fircg.ext_embox %[[arg]](%[[zero]]) origin %[[zero]][%[[zero]], %[[zero]], %[[zero]]] substr %[[zero]], %[[zero]] : (!fir.ref<!fir.array<?xi32>>, index, index, index, index, index, index, index) -> !fir.box<!fir.array<?xi32>>
+  %3 = fir.embox %addr (%1) [%2] : (!fir.ref<!fir.array<?xi32>>, !fir.shapeshift<1>, !fir.slice<1>) -> !fir.box<!fir.array<?xi32>>
+  return
+}
+
+// -----
+
+// fir.rebox with slice with substr
+
+// CHECK-LABEL: func @codegen(
+// CHECK-SAME: %[[arg:.*]]: !fir
+func @codegen(%addr : !fir.box<!fir.array<?xf32>>) {
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 1 : index
+  %c10 = arith.constant 10 : index
+  %0 = fir.slice %c10, %c1, %c1 substr %c1, %c1: (index, index, index, index, index) -> !fir.slice<1>
+  %1 = fir.shift %c0 : (index) -> !fir.shift<1>
+  %2 = fir.rebox %addr(%1) [%0] : (!fir.box<!fir.array<?xf32>>, !fir.shift<1>, !fir.slice<1>) -> !fir.box<!fir.array<?xf32>>
+  return
+}
+
+// CHECK: %{{.*}} = fircg.ext_rebox %[[arg]] origin %{{.*}}[%{{.*}}, %{{.*}}, %{{.*}}] substr %{{.*}}, %{{.*}} : (!fir.box<!fir.array<?xf32>>, index, index, index, index, index, index) -> !fir.box<!fir.array<?xf32>>