#include "Dialect/TritonAMDGPU/IR/Dialect.h"
#include "TritonAMDGPUToLLVM/Passes.h"
#include "TritonAMDGPUTransforms/Passes.h"
#include "amd/include/hipblas_instance.h"
#include "amd/include/hipblas_types.h"
#include "lib/TritonAMDGPUToLLVM/TargetInfo.h"
#include "lld/Common/Driver.h"
#include "mlir/Pass/PassManager.h"
#include "mlir/Target/LLVMIR/Dialect/ROCDL/ROCDLToLLVMIRTranslation.h"
#include "passes.h"
#include "triton/Dialect/TritonInstrument/Transforms/Passes.h"
#include "llvm/ADT/SmallString.h"
#include "llvm/IR/Constants.h"
#include "llvm/IR/GlobalVariable.h"
#include "llvm/IR/Module.h"
#include "llvm/MC/MCAsmBackend.h"
#include "llvm/MC/MCAsmInfo.h"
#include "llvm/MC/MCCodeEmitter.h"
#include "llvm/MC/MCContext.h"
#include "llvm/MC/MCInstrInfo.h"
#include "llvm/MC/MCObjectFileInfo.h"
#include "llvm/MC/MCObjectWriter.h"
#include "llvm/MC/MCParser/MCAsmParser.h"
#include "llvm/MC/MCParser/MCTargetAsmParser.h"
#include "llvm/MC/MCRegisterInfo.h"
#include "llvm/MC/MCSection.h"
#include "llvm/MC/MCStreamer.h"
#include "llvm/MC/MCSubtargetInfo.h"
#include "llvm/MC/MCTargetOptions.h"
#include "llvm/MC/TargetRegistry.h"
#include "llvm/Support/FileSystem.h"
#include "llvm/Support/SourceMgr.h"
#include "llvm/TargetParser/AMDGPUTargetParser.h"
#include <array>
#include <nanobind/nanobind.h>
#include <nanobind/stl/optional.h>
#include <nanobind/stl/pair.h>
#include <nanobind/stl/string.h>
#include <nanobind/stl/vector.h>
#include <optional>
#include <sstream>
#include <stdexcept>

namespace py = nanobind;

namespace {
const char *const amdTargetTriple = "amdgcn-amd-amdhsa";

void init_triton_amd_passes_ttgpuir(py::module_ &m) {
  using namespace mlir::triton;
  m.def("add_to_llvmir",
        [](mlir::PassManager &pm, const std::string &arch, bool ftz) {
          pm.addPass(createConvertTritonAMDGPUToLLVMPass(arch, ftz));
        });
  m.def("add_builtin_func_to_llvmir",
        [](mlir::PassManager &pm, const std::string &arch, bool ftz) {
          pm.addPass(createConvertBuiltinFuncToLLVMPass(arch, ftz));
        });
  m.def("add_prepare_consan_captures", [](mlir::PassManager &pm) {
    mlir::triton::instrument::TritonInstrumentPrepareConSanCapturesOptions
        options;
    options.target = "amd";
    pm.addPass(
        mlir::triton::instrument::createTritonInstrumentPrepareConSanCaptures(
            options));
  });
  ADD_PASS_OPTION_WRAPPER_1("add_allocate_shared_memory",
                            mlir::triton::createAllocateAMDGPUSharedMemoryPass,
                            const std::string &);
  ADD_PASS_OPTION_WRAPPER_3("add_accelerate_matmul",
                            mlir::createTritonAMDGPUAccelerateMatmul,
                            const std::string, int, int);
  ADD_PASS_WRAPPER_0("add_optimize_epilogue",
                     mlir::createTritonAMDGPUOptimizeEpilogue);
  ADD_PASS_WRAPPER_0("add_warp_pipeline", mlir::createTritonAMDGPUWarpPipeline);
  ADD_PASS_OPTION_WRAPPER_1("add_warp_pipeline_conversion",
                            mlir::triton::AMD::createConvertWarpPipelinePass,
                            const std::string &);
  ADD_PASS_OPTION_WRAPPER_1(
      "add_optimize_dot_operands",
      mlir::triton::amdgpu::createTritonAMDGPUOptimizeDotOperands,
      const std::string &);
  m.def("add_hoist_layout_conversions", [](mlir::PassManager &pm) {
    pm.addNestedPass<mlir::triton::FuncOp>(
        mlir::createTritonAMDGPUHoistLayoutConversions());
  });
  m.def("add_sink_layout_conversions", [](mlir::PassManager &pm) {
    pm.addNestedPass<mlir::triton::FuncOp>(
        mlir::createTritonAMDGPUSinkLayoutConversions());
  });
  m.def("add_prepare_if_combining", [](mlir::PassManager &pm) {
    pm.addNestedPass<mlir::triton::FuncOp>(
        mlir::createTritonAMDGPUPrepareIfCombining());
  });
  m.def("add_canonicalize_pointers", [](mlir::PassManager &pm) {
    pm.addNestedPass<mlir::triton::FuncOp>(
        mlir::createTritonAMDGPUCanonicalizePointers());
  });
  m.def("add_move_up_prologue_loads", [](mlir::PassManager &pm) {
    pm.addNestedPass<mlir::triton::FuncOp>(
        mlir::createTritonAMDGPUMoveUpPrologueLoads());
  });
  ADD_PASS_OPTION_WRAPPER_3("add_convert_to_buffer_ops",
                            mlir::createTritonAMDGPUConvertToBufferOps,
                            const std::string &, bool, bool);
  ADD_FUNC_PASS_WRAPPER_0("add_optimize_buffer_op_ptr",
                          mlir::createTritonAMDGPUOptimizeBufferOpPtr);
  ADD_PASS_WRAPPER_0("add_fold_true_cmpi", mlir::createTritonAMDFoldTrueCmpI);
  ADD_PASS_WRAPPER_0("add_fp_sanitizer", mlir::createTritonAMDGPUFpSanitizer);

  ADD_PASS_OPTION_WRAPPER_1("add_block_pingpong",
                            mlir::createTritonAMDGPUBlockPingpong, int32_t);
  ADD_PASS_OPTION_WRAPPER_1("add_schedule_loops",
                            mlir::createTritonAMDGPUScheduleLoops, int);
  ADD_PASS_OPTION_WRAPPER_2("add_pipeline", mlir::createTritonAMDGPUPipeline,
                            bool, bool);
  ADD_PASS_OPTION_WRAPPER_1("add_coalesce_async_copy",
                            mlir::createTritonAMDGPUCoalesceAsyncCopy,
                            std::string);
  ADD_PASS_OPTION_WRAPPER_1("add_update_async_wait_count",
                            mlir::createTritonAMDGPUUpdateAsyncWaitCount,
                            std::string);
  ADD_PASS_WRAPPER_0("add_optimize_descriptor_encoding",
                     mlir::createTritonAMDGPUOptimizeDescriptorEncoding);
  ADD_PASS_WRAPPER_0("add_convert_to_tensor_ops",
                     mlir::createTritonAMDGPUConvertToTensorOps);
  mlir::registerConSanAMDHooks();
  m.def("add_in_thread_transpose", [](mlir::PassManager &pm) {
    pm.addNestedPass<mlir::triton::FuncOp>(
        mlir::createTritonAMDGPUInThreadTranspose());
  });
  ADD_PASS_WRAPPER_1(
      "add_warp_specialize_to_llvm",
      mlir::triton::AMD::createTritonAMDGPUConvertWarpSpecializeToLLVMPass,
      const std::string &);
}

void addControlConstant(llvm::Module *module, const char *name,
                        uint32_t bitwidth, uint32_t value) {
  using llvm::GlobalVariable;

  llvm::IntegerType *type =
      llvm::IntegerType::getIntNTy(module->getContext(), bitwidth);
  auto *initializer = llvm::ConstantInt::get(type, value, /*isSigned=*/false);
  auto *constant = new llvm::GlobalVariable(
      *module, type, /*isConstant=*/true,
      GlobalVariable::LinkageTypes::LinkOnceODRLinkage, initializer, name,
      /*before=*/nullptr, GlobalVariable::ThreadLocalMode::NotThreadLocal,
      /*addressSpace=*/4);
  constant->setAlignment(llvm::MaybeAlign(bitwidth / 8));
  constant->setUnnamedAddr(GlobalVariable::UnnamedAddr::Local);
  constant->setVisibility(GlobalVariable::VisibilityTypes::ProtectedVisibility);
}

} // namespace

LLD_HAS_DRIVER(elf)

static void checkMatmulConstraints(const std::string &A_dtype,
                                   const std::string &B_dtype,
                                   const std::string &C_dtype,
                                   const std::vector<int> &A_shape,
                                   const std::vector<int> &B_shape,
                                   const std::vector<int> &C_shape) {
  // Support FP32/FP16/BF16 and 8-bit FP8 (e4m3fn/e4m3fnuz) and BF8
  // (e5m2fn/e5m2fnuz).
  auto is_fp8 = [](const std::string &dtype) {
    return dtype == "torch.float8_e4m3fn" || dtype == "torch.float8_e5m2fn" ||
           dtype == "torch.float8_e4m3fnuz" || dtype == "torch.float8_e5m2fnuz";
  };
  auto is_fp16_family = [](const std::string &dtype) {
    return dtype == "torch.float16" || dtype == "torch.bfloat16";
  };
  const bool A_is_fp8 = is_fp8(A_dtype);
  const bool B_is_fp8 = is_fp8(B_dtype);
  const bool A_supported =
      (A_is_fp8 || is_fp16_family(A_dtype) || A_dtype == "torch.float32");
  const bool B_supported =
      (B_is_fp8 || is_fp16_family(B_dtype) || B_dtype == "torch.float32");
  const bool C_supported = (is_fp16_family(C_dtype) ||
                            C_dtype == "torch.float32" || is_fp8(C_dtype));

  if (!A_supported || !B_supported || !C_supported) {
    std::ostringstream oss;
    oss << "Unsupported data type. Got A=" << A_dtype << ", B=" << B_dtype
        << ", C=" << C_dtype
        << ". Supported: float32, float16, bfloat16, float8_e4m3fn, "
           "float8_e5m2fn, float8_e4m3fnuz, float8_e5m2fnuz.";
    throw std::runtime_error(oss.str());
  }

  if (A_is_fp8 && B_is_fp8) {
    if (C_dtype != "torch.float16" && C_dtype != "torch.float32" &&
        C_dtype != "torch.bfloat16") {
      std::ostringstream oss;
      oss << "When A/B are 8-bit (float8_e4m3fn/e4m3fnuz or "
             "float8_e5m2fn/e5m2fnuz), C must"
          << " be torch.float16, torch.float32, or torch.bfloat16.";
      throw std::runtime_error(oss.str());
    }
  } else {
    if (!(A_dtype == B_dtype && A_dtype == C_dtype)) {
      std::ostringstream oss;
      oss << "Data types do not match: A=" << A_dtype << ", B=" << B_dtype
          << ", C=" << C_dtype << ". Expected all equal when not using 8-bit"
          << " inputs.";
      throw std::runtime_error(oss.str());
    }
  }

  if (A_shape.size() != 2 || B_shape.size() != 2 || C_shape.size() != 2) {
    throw std::runtime_error("Only 2D matrices are supported.");
  }

  int k = A_shape[1];
  if (k != B_shape[1]) {
    std::ostringstream oss;
    oss << "Matrix dimensions do not match. A is [" << A_shape[0] << ", "
        << A_shape[1] << "], B is [" << B_shape[0] << ", " << B_shape[1]
        << "]. Expected A.shape[1] == B.shape[1]. Note that B needs to be "
           "transposed.";
    throw std::runtime_error(oss.str());
  }

  int m = A_shape[0];
  if (m != C_shape[0]) {
    std::ostringstream oss;
    oss << "Matrix dimensions do not match. A is [" << A_shape[0] << ", "
        << A_shape[1] << "], C is [" << C_shape[0] << ", " << C_shape[1]
        << "]. Expected A.shape[0] == C.shape[0].";
    throw std::runtime_error(oss.str());
  }

  int n = B_shape[0];
  if (n != C_shape[1]) {
    std::ostringstream oss;
    oss << "Matrix dimensions do not match. B is [" << B_shape[0] << ", "
        << B_shape[1] << "], C is [" << C_shape[0] << ", " << C_shape[1]
        << "]. Expected B.shape[0] == C.shape[1]. Note that B needs to be "
           "transposed.";
    throw std::runtime_error(oss.str());
  }
}

struct HipBlasInit {
  int m;
  int n;
  int k;
  hipDataType dtype;
  hipDataType out_dtype;
};

static HipBlasInit initialize_hipblas_op(py::object &A, py::object &B,
                                         py::object &out,
                                         std::optional<py::object> accumOpt) {
  auto A_shape = py::cast<std::vector<int>>(A.attr("shape"));
  auto B_shape = py::cast<std::vector<int>>(B.attr("shape"));
  auto OUT_shape = py::cast<std::vector<int>>(out.attr("shape"));

  auto A_dtype = py::cast<std::string>(A.attr("dtype").attr("__str__")());
  auto B_dtype = py::cast<std::string>(B.attr("dtype").attr("__str__")());
  auto OUT_dtype = py::cast<std::string>(out.attr("dtype").attr("__str__")());

  if (accumOpt.has_value()) {
    auto C = accumOpt.value();
    auto C_shape = py::cast<std::vector<int>>(C.attr("shape"));
    auto C_dtype = py::cast<std::string>(C.attr("dtype").attr("__str__")());

    checkMatmulConstraints(A_dtype, B_dtype, OUT_dtype, A_shape, B_shape,
                           OUT_shape);
    if (C_dtype != OUT_dtype) {
      throw std::runtime_error("C dtype must match output dtype, got C=" +
                               C_dtype + ", D=" + OUT_dtype);
    }
    if (C_shape != OUT_shape) {
      throw std::runtime_error("C and D shapes must match");
    }
  } else {
    checkMatmulConstraints(A_dtype, B_dtype, OUT_dtype, A_shape, B_shape,
                           OUT_shape);
  }

  hipDataType dtype;
  if (A_dtype == "torch.float8_e4m3fn") {
    // Supported for GFX950.
    dtype = HIP_R_8F_E4M3;
  } else if (A_dtype == "torch.float8_e5m2fn") {
    // supported for GFX950.
    dtype = HIP_R_8F_E5M2;
  } else if (A_dtype == "torch.float8_e4m3fnuz") {
    // Supported for GFX942.
    dtype = HIP_R_8F_E4M3_FNUZ;
  } else if (A_dtype == "torch.float8_e5m2fnuz") {
    // Supported for GFX942.
    dtype = HIP_R_8F_E5M2_FNUZ;
  } else if (A_dtype == "torch.float16") {
    dtype = HIP_R_16F;
  } else if (A_dtype == "torch.float32") {
    dtype = HIP_R_32F;
  } else if (A_dtype == "torch.bfloat16") {
    dtype = HIP_R_16BF;
  } else {
    throw std::runtime_error("Unsupported dtype for hipblasLt: " + A_dtype);
  }

  hipDataType out_dtype;
  if (OUT_dtype == "torch.float16") {
    out_dtype = HIP_R_16F;
  } else if (OUT_dtype == "torch.float32") {
    out_dtype = HIP_R_32F;
  } else if (OUT_dtype == "torch.bfloat16") {
    out_dtype = HIP_R_16BF;
  } else {
    throw std::runtime_error("Unsupported output dtype for hipblasLt: " +
                             OUT_dtype);
  }

  int m = A_shape[0];
  int n = B_shape[0];
  int k = A_shape[1];

  return HipBlasInit{m, n, k, dtype, out_dtype};
}

static std::optional<std::string> lldInvoke(const char *inPath,
                                            const char *outPath) {
  // Workaround: Disable parallelism to avoid hangs caused by LLVM's thread pool
  // when the following code is executed in a forked child process.
  // Context: lld::elf::LinkerDriver::link uses parallelFor which uses the
  // LLVM's thread pool. During cleanup at ~TaskGroup() the child process hangs
  // waiting.
  std::array args{"ld.lld", "--threads=1", "-shared", inPath, "-o", outPath};
  std::string errString;
  llvm::raw_string_ostream errStream(errString);
  auto lldRes = lld::lldMain(args, llvm::outs(), errStream,
                             {{lld::Gnu, &lld::elf::link}});
  bool noErrors = (!lldRes.retCode && lldRes.canRunAgain);
  if (!noErrors) {
    errStream.flush();
    return errString;
  }
  return {};
}

void init_triton_amd(py::module_ &m) {
  m.doc() = "Python bindings to the AMD Triton backend";

  auto passes = m.def_submodule("passes");
  auto ttgpuir_m = passes.def_submodule("ttgpuir");
  init_triton_amd_passes_ttgpuir(ttgpuir_m);

  m.attr("TARGET_TRIPLE") = amdTargetTriple;
  m.attr("CALLING_CONV_AMDGPU_KERNEL") =
      (unsigned)llvm::CallingConv::AMDGPU_KERNEL;

  m.def("load_dialects", [](mlir::MLIRContext &context) {
    mlir::DialectRegistry registry;
    registry.insert<mlir::triton::amdgpu::TritonAMDGPUDialect>();
    // registry.insert<mlir::ROCDL::ROCDLDialect>();
    mlir::registerROCDLDialectTranslation(registry);
    context.appendDialectRegistry(registry);
    context.loadAllAvailableDialects();
  });

  m.def("attach_target_triple", [](llvm::Module *module) {
    module->setTargetTriple(llvm::Triple(amdTargetTriple));
  });

  // Set target architecture ISA version
  m.def("set_isa_version", [](llvm::Module *module, const std::string &arch) {
    llvm::AMDGPU::IsaVersion version = llvm::AMDGPU::getIsaVersion(arch);
    addControlConstant(module, "__oclc_ISA_version", /*bitwidth=*/32,
                       version.Major * 1000 + version.Minor * 100 +
                           version.Stepping);
  });

  // Set boolean control constant
  m.def("set_bool_control_constant",
        [](llvm::Module *module, const std::string &name, bool enable) {
          addControlConstant(module, name.c_str(), /*bitwidth=*/8, enable);
        });

  // Set code object ABI version
  m.def("set_abi_version", [](llvm::Module *module, int version) {
    // Inject the control constant into the LLVM module so that device libraries
    // linked against module can resolve their references to it.
    llvm::Type *i32Ty = llvm::Type::getInt32Ty(module->getContext());
    llvm::GlobalVariable *abi = new llvm::GlobalVariable(
        *module, i32Ty, /*isConstant=*/true,
        llvm::GlobalValue::LinkageTypes::LinkOnceODRLinkage,
        llvm::ConstantInt::get(i32Ty, version), "__oclc_ABI_version", nullptr,
        llvm::GlobalValue::ThreadLocalMode::NotThreadLocal, 4);
    abi->setVisibility(llvm::GlobalValue::VisibilityTypes::ProtectedVisibility);
    abi->setAlignment(llvm::MaybeAlign(4));
    abi->setUnnamedAddr(llvm::GlobalValue::UnnamedAddr::Local);

    // Also attach the control attribute on the LLVM module. This is also needed
    // in addition to the above for various transformations to know what code
    // object version we are targeting at.
    module->addModuleFlag(llvm::Module::Error, "amdhsa_code_object_version",
                          version);
  });

  m.def("cleanup_bitcode_metadata", [](llvm::Module *module) {
    // We can have Clang version metadata from device libraries linked in. We
    // don't care about them so drop them.
    if (auto *ident = module->getNamedMetadata("llvm.ident"))
      module->eraseNamedMetadata(ident);
    // Also various OpenCL version details.
    if (auto *openclVersion = module->getNamedMetadata("opencl.ocl.version"))
      module->eraseNamedMetadata(openclVersion);
  });

  m.def("disable_print_inline", [](llvm::Module *module) {
    // List of functions name prefixes we want to forbid inline.
    std::array<const char *, 2> prefixes = {"__ockl_fprintf", "__ockl_printf"};

    for (llvm::Function &f : module->functions()) {
      if (!f.hasName())
        continue;
      llvm::StringRef name = f.getName();

      auto isNamePrefixed = [&name](const char *prefix) {
        return name.starts_with(prefix);
      };

      if (llvm::any_of(prefixes, isNamePrefixed))
        f.addFnAttr(llvm::Attribute::NoInline);
    }
  });

  m.def("assemble_amdgcn", [](const std::string &assembly,
                              const std::string &arch,
                              const std::string &features) {
    std::string error;

    llvm::Triple triple(amdTargetTriple);
    const llvm::Target *target =
        llvm::TargetRegistry::lookupTarget(triple, error);
    if (!target)
      throw std::runtime_error("target lookup error: " + error);

    llvm::SourceMgr srcMgr;
    srcMgr.AddNewSourceBuffer(llvm::MemoryBuffer::getMemBuffer(assembly),
                              llvm::SMLoc());

    const llvm::MCTargetOptions mcOptions;
    std::unique_ptr<llvm::MCRegisterInfo> mri(target->createMCRegInfo(triple));
    std::unique_ptr<llvm::MCAsmInfo> mai(
        target->createMCAsmInfo(*mri, triple, mcOptions));
    std::unique_ptr<llvm::MCSubtargetInfo> sti(
        target->createMCSubtargetInfo(triple, arch, features));

    llvm::MCContext ctx(triple, *mai, *mri, *sti, &srcMgr);
    std::unique_ptr<llvm::MCObjectFileInfo> mofi(
        target->createMCObjectFileInfo(ctx, /*PIC=*/false,
                                       /*LargeCodeModel=*/false));
    ctx.setObjectFileInfo(mofi.get());

    llvm::SmallString<128> cwd;
    if (!llvm::sys::fs::current_path(cwd))
      ctx.setCompilationDir(cwd);

    llvm::SmallVector<char, 0> result;
    llvm::raw_svector_ostream svos(result);

    std::unique_ptr<llvm::MCStreamer> mcStreamer;
    std::unique_ptr<llvm::MCInstrInfo> mcii(target->createMCInstrInfo());

    std::unique_ptr<llvm::MCCodeEmitter> ce(
        target->createMCCodeEmitter(*mcii, ctx));
    std::unique_ptr<llvm::MCAsmBackend> mab(
        target->createMCAsmBackend(*sti, *mri, mcOptions));
    std::unique_ptr<llvm::MCObjectWriter> ow(mab->createObjectWriter(svos));
    mcStreamer.reset(target->createMCObjectStreamer(
        triple, ctx, std::move(mab), std::move(ow), std::move(ce), *sti));

    std::unique_ptr<llvm::MCAsmParser> parser(
        createMCAsmParser(srcMgr, ctx, *mcStreamer, *mai));
    std::unique_ptr<llvm::MCTargetAsmParser> tap(
        target->createMCAsmParser(*sti, *parser, *mcii));
    if (!tap)
      throw std::runtime_error("assembler initializtion error");

    parser->setTargetParser(*tap);
    parser->Run(/*NoInitialTextSection=*/false);

    return py::bytes(result.data(), result.size());
  });

  m.def("has_architected_sgprs", [](const std::string &arch) {
    std::string error;
    llvm::Triple triple(amdTargetTriple);
    const llvm::Target *target =
        llvm::TargetRegistry::lookupTarget(triple, error);
    if (!target)
      throw std::runtime_error("target lookup error: " + error);
    std::unique_ptr<llvm::MCSubtargetInfo> sti(
        target->createMCSubtargetInfo(triple, arch, ""));
    return sti->checkFeatures("+architected-sgprs");
  });

  m.def("supports_multi_cta_launch", [](const std::string &arch) {
    return mlir::triton::AMD::TargetInfo(arch).supportsMultiCTALaunch();
  });

  m.def("supports_tdm", [](const std::string &arch) {
    return mlir::triton::AMD::TargetInfo(arch).supportsTDM();
  });

  m.def("need_extern_lib", [](llvm::Module *module, const std::string &lib) {
    for (llvm::Function &f : module->functions()) {
      if (f.hasExternalLinkage() && f.hasName() && !f.hasExactDefinition()) {
        llvm::StringRef funcName = f.getName();
        // The rule for linking the extern lib:
        //    if the function name includes ocml or ockl, link
        //    ocml or ockl accordingly.
        if (funcName.contains(lib))
          return true;
        if (funcName.contains("__nv_")) {
          std::stringstream message;
          message << "Implicit conversion of CUDA " << funcName.str()
                  << " device function has been dropped; "
                  << "please, update your source program to use "
                     "triton.language.extra.<op> "
                  << "to replace triton.language.extra.cuda.<op>";
          throw std::runtime_error(message.str());
        }
      }
    }
    return false;
  });

  m.def("set_all_fn_arg_inreg", [](llvm::Function *fn) {
    for (llvm::Argument &arg : fn->args()) {
      // Check for incompatible attributes.
      if (arg.hasByRefAttr() || arg.hasNestAttr())
        continue;
      arg.addAttr(llvm::Attribute::InReg);
    }
  });

  m.def("link_hsaco",
        [](const std::string &inPath, const std::string &outPath) {
          if (auto errString = lldInvoke(inPath.c_str(), outPath.c_str()))
            throw std::runtime_error("LLD failed to link hsaco source " +
                                     inPath + " into object file " + outPath +
                                     " because " + errString.value());
        });

  m.def("add_scalarize_packed_fops_llvm_pass", [](llvm::Function *fn) {
    mlir::triton::AMD::runScalarizePackedFOpsPass(*fn);
  });

  auto hipBlas = m.def_submodule("hipblas");
  // For ROCm installed via TheRock wheels: Preload hipblaslt library via
  // rocm_sdk if available. When using TheRock wheel installs, libhipblaslt
  // resides within the Python wheel package rather than in the standard
  // /opt/rocm/lib location. This preload ensures the library is properly
  // loaded before HipblasLtInstance tries to dlopen it, allowing the dynamic
  // linker to find it from the ROCm wheel's bundled libraries.
  try {
    py::module_::import_("rocm_sdk").attr("preload_libraries")("hipblaslt");
  } catch (...) {
  }
  py::class_<HipblasLtInstance>(hipBlas, "HipblasLt")
      .def(py::new_([](py::object workspace) {
        auto wrk_ptr = py::cast<uint64_t>(workspace.attr("data_ptr")());
        auto wrk_size = py::cast<size_t>(workspace.attr("numel")()) *
                        py::cast<size_t>(workspace.attr("element_size")());
        return new HipblasLtInstance(wrk_ptr, wrk_size);
      }))
      .def("matmul",
           [](HipblasLtInstance &self, py::object &A, py::object &B,
              py::object &C) {
             auto A_ptr = py::cast<uint64_t>(A.attr("data_ptr")());
             auto B_ptr = py::cast<uint64_t>(B.attr("data_ptr")());
             auto C_ptr = py::cast<uint64_t>(C.attr("data_ptr")());
             auto init = initialize_hipblas_op(A, B, C, std::nullopt);
             self.matmul(init.m, init.n, init.k, A_ptr, B_ptr, C_ptr,
                         init.dtype, init.out_dtype);
           })
      .def("gemm", [](HipblasLtInstance &self, py::object &A, py::object &B,
                      py::object &C, py::object &D, float alpha, float beta) {
        auto A_ptr = py::cast<uint64_t>(A.attr("data_ptr")());
        auto B_ptr = py::cast<uint64_t>(B.attr("data_ptr")());
        auto C_ptr = py::cast<uint64_t>(C.attr("data_ptr")());
        auto D_ptr = py::cast<uint64_t>(D.attr("data_ptr")());
        auto init = initialize_hipblas_op(A, B, D, C);
        self.gemm(init.m, init.n, init.k, A_ptr, B_ptr, C_ptr, D_ptr,
                  init.dtype, init.out_dtype, alpha, beta);
      });
}
