//===-- CodeGenOpenMP.cpp -------------------------------------------------===//
//
// 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
//
//===----------------------------------------------------------------------===//
//
// Coding style: https://mlir.llvm.org/getting_started/DeveloperGuide/
//
//===----------------------------------------------------------------------===//

#include "flang/Optimizer/CodeGen/CodeGenOpenMP.h"

#include "flang/Optimizer/Builder/FIRBuilder.h"
#include "flang/Optimizer/CodeGen/CodeGen.h"
#include "flang/Optimizer/Dialect/FIRType.h"
#include "flang/Optimizer/Dialect/Support/FIRContext.h"
#include "flang/Optimizer/Support/FatalError.h"
#include "flang/Optimizer/Support/Utils.h"
#include "mlir/Conversion/LLVMCommon/Pattern.h"
#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
#include "mlir/Dialect/OpenMP/OpenMPDialect.h"
#include "mlir/IR/PatternMatch.h"
#include "mlir/Transforms/DialectConversion.h"

using namespace fir;

#define DEBUG_TYPE "flang-codegen-openmp"

// fir::LLVMTypeConverter for converting to LLVM IR dialect types.
#include "flang/Optimizer/CodeGen/TypeConverter.h"

namespace {
/// A pattern that converts the region arguments in a single-region OpenMP
/// operation to the LLVM dialect. The body of the region is not modified and is
/// expected to either be processed by the conversion infrastructure or already
/// contain ops compatible with LLVM dialect types.
template <typename OpType>
class OpenMPFIROpConversion : public mlir::ConvertOpToLLVMPattern<OpType> {
public:
  explicit OpenMPFIROpConversion(const fir::LLVMTypeConverter &lowering)
      : mlir::ConvertOpToLLVMPattern<OpType>(lowering) {}

  const fir::LLVMTypeConverter &lowerTy() const {
    return *static_cast<const fir::LLVMTypeConverter *>(
        this->getTypeConverter());
  }
};

// FIR Op specific conversion for MapInfoOp that overwrites the default OpenMP
// Dialect lowering, this allows FIR specific lowering of types, required for
// descriptors of allocatables currently.
struct MapInfoOpConversion
    : public OpenMPFIROpConversion<mlir::omp::MapInfoOp> {
  using OpenMPFIROpConversion::OpenMPFIROpConversion;

  mlir::omp::MapBoundsOp
  createBoundsForCharString(mlir::ConversionPatternRewriter &rewriter,
                            unsigned int len, mlir::Location loc) const {
    mlir::Type i64Ty = rewriter.getIntegerType(64);
    auto lBound = mlir::LLVM::ConstantOp::create(rewriter, loc, i64Ty, 0);
    auto uBoundAndExt =
        mlir::LLVM::ConstantOp::create(rewriter, loc, i64Ty, len - 1);
    auto stride = mlir::LLVM::ConstantOp::create(rewriter, loc, i64Ty, 1);
    auto baseLb = mlir::LLVM::ConstantOp::create(rewriter, loc, i64Ty, 1);
    auto mapBoundType = rewriter.getType<mlir::omp::MapBoundsType>();
    return mlir::omp::MapBoundsOp::create(rewriter, loc, mapBoundType, lBound,
                                          uBoundAndExt, uBoundAndExt, stride,
                                          /*strideInBytes*/ false, baseLb);
  }

  llvm::LogicalResult
  matchAndRewrite(mlir::omp::MapInfoOp curOp, OpAdaptor adaptor,
                  mlir::ConversionPatternRewriter &rewriter) const override {
    const mlir::TypeConverter *converter = getTypeConverter();
    llvm::SmallVector<mlir::Type> resTypes;
    if (failed(converter->convertTypes(curOp->getResultTypes(), resTypes)))
      return mlir::failure();

    auto newProperties = curOp.getProperties();
    mlir::omp::MapBoundsOp mapBoundsOp;
    auto convertTypeAttr = [&](mlir::TypeAttr typeAttr) {
      mlir::Type newAttr;
      if (fir::isTypeWithDescriptor(typeAttr.getValue())) {
        newAttr = lowerTy().convertBoxTypeAsStruct(
            mlir::cast<fir::BaseBoxType>(typeAttr.getValue()));
      } else if (fir::isa_char_string(fir::unwrapSequenceType(
                     fir::unwrapPassByRefType(typeAttr.getValue()))) &&
                 !characterWithDynamicLen(
                     fir::unwrapPassByRefType(typeAttr.getValue()))) {
        // Characters with a LEN param are represented as strings
        // (array of characters), the lowering to LLVM dialect
        // doesn't generate bounds for these (and this is not
        // done at the initial lowering either) and there is
        // minor inconsistencies in the variable types we
        // create for the map without this step when converting
        // to the LLVM dialect.
        //
        // For example, given the types:
        //
        //  1) CHARACTER(LEN=16), dimension(:,:), allocatable :: char_arr
        //  2) CHARACTER(LEN=16), dimension(10,10) :: char_arr
        //
        // We get the FIR types (note for 1: we already peeled off the
        // dynamic extents from the type at this stage, but the conversion
        // to llvm dialect does that in any case, so the final result
        // is the same):
        //
        //  1) !fir.char<1,16>
        //  2) !fir.array<10x10x!fir.char<1,16>>
        //
        // Which are converted to the LLVM dialect types:
        //
        // 1) !llvm.array<16 x i8>
        // 2) llvm.array<10 x array<10 x array<16 x i8>>
        //
        // And in both cases, we are missing the innermost bounds for
        // the !fir.char<1,16> which is expanded into a 16 x i8 array
        // in the conversion to LLVM dialect.
        //
        // The problem with this is that we would like to treat these
        // cases identically and not have to create specialised
        // lowerings for either of these in the lowering to LLVM-IR
        // and treat them like any other array that passes through.
        //
        // To do so below, we generate an extra bound for the
        // innermost array (the char type/string) using the LEN
        // parameter of the character type. And we "canonicalize"
        // the type, stripping it down to the base element type,
        // which in this case is an i8. This effectively allows
        // the lowering to treat this as a 1-D array with multiple
        // bounds which it is capable of handling without any special
        // casing.
        // TODO: Handle dynamic LEN characters.
        if (auto ct = mlir::dyn_cast_or_null<fir::CharacterType>(
                fir::unwrapSequenceType(typeAttr.getValue()))) {
          newAttr = converter->convertType(
              fir::unwrapSequenceType(typeAttr.getValue()));
          if (auto type = mlir::dyn_cast<mlir::LLVM::LLVMArrayType>(newAttr))
            newAttr = type.getElementType();
          // We do not generate MapBoundsOps for the device pass, as
          // MapBoundsOps are not generated for the device pass, as
          // they're unused in the device lowering.
          auto offloadMod =
              llvm::dyn_cast_or_null<mlir::omp::OffloadModuleInterface>(
                  *curOp->getParentOfType<mlir::ModuleOp>());
          if (!offloadMod.getIsTargetDevice())
            mapBoundsOp = createBoundsForCharString(rewriter, ct.getLen(),
                                                    curOp.getLoc());
        } else {
          newAttr = converter->convertType(typeAttr.getValue());
        }
      } else {
        newAttr = converter->convertType(typeAttr.getValue());
      }
      return mlir::TypeAttr::get(newAttr);
    };
    newProperties.setVarPtrType(convertTypeAttr(curOp.getVarPtrTypeAttr()));
    if (mlir::TypeAttr varPtrPtrType = curOp.getVarPtrPtrTypeAttr())
      newProperties.setVarPtrPtrType(convertTypeAttr(varPtrPtrType));

    auto newOp = rewriter.replaceOpWithNewOp<mlir::omp::MapInfoOp>(
        curOp, resTypes, adaptor.getOperands(), newProperties,
        curOp->getDiscardableAttrDictionary().getValue());
    if (mapBoundsOp) {
      rewriter.startOpModification(newOp);
      newOp.getBoundsMutable().append(mlir::ValueRange{mapBoundsOp});
      rewriter.finalizeOpModification(newOp);
    }

    return mlir::success();
  }
};

// FIR op specific conversion for PrivateClauseOp that overwrites the default
// OpenMP Dialect lowering, this allows FIR-aware lowering of types, required
// for boxes because the OpenMP dialect conversion doesn't know anything about
// FIR types.
struct PrivateClauseOpConversion
    : public OpenMPFIROpConversion<mlir::omp::PrivateClauseOp> {
  using OpenMPFIROpConversion::OpenMPFIROpConversion;

  llvm::LogicalResult
  matchAndRewrite(mlir::omp::PrivateClauseOp curOp, OpAdaptor adaptor,
                  mlir::ConversionPatternRewriter &rewriter) const override {
    const fir::LLVMTypeConverter &converter = lowerTy();
    mlir::Type convertedAllocType;
    if (auto box = mlir::dyn_cast<fir::BaseBoxType>(curOp.getType())) {
      // In LLVM codegen fir.box<> == fir.ref<fir.box<>> == llvm.ptr
      // Here we really do want the actual structure
      if (box.isAssumedRank())
        TODO(curOp->getLoc(), "Privatize an assumed rank array");
      unsigned rank = 0;
      if (auto seqTy = mlir::dyn_cast<fir::SequenceType>(
              fir::unwrapRefType(box.getEleTy())))
        rank = seqTy.getShape().size();
      convertedAllocType = converter.convertBoxTypeAsStruct(box, rank);
    } else {
      convertedAllocType = converter.convertType(adaptor.getType());
    }
    if (!convertedAllocType)
      return mlir::failure();
    rewriter.startOpModification(curOp);
    curOp.setType(convertedAllocType);
    rewriter.finalizeOpModification(curOp);
    return mlir::success();
  }
};

// Convert FIR type to LLVM without turning fir.box<T> into memory
// reference.
static mlir::Type convertObjectType(const fir::LLVMTypeConverter &converter,
                                    mlir::Type firType) {
  if (auto boxTy = mlir::dyn_cast<fir::BaseBoxType>(firType))
    return converter.convertBoxTypeAsStruct(boxTy);
  return converter.convertType(firType);
}

// FIR Op specific conversion for TargetAllocMemOp
struct TargetAllocMemOpConversion
    : public OpenMPFIROpConversion<mlir::omp::TargetAllocMemOp> {
  using OpenMPFIROpConversion::OpenMPFIROpConversion;

  llvm::LogicalResult
  matchAndRewrite(mlir::omp::TargetAllocMemOp allocmemOp, OpAdaptor adaptor,
                  mlir::ConversionPatternRewriter &rewriter) const override {
    mlir::Type heapTy = allocmemOp.getAllocatedType();
    mlir::Location loc = allocmemOp.getLoc();
    auto ity = lowerTy().indexType();
    mlir::Type dataTy = fir::unwrapRefType(heapTy);
    mlir::Type llvmObjectTy = convertObjectType(lowerTy(), dataTy);
    if (fir::isRecordWithTypeParameters(fir::unwrapSequenceType(dataTy)))
      TODO(loc, "omp.target_allocmem codegen of derived type with length "
                "parameters");
    mlir::Value size = fir::computeElementDistance(
        loc, llvmObjectTy, ity, rewriter, lowerTy().getDataLayout());
    if (auto scaleSize = fir::genAllocationScaleSize(
            loc, allocmemOp.getInType(), ity, rewriter))
      size = mlir::LLVM::MulOp::create(rewriter, loc, ity, size, scaleSize);
    for (mlir::Value opnd : adaptor.getOperands().drop_front())
      size = mlir::LLVM::MulOp::create(
          rewriter, loc, ity, size,
          integerCast(lowerTy(), loc, rewriter, ity, opnd));
    auto mallocTyWidth = lowerTy().getIndexTypeBitwidth();
    auto mallocTy =
        mlir::IntegerType::get(rewriter.getContext(), mallocTyWidth);
    if (mallocTyWidth != ity.getIntOrFloatBitWidth())
      size = integerCast(lowerTy(), loc, rewriter, mallocTy, size);
    rewriter.modifyOpInPlace(allocmemOp, [&]() {
      allocmemOp.setInType(rewriter.getI8Type());
      allocmemOp.getTypeparamsMutable().clear();
      allocmemOp.getTypeparamsMutable().append(size);
    });
    return mlir::success();
  }
};

struct DeclareMapperOpConversion
    : public OpenMPFIROpConversion<mlir::omp::DeclareMapperOp> {
  using OpenMPFIROpConversion::OpenMPFIROpConversion;

  llvm::LogicalResult
  matchAndRewrite(mlir::omp::DeclareMapperOp curOp, OpAdaptor adaptor,
                  mlir::ConversionPatternRewriter &rewriter) const override {
    rewriter.startOpModification(curOp);
    curOp.setType(convertObjectType(lowerTy(), curOp.getType()));
    rewriter.finalizeOpModification(curOp);
    return mlir::success();
  }
};

} // namespace

void fir::populateOpenMPFIRToLLVMConversionPatterns(
    const LLVMTypeConverter &converter, mlir::RewritePatternSet &patterns) {
  patterns.add<MapInfoOpConversion>(converter);
  patterns.add<PrivateClauseOpConversion>(converter);
  patterns.add<TargetAllocMemOpConversion>(converter);
  patterns.add<DeclareMapperOpConversion>(converter);
}