Skip to content

Commit

Permalink
Add astc encode support
Browse files Browse the repository at this point in the history
  • Loading branch information
@wasimabbas-arm
wasimabbas-arm committed Oct 19, 2024
1 parent f8b4c66 commit 4faf9ea
Show file tree
Hide file tree
Showing 4 changed files with 343 additions and 8 deletions.
3 changes: 3 additions & 0 deletions include/ktx.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1307,6 +1307,9 @@ ktxTexture2_CompressAstcEx(ktxTexture2* This, ktxAstcParams* params);
KTX_API KTX_error_code KTX_APIENTRY
ktxTexture2_CompressAstc(ktxTexture2* This, ktx_uint32_t quality);

KTX_API KTX_error_code KTX_APIENTRY
ktxTexture2_DecodeAstc(ktxTexture2* This, ktx_uint32_t vkformat);

/**
* @memberof ktxTexture2
* @~English
Expand Down
318 changes: 315 additions & 3 deletions lib/astc_encode.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -360,7 +360,7 @@ astcSwizzle(const ktxAstcParams &params) {

static void
astcBlockDimensions(ktx_uint32_t block_size,
uint32_t& block_x, uint32_t& block_y, uint32_t& block_z) {
uint32_t& block_x, uint32_t& block_y, uint32_t& block_z) {
switch (block_size) {
case KTX_PACK_ASTC_BLOCK_DIMENSION_4x4 : block_x = 4; block_y = 4; block_z = 1; break;
case KTX_PACK_ASTC_BLOCK_DIMENSION_5x4 : block_x = 5; block_y = 4; block_z = 1; break;
Expand Down Expand Up @@ -391,6 +391,88 @@ astcBlockDimensions(ktx_uint32_t block_size,
}
}


static void
astcBlockDimensions(VkFormat format,
uint32_t &x, uint32_t &y, uint32_t &z) noexcept {
switch (format) {
case VK_FORMAT_ASTC_4x4_UNORM_BLOCK: x = 4; y = 4; z = 0; break;
case VK_FORMAT_ASTC_4x4_SRGB_BLOCK: x = 4; y = 4; z = 0; break;
case VK_FORMAT_ASTC_5x4_UNORM_BLOCK: x = 5; y = 4; z = 0; break;
case VK_FORMAT_ASTC_5x4_SRGB_BLOCK: x = 5; y = 4; z = 0; break;
case VK_FORMAT_ASTC_5x5_UNORM_BLOCK: x = 5; y = 5; z = 0; break;
case VK_FORMAT_ASTC_5x5_SRGB_BLOCK: x = 5; y = 5; z = 0; break;
case VK_FORMAT_ASTC_6x5_UNORM_BLOCK: x = 6; y = 5; z = 0; break;
case VK_FORMAT_ASTC_6x5_SRGB_BLOCK: x = 6; y = 5; z = 0; break;
case VK_FORMAT_ASTC_6x6_UNORM_BLOCK: x = 6; y = 6; z = 0; break;
case VK_FORMAT_ASTC_6x6_SRGB_BLOCK: x = 6; y = 6; z = 0; break;
case VK_FORMAT_ASTC_8x5_UNORM_BLOCK: x = 8; y = 5; z = 0; break;
case VK_FORMAT_ASTC_8x5_SRGB_BLOCK: x = 8; y = 5; z = 0; break;
case VK_FORMAT_ASTC_8x6_UNORM_BLOCK: x = 8; y = 6; z = 0; break;
case VK_FORMAT_ASTC_8x6_SRGB_BLOCK: x = 8; y = 6; z = 0; break;
case VK_FORMAT_ASTC_8x8_UNORM_BLOCK: x = 8; y = 8; z = 0; break;
case VK_FORMAT_ASTC_8x8_SRGB_BLOCK: x = 8; y = 8; z = 0; break;
case VK_FORMAT_ASTC_10x5_UNORM_BLOCK: x = 10; y = 5; z = 0; break;
case VK_FORMAT_ASTC_10x5_SRGB_BLOCK: x = 10; y = 5; z = 0; break;
case VK_FORMAT_ASTC_10x6_UNORM_BLOCK: x = 10; y = 6; z = 0; break;
case VK_FORMAT_ASTC_10x6_SRGB_BLOCK: x = 10; y = 6; z = 0; break;
case VK_FORMAT_ASTC_10x8_UNORM_BLOCK: x = 10; y = 8; z = 0; break;
case VK_FORMAT_ASTC_10x8_SRGB_BLOCK: x = 10; y = 8; z = 0; break;
case VK_FORMAT_ASTC_10x10_UNORM_BLOCK: x = 10; y = 10; z = 0; break;
case VK_FORMAT_ASTC_10x10_SRGB_BLOCK: x = 10; y = 10; z = 0; break;
case VK_FORMAT_ASTC_12x10_UNORM_BLOCK: x = 12; y = 10; z = 0; break;
case VK_FORMAT_ASTC_12x10_SRGB_BLOCK: x = 12; y = 10; z = 0; break;
case VK_FORMAT_ASTC_12x12_UNORM_BLOCK: x = 12; y = 12; z = 0; break;
case VK_FORMAT_ASTC_12x12_SRGB_BLOCK: x = 12; y = 12; z = 0; break;
case VK_FORMAT_ASTC_4x4_SFLOAT_BLOCK: x = 4; y = 4; z = 0; break;
case VK_FORMAT_ASTC_5x4_SFLOAT_BLOCK: x = 5; y = 4; z = 0; break;
case VK_FORMAT_ASTC_5x5_SFLOAT_BLOCK: x = 5; y = 5; z = 0; break;
case VK_FORMAT_ASTC_6x5_SFLOAT_BLOCK: x = 6; y = 5; z = 0; break;
case VK_FORMAT_ASTC_6x6_SFLOAT_BLOCK: x = 6; y = 6; z = 0; break;
case VK_FORMAT_ASTC_8x5_SFLOAT_BLOCK: x = 8; y = 5; z = 0; break;
case VK_FORMAT_ASTC_8x6_SFLOAT_BLOCK: x = 8; y = 6; z = 0; break;
case VK_FORMAT_ASTC_8x8_SFLOAT_BLOCK: x = 8; y = 8; z = 0; break;
case VK_FORMAT_ASTC_10x5_SFLOAT_BLOCK: x = 10; y = 5; z = 0; break;
case VK_FORMAT_ASTC_10x6_SFLOAT_BLOCK: x = 10; y = 6; z = 0; break;
case VK_FORMAT_ASTC_10x8_SFLOAT_BLOCK: x = 10; y = 8; z = 0; break;
case VK_FORMAT_ASTC_10x10_SFLOAT_BLOCK: x = 10; y = 10; z = 0; break;
case VK_FORMAT_ASTC_12x10_SFLOAT_BLOCK: x = 12; y = 10; z = 0; break;
case VK_FORMAT_ASTC_12x12_SFLOAT_BLOCK: x = 12; y = 12; z = 0; break;
case VK_FORMAT_ASTC_3x3x3_UNORM_BLOCK_EXT: x = 3; y = 3; z = 3; break;
case VK_FORMAT_ASTC_3x3x3_SRGB_BLOCK_EXT: x = 3; y = 3; z = 3; break;
case VK_FORMAT_ASTC_3x3x3_SFLOAT_BLOCK_EXT: x = 3; y = 3; z = 3; break;
case VK_FORMAT_ASTC_4x3x3_UNORM_BLOCK_EXT: x = 4; y = 3; z = 3; break;
case VK_FORMAT_ASTC_4x3x3_SRGB_BLOCK_EXT: x = 4; y = 3; z = 3; break;
case VK_FORMAT_ASTC_4x3x3_SFLOAT_BLOCK_EXT: x = 4; y = 3; z = 3; break;
case VK_FORMAT_ASTC_4x4x3_UNORM_BLOCK_EXT: x = 4; y = 4; z = 3; break;
case VK_FORMAT_ASTC_4x4x3_SRGB_BLOCK_EXT: x = 4; y = 4; z = 3; break;
case VK_FORMAT_ASTC_4x4x3_SFLOAT_BLOCK_EXT: x = 4; y = 4; z = 3; break;
case VK_FORMAT_ASTC_4x4x4_UNORM_BLOCK_EXT: x = 4; y = 4; z = 4; break;
case VK_FORMAT_ASTC_4x4x4_SRGB_BLOCK_EXT: x = 4; y = 4; z = 4; break;
case VK_FORMAT_ASTC_4x4x4_SFLOAT_BLOCK_EXT: x = 4; y = 4; z = 4; break;
case VK_FORMAT_ASTC_5x4x4_UNORM_BLOCK_EXT: x = 5; y = 4; z = 4; break;
case VK_FORMAT_ASTC_5x4x4_SRGB_BLOCK_EXT: x = 5; y = 4; z = 4; break;
case VK_FORMAT_ASTC_5x4x4_SFLOAT_BLOCK_EXT: x = 5; y = 4; z = 4; break;
case VK_FORMAT_ASTC_5x5x4_UNORM_BLOCK_EXT: x = 5; y = 5; z = 4; break;
case VK_FORMAT_ASTC_5x5x4_SRGB_BLOCK_EXT: x = 5; y = 5; z = 4; break;
case VK_FORMAT_ASTC_5x5x4_SFLOAT_BLOCK_EXT: x = 5; y = 5; z = 4; break;
case VK_FORMAT_ASTC_5x5x5_UNORM_BLOCK_EXT: x = 5; y = 5; z = 5; break;
case VK_FORMAT_ASTC_5x5x5_SRGB_BLOCK_EXT: x = 5; y = 5; z = 5; break;
case VK_FORMAT_ASTC_5x5x5_SFLOAT_BLOCK_EXT: x = 5; y = 5; z = 5; break;
case VK_FORMAT_ASTC_6x5x5_UNORM_BLOCK_EXT: x = 6; y = 5; z = 5; break;
case VK_FORMAT_ASTC_6x5x5_SRGB_BLOCK_EXT: x = 6; y = 5; z = 5; break;
case VK_FORMAT_ASTC_6x5x5_SFLOAT_BLOCK_EXT: x = 6; y = 5; z = 5; break;
case VK_FORMAT_ASTC_6x6x5_UNORM_BLOCK_EXT: x = 6; y = 6; z = 5; break;
case VK_FORMAT_ASTC_6x6x5_SRGB_BLOCK_EXT: x = 6; y = 6; z = 5; break;
case VK_FORMAT_ASTC_6x6x5_SFLOAT_BLOCK_EXT: x = 6; y = 6; z = 5; break;
case VK_FORMAT_ASTC_6x6x6_UNORM_BLOCK_EXT: x = 6; y = 6; z = 6; break;
case VK_FORMAT_ASTC_6x6x6_SRGB_BLOCK_EXT: x = 6; y = 6; z = 6; break;
case VK_FORMAT_ASTC_6x6x6_SFLOAT_BLOCK_EXT: x = 6; y = 6; z = 6; break;
default:
x = 0; y = 0; z = 0;
}
}

static float
astcQuality(ktx_uint32_t quality_level) {
switch (quality_level) {
Expand Down Expand Up @@ -716,8 +798,9 @@ ktxTexture2_CompressAstcEx(ktxTexture2* This, ktxAstcParams* params) {
assert(KHR_DFDVAL(prototype->pDfd+1, MODEL) == KHR_DF_MODEL_ASTC
&& "Invalid dfd generated for ASTC image\n");
assert((transfer == KHR_DF_TRANSFER_SRGB
? KHR_DFDVAL(prototype->pDfd+1, PRIMARIES) == KHR_DF_PRIMARIES_SRGB
: true) && "Not a valid sRGB image\n");
? KHR_DFDVAL(prototype->pDfd+1, TRANSFER) == KHR_DF_TRANSFER_SRGB &&
KHR_DFDVAL(prototype->pDfd+1, PRIMARIES) == KHR_DF_PRIMARIES_SRGB
: true) && "Not a valid sRGB image\n");

// Fix up the current (This) texture
#undef DECLARE_PRIVATE
Expand Down Expand Up @@ -807,3 +890,232 @@ ktxTexture2_CompressAstc(ktxTexture2* This, ktx_uint32_t quality) {

return ktxTexture2_CompressAstcEx(This, &params);
}

struct decompression_workload
{
astcenc_context* context;
uint8_t* data;
size_t data_len;
astcenc_image* image_out;
astcenc_swizzle swizzle;
astcenc_error error;
};

/**
* @brief Runner callback function for a decompression worker thread.
*
* @param thread_count The number of threads in the worker pool.
* @param thread_id The index of this thread in the worker pool.
* @param payload The parameters for this thread.
*/
static void decompression_workload_runner(int thread_count, int thread_id, void* payload) {
(void)thread_count;

decompression_workload* work = static_cast<decompression_workload*>(payload);
astcenc_error error = astcenc_decompress_image(work->context, work->data, work->data_len,
work->image_out, &work->swizzle, thread_id);
// This is a racy update, so which error gets returned is a random, but it
// will reliably report an error if an error occurs
if (error != ASTCENC_SUCCESS)
{
work->error = error;
}
}

/**
* @brief Decodes the provided ktx2 texture if its astc encoded
*
* @param This The texture to decode
* @param vkformat The decoding format to use
*/
// will update "This" with the uncompressed copy
KTX_API KTX_error_code KTX_APIENTRY
ktxTexture2_DecodeAstc(ktxTexture2 *This, ktx_uint32_t vkformat) {
// Decompress This using astc-decoder
uint32_t* BDB = This->pDfd + 1;
khr_df_model_e colorModel = (khr_df_model_e)KHR_DFDVAL(BDB, MODEL);
if (colorModel != KHR_DF_MODEL_ASTC && This->supercompressionScheme != KTX_SS_NONE) // No supercompression supported yet
{
return KTX_INVALID_OPERATION; // Not in valid astc decodable format
}

DECLARE_PRIVATE(priv, This);

uint32_t channelId = KHR_DFDSVAL(BDB, 0, CHANNELID);
if (channelId == KHR_DF_CHANNEL_ASTC_DATA) {
// Found astc data
}
else
return KTX_FILE_DATA_ERROR;

// Create a prototype texture to use for calculating sizes in the target
// format and, as useful side effects, provide us with a properly sized
// data allocation and the DFD for the target format.
ktxTextureCreateInfo createInfo;
createInfo.glInternalformat = 0;
createInfo.vkFormat = vkformat;
createInfo.baseWidth = This->baseWidth;
createInfo.baseHeight = This->baseHeight;
createInfo.baseDepth = This->baseDepth;
createInfo.generateMipmaps = This->generateMipmaps;
createInfo.isArray = This->isArray;
createInfo.numDimensions = This->numDimensions;
createInfo.numFaces = This->numFaces;
createInfo.numLayers = This->numLayers;
createInfo.numLevels = This->numLevels;
createInfo.pDfd = nullptr;

KTX_error_code result;
ktxTexture2* prototype;
result = ktxTexture2_Create(&createInfo, KTX_TEXTURE_CREATE_ALLOC_STORAGE, &prototype);

if (result != KTX_SUCCESS) {
assert(result == KTX_OUT_OF_MEMORY); // The only run time error
return result;
}

if (!This->pData) {
if (ktxTexture_isActiveStream((ktxTexture*)This)) {
// Load pending. Complete it.
result = ktxTexture2_LoadImageData(This, NULL, 0);
if (result != KTX_SUCCESS)
{
ktxTexture2_Destroy(prototype);
return result;
}
} else {
// No data to transcode.
ktxTexture2_Destroy(prototype);
return KTX_INVALID_OPERATION;
}
}

// This is where I do the decompression from "This" to prototype target
astcenc_profile profile{ASTCENC_PRF_LDR_SRGB};
astcenc_swizzle swizzle{ASTCENC_SWZ_R, ASTCENC_SWZ_G, ASTCENC_SWZ_B, ASTCENC_SWZ_A};

uint32_t block_size_x{4}; // Get the right blocks from vkformat
uint32_t block_size_y{4};
uint32_t block_size_z{1};
float quality{ASTCENC_PRE_MEDIUM};
uint32_t flags{0}; // TODO: Use normals mode to reconstruct normals params->normalMap ? ASTCENC_FLG_MAP_NORMAL : 0};

astcBlockDimensions((VkFormat)This->vkFormat, block_size_x, block_size_y, block_size_z);

// quality = astcQuality(params->qualityLevel);
// profile = astcEncoderAction(*params, BDB);
// swizzle = astcSwizzle(*params);

// if(params->perceptual) flags |= ASTCENC_FLG_USE_PERCEPTUAL;

uint32_t threadCount{1}; // Decompression isn't the bottlneck and only used when checking for psnr and ssim
astcenc_config astc_config;
astcenc_context *astc_context;
astcenc_error astc_error = astcenc_config_init(profile,
block_size_x, block_size_y, block_size_z,
quality, flags,
&astc_config);

if (astc_error != ASTCENC_SUCCESS)
return KTX_INVALID_OPERATION;

astc_error = astcenc_context_alloc(&astc_config, threadCount, &astc_context);

if (astc_error != ASTCENC_SUCCESS)
return KTX_INVALID_OPERATION;

decompression_workload work;
work.context = astc_context;
work.swizzle = swizzle;
work.error = ASTCENC_SUCCESS;

for (uint32_t levelIndex = 0; levelIndex < This->numLevels; ++levelIndex) {
const uint32_t imageWidth = std::max(This->baseWidth >> levelIndex, 1u);
const uint32_t imageHeight = std::max(This->baseHeight >> levelIndex, 1u);
const uint32_t imageDepths = std::max(This->baseDepth >> levelIndex, 1u);

for (uint32_t layerIndex = 0; layerIndex < This->numLayers; ++layerIndex) {
for (uint32_t faceIndex = 0; faceIndex < This->numFaces; ++faceIndex) {
for (uint32_t depthSliceIndex = 0; depthSliceIndex < imageDepths; ++depthSliceIndex) {

ktx_size_t levelImageSizeIn = ktxTexture_calcImageSize(ktxTexture(This), levelIndex, KTX_FORMAT_VERSION_TWO);

ktx_size_t imageOffsetIn;
ktx_size_t imageOffsetOut;

ktxTexture2_GetImageOffset(This, levelIndex, layerIndex, faceIndex + depthSliceIndex, &imageOffsetIn);
ktxTexture2_GetImageOffset(prototype, levelIndex, layerIndex, faceIndex + depthSliceIndex, &imageOffsetOut);

auto* imageDataIn = This->pData + imageOffsetIn;
auto* imageDataOut = prototype->pData + imageOffsetOut;

astcenc_image imageOut;
imageOut.dim_x = imageWidth;
imageOut.dim_y = imageHeight;
imageOut.dim_z = imageDepths;
imageOut.data_type = ASTCENC_TYPE_U8; // TODO: Fix for HDR types
imageOut.data = (void**)&imageDataOut; // TODO: Fix for HDR types

work.data = imageDataIn;
work.data_len = levelImageSizeIn;
work.image_out = &imageOut;

// Only launch worker threads for multi-threaded use - it makes basic
// single-threaded profiling and debugging a little less convoluted
if (threadCount > 1) {
launchThreads(threadCount, decompression_workload_runner, &work);
} else {
work.error = astcenc_decompress_image(work.context, work.data, work.data_len,
work.image_out, &work.swizzle, 0);
}

// Reset ASTC context for next image
astcenc_decompress_reset(astc_context);

if (work.error != ASTCENC_SUCCESS) {
std::cout << "ASTC decompressor failed\n" << astcenc_get_error_string(work.error) << std::endl;

astcenc_context_free(astc_context);
return KTX_INVALID_OPERATION;
}
}
}
}
}

// We are done with astcdecoder
astcenc_context_free(astc_context);

if (result == KTX_SUCCESS) {
// Fix up the current texture
DECLARE_PROTECTED(thisPrtctd, This);
DECLARE_PRIVATE(protoPriv, prototype);
DECLARE_PROTECTED(protoPrtctd, prototype);
memcpy(&thisPrtctd._formatSize, &protoPrtctd._formatSize,
sizeof(ktxFormatSize));
This->vkFormat = vkformat;
This->isCompressed = prototype->isCompressed;
This->supercompressionScheme = KTX_SS_NONE;
priv._requiredLevelAlignment = protoPriv._requiredLevelAlignment;
// Copy the levelIndex from the prototype to This.
memcpy(priv._levelIndex, protoPriv._levelIndex,
This->numLevels * sizeof(ktxLevelIndexEntry));
// Move the DFD and data from the prototype to This.
free(This->pDfd);
This->pDfd = prototype->pDfd;
prototype->pDfd = 0;
free(This->pData);
This->pData = prototype->pData;
This->dataSize = prototype->dataSize;
prototype->pData = 0;
prototype->dataSize = 0;
// Free SGD data
This->_private->_sgdByteLength = 0;
if (This->_private->_supercompressionGlobalData) {
free(This->_private->_supercompressionGlobalData);
This->_private->_supercompressionGlobalData = NULL;
}
}
ktxTexture2_Destroy(prototype);
return result;
}
15 changes: 12 additions & 3 deletions tools/ktx/command_encode.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -251,9 +251,18 @@ void CommandEncode::executeEncode() {

MetricsCalculator metrics;
metrics.saveReferenceImages(texture, options, *this);
ret = ktxTexture2_CompressBasisEx(texture, &options);
if (ret != KTX_SUCCESS)
fatal(rc::IO_FAILURE, "Failed to encode KTX2 file with codec \"{}\". KTX Error: {}", ktxErrorString(ret));

if (options.vkFormat != VK_FORMAT_UNDEFINED) {
options.mode = KTX_PACK_ASTC_ENCODER_MODE_LDR; // TODO: Fix me for HDR textures
ret = ktxTexture2_CompressAstcEx(texture, &options);
if (ret != KTX_SUCCESS)
fatal(rc::IO_FAILURE, "Failed to encode KTX2 file with codec \"{}\". KTX Error: {}", ktxErrorString(ret));
} else {
ret = ktxTexture2_CompressBasisEx(texture, &options);
if (ret != KTX_SUCCESS)
fatal(rc::IO_FAILURE, "Failed to encode KTX2 file with codec \"{}\". KTX Error: {}", ktxErrorString(ret));
}

metrics.decodeAndCalculateMetrics(texture, options, *this);

if (options.zstd) {
Expand Down
15 changes: 13 additions & 2 deletions tools/ktx/metrics_utils.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -101,11 +101,22 @@ class MetricsCalculator {
KTXTexture2 texture{static_cast<ktxTexture2*>(malloc(sizeof(ktxTexture2)))};
ktxTexture2_constructCopy(texture, encodedTexture);

const auto tSwizzleInfo = determineTranscodeSwizzle(texture, report);
// Start with a default swizzle
TranscodeSwizzleInfo tSwizzleInfo{};
tSwizzleInfo.defaultNumComponents = 4;
tSwizzleInfo.swizzle = "rgba";

ktx_error_code_e ec = KTX_SUCCESS;
// Decode the encoded texture to observe the compression losses
ec = ktxTexture2_TranscodeBasis(texture, KTX_TTF_RGBA32, 0);
if (isFormatAstc((VkFormat)texture->vkFormat))
{
ec = ktxTexture2_DecodeAstc(texture, VK_FORMAT_R8G8B8A8_UNORM);
}
else
{
ec = ktxTexture2_TranscodeBasis(texture, KTX_TTF_RGBA32, 0);
tSwizzleInfo = determineTranscodeSwizzle(texture, report);
}
if (ec != KTX_SUCCESS)
report.fatal(rc::KTX_FAILURE, "Failed to transcode KTX2 texture to calculate error metrics: {}", ktxErrorString(ec));

Expand Down

0 comments on commit 4faf9ea

Please sign in to comment.