PublicFunctions.cpp

/*
 * SPDX-FileCopyrightText: Copyright (c) 2023-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 * SPDX-License-Identifier: LicenseRef-NvidiaProprietary
 *
 * NVIDIA CORPORATION, its affiliates and licensors retain all intellectual
 * property and proprietary rights in and to this material, related
 * documentation and any modifications thereto. Any use, reproduction,
 * disclosure or distribution of this material and related documentation
 * without an express license agreement from NVIDIA CORPORATION or
 * its affiliates is strictly prohibited.
 */

#include "Context.h"
#include "Errors.h"
#include "FeatureGridMath.h"
#include "KnownLatentShapes.h"
#include "MlpDesc.h"
#include "TextureSetMetadata.h"
#include "VersionInfo.h" // Generated by CMake in the build directory
#include <cmath>
#include <limits>
#include <libntc/shaders/BlockCompressConstants.h>

// This file contains implementations for the public functions declared in ntc.h
// that do not belong to any class.

namespace ntc
{

uint32_t GetInterfaceVersion()
{
    return InterfaceVersion;
}

VersionInfo GetLibraryVersion()
{
    VersionInfo info;
    info.major = NTC_VERSION_MAJOR;
    info.minor = NTC_VERSION_MINOR;
    info.point = NTC_VERSION_POINT;
    info.branch = NTC_VERSION_BRANCH;
    info.commitHash = NTC_VERSION_HASH;
    return info;
}

class DefaultAllocator : public IAllocator
{
public:
    void* Allocate(size_t size) override
    {
        return malloc(size);
    }

    void Deallocate(void* ptr, size_t size) override
    {
        free(ptr);
    }
};

static DefaultAllocator g_DefaultAllocator;

Status CreateContext(IContext** pOutContext, ContextParameters const& params)
{
    ClearErrorMessage();

    if (params.interfaceVersion != InterfaceVersion)
    {
        SetErrorMessage("Invalid interface version: provided 0x%08x, expected 0x%08x.",
            params.interfaceVersion, InterfaceVersion);
        return Status::InterfaceMismatch;
    }

    if (!pOutContext)
    {
        SetErrorMessage("pOutContext is NULL.");
        return Status::InvalidArgument;
    }

    ContextParameters paramsCopy = params;

    bool cudaAvailable = false;

#if NTC_WITH_CUDA
    if (params.cudaDevice >= 0)
    {
        cudaDeviceProp prop{};
        cudaError_t err = cudaGetDeviceProperties(&prop, params.cudaDevice);
        
        if (err != cudaSuccess)
        {
            SetErrorMessage("Invalid CUDA device %d, error code = %s.",
                params.cudaDevice, cudaGetErrorName(err));
        }
        else if (prop.major < NTC_CUDA_MIN_COMPUTE_MAJOR ||
                 prop.major == NTC_CUDA_MIN_COMPUTE_MAJOR && prop.minor < NTC_CUDA_MIN_COMPUTE_MINOR)
        {
            SetErrorMessage("The CUDA device compute capability (%d.%d) is less than the minimum required "
                "by LibNTC (%d.%d).", prop.major, prop.minor, NTC_CUDA_MIN_COMPUTE_MAJOR, NTC_CUDA_MIN_COMPUTE_MINOR);
        }
        else
        {
            err = cudaInitDevice(params.cudaDevice, 0, 0);
            if (err == cudaSuccess)
            {
                cudaAvailable = true;
            }
            else
            {
                SetCudaErrorMessage("cudaInitDevice", err);
            }
        }
    }
#endif

    if (!cudaAvailable)
        paramsCopy.cudaDevice = -1;

    // Validate the graphicsAPI value and the device objects
    switch (params.graphicsApi)
    {
        case GraphicsAPI::None:
            break;
            
        case GraphicsAPI::D3D12:
#if NTC_WITH_DX12
            if (!params.d3d12Device)
            {
                SetErrorMessage("d3d12Device is NULL.");
                return Status::InvalidArgument;
            }
            break;
#else
            SetErrorMessage("This version of LibNTC doesn't support D3D12.");
            return Status::NotImplemented;
#endif

        case GraphicsAPI::Vulkan:
#if NTC_WITH_VULKAN
            if (!params.vkDevice || !params.vkInstance || !params.vkPhysicalDevice)
            {
                SetErrorMessage("One of (vkDevice, vkInstance, vkPhysicalDevice) is NULL.");
                return Status::InvalidArgument;
            }
#else
            SetErrorMessage("This version of LibNTC doesn't support Vulkan.");
            return Status::NotImplemented;
#endif
            break;

        default:
            SetErrorMessage("Invalid graphicsApi (%d).", int(params.graphicsApi));
            return Status::InvalidArgument;
    }

    if (!paramsCopy.pAllocator)
        paramsCopy.pAllocator = &g_DefaultAllocator;

    Context* context = new(paramsCopy.pAllocator->Allocate(sizeof(Context))) Context(paramsCopy);
    *pOutContext = context;
    
    return cudaAvailable ? Status::Ok : Status::CudaUnavailable;
}

void DestroyContext(IContext* context)
{
    if (!context)
        return;

    auto allocator = ((Context*)context)->GetAllocator();
    context->~IContext();
    allocator->Deallocate(context, sizeof(Context));
}

const char* GetLastErrorMessage()
{
    return GetErrorBuffer();
}

const char* StatusToString(Status status)
{
    switch(status)
    {
    case Status::Ok:
        return "Ok";
    case Status::InterfaceMismatch:
        return "InterfaceMismatch";
    case Status::Incomplete:
        return "Incomplete";
    case Status::Unsupported:
        return "Unsupported";
    case Status::InvalidArgument:
        return "InvalidArgument";
    case Status::InvalidData:
        return "InvalidData";
    case Status::InvalidState:
        return "InvalidState";
    case Status::UnknownError:
        return "UnknownError";
    case Status::FileUnavailable:
        return "FileUnavailable";
    case Status::FileUnrecognized:
        return "FileUnrecognized";
    case Status::FileIncompatible:
        return "FileIncompatible";
    case Status::IOError:
        return "IOError";
    case Status::OutOfRange:
        return "OutOfRange";
    case Status::CudaUnavailable:
        return "CudaUnavailable";
    case Status::CudaError:
        return "CudaError";
    case Status::NotImplemented:
        return "NotImplemented";
    case Status::InternalError:
        return "InternalError";
    case Status::OutOfMemory:
        return "OutOfMemory";
    case Status::ShaderUnavailable:
        return "ShaderUnavailable";
    default:
        static char string[16];
        snprintf(string, sizeof(string), "%d", int(status));
        return string;
    }
}

const char* ChannelFormatToString(ChannelFormat format)
{
    switch (format)
    {
    case ChannelFormat::UNKNOWN:
        return "UNKNOWN";
    case ChannelFormat::UNORM8:
        return "UNORM8";
    case ChannelFormat::UNORM16:
        return "UNORM16";
    case ChannelFormat::FLOAT16:
        return "FLOAT16";
    case ChannelFormat::FLOAT32:
        return "FLOAT32";
    case ChannelFormat::UINT32:
        return "UINT32";
    default:
        static char string[16];
        snprintf(string, sizeof(string), "%d", int(format));
        return string;
    }
}

const char* BlockCompressedFormatToString(BlockCompressedFormat format)
{
    switch (format)
    {
    case BlockCompressedFormat::None:
        return "None";
    case BlockCompressedFormat::BC1:
        return "BC1";
    case BlockCompressedFormat::BC2:
        return "BC2";
    case BlockCompressedFormat::BC3:
        return "BC3";
    case BlockCompressedFormat::BC4:
        return "BC4";
    case BlockCompressedFormat::BC5:
        return "BC5";
    case BlockCompressedFormat::BC6:
        return "BC6";
    case BlockCompressedFormat::BC7:
        return "BC7";
    default:
        static char string[16];
        snprintf(string, sizeof(string), "%d", int(format));
        return string;
    }
}

const char* ColorSpaceToString(ColorSpace colorSpace)
{
    switch (colorSpace)
    {
    case ColorSpace::Linear:
        return "Linear";
    case ColorSpace::sRGB:
        return "sRGB";
    case ColorSpace::HLG:
        return "HLG";
    default:
        static char string[16];
        snprintf(string, sizeof(string), "%d", int(colorSpace));
        return string;
    }
}

const char* InferenceWeightTypeToString(InferenceWeightType weightType)
{
    switch(weightType)
    {
    case InferenceWeightType::Unknown:
        return "Unknown";
    case InferenceWeightType::GenericInt8:
        return "GenericInt8";
    case InferenceWeightType::GenericFP8:
        return "GenericFP8";
    case InferenceWeightType::CoopVecFP8:
        return "CoopVecFP8";
    default:
        static char string[16];
        snprintf(string, sizeof(string), "%d", int(weightType));
        return string;
    }
}

const char* CompressionTypeToString(CompressionType type)
{
    switch(type)
    {
    case CompressionType::None:
        return "None";
    case CompressionType::GDeflate:
        return "GDeflate";
    default:
        static char string[16];
        snprintf(string, sizeof(string), "%d", int(type));
        return string;
    }
}

size_t GetBytesPerPixelComponent(ChannelFormat format)
{
    switch (format)
    {
    case ChannelFormat::UNORM8:
        return 1;
    case ChannelFormat::UNORM16:
    case ChannelFormat::FLOAT16:
        return 2;
    case ChannelFormat::FLOAT32:
    case ChannelFormat::UINT32:
        return 4;
    default:
        return 0;
    }
}

float LossToPSNR(float loss)
{
    // Math quirk: when loss is 0, log10f(0) returns inf, which turns into -inf. 
    // Returning positive inf for zero loss makes more sense.
    // Also, this function assumes that max value of the signal is 1.0, which is not true for HDR...
    return (loss == 0.f) ? std::numeric_limits<float>::infinity() : -10.f * log10f(loss);
}

Status EstimateCompressedTextureSetSize(TextureSetDesc const& textureSetDesc,
    LatentShape const& latentShape, size_t& outSize)
{
    Status status = TextureSetMetadata::ValidateTextureSetDesc(textureSetDesc);
    if (status != Status::Ok)
        return status;

    status = TextureSetMetadata::ValidateLatentShape(latentShape);
    if (status != Status::Ok)
        return status;

    size_t latentSize = FeatureGridMath::CalculateQuantizedLatentsSize(
        textureSetDesc.width,
        textureSetDesc.height,
        textureSetDesc.mips,
        latentShape.gridSizeScale,
        latentShape.numFeatures);

    // Int8 per weight, 2x float per output (scale and bias)
    size_t const mlpSize = MlpDesc::GetTotalWeightCount() + MlpDesc::GetTotalOutputCount() * 2 * sizeof(float);

    outSize = latentSize + mlpSize;

    return Status::Ok;
}

// The table of statistically best latent shapes for a set of BPP values.
KnownLatentShape const g_KnownLatentShapes[KnownLatentShapeCount] = {
    {  1.11f, { 6, 8 } },
    {  1.67f, { 6, 12 } },
    {  2.50f, { 4, 8 } },
    {  3.75f, { 4, 12 } },
    {  5.00f, { 4, 16 } },
    {  6.67f, { 3, 12 } },
    {  10.0f, { 2, 8 } },
    {  15.0f, { 2, 12 } },
    {  20.0f, { 2, 16 } }
};

int GetKnownLatentShapeCount()
{
    return KnownLatentShapeCount;
}

Status EnumerateKnownLatentShapes(int index, float& outBitsPerPixel, LatentShape& outShape)
{
    ClearErrorMessage();

    if (index < 0 || index >= KnownLatentShapeCount)
    {
        SetErrorMessage("Invalid index (%d), must be 0-%d.", index, KnownLatentShapeCount - 1);
        return Status::OutOfRange;
    }

    outBitsPerPixel = g_KnownLatentShapes[index].bitsPerPixel;
    outShape = g_KnownLatentShapes[index].shape;
    return Status::Ok;
}

Status PickLatentShape(float requestedBitsPerPixel, float& outBitsPerPixel, LatentShape& outShape)
{
    ClearErrorMessage();

    if (requestedBitsPerPixel <= 0.f)
    {
        SetErrorMessage("requestedBitsPerPixel (%f) must be positive.", requestedBitsPerPixel);
        return Status::OutOfRange;
    }
    
    bool found = false;
    float currentDiff = std::numeric_limits<float>::max();
    for (int index = 0; index < KnownLatentShapeCount; ++index)
    {
        KnownLatentShape const& shape = g_KnownLatentShapes[index];
        float const diff = fabsf(shape.bitsPerPixel - requestedBitsPerPixel);
        if (diff < requestedBitsPerPixel * 0.25f)
        {
            if (diff < currentDiff)
            {
                found = true;
                currentDiff = diff;
                outBitsPerPixel = shape.bitsPerPixel;
                outShape = shape.shape;
            }
        }
        else if (found)
        {
            // The known shape list is sorted, and we're past the matching range,
            // so there won't be anything interesting anymore.
            return Status::Ok;
        }
    }

    if (found)
        return Status::Ok;


    SetErrorMessage("The provided requestedBitsPerPixel (%f) is out of the supported range.", requestedBitsPerPixel);
    return Status::OutOfRange;
}

float GetLatentShapeBitsPerPixel(LatentShape const& shape)
{
    int const numLayers = FeatureGridMath::GetNumLayers(shape.numFeatures);
    
    // We use 2 adjacent mip levels of the latent texture together, so the effective rate
    // is calculated as (BytesPerLatentPixel * 8 bits per byte) * 1.25 = BytesPerLatentPixel * 10
    return float(numLayers) / float(shape.gridSizeScale * shape.gridSizeScale) * float(FeatureGridMath::BytesPerLatentPixel) * 10.f;
}

double DecodeImageDifferenceResult(uint64_t value)
{
    // Convert from 16.48 fixed point to double
    return double(value) * 0x1p-48;
}

size_t GetBC7ModeBufferSize(int widthInBlocks, int heightInBlocks)
{
    size_t blockCount = size_t(widthInBlocks) * size_t(heightInBlocks);
    return ((blockCount * BLOCK_COMPRESS_BC7_MP_BITS + 31) >> 5) << 2;
}

static void ExtractBC7ModeAndPartition(uint32_t firstWord, uint32_t& outMode, uint32_t& outPartition)
{
    // Extract the mode and partition indices from the block
#ifdef _MSC_VER
    outMode = std::min(7u, _tzcnt_u32(firstWord));
#else
    outMode = std::min(7, __builtin_ctz(firstWord));
#endif

    outPartition = firstWord >> (outMode + 1);
    static const uint32_t partitionMask[8] = { 15, 63, 63, 63, 7, 3, 0, 63 };
    outPartition &= partitionMask[outMode];
}

Status MakeBC7ModeBuffer(int widthInBlocks, int heightInBlocks, void const* blockData, size_t blockDataSize,
    size_t rowPitch, void* modeBuffer, size_t modeBufferSize)
{
    if (blockData == nullptr)
    {
        SetErrorMessage("blockData is NULL.");
        return Status::InvalidArgument;
    }

    if (modeBuffer == nullptr)
    {
        SetErrorMessage("modeBuffer is NULL.");
        return Status::InvalidArgument;
    }

    size_t const bytesPerBlock = 16;
    size_t const expectedRowPitch = size_t(widthInBlocks) * bytesPerBlock;
    if (rowPitch < expectedRowPitch || (rowPitch % bytesPerBlock) != 0)
    {
        SetErrorMessage("Invalid rowPitch (%zu), must be at least %zu and a multiple of %zu.",
            rowPitch, expectedRowPitch, bytesPerBlock);
        return Status::InvalidArgument;
    }

    size_t const expectedSize = size_t(heightInBlocks) * rowPitch;
    if (blockDataSize < expectedSize)
    {
        SetErrorMessage("Size of BC7 block data (%zu bytes) does not match the expected size (%zu bytes).",
            blockDataSize, expectedSize);
        return Status::InvalidArgument;
    }
    
    size_t const expectedModeBufferSize = GetBC7ModeBufferSize(widthInBlocks, heightInBlocks);
    if (modeBufferSize < expectedModeBufferSize)
    {
        SetErrorMessage("Size of BC7 mode buffer (%zu bytes) is smaller than the expected size (%zu bytes).",
            modeBufferSize, expectedModeBufferSize);
        return Status::InvalidArgument;
    }

    uint8_t const* blockDataBytes = reinterpret_cast<uint8_t const*>(blockData);
    uint8_t* modeDataBytes = reinterpret_cast<uint8_t*>(modeBuffer);

    for (int row = 0; row < heightInBlocks; ++row)
    {
        for (int col = 0; col < widthInBlocks; ++col)
        {
            size_t const blockId = size_t(row) * size_t(widthInBlocks) + size_t(col);
            uint32_t const firstWord = *reinterpret_cast<uint32_t const*>(blockDataBytes
                + row * rowPitch + col * bytesPerBlock);

            uint32_t mode, partition;
            ExtractBC7ModeAndPartition(firstWord, mode, partition);

            uint32_t modeAndPartition = (mode << 6) | partition;
            // Reserve zero value for integration errors, map mode 0 partition 0 to a different unused value
            if (modeAndPartition == 0)
                modeAndPartition = BLOCK_COMPRESS_BC7_MODE0_PART0_VALUE;
            size_t bitAddress = blockId * BLOCK_COMPRESS_BC7_MP_BITS;
            size_t byteOffset = bitAddress / 8;
            uint32_t bitOffset = bitAddress & 7;
            modeDataBytes[byteOffset] |= static_cast<uint8_t>(modeAndPartition << bitOffset);
            modeDataBytes[byteOffset + 1] |= static_cast<uint8_t>(modeAndPartition >> (8 - bitOffset));
        }
    }

    ClearErrorMessage();
    return Status::Ok;
}

}