Rewrite index types according to target

Co-authored-by: Lionel Henry <[email protected]>

Author: georgestagg

flang/include/flang/Optimizer/CodeGen/Target.h

@@ -87,6 +87,9 @@ class CodeGenSpecifics {
   virtual Marshalling complexReturnType(mlir::Location loc,
                                         mlir::Type eleTy) const = 0;
 
+  // Type presentation of a `index` type value in memory.
+  virtual mlir::Type indexMemoryType(mlir::Type eleTy) const = 0;
+
   /// Type presentation of a `boxchar<n>` type value in memory.
   virtual mlir::Type boxcharMemoryType(mlir::Type eleTy) const = 0;
 

flang/lib/Optimizer/CodeGen/Target.cpp

@@ -58,6 +58,10 @@ struct GenericTarget : public CodeGenSpecifics {
                                 mlir::TypeRange{eleTy, eleTy});
   }
 
+  mlir::Type indexMemoryType(mlir::Type eleTy) const override {
+    return mlir::IntegerType::get(eleTy.getContext(), S::defaultWidth);
+  }
+
   mlir::Type boxcharMemoryType(mlir::Type eleTy) const override {
     auto idxTy = mlir::IntegerType::get(eleTy.getContext(), S::defaultWidth);
     auto ptrTy = fir::ReferenceType::get(eleTy);

flang/lib/Optimizer/CodeGen/TargetRewrite.cpp

@@ -272,6 +272,12 @@ class TargetRewrite : public fir::impl::TargetRewritePassBase<TargetRewrite> {
       mlir::Value oper = std::get<1>(e.value());
       unsigned index = e.index();
       llvm::TypeSwitch<mlir::Type>(ty)
+          .template Case<mlir::IndexType>([&](mlir::IndexType indexTy) {
+            auto newTy = specifics->indexMemoryType(indexTy);
+            auto newOper = rewriter->create<mlir::arith::IndexCastOp>(loc, newTy, oper);
+            newInTys.push_back(newTy);
+            newOpers.push_back(newOper);
+          })
           .template Case<fir::BoxCharType>([&](fir::BoxCharType boxTy) {
             bool sret;
             if constexpr (std::is_same_v<std::decay_t<A>, fir::CallOp>) {
@@ -534,6 +540,7 @@ class TargetRewrite : public fir::impl::TargetRewritePassBase<TargetRewrite> {
     for (auto ty : func.getResults())
       if ((ty.isa<fir::BoxCharType>() && !noCharacterConversion) ||
           (fir::isa_complex(ty) && !noComplexConversion) ||
+          (ty.isa<mlir::IndexType>()) ||
           (ty.isa<mlir::IntegerType>() && hasCCallingConv)) {
         LLVM_DEBUG(llvm::dbgs() << "rewrite " << signature << " for target\n");
         return false;
@@ -542,6 +549,7 @@ class TargetRewrite : public fir::impl::TargetRewritePassBase<TargetRewrite> {
       if (((ty.isa<fir::BoxCharType>() || fir::isCharacterProcedureTuple(ty)) &&
            !noCharacterConversion) ||
           (fir::isa_complex(ty) && !noComplexConversion) ||
+          (ty.isa<mlir::IndexType>()) ||
           (ty.isa<mlir::IntegerType>() && hasCCallingConv)) {
         LLVM_DEBUG(llvm::dbgs() << "rewrite " << signature << " for target\n");
         return false;
@@ -624,6 +632,9 @@ class TargetRewrite : public fir::impl::TargetRewritePassBase<TargetRewrite> {
       auto ty = e.value();
       unsigned index = e.index();
       llvm::TypeSwitch<mlir::Type>(ty)
+          .Case<mlir::IndexType>([&](mlir::IndexType indexTy) {
+            newInTys.push_back(specifics->indexMemoryType(indexTy));
+          })
           .Case<fir::BoxCharType>([&](fir::BoxCharType boxTy) {
             if (noCharacterConversion) {
               newInTys.push_back(boxTy);