// SPDX-License-Identifier: Apache-2.0

// ----------------------------------------------------------------------------

// Copyright 2011-2026 Arm Limited

//

// Licensed under the Apache License, Version 2.0 (the "License"); you may not

// use this file except in compliance with the License. You may obtain a copy

// of the License at:

//

// http://www.apache.org/licenses/LICENSE-2.0

//

// Unless required by applicable law or agreed to in writing, software

// distributed under the License is distributed on an "AS IS" BASIS, WITHOUT

// WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the

// License for the specific language governing permissions and limitations

// under the License.

// ----------------------------------------------------------------------------

#if !defined(ASTCENC_DECOMPRESS_ONLY)

/**

* @brief Functions for angular-sum algorithm for weight alignment.

*

* This algorithm works as follows:

* - we compute a complex number P as (cos s*i, sin s*i) for each weight,

* where i is the input value and s is a scaling factor based on the spacing between the weights.

* - we then add together complex numbers for all the weights.

* - we then compute the length and angle of the resulting sum.

*

* This should produce the following results:

* - perfect alignment results in a vector whose length is equal to the sum of lengths of all inputs

* - even distribution results in a vector of length 0.

* - all samples identical results in perfect alignment for every scaling.

*

* For each scaling factor within a given set, we compute an alignment factor from 0 to 1. This

* should then result in some scalings standing out as having particularly good alignment factors;

* we can use this to produce a set of candidate scale/shift values for various quantization levels;

* we should then actually try them and see what happens.

*/

#include "astcenc_internal.h"

#include "astcenc_vecmathlib.h"

#include <stdio.h>

#include <cassert>

#include <cstring>

#include <cfloat>

static constexpr unsigned int ANGULAR_STEPS { 32 };

static_assert((ANGULAR_STEPS % ASTCENC_SIMD_WIDTH) == 0,

"ANGULAR_STEPS must be multiple of ASTCENC_SIMD_WIDTH");

static_assert(ANGULAR_STEPS >= 32,

"ANGULAR_STEPS must be at least max(steps_for_quant_level)");

// Store a reduced sin/cos table for 64 possible weight values; this causes

// slight quality loss compared to using sin() and cos() directly. Must be 2^N.

static constexpr unsigned int SINCOS_STEPS { 64 };

static const uint8_t steps_for_quant_level[12] {

2, 3, 4, 5, 6, 8, 10, 12, 16, 20, 24, 32

};

ASTCENC_ALIGNAS static float sin_table[SINCOS_STEPS][ANGULAR_STEPS];

ASTCENC_ALIGNAS static float cos_table[SINCOS_STEPS][ANGULAR_STEPS];

#if defined(ASTCENC_DIAGNOSTICS)

static bool print_once { true };

#endif

/* See header for documentation. */

void prepare_angular_tables()

{

for (unsigned int i = 0; i < ANGULAR_STEPS; i++)

{

float angle_step = static_cast<float>(i + 1);

for (unsigned int j = 0; j < SINCOS_STEPS; j++)

{

sin_table[j][i] = static_cast<float>(sinf((2.0f * astc::PI / (SINCOS_STEPS - 1.0f)) * angle_step * static_cast<float>(j)));

cos_table[j][i] = static_cast<float>(cosf((2.0f * astc::PI / (SINCOS_STEPS - 1.0f)) * angle_step * static_cast<float>(j)));

}

}

}

/**

* @brief Compute the angular alignment factors and offsets.

*

* @param weight_count The number of (decimated) weights.

* @param dec_weight_ideal_value The ideal decimated unquantized weight values.

* @param max_angular_steps The maximum number of steps to be tested.

* @param[out] offsets The output angular offsets array.

*/

static void compute_angular_offsets(

unsigned int weight_count,

const float* dec_weight_ideal_value,

unsigned int max_angular_steps,

float* offsets

) {

promise(weight_count > 0);

promise(max_angular_steps > 0);

ASTCENC_ALIGNAS int isamplev[BLOCK_MAX_WEIGHTS];

// Precompute isample; arrays are always allocated 64 elements long

for (unsigned int i = 0; i < weight_count; i += ASTCENC_SIMD_WIDTH)

{

// Ideal weight can be outside [0, 1] range, so clamp to fit table

vfloat ideal_weight = clampzo(loada(dec_weight_ideal_value + i));

// Convert a weight to a sincos table index

vfloat sample = ideal_weight * (SINCOS_STEPS - 1.0f);

vint isample = float_to_int_rtn(sample);

storea(isample, isamplev + i);

}

// Arrays are multiple of SIMD width (ANGULAR_STEPS), safe to overshoot max

vfloat mult(1.0f / (2.0f * astc::PI));

for (unsigned int i = 0; i < max_angular_steps; i += ASTCENC_SIMD_WIDTH)

{

vfloat anglesum_x = vfloat::zero();

vfloat anglesum_y = vfloat::zero();

for (unsigned int j = 0; j < weight_count; j++)

{

int isample = isamplev[j];

anglesum_x += loada(cos_table[isample] + i);

anglesum_y += loada(sin_table[isample] + i);

}

vfloat angle = atan2(anglesum_y, anglesum_x);

// Suppress NaNs generated if anglesums are both zero

angle = select(vfloat::zero(), angle, angle == angle);

vfloat ofs = angle * mult;

storea(ofs, offsets + i);

}

}

/**

* @brief For a given step size compute the lowest and highest weight.

*

* Compute the lowest and highest weight that results from quantizing using the given stepsize and

* offset, and then compute the resulting error. The cut errors indicate the error that results from

* forcing samples that should have had one weight value one step up or down.

*

* @param weight_count The number of (decimated) weights.

* @param dec_weight_ideal_value The ideal decimated unquantized weight values.

* @param max_angular_steps The maximum number of steps to be tested.

* @param max_quant_steps The maximum quantization level to be tested.

* @param offsets The angular offsets array.

* @param[out] lowest_weight Per angular step, the lowest weight.

* @param[out] weight_span Per angular step, the span between lowest and highest weight.

* @param[out] error Per angular step, the error.

* @param[out] cut_low_weight_error Per angular step, the low weight cut error.

* @param[out] cut_high_weight_error Per angular step, the high weight cut error.

*/

static void compute_lowest_and_highest_weight(

unsigned int weight_count,

const float* dec_weight_ideal_value,

unsigned int max_angular_steps,

unsigned int max_quant_steps,

const float* offsets,

float* lowest_weight,

int* weight_span,

float* error,

float* cut_low_weight_error,

float* cut_high_weight_error

) {

promise(weight_count > 0);

promise(max_angular_steps > 0);

vfloat rcp_stepsize = int_to_float(vint::lane_id()) + vfloat(1.0f);

// Compute minimum/maximum weights in the weight array. Our remapping

// is monotonic, so the min/max rounded weights relate to the min/max

// unrounded weights in a straightforward way.

vfloat min_weight(FLT_MAX);

vfloat max_weight(-FLT_MAX);

vint lane_id = vint::lane_id();

for (unsigned int i = 0; i < weight_count; i += ASTCENC_SIMD_WIDTH)

{

vmask active = lane_id < vint(weight_count);

lane_id += vint(ASTCENC_SIMD_WIDTH);

vfloat weights = loada(dec_weight_ideal_value + i);

min_weight = min(select(min_weight, weights, active), min_weight);

max_weight = max(select(max_weight, weights, active), max_weight);

}

min_weight = hmin(min_weight);

max_weight = hmax(max_weight);

// Arrays are ANGULAR_STEPS long, so always safe to run full vectors

for (unsigned int sp = 0; sp < max_angular_steps; sp += ASTCENC_SIMD_WIDTH)

{

vfloat errval = vfloat::zero();

vfloat cut_low_weight_err = vfloat::zero();

vfloat cut_high_weight_err = vfloat::zero();

vfloat offset = loada(offsets + sp);

// We know the min and max weight values, so we can figure out

// the corresponding indices before we enter the loop.

vfloat minidx = round(min_weight * rcp_stepsize - offset);

vfloat maxidx = round(max_weight * rcp_stepsize - offset);

for (unsigned int j = 0; j < weight_count; j++)

{

vfloat sval = load1(dec_weight_ideal_value + j) * rcp_stepsize - offset;

vfloat svalrte = round(sval);

vfloat diff = sval - svalrte;

errval += diff * diff;

// Accumulate errors for minimum index

vmask mask = svalrte == minidx;

vfloat accum = cut_low_weight_err + vfloat(1.0f) - vfloat(2.0f) * diff;

cut_low_weight_err = select(cut_low_weight_err, accum, mask);

// Accumulate errors for maximum index

mask = svalrte == maxidx;

accum = cut_high_weight_err + vfloat(1.0f) + vfloat(2.0f) * diff;

cut_high_weight_err = select(cut_high_weight_err, accum, mask);

}

// Write out min weight and weight span; clamp span to a usable range

vint span = float_to_int(maxidx - minidx + vfloat(1));

span = min(span, vint(max_quant_steps + 3));

span = max(span, vint(2));

storea(minidx, lowest_weight + sp);

storea(span, weight_span + sp);

// The cut_(lowest/highest)_weight_error indicate the error that results from forcing

// samples that should have had the weight value one step (up/down).

vfloat ssize = 1.0f / rcp_stepsize;

vfloat errscale = ssize * ssize;

storea(errval * errscale, error + sp);

storea(cut_low_weight_err * errscale, cut_low_weight_error + sp);

storea(cut_high_weight_err * errscale, cut_high_weight_error + sp);

rcp_stepsize = rcp_stepsize + vfloat(ASTCENC_SIMD_WIDTH);

}

}

/**

* @brief The main function for the angular algorithm.

*

* @param weight_count The number of (decimated) weights.

* @param dec_weight_ideal_value The ideal decimated unquantized weight values.

* @param max_quant_level The maximum quantization level to be tested.

* @param[out] low_value Per angular step, the lowest weight value.

* @param[out] high_value Per angular step, the highest weight value.

*/

static void compute_angular_endpoints_for_quant_levels(

unsigned int weight_count,

const float* dec_weight_ideal_value,

unsigned int max_quant_level,

float low_value[TUNE_MAX_ANGULAR_QUANT + 1],

float high_value[TUNE_MAX_ANGULAR_QUANT + 1]

) {

unsigned int max_quant_steps = steps_for_quant_level[max_quant_level];

unsigned int max_angular_steps = steps_for_quant_level[max_quant_level];

ASTCENC_ALIGNAS float angular_offsets[ANGULAR_STEPS];

compute_angular_offsets(weight_count, dec_weight_ideal_value,

max_angular_steps, angular_offsets);

ASTCENC_ALIGNAS float lowest_weight[ANGULAR_STEPS];

ASTCENC_ALIGNAS int32_t weight_span[ANGULAR_STEPS];

ASTCENC_ALIGNAS float error[ANGULAR_STEPS];

ASTCENC_ALIGNAS float cut_low_weight_error[ANGULAR_STEPS];

ASTCENC_ALIGNAS float cut_high_weight_error[ANGULAR_STEPS];

compute_lowest_and_highest_weight(weight_count, dec_weight_ideal_value,

max_angular_steps, max_quant_steps,

angular_offsets, lowest_weight, weight_span, error,

cut_low_weight_error, cut_high_weight_error);

// For each quantization level, find the best error terms. Use packed vectors so data-dependent

// branches can become selects. This involves some integer to float casts, but the values are

// small enough so they never round the wrong way.

vfloat4 best_results[36];

// Initialize the array to some safe defaults

promise(max_quant_steps > 0);

for (unsigned int i = 0; i < (max_quant_steps + 4); i++)

{

// Lane<0> = Best error

// Lane<1> = Best scale; -1 indicates no solution found

// Lane<2> = Cut low weight

best_results[i] = vfloat4(ERROR_CALC_DEFAULT, -1.0f, 0.0f, 0.0f);

}

promise(max_angular_steps > 0);

for (unsigned int i = 0; i < max_angular_steps; i++)

{

float i_flt = static_cast<float>(i);

int idx_span = weight_span[i];

float error_cut_low = error[i] + cut_low_weight_error[i];

float error_cut_high = error[i] + cut_high_weight_error[i];

float error_cut_low_high = error[i] + cut_low_weight_error[i] + cut_high_weight_error[i];

// Check best error against record N

vfloat4 best_result = best_results[idx_span];

vfloat4 new_result = vfloat4(error[i], i_flt, 0.0f, 0.0f);

vmask4 mask = vfloat4(best_result.lane<0>()) > vfloat4(error[i]);

best_results[idx_span] = select(best_result, new_result, mask);

// Check best error against record N-1 with either cut low or cut high

best_result = best_results[idx_span - 1];

new_result = vfloat4(error_cut_low, i_flt, 1.0f, 0.0f);

mask = vfloat4(best_result.lane<0>()) > vfloat4(error_cut_low);

best_result = select(best_result, new_result, mask);

new_result = vfloat4(error_cut_high, i_flt, 0.0f, 0.0f);

mask = vfloat4(best_result.lane<0>()) > vfloat4(error_cut_high);

best_results[idx_span - 1] = select(best_result, new_result, mask);

// Check best error against record N-2 with both cut low and high

best_result = best_results[idx_span - 2];

new_result = vfloat4(error_cut_low_high, i_flt, 1.0f, 0.0f);

mask = vfloat4(best_result.lane<0>()) > vfloat4(error_cut_low_high);

best_results[idx_span - 2] = select(best_result, new_result, mask);

}

for (unsigned int i = 0; i <= max_quant_level; i++)

{

unsigned int q = steps_for_quant_level[i];

int bsi = static_cast<int>(best_results[q].lane<1>());

// Did we find anything?

#if defined(ASTCENC_DIAGNOSTICS)

if ((bsi < 0) && print_once)

{

print_once = false;

printf("INFO: Unable to find full encoding within search error limit.\n\n");

}

#endif

bsi = astc::max(0, bsi);

float lwi = lowest_weight[bsi] + best_results[q].lane<2>();

float hwi = lwi + static_cast<float>(q) - 1.0f;

float stepsize = 1.0f / (1.0f + static_cast<float>(bsi));

low_value[i] = (angular_offsets[bsi] + lwi) * stepsize;

high_value[i] = (angular_offsets[bsi] + hwi) * stepsize;

}

}

/* See header for documentation. */

void compute_angular_endpoints_1plane(

bool only_always,

const block_size_descriptor& bsd,

const float* dec_weight_ideal_value,

unsigned int max_weight_quant,

compression_working_buffers& tmpbuf

) {

float (&low_value)[WEIGHTS_MAX_BLOCK_MODES] = tmpbuf.weight_low_value1;

float (&high_value)[WEIGHTS_MAX_BLOCK_MODES] = tmpbuf.weight_high_value1;

float (&low_values)[WEIGHTS_MAX_DECIMATION_MODES][TUNE_MAX_ANGULAR_QUANT + 1] = tmpbuf.weight_low_values1;

float (&high_values)[WEIGHTS_MAX_DECIMATION_MODES][TUNE_MAX_ANGULAR_QUANT + 1] = tmpbuf.weight_high_values1;

unsigned int max_decimation_modes = only_always ? bsd.decimation_mode_count_always

: bsd.decimation_mode_count_selected;

promise(max_decimation_modes > 0);

for (unsigned int i = 0; i < max_decimation_modes; i++)

{

const decimation_mode& dm = bsd.decimation_modes[i];

if (!dm.is_ref_1plane(static_cast<quant_method>(max_weight_quant)))

{

continue;

}

unsigned int weight_count = bsd.get_decimation_info(i).weight_count;

unsigned int max_precision = dm.maxprec_1plane;

if (max_precision > TUNE_MAX_ANGULAR_QUANT)

{

max_precision = TUNE_MAX_ANGULAR_QUANT;

}

if (max_precision > max_weight_quant)

{

max_precision = max_weight_quant;

}

compute_angular_endpoints_for_quant_levels(

weight_count,

dec_weight_ideal_value + i * BLOCK_MAX_WEIGHTS,

max_precision, low_values[i], high_values[i]);

}

unsigned int max_block_modes = only_always ? bsd.block_mode_count_1plane_always

: bsd.block_mode_count_1plane_selected;

promise(max_block_modes > 0);

for (unsigned int i = 0; i < max_block_modes; i++)

{

const block_mode& bm = bsd.block_modes[i];

assert(!bm.is_dual_plane);

unsigned int quant_mode = bm.quant_mode;

unsigned int decim_mode = bm.decimation_mode;

if (quant_mode <= TUNE_MAX_ANGULAR_QUANT)

{

low_value[i] = low_values[decim_mode][quant_mode];

high_value[i] = high_values[decim_mode][quant_mode];

}

else

{

low_value[i] = 0.0f;

high_value[i] = 1.0f;

}

}

}

/* See header for documentation. */

void compute_angular_endpoints_2planes(

const block_size_descriptor& bsd,

const float* dec_weight_ideal_value,

unsigned int max_weight_quant,

compression_working_buffers& tmpbuf

) {

float (&low_value1)[WEIGHTS_MAX_BLOCK_MODES] = tmpbuf.weight_low_value1;

float (&high_value1)[WEIGHTS_MAX_BLOCK_MODES] = tmpbuf.weight_high_value1;

float (&low_value2)[WEIGHTS_MAX_BLOCK_MODES] = tmpbuf.weight_low_value2;

float (&high_value2)[WEIGHTS_MAX_BLOCK_MODES] = tmpbuf.weight_high_value2;

float (&low_values1)[WEIGHTS_MAX_DECIMATION_MODES][TUNE_MAX_ANGULAR_QUANT + 1] = tmpbuf.weight_low_values1;

float (&high_values1)[WEIGHTS_MAX_DECIMATION_MODES][TUNE_MAX_ANGULAR_QUANT + 1] = tmpbuf.weight_high_values1;

float (&low_values2)[WEIGHTS_MAX_DECIMATION_MODES][TUNE_MAX_ANGULAR_QUANT + 1] = tmpbuf.weight_low_values2;

float (&high_values2)[WEIGHTS_MAX_DECIMATION_MODES][TUNE_MAX_ANGULAR_QUANT + 1] = tmpbuf.weight_high_values2;

promise(bsd.decimation_mode_count_selected > 0);

for (unsigned int i = 0; i < bsd.decimation_mode_count_selected; i++)

{

const decimation_mode& dm = bsd.decimation_modes[i];

if (!dm.is_ref_2plane(static_cast<quant_method>(max_weight_quant)))

{

continue;

}

unsigned int weight_count = bsd.get_decimation_info(i).weight_count;

unsigned int max_precision = dm.maxprec_2planes;

if (max_precision > TUNE_MAX_ANGULAR_QUANT)

{

max_precision = TUNE_MAX_ANGULAR_QUANT;

}

if (max_precision > max_weight_quant)

{

max_precision = max_weight_quant;

}

compute_angular_endpoints_for_quant_levels(

weight_count,

dec_weight_ideal_value + i * BLOCK_MAX_WEIGHTS,

max_precision, low_values1[i], high_values1[i]);

compute_angular_endpoints_for_quant_levels(

weight_count,

dec_weight_ideal_value + i * BLOCK_MAX_WEIGHTS + WEIGHTS_PLANE2_OFFSET,

max_precision, low_values2[i], high_values2[i]);

}

unsigned int start = bsd.block_mode_count_1plane_selected;

unsigned int end = bsd.block_mode_count_1plane_2plane_selected;

for (unsigned int i = start; i < end; i++)

{

const block_mode& bm = bsd.block_modes[i];

unsigned int quant_mode = bm.quant_mode;

unsigned int decim_mode = bm.decimation_mode;

if (quant_mode <= TUNE_MAX_ANGULAR_QUANT)

{

low_value1[i] = low_values1[decim_mode][quant_mode];

high_value1[i] = high_values1[decim_mode][quant_mode];

low_value2[i] = low_values2[decim_mode][quant_mode];

high_value2[i] = high_values2[decim_mode][quant_mode];

}

else

{

low_value1[i] = 0.0f;

high_value1[i] = 1.0f;

low_value2[i] = 0.0f;

high_value2[i] = 1.0f;

}

}

}

#endif