/* Copyright 2021 The JAX Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/

#include <fcntl.h>
#include <spawn.h>
#include <sys/wait.h>
#include <unistd.h>

#include <algorithm>
#include <array>
#include <cstddef>
#include <cstdint>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <functional>
#include <memory>
#include <optional>
#include <string>
#include <string_view>
#include <system_error>  // NOLINT
#include <utility>
#include <variant>
#include <vector>

#include "absl/algorithm/container.h"
#include "absl/base/call_once.h"
#include "absl/base/no_destructor.h"
#include "absl/base/nullability.h"
#include "absl/base/thread_annotations.h"
#include "absl/cleanup/cleanup.h"
#include "absl/container/flat_hash_map.h"
#include "absl/container/flat_hash_set.h"
#include "absl/functional/function_ref.h"
#include "absl/log/check.h"
#include "absl/log/log.h"
#include "absl/status/status.h"
#include "absl/status/status_macros.h"
#include "absl/status/statusor.h"
#include "absl/strings/match.h"
#include "absl/strings/numbers.h"
#include "absl/strings/str_cat.h"
#include "absl/strings/str_format.h"
#include "absl/strings/str_join.h"
#include "absl/strings/str_split.h"
#include "absl/strings/string_view.h"
#include "absl/synchronization/mutex.h"
#include "third_party/gpus/cuda/include/cuda.h"
#include "third_party/gpus/cuda/include/driver_types.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/ADT/StringRef.h"
#include "llvm/ExecutionEngine/Orc/Core.h"
#include "llvm/ExecutionEngine/Orc/ExecutionUtils.h"
#include "llvm/ExecutionEngine/Orc/JITTargetMachineBuilder.h"
#include "llvm/ExecutionEngine/Orc/LLJIT.h"
#include "llvm/ExecutionEngine/Orc/RTDyldObjectLinkingLayer.h"
#include "llvm/ExecutionEngine/SectionMemoryManager.h"
#include "llvm/IR/DataLayout.h"
#include "llvm/IR/LegacyPassManager.h"
#include "llvm/MC/TargetRegistry.h"
#include "llvm/Support/CodeGen.h"
#include "llvm/Support/Debug.h"
#include "llvm/Support/DynamicLibrary.h"
#include "llvm/Support/Error.h"
#include "llvm/Support/FileSystem.h"
#include "llvm/Support/MemoryBuffer.h"
#include "llvm/Support/MemoryBufferRef.h"
#include "llvm/Support/TargetSelect.h"
#include "llvm/Support/raw_ostream.h"
#include "llvm/Target/TargetMachine.h"
#include "llvm/TargetParser/Host.h"
#include "llvm/TargetParser/Triple.h"
#include "mlir/Conversion/ArithToLLVM/ArithToLLVM.h"
#include "mlir/Conversion/ComplexToLLVM/ComplexToLLVM.h"
#include "mlir/Conversion/ControlFlowToLLVM/ControlFlowToLLVM.h"
#include "mlir/Conversion/FuncToLLVM/ConvertFuncToLLVM.h"
#include "mlir/Conversion/IndexToLLVM/IndexToLLVM.h"
#include "mlir/Conversion/MathToLLVM/MathToLLVM.h"
#include "mlir/Conversion/MemRefToLLVM/MemRefToLLVM.h"
#include "mlir/Conversion/NVVMToLLVM/NVVMToLLVM.h"
#include "mlir/Conversion/Passes.h"
#include "mlir/Conversion/UBToLLVM/UBToLLVM.h"
#include "mlir/Conversion/VectorToLLVM/ConvertVectorToLLVM.h"
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/Arith/Transforms/Passes.h"
#include "mlir/Dialect/ControlFlow/IR/ControlFlow.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Dialect/GPU/IR/GPUDialect.h"
#include "mlir/Dialect/GPU/Transforms/Passes.h"
#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
#include "mlir/Dialect/LLVMIR/NVVMDialect.h"
#include "mlir/Dialect/LLVMIR/Transforms/Passes.h"
#include "mlir/Dialect/Math/IR/Math.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/Dialect/MemRef/Transforms/Passes.h"
#include "mlir/Dialect/NVGPU/IR/NVGPUDialect.h"
#include "mlir/Dialect/SCF/IR/SCF.h"
#include "mlir/Dialect/Vector/IR/VectorOps.h"
#include "mlir/ExecutionEngine/ExecutionEngine.h"
#include "mlir/ExecutionEngine/OptUtils.h"
#include "mlir/IR/AsmState.h"
#include "mlir/IR/BuiltinAttributes.h"
#include "mlir/IR/BuiltinOps.h"
#include "mlir/IR/Diagnostics.h"
#include "mlir/IR/DialectRegistry.h"
#include "mlir/IR/MLIRContext.h"
#include "mlir/IR/OperationSupport.h"
#include "mlir/Parser/Parser.h"
#include "mlir/Pass/PassManager.h"
#include "mlir/Pass/PassRegistry.h"
#include "mlir/Support/LLVM.h"
#include "mlir/Support/LogicalResult.h"
#include "mlir/Target/LLVM/NVVM/Target.h"
#include "mlir/Target/LLVMIR/Dialect/Builtin/BuiltinToLLVMIRTranslation.h"
#include "mlir/Target/LLVMIR/Dialect/GPU/GPUToLLVMIRTranslation.h"
#include "mlir/Target/LLVMIR/Dialect/LLVMIR/LLVMToLLVMIRTranslation.h"
#include "mlir/Target/LLVMIR/Dialect/NVVM/NVVMToLLVMIRTranslation.h"
#include "mlir/Target/LLVMIR/Export.h"
#include "mlir/Transforms/Passes.h"
#include "jaxlib/gpu/vendor.h"
#include "jaxlib/mosaic/dialect/gpu/mosaic_gpu.h"
#include "jaxlib/mosaic/gpu/assembly_to_binary.h"
#include "jaxlib/mosaic/gpu/dump.h"
#include "jaxlib/mosaic/gpu/gpu_module_to_assembly.h"
#include "jaxlib/mosaic/gpu/launch_config.h"
#include "jaxlib/mosaic/gpu/launch_lowering.h"
#include "jaxlib/mosaic/gpu/mosaic_gpu.pb.h"
#include "jaxlib/mosaic/gpu/nvshmem.h"
#include "jaxlib/mosaic/gpu/passes.h"
#include "jaxlib/mosaic/gpu/serde.h"
#include "jaxlib/mosaic/gpu/target.h"
#include "xla/backends/cpu/target_machine_options.h"
#include "xla/backends/gpu/collectives/gpu_clique_key.h"
#include "xla/backends/gpu/collectives/gpu_collectives.h"
#include "xla/backends/gpu/collectives/gpu_communicator.h"
#include "xla/backends/gpu/ffi.h"
#include "xla/backends/gpu/runtime/collective_clique_requests.h"
#include "xla/backends/gpu/runtime/collective_cliques.h"
#include "xla/backends/gpu/runtime/collective_execution.h"
#include "xla/backends/gpu/runtime/collective_memory.h"
#include "xla/backends/gpu/runtime/collective_memory_requests.h"
#include "xla/backends/gpu/runtime/collective_params.h"
#include "xla/core/collectives/rank_id.h"
#include "xla/core/collectives/symmetric_memory.h"
#include "xla/executable_run_options.h"
#include "xla/ffi/api/record_api.h"
#include "xla/ffi/api/record_c_api.h"
#include "xla/ffi/ffi.h"
#include "xla/ffi/record_ffi.h"
#include "xla/ffi/type_registry.h"
#include "xla/service/gpu/llvm_gpu_backend/nvptx_libdevice_path.h"
#include "xla/service/llvm_ir/llvm_command_line_options.h"
#include "xla/status_macros.h"
#include "xla/stream_executor/cuda/compilation_provider.h"
#include "xla/stream_executor/cuda/cuda_compute_capability.h"
#include "xla/stream_executor/device_address_allocator.h"
#include "xla/stream_executor/device_description.h"
#include "xla/stream_executor/gpu/collective_kernel_metadata.h"
#include "xla/util.h"
#include "tsl/platform/path.h"
#include "tsl/profiler/lib/traceme.h"

namespace {

using ::mosaic::gpu::NvshmemApi;

namespace ffi = xla::ffi;
namespace se = ::stream_executor;
// Emitted host function that fills/builds the kernel spec but does not launch.
using MosaicHostFunc = void(mosaic::gpu::MosaicKernelSpec*, void**);
using KernelHash = std::array<uint64_t, 4>;

// Returns the latest PTX ISA version supported by both LLVM and the underlying
// PTX compiler.
absl::StatusOr<std::string> GetPtxIsaVersion(
    const se::cuda::CompilationProvider& compilation_provider) {
  ASSIGN_OR_RETURN(int ptxas_latest_version,
                   compilation_provider.GetLatestPtxIsaVersion());
  // We'd like to target the latest PTX ISA version supported by
  // ptxas. However, it doesn't make sense to ask LLVM to target a PTX
  // ISA that it isn't aware of yet. Find the latest version supported
  // by LLVM and return the minimum of the two versions, one from
  // ptxas and the other from LLVM.
  ASSIGN_OR_RETURN(int llvm_latest_version,
                   mosaic::gpu::GetLatestLlvmPtxIsaVersion());
  int final_version = std::min(ptxas_latest_version, llvm_latest_version);
  return absl::StrFormat("ptx%d", final_version);
}

void EnsureLLVMisInitialized() {
  static absl::once_flag init_flag;
  absl::call_once(init_flag, []() {
    llvm::InitializeNativeTarget();
    llvm::InitializeNativeTargetAsmPrinter();
    // Initialize X86 and AArch64 targets for cross-compilation.
    LLVMInitializeX86TargetInfo();
    LLVMInitializeX86Target();
    LLVMInitializeX86TargetMC();
    LLVMInitializeX86AsmPrinter();
    LLVMInitializeAArch64TargetInfo();
    LLVMInitializeAArch64Target();
    LLVMInitializeAArch64TargetMC();
    LLVMInitializeAArch64AsmPrinter();
  });
}

auto AppendDiagnosticTo(std::string* diagnostic) {
  return [diagnostic](mlir::Diagnostic& diag) {
    absl::StrAppend(diagnostic, diag.str(), "\n");
    return mlir::LogicalResult::failure();
  };
}

absl::StatusOr<mlir::OpPassManager> GetPassPipeline(
    mlir::MLIRContext* ctx,
    const se::cuda::CompilationProvider* compilation_provider,
    const se::CudaComputeCapability& cc, const std::string& sm,
    const std::string& ptx_isa, const std::string& nvshmem_path,
    bool verify_target) {
  static absl::once_flag register_passes_flag;
  absl::call_once(register_passes_flag, [&compilation_provider, &cc]() {
    mosaic::gpu::EnsureLLVMNVPTXTargetIsRegistered();
    EnsureLLVMisInitialized();
    mlir::registerCanonicalizerPass();
    mlir::registerCSEPass();
    mlir::registerStripDebugInfoPass();
    mlir::registerConvertNVGPUToNVVMPass();
    mlir::registerConvertVectorToSCF();
    mlir::registerSCFToControlFlowPass();
    mlir::registerConvertNVVMToLLVMPass();
    mlir::registerArithToLLVMConversionPass();
    mlir::registerConvertIndexToLLVMPass();
    mlir::registerConvertGpuOpsToNVVMOps();
    mlir::registerConvertMathToLLVMPass();
    mlir::registerConvertFuncToLLVMPass();
    mlir::registerLowerAffinePass();
    mlir::registerReconcileUnrealizedCastsPass();
    // TODO(apaszke): Only register the passes we actually use.
    mlir::memref::registerMemRefPasses();
    mlir::registerConvertToLLVMPass();
    mlir::registerGPUPasses();
    mlir::registerGpuLaunchSinkIndexComputationsPass();
    mosaic::gpu::registerGpuModuleToAssemblyPass();
    mosaic::gpu::registerAssemblyToBinaryPass(compilation_provider, cc);
    mosaic::gpu::registerGpuLaunchLoweringPass();
    mosaic::gpu::registerConvertGpuToLLVMPass();
    mosaic::gpu::registerByvalInsertionPass();
    mosaic::gpu::registerLLVMAttrInsertionPass();
    mosaic::gpu::registerResolveTrivialLocationsPass();
    mosaic::gpu::registerGpuSinkMemRefDescriptorsPass();
    mlir::arith::registerArithExpandOpsPass();
    mlir::LLVM::registerDIScopeForLLVMFuncOpPass();
  });

  std::vector<std::string> libraries_to_link{
      ::xla::gpu::nvptx::LibDevicePath(mosaic::gpu::kDefaultCudaDataDir)};
  if (!nvshmem_path.empty()) {
    libraries_to_link.push_back(nvshmem_path);
  }
  std::string error;
  llvm::raw_string_ostream error_stream(error);
  auto passes = mlir::parsePassPipeline(
      absl::StrFormat(
          R"(
        builtin.module(
          mosaic-gpu-resolve-trivial-locations,
          arith-expand,
          canonicalize,
          mosaic-gpu-sink-memref-descriptors,
          convert-nvgpu-to-nvvm,
          gpu-kernel-outlining{data-layout-str=},
          convert-vector-to-scf{full-unroll=false lower-tensors=false target-rank=1},
          convert-scf-to-cf,
          convert-nvvm-to-llvm,
          expand-strided-metadata,
          nvvm-attach-target{O=3 chip=%1$s fast=false features=+%2$s ftz=false  module= triple=nvptx64-nvidia-cuda verify-target-arch=%4$v},
          lower-affine,
          convert-arith-to-llvm{index-bitwidth=0},
          convert-index-to-llvm{index-bitwidth=64},
          canonicalize{max-iterations=10 max-num-rewrites=-1 region-simplify=normal test-convergence=false top-down=true},
          cse,
          gpu.module(convert-gpu-to-nvvm{has-redux=false index-bitwidth=64 use-bare-ptr-memref-call-conv=false}),
          gpu.module(canonicalize{max-iterations=10 max-num-rewrites=-1 region-simplify=normal test-convergence=false top-down=true}),
          gpu.module(cse),
          gpu.module(mosaic-byval-insertion),
          gpu.module(mosaic-llvm-attr-insertion),
          mosaic-convert-gpu-to-llvm,
          gpu.module(reconcile-unrealized-casts),
          ensure-debug-info-scope-on-llvm-func{emission-kind=DebugDirectivesOnly},
          convert-math-to-llvm{approximate-log1p=true},
          mosaic-gpu-module-to-assembly{libraries-to-link=%3$s},
          canonicalize{max-iterations=10 max-num-rewrites=-1 region-simplify=normal test-convergence=false top-down=true},
          cse,
          mosaic-gpu-assembly-to-binary,
          gpu-launch-lowering,
          convert-to-llvm,
          reconcile-unrealized-casts
        )
      )",
          sm, ptx_isa, absl::StrJoin(libraries_to_link, ","), verify_target),
      error_stream);
  if (mlir::failed(passes)) {
    return absl::InternalError(
        absl::StrCat("Failed to construct Mosaic GPU pass pipeline: ", error));
  }
  return std::move(*passes);
}

absl::Status RunPasses(mlir::OpPassManager&& passes,
                       mlir::ModuleOp module,
                       const mosaic::gpu::DumpOptions& dump_opts) {
  mlir::PassManager pm(module.getContext());
  pm.enableVerifier(false);
  *static_cast<mlir::OpPassManager*>(&pm) = std::move(passes);
  std::string diagnostic;
  mlir::ScopedDiagnosticHandler diagnostic_handler(
      module.getContext(), AppendDiagnosticTo(&diagnostic));
  std::optional<llvm::raw_fd_ostream> dump_stream;
  if (dump_opts.mlir_passes) {
    if (!dump_opts.dump_path.empty()) {
      std::string path = tsl::io::JoinPath(
          dump_opts.dump_path, dump_opts.module_basename + ".mlir-passes.log");
      std::error_code error;
      dump_stream.emplace(path, error, llvm::sys::fs::OF_None);
      if (error) {
        dump_stream.reset();
        LOG(ERROR) << error.message();
        LOG(ERROR) << "Output will be written to stdout instead.";
        dump_stream = std::nullopt;
      }
    }
    pm.getContext()->disableMultithreading();
    auto print_always = [](mlir::Pass*, mlir::Operation*) { return true; };
    pm.enableIRPrinting(
        /*shouldPrintBeforePass=*/print_always,
        /*shouldPrintAfterPass=*/print_always,
        /*printModuleScope=*/true,
        /*printAfterOnlyOnChange=*/false,
        /*printAfterOnlyOnFailure=*/true,
        dump_stream.has_value() ? *dump_stream : llvm::outs(),
        mlir::OpPrintingFlags().enableDebugInfo().printNameLocAsPrefix());
  }
  if (mlir::failed(pm.run(module))) {
    return absl::InternalError(
        absl::StrCat("Mosaic GPU pass pipeline failed: ", diagnostic));
  }
  return absl::OkStatus();
}

void InitContext(mlir::MLIRContext* context) {
  mlir::DialectRegistry registry;
  registry.insert<mlir::arith::ArithDialect, mlir::cf::ControlFlowDialect,
                  mlir::func::FuncDialect, mlir::math::MathDialect,
                  mlir::memref::MemRefDialect, mlir::scf::SCFDialect,
                  mlir::vector::VectorDialect, mlir::gpu::GPUDialect,
                  mlir::LLVM::LLVMDialect, mosaic_gpu::MosaicGPUDialect>();
  mlir::registerConvertComplexToLLVMInterface(registry);
  mlir::registerConvertMemRefToLLVMInterface(registry);
  mlir::registerConvertMathToLLVMInterface(registry);
  mlir::registerConvertFuncToLLVMInterface(registry);
  mlir::vector::registerConvertVectorToLLVMInterface(registry);
  mlir::index::registerConvertIndexToLLVMInterface(registry);
  mlir::cf::registerConvertControlFlowToLLVMInterface(registry);
  mlir::ub::registerConvertUBToLLVMInterface(registry);
  mlir::arith::registerConvertArithToLLVMInterface(registry);
  mlir::registerConvertMemRefToLLVMInterface(registry);
  mlir::gpu::registerOffloadingLLVMTranslationInterfaceExternalModels(registry);
  mlir::registerBuiltinDialectTranslation(registry);
  mlir::registerGPUDialectTranslation(registry);
  mlir::registerLLVMDialectTranslation(registry);
  context->appendDialectRegistry(registry);
  context->loadAllAvailableDialects();
}

void LoadNvDialects(mlir::MLIRContext* context) {
  mlir::DialectRegistry registry;
  // NVGPUDialect depends on NVVMDialect, so both must be loaded here after
  // parsing to prevent NVVM from eagerly parsing legacy bytecode attributes.
  registry.insert<mlir::nvgpu::NVGPUDialect, mlir::NVVM::NVVMDialect>();
  mlir::registerConvertNVVMToLLVMInterface(registry);
  mlir::NVVM::registerNVVMTargetInterfaceExternalModels(registry);
  mlir::registerNVVMDialectTranslation(registry);
  context->appendDialectRegistry(registry);
  context->loadAllAvailableDialects();
}

bool is_nvshmem_used(mlir::ModuleOp module) {
  constexpr std::string_view prefix1 = "nvshmem_";
  constexpr std::string_view prefix2 = "nvshmemx_";
  for (mlir::LLVM::LLVMFuncOp llvm_func :
       module.getOps<mlir::LLVM::LLVMFuncOp>()) {
    const auto& func_name = llvm_func.getName();
    if (!func_name.starts_with(prefix1) && !func_name.starts_with(prefix2)) {
      continue;
    }
    auto uses =
        mlir::SymbolTable::getSymbolUses(llvm_func, module.getOperation());
    if (uses && !uses->empty()) {
      return true;
    }
  }
  return false;
}

bool is_multimem_used(mlir::ModuleOp mod) {
  return static_cast<bool>(mod->getAttr("mosaic_gpu.multimem_used"));
}

absl::StatusOr<std::string> get_nvshmem_llvm_lib_path() {
  const char* nvshmem_path_ptr = getenv("MOSAIC_GPU_NVSHMEM_BC_PATH");
  if (!nvshmem_path_ptr)
    return absl::InternalError("Failed to get MOSAIC_GPU_NVSHMEM_BC_PATH");
  return nvshmem_path_ptr;
}

std::string CUDAErrorString(CUresult result) {
  const char* error;
  cuGetErrorString(result, &error);
  return error;
}
// Returns if the CUDA expression returns an error.
#define CUDA_RETURN_IF_ERROR(stmt)                         \
  do {                                                     \
    if (CUresult result = stmt; result != CUDA_SUCCESS) {  \
      return absl::InternalError(CUDAErrorString(result)); \
    }                                                      \
  } while (0)

absl::StatusOr<se::CudaComputeCapability> GetCudaComputeCapability() {
  // Assumes driver has been initialized and a context exists. XLA already has
  // some utilities to query this, but we try to stay runtime-agnostic, so we
  // build our own here.
  CUdevice device;
  CUDA_RETURN_IF_ERROR(cuCtxGetDevice(&device));
  int major = 0;
  CUDA_RETURN_IF_ERROR(cuDeviceGetAttribute(
      &major, CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MAJOR, device));
  int minor = 0;
  CUDA_RETURN_IF_ERROR(cuDeviceGetAttribute(
      &minor, CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MINOR, device));

  ASSIGN_OR_RETURN(std::string sm, mosaic::gpu::GetSmVersion(major, minor));
  bool has_accelerated_features = absl::EndsWith(sm, "a");

  using FeatureExtension = se::CudaComputeCapability::FeatureExtension;
  return se::CudaComputeCapability(major, minor,
                                   has_accelerated_features
                                       ? FeatureExtension::kAcceleratedFeatures
                                       : FeatureExtension::kNone);
}

absl::StatusOr<std::string> GetHostFuncName(mlir::ModuleOp module_op) {
  // We look for top level host C-interface function "_mlir_ciface_<foo>"
  constexpr std::string_view prefix = "_mlir_ciface_";
  std::vector<std::string> names;
  for (mlir::LLVM::LLVMFuncOp llvm_func :
       module_op.getOps<mlir::LLVM::LLVMFuncOp>()) {
    if (llvm_func.getName().starts_with(prefix)) {
      names.push_back(llvm_func.getName().str());
    }
  }
  if (auto size = names.size(); size != 1) {
    return absl::InternalError(absl::StrFormat(
        "Expected to locate exactly 1 symbol with %s prefix in the MLIR "
        "module, found %d instead.",
        prefix, size));
  }
  return names[0];
}

struct CompiledKernel {
  CompiledKernel(std::unique_ptr<llvm::orc::LLJIT> lljit,
                 MosaicHostFunc* host_func, bool is_nvshmem_used,
                 bool is_multimem_used, std::string object_file,
                 std::string host_func_name, std::string gpu_binary,
                 std::string kernel_name, int32_t smem_bytes,
                 int32_t cluster_size)
      : lljit(std::move(lljit)),
        host_func(host_func),
        is_nvshmem_used(is_nvshmem_used),
        is_multimem_used(is_multimem_used),
        gpu_binary(std::move(gpu_binary)),
        kernel_name(std::move(kernel_name)),
        smem_bytes(smem_bytes),
        cluster_size(cluster_size),
        object_file(std::move(object_file)),
        host_func_name(std::move(host_func_name)) {}

  // CompiledKernel is neither copyable nor movable. We use CompiledKernel* as a
  // key in a cache, so we require pointer stability.
  CompiledKernel(const CompiledKernel&) = delete;
  CompiledKernel(CompiledKernel&& other) = delete;

  std::unique_ptr<llvm::orc::LLJIT> lljit;
  const MosaicHostFunc* host_func;
  bool is_nvshmem_used = false;
  bool is_multimem_used = false;
  std::string gpu_binary;
  std::string kernel_name;
  int32_t smem_bytes = 0;
  int32_t cluster_size = 1;
  // The following fields are used for de/serialization of CompiledKernel.
  std::string object_file;
  std::string host_func_name;
};

absl::Status CUresultAsStatus(CUresult launch_result,
                              const CompiledKernel& kernel,
                              int32_t device_ordinal) {
  if (launch_result != CUDA_SUCCESS) {
    const char* error_name = nullptr;
    const char* error_str = nullptr;
    cuGetErrorName(launch_result, &error_name);
    cuGetErrorString(launch_result, &error_str);
    return absl::InternalError(absl::StrFormat(
        "Mosaic GPU kernel '%s' (host func: '%s', device: %d, smem: %d bytes, "
        "cluster: %d) launch failed: %s [%s] (CUDA error %d)",
        kernel.kernel_name, kernel.host_func_name, device_ordinal,
        kernel.smem_bytes, kernel.cluster_size,
        error_name ? error_name : "UNKNOWN_ERROR",
        error_str ? error_str : "unknown", launch_result));
  }
  return absl::OkStatus();
}

absl::Status MosaicGpuLaunchKernel(CUfunction function, CUstream stream,
                                   mosaic::gpu::MosaicKernelSpec* cfg,
                                   const CompiledKernel& kernel,
                                   int32_t device_ordinal) {
  CUlaunchConfig config{
      .gridDimX = cfg->grid.x,
      .gridDimY = cfg->grid.y,
      .gridDimZ = cfg->grid.z,
      .blockDimX = cfg->block.x,
      .blockDimY = cfg->block.y,
      .blockDimZ = cfg->block.z,
      .sharedMemBytes = cfg->smem_bytes,
      .hStream = stream,
      .attrs = nullptr,
      .numAttrs = 0,
  };
  CUlaunchAttribute attrs[2];
  int32_t num_attrs = 0;
  if (cfg->cluster.x != 0) {
    attrs[num_attrs].id = CU_LAUNCH_ATTRIBUTE_CLUSTER_DIMENSION;
    attrs[num_attrs].value.clusterDim = {
        .x = cfg->cluster.x,
        .y = cfg->cluster.y,
        .z = cfg->cluster.z,
    };
    num_attrs++;
  }
  if (cfg->uses_pdl) {
    attrs[num_attrs].id = CU_LAUNCH_ATTRIBUTE_PROGRAMMATIC_STREAM_SERIALIZATION;
    attrs[num_attrs].value.programmaticStreamSerializationAllowed = 1;
    num_attrs++;
  }
  if (num_attrs > 0) {
    config.attrs = attrs;
    config.numAttrs = num_attrs;
  }
  CUresult result =
      cuLaunchKernelEx(&config, function, cfg->kernel_params().data(), nullptr);
  if (result == CUDA_ERROR_INVALID_CLUSTER_SIZE) {
    int max_cluster_size;
    if (cuOccupancyMaxPotentialClusterSize(&max_cluster_size, function,
                                           &config) == CUDA_SUCCESS) {
      fprintf(stderr,
              "cuLaunchKernel failed with invalid cluster size (%d, %d, %d)"
              ": maximum is %d\n",
              cfg->cluster.x, cfg->cluster.y, cfg->cluster.z, max_cluster_size);
    }
  }
  return CUresultAsStatus(result, kernel, device_ordinal);
}

absl::Status RunMlirPasses(mlir::ModuleOp module, se::CudaComputeCapability cc,
                           bool is_nvshmem_used,
                           const mosaic::gpu::DumpOptions& dump_opts) {
  ASSIGN_OR_RETURN(se::cuda::CompilationProvider * compilation_provider,
                   mosaic::gpu::GetAssemblyToBinaryCompilationProvider());
  ASSIGN_OR_RETURN(std::string sm,
                   mosaic::gpu::GetSmVersion(cc.major, cc.minor));
  // Here, it is important to use a PTX ISA version that is supported by both
  // the underlying compilation provider, and by LLVM. Using a version that is
  // newer than what LLVM supports will lead to the indication being ignored by
  // LLVM (potentially causing a version downgrade), while using a version that
  // is newer than what the compilation provider supports will lead to LLVM
  // potentially generating PTX that the compilation provider cannot handle.
  ASSIGN_OR_RETURN(std::string llvm_ptx_isa,
                   GetPtxIsaVersion(*compilation_provider));
  std::string nvshmem_path = "";
  if (is_nvshmem_used) {
    ASSIGN_OR_RETURN(nvshmem_path, get_nvshmem_llvm_lib_path());
  }
  // nvbug/5809460: spurious LLVM/MLIR errors with tcgen05+sm_103a; disable
  // verification on sm_103a, sm_110a etc. where we see spurious failures.
  bool verify_target = !((cc.major == 10 && cc.minor > 0) || cc.major == 11);
  ASSIGN_OR_RETURN(
      auto passes,
      GetPassPipeline(module.getContext(), compilation_provider, cc, sm,
                      llvm_ptx_isa, nvshmem_path, verify_target));
  return RunPasses(std::move(passes), module, dump_opts);
}

// This function was inspired by mlir::ExecutionEngine::create. It takes an MLIR
// module, compiles the module to LLVM IR, optimizes it using LLVM, and returns
// the compiled object file.
absl::StatusOr<std::unique_ptr<llvm::MemoryBuffer>> CompileModuleToObject(
    mlir::ModuleOp module, const mosaic::gpu::DumpOptions& dump_opts,
    const xla::cpu::TargetMachineOptions& cpu_target_machine_options) {
  llvm::LLVMContext llvm_context;
  auto llvm_module = mlir::translateModuleToLLVMIR(module, llvm_context);
  if (!llvm_module) {
    return absl::InternalError("Failed to translate module to LLVM IR");
  }

  // Create the target machine builder using the provided
  // cpu_target_machine_options.
  llvm::orc::JITTargetMachineBuilder tm_builder(
      llvm::Triple(cpu_target_machine_options.triple()));
  tm_builder.setCPU(cpu_target_machine_options.cpu());
  tm_builder.addFeatures(
      cpu_target_machine_options.GetTargetMachineFeaturesVector());

  tm_builder.setCodeGenOptLevel(llvm::CodeGenOptLevel::Aggressive);

  auto tm_or_error = tm_builder.createTargetMachine();
  if (!tm_or_error) {
    return absl::InternalError(
        absl::StrFormat("Failed to create target machine: %s",
                        llvm::toString(tm_or_error.takeError())));
  }
  std::unique_ptr<llvm::TargetMachine> target_machine =
      std::move(tm_or_error.get());

  mlir::ExecutionEngine::setupTargetTripleAndDataLayout(llvm_module.get(),
                                                        target_machine.get());

  if (getenv("MOSAIC_GPU_DUMP_HOST_LLVM")) {
    std::string ll_str;
    llvm::raw_string_ostream os(ll_str);
    llvm_module->print(os, nullptr);
    os.flush();
    mosaic::gpu::DumpToFileOrStdout(ll_str, dump_opts.module_basename + ".ll",
                                    dump_opts.dump_path);
  }

  auto transformer = mlir::makeOptimizingTransformer(
      /*optLevel=*/3, /*sizeLevel=*/0, /*targetMachine=*/target_machine.get());
  if (auto err = transformer(llvm_module.get())) {
    return absl::InternalError(absl::StrFormat("Transformer failed: %s",
                                               llvm::toString(std::move(err))));
  }

  std::string result;
  llvm::raw_string_ostream stream(result);
  {
    llvm::buffer_ostream bstream(stream);
    llvm::legacy::PassManager codegen_pm;
    if (target_machine->addPassesToEmitFile(
            codegen_pm, bstream, nullptr, llvm::CodeGenFileType::ObjectFile)) {
      return absl::InternalError(
          "TargetMachine can't emit a file of this type");
    }
    codegen_pm.run(*llvm_module);
  }

  return llvm::MemoryBuffer::getMemBufferCopy(result, "kernel");
}

absl::StatusOr<std::unique_ptr<CompiledKernel>> CreateAndInitJIT(
    std::unique_ptr<llvm::MemoryBuffer> object_file, std::string host_func_name,
    bool is_nvshmem_used, bool is_multimem_used, std::string gpu_binary,
    std::string kernel_name, int32_t smem_bytes, int32_t cluster_size) {
  EnsureLLVMisInitialized();
  std::string object_file_str = object_file->getBuffer().str();
  auto lljit_builder = llvm::orc::LLJITBuilder();

  auto tm_builder_or_error = llvm::orc::JITTargetMachineBuilder::detectHost();
  if (!tm_builder_or_error) {
    return absl::InternalError(
        absl::StrFormat("Failed to detect host: %s",
                        llvm::toString(tm_builder_or_error.takeError())));
  }
  char global_prefix =
      llvm::DataLayout(
          tm_builder_or_error->getTargetTriple().computeDataLayout())
          .getGlobalPrefix();

  lljit_builder.setObjectLinkingLayerCreator(
      [global_prefix](llvm::orc::ExecutionSession& session,
                      llvm::jitlink::JITLinkMemoryManager& /*MemMgr*/)
          -> llvm::Expected<std::unique_ptr<llvm::orc::ObjectLayer>> {
        auto objectLayer =
            std::make_unique<llvm::orc::RTDyldObjectLinkingLayer>(
                session, [](llvm::MemoryBufferRef) {
                  return std::make_unique<llvm::SectionMemoryManager>();
                });

        llvm::SmallVector<llvm::StringRef> runtime_libs;
        if (const char* runtime_lib_path =
                getenv("MOSAIC_GPU_RUNTIME_LIB_PATH")) {
          runtime_libs.emplace_back(runtime_lib_path);
        }
        if (const char* nvshmem_path = getenv("MOSAIC_GPU_NVSHMEM_SO_PATH")) {
          runtime_libs.emplace_back(nvshmem_path);
        }

        for (const auto& lib : runtime_libs) {
          llvm::SmallString<128> abs_lib_path(lib);
          if (std::error_code ec = llvm::sys::fs::make_absolute(abs_lib_path)) {
            return llvm::make_error<llvm::StringError>(
                absl::StrFormat("Failed to get absolute path for %s: %s",
                                std::string_view(lib.data(), lib.size()),
                                ec.message()),
                llvm::inconvertibleErrorCode());
          }

          // inspired by mlir::ExecutionEngine::create
          auto mb = llvm::MemoryBuffer::getFile(abs_lib_path);
          if (!mb) {
            return llvm::make_error<llvm::StringError>(
                absl::StrFormat("Failed to create MemoryBuffer for: %s: %s",
                                abs_lib_path.c_str(), mb.getError().message()),
                llvm::inconvertibleErrorCode());
          }

          auto& jd = session.createBareJITDylib(std::string(abs_lib_path));
          auto loaded = llvm::orc::DynamicLibrarySearchGenerator::Load(
              abs_lib_path.c_str(), global_prefix);
          if (!loaded) {
            return loaded.takeError();
          }
          jd.addGenerator(std::move(*loaded));

          if (auto err = objectLayer->add(jd, std::move(mb.get()))) {
            return std::move(err);
          }
        }

        return objectLayer;
      });

  auto lljit_or_error = lljit_builder.create();
  if (auto err = lljit_or_error.takeError()) {
    return absl::InternalError(absl::StrFormat("Failed to create LLJIT: %s",
                                               llvm::toString(std::move(err))));
  }
  std::unique_ptr<llvm::orc::LLJIT> lljit = std::move(*lljit_or_error);

  // Resolve symbols that are linked in the current process. This is necessary
  // to find the host and init functions and resolve all library calls.
  {
    auto& main_jd = lljit->getMainJITDylib();
    auto generator_or_error =
        llvm::orc::DynamicLibrarySearchGenerator::GetForCurrentProcess(
            global_prefix);
    if (auto err = generator_or_error.takeError()) {
      return absl::InternalError(absl::StrFormat(
          "Failed to create generator: %s", llvm::toString(std::move(err))));
    }
    main_jd.addGenerator(std::move(*generator_or_error));
  }

  // This is the actual linking step. Returns an error if it fails to resolve
  // all symbols from the object file.
  if (auto err = lljit->addObjectFile(std::move(object_file))) {
    return absl::InternalError(absl::StrFormat("Failed to add object file: %s",
                                               llvm::toString(std::move(err))));
  }

  // Run all the static initializers in the object file if there are any.
  if (auto err = lljit->initialize(lljit->getMainJITDylib())) {
    return absl::InternalError(absl::StrFormat("Failed to initialize LLJIT: %s",
                                               llvm::toString(std::move(err))));
  }

  auto host_sym = lljit->lookup(host_func_name);
  if (auto err = host_sym.takeError()) {
    return absl::InternalError(absl::StrFormat(
        "Failed to lookup host symbol: %s", llvm::toString(std::move(err))));
  }

  VLOG(5) << "Successfully JIT-linked Mosaic GPU kernel";
  MosaicHostFunc* host_func = host_sym->toPtr<MosaicHostFunc*>();
  return std::make_unique<CompiledKernel>(
      std::move(lljit), host_func, is_nvshmem_used, is_multimem_used,
      std::move(object_file_str), std::move(host_func_name),
      std::move(gpu_binary), std::move(kernel_name), smem_bytes, cluster_size);
}

absl::StatusOr<std::unique_ptr<CompiledKernel>> Compile(
    llvm::StringRef module_str, se::CudaComputeCapability cc,
    const xla::cpu::TargetMachineOptions& cpu_target_machine_options) {
  tsl::profiler::TraceMe trace("Compile");
  mlir::MLIRContext context(mlir::MLIRContext::Threading::DISABLED);
  context.allowUnregisteredDialects(true);
  InitContext(&context);
  mlir::ParserConfig parse_config(&context);
  // TODO(bchetioui): remove 6 months after IR version 9 is the forward
  // compatible IR version.
  std::string upgraded_module_str =
      mosaic::gpu::UpgradeLegacyAttributeSyntax(module_str);
  auto module = mlir::parseSourceString<mlir::ModuleOp>(upgraded_module_str,
                                                        parse_config);
  if (!module) {
    return absl::InternalError("Failed to parse Mosaic GPU module");
  }
  // NVVM is loaded after `InitContext` because `mlir::parseSourceString`
  // chokes on parsing NVVM attributes after
  // https://github.com/llvm/llvm-project/pull/196289 otherwise.
  //
  // TODO(bchetioui): remove dependency on NVVM to lighten the maintenance
  // burden.
  LoadNvDialects(&context);

  {
    std::string diagnostic;
    mlir::ScopedDiagnosticHandler diagnostic_handler(
        &context, AppendDiagnosticTo(&diagnostic));
    auto manager = mlir::PassManager::on<mlir::ModuleOp>(module->getContext());
    manager.addPass(mosaic::gpu::createSerdePass(
        mosaic::gpu::SerdePassOptions{.serialize = false}));
    if (manager.run(module.get()).failed()) {
      return absl::InternalError(absl::StrCat(
          "Failed to deserialize Mosaic GPU module: ", diagnostic));
    }
  }

  const char* dump_llvm = getenv("MOSAIC_GPU_DUMP_LLVM");
  const char* llvm_debug_only = getenv("MOSAIC_GPU_LLVM_DEBUG_ONLY");
#ifndef NDEBUG
  bool old_debug_state = false;
  std::vector<std::string_view> debug_only_types;
  if (llvm_debug_only) {
    debug_only_types = absl::StrSplit(llvm_debug_only, ',');
  }
  if (dump_llvm) {
    debug_only_types.push_back("serialize-to-llvm");
  }
  if (!debug_only_types.empty()) {
    old_debug_state = llvm::DebugFlag;
    std::vector<const char*> debug_only_types_ptrs;
    debug_only_types_ptrs.reserve(debug_only_types.size());
    for (std::string_view debug_only_type : debug_only_types) {
      debug_only_types_ptrs.push_back(debug_only_type.data());
    }
    llvm::setCurrentDebugTypes(debug_only_types_ptrs.data(),
                               debug_only_types_ptrs.size());
    llvm::DebugFlag = true;
  }
#else
  if (llvm_debug_only || dump_llvm) {
    fprintf(
        stderr,
        "MOSAIC_GPU_LLVM_DEBUG_ONLY or MOSAIC_GPU_DUMP_LLVM is set but LLVM "
        "was built with NDEBUG\n");
    abort();
  }
#endif

  // Use `div.full` for float32 division---this generates better SASS.
  const std::vector<std::string> llvm_cl_options{"-nvptx-prec-divf32=1"};
  // Acquire a lock over the LLVM command line options here. XLA uses this
  // lock to override the default LLVM command line options on a per-client
  // basis. This means that failing to acquire this lock and explicitly
  // setting our own command line options makes compilation dependent on
  // outside state/non-deterministic.
  xla::llvm_ir::LLVMCommandLineOptionsLock llvm_lock(llvm_cl_options);
  mosaic::gpu::EnsureLLVMNVPTXTargetIsRegistered();

  bool use_nvshmem = is_nvshmem_used(*module);
  bool multimem_used = is_multimem_used(*module);
  mosaic::gpu::DumpOptions dump_opts =
      mosaic::gpu::GetOrSetDumpOptionsForModule(*module);
  RETURN_IF_ERROR(RunMlirPasses(*module, cc, use_nvshmem, dump_opts));

  auto gpu_binary_attr =
      (*module)->getAttrOfType<mlir::StringAttr>("mosaic_gpu.gpu_binary");
  auto kernel_name_attr =
      (*module)->getAttrOfType<mlir::StringAttr>("mosaic_gpu.kernel_name");
  auto smem_bytes_attr =
      (*module)->getAttrOfType<mlir::IntegerAttr>("mosaic_gpu.smem_bytes");
  auto cluster_size_attr =
      (*module)->getAttrOfType<mlir::IntegerAttr>("mosaic_gpu.cluster_size");

  if (!gpu_binary_attr || !kernel_name_attr || !smem_bytes_attr ||
      !cluster_size_attr) {
    return absl::InternalError(
        "Failed to retrieve kernel metadata attributes after running MLIR "
        "passes");
  }

  std::string gpu_binary = gpu_binary_attr.getValue().str();
  std::string kernel_name = kernel_name_attr.getValue().str();
  int32_t smem_bytes = smem_bytes_attr.getInt();
  int32_t cluster_size = cluster_size_attr.getInt();

  ASSIGN_OR_RETURN(
      auto object_file,
      CompileModuleToObject(*module, dump_opts, cpu_target_machine_options));
  VLOG(5) << "Successfully compiled Mosaic GPU kernel to object file";

#ifndef NDEBUG
  if (llvm_debug_only || dump_llvm) {
    llvm::DebugFlag = old_debug_state;
  }
#endif

  ASSIGN_OR_RETURN(auto host_func_name, GetHostFuncName(*module));

  return CreateAndInitJIT(std::move(object_file), std::move(host_func_name),
                          use_nvshmem, multimem_used, std::move(gpu_binary),
                          std::move(kernel_name), smem_bytes, cluster_size);
}

struct KernelCache {
  absl::Mutex mutex;
  absl::flat_hash_map<KernelHash, std::unique_ptr<CompiledKernel>> kernels
      ABSL_GUARDED_BY(mutex);
};

static KernelCache& GetKernelCache() {
  static absl::NoDestructor<KernelCache> cache;
  return *cache;
}

absl::StatusOr<CompiledKernel*> GetOrCreateKernel(
    const KernelHash& kernel_hash,
    absl::FunctionRef<absl::StatusOr<std::unique_ptr<CompiledKernel>>()>
        factory) {
  auto& cache = GetKernelCache();
  {
    absl::MutexLock lock(cache.mutex);
    auto it = cache.kernels.find(kernel_hash);
    if (it != cache.kernels.end()) {
      return it->second.get();
    }
  }
  // Release the lock while compiling the kernel. It is possible that multiple
  // threads compile the same kernel concurrently. In that case, we will discard
  // all but the first result.
  ASSIGN_OR_RETURN(auto kernel, factory());
  absl::MutexLock lock(cache.mutex);
  auto [iter, _] = cache.kernels.try_emplace(kernel_hash, std::move(kernel));
  return iter->second.get();
}

class KernelHandle {
 public:
  KernelHandle(CUmodule module, CUfunction function, CUcontext ctx)
      : module_(module), function_(function), ctx_(ctx) {}
  ~KernelHandle() {
    if (auto s = Unload(module_, ctx_); !s.ok()) {
      LOG(ERROR) << "Failed to unload GPU module: " << s;
    } else {
      VLOG(5) << "Successfully unloaded GPU module";
    }
  }

  // KernelHandle is neither copyable nor movable.
  KernelHandle(const KernelHandle&) = delete;
  KernelHandle& operator=(const KernelHandle&) = delete;

  CUmodule module() const { return module_; }
  CUfunction function() const { return function_; }

 private:
  static absl::Status Unload(CUmodule module, CUcontext ctx) {
    CUDA_RETURN_IF_ERROR(cuCtxPushCurrent(ctx));
    CUDA_RETURN_IF_ERROR(cuModuleUnload(module));
    CUcontext unused;
    CUDA_RETURN_IF_ERROR(cuCtxPopCurrent(&unused));
    return absl::OkStatus();
  }

  CUmodule module_;
  CUfunction function_;
  CUcontext ctx_;
};

struct InitResult {
  CUmodule module;
  CUfunction function;
};

absl::StatusOr<InitResult> InitKernel(const CompiledKernel& kernel) {
  if (kernel.is_nvshmem_used &&
      !NvshmemApi::Default(/*assert_ok=*/false).is_loaded()) {
    return absl::InternalError(
        "Failed to load the NVSHMEM library. Make sure it is installed (e.g. "
        "`pip install nvidia-nvshmem-cu12`).");
  }
  if (kernel.is_multimem_used) {
    CUdevice device;
    CUDA_RETURN_IF_ERROR(cuCtxGetDevice(&device));
    int supports_multicast;
    CUDA_RETURN_IF_ERROR(cuDeviceGetAttribute(
        &supports_multicast, CU_DEVICE_ATTRIBUTE_MULTICAST_SUPPORTED, device));
    if (supports_multicast == 0) {
      return absl::FailedPreconditionError(
          "System does not support multicast memory; multimem instructions "
          "cannot be used.");
    }

    if (kernel.is_nvshmem_used) {
      int nvshmem_world_size = NvshmemApi::Default().n_pes();
      // multimem instructions require multicast memory; mgpu will emit
      // device-side calls to nvshmemx_mc_ptr to translate unicast->multicast
      // addresses, which will return nullptr and lead to memory errors if
      // multicast is not supported on the device
      // https://docs.nvidia.com/nvshmem/api/using.html#communication-model
      // https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#data-movement-and-conversion-instructions-multimem
      // https://docs.nvidia.com/cuda/cuda-programming-guide/04-special-topics/virtual-memory-management.html#multicast-memory-sharing
      if (nvshmem_world_size == 1) {
        // There is only one device, so NVSHMEM will not configure multicast
        // mappings and nvshmemx_mc_ptr will return nullptr.
        return absl::FailedPreconditionError(
            "Multicast memory mappings are not configured with only one "
            "device; multimem instructions cannot be used.");
      }
    }
  }
  CUmodule module = nullptr;
  CUDA_RETURN_IF_ERROR(cuModuleLoadData(&module, kernel.gpu_binary.data()));
  if (kernel.is_nvshmem_used) {
    if (NvshmemApi::Default().cumodule_init(module) != NVSHMEM_SUCCESS) {
      return absl::InternalError("nvshmemx_cumodule_init failed.");
    }
  }
  CUfunction function = nullptr;
  // TODO(allanrenucci): We should unload the kernel if any of the following
  // calls fail.
  CUDA_RETURN_IF_ERROR(
      cuModuleGetFunction(&function, module, kernel.kernel_name.c_str()));
  if (kernel.smem_bytes) {
    CUDA_RETURN_IF_ERROR(cuFuncSetAttribute(
        function, CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
        kernel.smem_bytes));
  }
  if (kernel.cluster_size > 8) {
    CUDA_RETURN_IF_ERROR(cuFuncSetAttribute(
        function, CU_FUNC_ATTRIBUTE_NON_PORTABLE_CLUSTER_SIZE_ALLOWED, 1));
  }
  VLOG(5) << "Successfully initialized Mosaic GPU kernel";
  return InitResult{module, function};
}

// Initializes the kernel in the current CUDA context and return a handle to the
// kernel.
absl::StatusOr<std::shared_ptr<KernelHandle>> CachedInit(
    const CompiledKernel* absl_nonnull kernel) {
  using CacheKey = std::pair<const CompiledKernel*, uintptr_t>;
  struct Cache {
    absl::Mutex mutex;
    absl::flat_hash_map<CacheKey, std::weak_ptr<KernelHandle>> handles
        ABSL_GUARDED_BY(mutex);
  };
  static absl::NoDestructor<Cache> cache;

  CUcontext ctx;
  CUDA_RETURN_IF_ERROR(cuCtxGetCurrent(&ctx));
  CacheKey key(kernel, reinterpret_cast<uintptr_t>(ctx));

  absl::MutexLock lock(cache->mutex);
  if (auto it = cache->handles.find(key); it != cache->handles.end()) {
    if (auto handle = it->second.lock()) {
      VLOG(5) << "Found Mosaic GPU kernel in cache";
      return handle;
    }
  }
  ASSIGN_OR_RETURN(InitResult res, InitKernel(*kernel));
  auto handle = std::make_shared<KernelHandle>(res.module, res.function, ctx);
  cache->handles.insert_or_assign(key, handle);
  return handle;
}

// Structure stores data needed during the execution and filled during the
// initialization.
struct DeviceState {
  // Serialized collective kernel metadata.
  // Structure has the following layout:
  // [CollectiveKernelMetadata][param_to_peers][multimem_addresses]
  // multimem_addresses are set only for modules that use multimem
  // instructions.
  // Note: the collective metadata param to peers and multimem addresses are
  // pointing to the nullptr and should not be used during the lowering.
  std::vector<std::byte> metadata_bytes;

  // The RAII handle of the buffer on the device which stores the structure
  // above.
  se::ScopedDeviceAddress<uint8_t> metadata_handle;

  // Pointer (CUmodule + CUfunction) to the kernel loaded on the GPU.
  std::shared_ptr<KernelHandle> kernel_handle;
};

constexpr int kMaxLocalDevices = 8;

}  // namespace

namespace mosaic::gpu {
struct CustomCallResources {
  CompiledKernel* kernel = nullptr;

  // For each participating local device store device-specific metadata for the
  // custom call operation.
  std::array<DeviceState, kMaxLocalDevices> device_states;
  KernelHash hash;

  static absl::StatusOr<std::string> Serialize(
      const CustomCallResources& resources);
  static absl::StatusOr<std::unique_ptr<CustomCallResources>> Deserialize(
      absl::string_view data);
};
absl::StatusOr<std::string> CustomCallResources::Serialize(
    const CustomCallResources& resources) {
  mosaic::gpu::MosaicGpuKernelProto kernel_proto;
  CompiledKernel* kernel = resources.kernel;
  if (kernel == nullptr) {
    return absl::InternalError(
        "Failed to serialize CustomCallResources: CompiledKernel is null");
  }
  kernel_proto.set_version(4);
  kernel_proto.set_object_file(kernel->object_file);
  kernel_proto.set_is_nvshmem_used(kernel->is_nvshmem_used);
  kernel_proto.set_is_multimem_used(kernel->is_multimem_used);
  kernel_proto.set_kernel_hash(resources.hash.data(), sizeof(KernelHash));
  kernel_proto.set_host_func_name(kernel->host_func_name);
  kernel_proto.set_gpu_binary(kernel->gpu_binary);
  kernel_proto.set_kernel_name(kernel->kernel_name);
  kernel_proto.set_smem_bytes(kernel->smem_bytes);
  kernel_proto.set_cluster_size(kernel->cluster_size);
  return kernel_proto.SerializeAsString();
}

absl::StatusOr<std::unique_ptr<CustomCallResources>>
CustomCallResources::Deserialize(absl::string_view data) {
  mosaic::gpu::MosaicGpuKernelProto kernel_proto;
  if (!kernel_proto.ParseFromString(data)) {
    return absl::InternalError("Failed to parse MosaicGpuKernel proto");
  }
  if (kernel_proto.is_nvshmem_used()) {
    return absl::UnimplementedError("NVSHMEM is not supported in XLA.");
  }
  auto resources = std::make_unique<CustomCallResources>();
  if (kernel_proto.kernel_hash().size() != sizeof(KernelHash)) {
    return absl::InternalError("Invalid kernel hash size in proto");
  }
  std::memcpy(resources->hash.data(), kernel_proto.kernel_hash().data(),
              sizeof(KernelHash));

  if (kernel_proto.version() != 4) {
    return absl::InternalError(absl::StrCat(
        "Unsupported Mosaic GPU kernel version: ", kernel_proto.version()));
  }

  std::string host_func_name = kernel_proto.host_func_name();
  std::string gpu_binary = kernel_proto.gpu_binary();
  std::string kernel_name = kernel_proto.kernel_name();
  int32_t smem_bytes = kernel_proto.smem_bytes();
  int32_t cluster_size = kernel_proto.cluster_size();

  ASSIGN_OR_RETURN(
      resources->kernel,
      GetOrCreateKernel(
          resources->hash,
          [&]() -> absl::StatusOr<std::unique_ptr<CompiledKernel>> {
            return CreateAndInitJIT(
                llvm::MemoryBuffer::getMemBuffer(kernel_proto.object_file(),
                                                 "kernel"),
                std::move(host_func_name), kernel_proto.is_nvshmem_used(),
                kernel_proto.is_multimem_used(), std::move(gpu_binary),
                std::move(kernel_name), smem_bytes, cluster_size);
          }));
  return resources;
}

}  // namespace mosaic::gpu

namespace {

using ::mosaic::gpu::CustomCallResources;

// Validate custom call attributes and compile the kernel.
absl::StatusOr<std::unique_ptr<CustomCallResources>> InstantiateResources(
    const se::GpuComputeCapability* cc,
    const xla::cpu::TargetMachineOptions* cpu_target_machine_options,
    ffi::Dictionary attrs) {
  ASSIGN_OR_RETURN(bool use_custom_barrier,
                   attrs.get<bool>("use_custom_barrier"));
  ASSIGN_OR_RETURN(std::string_view kernel_hash,
                   attrs.get<std::string_view>("kernel_hash"));
  ASSIGN_OR_RETURN(std::string_view module,
                   attrs.get<std::string_view>("module"));
  if (use_custom_barrier) {
    return absl::UnimplementedError("Custom barrier is not supported on GPUs.");
  }
  if (kernel_hash.size() != sizeof(KernelHash)) {
    return absl::InvalidArgumentError(
        absl::StrFormat("Kernel hash size is %d bytes, expected %d bytes",
                        kernel_hash.size(), sizeof(KernelHash)));
  }
  KernelHash hash;
  std::memcpy(hash.data(), kernel_hash.data(), sizeof(KernelHash));
  CHECK(cc != nullptr);
  CHECK(cpu_target_machine_options != nullptr);
  ASSIGN_OR_RETURN(
      CompiledKernel * kernel,
      GetOrCreateKernel(
          hash, [&]() -> absl::StatusOr<std::unique_ptr<CompiledKernel>> {
            return Compile(module, *cc->cuda_compute_capability(),
                           *cpu_target_machine_options);
          }));
  if (kernel->is_nvshmem_used) {
    return absl::UnimplementedError("NVSHMEM is not supported in XLA.");
  }
  return std::make_unique<CustomCallResources>(
      CustomCallResources{.kernel = kernel, .hash = hash});
}

absl::StatusOr<std::vector<ffi::AnyBuffer>> GetBuffers(
    ffi::RemainingArgs inputs, ffi::RemainingRets results) {
  std::vector<ffi::AnyBuffer> buffers;
  buffers.reserve(inputs.size() + results.size());
  for (int i = 0; i < inputs.size(); ++i) {
    ASSIGN_OR_RETURN(auto input_buffer, inputs.get<ffi::AnyBuffer>(i));
    if (reinterpret_cast<uintptr_t>(input_buffer.untyped_data()) %
        mosaic::gpu::kExpectedHbmAlignment) {
      return absl::InvalidArgumentError(
          absl::StrFormat("Input buffer %d is not %d-byte aligned", i,
                          mosaic::gpu::kExpectedHbmAlignment));
    }
    buffers.push_back(input_buffer);
  }

  std::vector<ffi::AnyBuffer> result_buffers;
  result_buffers.reserve(results.size());
  for (int i = 0; i < results.size(); ++i) {
    ASSIGN_OR_RETURN(auto result_buffer, results.get<ffi::AnyBuffer>(i));
    if (reinterpret_cast<uintptr_t>(result_buffer->untyped_data()) %
        mosaic::gpu::kExpectedHbmAlignment) {
      return absl::InvalidArgumentError(
          absl::StrFormat("Output buffer %d is not %d-byte aligned", i,
                          mosaic::gpu::kExpectedHbmAlignment));
    }
    buffers.push_back(*result_buffer);
  }
  return buffers;
}

bool ModuleUsesCollectiveMetadata(const xla::ffi::Dictionary& attrs) {
  return attrs.get<bool>("uses_xla_collective_metadata").value_or(false);
}

bool ModuleSkipsDeviceBarrier(const xla::ffi::Dictionary& attrs) {
  return attrs.get<bool>("skip_device_barrier").value_or(false);
}

size_t GetCollectiveMetadataSize(size_t num_buffers, size_t num_devices) {
  const size_t param_to_peers_size_bytes =
      num_buffers * num_devices * sizeof(void*);
  const size_t param_to_multimem_addresses_size_bytes =
      num_buffers * sizeof(void*);
  return sizeof(CollectiveKernelMetadata) + param_to_peers_size_bytes +
         param_to_multimem_addresses_size_bytes;
}

absl::StatusOr<std::vector<int64_t>> ParseInts(std::string_view str) {
  std::vector<int64_t> result;
  std::vector<std::string> strs = absl::StrSplit(str, ',');
  result.reserve(strs.size());
  for (const std::string& s : strs) {
    int64_t result_int;
    if (!absl::SimpleAtoi(s, &result_int)) {
      return absl::InvalidArgumentError(
          absl::StrFormat("Failed to parse integer: %s", s));
    }
    result.push_back(result_int);
  }
  return result;
}

absl::StatusOr<std::vector<int64_t>> GetReplicaIds(
    const xla::ffi::Dictionary& attributes) {
  std::string_view replica_ids_str =
      attributes.get<std::string_view>("xla_replica_ids").value_or("");
  if (replica_ids_str.empty()) {
    return absl::InvalidArgumentError("No replica ids found in attributes.");
  }
  return ParseInts(replica_ids_str);
}

// Returns the vector of booleans indicating whether the corresponding buffer
// is used with multimem instruction or across several processes.
absl::StatusOr<std::vector<bool>> ParseCollectiveMemoryParameters(
    xla::ffi::Dictionary attributes, size_t num_buffers) {
  std::string_view attribute_value =
      attributes.get<std::string_view>("symmetric_memory_parameters")
          .value_or("");
  if (attribute_value.empty()) {
    // TODO(b/546817872): Remove multimem_parameters check once backward
    // compatibility period is over (should be after 01-FEB-2027).
    attribute_value =
        attributes.get<std::string_view>("multimem_parameters").value_or("");

    if (attribute_value.empty()) {
      return std::vector<bool>(num_buffers, false);
    }
  }
  ASSIGN_OR_RETURN(std::vector<int64_t> collective_memory_parameters,
                   ParseInts(attribute_value));
  if (collective_memory_parameters.size() != num_buffers) {
    return absl::InvalidArgumentError(absl::StrFormat(
        "Multimem arguments list size %d is not equal to number "
        "of buffers %d",
        collective_memory_parameters.size(), num_buffers));
  }
  std::vector<bool> result(num_buffers);
  for (int64_t arg = 0; arg < num_buffers; ++arg) {
    const int collective_memory = collective_memory_parameters[arg];
    switch (collective_memory) {
      case 0:
        result[arg] = false;
        break;
      case 1:
        result[arg] = true;
        break;
      default:
        return absl::InvalidArgumentError(absl::StrFormat(
            "Invalid collective memory parameter: %d", collective_memory));
    }
  }
  return result;
}

absl::StatusOr<xla::gpu::GpuCliqueKey> GetCliqueKey(
    const xla::gpu::CollectiveParams& collective_params,
    const xla::ffi::Dictionary& attributes) {
  ASSIGN_OR_RETURN(std::vector<int64_t> replica_ids, GetReplicaIds(attributes));

  xla::ReplicaGroup group;
  group.mutable_replica_ids()->Reserve(replica_ids.size());
  for (int64_t replica_id : replica_ids) {
    group.add_replica_ids(replica_id);
  }

  return GetGpuCliqueKey(
      collective_params, {group},
      xla::CollectiveOpGroupMode::COLLECTIVE_OP_GROUP_MODE_FLATTENED_ID);
}

// Returns device groups for a collective operation. Device groups used by
// XLA to decide if collective communicator can be safely split from one of
// the existing clique. Passing incorrect device groups can lead to deadlocks
// or wasted resources from allocating duplicate communicators.
absl::StatusOr<std::vector<std::vector<xla::GlobalDeviceId>>>
GetCliqueDeviceGroups(const xla::gpu::CollectiveParams& collective_params,
                      const xla::ffi::Dictionary& attributes) {
  ASSIGN_OR_RETURN(std::vector<int64_t> replica_ids, GetReplicaIds(attributes));

  std::vector<std::vector<xla::GlobalDeviceId>> device_groups(1);
  device_groups[0].resize(replica_ids.size());
  for (int32_t i = 0; i < replica_ids.size(); ++i) {
    device_groups[0][i] = xla::GlobalDeviceId(replica_ids[i]);
  }

  return device_groups;
}

struct PtrFormatter {
  void operator()(std::string* out, const void* ptr) const {
    absl::StrAppend(out, absl::StrFormat("%p", ptr));
  }
};

struct DeviceAddressFormatter {
  void operator()(std::string* out,
                  const se::DeviceAddressBase& address) const {
    absl::StrAppend(out, absl::StrFormat("DeviceAddress(ptr=%p, size=%d)",
                                         address.opaque(), address.size()));
  }
};

void* AddOffset(void* ptrs, int64_t offset) {
  return reinterpret_cast<void*>(reinterpret_cast<uint64_t>(ptrs) + offset);
}

void* SubtractOffset(void* ptrs, int64_t offset) {
  return reinterpret_cast<void*>(reinterpret_cast<uint64_t>(ptrs) - offset);
}

DeviceState& GetDeviceState(CustomCallResources* resources,
                            int32_t device_ordinal) {
  CHECK(0 <= device_ordinal && device_ordinal < resources->device_states.size())
      << "Device ordinal " << device_ordinal
      << " is out of collective metadata bounds: "
      << resources->device_states.size();
  return resources->device_states[device_ordinal];
}

absl::Status MosaicGpuPrepare(
    const xla::gpu::CollectiveParams* absl_nullable collective_params,
    xla::gpu::CollectiveMemoryRequests* absl_nullable
        collective_memory_requests,
    xla::gpu::CollectiveCliqueRequests* absl_nullable clique_requests,
    se::DeviceAddressAllocator* absl_nullable allocator,
    CustomCallResources* resources, ffi::RemainingArgs inputs,
    ffi::RemainingRets results, xla::ffi::Dictionary attributes) {
  int device_ordinal = collective_params->global_device_id.value();
  XLA_VLOG_DEVICE(5, device_ordinal) << "MosaicGpuPrepare";
  // Module initialization calls cuModuleLoadData to load the PTX into the GPU.
  // CUDA has a system-wide mutex around this call which can cause a deadlock
  // if user tries to load the module on one GPU while another GPU is executing
  // the same module. To prevent this we need to initialize all modules used
  // by several devices before the first execution.
  // See more details:
  // https://docs.nvidia.com/cuda/cuda-programming-guide/04-special-topics/lazy-loading.html#impact-on-concurrent-kernel-execution
  // This operation should be done at Prepare stage since XLA launches a
  // rendez-vous between Prepare and Initialize, which we need here to make sure
  // that modules were loaded on all devices before the first execution.
  DeviceState& device_state =
      GetDeviceState(resources, collective_params->local_device_id.value());
  ASSIGN_OR_RETURN(device_state.kernel_handle, CachedInit(resources->kernel));
  CHECK(device_state.kernel_handle != nullptr);

  if (!ModuleUsesCollectiveMetadata(attributes)) {
    return absl::OkStatus();
  }
  XLA_VLOG_DEVICE(5, device_ordinal)
      << "MosaicGpuPrepare uses collective metadata";

  if (collective_params == nullptr || clique_requests == nullptr ||
      collective_memory_requests == nullptr || allocator == nullptr) {
    return absl::InternalError(
        "collective_params, clique_requests, collective_memory_requests, and "
        "allocator must not be null in MosaicGpuPrepare");
  }

  ASSIGN_OR_RETURN(std::vector<ffi::AnyBuffer> buffers,
                   GetBuffers(inputs, results));

  ASSIGN_OR_RETURN(xla::gpu::GpuCliqueKey clique_key,
                   GetCliqueKey(*collective_params, attributes));
  ASSIGN_OR_RETURN(std::vector<std::vector<xla::GlobalDeviceId>> device_groups,
                   GetCliqueDeviceGroups(*collective_params, attributes));
  xla::gpu::CollectiveCliqueRequests::CliqueRequirements requirements;
  requirements.barrier_reqs =
      xla::gpu::CollectiveCliqueRequests::BarrierRequirements();
  requirements.barrier_reqs->use_cross_device_barrier = true;
  RETURN_IF_ERROR(
      clique_requests->RequestClique(clique_key, device_groups, requirements));
  ASSIGN_OR_RETURN(std::vector<bool> collective_parameters,
                   ParseCollectiveMemoryParameters(attributes, buffers.size()));
  for (int i = 0; i < buffers.size(); ++i) {
    if (collective_parameters[i]) {
      RETURN_IF_ERROR(collective_memory_requests->RequestSymmetricAddress(
          clique_key, buffers[i].device_memory()));
    } else {
      RETURN_IF_ERROR(collective_memory_requests->RequestPeerAddress(
          clique_key, buffers[i].device_memory()));
    }
  }

  const size_t metadata_size =
      GetCollectiveMetadataSize(buffers.size(), clique_key.num_devices());
  if (device_state.metadata_handle.is_null()) {
    XLA_VLOG_DEVICE(5, device_ordinal)
        << "Allocating device memory for Mosaic GPU collective metadata";
    ASSIGN_OR_RETURN(
        device_state.metadata_handle,
        allocator->Allocate(collective_params->executor->device_ordinal(),
                            metadata_size));
  }

  XLA_VLOG_DEVICE(5, device_ordinal)
      << "MosaicGpuPrepare is done for clique key: " << clique_key;
  return absl::OkStatus();
}

absl::Status MosaicGpuInitialize(
    se::Stream* stream, const xla::gpu::CollectiveParams* collective_params,
    xla::gpu::CollectiveCliques* collective_cliques,
    const xla::gpu::CollectiveMemory* collective_memory,
    ffi::RemainingArgs inputs, ffi::RemainingRets results,
    CustomCallResources* resources, xla::ffi::Dictionary attributes) {
  bool uses_collective_metadata = ModuleUsesCollectiveMetadata(attributes);
  if (!uses_collective_metadata) {
    // If the kernel does not use collective metadata, we can skip the
    // initialization.
    return absl::OkStatus();
  }

  const char* xla_flags = getenv("XLA_FLAGS");
  if (xla_flags &&
      absl::StrContains(xla_flags, "xla_gpu_experimental_enable_nvshmem") &&
      !absl::StrContains(xla_flags,
                         "xla_gpu_experimental_enable_nvshmem=false")) {
    return absl::InvalidArgumentError(
        "If you're using a single-process for multiple devices, you should "
        "remove --xla_gpu_experimental_enable_nvshmem from your XLA flags.");
  }

  if (stream == nullptr || collective_params == nullptr ||
      collective_cliques == nullptr || collective_memory == nullptr ||
      resources == nullptr) {
    return absl::InternalError(
        "stream, collective_params, collective_cliques, collective_memory, and "
        "resources must not be null in MosaicGpuInitialize");
  }
  int device_ordinal = collective_params->global_device_id.value();
  XLA_VLOG_DEVICE(5, device_ordinal) << "MosaicGpuInitialize start";
  ASSIGN_OR_RETURN(std::vector<ffi::AnyBuffer> buffers,
                   GetBuffers(inputs, results));

  std::vector<void*> parameter_multimem_addresses(buffers.size(), nullptr);
  ASSIGN_OR_RETURN(xla::gpu::GpuCliqueKey clique_key,
                   GetCliqueKey(*collective_params, attributes));

  std::vector<void*> param_to_peers(buffers.size() * clique_key.num_devices());
  ASSIGN_OR_RETURN(std::vector<bool> collective_memory_parameters,
                   ParseCollectiveMemoryParameters(attributes, buffers.size()));
  bool is_multimem_used =
      resources->kernel != nullptr && resources->kernel->is_multimem_used;

  const bool all_parameters_in_collective_memory = absl::c_all_of(
      collective_memory_parameters,
      [](bool is_collective_memory) { return is_collective_memory; });
  if (!all_parameters_in_collective_memory && !clique_key.is_local()) {
    return absl::InvalidArgumentError(
        "Non-local cliques are only supported when all parameters are allocated"
        " within the collective memory.");
  }

  xla::RankId rank =
      clique_key.rank(collective_params->global_device_id).value();
  for (int i = 0; i < buffers.size(); ++i) {
    XLA_VLOG_DEVICE(6, device_ordinal)
        << "MosaicGpuInitialize processing buffer: " << i;
    se::DeviceAddressBase device_address = buffers[i].device_memory();

    if (collective_memory_parameters[i]) {
      // The physical memory range contains several allocations
      // (for example we have separate allocations for the HLO module parameters
      // and temporary buffers). Each allocation can contain several HLO buffers
      // which can overlap within the lifetime of HLO execution.
      // FindMultimemAddress returns the address of an allocation within which
      // a given buffer is located and the offset of this buffer. Since we are
      // mapping the allocation range to multimem address space we also need to
      // substract the offset from the parameter base address to exchange the
      // addresses of allocation in which the parameter is located.
      auto [symmetric_memory, offset] =
          collective_memory->FindSymmetricMemory(clique_key, device_address);
      if (symmetric_memory == nullptr) {
        return absl::InternalError(absl::StrFormat(
            "Failed to find symmetric memory for buffer %d in clique %s", i,
            clique_key.ToString()));
      }

      if (is_multimem_used) {
        ASSIGN_OR_RETURN(se::DeviceAddressBase multimem_address,
                         symmetric_memory->multimem_addr());

        XLA_VLOG_DEVICE(6, device_ordinal)
            << "MosaicGpuInitialize buffer: " << i << " device_address: ("
            << device_address.opaque() << ", size: " << device_address.size()
            << ") found multimem_address: (" << multimem_address.opaque()
            << ", offset: " << offset << ")";

        parameter_multimem_addresses[i] = multimem_address.opaque();
      }

      // Use the allocated memory allocation instead to correctly calculate
      // the offset of the multimem parameter.
      const size_t parameter_offset = i * clique_key.num_devices();
      for (int device_rank = 0; device_rank < clique_key.num_devices();
           ++device_rank) {
        if (device_rank == rank.value()) {
          param_to_peers[parameter_offset + device_rank] =
              se::DeviceAddressBase(
                  SubtractOffset(device_address.opaque(), offset),
                  device_address.size() + offset)
                  .opaque();
        } else {
          // Peer address returns the address of XLA allocation on a peer
          // device This address corresponds to the multimem address space.
          ASSIGN_OR_RETURN(
              se::DeviceAddressBase peer_address,
              symmetric_memory->peer_addr(xla::RankId(device_rank)));
          param_to_peers[parameter_offset + device_rank] =
              peer_address.opaque();
        }
        XLA_VLOG_DEVICE(7, device_ordinal)
            << "MosaicGpuInitialize buffer: " << i << " device_address: ("
            << device_address.opaque() << ", size: " << device_address.size()
            << ") found peer [" << device_rank << "] address: ("
            << param_to_peers[parameter_offset + device_rank]
            << ", size: " << device_address.size() + offset << ")";
      }
    } else {
      for (int device_rank = 0; device_rank < clique_key.num_devices();
           ++device_rank) {
        const size_t parameter_offset = i * clique_key.num_devices();
        std::optional<se::DeviceAddressBase> peer_address =
            collective_memory->FindPeerAddress(
                clique_key, xla::RankId(device_rank), device_address);
        if (!peer_address.has_value()) {
          return absl::InternalError(absl::StrFormat(
              "Failed to find peer address for buffer %d and device rank %d "
              "in clique %s",
              i, device_rank, clique_key.ToString()));
        }
        XLA_VLOG_DEVICE(7, device_ordinal)
            << "MosaicGpuInitialize buffer: " << i << " device_address: ("
            << device_address.opaque() << ", size: " << device_address.size()
            << ") found peer [" << device_rank << "] address: ("
            << peer_address->opaque() << ", size: " << peer_address->size()
            << ")";
        param_to_peers[parameter_offset + device_rank] = peer_address->opaque();
      }
    }
  }

  DeviceState& device_state =
      GetDeviceState(resources, collective_params->local_device_id.value());

  // Construct the collective kernel metadata information.
  CollectiveKernelMetadata metadata;
  metadata.rank = rank.value();
  // See description of DeviceState for more details.
  metadata.param_to_peers = nullptr;
  metadata.param_to_multimem_addresses = nullptr;

  const size_t param_to_peers_size_bytes =
      param_to_peers.size() * sizeof(void*);
  const size_t param_to_multimem_addresses_size_bytes =
      parameter_multimem_addresses.size() * sizeof(void*);

  const size_t metadata_size =
      GetCollectiveMetadataSize(buffers.size(), clique_key.num_devices());
  device_state.metadata_bytes.resize(metadata_size);
  std::memcpy(device_state.metadata_bytes.data(), &metadata,
              sizeof(CollectiveKernelMetadata));
  void* param_to_peers_ptr = AddOffset(device_state.metadata_bytes.data(),
                                       sizeof(CollectiveKernelMetadata));
  std::memcpy(param_to_peers_ptr, param_to_peers.data(),
              param_to_peers_size_bytes);
  void* param_to_multimem_addresses_ptr =
      AddOffset(param_to_peers_ptr, param_to_peers_size_bytes);
  std::memcpy(param_to_multimem_addresses_ptr,
              parameter_multimem_addresses.data(),
              param_to_multimem_addresses_size_bytes);
  // Copy metadata to the device.
  se::DeviceAddressBase metadata_address = device_state.metadata_handle.cref();
  RETURN_IF_ERROR(stream->Memcpy(&metadata_address,
                                 device_state.metadata_bytes.data(),
                                 device_state.metadata_bytes.size()));

  XLA_VLOG_DEVICE(5, device_ordinal)
      << "Constructed device state {"
      << " metadata rank: " << metadata.rank << ", param_to_peers: ("
      << absl::StrJoin(param_to_peers, ", ", PtrFormatter{})
      << "), multimem address spaces: ("
      << absl::StrJoin(parameter_multimem_addresses, ", ", PtrFormatter{})
      << "), copied metadata to the device with address: "
      << metadata_address.opaque() << "}";
  return absl::OkStatus();
}

// Collects input/output buffer pointers (plus device and host collective
// metadata pointers when applicable), invokes `kernel->host_func` to populate
// the kernel launch specification, and resolves the communicator for a
// pre-launch multi-GPU barrier when required (or returns nullptr if no barrier
// is needed).
absl::StatusOr<
    std::pair<mosaic::gpu::MosaicKernelSpec, xla::gpu::GpuCommunicator*>>
BuildKernelSpec(const xla::gpu::CollectiveParams* collective_params,
                const xla::gpu::CollectiveCliques* collective_cliques,
                ffi::RemainingArgs inputs, ffi::RemainingRets results,
                CompiledKernel* kernel, DeviceState& device_state,
                xla::ffi::Dictionary attributes, int device_ordinal) {
  ASSIGN_OR_RETURN(std::vector<ffi::AnyBuffer> buffers,
                   GetBuffers(inputs, results));
  bool uses_collective_metadata = ModuleUsesCollectiveMetadata(attributes);
  std::vector<void*> buffer_ptrs;
  buffer_ptrs.reserve(buffers.size() + (uses_collective_metadata ? 2 : 0));
  for (const ffi::AnyBuffer& buffer : buffers) {
    buffer_ptrs.push_back(buffer.untyped_data());
  }

  xla::gpu::GpuCommunicator* comm = nullptr;
  // Adding a CPU version of the collective metadata for TMA initialization.
  if (uses_collective_metadata) {
    TF_RET_CHECK(collective_params != nullptr)
        << "InternalError: collective_params is null";
    ASSIGN_OR_RETURN(xla::gpu::GpuCliqueKey clique_key,
                     GetCliqueKey(*collective_params, attributes));
    auto current_rank =
        clique_key.rank(collective_params->global_device_id).value();

    se::DeviceAddressBase metadata_address =
        device_state.metadata_handle.cref();
    XLA_VLOG_DEVICE(6, device_ordinal)
        << "Using collective metadata address: " << metadata_address.opaque()
        << " clique_key: " << clique_key;

    // Appending both the device and the host-side collective metadata.
    // The host-side metadata is needed for TMA initialization.
    buffer_ptrs.push_back(metadata_address.opaque());
    buffer_ptrs.push_back(device_state.metadata_bytes.data());

    if (!ModuleSkipsDeviceBarrier(attributes)) {
      ASSIGN_OR_RETURN(comm,
                       collective_cliques->GetComm(clique_key, current_rank));
    }
  }

  TF_RET_CHECK(kernel->host_func != nullptr)
      << "InternalError: build_kernel_spec is null";
  mosaic::gpu::MosaicKernelSpec cfg;
  kernel->host_func(&cfg, buffer_ptrs.data());
  return std::make_pair(std::move(cfg), comm);
}

absl::Status MosaicGpuRecord(
    se::Stream* stream, const xla::gpu::CollectiveParams* collective_params,
    const xla::gpu::CollectiveCliques* collective_cliques,
    ffi::RecordContext record_ctx, ffi::RemainingArgs inputs,
    ffi::RemainingRets results, CustomCallResources* resources,
    xla::ffi::Dictionary attributes) {
  tsl::profiler::TraceMe trace("MosaicGpuRecord");
  CompiledKernel* kernel = resources->kernel;
  const int device_ordinal = stream->parent()->device_ordinal();
  XLA_VLOG_DEVICE(5, device_ordinal)
      << "MosaicGpuRecord called for " << kernel->kernel_name << " with "
      << record_ctx.action();
  absl::Cleanup cleanup = [&] {
    XLA_VLOG_DEVICE(5, device_ordinal)
        << "MosaicGpuRecord finished for " << kernel->kernel_name << " with "
        << record_ctx.action();
  };

  DeviceState& device_state = GetDeviceState(resources, device_ordinal);
  ASSIGN_OR_RETURN(
      (auto [cfg, comm]),
      BuildKernelSpec(collective_params, collective_cliques, inputs, results,
                      kernel, device_state, attributes, device_ordinal));

  // Translate from MosaicKernelSpec to ffi::KernelArg.
  llvm::SmallVector<ffi::KernelArg, 8> kernel_args;
  kernel_args.reserve(cfg.args.size());
  for (const auto& arg : cfg.args) {
    if (const auto* host =
            std::get_if<mosaic::gpu::MosaicKernelSpec::HostArg>(&arg)) {
      kernel_args.push_back(ffi::HostValue{cfg.host_bytes.data() + host->offset,
                                           static_cast<size_t>(host->size)});
    } else {
      kernel_args.push_back(ffi::DevicePointer{
          std::get<mosaic::gpu::MosaicKernelSpec::DeviceArg>(arg).ptr});
    }
  }
  XLA_FFI_LaunchDims dims{
      /*grid=*/{static_cast<int32_t>(cfg.grid.x),
                static_cast<int32_t>(cfg.grid.y),
                static_cast<int32_t>(cfg.grid.z)},
      /*block=*/
      {static_cast<int32_t>(cfg.block.x), static_cast<int32_t>(cfg.block.y),
       static_cast<int32_t>(cfg.block.z)},
      /*cluster=*/
      {static_cast<int32_t>(cfg.cluster.x), static_cast<int32_t>(cfg.cluster.y),
       static_cast<int32_t>(cfg.cluster.z)}};

  if (record_ctx.action() == ffi::RecordAction::kCreate) {
    TF_RET_CHECK(device_state.kernel_handle != nullptr)
        << "InternalError: kernel_handle is null; MosaicGpuPrepare must run "
           "before MosaicGpuRecord";
    llvm::SmallVector<const XLA_FFI_Command*, 1> dependencies;
    if (comm != nullptr) {
      ASSIGN_OR_RETURN(
          const XLA_FFI_Command* barrier_cmd,
          comm->RecordMultiGpuBarrier(record_ctx, /*cmd=*/nullptr));
      dependencies.push_back(barrier_cmd);
    }
    XLA_VLOG_DEVICE(5, device_ordinal)
        << "MosaicGpuRecord creating launch for " << kernel->kernel_name
        << " uses_pdl: " << cfg.uses_pdl << " grid: " << cfg.grid.x << ", "
        << cfg.grid.y << ", " << cfg.grid.z << " cluster: " << cfg.cluster.x
        << ", " << cfg.cluster.y << ", " << cfg.cluster.z
        << " block: " << cfg.block.x << ", " << cfg.block.y << ", "
        << cfg.block.z << " smem_bytes: " << cfg.smem_bytes;
    return record_ctx
        .CreateLaunch(kernel->kernel_name.c_str(),
                      device_state.kernel_handle->function(),
                      /*kernel_size=*/0, ffi::SourceFormat::kFunctionPtr, dims,
                      cfg.smem_bytes, /*uses_pdl=*/cfg.uses_pdl, kernel_args,
                      dependencies)
        .status();
  }
  TF_RET_CHECK(record_ctx.action() == ffi::RecordAction::kUpdate)
      << "InternalError: unexpected record action: "
      << static_cast<int>(record_ctx.action());
  if (comm != nullptr) {
    TF_RET_CHECK(record_ctx.commands().size() == 2)
        << "InternalError: expected 2 commands in record context, got "
        << record_ctx.commands().size();
    RETURN_IF_ERROR(
        comm->RecordMultiGpuBarrier(record_ctx, record_ctx.commands()[0])
            .status());
    return record_ctx.UpdateLaunch(record_ctx.commands()[1], kernel_args);
  }
  TF_RET_CHECK(record_ctx.commands().size() == 1)
      << "InternalError: expected 1 command in record context, got "
      << record_ctx.commands().size();
  return record_ctx.UpdateLaunch(record_ctx.commands()[0], kernel_args);
}

absl::Status MosaicGpuExecute(
    se::Stream* stream, const xla::gpu::CollectiveParams* collective_params,
    const xla::gpu::CollectiveCliques* collective_cliques,
    ffi::RemainingArgs inputs, ffi::RemainingRets results,
    CustomCallResources* resources, xla::ffi::Dictionary attributes) {
  CompiledKernel* kernel = resources->kernel;
  cudaStream_t cuda_stream =
      reinterpret_cast<cudaStream_t>(stream->platform_specific_handle().stream);
  DeviceState& device_state =
      GetDeviceState(resources, collective_params->local_device_id.value());
  int device_ordinal = collective_params->global_device_id.value();

  ASSIGN_OR_RETURN(
      (auto [cfg, comm]),
      BuildKernelSpec(collective_params, collective_cliques, inputs, results,
                      kernel, device_state, attributes, device_ordinal));

  if (comm != nullptr) {
    XLA_VLOG_DEVICE(6, device_ordinal) << "Starting multi-GPU barrier";
    xla::gpu::GpuCollectives::Executor executor(stream);
    RETURN_IF_ERROR(comm->LaunchMultiGpuBarrier(executor));
    XLA_VLOG_DEVICE(6, device_ordinal) << "Finished multi-GPU barrier";
  }

  XLA_VLOG_DEVICE(5, device_ordinal)
      << "MosaicGpuExecute launching kernel with name: " << kernel->kernel_name
      << " uses_pdl: " << cfg.uses_pdl << " grid: " << cfg.grid.x << ", "
      << cfg.grid.y << ", " << cfg.grid.z << " cluster: " << cfg.cluster.x
      << ", " << cfg.cluster.y << ", " << cfg.cluster.z
      << " block: " << cfg.block.x << ", " << cfg.block.y << ", " << cfg.block.z
      << " smem_bytes: " << cfg.smem_bytes;
  return MosaicGpuLaunchKernel(device_state.kernel_handle->function(),
                               cuda_stream, &cfg, *kernel, device_ordinal);
}

}  // namespace

namespace xla::ffi {
template <>
struct TypeRegistry::SerDes<mosaic::gpu::CustomCallResources> {
  static constexpr bool value = true;

  static absl::StatusOr<std::string> Serialize(
      const mosaic::gpu::CustomCallResources& resources) {
    return mosaic::gpu::CustomCallResources::Serialize(resources);
  }

  static absl::StatusOr<std::unique_ptr<mosaic::gpu::CustomCallResources>>
  Deserialize(absl::string_view data) {
    return mosaic::gpu::CustomCallResources::Deserialize(data);
  }
};
}  // namespace xla::ffi

namespace {
XLA_FFI_DEFINE_HANDLER(kMosaicGpuPrepare, MosaicGpuPrepare,
                       ffi::Ffi::BindPrepare()
                           .Ctx<ffi::CollectiveParams>()
                           .Ctx<ffi::CollectiveMemoryRequests>()
                           .Ctx<ffi::CollectiveCliqueRequests>()
                           .Ctx<ffi::Allocator>()
                           .Ctx<xla::ffi::State<CustomCallResources>>()
                           .RemainingArgs()
                           .RemainingRets()
                           .Attrs());

XLA_FFI_DEFINE_HANDLER(kMosaicGPUInstantiate, InstantiateResources,
                       ffi::Ffi::BindInstantiate()
                           .Ctx<ffi::TargetGpuComputeCapability>()
                           .Ctx<ffi::CpuTargetMachineOptions>()
                           .Attrs());

XLA_FFI_DEFINE_HANDLER(
    kMosaicGpuInitialize, MosaicGpuInitialize,
    ffi::Ffi::BindInitialize()
        .Ctx<ffi::Stream>()
        .Ctx<ffi::CollectiveParams>()
        .Ctx<ffi::CollectiveCliques>()
        .Ctx<ffi::CollectiveMemory>()
        .RemainingArgs()
        .RemainingRets()
        .Ctx<xla::ffi::State<mosaic::gpu::CustomCallResources>>()
        .Attrs(),
    {ffi::Traits::kCmdBufferCompatible});

XLA_FFI_DEFINE_HANDLER(
    kMosaicGpuRecord, MosaicGpuRecord,
    xla::ffi::Ffi::BindRecord()
        .Ctx<ffi::Stream>()
        .Ctx<ffi::CollectiveParams>()
        .Ctx<ffi::CollectiveCliques>()
        .Ctx<ffi::Extension<ffi::RecordExtension>>()
        .RemainingArgs()
        .RemainingRets()
        .Ctx<xla::ffi::State<mosaic::gpu::CustomCallResources>>()
        .Attrs());

//  We expect the following attributes:
// - kernel_hash: a hash of the kernel.
// - module: the serialized MLIR module.
// - use_custom_barrier
// - skip_device_barrier (optional)
// - uses_xla_collective_metadata (optional)
XLA_FFI_DEFINE_HANDLER(
    kMosaicGpuExecute, MosaicGpuExecute,
    ffi::Ffi::BindExecute()
        .Ctx<ffi::Stream>()
        .Ctx<ffi::CollectiveParams>()
        .Ctx<ffi::CollectiveCliques>()
        .RemainingArgs()
        .RemainingRets()
        .Ctx<xla::ffi::State<mosaic::gpu::CustomCallResources>>()
        .Attrs(),
    {ffi::Traits::kCmdBufferCompatible});

XLA_FFI_REGISTER_HANDLER(ffi::GetXlaFfiApi(), "mosaic_gpu_v2", "CUDA",
                         {
                             /*instantiate=*/kMosaicGPUInstantiate,
                             /*prepare=*/kMosaicGpuPrepare,
                             /*initialize=*/kMosaicGpuInitialize,
                             /*execute=*/kMosaicGpuExecute,
                             /*record=*/kMosaicGpuRecord,
                         });

}  // namespace

extern "C" {

// C-Wrapper that both populates the kernel spec and launches using it.
void MosaicGpuLaunch(void* kernel, CUfunction function, void* stream,
                     void** buffers, int device_ordinal) {
  auto* k = reinterpret_cast<CompiledKernel*>(kernel);
  mosaic::gpu::MosaicKernelSpec cfg;
  k->host_func(&cfg, buffers);
  absl::Status launch_status = MosaicGpuLaunchKernel(
      function, reinterpret_cast<CUstream>(stream), &cfg, *k, device_ordinal);
  if (!launch_status.ok()) {
    LOG(ERROR) << "MosaicGpuLaunch failed: " << launch_status;
  }
}

__attribute__((visibility("default"))) void** MosaicGpuCompile(
    const char* module, int num_module_bytes) {
  // We should make sure that NVPTX target is initialized before
  // getting cuda compute capability
  mosaic::gpu::EnsureLLVMNVPTXTargetIsRegistered();
  std::string module_str(module, num_module_bytes);
  auto cc = GetCudaComputeCapability();
  if (!cc.ok()) {
    return nullptr;
  }
  auto kernel =
      Compile(module_str, *cc, xla::cpu::TargetMachineOptions::Native());
  if (!kernel.ok()) {
    return nullptr;
  }
  absl::StatusOr<InitResult> init_res = InitKernel(**kernel);
  if (!init_res.ok()) {
    return nullptr;
  }
  auto tuple_ptr = new void*[3];
  tuple_ptr[0] = init_res->function;
  tuple_ptr[1] = reinterpret_cast<void*>(&MosaicGpuLaunch);
  tuple_ptr[2] = (*kernel).release();
  return tuple_ptr;
}

__attribute__((visibility("default"))) void MosaicGpuUnload(void** tuple_ptr) {
  delete reinterpret_cast<CompiledKernel*>(tuple_ptr[2]);
  delete[] tuple_ptr;
}

__attribute__((visibility("default"))) void MosaicGpuClearKernelCache() {
  auto& cache = GetKernelCache();
  absl::MutexLock lock(cache.mutex);
  cache.kernels.clear();
}

}  // extern "C"