[flang][lowering] Add support for lowering of the `index` intrinsics

[awarzynski]

Summary

This patch adds support for lowering of the index intrinsics from
Fortran to the FIR dialect of MLIR.

This is part of the upstreaming effort from the fir-dev branch in [1].

[1] https://github.com/flang-compiler/f18-llvm-project

Co-authored-by: Jean Perier <[e-mail]>
Co-authored-by: Valentin Clement <[e-mail]>
Co-authored-by: V Donaldson <[e-mail]>
Co-authored-by: Eric Schweitz <[e-mail]>

[clementval]

LGTM

flang/lib/Lower/IntrinsicCall.cpp

struct IntrinsicLibrary {
   /// Lowering for the IAND intrinsic. The IAND intrinsic expects two arguments
   /// in the llvm::ArrayRef.
   mlir::Value genIand(mlir::Type, llvm::ArrayRef<mlir::Value>);
   mlir::Value genIbclr(mlir::Type, llvm::ArrayRef<mlir::Value>);
   mlir::Value genIbits(mlir::Type, llvm::ArrayRef<mlir::Value>);
   mlir::Value genIbset(mlir::Type, llvm::ArrayRef<mlir::Value>);
   fir::ExtendedValue genIchar(mlir::Type, llvm::ArrayRef<fir::ExtendedValue>);
   mlir::Value genIeor(mlir::Type, llvm::ArrayRef<mlir::Value>);
+  fir::ExtendedValue genIndex(mlir::Type, llvm::ArrayRef<fir::ExtendedValue>);
   mlir::Value genIshft(mlir::Type, llvm::ArrayRef<mlir::Value>);
   mlir::Value genIshftc(mlir::Type, llvm::ArrayRef<mlir::Value>);
   fir::ExtendedValue genLbound(mlir::Type, llvm::ArrayRef<fir::ExtendedValue>);
   fir::ExtendedValue genNull(mlir::Type, llvm::ArrayRef<fir::ExtendedValue>);
   fir::ExtendedValue genLen(mlir::Type, llvm::ArrayRef<fir::ExtendedValue>);
   fir::ExtendedValue genLenTrim(mlir::Type, llvm::ArrayRef<fir::ExtendedValue>);
   fir::ExtendedValue genMaxloc(mlir::Type, llvm::ArrayRef<fir::ExtendedValue>);
   fir::ExtendedValue genMaxval(mlir::Type, llvm::ArrayRef<fir::ExtendedValue>);
static constexpr IntrinsicHandler handlers[]{
      /*isElemental=*/false},
     {"iachar", &I::genIchar},
     {"iand", &I::genIand},
     {"ibclr", &I::genIbclr},
     {"ibits", &I::genIbits},
     {"ibset", &I::genIbset},
     {"ichar", &I::genIchar},
     {"ieor", &I::genIeor},
+    {"index",
+     &I::genIndex,
+     {{{"string", asAddr},
+       {"substring", asAddr},
+       {"back", asValue, handleDynamicOptional},
+       {"kind", asValue}}}},
     {"ishft", &I::genIshft},
     {"ishftc", &I::genIshftc},
     {"len",
      &I::genLen,
      {{{"string", asInquired}, {"kind", asValue}}},
      /*isElemental=*/false},
     {"len_trim", &I::genLenTrim},
     {"lge", &I::genCharacterCompare<mlir::arith::CmpIPredicate::sge>},

 
 // IEOR
 mlir::Value IntrinsicLibrary::genIeor(mlir::Type resultType,
                                       llvm::ArrayRef<mlir::Value> args) {
   assert(args.size() == 2);
   return builder.create<mlir::arith::XOrIOp>(loc, args[0], args[1]);
 }
 
+// INDEX
+fir::ExtendedValue
+IntrinsicLibrary::genIndex(mlir::Type resultType,
+                           llvm::ArrayRef<fir::ExtendedValue> args) {
+  assert(args.size() >= 2 && args.size() <= 4);
+
+  mlir::Value stringBase = fir::getBase(args[0]);
+  fir::KindTy kind =
+      fir::factory::CharacterExprHelper{builder, loc}.getCharacterKind(
+          stringBase.getType());
+  mlir::Value stringLen = fir::getLen(args[0]);
+  mlir::Value substringBase = fir::getBase(args[1]);
+  mlir::Value substringLen = fir::getLen(args[1]);
+  mlir::Value back =
+      isAbsent(args, 2)
+          ? builder.createIntegerConstant(loc, builder.getI1Type(), 0)
+          : fir::getBase(args[2]);
+  if (isAbsent(args, 3))
+    return builder.createConvert(
+        loc, resultType,
+        fir::runtime::genIndex(builder, loc, kind, stringBase, stringLen,
+                               substringBase, substringLen, back));
+
+  // Call the descriptor-based Index implementation
+  mlir::Value string = builder.createBox(loc, args[0]);
+  mlir::Value substring = builder.createBox(loc, args[1]);
+  auto makeRefThenEmbox = [&](mlir::Value b) {
+    fir::LogicalType logTy = fir::LogicalType::get(
+        builder.getContext(), builder.getKindMap().defaultLogicalKind());
+    mlir::Value temp = builder.createTemporary(loc, logTy);
+    mlir::Value castb = builder.createConvert(loc, logTy, b);
+    builder.create<fir::StoreOp>(loc, castb, temp);
+    return builder.createBox(loc, temp);
+  };
+  mlir::Value backOpt = isAbsent(args, 2)
+                            ? builder.create<fir::AbsentOp>(
+                                  loc, fir::BoxType::get(builder.getI1Type()))
+                            : makeRefThenEmbox(fir::getBase(args[2]));
+  mlir::Value kindVal = isAbsent(args, 3)
+                            ? builder.createIntegerConstant(
+                                  loc, builder.getIndexType(),
+                                  builder.getKindMap().defaultIntegerKind())
+                            : fir::getBase(args[3]);
+  // Create mutable fir.box to be passed to the runtime for the result.
+  fir::MutableBoxValue mutBox =
+      fir::factory::createTempMutableBox(builder, loc, resultType);
+  mlir::Value resBox = fir::factory::getMutableIRBox(builder, loc, mutBox);
+  // Call runtime. The runtime is allocating the result.
+  fir::runtime::genIndexDescriptor(builder, loc, resBox, string, substring,
+                                   backOpt, kindVal);
+  // Read back the result from the mutable box.
+  return readAndAddCleanUp(mutBox, resultType, "INDEX");
+}
+
 // ISHFT
 mlir::Value IntrinsicLibrary::genIshft(mlir::Type resultType,
                                        llvm::ArrayRef<mlir::Value> args) {
   // A conformant ISHFT(I,SHIFT) call satisfies:
   //     abs(SHIFT) <= BIT_SIZE(I)
   // Return:  abs(SHIFT) >= BIT_SIZE(I)
   //              ? 0
   //              : SHIFT < 0

flang/test/Lower/Intrinsics/index.f90

+! RUN: bbc -emit-fir %s -o - | FileCheck %s
+
+! CHECK-LABEL: func @_QPindex_test(
+! CHECK-SAME: %[[s:[^:]+]]: !fir.boxchar<1>{{.*}}, %[[ss:[^:]+]]: !fir.boxchar<1>{{.*}}) -> i32
+integer function index_test(s1, s2)
+  character(*) :: s1, s2
+  ! CHECK: %[[st:[^:]*]]:2 = fir.unboxchar %[[s]] : (!fir.boxchar<1>) -> (!fir.ref<!fir.char<1,?>>, index)
+  ! CHECK: %[[sst:[^:]*]]:2 = fir.unboxchar %[[ss]] : (!fir.boxchar<1>) -> (!fir.ref<!fir.char<1,?>>, index)
+  ! CHECK: %[[a1:.*]] = fir.convert %[[st]]#0 : (!fir.ref<!fir.char<1,?>>) -> !fir.ref<i8>
+  ! CHECK: %[[a2:.*]] = fir.convert %[[st]]#1 : (index) -> i64
+  ! CHECK: %[[a3:.*]] = fir.convert %[[sst]]#0 : (!fir.ref<!fir.char<1,?>>) -> !fir.ref<i8>
+  ! CHECK: %[[a4:.*]] = fir.convert %[[sst]]#1 : (index) -> i64
+  ! CHECK: = fir.call @_FortranAIndex1(%[[a1]], %[[a2]], %[[a3]], %[[a4]], %{{.*}}) : (!fir.ref<i8>, i64, !fir.ref<i8>, i64, i1) -> i64
+  index_test = index(s1, s2)
+end function index_test
+
+! CHECK-LABEL: func @_QPindex_test2(
+! CHECK-SAME: %[[s:[^:]+]]: !fir.boxchar<1>{{.*}}, %[[ss:[^:]+]]: !fir.boxchar<1>{{.*}}) -> i32
+integer function index_test2(s1, s2)
+  character(*) :: s1, s2
+  ! CHECK: %[[mut:.*]] = fir.alloca !fir.box<!fir.heap<i32>>
+  ! CHECK: %[[st:[^:]*]]:2 = fir.unboxchar %[[s]] : (!fir.boxchar<1>) -> (!fir.ref<!fir.char<1,?>>, index)
+  ! CHECK: %[[sst:[^:]*]]:2 = fir.unboxchar %[[ss]] : (!fir.boxchar<1>) -> (!fir.ref<!fir.char<1,?>>, index)
+  ! CHECK: %[[sb:.*]] = fir.embox %[[st]]#0 typeparams %[[st]]#1 : (!fir.ref<!fir.char<1,?>>, index) -> !fir.box<!fir.char<1,?>>
+  ! CHECK: %[[ssb:.*]] = fir.embox %[[sst]]#0 typeparams %[[sst]]#1 : (!fir.ref<!fir.char<1,?>>, index) -> !fir.box<!fir.char<1,?>>
+  ! CHECK: %[[back:.*]] = fir.embox %{{.*}} : (!fir.ref<!fir.logical<4>>) -> !fir.box<!fir.logical<4>>
+  ! CHECK: %[[hb:.*]] = fir.embox %{{.*}} : (!fir.heap<i32>) -> !fir.box<!fir.heap<i32>>
+  ! CHECK: %[[a0:.*]] = fir.convert %[[mut]] : (!fir.ref<!fir.box<!fir.heap<i32>>>) -> !fir.ref<!fir.box<none>>
+  ! CHECK: %[[a1:.*]] = fir.convert %[[sb]] : (!fir.box<!fir.char<1,?>>) -> !fir.box<none>
+  ! CHECK: %[[a2:.*]] = fir.convert %[[ssb]] : (!fir.box<!fir.char<1,?>>) -> !fir.box<none>
+  ! CHECK: %[[a3:.*]] = fir.convert %[[back]] : (!fir.box<!fir.logical<4>>) -> !fir.box<none>
+  ! CHECK: %[[a5:.*]] = fir.convert %{{.*}} : (!fir.ref<!fir.char<1,{{.*}}>>) -> !fir.ref<i8>
+  ! CHECK:  fir.call @_FortranAIndex(%[[a0]], %[[a1]], %[[a2]], %[[a3]], %{{.*}}, %[[a5]], %{{.*}}) : (!fir.ref<!fir.box<none>>, !fir.box<none>, !fir.box<none>, !fir.box<none>, i32, !fir.ref<i8>, i32) -> none
+  index_test2 = index(s1, s2, .true., 4)
+  ! CHECK: %[[ld1:.*]] = fir.load %[[mut]] : !fir.ref<!fir.box<!fir.heap<i32>>>
+  ! CHECK: %[[ad1:.*]] = fir.box_addr %[[ld1]] : (!fir.box<!fir.heap<i32>>) -> !fir.heap<i32>
+  ! CHECK: %[[ld2:.*]] = fir.load %[[ad1]] : !fir.heap<i32>
+  ! CHECK: fir.freemem %[[ad1]]
+end function index_test2
+
+! CHECK-LABEL: func @_QPindex_test3
+integer function index_test3(s, i)
+  character(*) :: s
+  integer :: i
+  ! CHECK: %[[tmpChar:.*]] = fir.alloca !fir.char<1>
+  ! CHECK: fir.store %{{.*}} to %[[tmpChar]] : !fir.ref<!fir.char<1>>
+  ! CHECK: %[[tmpCast:.*]] = fir.convert %[[tmpChar]] : (!fir.ref<!fir.char<1>>) -> !fir.ref<i8>
+  ! CHECK: fir.call @_FortranAIndex1(%{{.*}}, %{{.*}}, %[[tmpCast]], %{{.*}}, %{{.*}})
+  index_test3 = index(s, char(i))
+end function
+
+! CHECK-LABEL: func @_QPtest_optional(
+! CHECK-SAME:  %[[VAL_0:.*]]: !fir.box<!fir.array<?x!fir.char<1,?>>>
+! CHECK-SAME:  %[[VAL_1:.*]]: !fir.boxchar<1>
+! CHECK-SAME:  %[[VAL_2:.*]]: !fir.box<!fir.array<?x!fir.logical<4>>>
+subroutine test_optional(string, substring, back)
+  character (*) :: string(:), substring
+  logical, optional :: back(:)
+  print *, index(string, substring, back)
+! CHECK:  %[[VAL_11:.*]] = fir.is_present %[[VAL_2]] : (!fir.box<!fir.array<?x!fir.logical<4>>>) -> i1
+! CHECK:  %[[VAL_12:.*]] = fir.zero_bits !fir.ref<!fir.array<?x!fir.logical<4>>>
+! CHECK:  %[[VAL_13:.*]] = arith.constant 0 : index
+! CHECK:  %[[VAL_14:.*]] = fir.shape %[[VAL_13]] : (index) -> !fir.shape<1>
+! CHECK:  %[[VAL_15:.*]] = fir.embox %[[VAL_12]](%[[VAL_14]]) : (!fir.ref<!fir.array<?x!fir.logical<4>>>, !fir.shape<1>) -> !fir.box<!fir.array<?x!fir.logical<4>>>
+! CHECK:  %[[VAL_16:.*]] = arith.select %[[VAL_11]], %[[VAL_2]], %[[VAL_15]] : !fir.box<!fir.array<?x!fir.logical<4>>>
+! CHECK:  %[[VAL_17:.*]] = fir.array_load %[[VAL_16]] {fir.optional} : (!fir.box<!fir.array<?x!fir.logical<4>>>) -> !fir.array<?x!fir.logical<4>>
+! CHECK:  %[[VAL_24:.*]] = fir.do_loop %[[VAL_25:.*]] = %{{.*}} to %{{.*}} step %{{.*}} unordered iter_args(%[[VAL_26:.*]] = %{{.*}}) -> (!fir.array<?xi32>) {
+  ! CHECK:  %[[VAL_31:.*]] = fir.if %[[VAL_11]] -> (!fir.logical<4>) {
+    ! CHECK:  %[[VAL_32:.*]] = fir.array_fetch %[[VAL_17]], %[[VAL_25]] : (!fir.array<?x!fir.logical<4>>, index) -> !fir.logical<4>
+    ! CHECK:  fir.result %[[VAL_32]] : !fir.logical<4>
+  ! CHECK:  } else {
+    ! CHECK:  %[[VAL_33:.*]] = arith.constant false
+    ! CHECK:  %[[VAL_34:.*]] = fir.convert %[[VAL_33]] : (i1) -> !fir.logical<4>
+    ! CHECK:  fir.result %[[VAL_34]] : !fir.logical<4>
+  ! CHECK:  }
+  ! CHECK:  %[[VAL_39:.*]] = fir.convert %[[VAL_31]] : (!fir.logical<4>) -> i1
+  ! CHECK:  fir.call @_FortranAIndex1(%{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %[[VAL_39]]) : (!fir.ref<i8>, i64, !fir.ref<i8>, i64, i1) -> i64
+! CHECK:  }
+! CHECK:  fir.array_merge_store
+end subroutine