XFESystem

using System.Buffers.Binary;
using System.Diagnostics;
using System.Security.Cryptography;
using System.Text;

namespace XFESystem.DriverOracle;

internal sealed unsafe class HostExecution
{
    public bool ContextCreated { get; private set; }
    public bool KernelSubmitted { get; private set; }
    public bool SynchronizationPassed { get; private set; }
    public bool ResultAndGuardsPassed { get; private set; }
    public bool CleanupPassed { get; private set; }
    public bool ContextDestroyed { get; private set; }
    public bool OriginalContextRestored { get; private set; }
    public bool ActualHostExecutionPassed { get; private set; }
    public long DeviceAllocationBytes { get; private set; }
    public string? Error { get; private set; }
    public string? CubinSha256 { get; private set; }
    public int CubinBytes { get; private set; }
    public string? CubinElfFlags { get; private set; }
    public uint? CubinSm { get; private set; }
    public uint Seed { get; private set; }
    public object? Verification { get; private set; }
    public double ExecutionMilliseconds { get; private set; }
    public List<object> Allocations { get; } = [];
    public uint GridBlocks => (VectorCase.Count + VectorCase.BlockThreads - 1) / VectorCase.BlockThreads;
    public int ThreadsPerBlock => VectorCase.BlockThreads;
    public string JitTarget => "CU_JIT_TARGET=120 / sm_120";
    public string Kernel => "XfeVectorAdd: unsigned 32-bit addition modulo 2^32";

    public void Run(CudaDriver driver, CudaDevice device, byte[] ptx, string output)
    {
        nint context = 0, prior = 0, module = 0, link = 0;
        var owned = new List<ulong>(); bool moduleOwned = false, linkOwned = false;
        byte[] info = new byte[8192], errors = new byte[8192];
        var timer = Stopwatch.StartNew();
        try
        {
            driver.Check("cuCtxGetCurrent(before)", driver.Api<CudaDriver.PointerOutD>("cuCtxGetCurrent")(out prior));
            if (prior != 0) throw new InvalidOperationException("The dedicated oracle process unexpectedly already owns a CUDA context.");
            driver.Check("cuCtxCreate_v2", driver.Api<CudaDriver.ContextCreateD>("cuCtxCreate_v2")(out context, 4 /* CU_CTX_SCHED_BLOCKING_SYNC */, device.Handle));
            ContextCreated = true;
            if (context == 0) throw new InvalidOperationException("CUDA returned a null owned context.");

            byte[] terminated = new byte[ptx.Length + 1]; ptx.CopyTo(terminated, 0);
            byte[] cubin;
            fixed (byte* infoPointer = info, errorPointer = errors, ptxPointer = terminated)
            {
                // Values of these CUDA ABI enums are fixed by the public Driver
                // API: INFO_LOG_BUFFER=3, INFO_SIZE=4, ERROR_BUFFER=5,
                // ERROR_SIZE=6, TARGET=9, CU_TARGET_COMPUTE_120=120.
                int* options = stackalloc int[5] { 3, 4, 5, 6, 9 };
                nint* values = stackalloc nint[5] { (nint)infoPointer, info.Length, (nint)errorPointer, errors.Length, 120 };
                driver.Check("cuLinkCreate_v2", driver.Api<CudaDriver.LinkCreateD>("cuLinkCreate_v2")(5, options, values, out link)); linkOwned = true;
                if (link == 0) throw new InvalidOperationException("CUDA returned a null linker state.");
                driver.Check("cuLinkAddData_v2(PTX)", driver.Api<CudaDriver.LinkAddD>("cuLinkAddData_v2")(link, 1 /* CU_JIT_INPUT_PTX */, ptxPointer, (nuint)terminated.Length, "XfeVectorAdd.ptx", 0, null, null));
                driver.Check("cuLinkComplete", driver.Api<CudaDriver.LinkCompleteD>("cuLinkComplete")(link, out nint image, out nuint length));
                if (image == 0 || length < 64 || length > 2 * 1024 * 1024) throw new InvalidOperationException("The generated CUDA image exceeds its bounded ELF geometry.");
                cubin = new ReadOnlySpan<byte>((void*)image, (int)length).ToArray();
            }
            CubinBytes = cubin.Length; CubinSha256 = Program.Hash(cubin);
            if (!cubin.AsSpan(0, 7).SequenceEqual(new byte[] { 0x7F, (byte)'E', (byte)'L', (byte)'F', 2, 1, 1 }) ||
                BinaryPrimitives.ReadUInt16LittleEndian(cubin.AsSpan(18)) != 190 /* EM_CUDA */)
                throw new InvalidOperationException("The real Driver JIT did not return a little-endian CUDA ELF64 image.");
            uint flags = BinaryPrimitives.ReadUInt32LittleEndian(cubin.AsSpan(48));
            // Keep ELF flags raw: CUDA ELF ABI revisions use different bit
            // layouts. The actual loaded function's Driver API attribute below
            // is the architecture gate, rather than a guessed mask.
            CubinElfFlags = "0x" + flags.ToString("x8");
            Program.Save(output, "XfeVectorAdd-sm120.cubin", cubin);
            driver.Check("cuLinkDestroy", driver.Api<CudaDriver.PointerD>("cuLinkDestroy")(link)); linkOwned = false;
            fixed (byte* image = cubin) driver.Check("cuModuleLoadData", driver.Api<CudaDriver.ModuleLoadD>("cuModuleLoadData")(out module, image));
            moduleOwned = true;
            driver.Check("cuModuleGetFunction", driver.Api<CudaDriver.FunctionD>("cuModuleGetFunction")(out nint function, module, "XfeVectorAdd"));
            if (module == 0 || function == 0) throw new InvalidOperationException("CUDA returned a null module/function.");
            driver.Check("cuFuncGetAttribute(BINARY_VERSION)", driver.Api<CudaDriver.FunctionAttributeD>("cuFuncGetAttribute")(out int binaryVersion, 6 /* CU_FUNC_ATTRIBUTE_BINARY_VERSION */, function));
            CubinSm = checked((uint)binaryVersion);
            if (CubinSm != 120) throw new InvalidOperationException("The real loaded function is not SM120.");

            Seed = BinaryPrimitives.ReadUInt32LittleEndian(RandomNumberGenerator.GetBytes(4));
            var vector = new VectorCase(Seed);
            ulong a = Allocate(vector.A.Length), b = Allocate(vector.B.Length), guarded = Allocate(vector.Initial.Length);
            DeviceAllocationBytes = vector.DeviceBytes;
            Upload(a, vector.A); Upload(b, vector.B); Upload(guarded, vector.Initial);
            ulong target = checked(guarded + VectorCase.GuardWords * sizeof(uint)); uint count = VectorCase.Count;
            void** parameters = stackalloc void*[4] { &a, &b, &target, &count };
            driver.Check("cuLaunchKernel", driver.Api<CudaDriver.LaunchD>("cuLaunchKernel")(function, GridBlocks, 1, 1, VectorCase.BlockThreads, 1, 1, 0, 0, parameters, null));
            KernelSubmitted = true;
            driver.Check("cuCtxSynchronize", () => driver.Api<CudaDriver.NoArgsD>("cuCtxSynchronize")()); SynchronizationPassed = true;
            uint[] actual = new uint[vector.Initial.Length];
            fixed (uint* actualPointer = actual)
                driver.Check("cuMemcpyDtoH_v2", driver.Api<CudaDriver.DownloadD>("cuMemcpyDtoH_v2")(actualPointer, guarded, checked((nuint)actual.Length * sizeof(uint))));
            Verification = vector.Verify(actual); ResultAndGuardsPassed = true;

            ulong Allocate(int words)
            {
                nuint bytes = checked((nuint)words * sizeof(uint));
                if ((ulong)bytes + sizes.Values.Aggregate(0UL, (total, value) => checked(total + value)) >= (ulong)VectorCase.MaximumDeviceBytes)
                    throw new InvalidOperationException("The owned allocation budget would be exceeded.");
                driver.Check("cuMemAlloc_v2", driver.Api<CudaDriver.AllocateD>("cuMemAlloc_v2")(out ulong address, bytes));
                if (address == 0) throw new InvalidOperationException("CUDA returned a null device allocation.");
                owned.Add(address); sizes.Add(address, (ulong)bytes);
                Allocations.Add(new { Address = "0x" + address.ToString("x"), Bytes = (ulong)bytes });
                return address;
            }
            void Upload(ulong address, uint[] data)
            {
                fixed (uint* pointer = data)
                    driver.Check("cuMemcpyHtoD_v2", driver.Api<CudaDriver.UploadD>("cuMemcpyHtoD_v2")(address, pointer, checked((nuint)data.Length * sizeof(uint))));
            }
        }
        catch (Exception error) { Error = error.ToString(); }
        finally
        {
            bool cleanup = true;
            // These are this context's allocations only. Never reset a device,
            // primary context, host graphics driver, or somebody else's memory.
            for (int index = owned.Count - 1; index >= 0; index--)
            {
                ulong address = owned[index];
                cleanup &= driver.Cleanup("cuMemFree_v2", () => driver.Api<CudaDriver.FreeD>("cuMemFree_v2")(address));
            }
            if (moduleOwned) cleanup &= driver.Cleanup("cuModuleUnload", () => driver.Api<CudaDriver.PointerD>("cuModuleUnload")(module));
            if (linkOwned) cleanup &= driver.Cleanup("cuLinkDestroy(cleanup)", () => driver.Api<CudaDriver.PointerD>("cuLinkDestroy")(link));
            if (ContextCreated)
            {
                ContextDestroyed = driver.Cleanup("cuCtxDestroy_v2", () => driver.Api<CudaDriver.PointerD>("cuCtxDestroy_v2")(context));
                cleanup &= ContextDestroyed;
                nint current = -1;
                bool read = driver.Cleanup("cuCtxGetCurrent(after)", () => driver.Api<CudaDriver.PointerOutD>("cuCtxGetCurrent")(out current));
                OriginalContextRestored = read && current == prior; cleanup &= OriginalContextRestored;
            }
            CleanupPassed = cleanup;
            ActualHostExecutionPassed = ContextCreated && KernelSubmitted && SynchronizationPassed && ResultAndGuardsPassed && ContextDestroyed && OriginalContextRestored && cleanup && Error is null;
            ExecutionMilliseconds = timer.Elapsed.TotalMilliseconds;
            Program.Save(output, "JitInfo.log", Encoding.UTF8.GetBytes(Encoding.UTF8.GetString(info).TrimEnd('\0')));
            Program.Save(output, "JitError.log", Encoding.UTF8.GetBytes(Encoding.UTF8.GetString(errors).TrimEnd('\0')));
        }
    }

    private readonly Dictionary<ulong, ulong> sizes = [];
}