Factory.h

//===-- Optimizer/Builder/Factory.h -----------------------------*- C++ -*-===//
//
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
//===----------------------------------------------------------------------===//
//
// Templates to generate more complex code patterns in transformation passes.
// In transformation passes, front-end information such as is available in
// lowering is not available.
//
//===----------------------------------------------------------------------===//
 
#ifndef FORTRAN_OPTIMIZER_BUILDER_FACTORY_H
#define FORTRAN_OPTIMIZER_BUILDER_FACTORY_H
 
#include "flang/Optimizer/Dialect/FIROps.h"
#include "flang/Optimizer/Dialect/FIRType.h"
 
namespace mlir {
class Location;
class Value;
} // namespace mlir
 
namespace fir::factory {
 
constexpr llvm::StringRef attrFortranArrayOffsets() {
  return "Fortran.offsets";
}

template <typename B>
void genCharacterCopy(mlir::Value src, mlir::Value srcLen, mlir::Value dst,
                      mlir::Value dstLen, B &builder, mlir::Location loc) {
  auto srcTy =
      mlir::cast<fir::CharacterType>(fir::dyn_cast_ptrEleTy(src.getType()));
  auto dstTy =
      mlir::cast<fir::CharacterType>(fir::dyn_cast_ptrEleTy(dst.getType()));
  if (!srcLen && !dstLen && srcTy.getFKind() == dstTy.getFKind() &&
      srcTy.getLen() == dstTy.getLen()) {
    // same size, so just use load and store
    auto load = fir::LoadOp::create(builder, loc, src);
    fir::StoreOp::create(builder, loc, load, dst);
    return;
  }
  auto zero = mlir::arith::ConstantIndexOp::create(builder, loc, 0);
  auto one = mlir::arith::ConstantIndexOp::create(builder, loc, 1);
  auto toArrayTy = [&](fir::CharacterType ty) {
    return fir::ReferenceType::get(fir::SequenceType::get(
        fir::SequenceType::ShapeRef{fir::SequenceType::getUnknownExtent()},
        fir::CharacterType::getSingleton(ty.getContext(), ty.getFKind())));
  };
  auto toEleTy = [&](fir::ReferenceType ty) {
    auto seqTy = mlir::cast<fir::SequenceType>(ty.getEleTy());
    return mlir::cast<fir::CharacterType>(seqTy.getEleTy());
  };
  auto toCoorTy = [&](fir::ReferenceType ty) {
    return fir::ReferenceType::get(toEleTy(ty));
  };
  if (!srcLen && !dstLen && srcTy.getLen() >= dstTy.getLen()) {
    auto upper =
        mlir::arith::ConstantIndexOp::create(builder, loc, dstTy.getLen() - 1);
    auto loop = fir::DoLoopOp::create(builder, loc, zero, upper, one);
    auto insPt = builder.saveInsertionPoint();
    builder.setInsertionPointToStart(loop.getBody());
    auto csrcTy = toArrayTy(srcTy);
    auto csrc = fir::ConvertOp::create(builder, loc, csrcTy, src);
    auto in = fir::CoordinateOp::create(builder, loc, toCoorTy(csrcTy), csrc,
                                        loop.getInductionVar());
    auto load = fir::LoadOp::create(builder, loc, in);
    auto cdstTy = toArrayTy(dstTy);
    auto cdst = fir::ConvertOp::create(builder, loc, cdstTy, dst);
    auto out = fir::CoordinateOp::create(builder, loc, toCoorTy(cdstTy), cdst,
                                         loop.getInductionVar());
    mlir::Value cast =
        srcTy.getFKind() == dstTy.getFKind()
            ? load.getResult()
            : fir::ConvertOp::create(builder, loc, toEleTy(cdstTy), load)
                  .getResult();
    fir::StoreOp::create(builder, loc, cast, out);
    builder.restoreInsertionPoint(insPt);
    return;
  }
  auto minusOne = [&](mlir::Value v) -> mlir::Value {
    return mlir::arith::SubIOp::create(
        builder, loc, fir::ConvertOp::create(builder, loc, one.getType(), v),
        one);
  };
  mlir::Value len = dstLen ? minusOne(dstLen)
                           : mlir::arith::ConstantIndexOp::create(
                                 builder, loc, dstTy.getLen() - 1)
                                 .getResult();
  auto loop = fir::DoLoopOp::create(builder, loc, zero, len, one);
  auto insPt = builder.saveInsertionPoint();
  builder.setInsertionPointToStart(loop.getBody());
  mlir::Value slen =
      srcLen
          ? fir::ConvertOp::create(builder, loc, one.getType(), srcLen)
                .getResult()
          : mlir::arith::ConstantIndexOp::create(builder, loc, srcTy.getLen())
                .getResult();
  auto cond =
      mlir::arith::CmpIOp::create(builder, loc, mlir::arith::CmpIPredicate::slt,
                                  loop.getInductionVar(), slen);
  auto ifOp = fir::IfOp::create(builder, loc, cond, /*withElse=*/true);
  builder.setInsertionPointToStart(&ifOp.getThenRegion().front());
  auto csrcTy = toArrayTy(srcTy);
  auto csrc = fir::ConvertOp::create(builder, loc, csrcTy, src);
  auto in = fir::CoordinateOp::create(builder, loc, toCoorTy(csrcTy), csrc,
                                      loop.getInductionVar());
  auto load = fir::LoadOp::create(builder, loc, in);
  auto cdstTy = toArrayTy(dstTy);
  auto cdst = fir::ConvertOp::create(builder, loc, cdstTy, dst);
  auto out = fir::CoordinateOp::create(builder, loc, toCoorTy(cdstTy), cdst,
                                       loop.getInductionVar());
  mlir::Value cast =
      srcTy.getFKind() == dstTy.getFKind()
          ? load.getResult()
          : fir::ConvertOp::create(builder, loc, toEleTy(cdstTy), load)
                .getResult();
  fir::StoreOp::create(builder, loc, cast, out);
  builder.setInsertionPointToStart(&ifOp.getElseRegion().front());
  auto space = fir::StringLitOp::create(builder, loc, toEleTy(cdstTy),
                                        llvm::ArrayRef<char>{' '});
  auto cdst2 = fir::ConvertOp::create(builder, loc, cdstTy, dst);
  auto out2 = fir::CoordinateOp::create(builder, loc, toCoorTy(cdstTy), cdst2,
                                        loop.getInductionVar());
  fir::StoreOp::create(builder, loc, space, out2);
  builder.restoreInsertionPoint(insPt);
}

inline llvm::SmallVector<mlir::Value> getExtents(mlir::Value shapeVal) {
  if (shapeVal)
    if (auto *shapeOp = shapeVal.getDefiningOp()) {
      if (auto shOp = mlir::dyn_cast<fir::ShapeOp>(shapeOp)) {
        auto operands = shOp.getExtents();
        return {operands.begin(), operands.end()};
      }
      if (auto shOp = mlir::dyn_cast<fir::ShapeShiftOp>(shapeOp)) {
        auto operands = shOp.getExtents();
        return {operands.begin(), operands.end()};
      }
    }
  return {};
}

inline llvm::SmallVector<mlir::Value> getOrigins(mlir::Value shapeVal) {
  if (shapeVal)
    if (auto *shapeOp = shapeVal.getDefiningOp()) {
      if (auto shOp = mlir::dyn_cast<fir::ShapeShiftOp>(shapeOp)) {
        auto operands = shOp.getOrigins();
        return {operands.begin(), operands.end()};
      }
      if (auto shOp = mlir::dyn_cast<fir::ShiftOp>(shapeOp)) {
        auto operands = shOp.getOrigins();
        return {operands.begin(), operands.end()};
      }
    }
  return {};
}

template <typename B>
llvm::SmallVector<mlir::Value>
originateIndices(mlir::Location loc, B &builder, mlir::Type memTy,
                 mlir::Value shapeVal, mlir::ValueRange indices) {
  llvm::SmallVector<mlir::Value> result;
  auto origins = getOrigins(shapeVal);
  if (origins.empty()) {
    assert(!shapeVal || mlir::isa<fir::ShapeOp>(shapeVal.getDefiningOp()));
    auto ty = fir::dyn_cast_ptrOrBoxEleTy(memTy);
    assert(ty && mlir::isa<fir::SequenceType>(ty));
    auto seqTy = mlir::cast<fir::SequenceType>(ty);
    auto one = mlir::arith::ConstantIndexOp::create(builder, loc, 1);
    const auto dimension = seqTy.getDimension();
    if (shapeVal) {
      assert(dimension == mlir::cast<fir::ShapeOp>(shapeVal.getDefiningOp())
                              .getType()
                              .getRank());
    }
    for (auto i : llvm::enumerate(indices)) {
      if (i.index() < dimension) {
        assert(fir::isa_integer(i.value().getType()));
        result.push_back(
            mlir::arith::AddIOp::create(builder, loc, i.value(), one));
      } else {
        result.push_back(i.value());
      }
    }
    return result;
  }
  const auto dimension = origins.size();
  unsigned origOff = 0;
  for (auto i : llvm::enumerate(indices)) {
    if (i.index() < dimension)
      result.push_back(mlir::arith::AddIOp::create(builder, loc, i.value(),
                                                   origins[origOff++]));
    else
      result.push_back(i.value());
  }
  return result;
}
 
} // namespace fir::factory
 
#endif // FORTRAN_OPTIMIZER_BUILDER_FACTORY_H