texture.cpp

cross platform rendering playground

src/asset/texture.cpp

12.81 KB
#include "texture.h"

#include <ktx.h>
#include <stb_image.h>

#include <cstdio>
#include <filesystem>
#include <string>

#include "core/globals.h"
#include "core/gpu_arena.h"
#include "quick_submit.h"

static u32 calculate_mip_levels(u32 width, u32 height) {
    u32 levels = 1;
    while (width > 1 || height > 1) {
        width >>= 1;
        height >>= 1;
        levels++;
    }
    return levels;
}

static ktx_transcode_fmt_e to_ktx_transcode_format(rhi::ImageFormat fmt) {
    switch (fmt) {
    case rhi::ImageFormat::BC7_UNORM:
    case rhi::ImageFormat::BC7_SRGB:
        return KTX_TTF_BC7_RGBA;
    default:
        return KTX_TTF_RGBA32;
    }
}

static rhi::ImageRange whole_range(u32 mip_levels, u32 layers) {
    rhi::ImageRange r{};
    r.mip_count = mip_levels;
    r.layer_count = layers;
    return r;
}

static rhi::ImageRange mip_range(u32 mip) {
    rhi::ImageRange r{};
    r.mip = mip;
    r.mip_count = 1;
    r.layer_count = 1;
    return r;
}

static void barrier_image_to_transfer_dst(rhi::CmdBuffer &cmd, rhi::Image *image, u32 mip_levels, u32 layers) {
    rhi::barrier(
        cmd, *image, whole_range(mip_levels, layers), rhi::ResourceState::Idle, rhi::ResourceState::TransferTo
    );
}

static void barrier_image_to_shader_read(rhi::CmdBuffer &cmd, rhi::Image *image, u32 mip_levels, u32 layers) {
    rhi::barrier(
        cmd, *image, whole_range(mip_levels, layers), rhi::ResourceState::TransferTo, rhi::ResourceState::TextureSample
    );
}

static void barrier_mip_to_src(rhi::CmdBuffer &cmd, rhi::Image *image, u32 mip) {
    rhi::barrier(cmd, *image, mip_range(mip), rhi::ResourceState::TransferTo, rhi::ResourceState::TransferFrom);
}

static void barrier_mip_to_shader_read(rhi::CmdBuffer &cmd, rhi::Image *image, u32 mip) {
    rhi::barrier(cmd, *image, mip_range(mip), rhi::ResourceState::TransferFrom, rhi::ResourceState::TextureSample);
}

static void set_view_shader_read(rhi::ImageView *view) {
    if (view == nullptr || view->desc.image == nullptr) {
        return;
    }
    view->desc.image->state = rhi::ResourceState::TextureSample;
}

static void generate_mips_2d(rhi::CmdBuffer &cmd, rhi::Image *image, u32 width, u32 height, u32 mip_levels) {
    u32 mip_w = width;
    u32 mip_h = height;
    for (u32 i = 0; i < mip_levels - 1; ++i) {
        u32 next_w = std::max(mip_w >> 1, 1u);
        u32 next_h = std::max(mip_h >> 1, 1u);

        barrier_mip_to_src(cmd, image, i);

        rhi::ImageBlit blit{
            .src_mip = i,
            .dst_mip = i + 1,
            .src_x0 = 0,
            .src_y0 = 0,
            .src_x1 = static_cast<i32>(mip_w),
            .src_y1 = static_cast<i32>(mip_h),
            .dst_x0 = 0,
            .dst_y0 = 0,
            .dst_x1 = static_cast<i32>(next_w),
            .dst_y1 = static_cast<i32>(next_h),
        };
        rhi::blit_image(cmd, *image, &blit, 1);

        barrier_mip_to_shader_read(cmd, image, i);

        mip_w = next_w;
        mip_h = next_h;
    }
    // previous mip is still in TransferTo — transition it now
    rhi::barrier(
        cmd, *image, mip_range(mip_levels - 1), rhi::ResourceState::TransferTo, rhi::ResourceState::TextureSample
    );
}

static bool upload_stb_pixels(
    rhi::Device &device,
    QuickSubmit &qs,
    Arena &staging,
    rhi::Image &out_image,
    rhi::ImageView &out_view,
    const void *pixels,
    u32 width,
    u32 height,
    u64 pixel_stride, // bytes per pixel × channels (4 or 16)
    rhi::ImageFormat format,
    const char *name = nullptr
) {
    u32 mip_levels = calculate_mip_levels(width, height);
    u64 image_size = static_cast<u64>(width) * height * pixel_stride;
    u64 staging_off = stage_to_arena(staging, pixels, image_size);
    if (staging_off == SIZE_MAX) {
        return false;
    }
    rhi::buffer_flush(device, staging.buffer, staging_off, image_size);

    rhi::ImageDesc img_desc{};
    img_desc.width = width;
    img_desc.height = height;
    img_desc.mip_levels = mip_levels;
    if (name != nullptr) {
        img_desc.name = name;
    }
    img_desc.usage = rhi::ImageUsage::Sampled | rhi::ImageUsage::TransferDst | rhi::ImageUsage::TransferSrc;
    img_desc.format = format;
    rhi::create_image(device, img_desc, out_image);

    rhi::ImageViewDesc view_desc{};
    view_desc.image = &out_image;
    if (!rhi::create_image_view(device, view_desc, out_view)) {
        rhi::destroy_image(device, out_image);
        return false;
    }

    QsCmd *cmd = qs_acquire(qs);

    barrier_image_to_transfer_dst(cmd->cmd, &out_image, mip_levels, 1);

    rhi::BufferTextureCopyRegion region{};
    region.buffer_offset = staging_off;
    region.texture_extent_width = width;
    region.texture_extent_height = height;
    region.texture_extent_depth = 1;
    region.base_array_layer = 0;
    region.mip_level = 0;
    region.layer_count = 1;
    rhi::copy_buffer_to_texture(cmd->cmd, staging.buffer, out_view, region);

    if (mip_levels > 1) {
        generate_mips_2d(cmd->cmd, &out_image, width, height, mip_levels);
    } else {
        barrier_image_to_shader_read(cmd->cmd, &out_image, 1, 1);
    }

    set_view_shader_read(&out_view);

    qs_submit(qs, *cmd);
    return true;
}

bool ktx_texture_needs_to_flip_to_top_left(ktxTexture *kTexture) {
    if (!kTexture) {
        return false;
    }

    if (kTexture->orientation.y == 'u') {
        return true;
    }

    return false;
}

static bool upload_ktx2_texture(
    rhi::Device &device,
    QuickSubmit &qs,
    Arena &staging,
    rhi::Image &out_image,
    rhi::ImageView &out_view,
    ktxTexture2 *ktx_tex,
    const char *name = nullptr
) {
    ktxTexture *base = ktxTexture(ktx_tex);
    u32 width = base->baseWidth;
    u32 height = base->baseHeight;
    u32 mip_levels = base->numLevels;
    u32 layers = base->numLayers;
    u32 faces = base->numFaces;
    u32 array_layers = (layers > 1) ? layers : faces;
    bool is_cubemap = (faces == 6);

    rhi::ImageDesc image_desc{};
    if (name != nullptr) {
        image_desc.name = name;
    }
    image_desc.format = rhi::image_format_from_vk(ktx_tex->vkFormat);
    image_desc.is_cubemap = is_cubemap;
    image_desc.width = width;
    image_desc.height = height;
    image_desc.mip_levels = mip_levels;
    image_desc.layers = array_layers;
    image_desc.usage = rhi::ImageUsage::TransferDst | rhi::ImageUsage::Sampled;
    rhi::create_image(device, image_desc, out_image);

    rhi::ImageViewDesc view_desc{};
    view_desc.image = &out_image;
    view_desc.aspect = rhi::ImageAspect::Color;
    view_desc.dimension = is_cubemap ? rhi::TextureViewDimension::TEXTURE_CUBE : rhi::TextureViewDimension::TEXTURE_2D;
    view_desc.layer_count = array_layers;
    if (!rhi::create_image_view(device, view_desc, out_view)) {
        rhi::destroy_image(device, out_image);
        ktxTexture_Destroy(base);
        return false;
    }

    ktx_size_t data_size = ktxTexture_GetDataSize(base);
    u64 staging_off = stage_to_arena(staging, base->pData, data_size);
    if (staging_off == SIZE_MAX) {
        ktxTexture_Destroy(base);
        return false;
    }
    rhi::buffer_flush(device, staging.buffer, staging_off, data_size);

    QsCmd *cmd = qs_acquire(qs);

    barrier_image_to_transfer_dst(cmd->cmd, &out_image, mip_levels, array_layers);

    for (u32 level = 0; level < mip_levels; ++level) {
        for (u32 layer = 0; layer < array_layers; ++layer) {
            ktx_size_t offset;
            ktxTexture_GetImageOffset(base, level, layer, 0, &offset);

            rhi::BufferTextureCopyRegion region{};
            region.buffer_offset = staging_off + offset;
            region.mip_level = level;
            region.base_array_layer = layer;
            region.layer_count = 1;
            region.texture_extent_width = std::max(width >> level, 1u);
            region.texture_extent_height = std::max(height >> level, 1u);
            region.texture_extent_depth = 1;
            rhi::copy_buffer_to_texture(cmd->cmd, staging.buffer, out_view, region);
        }
    }

    barrier_image_to_shader_read(cmd->cmd, &out_image, mip_levels, array_layers);

    set_view_shader_read(&out_view);

    qs_submit(qs, *cmd);
    ktxTexture_Destroy(base);
    return true;
}

////////////
bool load_texture_from_memory(
    rhi::Device &device,
    QuickSubmit &qs,
    Arena &staging,
    rhi::Image &out_image,
    rhi::ImageView &out_view,
    const u8 *data,
    usize size,
    ImageColorSpace color_space,
    const char *name
) {
    int w, h, channels;
    stbi_uc *pixels = stbi_load_from_memory(data, static_cast<int>(size), &w, &h, &channels, STBI_rgb_alpha);
    if (!pixels) {
        (void)std::fprintf(stderr, "Failed to load texture from memory (size: %zu)\n", size);
        return false;
    }

    rhi::ImageFormat fmt =
        (color_space == ImageColorSpace::SRGB) ? rhi::ImageFormat::RGBA8_SRGB : rhi::ImageFormat::RGBA8_UNORM;
    bool ok = upload_stb_pixels(
        device, qs, staging, out_image, out_view, pixels, static_cast<u32>(w), static_cast<u32>(h), 4, fmt, name
    );
    stbi_image_free(pixels);
    return ok;
}

bool create_texture_from_raw_rgba(
    rhi::Device &device,
    QuickSubmit &qs,
    Arena &staging,
    rhi::Image &out_image,
    rhi::ImageView &out_view,
    const void *pixels,
    u32 width,
    u32 height,
    u32 pixel_stride,
    rhi::ImageFormat cs,
    const char *name
) {
    return upload_stb_pixels(device, qs, staging, out_image, out_view, pixels, width, height, pixel_stride, cs, name);
}

bool load_texture_ktx2_from_memory(
    rhi::Device &device,
    QuickSubmit &qs,
    Arena &staging,
    rhi::Image &out_image,
    rhi::ImageView &out_view,
    const u8 *data,
    usize size,
    ImageColorSpace color_space,
    const char *name
) {
    ktxTexture2 *ktx_tex = nullptr;
    if (ktxTexture2_CreateFromMemory(data, size, KTX_TEXTURE_CREATE_LOAD_IMAGE_DATA_BIT, &ktx_tex) != KTX_SUCCESS) {
        return false;
    }

    if (ktxTexture2_NeedsTranscoding(ktx_tex)) {
        ktx_transcode_fmt_e tf =
            to_ktx_transcode_format(rhi::preferred_compressed_format(device, color_space == ImageColorSpace::SRGB));

        if (ktxTexture2_TranscodeBasis(ktx_tex, tf, 0) != KTX_SUCCESS) {
            (void)std::fprintf(stderr, "KTX2 transcode failed\n");
            ktxTexture_Destroy(ktxTexture(ktx_tex));
            return false;
        }
    }

    return upload_ktx2_texture(device, qs, staging, out_image, out_view, ktx_tex, name);
}

bool load_texture(
    rhi::Device &device,
    QuickSubmit &qs,
    Arena &staging,
    const char *path,
    rhi::Image &out_image,
    rhi::ImageView &out_view,
    ImageColorSpace color_space
) {
    std::filesystem::path p(path);
    std::string ext = p.extension().string();

    if (ext == ".ktx2") {
        ktxTexture2 *ktx_tex = nullptr;
        KTX_error_code err = ktxTexture2_CreateFromNamedFile(path, KTX_TEXTURE_CREATE_LOAD_IMAGE_DATA_BIT, &ktx_tex);

        if (err != KTX_SUCCESS) {
            (void)std::fprintf(stderr, "Failed to load KTX2: %s (error: %d)\n", path, err);
            return false;
        }

        if (ktxTexture2_NeedsTranscoding(ktx_tex)) {
            ktx_transcode_fmt_e tf =
                to_ktx_transcode_format(rhi::preferred_compressed_format(device, color_space == ImageColorSpace::SRGB));

            if (ktxTexture2_TranscodeBasis(ktx_tex, tf, 0) != KTX_SUCCESS) {
                (void)std::fprintf(stderr, "KTX2 transcode failed: %s\n", path);
                ktxTexture_Destroy(ktxTexture(ktx_tex));
                return false;
            }
        }

        return upload_ktx2_texture(device, qs, staging, out_image, out_view, ktx_tex, path);
    }

    if (ext == ".hdr") {
        int w, h, channels;
        f32 *pixels = stbi_loadf(path, &w, &h, &channels, STBI_rgb_alpha);
        if (!pixels) {
            (void)std::fprintf(stderr, "Failed to load HDR texture: %s\n", path);
            return false;
        }

        bool ok = upload_stb_pixels(
            device,
            qs,
            staging,
            out_image,
            out_view,
            pixels,
            static_cast<u32>(w),
            static_cast<u32>(h),
            sizeof(f32) * 4,
            rhi::ImageFormat::RGBA32_FLOAT,
            path
        );
        stbi_image_free(pixels);
        return ok;
    }

    // fallback: assume PNG / 8-bit
    int w, h, channels;
    stbi_uc *pixels = stbi_load(path, &w, &h, &channels, STBI_rgb_alpha);
    if (!pixels) {
        (void)std::fprintf(stderr, "Failed to load texture: %s\n", path);
        return false;
    }

    rhi::ImageFormat fmt =
        (color_space == ImageColorSpace::SRGB) ? rhi::ImageFormat::RGBA8_SRGB : rhi::ImageFormat::RGBA8_UNORM;
    bool ok = upload_stb_pixels(
        device, qs, staging, out_image, out_view, pixels, static_cast<u32>(w), static_cast<u32>(h), 4, fmt, path
    );
    stbi_image_free(pixels);
    return ok;
}

void destroy_texture(rhi::Device &device, rhi::Image &image, rhi::ImageView &view) {
    rhi::destroy_image_view(device, view);
    rhi::destroy_image(device, image);
    image = {};
    view = {};
}