# implementation of the GPUCompiler interfaces for generating PTX code

const NVPTX_LLVM_Backend_jll =
    LazyModule("NVPTX_LLVM_Backend_jll",
               UUID("ef6e0fe3-e6ef-59c0-bde6-4989574699e0"))


## target

export PTXCompilerTarget

# Wire-format encoding of the feature set, stamped into the `sm_features` LLVM
# global by `finish_module!` and read back by host-side runtime intrinsics (e.g.
# CUDA.jl's `target_feature_set()`).
@enum TargetFeatureSet::UInt32 begin
    BaselineFeatures = 0
    FamilyFeatures   = 1
    ArchFeatures     = 2
end

Base.@kwdef struct PTXCompilerTarget <: AbstractCompilerTarget
    cap::VersionNumber
    ptx::VersionNumber = v"6.0" # for compatibility with older versions of CUDA.jl

    # subtarget feature set, selecting the suffix on the LLVM CPU name (and `.target`):
    #   :baseline (no suffix)   - forward-compatible (sm_X for any sm_Y >= X)
    #   :family   ('f' suffix)  - same-major-family-portable; gates 'f'-tier intrinsics
    #   :arch     ('a' suffix)  - locked to one exact CC; unlocks all arch-accel intrinsics
    feature_set::Symbol = :baseline

    # codegen quirks
    ## can we emit debug info in the PTX assembly?
    debuginfo::Bool = false

    # optional properties
    minthreads::Union{Nothing,Int,NTuple{<:Any,Int}} = nothing
    maxthreads::Union{Nothing,Int,NTuple{<:Any,Int}} = nothing
    blocks_per_sm::Union{Nothing,Int} = nothing
    maxregs::Union{Nothing,Int} = nothing

    fastmath::Bool = Base.JLOptions().fast_math == 1

    # deprecated; remove with next major version
    exitable::Union{Nothing,Bool} = nothing
    unreachable::Union{Nothing,Bool} = nothing
end

function Base.hash(target::PTXCompilerTarget, h::UInt)
    h = hash(target.cap, h)
    h = hash(target.ptx, h)
    h = hash(target.feature_set, h)

    h = hash(target.debuginfo, h)

    h = hash(target.minthreads, h)
    h = hash(target.maxthreads, h)
    h = hash(target.blocks_per_sm, h)
    h = hash(target.maxregs, h)
    h = hash(target.fastmath, h)

    h
end

# format the LLVM CPU / PTX `.target` name for this target
function cpu_name(target::PTXCompilerTarget)
    suffix = target.feature_set === :arch    ? "a" :
             target.feature_set === :family  ? "f" :
             target.feature_set === :baseline ? "" :
             error("PTXCompilerTarget.feature_set must be one of :baseline, :family, :arch; got $(repr(target.feature_set))")
    return "sm_$(target.cap.major)$(target.cap.minor)$suffix"
end

source_code(target::PTXCompilerTarget) = "ptx"

llvm_triple(target::PTXCompilerTarget) = Int===Int64 ? "nvptx64-nvidia-cuda" : "nvptx-nvidia-cuda"


function llvm_machine(target::PTXCompilerTarget)
    @static if :NVPTX ∉ LLVM.backends()
        return nothing
    end
    triple = llvm_triple(target)
    t = Target(triple=triple)

    # this machine only drives the middle end; machine code is generated by an external,
    # up-to-date back-end that sees the true target. as the in-process LLVM may not even
    # know the targeted device (unknown CPUs noisily degrade to the sm_20 subtarget),
    # use the generic subtarget: subtarget-dependent middle-end behavior is handled by
    # GPUCompiler itself (`nvvm_reflect!`), and differences in cost models are immaterial.
    return TargetMachine(t, triple)
end

# match the layout computed by the NVPTX back-end (computeNVPTXDataLayout); components
# that match LLVM's defaults (small integers, floats, common vectors) are left implicit.
# note that the external back-end ignores the module datalayout and recomputes it from
# the target, so the middle end needs to match it exactly, or offsets folded early would
# disagree with the final layout. this notably affects i128, which Julia did not even
# align to 16 bytes before 1.12 (users may have to reject affected types, as kernel
# argument layouts will differ between such hosts and the device).
llvm_datalayout(target::PTXCompilerTarget) =
    # little endian
    "e-" *
    # on 32-bit systems, use 32-bit pointers, including for distributed shared
    # memory (addrspace 7). on 64-bit systems, use 64-bit pointers (the default).
    (Int === Int64 ? "" :  "p:32:32-p7:32:32-") *
    # tensor memory (addrspace 6) always uses 32-bit pointers
    "p6:32:32-" *
    # alignment of integer types
    "i64:64-i128:128-i256:256-" *
    # alignment of vector types
    "v16:16-v32:32-" *
    # native integer widths
    "n16:32:64"

have_fma(@nospecialize(target::PTXCompilerTarget), T::Type) = true

dwarf_version(target::PTXCompilerTarget) = Int32(2) # Cuda only supports dwarfv2

can_vectorize(job::CompilerJob{PTXCompilerTarget}) = true

## job

function Base.show(io::IO, @nospecialize(job::CompilerJob{PTXCompilerTarget}))
    print(io, "PTX CompilerJob of ", job.source)
    print(io, " for ", cpu_name(job.config.target))

    job.config.target.minthreads !== nothing && print(io, ", minthreads=$(job.config.target.minthreads)")
    job.config.target.maxthreads !== nothing && print(io, ", maxthreads=$(job.config.target.maxthreads)")
    job.config.target.blocks_per_sm !== nothing && print(io, ", blocks_per_sm=$(job.config.target.blocks_per_sm)")
    job.config.target.maxregs !== nothing && print(io, ", maxregs=$(job.config.target.maxregs)")
    job.config.target.fastmath && print(io, ", fast math enabled")
end

const ptx_intrinsics = ("vprintf", "__assertfail", "malloc", "free")
isintrinsic(@nospecialize(job::CompilerJob{PTXCompilerTarget}), fn::String) =
    in(fn, ptx_intrinsics) ||
    # NVVM intrinsics that the in-process LLVM does not know about, e.g. selected by
    # libdevice's __CUDA_ARCH dispatch, are still supported by the external back-end.
    startswith(fn, "llvm.nvvm.")

# XXX: the debuginfo part should be handled by GPUCompiler as it applies to all back-ends.
runtime_slug(@nospecialize(job::CompilerJob{PTXCompilerTarget})) =
    "ptx$(job.config.target.ptx.major)$(job.config.target.ptx.minor)" *
    "-$(cpu_name(job.config.target))" *
    "-debuginfo=$(Int(llvm_debug_info(job)))"

function finish_module!(@nospecialize(job::CompilerJob{PTXCompilerTarget}),
                        mod::LLVM.Module, entry::LLVM.Function)
    # tell NVVMReflect whether to flush denormals; this mirrors what Clang does
    # for `-fcuda-flush-denormals-to-zero` and is the only `__nvvm_reflect` key
    # LLVM's NVVMReflectPass honors besides `__CUDA_ARCH`. only emit it on the
    # toplevel module that runs through `optimize!`, as sub-modules (the cached
    # runtime, deferred jobs) don't need it, and the cached runtime in
    # particular would otherwise conflict on link if it was built with a
    # different `fastmath` setting (which isn't part of `runtime_slug`).
    if job.config.toplevel
        flags(mod)["nvvm-reflect-ftz", LLVM.API.LLVMModuleFlagBehaviorOverride] =
            Metadata(ConstantInt(Int32(job.config.target.fastmath ? 1 : 0)))
    end

    # emit the device capability and ptx isa version as constants in the module. this makes
    # it possible to 'query' these in device code, relying on LLVM to optimize the checks
    # away and generate static code. note that we only do so if there's actual uses of these
    # variables; unconditionally creating a gvar would result in duplicate declarations.
    sm_features = job.config.target.feature_set === :arch    ? ArchFeatures :
                  job.config.target.feature_set === :family  ? FamilyFeatures :
                                                               BaselineFeatures
    for (name, value) in ["sm_major"    => job.config.target.cap.major,
                          "sm_minor"    => job.config.target.cap.minor,
                          "sm_features" => UInt32(sm_features),
                          "ptx_major"   => job.config.target.ptx.major,
                          "ptx_minor"   => job.config.target.ptx.minor]
        if haskey(globals(mod), name)
            gv = globals(mod)[name]
            initializer!(gv, ConstantInt(LLVM.Int32Type(), value))
            # change the linkage so that we can inline the value
            linkage!(gv, LLVM.API.LLVMPrivateLinkage)
        end
    end

    # update calling convention
    if LLVM.version() >= v"8"
        for f in functions(mod)
            # JuliaGPU/GPUCompiler.jl#97
            #callconv!(f, LLVM.API.LLVMPTXDeviceCallConv)
        end
    end
    if job.config.kernel && LLVM.version() >= v"8"
        callconv!(entry, LLVM.API.LLVMPTXKernelCallConv)
    end

    if job.config.kernel
        # work around bad byval codegen (JuliaGPU/GPUCompiler.jl#92)
        entry = lower_byval(job, mod, entry)

        # emit kernel property annotations into the module. these have to be in
        # place before optimization runs: LLVM's NVPTX target machine registers a
        # PipelineStart EP callback that schedules NVVMIntrRangePass, which calls
        # `getMaxNTID` on every function. That populates a module-keyed
        # `AnnotationCache` entry (empty, because `nvvm.annotations` isn't there
        # yet), and subsequent lookups by the asm printer hit the stale empty
        # entry instead of re-reading the metadata.
        annotations = Metadata[entry]

        # note that the entry is only marked as a kernel through its ptx_kernel calling
        # convention; the "kernel" nvvm.annotation is redundant (the back-end upgrades
        # it to the calling convention), and on LLVM >= 20 even causes miscompilations.

        # launch bounds are emitted in two forms: as nvvm.annotations metadata, consumed
        # by the in-process LLVM's passes (and auto-upgraded by the back-end), and as
        # function attributes, the form modern LLVM consumes directly.

        # expected CTA sizes
        if job.config.target.minthreads !== nothing
            bounds = ntuple(i -> i <= length(job.config.target.minthreads) ?
                                 job.config.target.minthreads[i] : 1, 3)
            for (bound, name) in zip(bounds, (:x, :y, :z))
                append!(annotations, [MDString("reqntid$name"),
                                      ConstantInt(Int32(bound))])
            end
            push!(function_attributes(entry),
                  StringAttribute("nvvm.reqntid", join(bounds, ",")))
        end
        if job.config.target.maxthreads !== nothing
            bounds = ntuple(i -> i <= length(job.config.target.maxthreads) ?
                                 job.config.target.maxthreads[i] : 1, 3)
            for (bound, name) in zip(bounds, (:x, :y, :z))
                append!(annotations, [MDString("maxntid$name"),
                                      ConstantInt(Int32(bound))])
            end
            push!(function_attributes(entry),
                  StringAttribute("nvvm.maxntid", join(bounds, ",")))
        end

        if job.config.target.blocks_per_sm !== nothing
            append!(annotations, [MDString("minctasm"),
                                  ConstantInt(Int32(job.config.target.blocks_per_sm))])
            push!(function_attributes(entry),
                  StringAttribute("nvvm.minctasm", string(job.config.target.blocks_per_sm)))
        end

        if job.config.target.maxregs !== nothing
            append!(annotations, [MDString("maxnreg"),
                                  ConstantInt(Int32(job.config.target.maxregs))])
            push!(function_attributes(entry),
                  StringAttribute("nvvm.maxnreg", string(job.config.target.maxregs)))
        end

        if length(annotations) > 1
            push!(metadata(mod)["nvvm.annotations"], MDNode(annotations))
        end
    end

    # we emit properties (of the device and ptx isa) as private global constants,
    # so run the optimizer so that they are inlined before the rest of the optimizer runs.
    @dispose pb=NewPMPassBuilder() begin
        add!(pb, RecomputeGlobalsAAPass())
        add!(pb, GlobalOptPass())
        run!(pb, mod, llvm_machine(job.config.target))
    end

    return entry
end

function finish_linked_module!(@nospecialize(job::CompilerJob{PTXCompilerTarget}),
                               mod::LLVM.Module)
    # fold `__nvvm_reflect` calls, now that libdevice has been linked. we do this
    # ourselves, and not from the in-process LLVM's NVVMReflectPass (scheduled by the
    # NVPTX TargetMachine's pipeline start callbacks since LLVM 17), because that pass
    # derives `__CUDA_ARCH` from the middle-end subtarget, which cannot represent
    # devices newer than the in-process LLVM. by folding before optimization, at the
    # true target, the built-in pass is left with nothing to do.
    nvvm_reflect!(job, mod)

    # propagate `target.fastmath` as `@fastmath`-everywhere semantics
    # (mirrors nvcc's `--use_fast_math`). post-link so that bodies pulled in
    # from libdevice and the runtime also get the flags.
    if job.config.target.fastmath
        apply_fastmath!(mod)
        # additionally request FTZ on f32: NVPTX' `useF32FTZ` reads
        # `denormal-fp-math-f32` to pick the FTZ variants for
        # fdiv/fsqrt/etc.
        for f in functions(mod)
            isdeclaration(f) && continue
            push!(function_attributes(f),
                  StringAttribute("denormal-fp-math-f32",
                                  "preserve-sign,preserve-sign"))
        end
    end
    return
end

function optimize_module!(@nospecialize(job::CompilerJob{PTXCompilerTarget}),
                          mod::LLVM.Module)
    tm = llvm_machine(job.config.target)
    # TODO: Use the registered target passes (JuliaGPU/GPUCompiler.jl#450)
    @dispose pb=NewPMPassBuilder() begin
        register!(pb, PTXRSqrtFastPass())
        register!(pb, PTXFDivFastPass())
        register!(pb, PTXFSqrtFastPass())
        if get(optimization_options(job), :fastmath, true)
            add!(pb, PTXRSqrtFastPass())
            add!(pb, PTXFDivFastPass())
            add!(pb, PTXFSqrtFastPass())
        end

        add!(pb, NewPMFunctionPassManager()) do fpm
            # needed by GemmKernels.jl-like code
            add!(fpm, SpeculativeExecutionPass())

            # NVPTX's target machine info enables runtime unrolling,
            # but Julia's pass sequence only invokes the simple unroller.
            add!(fpm, LoopUnrollPass(; job.config.opt_level))
            add!(fpm, instcombine_pass(job))        # clean-up redundancy
            add!(fpm, NewPMLoopPassManager(; use_memory_ssa=true)) do lpm
                add!(lpm, LICMPass())           # the inner runtime check might be
                                                # outer loop invariant
            end

            # the above loop unroll pass might have unrolled regular, non-runtime nested loops.
            # that code still needs to be optimized (arguably, multiple unroll passes should be
            # scheduled by the Julia optimizer). do so here, instead of re-optimizing entirely.
            if job.config.opt_level == 2
                add!(fpm, GVNPass())
            elseif job.config.opt_level == 1
                add!(fpm, EarlyCSEPass())
            end
            add!(fpm, DSEPass())

            add!(fpm, SimplifyCFGPass())
        end

        # get rid of the internalized functions; now possible unused
        add!(pb, GlobalDCEPass())

        run!(pb, mod, tm)
    end
end

@unlocked function mcgen(@nospecialize(job::CompilerJob{PTXCompilerTarget}),
                         mod::LLVM.Module, format=LLVM.API.LLVMAssemblyFile)
    if !isavailable(NVPTX_LLVM_Backend_jll) || !NVPTX_LLVM_Backend_jll.is_available()
        error("NVPTX LLVM back-end not loaded; cannot compile to PTX.")
    end

    target = job.config.target
    filetype = if format == LLVM.API.LLVMAssemblyFile
        "asm"
    elseif format == LLVM.API.LLVMObjectFile
        "obj"
    else
        error("Unsupported PTX output format $format")
    end

    input  = tempname(cleanup=false) * ".bc"
    output = tempname(cleanup=false) * (filetype == "asm" ? ".ptx" : ".cubin")
    write(input, mod)

    cmd = `$(NVPTX_LLVM_Backend_jll.llc()) $input
              -mtriple=$(llvm_triple(target))
              -mcpu=$(cpu_name(target))
              -mattr=+ptx$(target.ptx.major)$(target.ptx.minor)
              -filetype=$filetype
              -o $output`
    out = Pipe()
    proc = run(pipeline(ignorestatus(cmd); stdout=out, stderr=out); wait=false)
    close(out.in)
    log = strip(read(out, String))
    wait(proc)
    if !success(proc)
        # keep the input around for debugging
        msg = "Failed to compile to PTX with external llc"
        isempty(log) || (msg *= ":\n" * log)
        msg *= "\nIf you think this is a bug, please file an issue and attach $(input)."
        isfile(output) && rm(output)
        error(msg)
    elseif !isempty(log)
        # llc only diagnoses on stderr; even successful compilation may e.g. have
        # ignored an unrecognized CPU or feature, so make sure this surfaces.
        @warn "External llc reported:\n$log"
    end

    code = filetype == "asm" ? read(output, String) : String(read(output))
    rm(input)
    rm(output)
    return code
end

function llvm_debug_info(@nospecialize(job::CompilerJob{PTXCompilerTarget}))
    # allow overriding the debug info from CUDA.jl
    if job.config.target.debuginfo
        invoke(llvm_debug_info, Tuple{CompilerJob}, job)
    else
        LLVM.API.LLVMDebugEmissionKindNoDebug
    end
end


## LLVM passes

# Replace occurrences of __nvvm_reflect("foo") and llvm.nvvm.reflect with an integer.
#
# This matches LLVM's NVVMReflectPass: `__CUDA_ARCH` is derived from the target
# capability, `__CUDA_FTZ` is read from the `nvvm-reflect-ftz` module flag, and every
# other key folds to 0. Knobs like denormal flushing or FMAD contraction must be
# configured through module flags or LLVM fast-math flags, not here.
const NVVM_REFLECT_FUNCTION = "__nvvm_reflect"
function nvvm_reflect!(@nospecialize(job::CompilerJob), mod::LLVM.Module)
    changed = false
    @tracepoint "nvvmreflect" begin

    # find and sanity check the nnvm-reflect function
    # TODO: also handle the llvm.nvvm.reflect intrinsic
    haskey(LLVM.functions(mod), NVVM_REFLECT_FUNCTION) || return false
    reflect_function = functions(mod)[NVVM_REFLECT_FUNCTION]
    isdeclaration(reflect_function) || error("_reflect function should not have a body")
    reflect_typ = return_type(function_type(reflect_function))
    isa(reflect_typ, LLVM.IntegerType) || error("_reflect's return type should be integer")

    # pull __CUDA_FTZ from the nvvm-reflect-ftz module flag (same source LLVM uses)
    ftz_val = 0
    if haskey(flags(mod), "nvvm-reflect-ftz")
        flag = flags(mod)["nvvm-reflect-ftz"]
        if flag isa LLVM.ConstantAsMetadata
            c = LLVM.Value(flag)
            if c isa ConstantInt
                ftz_val = Int(convert(Int64, c))
            end
        end
    end

    to_remove = []
    for use in uses(reflect_function)
        call = user(use)
        isa(call, LLVM.CallInst) || continue
        if length(operands(call)) != 2
            @error """Unrecognized format of __nvvm_reflect call:
                      $(string(call))
                      Wrong number of operands: expected 2, got $(length(operands(call)))."""
            continue
        end

        # decode the string argument
        if LLVM.version() >= v"17"
            sym = operands(call)[1]
        else
            str = operands(call)[1]
            if !isa(str, LLVM.ConstantExpr) || opcode(str) != LLVM.API.LLVMGetElementPtr
                @safe_error """Unrecognized format of __nvvm_reflect call:
                               $(string(call))
                               Operand should be a GEP instruction, got a $(typeof(str)). Please file an issue."""
                continue
            end
            sym = operands(str)[1]
            if isa(sym, LLVM.ConstantExpr) && opcode(sym) == LLVM.API.LLVMGetElementPtr
                # CUDA 11.0 or below
                sym = operands(sym)[1]
            end
        end
        if !isa(sym, LLVM.GlobalVariable)
            @safe_error """Unrecognized format of __nvvm_reflect call:
                           $(string(call))
                           Operand should be a global variable, got a $(typeof(sym)). Please file an issue."""
            continue
        end
        sym_op = operands(sym)[1]
        if !isa(sym_op, LLVM.ConstantArray) && !isa(sym_op, LLVM.ConstantDataArray)
            @safe_error """Unrecognized format of __nvvm_reflect call:
                           $(string(call))
                           Operand should be a constant array, got a $(typeof(sym_op)). Please file an issue."""
        end
        chars = convert.(Ref(UInt8), collect(sym_op))
        reflect_arg = String(chars[1:end-1])

        # match LLVM's NVVMReflectPass: unknown keys fold to 0.
        reflect_val = if reflect_arg == "__CUDA_ARCH"
            ConstantInt(reflect_typ,
                        job.config.target.cap.major*100 + job.config.target.cap.minor*10)
        elseif reflect_arg == "__CUDA_FTZ"
            ConstantInt(reflect_typ, ftz_val)
        else
            ConstantInt(reflect_typ, 0)
        end

        replace_uses!(call, reflect_val)
        push!(to_remove, call)
        changed = true
    end

    # remove the calls to the function
    for val in to_remove
        @assert isempty(uses(val))
        erase!(val)
    end

    # maybe also delete the function
    if isempty(uses(reflect_function))
        erase!(reflect_function)
    end

    end
    return changed
end

# The passes below rewrite `afn`-flagged f64 operations to NVPTX' fast lowerings.
# `apply_fastmath!` propagates job-wide `target.fastmath=true` as per-instruction
# `afn`, so a single flag check covers both per-call `@fastmath` and the job toggle.
# We emit NVPTX intrinsics by name (rather than libdevice `__nv_*`) so this doesn't
# depend on which libdevice symbols got linked in.
#
# Only f64 needs this: the back-end natively honors `afn` for f32 fdiv/sqrt/rsqrt
# (selecting the FTZ variants based on the `denormal-fp-math-f32` attribute), but
# lacks fast f64 lowerings. There is no hardware fast f64 div/sqrt, and the f64
# `rsqrt.approx.d` instruction is not selected for `1/sqrt(x)` upstream (NVPTX'
# `getSqrtEstimate` does not consult per-instruction `afn`, only the
# "reciprocal-estimates" function attribute).
#
# `PTXRSqrtFastPass` must run first: the rsqrt pattern (`fdiv afn 1.0, sqrt afn x`)
# spans an fdiv and a sqrt, so it has to claim both operands before the per-op
# passes below eat them.

# Rewrite `fdiv afn 1.0, sqrt afn(x)` on f64 to `nvvm.rsqrt.approx.d(x)`.
function ptx_rsqrt_fast!(mod::LLVM.Module)
    changed = false
    @tracepoint "ptx-rsqrt-fast" begin

    f64 = LLVM.DoubleType()

    # collect first to avoid mutation-during-iteration
    to_replace = Tuple{LLVM.FDivInst, LLVM.CallInst}[]
    for f in functions(mod), bb in blocks(f), inst in instructions(bb)
        inst isa LLVM.FDivInst || continue
        LLVM.value_type(inst) == f64 || continue
        LLVM.fast_math(inst).afn || continue

        # numerator must be the constant 1.0
        lhs = operands(inst)[1]
        lhs isa LLVM.ConstantFP || continue
        convert(Float64, lhs) == 1.0 || continue

        # denominator must be an `afn`-flagged `llvm.sqrt.f64` call
        rhs = operands(inst)[2]
        rhs isa LLVM.CallInst || continue
        callee = LLVM.called_operand(rhs)
        callee isa LLVM.Function || continue
        LLVM.name(callee) == "llvm.sqrt.f64" || continue
        LLVM.fast_math(rhs).afn || continue

        push!(to_replace, (inst, rhs))
    end
    isempty(to_replace) && return false

    # declare intrinsics by name so LLVM keeps the exact non-overloaded names;
    # LLVM.Intrinsic + type params would mangle to *.f64, unrecognized by NVPTX.
    fns = functions(mod)
    declare(name, ft) = haskey(fns, name) ? fns[name] : LLVM.Function(mod, name, ft)
    f64_ft = LLVM.FunctionType(f64, [f64])
    rsqrt_f64 = declare("llvm.nvvm.rsqrt.approx.d", f64_ft)

    @dispose builder=IRBuilder() begin
        for (fdiv, sqrt_call) in to_replace
            x = operands(sqrt_call)[1]
            position!(builder, fdiv)

            replacement = call!(builder, f64_ft, rsqrt_f64, [x])

            replace_uses!(fdiv, replacement)
            erase!(fdiv)
            # sqrt may still be used elsewhere; only clean it up if dead now.
            if isempty(uses(sqrt_call))
                erase!(sqrt_call)
            end
            changed = true
        end
    end

    end # @tracepoint
    return changed
end
PTXRSqrtFastPass() = NewPMModulePass("ptx-rsqrt-fast", ptx_rsqrt_fast!)

# Rewrite `afn`-flagged f64 `fdiv` to `rcp.approx.ftz.d` + one Newton step
# (there is no native fast f64 fdiv).
function ptx_fdiv_fast!(mod::LLVM.Module)
    changed = false
    @tracepoint "ptx-fdiv-fast" begin

    f64 = LLVM.DoubleType()

    # collect first to avoid mutation-during-iteration
    to_replace = LLVM.FDivInst[]
    for f in functions(mod), bb in blocks(f), inst in instructions(bb)
        inst isa LLVM.FDivInst || continue
        LLVM.value_type(inst) == f64 || continue
        LLVM.fast_math(inst).afn || continue
        push!(to_replace, inst)
    end
    isempty(to_replace) && return false

    fns = functions(mod)
    declare(name, ft) = haskey(fns, name) ? fns[name] : LLVM.Function(mod, name, ft)
    f64_ft1 = LLVM.FunctionType(f64, [f64])
    rcp_f64 = declare("llvm.nvvm.rcp.approx.ftz.d", f64_ft1)
    fma_ft  = LLVM.FunctionType(f64, [f64, f64, f64])
    fma_f64 = declare("llvm.fma.f64", fma_ft)
    one_f64 = ConstantFP(f64, 1.0)

    @dispose builder=IRBuilder() begin
        for inst in to_replace
            lhs, rhs = operands(inst)[1], operands(inst)[2]
            position!(builder, inst)

            inv_y   = call!(builder, f64_ft1, rcp_f64, [rhs])
            neg_rhs = fneg!(builder, rhs)
            # Newton refinement, matching CUDA.jl's `FastMath.inv_fast(::Float64)`
            e       = call!(builder, fma_ft, fma_f64, [inv_y, neg_rhs, one_f64])
            e       = call!(builder, fma_ft, fma_f64, [e, e, e])
            inv_ref = call!(builder, fma_ft, fma_f64, [e, inv_y, inv_y])
            replacement = fmul!(builder, lhs, inv_ref)

            replace_uses!(inst, replacement)
            erase!(inst)
            changed = true
        end
    end

    end # @tracepoint
    return changed
end
PTXFDivFastPass() = NewPMModulePass("ptx-fdiv-fast", ptx_fdiv_fast!)

# Rewrite `afn`-flagged `llvm.sqrt.f64` to `rcp.approx.ftz.d(rsqrt.approx.d(x))`
# (there is no native fast f64 sqrt; this matches NVPTX' `getSqrtEstimate`).
function ptx_fsqrt_fast!(mod::LLVM.Module)
    changed = false
    @tracepoint "ptx-fsqrt-fast" begin

    f64 = LLVM.DoubleType()

    to_replace = LLVM.CallInst[]
    for f in functions(mod), bb in blocks(f), inst in instructions(bb)
        inst isa LLVM.CallInst || continue
        callee = LLVM.called_operand(inst)
        callee isa LLVM.Function || continue
        LLVM.name(callee) == "llvm.sqrt.f64" || continue
        LLVM.fast_math(inst).afn || continue
        push!(to_replace, inst)
    end
    isempty(to_replace) && return false

    fns = functions(mod)
    declare(name, ft) = haskey(fns, name) ? fns[name] : LLVM.Function(mod, name, ft)
    f64_ft = LLVM.FunctionType(f64, [f64])
    rcp_f64   = declare("llvm.nvvm.rcp.approx.ftz.d", f64_ft)
    rsqrt_f64 = declare("llvm.nvvm.rsqrt.approx.d",   f64_ft)

    @dispose builder=IRBuilder() begin
        for inst in to_replace
            x = operands(inst)[1]
            position!(builder, inst)

            rsqrt = call!(builder, f64_ft, rsqrt_f64, [x])
            replacement = call!(builder, f64_ft, rcp_f64, [rsqrt])

            replace_uses!(inst, replacement)
            erase!(inst)
            changed = true
        end
    end

    end # @tracepoint
    return changed
end
PTXFSqrtFastPass() = NewPMModulePass("ptx-fsqrt-fast", ptx_fsqrt_fast!)
