//===- GPUOpsLowering.cpp - GPU FuncOp / ReturnOp lowering ----------------===//
//
// 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
//
//===----------------------------------------------------------------------===//

#include "GPUOpsLowering.h"

#include "mlir/Conversion/GPUCommon/GPUCommonPass.h"
#include "mlir/Conversion/LLVMCommon/VectorPattern.h"
#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
#include "mlir/IR/Attributes.h"
#include "mlir/IR/Builders.h"
#include "mlir/IR/BuiltinTypes.h"
#include "mlir/IR/SymbolTable.h"
#include "llvm/ADT/SmallVectorExtras.h"
#include "llvm/ADT/StringSet.h"
#include "llvm/Support/DebugLog.h"
#include "llvm/Support/FormatVariadic.h"

#define DEBUG_TYPE "gpu-lowering"

using namespace mlir;

LLVM::LLVMFuncOp mlir::getOrDefineFunction(Operation *moduleOp, Location loc,
                                           OpBuilder &b, StringRef name,
                                           LLVM::LLVMFunctionType type) {
  auto existing = dyn_cast_or_null<LLVM::LLVMFuncOp>(
      SymbolTable::lookupSymbolIn(moduleOp, name));
  if (existing)
    return existing;

  OpBuilder::InsertionGuard guard(b);
  b.setInsertionPointToStart(&moduleOp->getRegion(0).front());
  return LLVM::LLVMFuncOp::create(b, loc, name, type, LLVM::Linkage::External);
}

static SmallString<16> getUniqueSymbolName(Operation *moduleOp,
                                           StringRef prefix) {
  // Get a unique global name.
  unsigned stringNumber = 0;
  SmallString<16> stringConstName;
  do {
    stringConstName.clear();
    (prefix + Twine(stringNumber++)).toStringRef(stringConstName);
  } while (SymbolTable::lookupSymbolIn(moduleOp, stringConstName));
  return stringConstName;
}

LLVM::GlobalOp mlir::getOrCreateStringConstant(OpBuilder &b, Location loc,
                                               Operation *moduleOp, Type llvmI8,
                                               StringRef namePrefix,
                                               StringRef str,
                                               uint64_t alignment,
                                               unsigned addrSpace) {
  llvm::SmallString<20> nullTermStr(str);
  nullTermStr.push_back('\0'); // Null terminate for C
  auto globalType =
      LLVM::LLVMArrayType::get(llvmI8, nullTermStr.size_in_bytes());
  StringAttr attr = b.getStringAttr(nullTermStr);

  // Try to find existing global.
  for (auto globalOp : moduleOp->getRegion(0).getOps<LLVM::GlobalOp>())
    if (globalOp.getGlobalType() == globalType && globalOp.getConstant() &&
        globalOp.getValueAttr() == attr &&
        globalOp.getAlignment().value_or(0) == alignment &&
        globalOp.getAddrSpace() == addrSpace)
      return globalOp;

  // Not found: create new global.
  OpBuilder::InsertionGuard guard(b);
  b.setInsertionPointToStart(&moduleOp->getRegion(0).front());
  SmallString<16> name = getUniqueSymbolName(moduleOp, namePrefix);
  return LLVM::GlobalOp::create(b, loc, globalType,
                                /*isConstant=*/true, LLVM::Linkage::Internal,
                                name, attr, alignment, addrSpace);
}

FailureOr<LoweredLLVMFuncAttrs> GPUFuncOpLowering::buildLoweredGPULLVMFuncAttrs(
    gpu::GPUFuncOp gpuFuncOp, Type llvmFuncType, OpBuilder &rewriter) const {
  FailureOr<LoweredLLVMFuncAttrs> loweredAttrs =
      lowerDiscardableAttrsForLLVMFunc(gpuFuncOp, llvmFuncType);
  if (failed(loweredAttrs))
    return failure();

  MLIRContext *ctx = rewriter.getContext();
  LLVM::LLVMFuncOp::Properties &props = loweredAttrs->properties;
  props.sym_name = rewriter.getStringAttr(gpuFuncOp.getName());
  props.function_type = TypeAttr::get(llvmFuncType);
  const bool isKernelFunc = gpuFuncOp.isKernel();
  props.setCConv(LLVM::CConvAttr::get(ctx, isKernelFunc
                                               ? kernelCallingConvention
                                               : nonKernelCallingConvention));

  NamedAttrList &discardable = loweredAttrs->discardableAttrs;
  auto *gpuDialect = cast<gpu::GPUDialect>(gpuFuncOp->getDialect());

  auto appendIfNameAndValue = [&](StringAttr name, Attribute value) {
    if (name && value)
      discardable.append(name, value);
  };

  DenseI32ArrayAttr knownBlockSize = gpuFuncOp.getKnownBlockSizeAttr();
  DenseI32ArrayAttr knownGridSize = gpuFuncOp.getKnownGridSizeAttr();
  DenseI32ArrayAttr knownClusterSize = gpuFuncOp.getKnownClusterSizeAttr();

  appendIfNameAndValue(gpuDialect->getKnownBlockSizeAttrHelper().getName(),
                       knownBlockSize);
  appendIfNameAndValue(gpuDialect->getKnownGridSizeAttrHelper().getName(),
                       knownGridSize);
  appendIfNameAndValue(gpuDialect->getKnownClusterSizeAttrHelper().getName(),
                       knownClusterSize);

  if (isKernelFunc) {
    discardable.append(gpuDialect->getKernelFuncAttrName(),
                       rewriter.getUnitAttr());
    // Add a dialect specific kernel attribute in addition to GPU kernel
    // attribute. The former is necessary for further translation while the
    // latter is expected by gpu.launch_func.
    appendIfNameAndValue(kernelAttributeName, rewriter.getUnitAttr());
    appendIfNameAndValue(kernelBlockSizeAttributeName, knownBlockSize);
    appendIfNameAndValue(kernelClusterSizeAttributeName, knownClusterSize);
  }

  return loweredAttrs;
}

LogicalResult
GPUFuncOpLowering::matchAndRewrite(gpu::GPUFuncOp gpuFuncOp, OpAdaptor adaptor,
                                   ConversionPatternRewriter &rewriter) const {
  Location loc = gpuFuncOp.getLoc();

  SmallVector<LLVM::GlobalOp, 3> workgroupBuffers;
  if (encodeWorkgroupAttributionsAsArguments) {
    // Append an `llvm.ptr` argument to the function signature to encode
    // workgroup attributions.

    ArrayRef<BlockArgument> workgroupAttributions =
        gpuFuncOp.getWorkgroupAttributionBBArgs();
    size_t numAttributions = workgroupAttributions.size();

    // Insert all arguments at the end.
    unsigned index = gpuFuncOp.getNumArguments();
    SmallVector<unsigned> argIndices(numAttributions, index);

    // New arguments will simply be `llvm.ptr` with the correct address space
    Type workgroupPtrType =
        rewriter.getType<LLVM::LLVMPointerType>(workgroupAddrSpace);
    Repeated<Type> argTypes(numAttributions, workgroupPtrType);

    // Attributes: noalias, llvm.mlir.workgroup_attribution(<size>, <type>)
    std::array attrs{
        rewriter.getNamedAttr(LLVM::LLVMDialect::getNoAliasAttrName(),
                              rewriter.getUnitAttr()),
        rewriter.getNamedAttr(
            getDialect().getWorkgroupAttributionAttrHelper().getName(),
            rewriter.getUnitAttr()),
    };
    SmallVector<DictionaryAttr> argAttrs;
    for (BlockArgument attribution : workgroupAttributions) {
      auto attributionType = cast<MemRefType>(attribution.getType());
      IntegerAttr numElements =
          rewriter.getI64IntegerAttr(attributionType.getNumElements());
      Type llvmElementType =
          getTypeConverter()->convertType(attributionType.getElementType());
      if (!llvmElementType)
        return failure();
      TypeAttr type = TypeAttr::get(llvmElementType);
      attrs.back().setValue(
          rewriter.getAttr<LLVM::WorkgroupAttributionAttr>(numElements, type));
      argAttrs.push_back(rewriter.getDictionaryAttr(attrs));
    }

    // Location match function location
    SmallVector<Location> argLocs(numAttributions, gpuFuncOp.getLoc());

    // Perform signature modification
    rewriter.modifyOpInPlace(
        gpuFuncOp, [gpuFuncOp, &argIndices, &argTypes, &argAttrs, &argLocs]() {
          LogicalResult inserted =
              static_cast<FunctionOpInterface>(gpuFuncOp).insertArguments(
                  argIndices, argTypes, argAttrs, argLocs);
          (void)inserted;
          assert(succeeded(inserted) &&
                 "expected GPU funcs to support inserting any argument");
        });
  } else {
    workgroupBuffers.reserve(gpuFuncOp.getNumWorkgroupAttributions());
    for (auto [idx, attribution] :
         llvm::enumerate(gpuFuncOp.getWorkgroupAttributionBBArgs())) {
      auto type = dyn_cast<MemRefType>(attribution.getType());
      assert(type && type.hasStaticShape() && "unexpected type in attribution");

      uint64_t numElements = type.getNumElements();

      auto elementType =
          cast<Type>(typeConverter->convertType(type.getElementType()));
      auto arrayType = LLVM::LLVMArrayType::get(elementType, numElements);
      std::string name =
          std::string(llvm::formatv("__wg_{0}_{1}", gpuFuncOp.getName(), idx));
      uint64_t alignment = 0;
      if (auto alignAttr = dyn_cast_or_null<IntegerAttr>(
              gpuFuncOp.getWorkgroupAttributionAttr(
                  idx, LLVM::LLVMDialect::getAlignAttrName())))
        alignment = alignAttr.getInt();
      auto globalOp = LLVM::GlobalOp::create(
          rewriter, gpuFuncOp.getLoc(), arrayType, /*isConstant=*/false,
          LLVM::Linkage::Internal, name, /*value=*/Attribute(), alignment,
          workgroupAddrSpace);
      workgroupBuffers.push_back(globalOp);
    }
  }

  // Remap proper input types.
  TypeConverter::SignatureConversion signatureConversion(
      gpuFuncOp.front().getNumArguments());

  Type funcType = getTypeConverter()->convertFunctionSignature(
      gpuFuncOp.getFunctionType(), /*isVariadic=*/false,
      getTypeConverter()->getOptions().useBarePtrCallConv, signatureConversion);
  if (!funcType) {
    return rewriter.notifyMatchFailure(gpuFuncOp, [&](Diagnostic &diag) {
      diag << "failed to convert function signature type for: "
           << gpuFuncOp.getFunctionType();
    });
  }

  ArrayAttr argAttrs = gpuFuncOp.getArgAttrsAttr();

  FailureOr<LoweredLLVMFuncAttrs> loweredAttrs =
      buildLoweredGPULLVMFuncAttrs(gpuFuncOp, funcType, rewriter);
  if (failed(loweredAttrs))
    return rewriter.notifyMatchFailure(gpuFuncOp,
                                       "failed to lower func attributes");

  auto llvmFuncOp = LLVM::LLVMFuncOp::create(rewriter, gpuFuncOp.getLoc(),
                                             loweredAttrs->properties,
                                             loweredAttrs->discardableAttrs);

  {
    // Insert operations that correspond to converted workgroup and private
    // memory attributions to the body of the function. This must operate on
    // the original function, before the body region is inlined in the new
    // function to maintain the relation between block arguments and the
    // parent operation that assigns their semantics.
    OpBuilder::InsertionGuard guard(rewriter);

    // Rewrite workgroup memory attributions to addresses of global buffers.
    rewriter.setInsertionPointToStart(&gpuFuncOp.front());
    unsigned numProperArguments = gpuFuncOp.getNumArguments();

    if (encodeWorkgroupAttributionsAsArguments) {
      // Build a MemRefDescriptor with each of the arguments added above.

      unsigned numAttributions = gpuFuncOp.getNumWorkgroupAttributions();
      assert(numProperArguments >= numAttributions &&
             "Expecting attributions to be encoded as arguments already");

      // Arguments encoding workgroup attributions will be in positions
      // [numProperArguments, numProperArguments+numAttributions)
      ArrayRef<BlockArgument> attributionArguments =
          gpuFuncOp.getArguments().slice(numProperArguments - numAttributions,
                                         numAttributions);
      for (auto [idx, vals] : llvm::enumerate(
               llvm::zip_equal(gpuFuncOp.getWorkgroupAttributionBBArgs(),
                               attributionArguments))) {
        auto [attribution, arg] = vals;
        auto type = cast<MemRefType>(attribution.getType());

        // Arguments are of llvm.ptr type and attributions are of memref type:
        // we need to wrap them in memref descriptors.
        Value descr = MemRefDescriptor::fromStaticShape(
            rewriter, loc, *getTypeConverter(), type, arg);

        // And remap the arguments
        signatureConversion.remapInput(numProperArguments + idx, descr);
      }
    } else {
      for (const auto [idx, global] : llvm::enumerate(workgroupBuffers)) {
        auto ptrType = LLVM::LLVMPointerType::get(rewriter.getContext(),
                                                  global.getAddrSpace());
        Value address = LLVM::AddressOfOp::create(rewriter, loc, ptrType,
                                                  global.getSymNameAttr());
        Value memory =
            LLVM::GEPOp::create(rewriter, loc, ptrType, global.getType(),
                                address, ArrayRef<LLVM::GEPArg>{0, 0});

        // Build a memref descriptor pointing to the buffer to plug with the
        // existing memref infrastructure. This may use more registers than
        // otherwise necessary given that memref sizes are fixed, but we can try
        // and canonicalize that away later.
        Value attribution = gpuFuncOp.getWorkgroupAttributionBBArgs()[idx];
        auto type = cast<MemRefType>(attribution.getType());
        Value descr = MemRefDescriptor::fromStaticShape(
            rewriter, loc, *getTypeConverter(), type, memory);
        signatureConversion.remapInput(numProperArguments + idx, descr);
      }
    }

    // Rewrite private memory attributions to alloca'ed buffers.
    unsigned numWorkgroupAttributions = gpuFuncOp.getNumWorkgroupAttributions();
    auto int64Ty = IntegerType::get(rewriter.getContext(), 64);
    for (const auto [idx, attribution] :
         llvm::enumerate(gpuFuncOp.getPrivateAttributions())) {
      auto type = cast<MemRefType>(attribution.getType());
      assert(type && type.hasStaticShape() && "unexpected type in attribution");

      // Explicitly drop memory space when lowering private memory
      // attributions since NVVM models it as `alloca`s in the default
      // memory space and does not support `alloca`s with addrspace(5).
      Type elementType = typeConverter->convertType(type.getElementType());
      auto ptrType =
          LLVM::LLVMPointerType::get(rewriter.getContext(), allocaAddrSpace);
      Value numElements = LLVM::ConstantOp::create(
          rewriter, gpuFuncOp.getLoc(), int64Ty, type.getNumElements());
      uint64_t alignment = 0;
      if (auto alignAttr =
              dyn_cast_or_null<IntegerAttr>(gpuFuncOp.getPrivateAttributionAttr(
                  idx, LLVM::LLVMDialect::getAlignAttrName())))
        alignment = alignAttr.getInt();
      Value allocated =
          LLVM::AllocaOp::create(rewriter, gpuFuncOp.getLoc(), ptrType,
                                 elementType, numElements, alignment);
      Value descr = MemRefDescriptor::fromStaticShape(
          rewriter, loc, *getTypeConverter(), type, allocated);
      signatureConversion.remapInput(
          numProperArguments + numWorkgroupAttributions + idx, descr);
    }
  }

  // Move the region to the new function, update the entry block signature.
  rewriter.inlineRegionBefore(gpuFuncOp.getBody(), llvmFuncOp.getBody(),
                              llvmFuncOp.end());
  if (failed(rewriter.convertRegionTypes(&llvmFuncOp.getBody(), *typeConverter,
                                         &signatureConversion)))
    return failure();

  // Get memref type from function arguments and set the noalias to
  // pointer arguments.
  for (const auto [idx, argTy] :
       llvm::enumerate(gpuFuncOp.getArgumentTypes())) {
    auto remapping = signatureConversion.getInputMapping(idx);
    NamedAttrList argAttr =
        argAttrs ? cast<DictionaryAttr>(argAttrs[idx]) : NamedAttrList();
    auto copyAttribute = [&](StringRef attrName) {
      Attribute attr = argAttr.erase(attrName);
      if (!attr)
        return;
      for (size_t i = 0, e = remapping->size; i < e; ++i)
        llvmFuncOp.setArgAttr(remapping->inputNo + i, attrName, attr);
    };
    auto copyPointerAttribute = [&](StringRef attrName) {
      Attribute attr = argAttr.erase(attrName);

      if (!attr)
        return;
      if (remapping->size > 1 &&
          attrName == LLVM::LLVMDialect::getNoAliasAttrName()) {
        emitWarning(llvmFuncOp.getLoc(),
                    "Cannot copy noalias with non-bare pointers.\n");
        return;
      }
      for (size_t i = 0, e = remapping->size; i < e; ++i) {
        if (isa<LLVM::LLVMPointerType>(
                llvmFuncOp.getArgument(remapping->inputNo + i).getType())) {
          llvmFuncOp.setArgAttr(remapping->inputNo + i, attrName, attr);
        }
      }
    };

    if (argAttr.empty())
      continue;

    copyAttribute(LLVM::LLVMDialect::getReturnedAttrName());
    copyAttribute(LLVM::LLVMDialect::getNoUndefAttrName());
    copyAttribute(LLVM::LLVMDialect::getInRegAttrName());
    bool lowersToPointer = false;
    for (size_t i = 0, e = remapping->size; i < e; ++i) {
      lowersToPointer |= isa<LLVM::LLVMPointerType>(
          llvmFuncOp.getArgument(remapping->inputNo + i).getType());
    }

    if (lowersToPointer) {
      copyPointerAttribute(LLVM::LLVMDialect::getNoAliasAttrName());
      copyPointerAttribute(LLVM::LLVMDialect::getNoCaptureAttrName());
      copyPointerAttribute(LLVM::LLVMDialect::getNoFreeAttrName());
      copyPointerAttribute(LLVM::LLVMDialect::getAlignAttrName());
      copyPointerAttribute(LLVM::LLVMDialect::getReadonlyAttrName());
      copyPointerAttribute(LLVM::LLVMDialect::getWriteOnlyAttrName());
      copyPointerAttribute(LLVM::LLVMDialect::getReadnoneAttrName());
      copyPointerAttribute(LLVM::LLVMDialect::getNonNullAttrName());
      copyPointerAttribute(LLVM::LLVMDialect::getDereferenceableAttrName());
      copyPointerAttribute(
          LLVM::LLVMDialect::getDereferenceableOrNullAttrName());
      copyPointerAttribute(
          LLVM::LLVMDialect::WorkgroupAttributionAttrHelper::getNameStr());
    }
  }
  rewriter.eraseOp(gpuFuncOp);
  return success();
}

LogicalResult GPUPrintfOpToHIPLowering::matchAndRewrite(
    gpu::PrintfOp gpuPrintfOp, gpu::PrintfOpAdaptor adaptor,
    ConversionPatternRewriter &rewriter) const {
  Location loc = gpuPrintfOp->getLoc();

  mlir::Type llvmI8 = typeConverter->convertType(rewriter.getI8Type());
  auto ptrType = LLVM::LLVMPointerType::get(rewriter.getContext());
  mlir::Type llvmI32 = typeConverter->convertType(rewriter.getI32Type());
  mlir::Type llvmI64 = typeConverter->convertType(rewriter.getI64Type());

  Operation *moduleOp = gpuPrintfOp->getParentWithTrait<OpTrait::SymbolTable>();
  if (!moduleOp)
    return rewriter.notifyMatchFailure(gpuPrintfOp,
                                       "Couldn't find a parent module");

  auto ocklBegin =
      getOrDefineFunction(moduleOp, loc, rewriter, "__ockl_printf_begin",
                          LLVM::LLVMFunctionType::get(llvmI64, {llvmI64}));
  LLVM::LLVMFuncOp ocklAppendArgs;
  if (!adaptor.getArgs().empty()) {
    ocklAppendArgs = getOrDefineFunction(
        moduleOp, loc, rewriter, "__ockl_printf_append_args",
        LLVM::LLVMFunctionType::get(
            llvmI64, {llvmI64, /*numArgs*/ llvmI32, llvmI64, llvmI64, llvmI64,
                      llvmI64, llvmI64, llvmI64, llvmI64, /*isLast*/ llvmI32}));
  }
  auto ocklAppendStringN = getOrDefineFunction(
      moduleOp, loc, rewriter, "__ockl_printf_append_string_n",
      LLVM::LLVMFunctionType::get(
          llvmI64,
          {llvmI64, ptrType, /*length (bytes)*/ llvmI64, /*isLast*/ llvmI32}));

  /// Start the printf hostcall
  Value zeroI64 = LLVM::ConstantOp::create(rewriter, loc, llvmI64, 0);
  auto printfBeginCall =
      LLVM::CallOp::create(rewriter, loc, ocklBegin, zeroI64);
  Value printfDesc = printfBeginCall.getResult();

  // Create the global op or find an existing one.
  LLVM::GlobalOp global = getOrCreateStringConstant(
      rewriter, loc, moduleOp, llvmI8, "printfFormat_", adaptor.getFormat());

  // Get a pointer to the format string's first element and pass it to printf()
  Value globalPtr = LLVM::AddressOfOp::create(
      rewriter, loc,
      LLVM::LLVMPointerType::get(rewriter.getContext(), global.getAddrSpace()),
      global.getSymNameAttr());
  Value stringStart =
      LLVM::GEPOp::create(rewriter, loc, ptrType, global.getGlobalType(),
                          globalPtr, ArrayRef<LLVM::GEPArg>{0, 0});
  Value stringLen = LLVM::ConstantOp::create(
      rewriter, loc, llvmI64, cast<StringAttr>(global.getValueAttr()).size());

  Value oneI32 = LLVM::ConstantOp::create(rewriter, loc, llvmI32, 1);
  Value zeroI32 = LLVM::ConstantOp::create(rewriter, loc, llvmI32, 0);

  auto appendFormatCall = LLVM::CallOp::create(
      rewriter, loc, ocklAppendStringN,
      ValueRange{printfDesc, stringStart, stringLen,
                 adaptor.getArgs().empty() ? oneI32 : zeroI32});
  printfDesc = appendFormatCall.getResult();

  // __ockl_printf_append_args takes 7 values per append call
  constexpr size_t argsPerAppend = 7;
  size_t nArgs = adaptor.getArgs().size();
  for (size_t group = 0; group < nArgs; group += argsPerAppend) {
    size_t bound = std::min(group + argsPerAppend, nArgs);
    size_t numArgsThisCall = bound - group;

    SmallVector<mlir::Value, 2 + argsPerAppend + 1> arguments;
    arguments.push_back(printfDesc);
    arguments.push_back(
        LLVM::ConstantOp::create(rewriter, loc, llvmI32, numArgsThisCall));
    for (size_t i = group; i < bound; ++i) {
      Value arg = adaptor.getArgs()[i];
      if (auto floatType = dyn_cast<FloatType>(arg.getType())) {
        if (!floatType.isF64())
          arg = LLVM::FPExtOp::create(
              rewriter, loc, typeConverter->convertType(rewriter.getF64Type()),
              arg);
        arg = LLVM::BitcastOp::create(rewriter, loc, llvmI64, arg);
      }
      if (arg.getType().getIntOrFloatBitWidth() != 64)
        arg = LLVM::ZExtOp::create(rewriter, loc, llvmI64, arg);

      arguments.push_back(arg);
    }
    // Pad out to 7 arguments since the hostcall always needs 7
    for (size_t extra = numArgsThisCall; extra < argsPerAppend; ++extra) {
      arguments.push_back(zeroI64);
    }

    auto isLast = (bound == nArgs) ? oneI32 : zeroI32;
    arguments.push_back(isLast);
    auto call = LLVM::CallOp::create(rewriter, loc, ocklAppendArgs, arguments);
    printfDesc = call.getResult();
  }
  rewriter.eraseOp(gpuPrintfOp);
  return success();
}

LogicalResult GPUPrintfOpToLLVMCallLowering::matchAndRewrite(
    gpu::PrintfOp gpuPrintfOp, gpu::PrintfOpAdaptor adaptor,
    ConversionPatternRewriter &rewriter) const {
  Location loc = gpuPrintfOp->getLoc();

  mlir::Type llvmI8 = typeConverter->convertType(rewriter.getIntegerType(8));
  mlir::Type ptrType =
      LLVM::LLVMPointerType::get(rewriter.getContext(), addressSpace);

  Operation *moduleOp = gpuPrintfOp->getParentWithTrait<OpTrait::SymbolTable>();
  if (!moduleOp)
    return rewriter.notifyMatchFailure(gpuPrintfOp,
                                       "Couldn't find a parent module");

  auto printfType =
      LLVM::LLVMFunctionType::get(rewriter.getI32Type(), {ptrType},
                                  /*isVarArg=*/true);
  LLVM::LLVMFuncOp printfDecl =
      getOrDefineFunction(moduleOp, loc, rewriter, funcName, printfType);
  printfDecl.setCConv(callingConvention);

  // Create the global op or find an existing one.
  LLVM::GlobalOp global = getOrCreateStringConstant(
      rewriter, loc, moduleOp, llvmI8, "printfFormat_", adaptor.getFormat(),
      /*alignment=*/0, addressSpace);

  // Get a pointer to the format string's first element
  Value globalPtr = LLVM::AddressOfOp::create(
      rewriter, loc,
      LLVM::LLVMPointerType::get(rewriter.getContext(), global.getAddrSpace()),
      global.getSymNameAttr());
  Value stringStart =
      LLVM::GEPOp::create(rewriter, loc, ptrType, global.getGlobalType(),
                          globalPtr, ArrayRef<LLVM::GEPArg>{0, 0});

  // Construct arguments and function call
  auto argsRange = adaptor.getArgs();
  SmallVector<Value, 4> printfArgs;
  printfArgs.reserve(argsRange.size() + 1);
  printfArgs.push_back(stringStart);
  printfArgs.append(argsRange.begin(), argsRange.end());

  auto call = LLVM::CallOp::create(rewriter, loc, printfDecl, printfArgs);
  call.setCConv(callingConvention);
  rewriter.eraseOp(gpuPrintfOp);
  return success();
}

LogicalResult GPUPrintfOpToVPrintfLowering::matchAndRewrite(
    gpu::PrintfOp gpuPrintfOp, gpu::PrintfOpAdaptor adaptor,
    ConversionPatternRewriter &rewriter) const {
  Location loc = gpuPrintfOp->getLoc();

  mlir::Type llvmI8 = typeConverter->convertType(rewriter.getIntegerType(8));
  mlir::Type ptrType = LLVM::LLVMPointerType::get(rewriter.getContext());

  Operation *moduleOp = gpuPrintfOp->getParentWithTrait<OpTrait::SymbolTable>();
  if (!moduleOp)
    return rewriter.notifyMatchFailure(gpuPrintfOp,
                                       "Couldn't find a parent module");

  // Create a valid global location removing any metadata attached to the
  // location as debug info metadata inside of a function cannot be used outside
  // of that function.
  Location globalLoc = loc->findInstanceOfOrUnknown<FileLineColLoc>();

  auto vprintfType =
      LLVM::LLVMFunctionType::get(rewriter.getI32Type(), {ptrType, ptrType});
  LLVM::LLVMFuncOp vprintfDecl = getOrDefineFunction(
      moduleOp, globalLoc, rewriter, "vprintf", vprintfType);

  // Create the global op or find an existing one.
  LLVM::GlobalOp global =
      getOrCreateStringConstant(rewriter, globalLoc, moduleOp, llvmI8,
                                "printfFormat_", adaptor.getFormat());

  // Get a pointer to the format string's first element
  Value globalPtr = LLVM::AddressOfOp::create(rewriter, loc, global);
  Value stringStart =
      LLVM::GEPOp::create(rewriter, loc, ptrType, global.getGlobalType(),
                          globalPtr, ArrayRef<LLVM::GEPArg>{0, 0});
  SmallVector<Type> types;
  SmallVector<Value> args;
  // Promote and pack the arguments into a stack allocation.
  for (Value arg : adaptor.getArgs()) {
    Type type = arg.getType();
    Value promotedArg = arg;
    assert(type.isIntOrFloat());
    if (isa<FloatType>(type)) {
      type = rewriter.getF64Type();
      promotedArg = LLVM::FPExtOp::create(rewriter, loc, type, arg);
    }
    types.push_back(type);
    args.push_back(promotedArg);
  }
  Type structType =
      LLVM::LLVMStructType::getLiteral(gpuPrintfOp.getContext(), types);
  Value one = LLVM::ConstantOp::create(rewriter, loc, rewriter.getI64Type(),
                                       rewriter.getIndexAttr(1));
  Value tempAlloc =
      LLVM::AllocaOp::create(rewriter, loc, ptrType, structType, one,
                             /*alignment=*/0);
  for (auto [index, arg] : llvm::enumerate(args)) {
    Value ptr = LLVM::GEPOp::create(
        rewriter, loc, ptrType, structType, tempAlloc,
        ArrayRef<LLVM::GEPArg>{0, static_cast<int32_t>(index)});
    LLVM::StoreOp::create(rewriter, loc, arg, ptr);
  }
  std::array<Value, 2> printfArgs = {stringStart, tempAlloc};

  LLVM::CallOp::create(rewriter, loc, vprintfDecl, printfArgs);
  rewriter.eraseOp(gpuPrintfOp);
  return success();
}

/// Helper for impl::scalarizeVectorOp. Scalarizes vectors to elements.
/// Used either directly (for ops on 1D vectors) or as the callback passed to
/// detail::handleMultidimensionalVectors (for ops on higher-rank vectors).
static Value scalarizeVectorOpHelper(Operation *op, ValueRange operands,
                                     Type llvm1DVectorTy,
                                     ConversionPatternRewriter &rewriter,
                                     const LLVMTypeConverter &converter) {
  TypeRange operandTypes(operands);
  VectorType vectorType = cast<VectorType>(llvm1DVectorTy);
  Location loc = op->getLoc();
  Value result = LLVM::PoisonOp::create(rewriter, loc, vectorType);
  Type indexType = converter.convertType(rewriter.getIndexType());
  StringAttr name = op->getName().getIdentifier();
  Type elementType = vectorType.getElementType();

  for (int64_t i = 0; i < vectorType.getNumElements(); ++i) {
    Value index = LLVM::ConstantOp::create(rewriter, loc, indexType, i);
    auto extractElement = [&](Value operand) -> Value {
      if (!isa<VectorType>(operand.getType()))
        return operand;
      return LLVM::ExtractElementOp::create(rewriter, loc, operand, index);
    };
    auto scalarOperands = llvm::map_to_vector(operands, extractElement);
    Operation *scalarOp =
        rewriter.create(loc, name, scalarOperands, elementType, op->getAttrs());
    result = LLVM::InsertElementOp::create(rewriter, loc, result,
                                           scalarOp->getResult(0), index);
  }
  return result;
}

/// Unrolls op to array/vector elements.
LogicalResult impl::scalarizeVectorOp(Operation *op, ValueRange operands,
                                      ConversionPatternRewriter &rewriter,
                                      const LLVMTypeConverter &converter) {
  TypeRange operandTypes(operands);
  if (llvm::any_of(operandTypes, llvm::IsaPred<VectorType>)) {
    VectorType vectorType =
        cast<VectorType>(converter.convertType(op->getResultTypes()[0]));
    rewriter.replaceOp(op, scalarizeVectorOpHelper(op, operands, vectorType,
                                                   rewriter, converter));
    return success();
  }

  if (llvm::any_of(operandTypes, llvm::IsaPred<LLVM::LLVMArrayType>)) {
    return LLVM::detail::handleMultidimensionalVectors(
        op, operands, converter,
        [&](Type llvm1DVectorTy, ValueRange operands) -> Value {
          return scalarizeVectorOpHelper(op, operands, llvm1DVectorTy, rewriter,
                                         converter);
        },
        rewriter);
  }

  return rewriter.notifyMatchFailure(op, "no llvm.array or vector to unroll");
}

static IntegerAttr wrapNumericMemorySpace(MLIRContext *ctx, unsigned space) {
  return IntegerAttr::get(IntegerType::get(ctx, 64), space);
}

/// Generates a symbol with 0-sized array type for dynamic shared memory usage,
/// or uses existing symbol.
static LLVM::GlobalOp getDynamicSharedMemorySymbol(
    ConversionPatternRewriter &rewriter, gpu::GPUModuleOp moduleOp,
    gpu::DynamicSharedMemoryOp op, const LLVMTypeConverter *typeConverter,
    MemRefType memrefType, unsigned alignmentBit) {
  uint64_t alignmentByte = alignmentBit / memrefType.getElementTypeBitWidth();

  FailureOr<unsigned> addressSpace =
      typeConverter->getMemRefAddressSpace(memrefType);
  if (failed(addressSpace)) {
    op->emitError() << "conversion of memref memory space "
                    << memrefType.getMemorySpace()
                    << " to integer address space "
                       "failed. Consider adding memory space conversions.";
  }

  // Step 1. Collect symbol names of LLVM::GlobalOp Ops. Also if any of
  // LLVM::GlobalOp is suitable for shared memory, return it.
  llvm::StringSet<> existingGlobalNames;
  for (auto globalOp : moduleOp.getBody()->getOps<LLVM::GlobalOp>()) {
    existingGlobalNames.insert(globalOp.getSymName());
    if (auto arrayType = dyn_cast<LLVM::LLVMArrayType>(globalOp.getType())) {
      if (globalOp.getAddrSpace() == addressSpace.value() &&
          arrayType.getNumElements() == 0 &&
          globalOp.getAlignment().value_or(0) == alignmentByte) {
        return globalOp;
      }
    }
  }

  // Step 2. Find a unique symbol name
  unsigned uniquingCounter = 0;
  SmallString<128> symName = SymbolTable::generateSymbolName<128>(
      "__dynamic_shmem_",
      [&](StringRef candidate) {
        return existingGlobalNames.contains(candidate);
      },
      uniquingCounter);

  // Step 3. Generate a global op
  OpBuilder::InsertionGuard guard(rewriter);
  rewriter.setInsertionPointToStart(moduleOp.getBody());

  auto zeroSizedArrayType = LLVM::LLVMArrayType::get(
      typeConverter->convertType(memrefType.getElementType()), 0);

  return LLVM::GlobalOp::create(rewriter, op->getLoc(), zeroSizedArrayType,
                                /*isConstant=*/false, LLVM::Linkage::Internal,
                                symName, /*value=*/Attribute(), alignmentByte,
                                addressSpace.value());
}

LogicalResult GPUDynamicSharedMemoryOpLowering::matchAndRewrite(
    gpu::DynamicSharedMemoryOp op, OpAdaptor adaptor,
    ConversionPatternRewriter &rewriter) const {
  Location loc = op.getLoc();
  MemRefType memrefType = op.getResultMemref().getType();
  Type elementType = typeConverter->convertType(memrefType.getElementType());

  // Step 1: Generate a memref<0xi8> type
  MemRefLayoutAttrInterface layout = {};
  auto memrefType0sz =
      MemRefType::get({0}, elementType, layout, memrefType.getMemorySpace());

  // Step 2: Generate a global symbol or existing for the dynamic shared
  // memory with memref<0xi8> type
  auto moduleOp = op->getParentOfType<gpu::GPUModuleOp>();
  LLVM::GlobalOp shmemOp = getDynamicSharedMemorySymbol(
      rewriter, moduleOp, op, getTypeConverter(), memrefType0sz, alignmentBit);

  // Step 3. Get address of the global symbol
  OpBuilder::InsertionGuard guard(rewriter);
  rewriter.setInsertionPoint(op);
  auto basePtr = LLVM::AddressOfOp::create(rewriter, loc, shmemOp);
  Type baseType = basePtr->getResultTypes().front();

  // Step 4. Generate GEP using offsets
  SmallVector<LLVM::GEPArg> gepArgs = {0};
  Value shmemPtr = LLVM::GEPOp::create(rewriter, loc, baseType, elementType,
                                       basePtr, gepArgs);
  // Step 5. Create a memref descriptor
  SmallVector<Value> shape, strides;
  Value sizeBytes;
  getMemRefDescriptorSizes(loc, memrefType0sz, {}, rewriter, shape, strides,
                           sizeBytes);
  auto memRefDescriptor = this->createMemRefDescriptor(
      loc, memrefType0sz, shmemPtr, shmemPtr, shape, strides, rewriter);

  // Step 5. Replace the op with memref descriptor
  rewriter.replaceOp(op, {memRefDescriptor});
  return success();
}

LogicalResult GPUReturnOpLowering::matchAndRewrite(
    gpu::ReturnOp op, OpAdaptor adaptor,
    ConversionPatternRewriter &rewriter) const {
  Location loc = op.getLoc();
  unsigned numArguments = op.getNumOperands();
  SmallVector<Value, 4> updatedOperands;

  bool useBarePtrCallConv = getTypeConverter()->getOptions().useBarePtrCallConv;
  if (useBarePtrCallConv) {
    // For the bare-ptr calling convention, extract the aligned pointer to
    // be returned from the memref descriptor.
    for (auto it : llvm::zip(op->getOperands(), adaptor.getOperands())) {
      Type oldTy = std::get<0>(it).getType();
      Value newOperand = std::get<1>(it);
      if (isa<MemRefType>(oldTy) && getTypeConverter()->canConvertToBarePtr(
                                        cast<BaseMemRefType>(oldTy))) {
        MemRefDescriptor memrefDesc(newOperand);
        newOperand = memrefDesc.allocatedPtr(rewriter, loc);
      } else if (isa<UnrankedMemRefType>(oldTy)) {
        // Unranked memref is not supported in the bare pointer calling
        // convention.
        return failure();
      }
      updatedOperands.push_back(newOperand);
    }
  } else {
    updatedOperands = llvm::to_vector<4>(adaptor.getOperands());
    (void)copyUnrankedDescriptors(rewriter, loc, op.getOperands().getTypes(),
                                  updatedOperands,
                                  /*toDynamic=*/true);
  }

  // If ReturnOp has 0 or 1 operand, create it and return immediately.
  if (numArguments <= 1) {
    rewriter.replaceOpWithNewOp<LLVM::ReturnOp>(
        op, TypeRange(), updatedOperands, op->getAttrs());
    return success();
  }

  // Otherwise, we need to pack the arguments into an LLVM struct type before
  // returning.
  auto packedType = getTypeConverter()->packFunctionResults(
      op.getOperandTypes(), useBarePtrCallConv);
  if (!packedType) {
    return rewriter.notifyMatchFailure(op, "could not convert result types");
  }

  Value packed = LLVM::PoisonOp::create(rewriter, loc, packedType);
  for (auto [idx, operand] : llvm::enumerate(updatedOperands)) {
    packed = LLVM::InsertValueOp::create(rewriter, loc, packed, operand, idx);
  }
  rewriter.replaceOpWithNewOp<LLVM::ReturnOp>(op, TypeRange(), packed,
                                              op->getAttrs());
  return success();
}

void mlir::populateGpuMemorySpaceAttributeConversions(
    TypeConverter &typeConverter, const MemorySpaceMapping &mapping) {
  typeConverter.addTypeAttributeConversion(
      [mapping](BaseMemRefType type, gpu::AddressSpaceAttr memorySpaceAttr) {
        gpu::AddressSpace memorySpace = memorySpaceAttr.getValue();
        unsigned addressSpace = mapping(memorySpace);
        return wrapNumericMemorySpace(memorySpaceAttr.getContext(),
                                      addressSpace);
      });
}