/src/libjxl/lib/jxl/enc_ans_simd.cc

Source
// Copyright (c) the JPEG XL Project Authors. All rights reserved.
//
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

#include "lib/jxl/enc_ans_simd.h"

#include <cstdint>

#include "lib/jxl/base/status.h"
#include "lib/jxl/dec_ans.h"
#include "lib/jxl/memory_manager_internal.h"

#undef HWY_TARGET_INCLUDE
#define HWY_TARGET_INCLUDE "lib/jxl/enc_ans_simd.cc"
#include <hwy/foreach_target.h>
#include <hwy/highway.h>

HWY_BEFORE_NAMESPACE();
namespace jxl {
namespace HWY_NAMESPACE {

// These templates are not found via ADL.
using hwy::HWY_NAMESPACE::Add;
using hwy::HWY_NAMESPACE::And;
using hwy::HWY_NAMESPACE::Ge;
using hwy::HWY_NAMESPACE::GetLane;
using hwy::HWY_NAMESPACE::Gt;
using hwy::HWY_NAMESPACE::IfThenElse;
using hwy::HWY_NAMESPACE::IfThenElseZero;
using hwy::HWY_NAMESPACE::Iota;
using hwy::HWY_NAMESPACE::LoadU;
using hwy::HWY_NAMESPACE::Lt;
using hwy::HWY_NAMESPACE::Mul;
using hwy::HWY_NAMESPACE::Or;
using hwy::HWY_NAMESPACE::Set;
using hwy::HWY_NAMESPACE::ShiftRight;
using hwy::HWY_NAMESPACE::Store;
using hwy::HWY_NAMESPACE::Sub;
using hwy::HWY_NAMESPACE::Zero;

template <size_t E, size_t M, size_t L>
uint32_t EstimateTokenCostImpl(uint32_t* JXL_RESTRICT values, size_t len,
                               uint32_t* JXL_RESTRICT out) {
  const HWY_FULL(uint32_t) du;
  const HWY_FULL(float) df;
  const auto kZero = Zero(du);
  const auto kSplit = Set(du, 1 << E);
  const auto kExpOffset = Set(du, 127);
  const auto kEBOffset = Set(du, 127 + M + L);
  const auto kBase = Set(du, static_cast<uint32_t>((1 << E) - (E << (M + L))));
  const auto kMulN = Set(du, 1 << (M + L));
  const auto kMaskL = Set(du, (1 << L) - 1);
  const auto kMaskM = Set(du, ((1 << M) - 1) << L);
  const auto kLargeThreshold = Set(du, (1 << 22) - 1);
  constexpr size_t kLargeShiftVal = 10;
  const auto kLargeShift = Set(du, kLargeShiftVal);

  auto extra_bits = kZero;
  size_t last_full = Lanes(du) * (len / Lanes(du));
  for (size_t i = 0; i < last_full; i += Lanes(du)) {
    const auto val = LoadU(du, values + i);
    const auto is_large = Gt(val, kLargeThreshold);
    const auto val_shifted = ShiftRight<kLargeShiftVal>(val);
    const auto not_literal = Ge(val, kSplit);
    const auto val_fixed = IfThenElse(is_large, val_shifted, val);
    const auto b = BitCast(du, ConvertTo(df, val_fixed));
    const auto l = And(val, kMaskL);
    const auto exp = ShiftRight<23>(b);
    const auto exp_fixed = IfThenElse(is_large, Add(exp, kLargeShift), exp);
    const auto n = Sub(exp_fixed, kExpOffset);
    const auto eb = Sub(exp_fixed, kEBOffset);
    const auto m = ShiftRight<23 - M - L>(b);
    const auto a = Add(kBase, Mul(n, kMulN));
    const auto d = And(m, kMaskM);
    const auto eb_fixed = IfThenElseZero(not_literal, eb);
    const auto c = Or(a, l);
    extra_bits = Add(extra_bits, eb_fixed);
    const auto t = Or(c, d);
    const auto t_fixed = IfThenElse(not_literal, t, val);
    Store(t_fixed, du, out + i);
  }
  if (last_full < len) {
    const auto stop = Set(du, len);
    const auto fence = Iota(du, last_full);
    const auto take = Lt(fence, stop);
    const auto val = LoadU(du, values + last_full);
    const auto is_large = Gt(val, kLargeThreshold);
    const auto val_shifted = ShiftRight<kLargeShiftVal>(val);
    const auto not_literal = Ge(val, kSplit);
    const auto val_fixed = IfThenElse(is_large, val_shifted, val);
    const auto b = BitCast(du, ConvertTo(df, val_fixed));
    const auto l = And(val, kMaskL);
    const auto exp = ShiftRight<23>(b);
    const auto exp_fixed = IfThenElse(is_large, Add(exp, kLargeShift), exp);
    const auto n = Sub(exp_fixed, kExpOffset);
    const auto eb = Sub(exp_fixed, kEBOffset);
    const auto m = ShiftRight<23 - M - L>(b);
    const auto a = Add(kBase, Mul(n, kMulN));
    const auto d = And(m, kMaskM);
    const auto eb_fixed = IfThenElseZero(not_literal, eb);
    const auto eb_masked = IfThenElseZero(take, eb_fixed);
    const auto c = Or(a, l);
    extra_bits = Add(extra_bits, eb_masked);
    const auto t = Or(c, d);
    const auto t_fixed = IfThenElse(not_literal, t, val);
    Store(t_fixed, du, out + last_full);
  }
  return GetLane(SumOfLanes(du, extra_bits));
}

uint32_t EstimateTokenCost(uint32_t* JXL_RESTRICT values, size_t len,
                           HybridUintConfig cfg, AlignedMemory& tokens) {
  uint32_t* JXL_RESTRICT out = tokens.address<uint32_t>();
#if HWY_TARGET == HWY_SCALAR
  uint32_t extra_bits = 0;
  for (size_t i = 0; i < len; ++i) {
    uint32_t v = values[i];
    uint32_t tok, nbits, bits;
    cfg.Encode(v, &tok, &nbits, &bits);
    extra_bits += nbits;
    out[i] = tok;
  }
  return extra_bits;
#else
  if (cfg.split_exponent == 0) {
    return EstimateTokenCostImpl<0, 0, 0>(values, len, out);
  } else if (cfg.split_exponent == 2) {
    JXL_DASSERT((cfg.msb_in_token == 0) && (cfg.lsb_in_token == 1));
    return EstimateTokenCostImpl<2, 0, 1>(values, len, out);
  } else if (cfg.split_exponent == 3) {
    if (cfg.msb_in_token == 1) {
      if (cfg.lsb_in_token == 0) {
        return EstimateTokenCostImpl<3, 1, 0>(values, len, out);
      } else {
        JXL_DASSERT(cfg.lsb_in_token == 2);
        return EstimateTokenCostImpl<3, 1, 2>(values, len, out);
      }
    } else {
      JXL_DASSERT(cfg.msb_in_token == 2);
      if (cfg.lsb_in_token == 0) {
        return EstimateTokenCostImpl<3, 2, 0>(values, len, out);
      } else {
        JXL_DASSERT(cfg.lsb_in_token == 1);
        return EstimateTokenCostImpl<3, 2, 1>(values, len, out);
      }
    }
  } else if (cfg.split_exponent == 4) {
    if (cfg.msb_in_token == 1) {
      if (cfg.lsb_in_token == 0) {
        return EstimateTokenCostImpl<4, 1, 0>(values, len, out);
      } else if (cfg.lsb_in_token == 2) {
        return EstimateTokenCostImpl<4, 1, 2>(values, len, out);
      } else {
        JXL_DASSERT(cfg.lsb_in_token == 3);
        return EstimateTokenCostImpl<4, 1, 3>(values, len, out);
      }
    } else {
      JXL_DASSERT(cfg.msb_in_token == 2);
      if (cfg.lsb_in_token == 0) {
        return EstimateTokenCostImpl<4, 2, 0>(values, len, out);
      } else if (cfg.lsb_in_token == 1) {
        return EstimateTokenCostImpl<4, 2, 1>(values, len, out);
      } else {
        JXL_DASSERT(cfg.lsb_in_token == 2);
        return EstimateTokenCostImpl<4, 2, 2>(values, len, out);
      }
    }
  } else if (cfg.split_exponent == 5) {
    if (cfg.msb_in_token == 1) {
      if (cfg.lsb_in_token == 0) {
        return EstimateTokenCostImpl<5, 1, 0>(values, len, out);
      } else if (cfg.lsb_in_token == 2) {
        return EstimateTokenCostImpl<5, 1, 2>(values, len, out);
      } else {
        JXL_DASSERT(cfg.lsb_in_token == 4);
        return EstimateTokenCostImpl<5, 1, 4>(values, len, out);
      }
    } else {
      JXL_DASSERT(cfg.msb_in_token == 2);
      if (cfg.lsb_in_token == 0) {
        return EstimateTokenCostImpl<5, 2, 0>(values, len, out);
      } else if (cfg.lsb_in_token == 1) {
        return EstimateTokenCostImpl<5, 2, 1>(values, len, out);
      } else if (cfg.lsb_in_token == 2) {
        return EstimateTokenCostImpl<5, 2, 2>(values, len, out);
      } else {
        JXL_DASSERT(cfg.lsb_in_token == 3);
        return EstimateTokenCostImpl<5, 2, 3>(values, len, out);
      }
    }
  } else if (cfg.split_exponent == 6) {
    if (cfg.msb_in_token == 0) {
      JXL_DASSERT(cfg.lsb_in_token == 0);
      return EstimateTokenCostImpl<6, 0, 0>(values, len, out);
    } else if (cfg.msb_in_token == 1) {
      JXL_DASSERT(cfg.lsb_in_token == 5);
      return EstimateTokenCostImpl<6, 1, 5>(values, len, out);
    } else {
      JXL_DASSERT(cfg.msb_in_token == 2);
      JXL_DASSERT(cfg.lsb_in_token == 4);
      return EstimateTokenCostImpl<6, 2, 4>(values, len, out);
    }
  } else if (cfg.split_exponent >= 7 && cfg.split_exponent <= 12) {
    JXL_DASSERT(cfg.msb_in_token == 0);
    JXL_DASSERT(cfg.lsb_in_token == 0);
    if (cfg.split_exponent == 7) {
      return EstimateTokenCostImpl<7, 0, 0>(values, len, out);
    } else if (cfg.split_exponent == 8) {
      return EstimateTokenCostImpl<8, 0, 0>(values, len, out);
    } else if (cfg.split_exponent == 9) {
      return EstimateTokenCostImpl<9, 0, 0>(values, len, out);
    } else if (cfg.split_exponent == 10) {
      return EstimateTokenCostImpl<10, 0, 0>(values, len, out);
    } else if (cfg.split_exponent == 11) {
      return EstimateTokenCostImpl<11, 0, 0>(values, len, out);
    } else {
      return EstimateTokenCostImpl<12, 0, 0>(values, len, out);
    }
  } else {
    JXL_DASSERT(false);
  }
  return ~0;
#endif
}

// NOLINTNEXTLINE(google-readability-namespace-comments)
}  // namespace HWY_NAMESPACE
}  // namespace jxl
HWY_AFTER_NAMESPACE();

#if HWY_ONCE
namespace jxl {

HWY_EXPORT(EstimateTokenCost);

uint32_t EstimateTokenCost(uint32_t* JXL_RESTRICT values, size_t len,
                           HybridUintConfig cfg, AlignedMemory& tokens) {
  JXL_DASSERT(cfg.lsb_in_token + cfg.msb_in_token <= cfg.split_exponent);
  return HWY_DYNAMIC_DISPATCH(EstimateTokenCost)(values, len, cfg, tokens);
}

}  // namespace jxl
#endif

Coverage Report

Created: 2025-12-03 07:28