From 2a1bf6e763b1942afa782989d2f6040f26edb482 Mon Sep 17 00:00:00 2001 From: Jiseong Oh Date: Tue, 13 Oct 2026 19:35:59 +0900 Subject: [PATCH 01/17] Add prototype ENN executor runner for Llama Add a prototype runner that executes an ENN backend Llama program with its KV cache managed on the host. Signed-off-by: Jiseong Oh --- examples/models/llama/runner/CMakeLists.txt | 1 + examples/samsung/CMakeLists.txt | 158 ++++- examples/samsung/build.sh | 13 +- .../enn_executor_llama_runner.cpp | 128 ++++ .../llama/runner/CMakeLists.txt | 88 +++ .../executor_runner/llama/runner/decoder.cpp | 122 ++++ .../executor_runner/llama/runner/decoder.h | 110 ++++ .../executor_runner/llama/runner/io_manager.h | 227 +++++++ .../llama/runner/prefiller.cpp | 149 +++++ .../executor_runner/llama/runner/prefiller.h | 66 ++ .../executor_runner/llama/runner/runner.cpp | 568 ++++++++++++++++++ .../executor_runner/llama/runner/runner.h | 185 ++++++ .../llama/runner/token_generator.h | 171 ++++++ 13 files changed, 1983 insertions(+), 3 deletions(-) create mode 100644 examples/samsung/executor_runner/enn_executor_llama_runner.cpp create mode 100644 examples/samsung/executor_runner/llama/runner/CMakeLists.txt create mode 100644 examples/samsung/executor_runner/llama/runner/decoder.cpp create mode 100644 examples/samsung/executor_runner/llama/runner/decoder.h create mode 100644 examples/samsung/executor_runner/llama/runner/io_manager.h create mode 100644 examples/samsung/executor_runner/llama/runner/prefiller.cpp create mode 100644 examples/samsung/executor_runner/llama/runner/prefiller.h create mode 100644 examples/samsung/executor_runner/llama/runner/runner.cpp create mode 100644 examples/samsung/executor_runner/llama/runner/runner.h create mode 100644 examples/samsung/executor_runner/llama/runner/token_generator.h diff --git a/examples/models/llama/runner/CMakeLists.txt b/examples/models/llama/runner/CMakeLists.txt index 31bce16dd12..b459bf313d5 100644 --- a/examples/models/llama/runner/CMakeLists.txt +++ b/examples/models/llama/runner/CMakeLists.txt @@ -54,6 +54,7 @@ set(llama_runner_deps target_link_libraries(llama_runner PUBLIC ${llama_runner_deps}) target_link_libraries(llama_runner PUBLIC tokenizers::tokenizers) +target_link_libraries(llama_runner PUBLIC enn_backend) target_include_directories( llama_runner PUBLIC ${EXECUTORCH_ROOT}/extension/llm/tokenizers/include diff --git a/examples/samsung/CMakeLists.txt b/examples/samsung/CMakeLists.txt index 28a6dcc6660..6182fb7bfc0 100644 --- a/examples/samsung/CMakeLists.txt +++ b/examples/samsung/CMakeLists.txt @@ -4,11 +4,13 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. -cmake_minimum_required(VERSION 3.15) +cmake_minimum_required(VERSION 3.24) set(CMAKE_CXX_STANDARD 17) set(CMAKE_CXX_STANDARD_REQUIRED ON) project(samsung_example) +message(EXECUTORCH_ROOT ${EXECUTORCH_ROOT}) + # Source root directory for executorch. if(NOT EXECUTORCH_ROOT) set(EXECUTORCH_ROOT ${CMAKE_CURRENT_SOURCE_DIR}/../..) @@ -31,16 +33,49 @@ endif() # prebuilt libraries. executorch package should contain portable_ops_lib, # etdump, bundled_program. +find_package(tokenizers REQUIRED) find_package(executorch CONFIG REQUIRED) target_compile_options(executorch INTERFACE -DET_EVENT_TRACER_ENABLED) find_package(gflags REQUIRED) set(_common_compile_options -Wno-deprecated-declarations -fPIC) -add_compile_options(-Wall -Werror -fPIC) +# +# The `__srcs` lists are defined by executorch_load_build_variables. +# +executorch_load_build_variables() + +set(EXECUTORCH_SRCS_FILE + "${CMAKE_CURRENT_BINARY_DIR}/../../executorch_srcs.cmake" +) + +get_filename_component( + EXECUTORCH_SOURCE_DIR "${CMAKE_CURRENT_LIST_DIR}/../.." ABSOLUTE +) + +add_compile_options( + -Wall + -Werror + -Wno-deprecated-declarations + -Wno-unused-const-variable + -Wno-unused-variable + -Wno-inconsistent-missing-override + -fPIC +) message("Build Samsung Android Examples") +link_directories(${EXECUTORCH_ROOT}/build_samsung_android/lib) + +add_library(lib_portable_ops STATIC IMPORTED) + +set_target_properties( + lib_portable_ops + PROPERTIES + IMPORTED_LOCATION + "${EXECUTORCH_ROOT}/build_samsung_android/lib/libportable_ops_lib.a" +) + set(__enn_executor_runner__srcs ${CMAKE_CURRENT_LIST_DIR}/executor_runner/enn_executor_runner.cpp ) @@ -59,3 +94,122 @@ target_link_libraries( set_target_properties( enn_executor_runner PROPERTIES CXX_VISIBILITY_PRESET hidden ) + +# llama_runner library +add_subdirectory(executor_runner/llama/runner) + +target_link_directories(enn_llama_runner PRIVATE) + +add_library(lib_enn_llama_runner STATIC IMPORTED) +set_target_properties( + lib_enn_llama_runner + PROPERTIES + IMPORTED_LOCATION + "${EXECUTORCH_ROOT}/build_samsung_android/examples/samsung/executor_runner/llama/runner/libenn_llama_runner.a" +) + +add_library(lib_tokenizer STATIC IMPORTED) +set_target_properties( + lib_tokenizer + PROPERTIES + IMPORTED_LOCATION + "${EXECUTORCH_ROOT}/build_samsung_android/extension/llm/tokenizers/libtokenizers.a" +) + +add_library(lib_extension_llm_runner STATIC IMPORTED) +set_target_properties( + lib_extension_llm_runner + PROPERTIES + IMPORTED_LOCATION + "${EXECUTORCH_ROOT}/build_samsung_android/extension/llm/runner/libextension_llm_runner.a" +) + +add_library(lib_enn_backend STATIC IMPORTED) +set_target_properties( + lib_enn_backend + PROPERTIES IMPORTED_LOCATION + "${EXECUTORCH_ROOT}/build_samsung_android/lib/libenn_backend.a" +) + +add_library(lib_executorch STATIC IMPORTED) +set_target_properties( + lib_executorch + PROPERTIES IMPORTED_LOCATION + "${EXECUTORCH_ROOT}/build_samsung_android/lib/libexecutorch.a" +) + +add_library(lib_custom_op STATIC IMPORTED) +set_target_properties( + lib_custom_op + PROPERTIES IMPORTED_LOCATION + "${EXECUTORCH_ROOT}/build_samsung_android/lib/libcustom_ops.a" +) + +set(__enn_executor_llama_runner__srcs ${_executor_runner__srcs}) +list(TRANSFORM __enn_executor_llama_runner__srcs + PREPEND "${EXECUTORCH_SOURCE_DIR}/" +) +list(FILTER __enn_executor_llama_runner__srcs EXCLUDE REGEX + ".*executor_runner.cpp$" +) +list(PREPEND __enn_executor_llama_runner_srcs + ${CMAKE_CURRENT_LIST_DIR}/executor_runner/enn_executor_llama_runner.cpp +) + +set(enn_executor_llama_runner_deps + executorch_core extension_data_loader extension_module extension_tensor + extension_flat_tensor lib_extension_llm_runner +) + +add_executable(enn_executor_llama_runner ${__enn_executor_llama_runner_srcs}) + +target_include_directories( + enn_executor_llama_runner + PUBLIC ${_common_include_directories} + ${EXECUTORCH_ROOT}/backends/samsung/runtime + ${EXECUTORCH_ROOT}/extension/llm/tokenizers/include +) + +set(xnnpack_backend_libs xnnpack_backend XNNPACK xnnpack-microkernels-prod + extension_threadpool pthreadpool kleidiai +) + +list( + APPEND + xnnpack_backend_libs + optimized_native_cpu_ops_lib + optimized_kernels + portable_kernels + cpublas + eigen_blas +) + +# Extra sources for cpuinfo +set(XNNPACK_ROOT ${EXECUTORCH_ROOT}/backends/xnnpack) +list(APPEND xnnpack_backend_libs cpuinfo) +list(APPEND _common_include_directories + ${XNNPACK_ROOT}/third-party/cpuinfo/include +) + +# quantized_ops_lib: Register quantized op kernels into the runtime +list(APPEND xnnpack_backend_libs quantized_kernels quantized_ops_lib) + +target_link_libraries( + enn_executor_llama_runner + PUBLIC enn_llama_runner + PRIVATE log + gflags + android + lib_portable_ops + ${enn_executor_llama_runner_deps} + ${xnnpack_backend_libs} + sentencepiece + re2 + lib_tokenizer + -Wl,--whole-archive + lib_executorch + lib_custom_op + lib_enn_backend + enn_logging + -Wl,--no-whole-archive +) diff --git a/examples/samsung/build.sh b/examples/samsung/build.sh index 4568180777a..0a0a8b0e1cf 100755 --- a/examples/samsung/build.sh +++ b/examples/samsung/build.sh @@ -21,6 +21,17 @@ echo ANDROID_PLATFORM=${ANDROID_PLATFORM} main() { cd "$PROJECT_DIR" + + if [[ -z ${ANDROID_NDK_ROOT} ]]; then + echo "Please export ANDROID_NDK_ROOT by setting andrid NDK path" + echo "$ export ANDROID_NDK_ROOT = path of android NDK" + exit 1 + fi + + if [[ -z ${ANDROID_NDK} ]]; then + ANDROID_NDK=${ANDROID_NDK_ROOT} + fi + local build_dir_root="build_samsung_android" local example_root="examples/samsung" local build_dir_example="$PROJECT_DIR/${build_dir_root}/${example_root}" @@ -39,7 +50,7 @@ main() { -DCMAKE_BUILD_TYPE=Release \ -B"${build_dir_example}" \ "$PROJECT_DIR/${example_root}" - cmake --build build_samsung_android/examples/samsung/ --config Release + cmake --build ${build_dir_example} --config Release } main "$@" diff --git a/examples/samsung/executor_runner/enn_executor_llama_runner.cpp b/examples/samsung/executor_runner/enn_executor_llama_runner.cpp new file mode 100644 index 00000000000..54b9d1eb868 --- /dev/null +++ b/examples/samsung/executor_runner/enn_executor_llama_runner.cpp @@ -0,0 +1,128 @@ +/* + * Copyright (c) 2025 Samsung Electronics Co. LTD + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#include + +#include + +#if defined(ET_USE_THREADPOOL) +#include +#include +#endif + +namespace llm = ::executorch::extension::llm; + +DEFINE_string( + model_path, + "llama2.pte", + "Model serialized in flatbuffer format."); + +DEFINE_string(data_path, "", "Data file for the model."); + +DEFINE_string(tokenizer_path, "tokenizer.bin", "Tokenizer stuff."); + +DEFINE_string(prompt, "The answer to the ultimate question is", "Prompt."); + +DEFINE_double( + temperature, + 0.8f, + "Temperature; Default is 0.8f. 0 = greedy argmax sampling (deterministic). Lower temperature = more deterministic"); + +DEFINE_int32( + seq_len, + 128, + "Total number of tokens to generate (prompt + output). Defaults to max_seq_len. If the number of input tokens + seq_len > max_seq_len, the output will be truncated to max_seq_len tokens."); + +DEFINE_int32( + cpu_threads, + -1, + "Number of CPU threads for inference. Defaults to -1, which implies we'll use a heuristic to derive the # of performant cores for a specific device."); + +DEFINE_int32( + num_bos, + 0, + "Number of BOS tokens to prepend to the prompt. Defaults to 0. If > 0, the prompt will be prepended with BOS tokens. This is useful for models that expect one or more BOS token at the start."); + +DEFINE_int32( + num_eos, + 0, + "Number of EOS tokens to append to the prompt. Defaults to 0. If > 0, the prompt will be appended with EOS tokens. This is useful for models that expect one or more EOS token at the end."); + +DEFINE_bool(warmup, false, "Whether to run a warmup run."); + +int32_t main(int32_t argc, char** argv) { + gflags::ParseCommandLineFlags(&argc, &argv, true); + + // Create a loader to get the data of the program file. There are other + // DataLoaders that use mmap() or point32_t to data that's already in memory, + // and users can create their own DataLoaders to load from arbitrary sources. + const char* model_path = FLAGS_model_path.c_str(); + + std::optional data_path = std::nullopt; + if (!FLAGS_data_path.empty()) { + data_path = FLAGS_data_path.c_str(); + } + + const char* tokenizer_path = FLAGS_tokenizer_path.c_str(); + + ET_LOG(Info, "prompt : %s", FLAGS_prompt.c_str()); + const char* prompt = FLAGS_prompt.c_str(); + + float temperature = FLAGS_temperature; + + int32_t seq_len = FLAGS_seq_len; + + int32_t cpu_threads = FLAGS_cpu_threads; + + bool warmup = FLAGS_warmup; + +#if defined(ET_USE_THREADPOOL) + uint32_t num_performant_cores = cpu_threads == -1 + ? ::executorch::extension::cpuinfo::get_num_performant_cores() + : static_cast(cpu_threads); + ET_LOG( + Info, "Resetting threadpool with num threads = %d", num_performant_cores); + if (num_performant_cores > 0) { + ::executorch::extension::threadpool::get_threadpool() + ->_unsafe_reset_threadpool(num_performant_cores); + } +#endif + + // create llama runner + std::unique_ptr runner = + example::create_llama_runner(model_path, tokenizer_path, data_path); + + if (runner == nullptr) { + ET_LOG(Error, "Failed to create llama runner"); + return 1; + } + + if (warmup) { + ET_LOG( + Info, + "Can not support warm-up (TODO : IRunner doesn't have warmup interface (pure virtual)"); +#if 0 + auto error = runner->warmup(prompt, /*max_new_tokens=*/seq_len); + if (error != executorch::runtime::Error::Ok) { + ET_LOG(Error, "Failed to warmup llama runner"); + return 1; + } +#endif + } + + // generate + executorch::extension::llm::GenerationConfig config{ + .seq_len = seq_len, .temperature = temperature}; + auto error = runner->generate(prompt, config, {}, {}); + if (error != executorch::runtime::Error::Ok) { + ET_LOG(Error, "Failed to generate token from llama runner"); + return 1; + } + + return 0; +} diff --git a/examples/samsung/executor_runner/llama/runner/CMakeLists.txt b/examples/samsung/executor_runner/llama/runner/CMakeLists.txt new file mode 100644 index 00000000000..648e8015749 --- /dev/null +++ b/examples/samsung/executor_runner/llama/runner/CMakeLists.txt @@ -0,0 +1,88 @@ +# Copyright (c) 2025 Samsung Electronics Co. LTD +# All rights reserved +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +message("###### Build ENN llama runner ######") +message(ROOT ${CMAKE_CURRENT_SOURCE_DIR}/../../../../..) +message(EXECUTORCH_ROOT ${EXECUTORCH_ROOT}) +message(CMAKE_CURRENT_SOURCE_DIR ${CMAKE_CURRENT_SOURCE_DIR}) +message(CMAKE_CURRENT_BINARY_DIR ${CMAKE_CURRENT_BINARY_DIR}) +message(CMAKE_CURRENT_LIST_DIR ${CMAKE_CURRENT_LIST_DIR}) + +include(${CMAKE_CURRENT_SOURCE_DIR}/../../../../../tools/cmake/Utils.cmake) +include(${CMAKE_CURRENT_SOURCE_DIR}/../../../../../tools/cmake/Codegen.cmake) + +list(TRANSFORM _enn_llama_runner__srcs + PREPEND "${CMAKE_CURRENT_SOURCE_DIR}/../../../../.." +) +list( + PREPEND + _enn_llama_runner__srcs + ${CMAKE_CURRENT_LIST_DIR}/runner.cpp + ${CMAKE_CURRENT_LIST_DIR}/decoder.cpp + ${CMAKE_CURRENT_LIST_DIR}/prefiller.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/../../../../../examples/models/llama/tokenizer/llama_tiktoken.cpp +) + +target_include_directories( + extension_module INTERFACE ${_common_include_directories} +) + +add_library(lib_tokenizer STATIC IMPORTED) +set_target_properties( + lib_tokenizer + PROPERTIES + IMPORTED_LOCATION + "${CMAKE_CURRENT_BINARY_DIR}/../../../../../../build_samsung_android/lib/libtokenizers.a" +) + +add_library(lib_extension_llm_runner STATIC IMPORTED) +set_target_properties( + lib_extension_llm_runner + PROPERTIES + IMPORTED_LOCATION + "${CMAKE_CURRENT_BINARY_DIR}/../../../../../../build_samsung_android/extension/llm/runner/libextension_llm_runner.a" +) + +if(CMAKE_TOOLCHAIN_IOS + OR ANDROID + OR APPLE +) + # Building a share library on iOS requires code signing On Android we see + # duplicated registration when using shared lib + add_library(enn_llama_runner STATIC ${_enn_llama_runner__srcs}) +else() + add_library(enn_llama_runner SHARED ${_enn_llama_runner__srcs}) +endif() + +# For extension_llm_runner +if(NOT TARGET extension_llm_runner) + add_subdirectory( + ${CMAKE_CURRENT_SOURCE_DIR}/../../../../../extension/llm/runner + ${CMAKE_CURRENT_BINARY_DIR}/../../../../../extension/llm/runner + ) +endif() + +set(enn_llama_runner_deps + executorch_core extension_data_loader extension_module extension_tensor + extension_flat_tensor lib_extension_llm_runner +) + +target_link_libraries(enn_llama_runner PUBLIC ${enn_llama_runner_deps}) +target_link_libraries(enn_llama_runner PUBLIC tokenizers::tokenizers) +target_link_libraries(enn_llama_runner PUBLIC lib_tokenizer) + +target_include_directories( + enn_llama_runner + PUBLIC + ${EXECUTORCH_ROOT}/.. + ${CMAKE_CURRENT_SOURCE_DIR}/../../../../../extension/llm/tokenizers/include + ${CMAKE_CURRENT_SOURCE_DIR}/../../../../../extension/llm/tokenizers/third-party/re2 + ${CMAKE_CURRENT_SOURCE_DIR}/../../../../../extension/llm/tokenizers/third-party/abseil-cpp + ${CMAKE_CURRENT_SOURCE_DIR}/../../../../../extension/llm/tokenizers/third-party/json/include + ${CMAKE_CURRENT_SOURCE_DIR}/../../../../../examples/models/llama/tokenizer + ${CMAKE_CURRENT_SOURCE_DIR}/../../../../../examples/samsung/executor_runner/llama/runner + ${CMAKE_CURRENT_SOURCE_DIR}/../../../../../extension/llm/runner +) diff --git a/examples/samsung/executor_runner/llama/runner/decoder.cpp b/examples/samsung/executor_runner/llama/runner/decoder.cpp new file mode 100644 index 00000000000..892d0f0863b --- /dev/null +++ b/examples/samsung/executor_runner/llama/runner/decoder.cpp @@ -0,0 +1,122 @@ +/* + * Copyright (c) 2025 Samsung Electronics Co. LTD + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +// Given inputs, run a text decoder and return logits. + +#include + +#include + +#include +#include + +#include + +namespace example { + +class IExecutionIOBinder { + virtual ~IExecutionIOBinder() = default; + virtual void beforeExecution() = 0; + virtual void afterExecution() = 0; +}; + +::executorch::runtime::Result Decoder::step( + executorch::extension::TensorPtr& tokens, + int64_t start_pos) { + // ET_LOG(Info, "Input token %" PRIu64, input_token); + auto method_meta = ET_UNWRAP(module_->method_meta("forward")); + // If only 1 input, we are not using kv cache + bool use_kv_cache = method_meta.num_inputs() > 1; + + if (use_kv_cache) { + // Size of the second argument. This could be either input_pos or + // cache_positions + + // Check if we are using cache positions instead of input pos. + auto second_input_info = ET_UNWRAP(method_meta.input_tensor_meta(1)); + // For input_pos, numel is 1, for cache_positions, numel is max_seq_len + auto sizes = second_input_info.sizes(); + // Assuming 1D tensor + ET_CHECK_OR_RETURN_ERROR( + sizes.size() == 1, + InvalidProgram, + "The second input tensor is not 1D tensor. Got dimension (%zu)", + sizes.size()); + auto numel = sizes[0]; + std::vector<::executorch::aten::SizesType> sizes_vec = {numel}; + + executorch::extension::TensorPtr start_pos_tensor; + + // TODO(sicheon.oh) : manage tensor for in/out + + int64_t dummy_pos = 1; + std::vector<::executorch::aten::SizesType> dummy_sizes_vec = {1}; + executorch::extension::TensorPtr dummy_io_tensor = + executorch::extension::empty( + dummy_sizes_vec, executorch::aten::ScalarType::Long); + + if (numel > 1) { + // If we are here, model is exported with cache_positions, create a tensor + // with the same length as input_ids. Assuming the last dimension is the + // one with the variable token length, for example [1, S] or [1, 1, S] + sizes_vec[sizes_vec.size() - 1] = tokens->numel(); + start_pos_tensor = executorch::extension::empty( + sizes_vec, ::executorch::aten::ScalarType::Long); + torch::executor::native::arange_out_impl( + start_pos, start_pos + tokens->numel(), 1.0, *start_pos_tensor); + } else { + // Assuming model is exported with input_pos, create a tensor with size 1 + start_pos_tensor = executorch::extension::from_blob( + &start_pos, sizes_vec, ::executorch::aten::ScalarType::Long); + } + ET_LOG( + Debug, + "Start pos tensor numel: %zu, tokens numel: %zu", + start_pos_tensor->numel(), + tokens->numel()); + + // TODO(sicheon.oh) : manage tensor for in/out + auto inputs_res = + io_manager_->prepare_decode(dummy_io_tensor, start_pos_tensor); + ET_CHECK_OK_OR_RETURN_ERROR(inputs_res.error()); + + auto outputs_res = module_->forward({tokens, start_pos_tensor}); + ET_CHECK_OK_OR_RETURN_ERROR(outputs_res.error()); + ET_CHECK_MSG( + outputs_res.get().size() == 1, + "More then one output returned from executing LLM."); + ET_CHECK_MSG( + outputs_res.get()[0].isTensor(), + "Non Tensor Output returned from executing LLM"); + + // TODO(sicheon.oh) : manage tensor for in/out + if (io_manager_->update_decode(outputs_res.get(), "forward") != + executorch::runtime::Error::Ok) { + ET_LOG(Error, "Failed to updated output by io_manager"); + } + + // Return the logits tensor + return outputs_res.get()[0].toTensor(); + } else { // no kv cache + (void)start_pos; // unused + + auto outputs_res = module_->forward(tokens); + ET_CHECK_OK_OR_RETURN_ERROR(outputs_res.error()); + ET_CHECK_MSG( + outputs_res.get().size() == 1, + "More then one output returned from executing LLM."); + ET_CHECK_MSG( + outputs_res.get()[0].isTensor(), + "Non Tensor Output returned from executing LLM"); + + // Return the logits tensor + return outputs_res.get()[0].toTensor(); + } +} + +} // namespace example diff --git a/examples/samsung/executor_runner/llama/runner/decoder.h b/examples/samsung/executor_runner/llama/runner/decoder.h new file mode 100644 index 00000000000..3f1f7e37088 --- /dev/null +++ b/examples/samsung/executor_runner/llama/runner/decoder.h @@ -0,0 +1,110 @@ +/* + * Copyright (c) 2025 Samsung Electronics Co. LTD + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +// Given inputs, run a text decoder in LLM and return the output. + +#pragma once + +#include +#include +#include + +#include + +namespace example { + +namespace llm = ::executorch::extension::llm; + +static constexpr auto default_method_name = "forward"; + +class ET_EXPERIMENTAL Decoder + : public executorch::extension::llm::TextDecoderRunner { + public: + explicit Decoder( + executorch::extension::Module* module, + llm::IOManager* io_manager) + : TextDecoderRunner(module, io_manager) {} + + virtual ::executorch::runtime::Result step( + executorch::extension::TensorPtr& input, + int64_t start_pos) override; + + inline ::executorch::runtime::Error load() override { + executorch::runtime::Error ret = executorch::runtime::Error::Ok; + + /** + * Dealing with multiple method. + */ + for (auto method : methods_.name) { + if (module_->is_method_loaded(method)) { + ET_LOG(Info, "Method \"%s\" is already loaded", method.c_str()); + continue; + } + + ret = module_->load_method(method); + if (ret != executorch::runtime::Error::Ok) { + ET_LOG(Error, "Failed to load method"); + return ret; + } + + ET_LOG(Info, "Loading Method (%s)", method.c_str()); + } + + ret = load_io_manager(); + if (ret != executorch::runtime::Error::Ok) { + ET_LOG(Error, "Failed to load io manager (%d)", ret); + return ret; + } + + return ret; + } + + inline bool is_method_loaded() override { + bool ret = false; + for (auto method : methods_.name) { + ret = module_->is_method_loaded(method); + if (ret == false) { + ET_LOG(Info, "method name \"%s\" is not loaded", method.c_str()); + return ret; + } + } + + return ret; + } + + inline void register_method(std::string& method) { + methods_.name[methods_.num++] = method; + for (auto method : methods_.name) { + ET_LOG(Info, "Registered method : \"%s\"", method.c_str()); + } + } + + private: + inline ::executorch::runtime::Error load_io_manager() { + executorch::runtime::Error ret = executorch::runtime::Error::Ok; + + return io_manager_->load( + methods_.name[static_cast(DecoderMode::PREFILL)], + methods_.name[static_cast(DecoderMode::DECODE)]); + } + + enum class DecoderMode : int { PREFILL = 0, DECODE = 1 }; + + struct RegisterdMethod { + uint32_t num = 0; + std::vector name = {default_method_name, default_method_name}; + std::vector objs = { + nullptr, + nullptr}; + }; + + private: + RegisterdMethod methods_; +}; + +} // namespace example diff --git a/examples/samsung/executor_runner/llama/runner/io_manager.h b/examples/samsung/executor_runner/llama/runner/io_manager.h new file mode 100644 index 00000000000..f86c9447cc3 --- /dev/null +++ b/examples/samsung/executor_runner/llama/runner/io_manager.h @@ -0,0 +1,227 @@ +/* + * Copyright (c) 2025 Samsung Electronics Co. LTD + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +#include + +#include + +namespace example { + +/** + * @brief Base class for managing input/output operations for LLM inference. + * + * IOManager provides an interface for handling the input preparation and + * output processing for both prefill and decode phases of LLM inference. + * Derived classes must implement the virtual methods to provide specific IO + * management functionality. + */ +class ET_EXPERIMENTAL LlmIOManager + : public executorch::extension::llm::IOManager { + public: + /** + * @brief Construct an IOManager bound to a Module. + * + * @param module The Module used for querying method metadata and execution. + */ + explicit LlmIOManager(executorch::extension::Module& module) + : IOManager(module), module_(module) {} + + /** + * @brief Virtual destructor to allow proper cleanup in derived classes. + */ + virtual ~LlmIOManager() = default; + + /** + * @brief Load the IO manager with method metadata for prefill and + * decode operations. + * + * @param prefill_method The prefill method to initialize with. + * @param decode_method The decode method to initialize with. + */ + ET_NODISCARD virtual executorch::runtime::Error load( + const std::string& prefill_method, + const std::string& decode_method) override { + return executorch::runtime::Error::Ok; + } + + /** + * @brief Reset the IO manager state. + * + * @param prefill_method The prefill method to reset with. + * @param decode_method The decode method to reset with. + */ + ET_NODISCARD virtual executorch::runtime::Error reset( + const std::string& prefill_method, + const std::string& decode_method) override { + (void)prefill_method; + (void)decode_method; + return executorch::runtime::Error::Ok; + } + + /** + * @brief Prepare inputs for the prefill phase of LLM inference. + * + * @param input The input tensor containing token IDs. + * @param start_pos The tensor containing the starting position of the current + * input within the context. + * @param prefill_method The prefill method to prepare inputs for. + * @return std::vector Vector of prepared inputs + * for the prefill method. + */ + virtual executorch::runtime::Result> + prepare_prefill( + const executorch::extension::TensorPtr& input, + const executorch::extension::TensorPtr& start_pos, + const std::string& prefill_method) override { + auto method_meta = module_.method_meta(prefill_method); + if (!method_meta.ok()) { + return method_meta.error(); + } + if (method_meta->num_inputs() != 2) { + ET_LOG( + Error, + "Expected 2 inputs for prefill method, got %zu. Likely the model takes the caches or mask as an argument which this IOManager does not support.", + method_meta->num_inputs()); + return executorch::runtime::Error::InvalidState; + } + // Cpu IO Manager supports dynamic shapes for prefill, so no work to be done + // here. + return std::vector{input, start_pos}; + } + + /** + * @brief Prepare inputs for the prefill phase using the default method name. + * + * Uses "forward" as the prefill method. + * + * @param input The input tensor containing token IDs. + * @param start_pos The tensor containing the starting position. + * @return Vector of prepared inputs for the prefill method. + */ + executorch::runtime::Result> + prepare_prefill( + const executorch::extension::TensorPtr& input, + const executorch::extension::TensorPtr& start_pos) { + return prepare_prefill(input, start_pos, "forward"); + } + + /** + * @brief Prepare inputs for the decode phase of LLM inference. + * + * @param input The input tensor containing token IDs. + * @param start_pos The tensor containing the starting position of the current + * input within the context. + * @param decode_method The decode method to prepare inputs for. + * @return std::vector Vector of prepared inputs + * for the decode method. + */ + virtual executorch::runtime::Result> + prepare_decode( + const executorch::extension::TensorPtr& input, + const executorch::extension::TensorPtr& start_pos, + const std::string& decode_method) override { + auto method_meta = module_.method_meta(decode_method); + if (!method_meta.ok()) { + return method_meta.error(); + } + if (method_meta->num_inputs() != 2) { + ET_LOG( + Error, + "Expected 2 inputs for decode method, got %zu. Likely the model takes the caches or mask as an argument which this IOManager does not support.", + method_meta->num_inputs()); + return executorch::runtime::Error::InvalidState; + } + // Cpu IO Manager supports dynamic shapes for prefill, so no work to be done + // here. + return std::vector{input, start_pos}; + } + + /** + * @brief Prepare inputs for the decode phase using the default method name. + * + * Uses "forward" as the decode method. + * + * @param input The input tensor containing token IDs. + * @param start_pos The tensor containing the starting position. + * @return Vector of prepared inputs for the decode method. + */ + executorch::runtime::Result> + prepare_decode( + const executorch::extension::TensorPtr& input, + const executorch::extension::TensorPtr& start_pos) { + return prepare_decode(input, start_pos, "forward"); + } + + /** + * @brief Process and update internal state with outputs from the prefill + * phase. + * + * @param prefill_method The prefill method to update with outputs. + * @param model_outputs Vector of outputs from the prefill method execution. + */ + ET_NODISCARD virtual executorch::runtime::Error update_prefill( + const std::vector& model_outputs, + const std::string& prefill_method) override { + (void)model_outputs; + (void)prefill_method; + // No post inference work to do. + return executorch::runtime::Error::Ok; + } + + /** + * @brief Process outputs from the prefill phase using the default method. + * + * Uses "forward" as the prefill method. + * + * @param model_outputs Vector of outputs from the prefill execution. + * @return Error code. + */ + ET_NODISCARD executorch::runtime::Error update_prefill( + const std::vector& model_outputs) { + return update_prefill(model_outputs, "forward"); + } + + /** + * @brief Process and update internal state with outputs from the decode + * phase. + * + * @param decode_method The decode method to update with outputs. + * @param model_outputs Vector of outputs from the decode method execution. + */ + ET_NODISCARD virtual executorch::runtime::Error update_decode( + const std::vector& model_outputs, + const std::string& decode_method) override { + (void)model_outputs; + (void)decode_method; + // No post inference work to do. + return executorch::runtime::Error::Ok; + } + + /** + * @brief Process outputs from the decode phase using the default method. + * + * Uses "forward" as the decode method. + * + * @param model_outputs Vector of outputs from the decode execution. + * @return Error code. + */ + ET_NODISCARD executorch::runtime::Error update_decode( + const std::vector& model_outputs) { + return update_decode(model_outputs, "forward"); + } + + private: + /** + * @brief Reference to the Module used for method metadata and execution. + */ + executorch::extension::module::Module& module_; +}; + +} // namespace example diff --git a/examples/samsung/executor_runner/llama/runner/prefiller.cpp b/examples/samsung/executor_runner/llama/runner/prefiller.cpp new file mode 100644 index 00000000000..911852a2c1d --- /dev/null +++ b/examples/samsung/executor_runner/llama/runner/prefiller.cpp @@ -0,0 +1,149 @@ +/* + * Copyright (c) 2025 Samsung Electronics Co. LTD + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +// Given a text prompt, encode it using tokenizer and prefill the KV cache of a +// LLM. + +#include +#include + +namespace example { + +Prefiller::Prefiller( + executorch::extension::llm::TextDecoderRunner* text_decoder_runner, + bool use_kv_cache, + bool enable_parallel_prefill, + int64_t max_seq_len) + : TextPrefiller( + text_decoder_runner, + use_kv_cache, + enable_parallel_prefill, + max_seq_len), + text_decoder_runner_(text_decoder_runner), + use_kv_cache_(use_kv_cache), + enable_parallel_prefill_(enable_parallel_prefill), + max_seq_len_(max_seq_len > 0 ? max_seq_len : 128) {} + +::executorch::runtime::Result Prefiller::prefill( + std::vector& prompt_tokens, + int64_t& start_pos) { + ET_CHECK_MSG(!prompt_tokens.empty(), "Prompt cannot be null"); + if (!text_decoder_runner_->is_method_loaded()) { + ET_CHECK_OK_OR_RETURN_ERROR(text_decoder_runner_->load()); + } + + // Check if we need to chunk the prompt tokens + int32_t num_prompt_tokens = prompt_tokens.size(); + + // If prompt tokens exceed max_seq_len_, we need to chunk them + ET_LOG( + Info, + "Micro Batch: num_prompt(%d), max_seq_len(%ld), start_pos(%ld)", + num_prompt_tokens, + max_seq_len_, + start_pos); + + if (num_prompt_tokens > max_seq_len_) { + uint64_t cur_token = 0; + int num_tokens_to_process = 0; + + while (num_tokens_to_process < num_prompt_tokens) { + auto num_tokens_to_prefill_with = std::min( + num_prompt_tokens - num_tokens_to_process, max_seq_len_); + + std::vector prompt_tokens_to_process( + num_tokens_to_prefill_with); + std::copy( + prompt_tokens.begin() + num_tokens_to_process, + prompt_tokens.begin() + num_tokens_to_process + + num_tokens_to_prefill_with, + prompt_tokens_to_process.begin()); + + // Process this chunk + auto chunk_result = prefill_chunk(prompt_tokens_to_process, start_pos); + ET_CHECK_OK_OR_RETURN_ERROR(chunk_result.error()); + cur_token = chunk_result.get(); + + start_pos += num_tokens_to_prefill_with; + num_tokens_to_process += num_tokens_to_prefill_with; + } + + return cur_token; + } else { + // If prompt tokens don't exceed max_seq_len_, process them directly + return prefill_chunk(prompt_tokens, start_pos); + } +} + +::executorch::runtime::Result Prefiller::prefill_chunk( + std::vector& prompt_tokens, + int64_t& start_pos) { + // enable_parallel_prefill_ maybe set even when not using kv cache + // When kv cache is not used, start pos is ignored + int32_t num_prompt_tokens = prompt_tokens.size(); + + // store the token + uint64_t cur_token; + if (enable_parallel_prefill_ || !use_kv_cache_) { + // TODO (sicheon.oh) : embedding token + + // initialize tensor wrappers + auto tokens = executorch::extension::from_blob( + prompt_tokens.data(), + {1, num_prompt_tokens}, + executorch::aten::ScalarType::Long); + + auto outputs_res = text_decoder_runner_->step(tokens, start_pos); + + ET_CHECK_OK_OR_RETURN_ERROR(outputs_res.error()); + ET_LOG( + Info, "Prefill token result numel(): %zu", outputs_res.get().numel()); + + start_pos += num_prompt_tokens; + cur_token = text_decoder_runner_->logits_to_token(outputs_res.get()); + } else { // sequential prefill + int64_t pos = 0; // position in the sequence + // NOLINTNEXTLINE(facebook-hte-ParameterUncheckedArrayBounds) + cur_token = prompt_tokens[0]; + + ET_LOG( + Info, + "num_prompt_token(%d), cur_token(%ld), start_pos(%ld)", + num_prompt_tokens, + cur_token, + start_pos); + + // initialize tensor wrappers + auto tokens = executorch::extension::from_blob( + &cur_token, {1, 1}, executorch::aten::ScalarType::Long); + + // run the first token and get back logits tensor. Assuming the first token + // is bos so don't callback. + auto logits_tensor = + ET_UNWRAP(text_decoder_runner_->step(tokens, start_pos)); + + pos += 1; // start the loop from index 1 + start_pos += 1; + + while (pos < num_prompt_tokens) { + // Run the model + // NOLINTNEXTLINE(facebook-hte-ParameterUncheckedArrayBounds) + cur_token = prompt_tokens[pos]; + + logits_tensor = ET_UNWRAP(text_decoder_runner_->step(tokens, start_pos)); + + pos++; + start_pos++; + } + + cur_token = text_decoder_runner_->logits_to_token(logits_tensor); + } + return cur_token; +} + +} // namespace example diff --git a/examples/samsung/executor_runner/llama/runner/prefiller.h b/examples/samsung/executor_runner/llama/runner/prefiller.h new file mode 100644 index 00000000000..42102cc7f25 --- /dev/null +++ b/examples/samsung/executor_runner/llama/runner/prefiller.h @@ -0,0 +1,66 @@ +/* + * Copyright (c) 2025 Samsung Electronics Co. LTD + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +// Given a text prompt, encode it using tokenizer and prefill the KV cache of a +// LLM. + +#pragma once + +#include +#include + +namespace example { + +class ET_EXPERIMENTAL Prefiller + : public executorch::extension::llm::TextPrefiller { + public: + Prefiller( + executorch::extension::llm::TextDecoderRunner* text_decoder_runner, + bool use_kv_cache, + bool enable_parallel_prefill, + int64_t max_seq_len = 128); + + virtual ~Prefiller() = default; + + /** + * Prefill an LLM Module with the given text input. + * @param prompt_tokens The text prompt tokens to the LLM Module. Encoded by + * tokenizer. + * @param start_pos The starting position in KV cache of the input in the LLM + * Module. + * @return The next token of the LLM Module after prefill. + */ + virtual ::executorch::runtime::Result prefill( + std::vector& prompt_tokens, + int64_t& start_pos) override; + + /** + * Helper method to prefill a chunk of tokens. + * @param prompt_tokens The chunk of text prompt tokens to process. + * @param start_pos The starting position in KV cache of the input in the LLM + * Module. + * @return The next token of the LLM Module after prefilling this chunk. + */ + virtual ::executorch::runtime::Result prefill_chunk( + std::vector& prompt_tokens, + int64_t& start_pos) override; + + private: + /** + * Note: TextPrefiller does not own the TextDecoderRunner instance. + * The responsibility of managing the lifecycle of TextDecoderRunner + * lies with the outer class or entity (likely Runner) that creates + * and passes the TextDecoderRunner instance to TextPrefiller. + */ + executorch::extension::llm::TextDecoderRunner* text_decoder_runner_; + bool use_kv_cache_; + bool enable_parallel_prefill_; + int64_t max_seq_len_; +}; + +} // namespace example diff --git a/examples/samsung/executor_runner/llama/runner/runner.cpp b/examples/samsung/executor_runner/llama/runner/runner.cpp new file mode 100644 index 00000000000..c09815d462a --- /dev/null +++ b/examples/samsung/executor_runner/llama/runner/runner.cpp @@ -0,0 +1,568 @@ +/* + * Copyright (c) Samsung Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +// A simple llama2 runner that includes preprocessing and post processing logic. +// The module takes in a string as input and emits a string as output. + +#include +#include +#include +#include + +#include +#include +#include + +#include +#include +#include +#include +#include + +namespace example { + +using ::executorch::extension::Module; +using ::executorch::runtime::Error; +using ::executorch::runtime::Result; + +namespace llm = ::executorch::extension::llm; + +/** + * @brief Loads a tokenizer from the specified path + * + * This function creates and initializes a tokenizer from a file, with options + * to customize special tokens and regex patterns. + * + * @param tokenizer_path Path to the tokenizer file + * @param special_tokens Optional list of special tokens to add to the tokenizer + * @param pattern Optional regex pattern for tokenization + * @param bos_token_index Index of the beginning-of-sequence token + * @param eos_token_index Index of the end-of-sequence token + * @return std::unique_ptr Initialized tokenizer instance + */ +ET_EXPERIMENTAL std::unique_ptr load_tokenizer( + const std::string& tokenizer_path, + std::unique_ptr> special_tokens = nullptr, + std::optional pattern = std::nullopt, + size_t bos_token_index = 0, + size_t eos_token_index = 1); + +/** + * @brief Creates a LLMRunner instance with the specified model and + * tokenizer + * + * This factory function creates and initializes a LLMRunner with all + * necessary components for text generation using the specified model and + * tokenizer. + * + * @param model_path Path to the model file + * @param tokenizer Initialized tokenizer instance + * @param data_path Optional path to additional data required by the model + * @param temperature Optional temperature parameter for controlling randomness + * (deprecated) + * @return std::unique_ptr Initialized LLMRunner instance + */ +ET_EXPERIMENTAL std::unique_ptr create_llm_runner( + const std::string& model_path, + std::unique_ptr<::tokenizers::Tokenizer> tokenizer, + std::optional data_path = std::nullopt, + float temperature = -1.0f); + +std::unique_ptr<::tokenizers::Tokenizer> load_llama_tokenizer( + const std::string& tokenizer_path, + Version version) { + auto special_tokens = get_special_tokens(version); + return load_tokenizer(tokenizer_path, std::move(special_tokens)); +} + +std::unique_ptr create_llama_runner( + const std::string& model_path, + const std::string& tokenizer_path, + std::optional data_path, + float temperature) { + ET_LOG( + Info, + "Creating LLaMa runner: model_path=%s, tokenizer_path=%s", + model_path.c_str(), + tokenizer_path.c_str()); + + // Create and load tokenizer + std::unique_ptr<::tokenizers::Tokenizer> tokenizer = + load_llama_tokenizer(tokenizer_path, Version::Default); + if (tokenizer == nullptr) { + ET_LOG( + Info, + "Failed to load %s as a Tiktoken, Sentencepiece or Llama2.c tokenizer, make sure the artifact is one of these types", + tokenizer_path.c_str()); + return nullptr; + } + ET_LOG( + Info, + "[Create_test_llm_runner] %s %s", + model_path.c_str(), + data_path->c_str()); + return create_llm_runner(model_path, std::move(tokenizer), data_path); +} + +static const std::unordered_map kMetadataMap{ + {llm::kEnableDynamicShape, "enable_dynamic_shape"}, + {llm::kBosId, "get_bos_id"}, + {llm::kEosIds, "get_eos_ids"}, + {llm::kMaxSeqLen, "get_max_seq_len"}, + {llm::kMaxContextLen, "get_max_context_len"}, + {llm::kVocabSize, "get_vocab_size"}, + {llm::kUseKVCache, "use_kv_cache"}, + {llm::kUseSDPAWithKVCache, "use_sdpa_with_kv_cache"}}; + +LLMRunner::LLMRunner( + std::unordered_map metadata, + std::unique_ptr<::tokenizers::Tokenizer> tokenizer, + std::unique_ptr<::executorch::extension::Module> module, + std::unique_ptr text_decoder_runner, + std::unique_ptr text_prefiller, + std::unique_ptr text_token_generator, + std::unique_ptr stats, + std::unique_ptr io_manager, + float temperature) + : tokenizer_(std::move(tokenizer)), + metadata_(std::move(metadata)), + module_(std::move(module)), + text_decoder_runner_(std::move(text_decoder_runner)), + text_prefiller_(std::move(text_prefiller)), + text_token_generator_(std::move(text_token_generator)), + io_manager_(std::move(io_manager)), + stats_(std::move(stats)), + temperature_(temperature) { + // Note: This constructor assumes that text_prefiller and text_token_generator + // already have references to the Module and TextDecoderRunner they need +} + +bool LLMRunner::is_loaded() const { + return text_prefiller_->is_loaded() && text_token_generator_->is_loaded(); +} + +Error LLMRunner::load() { + ET_CHECK_OK_OR_RETURN_ERROR(text_decoder_runner_->load()); + ET_CHECK_OK_OR_RETURN_ERROR(text_prefiller_->load()); + ET_CHECK_OK_OR_RETURN_ERROR(text_token_generator_->load()); + return Error::Ok; +} + +// Don't print with the same priority during warmup +#define RUNNER_ET_LOG(warmup, format, ...) \ + if (warmup) { \ + ET_LOG(Debug, format, __VA_ARGS__); \ + } else { \ + ET_LOG(Info, format, __VA_ARGS__); \ + } + +Error LLMRunner::generate_from_pos( + const std::string& prompt, + int64_t start_pos, + const llm::GenerationConfig& config, + std::function token_callback, + std::function stats_callback) { + // Prepare the inputs. + // Use ones-initialized inputs. + ET_CHECK_MSG(!prompt.empty(), "Prompt cannot be null"); + if (!is_loaded()) { + stats_->model_load_start_ms = llm::time_in_ms(); + ET_CHECK_OK_OR_RETURN_ERROR(load()); + stats_->model_load_end_ms = llm::time_in_ms(); + } + + if (config.warming) { + ET_LOG(Info, "Doing a warmup run..."); + } + + RUNNER_ET_LOG( + config.warming, + "RSS after loading model: %f MiB (0 if unsupported)", + llm::get_rss_bytes() / 1024.0 / 1024.0); + + // Wrap the token_callback with print function + std::function wrapped_callback = + [token_callback, config](const std::string& piece) { + if (!config.warming) { + llm::safe_printf(piece.c_str()); + fflush(stdout); + } + if (token_callback) { + token_callback(piece); + } + }; + // First token time only measures the time it takes to encode the prompt and + // return a response token. + + stats_->inference_start_ms = llm::time_in_ms(); + shouldStop_ = false; + + ET_LOG( + Info, + "################ input prompt ################\n %s\n", + prompt.c_str()); + + ::tokenizers::Result> encode_res = tokenizer_->encode( + prompt, + /*bos=*/config.num_bos, + /*eos=*/config.num_eos); + + ET_CHECK_TK_OK_OR_RETURN_ERROR( + encode_res.error(), "Failed to encode prompt %s", prompt.c_str()); + + // encode the (string) prompt into tokens sequence + std::vector prompt_tokens = encode_res.get(); + int num_prompt_tokens = prompt_tokens.size(); + + // Reduce max_context_len by start_pos + int64_t max_context_len = metadata_.at(llm::kMaxContextLen) - start_pos; + ET_CHECK_OR_RETURN_ERROR( + num_prompt_tokens >= 1, + InvalidArgument, + "Expected at least 1 prompt token"); + ET_CHECK_OR_RETURN_ERROR( + num_prompt_tokens < max_context_len, + InvalidArgument, + "num_prompt_tokens %d >= max_context_len %" PRId64 + ", Max seq length exceeded - please increase max seq len value in your export script", + num_prompt_tokens, + max_context_len); + + // Determine max_new_tokens using the GenerationConfig's resolve method, + // then subtract start_pos for max_new_tokens. + int max_new_tokens = + config.resolve_max_new_tokens(max_context_len, num_prompt_tokens); + + ET_LOG( + Info, + "Max new tokens resolved: %d, given start_pos %" PRId64 + ", num_prompt_tokens %zu, max_context_len %" PRId64, + max_new_tokens, + start_pos, + prompt_tokens.size(), + max_context_len); + ET_CHECK_OR_RETURN_ERROR( + max_new_tokens > 0, + InvalidArgument, + "Max new tokens %d is less than or equal to 0", + max_new_tokens); + // Prefill first + // Here feed all tokens to the model and get the next predicted token + // after the prompt. After that we will enter generate loop. + + // print prompts + if (config.echo) { + wrapped_callback(prompt); + } + int64_t pos = start_pos; + auto prefill_res = text_prefiller_->prefill(prompt_tokens, pos); + ET_CHECK_OK_OR_RETURN_ERROR(prefill_res.error()); + uint64_t cur_token = prefill_res.get(); + stats_->first_token_ms = llm::time_in_ms(); + stats_->prompt_eval_end_ms = llm::time_in_ms(); + + RUNNER_ET_LOG( + config.warming, + "\nRSS after prompt prefill: %f MiB (0 if unsupported)", + llm::get_rss_bytes() / 1024.0 / 1024.0); + + ::tokenizers::Result> newline_token = + tokenizer_->encode("\n"); + + ET_LOG(Info, "[ENN] first token %lu\n", cur_token); + + ET_LOG( + Info, + "################ The start of token generation ################\n"); + + // print the first token from prefill. No prev_token so use cur_token for it. + wrapped_callback( + ET_UNWRAP_TOKENIZER(tokenizer_->decode(cur_token, cur_token))); + + // start the main loop + prompt_tokens.push_back(cur_token); + + // Generate max_new_tokens - 1 because prefill already generated 1 token. + int64_t num_generated_tokens = ET_UNWRAP(text_token_generator_->generate( + prompt_tokens, + num_prompt_tokens, + max_new_tokens - 1, + temperature_ == -1.0f ? config.temperature : temperature_, + wrapped_callback)); + + // A temporary new line to separate stdout mixed with results. + wrapped_callback(ET_UNWRAP_TOKENIZER( + tokenizer_->decode(newline_token.get()[0], newline_token.get()[0]))); + + ET_LOG( + Info, "################ The end of token generation ################\n"); + + stats_->inference_end_ms = llm::time_in_ms(); + if (!config.warming) { + printf("\n"); + } + RUNNER_ET_LOG( + config.warming, + "RSS after finishing text generation: %f MiB (0 if unsupported)", + llm::get_rss_bytes() / 1024.0 / 1024.0); + + if (num_generated_tokens == max_new_tokens) { + RUNNER_ET_LOG(config.warming, "Max new tokens %i reached!", max_new_tokens); + } + + stats_->num_prompt_tokens = num_prompt_tokens; + stats_->num_generated_tokens = num_generated_tokens; + + if (config.warming) { + ET_LOG(Info, "Warmup run finished!"); + } else { + // Do not print report during warmup + print_report(*stats_); + } + if (stats_callback) { + stats_callback(*stats_); + } + + return Error::Ok; +} + +Error LLMRunner::generate( + const std::string& prompt, + const llm::GenerationConfig& config, + std::function token_callback, + std::function stats_callback) { + return generate_from_pos(prompt, 0, config, token_callback, stats_callback); +} + +Error LLMRunner::warmup(const std::string& prompt, int32_t max_new_tokens) { + // Create a GenerationConfig for warmup + llm::GenerationConfig config{ + .echo = false, .max_new_tokens = max_new_tokens, .warming = true}; + + // Call generate with the warmup config + Error err = generate(prompt, config); + + // Reset stats after warmup, not resetting the std::unique_ptr! + stats_->reset(); + return err; +} + +void LLMRunner::stop() { + if (is_loaded()) { + text_token_generator_->stop(); + } else { + ET_LOG(Error, "Token generator is not loaded, cannot stop"); + } +} + +std::unique_ptr load_tokenizer( + const std::string& tokenizer_path, + std::unique_ptr> special_tokens, + std::optional pattern, + size_t bos_token_index, + size_t eos_token_index) { + ::executorch::runtime::runtime_init(); + auto json_tokenizer = std::make_unique(); + if (json_tokenizer->load(tokenizer_path) == ::tokenizers::Error::Ok) { + ET_LOG(Info, "Loaded json tokenizer"); + return json_tokenizer; + } + std::unique_ptr<::tokenizers::Tiktoken> tiktoken_tokenizer; + if (special_tokens != nullptr && !pattern.has_value()) { + tiktoken_tokenizer = std::make_unique<::tokenizers::Tiktoken>( + std::move(special_tokens), bos_token_index, eos_token_index); + } else if (special_tokens != nullptr && pattern.has_value()) { + tiktoken_tokenizer = std::make_unique<::tokenizers::Tiktoken>( + pattern.value(), + std::move(special_tokens), + bos_token_index, + eos_token_index); + } else { + tiktoken_tokenizer = std::make_unique<::tokenizers::Tiktoken>(); + } + if (tiktoken_tokenizer->load(tokenizer_path) == ::tokenizers::Error::Ok) { + ET_LOG(Info, "Loaded TikToken tokenizer"); + return tiktoken_tokenizer; + } + + auto sp_tokenizer = std::make_unique<::tokenizers::SPTokenizer>(); + if (sp_tokenizer->load(tokenizer_path) == ::tokenizers::Error::Ok) { + ET_LOG(Info, "Loaded Sentencepiece tokenizer"); + return sp_tokenizer; + } + + auto bpe_tokenizer = std::make_unique<::tokenizers::Llama2cTokenizer>(); + if (bpe_tokenizer->load(tokenizer_path) == ::tokenizers::Error::Ok) { + ET_LOG(Info, "Loaded BPE tokenizer"); + return bpe_tokenizer; + } + + return nullptr; +} + +std::unordered_map get_llm_metadata( + tokenizers::Tokenizer* tokenizer, + Module* module) { + // Initialize metadata with default values + std::unordered_map metadata({ + {kMetadataMap.at(llm::kEnableDynamicShape), false}, + {kMetadataMap.at(llm::kMaxSeqLen), 128}, + {kMetadataMap.at(llm::kMaxContextLen), 128}, + {kMetadataMap.at(llm::kUseKVCache), true}, + {kMetadataMap.at(llm::kUseSDPAWithKVCache), false}, + }); + + // Read metadata from the model + auto method_names_result = module->method_names(); + if (method_names_result.error() != Error::Ok) { + ET_LOG(Error, "Failed reading method names"); + return metadata; + } + const auto method_names = method_names_result.get(); + + for (auto& pair : metadata) { + const auto& method_name = pair.first; + auto& value = pair.second; + + if (method_names.count(method_name)) { + auto get_result = module->get(method_name); + value = get_result.get().toScalar().to(); + } else { + ET_LOG( + Info, + "Method %s not found, using the default value %" PRId64, + method_name.c_str(), + value); + } + ET_LOG(Info, "Metadata: %s = %" PRId64, method_name.c_str(), value); + } + // Set tokenizer-related metadata + metadata[llm::kBosId] = tokenizer->bos_tok(); + metadata[llm::kVocabSize] = tokenizer->vocab_size(); + return metadata; +} + +std::unordered_set get_eos_ids( + tokenizers::Tokenizer* tokenizer, + Module* module) { + std::unordered_set eos_ids = {tokenizer->eos_tok()}; + // Get EOS IDs if available + auto method_names_result = module->method_names(); + if (method_names_result.error() != Error::Ok) { + ET_LOG(Error, "Failed reading method names"); + return eos_ids; + } + const auto method_names = method_names_result.get(); + + if (method_names.count(llm::kEosIds)) { + eos_ids.clear(); + auto execute_result = module->execute(llm::kEosIds); + if (execute_result.error() != Error::Ok) { + ET_LOG(Error, "Failed to execute %s", llm::kEosIds); + return eos_ids; + } + for (const auto& eos_id : execute_result.get()) { + auto value = eos_id.toScalar().to(); + eos_ids.emplace(value); + ET_LOG(Info, "eos_id = %" PRId64, value); + } + } + return eos_ids; +} + +std::unique_ptr create_llm_runner( + const std::string& model_path, + std::unique_ptr<::tokenizers::Tokenizer> tokenizer, + std::optional data_path, + float temperature) { + // Sanity check tokenizer + if (!tokenizer || !tokenizer->is_loaded()) { + ET_LOG(Error, "Tokenizer is null or not loaded"); + return nullptr; + } + + // Create the Module + std::unique_ptr module; + if (data_path.has_value()) { + module = std::make_unique( + model_path, data_path.value(), Module::LoadMode::File); + } else { + module = std::make_unique(model_path, Module::LoadMode::File); + } + + // Get methods + executorch::runtime::Result> methods_res = + module->method_names(); + if (methods_res.error() != executorch::runtime::Error::Ok) { + ET_CHECK_MSG(false, "Failed to get method names"); + } + + auto methods = methods_res.get(); + std::vector register_methods; + for (auto method : methods) { + if (kMetadataMap.find(method) != kMetadataMap.end()) { + ET_LOG(Info, "It is not method(%s)", method.c_str()); + break; + } + ET_LOG(Info, "Register method \"%s\" to decoder", method.c_str()); + register_methods.emplace_back(method); + } + + // Create IO Manager + auto io_manager = std::make_unique(*module.get()); + + // Get metadata from Module + ET_LOG(Info, "Reading metadata from model"); + auto metadata = get_llm_metadata(tokenizer.get(), module.get()); + + auto eos_ids = std::make_unique>( + get_eos_ids(tokenizer.get(), module.get())); + + // Create text_decoder_runner. Use a shared_ptr so that it can be shared with + // TextPrefiller and TextTokenGenerator + auto text_decoder_runner = + std::make_unique(module.get(), io_manager.get()); + + // Register methods to cope with multi-methods + for (auto method : register_methods) { + text_decoder_runner->register_method(method); + } + + // Create text_prefiller + auto text_prefiller = std::make_unique( + text_decoder_runner.get(), + metadata.at(llm::kUseKVCache), + metadata.at(llm::kEnableDynamicShape), + metadata.at(llm::kMaxSeqLen)); + + // Create text_token_generator with stats + auto stats = std::make_unique(); + auto text_token_generator = std::make_unique( + tokenizer.get(), + text_decoder_runner.get(), + metadata.at(llm::kUseKVCache), + std::move(eos_ids), + stats.get()); + + // Create and return the Runner instance + std::unique_ptr runner = std::make_unique( + std::move(metadata), + std::move(tokenizer), + std::move(module), + std::move(text_decoder_runner), + std::move(text_prefiller), + std::move(text_token_generator), + std::move(stats), + std::move(io_manager), + temperature); + + runner->load(); + + return runner; +} +} // namespace example diff --git a/examples/samsung/executor_runner/llama/runner/runner.h b/examples/samsung/executor_runner/llama/runner/runner.h new file mode 100644 index 00000000000..bd7443f9f61 --- /dev/null +++ b/examples/samsung/executor_runner/llama/runner/runner.h @@ -0,0 +1,185 @@ +/* + * Copyright (c) 2025 Samsung Electronics Co. LTD + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +#include +#include + +#include +#include +#include +#include + +#include +#include + +namespace example { + +namespace llm = ::executorch::extension::llm; + +class LLMRunner : public llm::IRunner { + public: + /** + * @brief Constructor for LLMRunner with dependency injection + * + * Creates a LLMRunner instance with all required components for text + * generation. + * + * @param metadata Key-value pairs containing model metadata (e.g., + * vocab_size, context_length) + * @param tokenizer Tokenizer for converting between text and token IDs + * @param module The underlying model module that performs inference + * @param text_decoder_runner Component responsible for running the decoder + * part of the model + * @param text_prefiller Component for handling the prefill phase of text + * generation + * @param text_token_generator Component for generating tokens during the + * decode phase + * @param stats Statistics tracking object for performance monitoring + * @param io_manager IOManager manages input/output operations for LLM + * inference. + * @param temperature Temperature parameter for controlling randomness in + * generation (deprecated). Please use GenerationConfig.temperature instead. + */ + explicit LLMRunner( + std::unordered_map metadata, + std::unique_ptr<::tokenizers::Tokenizer> tokenizer, + std::unique_ptr<::executorch::extension::Module> module, + std::unique_ptr text_decoder_runner, + std::unique_ptr text_prefiller, + std::unique_ptr text_token_generator, + std::unique_ptr stats, + std::unique_ptr io_manager, + float temperature = -1.0f); + + /** + * @brief Checks if the model is loaded and ready for inference + * + * @return bool True if the model is loaded, false otherwise + */ + bool is_loaded() const; + /** + * @brief Loads the model and prepares it for inference + * + * This method initializes all components and prepares the model for text + * generation. + * + * @return ::executorch::runtime::Error Success or error status + */ + ::executorch::runtime::Error load(); + /** + * @brief Generates text based on the provided prompt + * + * This method performs text generation using the loaded model. It processes + * the input prompt, runs the model in prefill and decode phases until max + * tokens to generate is reached or eos token is generated, then returns + * generated text and perf stats through callbacks. + * + * @param prompt The input text to generate from + * @param config Configuration parameters for text generation (e.g., + * max_new_tokens, temperature) + * @param token_callback Function called for each generated token with the + * decoded text + * @param stats_callback Function called with performance statistics + * @return ::executorch::runtime::Error Success or error status + */ + ::executorch::runtime::Error generate( + const std::string& prompt, + const llm::GenerationConfig& config, + std::function token_callback = {}, + std::function stats_callback = {}); + + /** + * @brief Generates text based on the provided prompt and start position + * + * This method performs text generation using the loaded model. It processes + * the input prompt, runs the model in prefill and decode phases using the + * start position until max tokens to generate is reached or eos token is + * generated, then returns generated text and perf stats through callbacks. + * + * @param prompt The input text to generate from + * @param start_pos The starting position in KV cache of the input + * @param config Configuration parameters for text generation (e.g., + * max_new_tokens, temperature) + * @param token_callback Function called for each generated token with the + * decoded text + * @param stats_callback Function called with performance statistics + * @return ::executorch::runtime::Error Success or error status + */ + ::executorch::runtime::Error generate_from_pos( + const std::string& prompt, + int64_t start_pos, + const llm::GenerationConfig& config, + std::function token_callback = {}, + std::function stats_callback = {}); + + /** + * @brief Warms up the model with a sample prompt + * + * This method runs a complete generation cycle without returning results, + * which helps initialize internal caches and optimize subsequent inferences. + * + * @param prompt The sample prompt to use for warmup + * @param max_new_tokens Maximum number of tokens to generate during warmup + * @return ::executorch::runtime::Error Success or error status + */ + ::executorch::runtime::Error warmup( + const std::string& prompt, + int32_t max_new_tokens); + /** + * @brief Stops the ongoing text generation process + * + * This method signals the generator to stop producing new tokens and + * terminate the current generation process. + */ + void stop(); + + /** + * @brief Force remove prefilled tokens and reset KV cache start position + * + * This mehtod removes the prefilled tokens from the KV cache and resets the + * start position to 0. + */ + void reset() {}; + + private: + bool shouldStop_{false}; + + // Components + std::unique_ptr<::tokenizers::Tokenizer> tokenizer_; + std::unordered_map metadata_; + std::unique_ptr<::executorch::extension::Module> + module_; // Manage module's lifecycle, make sure it outlives + // text_decoder_runner_llm::. + std::unique_ptr + text_decoder_runner_; // Manage text_decoder_runner_'s lifecycle, make + // sure it outlives text_prefiller_ & + // text_token_generator_. + std::unique_ptr text_prefiller_; + std::unique_ptr text_token_generator_; + std::unique_ptr io_manager_; + + // Stats + std::unique_ptr stats_; + + // temperature. + // Deprecated, we should rely on the temperature in GenerationConfig instead. + float temperature_ = -1.0f; +}; + +std::unique_ptr create_llama_runner( + const std::string& model_path, + const std::string& tokenizer_path, + std::optional data_path = std::nullopt, + float temperature = -1.0f); + +std::unique_ptr load_llama_tokenizer( + const std::string& tokenizer_path); + +} // namespace example diff --git a/examples/samsung/executor_runner/llama/runner/token_generator.h b/examples/samsung/executor_runner/llama/runner/token_generator.h new file mode 100644 index 00000000000..dcb3c8c5d75 --- /dev/null +++ b/examples/samsung/executor_runner/llama/runner/token_generator.h @@ -0,0 +1,171 @@ +/* + * Copyright (c) 2025 Samsung Electronics Co. LTD + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +// Generate tokens in a loop. +#pragma once + +#include +#include +#include +#include + +namespace example { + +class ET_EXPERIMENTAL TokenGenerator { + public: + TokenGenerator( + ::tokenizers::Tokenizer* tokenizer, + executorch::extension::llm::TextDecoderRunner* text_decoder_runner, + bool use_kv_cache, + std::unique_ptr>&& eos_ids, + executorch::extension::llm::Stats* stats) + : tokenizer_(tokenizer), + text_decoder_runner_(text_decoder_runner), + eos_ids_(std::move(eos_ids)), + use_kv_cache_(use_kv_cache), + stats_(stats) {} + + virtual ~TokenGenerator() = default; + + /** + * Token generation loop. + * @param tokens prompt tokens as well as the first token generated by + * prefill. + * @param start_pos the start position of the new tokens, based on how many + * prompt tokens is prefilled. + * @param max_new_tokens Maximum number of new tokens to generate. + * @param temperature controls the randomness of predictions by scaling the + * logits before applying softmax. A higher temperature results in more + * random predictions, while a lower temperature results in more deterministic + * predictions. + * @param token_callback what to do after a token is generated. + * @return how many tokens are generated. + */ + inline ::executorch::runtime::Result generate( + std::vector tokens, + int64_t start_pos, + int32_t max_new_tokens, + float temperature = 0.0f, + const std::function& token_callback = {}) { + ET_CHECK_MSG( + !tokens.empty(), "Token generation loop shouldn't take empty tokens"); + int64_t pos = start_pos; // position in the sequence + + std::vector token_data; // allocate space for the tokens + std::vector token_shape; + + // Token after prefill + uint64_t cur_token = tokens.back(); + uint64_t prev_token; + + if (use_kv_cache_) { + // hard code these to size 1 as kv cache is locked to static size right + // now. + token_data = {cur_token}; + token_shape = {1, 1}; + } else { + token_data = tokens; + token_shape = {1, static_cast(tokens.size())}; + } + + // initialize tensor wrappers + auto tokens_managed = executorch::extension::from_blob( + token_data.data(), token_shape, executorch::aten::ScalarType::Long); + + should_stop_ = false; + + // Generate our tokens + while (pos < start_pos + max_new_tokens) { + // Run the model + auto logits_res = text_decoder_runner_->step(tokens_managed, pos); + + ET_CHECK_OK_OR_RETURN_ERROR(logits_res.error()); + executorch::aten::Tensor& logits_tensor = logits_res.get(); + + prev_token = cur_token; + + stats_->on_sampling_begin(); + cur_token = + text_decoder_runner_->logits_to_token(logits_tensor, temperature); + stats_->on_sampling_end(); + + pos++; + + if (use_kv_cache_) { + // update the token tensor. token_data will not be empty. + // NOLINTNEXTLINE(facebook-hte-LocalUncheckedArrayBounds) + token_data[0] = cur_token; + } else { + // push it to the back + token_data.push_back(cur_token); + ET_CHECK_OK_OR_RETURN_ERROR(executorch::extension::resize_tensor_ptr( + tokens_managed, {1, static_cast(token_data.size())})); + } + + // print the token as string, decode it with the Tokenizer object + token_callback( + ET_UNWRAP_TOKENIZER(tokenizer_->decode(prev_token, cur_token))); + + if (should_stop_) { + break; + } + + // data-dependent terminating condition: we have n_eos_ number of EOS + if (eos_ids_->find(cur_token) != eos_ids_->end()) { + // ET_LOG(Info, "Reached to the end of generation"); + break; + } + } + return pos - start_pos; + } + + /** + * Stop the generation loop. + */ + inline void stop() { + should_stop_ = true; + } + + /** + * Load the necessary resources for TokenGenerator. + * This method should be called before using the generate() method. + */ + ::executorch::runtime::Error load() { + // return text_decoder_runner_->load(); + return ::executorch::runtime::Error::Ok; + } + + /** + * Check if the TokenGenerator has been successfully loaded. + * @return True if the resources are loaded, false otherwise. + */ + bool inline is_loaded() const { + // Implementation to check if resources are loaded + return tokenizer_->is_loaded() && text_decoder_runner_->is_method_loaded(); + } + + private: + /** + * Note: TokenGenerator does not own the tokenizer_ and + * text_decoder_runner_. The lifecycle of these objects should be managed + * externally, likely in the Runner. This class assumes that the provided + * pointers remain valid for the duration of its use. + */ + ::tokenizers::Tokenizer* tokenizer_; + executorch::extension::llm::TextDecoderRunner* text_decoder_runner_; + std::unique_ptr> eos_ids_; + bool use_kv_cache_; + + // state machine + bool should_stop_ = false; + + // stats + executorch::extension::llm::Stats* stats_; +}; + +} // namespace example From 0d5077aaa621b02e5cb65a02f8b188492e34bbec Mon Sep 17 00:00:00 2001 From: Jiseong Oh Date: Tue, 13 Oct 2026 19:36:31 +0900 Subject: [PATCH 02/17] Enable pte and ptd loading in the ENN LLM runner Load the program together with its .ptd weights so the runner can use programs whose weights are stored externally. Signed-off-by: Jiseong Oh Co-authored-by: brucekk.kim --- .../runtime/enn_api_implementation.cpp | 2 +- .../samsung/runtime/enn_api_implementation.h | 14 +++++ backends/samsung/runtime/enn_backend.cpp | 50 +++++++++++++-- backends/samsung/runtime/enn_executor.cpp | 63 ++++++++++++++++--- backends/samsung/runtime/enn_executor.h | 8 ++- 5 files changed, 119 insertions(+), 18 deletions(-) diff --git a/backends/samsung/runtime/enn_api_implementation.cpp b/backends/samsung/runtime/enn_api_implementation.cpp index 0cdf810d5c4..9ea5317b870 100644 --- a/backends/samsung/runtime/enn_api_implementation.cpp +++ b/backends/samsung/runtime/enn_api_implementation.cpp @@ -96,7 +96,7 @@ Error EnnApi::loadApiLib() { ENN_LOAD_API_FUNC( libenn_public_api_, EnnGetFileDescriptorFromEnnBuffer, this); ENN_LOAD_API_FUNC(libenn_public_api_, EnnOpenModelFromFd, this); - + ENN_LOAD_API_FUNC(libenn_public_api_, EnnOpenModelFromFdWithWeight, this); return Error::Ok; } diff --git a/backends/samsung/runtime/enn_api_implementation.h b/backends/samsung/runtime/enn_api_implementation.h index 3e0bf1c5f2a..675d877565b 100644 --- a/backends/samsung/runtime/enn_api_implementation.h +++ b/backends/samsung/runtime/enn_api_implementation.h @@ -39,6 +39,13 @@ class EnnApi { const uint32_t size, EnnModelId* model_id); EnnReturn (*EnnOpenModelFromFd)(int _fd, EnnModelId* model_id); + EnnReturn (*EnnOpenModelFromFdWithWeight)( + const int fd, + const uint32_t size, + const uint32_t offset, + EnnBufferPtr* weights, + uint32_t n_weights, + EnnModelId* model_id); EnnReturn (*EnnSetFastIpc)(void); EnnReturn (*EnnUnsetFastIpc)(void); EnnReturn (*EnnExecuteModelFastIpc)( @@ -136,6 +143,13 @@ typedef EnnReturn (*EnnReleaseBuffer_fn)(EnnBufferPtr buf); typedef EnnReturn ( *EnnGetFileDescriptorFromEnnBuffer_fn)(EnnBufferPtr buffer, int32_t* fd); typedef EnnReturn (*EnnOpenModelFromFd_fn)(int _fd, EnnModelId* model_id); +typedef EnnReturn (*EnnOpenModelFromFdWithWeight_fn)( + const int fd, + const uint32_t size, + const uint32_t offset, + EnnBufferPtr* weights, + uint32_t n_weights, + EnnModelId* model_id); } // namespace enn } // namespace executor diff --git a/backends/samsung/runtime/enn_backend.cpp b/backends/samsung/runtime/enn_backend.cpp index 81484afc44c..e7195749030 100644 --- a/backends/samsung/runtime/enn_backend.cpp +++ b/backends/samsung/runtime/enn_backend.cpp @@ -13,9 +13,11 @@ #include #include #include +#include #include #include +#include #include #pragma clang diagnostic ignored "-Wglobal-constructors" @@ -23,6 +25,10 @@ namespace torch { namespace executor { +using executorch::runtime::NamedDataMap; + +const char ennWeightBufferKey[] = "enn_weight_buffer_hash_id"; + class EnnBackend final : public PyTorchBackendInterface { public: ~EnnBackend() = default; @@ -31,6 +37,11 @@ class EnnBackend final : public PyTorchBackendInterface { return true; } + std::string decode_spec_value(const CompileSpec& spec) const { + const char* char_ptr = static_cast(spec.value.buffer); + return std::string(char_ptr, spec.value.nbytes); + } + Result init( BackendInitContext& context, FreeableBuffer* processed, @@ -41,12 +52,40 @@ class EnnBackend final : public PyTorchBackendInterface { const char* binary_buf_addr = reinterpret_cast(processed->data()); size_t buf_size = processed->size(); - Error err = executor->initialize(binary_buf_addr, buf_size); - if (err != Error::Ok) { - ENN_LOG_ERROR("Exynos backend initialize failed."); - executor->~EnnExecutor(); - return err; + const NamedDataMap* named_data_map = context.get_named_data_map(); + + auto it = std::find_if( + compile_specs.begin(), + compile_specs.end(), + [](const CompileSpec& spec) { + return std::strcmp(spec.key, ennWeightBufferKey) == 0; + }); + + if (it != compile_specs.end()) { + std::string weight_key = decode_spec_value(*it); + ET_LOG( + Info, "EnnBackend, Weight Buffer Enabled for %s", weight_key.c_str()); + Result weight = + named_data_map->get_data(weight_key.c_str()); + const char* weight_buf_addr = + reinterpret_cast(weight->data()); + size_t weight_size = weight->size(); + Error err = executor->initialize( + binary_buf_addr, buf_size, weight_buf_addr, weight_size, true); + if (err != Error::Ok) { + ENN_LOG_ERROR("Exynos backend initialize failed."); + executor->~EnnExecutor(); + return err; + } + } else { + Error err = executor->initialize(binary_buf_addr, buf_size); + if (err != Error::Ok) { + ENN_LOG_ERROR("Exynos backend initialize failed."); + executor->~EnnExecutor(); + return err; + } } + return executor; } @@ -88,6 +127,7 @@ class EnnBackend final : public PyTorchBackendInterface { executor->~EnnExecutor(); } } + }; // namespace executor namespace { diff --git a/backends/samsung/runtime/enn_executor.cpp b/backends/samsung/runtime/enn_executor.cpp index 99a2ce11e5e..204ef71f950 100644 --- a/backends/samsung/runtime/enn_executor.cpp +++ b/backends/samsung/runtime/enn_executor.cpp @@ -20,6 +20,15 @@ namespace executor { namespace enn { Error EnnExecutor::initialize(const char* binary_buf_addr, size_t buf_size) { + return initialize(binary_buf_addr, buf_size, nullptr, 0, false); +} + +Error EnnExecutor::initialize( + const char* binary_buf_addr, + size_t buf_size, + const char* weight_buf_addr, + size_t weight_size, + bool use_weight_fd) { EXYNOS_ATRACE_FUNCTION_LINE(); auto sm_instance = executorch::backends::enn::shared_memory_manager:: SharedMemoryManager::getInstance(); @@ -30,18 +39,44 @@ Error EnnExecutor::initialize(const char* binary_buf_addr, size_t buf_size) { ET_LOG(Info, "Start to open model %p, %zu", binary_buf_addr, buf_size); - EnnBufferPtr shared_buffer = nullptr; - if (sm_instance->query(&shared_buffer, binary_buf_addr, buf_size)) { + EnnBufferPtr pte_buffer = nullptr; + if (sm_instance->query(&pte_buffer, binary_buf_addr, buf_size)) { int32_t fd; - if (shared_buffer->va == binary_buf_addr && - !enn_api_inst->EnnGetFileDescriptorFromEnnBuffer(shared_buffer, &fd)) { - ret = enn_api_inst->EnnOpenModelFromFd(fd, &model_id_); - if (ret == ENN_RET_SUCCESS) { - ET_LOG(Info, "Opened model from file descriptor, so fd is closed"); - sm_instance->free(shared_buffer->va); + if (pte_buffer->va == binary_buf_addr && + !enn_api_inst->EnnGetFileDescriptorFromEnnBuffer(pte_buffer, &fd)) { + if (use_weight_fd) { + EnnBufferPtr ptd_buffer = nullptr; + ET_CHECK_OR_RETURN_ERROR( + sm_instance->query(&ptd_buffer, weight_buf_addr, weight_size), + Internal, + "PTD EnnBuffer not found."); + ET_LOG( + Info, + "Open FD & weights. fd: %d, model size %zu, offset: %u, weight buffer: %p", + fd, + buf_size, + pte_buffer->offset, + ptd_buffer->va); + ret = enn_api_inst->EnnOpenModelFromFdWithWeight( + fd, buf_size, pte_buffer->offset, &ptd_buffer, 1, &model_id_); + ET_CHECK_OR_RETURN_ERROR( + ret == ENN_RET_SUCCESS, Internal, "Open model with weights failed."); + ET_LOG(Info, "Opened model from file descriptor with weights"); + sm_instance->free(ptd_buffer->va); + sm_instance->free(pte_buffer->va); + } else { + ret = enn_api_inst->EnnOpenModelFromFd(fd, &model_id_); + if (ret == ENN_RET_SUCCESS) { + ET_LOG(Info, "Opened model from file descriptor, so fd is closed"); + sm_instance->free(pte_buffer->va); + } } } } + ET_CHECK_OR_RETURN_ERROR( + !use_weight_fd || model_id_, + Internal, + "Failed to open the model with external weights from file descriptor."); if (!model_id_) { ET_LOG(Info, "Open model from memory"); ret = enn_api_inst->EnnOpenModelFromMemory( @@ -96,10 +131,13 @@ Error EnnExecutor::eval( } EXYNOS_ATRACE_END(); - ENN_LOG_DEBUG("Start to execute model."); + ENN_LOG_DEBUG("Starting model execution"); auto ret = enn_api_inst->EnnExecuteModel(model_id_); if (ret != ENN_RET_SUCCESS) { - ENN_LOG_ERROR("EnnExecuteModel Failed"); + ENN_LOG_ERROR( + "EnnExecuteModel failed, ret=%d (model_id=0x%" PRIX64 ")", + ret, + model_id_); return Error::Internal; } @@ -124,6 +162,11 @@ EnnExecutor::~EnnExecutor() { enn_api_inst->EnnReleaseBuffers(alloc_buffer_, num_of_buffers); } enn_api_inst->EnnCloseModel(model_id_); + + ET_LOG( + Info, + "EnnExecutor destroyed: model 0x%" PRIX64 " closed, resources released", + model_id_); } } // namespace enn diff --git a/backends/samsung/runtime/enn_executor.h b/backends/samsung/runtime/enn_executor.h index 68ee1f56847..f3b7c3ad111 100644 --- a/backends/samsung/runtime/enn_executor.h +++ b/backends/samsung/runtime/enn_executor.h @@ -7,10 +7,8 @@ * */ #pragma once - #include #include - #include namespace torch { @@ -25,6 +23,12 @@ struct DataBuffer { class EnnExecutor { public: Error initialize(const char* binary_buf_addr, size_t buf_size); + Error initialize( + const char* binary_buf_addr, + size_t buf_size, + const char* weight_buf_addr, + size_t weight_size, + bool use_weight_fd); Error eval( const std::vector& inputs, From 5a23a94e87f3e3d6e9ba33c0813ce65c15241286 Mon Sep 17 00:00:00 2001 From: Jiseong Oh Date: Tue, 13 Oct 2026 19:37:01 +0900 Subject: [PATCH 03/17] Fix weight buffer lifetime with the debug logger Keep the weight buffer alive for as long as the debug logger reads it. Signed-off-by: Jiseong Oh --- backends/samsung/runtime/enn_backend.cpp | 10 ++++++++-- backends/samsung/runtime/enn_executor.cpp | 1 - backends/samsung/runtime/enn_executor.h | 3 +++ 3 files changed, 11 insertions(+), 3 deletions(-) diff --git a/backends/samsung/runtime/enn_backend.cpp b/backends/samsung/runtime/enn_backend.cpp index e7195749030..3ee8e2645c5 100644 --- a/backends/samsung/runtime/enn_backend.cpp +++ b/backends/samsung/runtime/enn_backend.cpp @@ -63,13 +63,19 @@ class EnnBackend final : public PyTorchBackendInterface { if (it != compile_specs.end()) { std::string weight_key = decode_spec_value(*it); - ET_LOG( - Info, "EnnBackend, Weight Buffer Enabled for %s", weight_key.c_str()); Result weight = named_data_map->get_data(weight_key.c_str()); const char* weight_buf_addr = reinterpret_cast(weight->data()); size_t weight_size = weight->size(); + + ET_LOG( + Info, "EnnBackend, [%s] Weight Buffer Enabled for %s size(%ld), addr(%p)", + context.get_method_name(), + weight_key.c_str(), weight_size, weight_buf_addr); + + FreeableBuffer* wgt = new FreeableBuffer(std::move(weight.get())); + Error err = executor->initialize( binary_buf_addr, buf_size, weight_buf_addr, weight_size, true); if (err != Error::Ok) { diff --git a/backends/samsung/runtime/enn_executor.cpp b/backends/samsung/runtime/enn_executor.cpp index 204ef71f950..7572f278510 100644 --- a/backends/samsung/runtime/enn_executor.cpp +++ b/backends/samsung/runtime/enn_executor.cpp @@ -62,7 +62,6 @@ Error EnnExecutor::initialize( ET_CHECK_OR_RETURN_ERROR( ret == ENN_RET_SUCCESS, Internal, "Open model with weights failed."); ET_LOG(Info, "Opened model from file descriptor with weights"); - sm_instance->free(ptd_buffer->va); sm_instance->free(pte_buffer->va); } else { ret = enn_api_inst->EnnOpenModelFromFd(fd, &model_id_); diff --git a/backends/samsung/runtime/enn_executor.h b/backends/samsung/runtime/enn_executor.h index f3b7c3ad111..b71dd827e0d 100644 --- a/backends/samsung/runtime/enn_executor.h +++ b/backends/samsung/runtime/enn_executor.h @@ -11,6 +11,8 @@ #include #include +#include + namespace torch { namespace executor { namespace enn { @@ -48,6 +50,7 @@ class EnnExecutor { EnnBufferPtr* alloc_buffer_ = nullptr; int32_t num_of_inputs_ = 0; int32_t num_of_outputs_ = 0; + executorch::runtime::FreeableBuffer* weight_ = nullptr; }; } // namespace enn From 3aa745d337297185f22add747975b23af11d0af5 Mon Sep 17 00:00:00 2001 From: Jiseong Oh Date: Tue, 13 Oct 2026 19:37:21 +0900 Subject: [PATCH 04/17] Enable LLM weight sharing via cache Methods that share weights reuse one buffer through a cache keyed by the weight-buffer hash id, and their models are opened from the file descriptor with the external weights. Signed-off-by: Jiseong Oh Co-authored-by: Hoon Choi --- backends/samsung/runtime/enn_backend.cpp | 74 ++++++++++++++++++++---- backends/samsung/runtime/enn_executor.h | 14 +++++ 2 files changed, 78 insertions(+), 10 deletions(-) diff --git a/backends/samsung/runtime/enn_backend.cpp b/backends/samsung/runtime/enn_backend.cpp index 3ee8e2645c5..98899f74255 100644 --- a/backends/samsung/runtime/enn_backend.cpp +++ b/backends/samsung/runtime/enn_backend.cpp @@ -18,6 +18,8 @@ #include #include +#include +#include #include #pragma clang diagnostic ignored "-Wglobal-constructors" @@ -29,6 +31,9 @@ using executorch::runtime::NamedDataMap; const char ennWeightBufferKey[] = "enn_weight_buffer_hash_id"; +// Weight cache LUT for reusing FreeableBuffer across multiple init calls +static std::unordered_map g_weight_cache_lut; + class EnnBackend final : public PyTorchBackendInterface { public: ~EnnBackend() = default; @@ -61,20 +66,55 @@ class EnnBackend final : public PyTorchBackendInterface { return std::strcmp(spec.key, ennWeightBufferKey) == 0; }); + const char* weight_buf_addr; + size_t weight_size; + if (it != compile_specs.end()) { std::string weight_key = decode_spec_value(*it); - Result weight = - named_data_map->get_data(weight_key.c_str()); - const char* weight_buf_addr = - reinterpret_cast(weight->data()); - size_t weight_size = weight->size(); - ET_LOG( - Info, "EnnBackend, [%s] Weight Buffer Enabled for %s size(%ld), addr(%p)", - context.get_method_name(), - weight_key.c_str(), weight_size, weight_buf_addr); + // Check if weight buffer is already cached + FreeableBuffer* wgt = nullptr; + auto cache_it = g_weight_cache_lut.find(weight_key); + + if (cache_it != g_weight_cache_lut.end()) { + weight_buf_addr = + reinterpret_cast((cache_it->second)->data()); + weight_size = (cache_it->second)->size(); + wgt = cache_it->second; + + ET_LOG( + Info, + "EnnBackend, [%s] Weight Buffer from Cache for %s size(%ld), addr(%p)", + context.get_method_name(), + weight_key.c_str(), + weight_size, + weight_buf_addr); + + // For cached buffers, we don't own it, so don't store weight_key + } else { + // Load weight from named_data_map and cache it + Result weight = + named_data_map->get_data(weight_key.c_str()); + + weight_buf_addr = reinterpret_cast(weight->data()); + weight_size = weight->size(); + + ET_LOG( + Info, + "EnnBackend, [%s] Weight Buffer Loaded for %s size(%ld), addr(%p)", + context.get_method_name(), + weight_key.c_str(), + weight_size, + weight_buf_addr); + + wgt = new FreeableBuffer(std::move(weight.get())); - FreeableBuffer* wgt = new FreeableBuffer(std::move(weight.get())); + // Cache the weight buffer + g_weight_cache_lut[weight_key] = wgt; + + // Store the weight_key for cleanup in destroy + executor->set_weight_key(weight_key); + } Error err = executor->initialize( binary_buf_addr, buf_size, weight_buf_addr, weight_size, true); @@ -130,6 +170,20 @@ class EnnBackend final : public PyTorchBackendInterface { EXYNOS_ATRACE_NAME_LINE("backend_destroy"); if (handle != nullptr) { auto executor = static_cast(handle); + + // Find and remove the cached buffer that belongs to this executor + const std::string& weight_key = executor->get_weight_key(); + if (!weight_key.empty()) { + auto it = g_weight_cache_lut.find(weight_key); + if (it != g_weight_cache_lut.end()) { + ET_LOG( + Info, "Deleting cached buffer for key: %s", weight_key.c_str()); + delete it->second; + g_weight_cache_lut.erase(it); + ET_LOG(Info, "Cached buffer removed from LUT"); + } + } + executor->~EnnExecutor(); } } diff --git a/backends/samsung/runtime/enn_executor.h b/backends/samsung/runtime/enn_executor.h index b71dd827e0d..f7429f4f8b4 100644 --- a/backends/samsung/runtime/enn_executor.h +++ b/backends/samsung/runtime/enn_executor.h @@ -9,6 +9,7 @@ #pragma once #include #include +#include #include #include @@ -43,6 +44,18 @@ class EnnExecutor { return num_of_outputs_; } + executorch::runtime::FreeableBuffer* get_weight() const { + return weight_; + } + + void set_weight_key(const std::string& key) { + weight_key_ = key; + } + + const std::string& get_weight_key() const { + return weight_key_; + } + ~EnnExecutor(); private: @@ -51,6 +64,7 @@ class EnnExecutor { int32_t num_of_inputs_ = 0; int32_t num_of_outputs_ = 0; executorch::runtime::FreeableBuffer* weight_ = nullptr; + std::string weight_key_; }; } // namespace enn From fb9867b5d0cfbaf44569cc529b8344eee47e3485 Mon Sep 17 00:00:00 2001 From: Jiseong Oh Date: Tue, 13 Oct 2026 19:37:21 +0900 Subject: [PATCH 05/17] Implement core LLM runner logic and IO management Add the runner flow and the IO manager that prepares the token, mask, RoPE and KV-cache inputs and folds the KV outputs back into the cache. Signed-off-by: Jiseong Oh --- .../enn_executor_llama_runner.cpp | 6 +- .../llama/runner/CMakeLists.txt | 1 + .../executor_runner/llama/runner/constants.h | 23 ++ .../executor_runner/llama/runner/decoder.cpp | 147 +++++-- .../executor_runner/llama/runner/decoder.h | 58 ++- .../llama/runner/io_manager.cpp | 364 ++++++++++++++++++ .../executor_runner/llama/runner/io_manager.h | 209 +++++++--- .../llama/runner/prefiller.cpp | 11 +- .../executor_runner/llama/runner/runner.cpp | 86 ++++- .../executor_runner/llama/runner/runner.h | 5 +- extension/llm/runner/text_decoder_runner.cpp | 1 + 11 files changed, 771 insertions(+), 140 deletions(-) create mode 100644 examples/samsung/executor_runner/llama/runner/constants.h create mode 100644 examples/samsung/executor_runner/llama/runner/io_manager.cpp diff --git a/examples/samsung/executor_runner/enn_executor_llama_runner.cpp b/examples/samsung/executor_runner/enn_executor_llama_runner.cpp index 54b9d1eb868..a14f5f033e6 100644 --- a/examples/samsung/executor_runner/enn_executor_llama_runner.cpp +++ b/examples/samsung/executor_runner/enn_executor_llama_runner.cpp @@ -55,6 +55,8 @@ DEFINE_int32( DEFINE_bool(warmup, false, "Whether to run a warmup run."); +DEFINE_bool(decode_only, false, "Whether to run with decode method only."); + int32_t main(int32_t argc, char** argv) { gflags::ParseCommandLineFlags(&argc, &argv, true); @@ -81,6 +83,8 @@ int32_t main(int32_t argc, char** argv) { bool warmup = FLAGS_warmup; + bool decode_only = FLAGS_decode_only; + #if defined(ET_USE_THREADPOOL) uint32_t num_performant_cores = cpu_threads == -1 ? ::executorch::extension::cpuinfo::get_num_performant_cores() @@ -95,7 +99,7 @@ int32_t main(int32_t argc, char** argv) { // create llama runner std::unique_ptr runner = - example::create_llama_runner(model_path, tokenizer_path, data_path); + example::create_llama_runner(model_path, tokenizer_path, data_path, decode_only); if (runner == nullptr) { ET_LOG(Error, "Failed to create llama runner"); diff --git a/examples/samsung/executor_runner/llama/runner/CMakeLists.txt b/examples/samsung/executor_runner/llama/runner/CMakeLists.txt index 648e8015749..016f3289142 100644 --- a/examples/samsung/executor_runner/llama/runner/CMakeLists.txt +++ b/examples/samsung/executor_runner/llama/runner/CMakeLists.txt @@ -23,6 +23,7 @@ list( ${CMAKE_CURRENT_LIST_DIR}/runner.cpp ${CMAKE_CURRENT_LIST_DIR}/decoder.cpp ${CMAKE_CURRENT_LIST_DIR}/prefiller.cpp + ${CMAKE_CURRENT_LIST_DIR}/io_manager.cpp ${CMAKE_CURRENT_SOURCE_DIR}/../../../../../examples/models/llama/tokenizer/llama_tiktoken.cpp ) diff --git a/examples/samsung/executor_runner/llama/runner/constants.h b/examples/samsung/executor_runner/llama/runner/constants.h new file mode 100644 index 00000000000..af0d3a52d2f --- /dev/null +++ b/examples/samsung/executor_runner/llama/runner/constants.h @@ -0,0 +1,23 @@ +/* + * Copyright (c) Samsung Electronics Co. LTD + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once +// constants for LLM runtime +namespace example { + +// Exynos Runtime metadata key constants +inline constexpr auto kPrefillArLen = "get_prefill_ar_len"; +inline constexpr auto kNumOfLayers = "get_n_layers"; +inline constexpr auto kNumOfHiddenDim = "get_hidden_dim"; +inline constexpr auto kNumOfHeads = "get_n_heads"; +inline constexpr auto kNumOfKVHeads = "get_n_kv_heads"; + +static constexpr auto default_method_name = "forward"; +static constexpr auto default_method_name_prefill = "prefill_forward"; +static constexpr auto default_method_name_decode = "decode_forward"; +} // namespace executorch::extension::llm diff --git a/examples/samsung/executor_runner/llama/runner/decoder.cpp b/examples/samsung/executor_runner/llama/runner/decoder.cpp index 892d0f0863b..0293c7c5abb 100644 --- a/examples/samsung/executor_runner/llama/runner/decoder.cpp +++ b/examples/samsung/executor_runner/llama/runner/decoder.cpp @@ -28,8 +28,16 @@ class IExecutionIOBinder { ::executorch::runtime::Result Decoder::step( executorch::extension::TensorPtr& tokens, int64_t start_pos) { + + std::string method_stage; + if (tokens->numel() > 1) { + method_stage = methods_.name[static_cast(DecoderMode::PREFILL)]; + } else { + method_stage = methods_.name[static_cast(DecoderMode::DECODE)]; + } + // ET_LOG(Info, "Input token %" PRIu64, input_token); - auto method_meta = ET_UNWRAP(module_->method_meta("forward")); + auto method_meta = ET_UNWRAP(module_->method_meta(method_stage.c_str())); // If only 1 input, we are not using kv cache bool use_kv_cache = method_meta.num_inputs() > 1; @@ -50,54 +58,51 @@ ::executorch::runtime::Result Decoder::step( auto numel = sizes[0]; std::vector<::executorch::aten::SizesType> sizes_vec = {numel}; + auto n_inputs = method_meta.num_inputs(); + auto input_info = ET_UNWRAP(method_meta.input_tensor_meta(2)); + sizes = input_info.sizes(); + if (sizes.size() != 2) ET_LOG(Error, "third input of %zu is invalid dimension : %zu", n_inputs, sizes.size()); + + input_info = ET_UNWRAP(method_meta.input_tensor_meta(3)); + sizes = input_info.sizes(); + if (sizes.size() != 4) ET_LOG(Error, "forth input of %zu is invalid dimension : %zu", n_inputs, sizes.size()); + executorch::extension::TensorPtr start_pos_tensor; - // TODO(sicheon.oh) : manage tensor for in/out - - int64_t dummy_pos = 1; - std::vector<::executorch::aten::SizesType> dummy_sizes_vec = {1}; - executorch::extension::TensorPtr dummy_io_tensor = - executorch::extension::empty( - dummy_sizes_vec, executorch::aten::ScalarType::Long); - - if (numel > 1) { - // If we are here, model is exported with cache_positions, create a tensor - // with the same length as input_ids. Assuming the last dimension is the - // one with the variable token length, for example [1, S] or [1, 1, S] - sizes_vec[sizes_vec.size() - 1] = tokens->numel(); - start_pos_tensor = executorch::extension::empty( - sizes_vec, ::executorch::aten::ScalarType::Long); - torch::executor::native::arange_out_impl( - start_pos, start_pos + tokens->numel(), 1.0, *start_pos_tensor); - } else { - // Assuming model is exported with input_pos, create a tensor with size 1 - start_pos_tensor = executorch::extension::from_blob( - &start_pos, sizes_vec, ::executorch::aten::ScalarType::Long); + // Assuming model is exported with input_pos, create a tensor with size 1 + start_pos_tensor = executorch::extension::from_blob( + &start_pos, sizes_vec, ::executorch::aten::ScalarType::Long); + + std::vector inputs; + + if (method_stage == methods_.name[static_cast(DecoderMode::PREFILL)]) { + auto inputs_res = + io_manager_->prepare_prefill(tokens, start_pos_tensor); + ET_CHECK_OK_OR_RETURN_ERROR(inputs_res.error()); + inputs = inputs_res.get(); + } else if (method_stage == methods_.name[static_cast(DecoderMode::DECODE)]) { + auto inputs_res = + io_manager_->prepare_decode(tokens, start_pos_tensor); + ET_CHECK_OK_OR_RETURN_ERROR(inputs_res.error()); + inputs = inputs_res.get(); } - ET_LOG( - Debug, - "Start pos tensor numel: %zu, tokens numel: %zu", - start_pos_tensor->numel(), - tokens->numel()); - - // TODO(sicheon.oh) : manage tensor for in/out - auto inputs_res = - io_manager_->prepare_decode(dummy_io_tensor, start_pos_tensor); - ET_CHECK_OK_OR_RETURN_ERROR(inputs_res.error()); - - auto outputs_res = module_->forward({tokens, start_pos_tensor}); + + auto outputs_res = module_->execute(method_stage.c_str(), inputs); ET_CHECK_OK_OR_RETURN_ERROR(outputs_res.error()); - ET_CHECK_MSG( - outputs_res.get().size() == 1, - "More then one output returned from executing LLM."); ET_CHECK_MSG( outputs_res.get()[0].isTensor(), "Non Tensor Output returned from executing LLM"); - // TODO(sicheon.oh) : manage tensor for in/out - if (io_manager_->update_decode(outputs_res.get(), "forward") != - executorch::runtime::Error::Ok) { - ET_LOG(Error, "Failed to updated output by io_manager"); + if (method_stage == methods_.name[static_cast(DecoderMode::PREFILL)]) { + if (io_manager_->update_prefill(outputs_res.get()) != + executorch::runtime::Error::Ok) { + ET_LOG(Error, "Failed to updated prefill output by io_manager"); + } + } else if (method_stage == methods_.name[static_cast(DecoderMode::DECODE)]) { + if (io_manager_->update_decode(outputs_res.get()) != + executorch::runtime::Error::Ok) { + ET_LOG(Error, "Failed to updated decode output by io_manager"); + } } // Return the logits tensor @@ -105,7 +110,8 @@ ::executorch::runtime::Result Decoder::step( } else { // no kv cache (void)start_pos; // unused - auto outputs_res = module_->forward(tokens); + auto outputs_res = module_->execute(method_stage.c_str(), tokens); + ET_LOG(Info, "After module forward"); ET_CHECK_OK_OR_RETURN_ERROR(outputs_res.error()); ET_CHECK_MSG( outputs_res.get().size() == 1, @@ -119,4 +125,61 @@ ::executorch::runtime::Result Decoder::step( } } +::executorch::runtime::Error Decoder::create_shared_planned_memory() { + executorch::runtime::Error ret = executorch::runtime::Error::Ok; + + if (shared_method_memory_ != nullptr) { + ET_LOG(Info, "Shared Planned memory is already created"); + return ret; + } + + auto prefill_meta = module_->method_meta( + methods_.name[static_cast(DecoderMode::PREFILL)].c_str()); + auto decode_meta = module_->method_meta( + methods_.name[static_cast(DecoderMode::DECODE)].c_str()); + if (!prefill_meta.ok() || !decode_meta.ok()) { + ET_LOG(Error, "Failed to get method metadata for shared memory setup"); + return executorch::runtime::Error::InvalidArgument; + } + + size_t prefill_bufs = prefill_meta->num_memory_planned_buffers(); + size_t decode_bufs = decode_meta->num_memory_planned_buffers(); + size_t num_buffers = std::max(prefill_bufs, decode_bufs); + + ET_LOG(Info, "Prefill_bufs(%ld), decode_bufs(%ld) num_buffers(%ld)", + prefill_bufs, decode_bufs, num_buffers); + + shared_method_memory_ = + std::make_unique(); + shared_method_memory_->buffers.reserve(num_buffers); + shared_method_memory_->spans.reserve(num_buffers); + + for (size_t i = 0; i < num_buffers; ++i) { + int64_t prefill_size = + (i < prefill_bufs) + ? prefill_meta->memory_planned_buffer_size(i).get() + : 0; + int64_t decode_size = + (i < decode_bufs) + ? decode_meta->memory_planned_buffer_size(i).get() + : 0; + size_t max_size = static_cast(std::max(prefill_size, decode_size)); + ET_LOG(Info, "Prefill_size(%ld), decode_size(%ld) max_size(%ld)", + prefill_size, decode_size, max_size); + + shared_method_memory_->buffers.emplace_back(max_size); + shared_method_memory_->spans.emplace_back( + shared_method_memory_->buffers.back().data(), max_size); + } + + shared_method_memory_->allocator = + std::make_unique<::executorch::runtime::HierarchicalAllocator>( + ::executorch::runtime::Span<::executorch::runtime::Span>( + shared_method_memory_->spans.data(), + shared_method_memory_->spans.size())); + return ret; +} + + + } // namespace example diff --git a/examples/samsung/executor_runner/llama/runner/decoder.h b/examples/samsung/executor_runner/llama/runner/decoder.h index 3f1f7e37088..59a1debdc6c 100644 --- a/examples/samsung/executor_runner/llama/runner/decoder.h +++ b/examples/samsung/executor_runner/llama/runner/decoder.h @@ -13,22 +13,27 @@ #include #include #include - +#include #include namespace example { namespace llm = ::executorch::extension::llm; -static constexpr auto default_method_name = "forward"; - class ET_EXPERIMENTAL Decoder : public executorch::extension::llm::TextDecoderRunner { public: + struct SharedMethodMemory { + std::vector> buffers; + std::vector> spans; + std::unique_ptr allocator; + }; + explicit Decoder( executorch::extension::Module* module, llm::IOManager* io_manager) - : TextDecoderRunner(module, io_manager) {} + : TextDecoderRunner(module, io_manager), + shared_method_memory_(nullptr) {} virtual ::executorch::runtime::Result step( executorch::extension::TensorPtr& input, @@ -37,6 +42,14 @@ class ET_EXPERIMENTAL Decoder inline ::executorch::runtime::Error load() override { executorch::runtime::Error ret = executorch::runtime::Error::Ok; + + if (methods_.num > 1) { + ret = create_shared_planned_memory(); + if (ret != executorch::runtime::Error::Ok) { + ET_LOG(Error, "Failed to create share planned memory"); + } + } + /** * Dealing with multiple method. */ @@ -46,10 +59,12 @@ class ET_EXPERIMENTAL Decoder continue; } - ret = module_->load_method(method); - if (ret != executorch::runtime::Error::Ok) { - ET_LOG(Error, "Failed to load method"); - return ret; + if (shared_method_memory_) { + ret = module_->load_method(method, shared_method_memory_->allocator.get()); + if (ret != executorch::runtime::Error::Ok) { + ET_LOG(Error, "Failed to load method"); + return ret; + } } ET_LOG(Info, "Loading Method (%s)", method.c_str()); @@ -79,8 +94,12 @@ class ET_EXPERIMENTAL Decoder inline void register_method(std::string& method) { methods_.name[methods_.num++] = method; - for (auto method : methods_.name) { - ET_LOG(Info, "Registered method : \"%s\"", method.c_str()); + + ET_LOG(Info, "method[%d] : %s", methods_.num - 1, method.c_str()); + + int i = 0; + for (int i = 0; i < methods_.num; i++) { + ET_LOG(Info, "Registered method[%d] : \"%s\"", i, methods_.name[i].c_str()); } } @@ -88,12 +107,21 @@ class ET_EXPERIMENTAL Decoder inline ::executorch::runtime::Error load_io_manager() { executorch::runtime::Error ret = executorch::runtime::Error::Ok; - return io_manager_->load( - methods_.name[static_cast(DecoderMode::PREFILL)], - methods_.name[static_cast(DecoderMode::DECODE)]); + if (methods_.num > 1) { + ret = io_manager_->load( + methods_.name[static_cast(DecoderMode::PREFILL)], + methods_.name[static_cast(DecoderMode::DECODE)]); + } else { + ret = io_manager_->load( + methods_.name[static_cast(DecoderMode::DECODE)], + methods_.name[static_cast(DecoderMode::DECODE)]); + } + return ret; } - enum class DecoderMode : int { PREFILL = 0, DECODE = 1 }; + ::executorch::runtime::Error create_shared_planned_memory(); + + enum class DecoderMode : int { PREFILL = 1, DECODE = 0 }; struct RegisterdMethod { uint32_t num = 0; @@ -105,6 +133,8 @@ class ET_EXPERIMENTAL Decoder private: RegisterdMethod methods_; + + std::unique_ptr shared_method_memory_; }; } // namespace example diff --git a/examples/samsung/executor_runner/llama/runner/io_manager.cpp b/examples/samsung/executor_runner/llama/runner/io_manager.cpp new file mode 100644 index 00000000000..d6f30435657 --- /dev/null +++ b/examples/samsung/executor_runner/llama/runner/io_manager.cpp @@ -0,0 +1,364 @@ +/* + * Copyright (c) 2025 Samsung Electronics Co. LTD + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#include +#include + +namespace example { + executorch::runtime::Result> + LlmIOManager::prepare_prefill( + const executorch::extension::TensorPtr& input, + const executorch::extension::TensorPtr& start_pos, + const std::string& prefill_method) { + (void)prefill_method; + + std::vector inputs; + inputs.reserve(3 + config_.n_layers * 2); + + //re-producessing input token tensor + const int64_t* prompt_tokens = input->mutable_data_ptr(); + size_t prompt_token_len = input->numel(); + + input_tokens_.clear(); + input_tokens_.assign(prompt_tokens, prompt_tokens + prompt_token_len); + input_tokens_.resize(config_.prefill_seq_len, input_tokens_.back()); + + auto input_tokens_tensor = executorch::extension::from_blob( + input_tokens_.data(), {1, static_cast(config_.prefill_seq_len)}, + executorch::aten::ScalarType::Long); + + inputs.emplace_back(input_tokens_tensor); + + //re-producessing start_pos tensor + const int64_t* start_position = start_pos->mutable_data_ptr(); + + pos_ids_.clear(); + pos_ids_.resize(config_.prefill_seq_len); + for(int64_t i = 0; i < config_.prefill_seq_len; ++i) { + if (i < prompt_token_len) pos_ids_[i] = i; + else pos_ids_[i] = prompt_token_len - 1; + } + auto start_pos_tensor = executorch::extension::from_blob( + pos_ids_.data(), {static_cast(config_.prefill_seq_len)}, executorch::aten::ScalarType::Long); + + inputs.emplace_back(start_pos_tensor); + + current_pos_ = *start_position; + + //attention tensor + update_attn_mask(current_pos_, prompt_token_len); + + attn_mask_tensor_ = executorch::extension::from_blob( + attn_mask_.data(), + {static_cast(config_.attn_mask_rows), + static_cast(config_.attn_mask_cols)}); + + inputs.emplace_back(attn_mask_tensor_); + + prefill_k_input_tensors_.clear(); + prefill_v_input_tensors_.clear(); + prefill_k_output_tensors_.clear(); + prefill_v_output_tensors_.clear(); + + for (size_t l = 0; l < config_.n_layers; l++) { + auto k_in = executorch::extension::from_blob( + prefill_k_input_[l].data(), + {1, + static_cast(config_.n_kv_heads), + static_cast(prefill_cache_len_), + static_cast(config_.head_dim)}, + executorch::aten::ScalarType::Half); + prefill_k_input_tensors_.push_back(k_in); + inputs.emplace_back(k_in); + } + + for (size_t l = 0; l < config_.n_layers; l++) { + auto v_in = executorch::extension::from_blob( + prefill_v_input_[l].data(), + {1, + static_cast(config_.n_kv_heads), + static_cast(prefill_cache_len_), + static_cast(config_.head_dim)}, + executorch::aten::ScalarType::Half); + prefill_v_input_tensors_.push_back(v_in); + inputs.emplace_back(v_in); + } + + dump_kv_input("prefill"); + + return inputs; + } + + executorch::runtime::Result> + LlmIOManager::prepare_decode( + const executorch::extension::TensorPtr& input, + const executorch::extension::TensorPtr& start_pos, + const std::string& decode_method) { + (void)decode_method; + + if (!prefill_done_) { + copy_prefill_to_decode_input(); + prefill_done_ = true; + } + + const int64_t* start_position = start_pos->mutable_data_ptr(); + current_pos_ = *start_position; + + update_attn_mask(current_pos_, 1); + + std::vector inputs; + inputs.reserve(3 + config_.n_layers * 2); + inputs.emplace_back(input); + inputs.emplace_back(start_pos); + + attn_mask_tensor_ = executorch::extension::from_blob( + attn_mask_.data(), + {1, + static_cast(config_.attn_mask_cols)}); + inputs.emplace_back(attn_mask_tensor_); + + decode_k_input_tensors_.clear(); + decode_v_input_tensors_.clear(); + decode_k_output_tensors_.clear(); + decode_v_output_tensors_.clear(); + + for (size_t l = 0; l < config_.n_layers; l++) { + auto k_in = executorch::extension::from_blob( + decode_k_input_[l].data(), + {1, + static_cast(config_.n_kv_heads), + static_cast(decode_cache_len_), + static_cast(config_.head_dim)}, + executorch::aten::ScalarType::Half); + decode_k_input_tensors_.push_back(k_in); + inputs.emplace_back(k_in); + } + for (size_t l = 0; l < config_.n_layers; l++) { + auto v_in = executorch::extension::from_blob( + decode_v_input_[l].data(), + {1, + static_cast(config_.n_kv_heads), + static_cast(decode_cache_len_), + static_cast(config_.head_dim)}, + executorch::aten::ScalarType::Half); + decode_v_input_tensors_.push_back(v_in); + inputs.emplace_back(v_in); + } + + dump_kv_input("decode"); + + return inputs; + } + + executorch::runtime::Error LlmIOManager::update_prefill( + const std::vector& model_outputs, + const std::string& prefill_method) { + (void)prefill_method; + // model_outputs: [logits, k0, ..., k_{L-1}, v0, ..., v_{L-1}] + // k_out shape: [1, n_kv_heads, prefill_seq_len, head_dim] + // Accumulate prefill output into prefill input buffers. The copy to + // decode input is deferred until the first prepare_decode call so that + // multiple micro-batch prefill iterations work correctly. + + for (size_t l = 0; l < config_.n_layers; l++) { + const auto& k_out = model_outputs[1 + l].toTensor(); + const _Float16* k_data = k_out.const_data_ptr<_Float16>(); + copy_to_cache( + k_data, + config_.prefill_seq_len, + prefill_k_input_[l].data(), + prefill_cache_len_, + current_pos_); + } + for (size_t l = 0; l < config_.n_layers; l++) { + const auto& v_out = + model_outputs[1 + config_.n_layers + l].toTensor(); + const _Float16* v_data = v_out.const_data_ptr<_Float16>(); + copy_to_cache( + v_data, + config_.prefill_seq_len, + prefill_v_input_[l].data(), + prefill_cache_len_, + current_pos_); + } + + dump_kv_output("prefill", model_outputs); + dump_step_++; + + return executorch::runtime::Error::Ok; + } + + executorch::runtime::Error LlmIOManager::update_decode( + const std::vector& model_outputs, + const std::string& decode_method) { + (void)decode_method; + // model_outputs: [logits, k0, ..., k_{L-1}, v0, ..., v_{L-1}] + // k_out shape: [1, n_kv_heads, 1, head_dim] + // Copy decode output into decode input cache at current_pos_. + + for (size_t l = 0; l < config_.n_layers; l++) { + const auto& k_out = model_outputs[1 + l].toTensor(); + const _Float16* k_data = k_out.const_data_ptr<_Float16>(); + copy_to_cache( + k_data, + 1, + decode_k_input_[l].data(), + decode_cache_len_, + current_pos_); + } + for (size_t l = 0; l < config_.n_layers; l++) { + const auto& v_out = + model_outputs[1 + config_.n_layers + l].toTensor(); + const _Float16* v_data = v_out.const_data_ptr<_Float16>(); + copy_to_cache( + v_data, + 1, + decode_v_input_[l].data(), + decode_cache_len_, + current_pos_); + } + + dump_kv_output("decode", model_outputs); + dump_step_++; + + return executorch::runtime::Error::Ok; + } + + void LlmIOManager::allocate_buffers() { + size_t prefill_in_size = config_.n_kv_heads * prefill_cache_len_ * + config_.head_dim; + size_t decode_in_size = config_.n_kv_heads * decode_cache_len_ * + config_.head_dim; + + prefill_k_input_.resize(config_.n_layers); + prefill_v_input_.resize(config_.n_layers); + decode_k_input_.resize(config_.n_layers); + decode_v_input_.resize(config_.n_layers); + + for (size_t l = 0; l < config_.n_layers; l++) { + prefill_k_input_[l].resize(prefill_in_size, 0.0f); + prefill_v_input_[l].resize(prefill_in_size, 0.0f); + decode_k_input_[l].resize(decode_in_size, 0.0f); + decode_v_input_[l].resize(decode_in_size, 0.0f); + } + + attn_mask_.resize( + config_.attn_mask_rows * config_.attn_mask_cols, -1.0f); + } + + void LlmIOManager::copy_prefill_to_decode_input() { + for (size_t l = 0; l < config_.n_layers; l++) { + for (size_t h = 0; h < config_.n_kv_heads; h++) { + const _Float16* src = prefill_k_input_[l].data() + + h * prefill_cache_len_ * config_.head_dim; + _Float16* dst = decode_k_input_[l].data() + + h * decode_cache_len_ * config_.head_dim; + std::memcpy( + dst, src, prefill_cache_len_ * config_.head_dim * sizeof(_Float16)); + } + for (size_t h = 0; h < config_.n_kv_heads; h++) { + const _Float16* src = prefill_v_input_[l].data() + + h * prefill_cache_len_ * config_.head_dim; + _Float16* dst = decode_v_input_[l].data() + + h * decode_cache_len_ * config_.head_dim; + std::memcpy( + dst, src, prefill_cache_len_ * config_.head_dim * sizeof(_Float16)); + } + } + } + + void LlmIOManager::copy_to_cache( + const _Float16* src, + size_t new_len, + _Float16* cache, + size_t cache_len, + size_t pos) { + for (size_t h = 0; h < config_.n_kv_heads; h++) { + const _Float16* src_head = src + h * new_len * config_.head_dim; + _Float16* dst_head = + cache + h * cache_len * config_.head_dim + pos * config_.head_dim; + std::memcpy( + dst_head, src_head, new_len * config_.head_dim * sizeof(_Float16)); + } + } + + void LlmIOManager::update_attn_mask(size_t pos, size_t seq_len) { + std::fill(attn_mask_.begin(), attn_mask_.end(), 0.0f); + size_t visible = 0; + for (size_t r = 0; r < seq_len && r < config_.attn_mask_rows; r++) { + visible = pos + r + 1; + if (visible > config_.attn_mask_cols) { + visible = config_.attn_mask_cols; + } + for (size_t c = 0; c < visible; c++) { + attn_mask_[r * config_.attn_mask_cols + c] = 1.0f; + } + } + + //Mask padding with the value of last token. + for (size_t r = seq_len; r < config_.attn_mask_rows; r++) { + for (size_t c = 0; c < visible; c++) { + attn_mask_[r * config_.attn_mask_cols + c] = 1.0f; + } + } + } + + void LlmIOManager::dump_buffer( + const std::string& tag, size_t layer, const _Float16* data, size_t size) { + if (!dump_enabled_) { + return; + } + char filename[512]; + std::snprintf( + filename, sizeof(filename), + "%s/%s_layer%zu_pos%zu_step%zu.bin", + dump_dir_.c_str(), tag.c_str(), layer, current_pos_, dump_step_); + + FILE* f = std::fopen(filename, "wb"); + if (f) { + std::fwrite(data, sizeof(_Float16), size, f); + std::fclose(f); + } + } + + void LlmIOManager::dump_kv_input(const std::string& phase) { + if (!dump_enabled_) { + return; + } + const bool is_prefill = (phase == "prefill"); + const auto& k_bufs = is_prefill ? prefill_k_input_ : decode_k_input_; + const auto& v_bufs = is_prefill ? prefill_v_input_ : decode_v_input_; + for (size_t l = 0; l < config_.n_layers; l++) { + dump_buffer(phase + "_k_input", l, k_bufs[l].data(), k_bufs[l].size()); + dump_buffer(phase + "_v_input", l, v_bufs[l].data(), v_bufs[l].size()); + } + } + + void LlmIOManager::dump_kv_output( + const std::string& phase, + const std::vector& model_outputs) { + if (!dump_enabled_) { + return; + } + for (size_t l = 0; l < config_.n_layers; l++) { + const auto& k_out = model_outputs[1 + l].toTensor(); + const auto& v_out = + model_outputs[1 + config_.n_layers + l].toTensor(); + dump_buffer( + phase + "_k_output", + l, + k_out.const_data_ptr<_Float16>(), + k_out.numel()); + dump_buffer( + phase + "_v_output", + l, + v_out.const_data_ptr<_Float16>(), + v_out.numel()); + } + } +} // namespace example diff --git a/examples/samsung/executor_runner/llama/runner/io_manager.h b/examples/samsung/executor_runner/llama/runner/io_manager.h index f86c9447cc3..0dd78e24f25 100644 --- a/examples/samsung/executor_runner/llama/runner/io_manager.h +++ b/examples/samsung/executor_runner/llama/runner/io_manager.h @@ -11,27 +11,67 @@ #include #include +#include namespace example { /** - * @brief Base class for managing input/output operations for LLM inference. + * IO manager for models where the runner manages KV-cache externally. * - * IOManager provides an interface for handling the input preparation and - * output processing for both prefill and decode phases of LLM inference. - * Derived classes must implement the virtual methods to provide specific IO - * management functionality. + * The model receives past KV-cache as input and produces new KV-cache entries + * as output. The runner is responsible for concatenating outputs into the + * cache and feeding the updated cache back on the next step. + * + * Supports micro-batch prefill: the runner may call prepare_prefill / + * update_prefill multiple times (each processing prefill_seq_len tokens) + * before switching to decode. The KV-cache accumulates in the prefill + * buffer across micro-batches and is copied to the decode buffer on the + * first prepare_decode call. + * + * Prefill (per micro-batch): + * Input: token[prefill_seq_len], start_pos, attn_mask[R,C], + * k_cache[1,H,past,D] x L, v_cache[1,H,past,D] x L + * Output: logits, k_cache[1,H,prefill_seq_len,D] x L, + * v_cache[1,H,prefill_seq_len,D] x L + * where past = max_seq_len - prefill_seq_len + * + * Decode: + * Input: token[1], start_pos, attn_mask[R,C], + * k_cache[1,H,max-1,D] x L, v_cache[1,H,max-1,D] x L + * Output: logits, k_cache[1,H,1,D] x L, v_cache[1,H,1,D] x L */ class ET_EXPERIMENTAL LlmIOManager : public executorch::extension::llm::IOManager { public: + struct Config { + size_t max_seq_len; + size_t prefill_seq_len; + size_t n_layers; + size_t n_kv_heads; + size_t head_dim; + size_t attn_mask_rows = 128; + size_t attn_mask_cols = 1024; + }; + /** * @brief Construct an IOManager bound to a Module. * * @param module The Module used for querying method metadata and execution. */ - explicit LlmIOManager(executorch::extension::Module& module) - : IOManager(module), module_(module) {} + LlmIOManager(executorch::extension::Module& module, Config config) + : IOManager(module), module_(module), config_(config) { + ET_LOG(Info, + "(%s) layers(%ld), kv_heads(%ld) max_sqe_len(%ld) \ + head_dim(%ld) prefill_seq_len(%ld)", + ET_FUNCTION, config_.n_layers, config_.n_kv_heads, + config_.max_seq_len, config_.head_dim, config_.prefill_seq_len); + (void)module_; + + prefill_cache_len_ = config_.max_seq_len - config_.prefill_seq_len; + decode_cache_len_ = config_.max_seq_len - 1; + + allocate_buffers(); + } /** * @brief Virtual destructor to allow proper cleanup in derived classes. @@ -48,6 +88,12 @@ class ET_EXPERIMENTAL LlmIOManager ET_NODISCARD virtual executorch::runtime::Error load( const std::string& prefill_method, const std::string& decode_method) override { + ET_LOG(Info, + "(%s) layers(%ld), kv_heads(%ld) max_sqe_len(%ld) \ + head_dim(%ld) prefill_seq_len(%ld)", + ET_FUNCTION, config_.n_layers, config_.n_kv_heads, + config_.max_seq_len, config_.head_dim, config_.prefill_seq_len); + return executorch::runtime::Error::Ok; } @@ -62,6 +108,24 @@ class ET_EXPERIMENTAL LlmIOManager const std::string& decode_method) override { (void)prefill_method; (void)decode_method; + + dump_step_ = 0; + + current_pos_ = 0; + prefill_done_ = false; + + for (size_t l = 0; l < config_.n_layers; l++) { + std::fill( + prefill_k_input_[l].begin(), prefill_k_input_[l].end(), 0.0f); + std::fill( + prefill_v_input_[l].begin(), prefill_v_input_[l].end(), 0.0f); + std::fill( + decode_k_input_[l].begin(), decode_k_input_[l].end(), 0.0f); + std::fill( + decode_v_input_[l].begin(), decode_v_input_[l].end(), 0.0f); + } + std::fill(attn_mask_.begin(), attn_mask_.end(), 0.0f); + return executorch::runtime::Error::Ok; } @@ -79,22 +143,7 @@ class ET_EXPERIMENTAL LlmIOManager prepare_prefill( const executorch::extension::TensorPtr& input, const executorch::extension::TensorPtr& start_pos, - const std::string& prefill_method) override { - auto method_meta = module_.method_meta(prefill_method); - if (!method_meta.ok()) { - return method_meta.error(); - } - if (method_meta->num_inputs() != 2) { - ET_LOG( - Error, - "Expected 2 inputs for prefill method, got %zu. Likely the model takes the caches or mask as an argument which this IOManager does not support.", - method_meta->num_inputs()); - return executorch::runtime::Error::InvalidState; - } - // Cpu IO Manager supports dynamic shapes for prefill, so no work to be done - // here. - return std::vector{input, start_pos}; - } + const std::string& prefill_method) override; /** * @brief Prepare inputs for the prefill phase using the default method name. @@ -109,7 +158,7 @@ class ET_EXPERIMENTAL LlmIOManager prepare_prefill( const executorch::extension::TensorPtr& input, const executorch::extension::TensorPtr& start_pos) { - return prepare_prefill(input, start_pos, "forward"); + return prepare_prefill(input, start_pos, example::default_method_name_prefill); } /** @@ -126,22 +175,7 @@ class ET_EXPERIMENTAL LlmIOManager prepare_decode( const executorch::extension::TensorPtr& input, const executorch::extension::TensorPtr& start_pos, - const std::string& decode_method) override { - auto method_meta = module_.method_meta(decode_method); - if (!method_meta.ok()) { - return method_meta.error(); - } - if (method_meta->num_inputs() != 2) { - ET_LOG( - Error, - "Expected 2 inputs for decode method, got %zu. Likely the model takes the caches or mask as an argument which this IOManager does not support.", - method_meta->num_inputs()); - return executorch::runtime::Error::InvalidState; - } - // Cpu IO Manager supports dynamic shapes for prefill, so no work to be done - // here. - return std::vector{input, start_pos}; - } + const std::string& decode_method) override; /** * @brief Prepare inputs for the decode phase using the default method name. @@ -156,7 +190,7 @@ class ET_EXPERIMENTAL LlmIOManager prepare_decode( const executorch::extension::TensorPtr& input, const executorch::extension::TensorPtr& start_pos) { - return prepare_decode(input, start_pos, "forward"); + return prepare_decode(input, start_pos, example::default_method_name_decode); } /** @@ -168,12 +202,7 @@ class ET_EXPERIMENTAL LlmIOManager */ ET_NODISCARD virtual executorch::runtime::Error update_prefill( const std::vector& model_outputs, - const std::string& prefill_method) override { - (void)model_outputs; - (void)prefill_method; - // No post inference work to do. - return executorch::runtime::Error::Ok; - } + const std::string& prefill_method) override; /** * @brief Process outputs from the prefill phase using the default method. @@ -185,7 +214,7 @@ class ET_EXPERIMENTAL LlmIOManager */ ET_NODISCARD executorch::runtime::Error update_prefill( const std::vector& model_outputs) { - return update_prefill(model_outputs, "forward"); + return update_prefill(model_outputs, example::default_method_name_prefill); } /** @@ -197,12 +226,7 @@ class ET_EXPERIMENTAL LlmIOManager */ ET_NODISCARD virtual executorch::runtime::Error update_decode( const std::vector& model_outputs, - const std::string& decode_method) override { - (void)model_outputs; - (void)decode_method; - // No post inference work to do. - return executorch::runtime::Error::Ok; - } + const std::string& decode_method) override; /** * @brief Process outputs from the decode phase using the default method. @@ -214,14 +238,89 @@ class ET_EXPERIMENTAL LlmIOManager */ ET_NODISCARD executorch::runtime::Error update_decode( const std::vector& model_outputs) { - return update_decode(model_outputs, "forward"); + return update_decode(model_outputs, example::default_method_name_decode); + } + + size_t current_pos() const { + return current_pos_; } private: + void dump_buffer( + const std::string& tag, + size_t layer, + const _Float16* data, + size_t size); + void dump_kv_input(const std::string& phase); + void dump_kv_output( + const std::string& phase, + const std::vector& model_outputs); + + void allocate_buffers(); + + bool dump_enabled_ = false; + mutable size_t dump_step_ = 0; + std::string dump_dir_ = "/data/local/tmp/kv_cache_dump"; + + /** + * Copy prefill output KV cache (layout [1,H,prefill_cache_len,D]) into + * the decode input buffer (layout [1,H,decode_cache_len,D]). + * + * Since the head stride differs between the two layouts, we copy per-head. + */ + void copy_prefill_to_decode_input(); + + /** + * Copy new cache entries from output (layout [1,H,new_len,D]) into the + * decode input buffer (layout [1,H,cache_len,D]) at the given position. + * + * For each head h: + * cache[h * cache_len * D + pos * D .. + (pos + new_len) * D] + * = src[h * new_len * D .. + new_len * D] + */ + void copy_to_cache(const _Float16* src, size_t new_len, + _Float16* cache, size_t cache_len, size_t pos); + + /** + * Update the attention mask for a causal (lower-triangular) pattern. + * For each query row r in [0, seq_len), the mask allows attending to + * positions [0, pos + r] (value 1) and masks future positions (value 0). + */ + void update_attn_mask(size_t pos, size_t seq_len); + /** * @brief Reference to the Module used for method metadata and execution. */ executorch::extension::module::Module& module_; + Config config_; + + size_t prefill_cache_len_; + size_t decode_cache_len_; + size_t current_pos_ = 0; + bool prefill_done_ = false; + + std::vector input_tokens_; + std::vector pos_ids_; + + // Per-layer input buffers: [1, n_kv_heads, cache_len, head_dim] flattened + std::vector> prefill_k_input_; + std::vector> prefill_v_input_; + std::vector> decode_k_input_; + std::vector> decode_v_input_; + + // Attention mask buffer: [attn_mask_rows, attn_mask_cols] flattened + std::vector attn_mask_; + executorch::extension::TensorPtr attn_mask_tensor_; + + // Tensor wrappers (kept alive between prepare and model execution) + std::vector prefill_k_input_tensors_; + std::vector prefill_v_input_tensors_; + std::vector prefill_k_output_tensors_; + std::vector prefill_v_output_tensors_; + std::vector decode_k_input_tensors_; + std::vector decode_v_input_tensors_; + std::vector decode_k_output_tensors_; + std::vector decode_v_output_tensors_; }; } // namespace example diff --git a/examples/samsung/executor_runner/llama/runner/prefiller.cpp b/examples/samsung/executor_runner/llama/runner/prefiller.cpp index 911852a2c1d..ca41cd274e2 100644 --- a/examples/samsung/executor_runner/llama/runner/prefiller.cpp +++ b/examples/samsung/executor_runner/llama/runner/prefiller.cpp @@ -43,7 +43,7 @@ ::executorch::runtime::Result Prefiller::prefill( // If prompt tokens exceed max_seq_len_, we need to chunk them ET_LOG( Info, - "Micro Batch: num_prompt(%d), max_seq_len(%ld), start_pos(%ld)", + "Micro Batch: num_prompt(%d), prefill::max_seq_len(%ld), start_pos(%ld)", num_prompt_tokens, max_seq_len_, start_pos); @@ -69,7 +69,9 @@ ::executorch::runtime::Result Prefiller::prefill( ET_CHECK_OK_OR_RETURN_ERROR(chunk_result.error()); cur_token = chunk_result.get(); - start_pos += num_tokens_to_prefill_with; + if (enable_parallel_prefill_) { + start_pos += num_tokens_to_prefill_with; + } num_tokens_to_process += num_tokens_to_prefill_with; } @@ -90,12 +92,10 @@ ::executorch::runtime::Result Prefiller::prefill_chunk( // store the token uint64_t cur_token; if (enable_parallel_prefill_ || !use_kv_cache_) { - // TODO (sicheon.oh) : embedding token - // initialize tensor wrappers auto tokens = executorch::extension::from_blob( prompt_tokens.data(), - {1, num_prompt_tokens}, + {1, static_cast(prompt_tokens.size())}, executorch::aten::ScalarType::Long); auto outputs_res = text_decoder_runner_->step(tokens, start_pos); @@ -119,6 +119,7 @@ ::executorch::runtime::Result Prefiller::prefill_chunk( start_pos); // initialize tensor wrappers + // TODO(sicheon.oh) : Set tensor size by metadata::get_prefill_ar_len auto tokens = executorch::extension::from_blob( &cur_token, {1, 1}, executorch::aten::ScalarType::Long); diff --git a/examples/samsung/executor_runner/llama/runner/runner.cpp b/examples/samsung/executor_runner/llama/runner/runner.cpp index c09815d462a..4a3b64a86ca 100644 --- a/examples/samsung/executor_runner/llama/runner/runner.cpp +++ b/examples/samsung/executor_runner/llama/runner/runner.cpp @@ -23,6 +23,8 @@ #include #include #include +#include + namespace example { @@ -63,6 +65,7 @@ ET_EXPERIMENTAL std::unique_ptr load_tokenizer( * @param model_path Path to the model file * @param tokenizer Initialized tokenizer instance * @param data_path Optional path to additional data required by the model + * @param decode_only Optional flag to run with decode method only * @param temperature Optional temperature parameter for controlling randomness * (deprecated) * @return std::unique_ptr Initialized LLMRunner instance @@ -71,6 +74,7 @@ ET_EXPERIMENTAL std::unique_ptr create_llm_runner( const std::string& model_path, std::unique_ptr<::tokenizers::Tokenizer> tokenizer, std::optional data_path = std::nullopt, + bool decode_only = false, float temperature = -1.0f); std::unique_ptr<::tokenizers::Tokenizer> load_llama_tokenizer( @@ -84,6 +88,7 @@ std::unique_ptr create_llama_runner( const std::string& model_path, const std::string& tokenizer_path, std::optional data_path, + bool decode_only, float temperature) { ET_LOG( Info, @@ -106,7 +111,7 @@ std::unique_ptr create_llama_runner( "[Create_test_llm_runner] %s %s", model_path.c_str(), data_path->c_str()); - return create_llm_runner(model_path, std::move(tokenizer), data_path); + return create_llm_runner(model_path, std::move(tokenizer), data_path, decode_only); } static const std::unordered_map kMetadataMap{ @@ -117,14 +122,20 @@ static const std::unordered_map kMetadataMap{ {llm::kMaxContextLen, "get_max_context_len"}, {llm::kVocabSize, "get_vocab_size"}, {llm::kUseKVCache, "use_kv_cache"}, - {llm::kUseSDPAWithKVCache, "use_sdpa_with_kv_cache"}}; + {llm::kUseSDPAWithKVCache, "use_sdpa_with_kv_cache"}, + {example::kPrefillArLen, "get_prefill_ar_len"}, + {example::kNumOfHiddenDim, "get_hidden_dim"}, + {example::kNumOfHeads, "get_n_heads"}, + {example::kNumOfKVHeads, "get_n_kv_heads"}, + {example::kNumOfLayers, "get_n_layers"} +}; LLMRunner::LLMRunner( std::unordered_map metadata, std::unique_ptr<::tokenizers::Tokenizer> tokenizer, std::unique_ptr<::executorch::extension::Module> module, std::unique_ptr text_decoder_runner, - std::unique_ptr text_prefiller, + std::unique_ptr text_prefiller, std::unique_ptr text_token_generator, std::unique_ptr stats, std::unique_ptr io_manager, @@ -367,12 +378,7 @@ std::unique_ptr load_tokenizer( size_t bos_token_index, size_t eos_token_index) { ::executorch::runtime::runtime_init(); - auto json_tokenizer = std::make_unique(); - if (json_tokenizer->load(tokenizer_path) == ::tokenizers::Error::Ok) { - ET_LOG(Info, "Loaded json tokenizer"); - return json_tokenizer; - } - std::unique_ptr<::tokenizers::Tiktoken> tiktoken_tokenizer; + std::unique_ptr<::tokenizers::Tiktoken> tiktoken_tokenizer; if (special_tokens != nullptr && !pattern.has_value()) { tiktoken_tokenizer = std::make_unique<::tokenizers::Tiktoken>( std::move(special_tokens), bos_token_index, eos_token_index); @@ -385,11 +391,18 @@ std::unique_ptr load_tokenizer( } else { tiktoken_tokenizer = std::make_unique<::tokenizers::Tiktoken>(); } + if (tiktoken_tokenizer->load(tokenizer_path) == ::tokenizers::Error::Ok) { ET_LOG(Info, "Loaded TikToken tokenizer"); return tiktoken_tokenizer; } + auto json_tokenizer = std::make_unique(); + if (json_tokenizer->load(tokenizer_path) == ::tokenizers::Error::Ok) { + ET_LOG(Info, "Loaded json tokenizer"); + return json_tokenizer; + } + auto sp_tokenizer = std::make_unique<::tokenizers::SPTokenizer>(); if (sp_tokenizer->load(tokenizer_path) == ::tokenizers::Error::Ok) { ET_LOG(Info, "Loaded Sentencepiece tokenizer"); @@ -415,6 +428,12 @@ std::unordered_map get_llm_metadata( {kMetadataMap.at(llm::kMaxContextLen), 128}, {kMetadataMap.at(llm::kUseKVCache), true}, {kMetadataMap.at(llm::kUseSDPAWithKVCache), false}, + {kMetadataMap.at(example::kPrefillArLen), 0}, + //Params.json for llama3.2-1B + {kMetadataMap.at(example::kNumOfHiddenDim), 2048}, + {kMetadataMap.at(example::kNumOfHeads), 32}, + {kMetadataMap.at(example::kNumOfKVHeads), 8}, + {kMetadataMap.at(example::kNumOfLayers), 16}, }); // Read metadata from the model @@ -479,6 +498,7 @@ std::unique_ptr create_llm_runner( const std::string& model_path, std::unique_ptr<::tokenizers::Tokenizer> tokenizer, std::optional data_path, + bool decode_only, float temperature) { // Sanity check tokenizer if (!tokenizer || !tokenizer->is_loaded()) { @@ -504,22 +524,42 @@ std::unique_ptr create_llm_runner( auto methods = methods_res.get(); std::vector register_methods; - for (auto method : methods) { - if (kMetadataMap.find(method) != kMetadataMap.end()) { - ET_LOG(Info, "It is not method(%s)", method.c_str()); - break; + if (decode_only) { + register_methods.emplace_back(default_method_name_decode); + } else { + for (auto method : methods) { + if (kMetadataMap.find(method) != kMetadataMap.end()) { + ET_LOG(Info, "It is not method(%s)", method.c_str()); + continue; + } + ET_LOG(Info, "Register method \"%s\" to decoder", method.c_str()); + register_methods.emplace_back(method); } - ET_LOG(Info, "Register method \"%s\" to decoder", method.c_str()); - register_methods.emplace_back(method); - } - - // Create IO Manager - auto io_manager = std::make_unique(*module.get()); + } // Get metadata from Module ET_LOG(Info, "Reading metadata from model"); auto metadata = get_llm_metadata(tokenizer.get(), module.get()); + // Create IO Manager + example::LlmIOManager::Config conf; + conf.n_layers = metadata.at(example::kNumOfLayers); + conf.n_kv_heads = metadata.at(example::kNumOfKVHeads); + conf.max_seq_len = metadata.at(llm::kMaxSeqLen); + conf.prefill_seq_len = (metadata.at(example::kPrefillArLen) > 1) ? + (metadata.at(example::kPrefillArLen)) : + (metadata.at(llm::kMaxSeqLen)); + + if (decode_only) { + conf.prefill_seq_len = 1; + } + + size_t dim = metadata.at(example::kNumOfHiddenDim); + size_t n_heads = metadata.at(example::kNumOfHeads); + conf.head_dim = dim / n_heads; + + auto io_manager = std::make_unique(*module.get(), conf); + auto eos_ids = std::make_unique>( get_eos_ids(tokenizer.get(), module.get())); @@ -533,12 +573,16 @@ std::unique_ptr create_llm_runner( text_decoder_runner->register_method(method); } + // TODO(sicheon.oh) : Need to find proper flag to distinguish separaged prefill/decode opion. + bool separated_prefill_decode = (conf.prefill_seq_len > 1); + ET_LOG(Info, "separated_prefill_decode(%s)", (separated_prefill_decode)?"True":"False"); + // Create text_prefiller auto text_prefiller = std::make_unique( text_decoder_runner.get(), metadata.at(llm::kUseKVCache), - metadata.at(llm::kEnableDynamicShape), - metadata.at(llm::kMaxSeqLen)); + metadata.at(llm::kEnableDynamicShape) | (separated_prefill_decode), + conf.prefill_seq_len); // Create text_token_generator with stats auto stats = std::make_unique(); diff --git a/examples/samsung/executor_runner/llama/runner/runner.h b/examples/samsung/executor_runner/llama/runner/runner.h index bd7443f9f61..a1a0d348a6d 100644 --- a/examples/samsung/executor_runner/llama/runner/runner.h +++ b/examples/samsung/executor_runner/llama/runner/runner.h @@ -52,7 +52,7 @@ class LLMRunner : public llm::IRunner { std::unique_ptr<::tokenizers::Tokenizer> tokenizer, std::unique_ptr<::executorch::extension::Module> module, std::unique_ptr text_decoder_runner, - std::unique_ptr text_prefiller, + std::unique_ptr text_prefiller, std::unique_ptr text_token_generator, std::unique_ptr stats, std::unique_ptr io_manager, @@ -161,7 +161,7 @@ class LLMRunner : public llm::IRunner { text_decoder_runner_; // Manage text_decoder_runner_'s lifecycle, make // sure it outlives text_prefiller_ & // text_token_generator_. - std::unique_ptr text_prefiller_; + std::unique_ptr text_prefiller_; std::unique_ptr text_token_generator_; std::unique_ptr io_manager_; @@ -177,6 +177,7 @@ std::unique_ptr create_llama_runner( const std::string& model_path, const std::string& tokenizer_path, std::optional data_path = std::nullopt, + bool decode_only = false, float temperature = -1.0f); std::unique_ptr load_llama_tokenizer( diff --git a/extension/llm/runner/text_decoder_runner.cpp b/extension/llm/runner/text_decoder_runner.cpp index bb64d79dd2e..6249719a7dd 100644 --- a/extension/llm/runner/text_decoder_runner.cpp +++ b/extension/llm/runner/text_decoder_runner.cpp @@ -46,6 +46,7 @@ ::executorch::runtime::Result TextDecoderRunner::step( return method_meta_result.error(); } auto method_meta = std::move(*method_meta_result); + // If only 1 input, we are not using kv cache bool use_kv_cache = method_meta.num_inputs() > 1; From dfd0dad84d0bbd519dcd5f07c44c604eac0a20df Mon Sep 17 00:00:00 2001 From: Jiseong Oh Date: Tue, 13 Oct 2026 19:37:22 +0900 Subject: [PATCH 06/17] Fix mask buffer generation in the LLM runner Rework how the IO manager fills the attention mask buffer. Signed-off-by: Jiseong Oh Co-authored-by: MJ Kim --- .../llama/runner/io_manager.cpp | 44 ++++++++++++------- 1 file changed, 29 insertions(+), 15 deletions(-) diff --git a/examples/samsung/executor_runner/llama/runner/io_manager.cpp b/examples/samsung/executor_runner/llama/runner/io_manager.cpp index d6f30435657..c919d8d8af9 100644 --- a/examples/samsung/executor_runner/llama/runner/io_manager.cpp +++ b/examples/samsung/executor_runner/llama/runner/io_manager.cpp @@ -288,23 +288,37 @@ namespace example { } void LlmIOManager::update_attn_mask(size_t pos, size_t seq_len) { - std::fill(attn_mask_.begin(), attn_mask_.end(), 0.0f); - size_t visible = 0; - for (size_t r = 0; r < seq_len && r < config_.attn_mask_rows; r++) { - visible = pos + r + 1; - if (visible > config_.attn_mask_cols) { - visible = config_.attn_mask_cols; - } - for (size_t c = 0; c < visible; c++) { - attn_mask_[r * config_.attn_mask_cols + c] = 1.0f; - } + static int num_prompt_token = 0; // temporary variable to share num_prompt_token info from prefill to decode + std::fill(attn_mask_.begin(), attn_mask_.end(), -std::numeric_limits::infinity()); + + const size_t start_col = config_.attn_mask_cols - config_.attn_mask_rows; + if (seq_len > 1) { // prefill : start from 896 and fill same mask as last_token to all empty rows. + size_t last_valid_row = 0; + num_prompt_token = seq_len; + + for (size_t r = 0; r < seq_len && r < config_.attn_mask_rows; r++) { + size_t visible_count = pos + r + 1; + for (size_t c = 0; c < visible_count; c++) { + size_t target_col = start_col + c; + if (target_col < config_.attn_mask_cols) { + attn_mask_[r * config_.attn_mask_cols + target_col] = 0.0f; + } + } + last_valid_row = r; + } + for (size_t r = last_valid_row + 1; r < config_.attn_mask_rows; r++) { + float* src = &attn_mask_[last_valid_row * config_.attn_mask_cols]; + float* dst = &attn_mask_[r * config_.attn_mask_cols]; + std::memcpy(dst, src, config_.attn_mask_cols * sizeof(float)); + } } + else { + int pos_decode = pos - num_prompt_token + 1; + for (int r = 0; r < (pos); r++) { + attn_mask_[r] = 0.0f; + } - //Mask padding with the value of last token. - for (size_t r = seq_len; r < config_.attn_mask_rows; r++) { - for (size_t c = 0; c < visible; c++) { - attn_mask_[r * config_.attn_mask_cols + c] = 1.0f; - } + attn_mask_[config_.attn_mask_cols - 1] = 0.0f; } } From 7703bcbde205cb32e831041ac2a24658af755421 Mon Sep 17 00:00:00 2001 From: Jiseong Oh Date: Tue, 13 Oct 2026 19:37:22 +0900 Subject: [PATCH 07/17] Fix a build issue in the LLM runner decoder decoder.h derives from TextDecoderRunner, so include its header directly instead of relying on a transitive include. Signed-off-by: Jiseong Oh --- examples/samsung/executor_runner/llama/runner/decoder.h | 1 + 1 file changed, 1 insertion(+) diff --git a/examples/samsung/executor_runner/llama/runner/decoder.h b/examples/samsung/executor_runner/llama/runner/decoder.h index 59a1debdc6c..8d9ef611a0f 100644 --- a/examples/samsung/executor_runner/llama/runner/decoder.h +++ b/examples/samsung/executor_runner/llama/runner/decoder.h @@ -15,6 +15,7 @@ #include #include #include +#include namespace example { From 32e7d7f6d3d241a7731637b6dbb3473804a22708 Mon Sep 17 00:00:00 2001 From: Jiseong Oh Date: Tue, 13 Oct 2026 19:37:42 +0900 Subject: [PATCH 08/17] Enable Gemma model support in the LLM runner Fix lint errors across the runner, then add Gemma support: a positional embedding class, attention mask and position id handling for sequences longer than 128 tokens, outputs aligned with the PyTorch reference, and sample run scripts for Gemma 3 and Llama 3. Signed-off-by: Jiseong Oh Co-authored-by: MJ Kim --- backends/samsung/build.sh | 26 +- .../enn_executor_llama_runner.cpp | 36 +- .../llama/runner/CMakeLists.txt | 1 + .../executor_runner/llama/runner/constants.h | 8 +- .../executor_runner/llama/runner/decoder.cpp | 125 ++-- .../executor_runner/llama/runner/decoder.h | 17 +- .../llama/runner/io_manager.cpp | 537 +++++++++++------- .../executor_runner/llama/runner/io_manager.h | 191 ++++++- .../llama/runner/prefiller.cpp | 6 +- .../llama/runner/rotary_embedding.cpp | 80 +++ .../llama/runner/rotary_embedding.h | 41 ++ .../executor_runner/llama/runner/runner.cpp | 231 ++++---- .../executor_runner/llama/runner/runner.h | 30 +- .../executor_runner/run_gemma3_sample.sh | 29 + .../executor_runner/run_llama3_sample.sh | 27 + extension/llm/runner/constants.h | 1 + extension/llm/runner/io_manager/io_manager.h | 13 + extension/llm/runner/text_decoder_runner.cpp | 23 +- 18 files changed, 979 insertions(+), 443 deletions(-) create mode 100644 examples/samsung/executor_runner/llama/runner/rotary_embedding.cpp create mode 100644 examples/samsung/executor_runner/llama/runner/rotary_embedding.h create mode 100644 examples/samsung/executor_runner/run_gemma3_sample.sh create mode 100644 examples/samsung/executor_runner/run_llama3_sample.sh diff --git a/backends/samsung/build.sh b/backends/samsung/build.sh index da8e211ee64..32b955ab765 100755 --- a/backends/samsung/build.sh +++ b/backends/samsung/build.sh @@ -89,36 +89,30 @@ function build_android() { -DCMAKE_BUILD_TYPE=Release \ -DEXECUTORCH_BUILD_ENN=ON \ -DEXYNOS_AI_LITECORE_ROOT=${EXYNOS_AI_LITECORE_ROOT} \ + -DCMAKE_FIND_ROOT_PATH_MODE_PACKAGE=BOTH \ -DEXECUTORCH_BUILD_EXTENSION_MODULE=ON \ -DEXECUTORCH_BUILD_EXTENSION_DATA_LOADER=ON \ -DEXECUTORCH_BUILD_EXTENSION_FLAT_TENSOR=ON \ - -DEXECUTORCH_BUILD_EXTENSION_LLM=ON \ - -DEXECUTORCH_BUILD_EXTENSION_LLM_RUNNER=ON \ -DEXECUTORCH_BUILD_EXTENSION_NAMED_DATA_MAP=ON \ -DEXECUTORCH_BUILD_EXTENSION_TENSOR=ON \ + -DEXECUTORCH_BUILD_EXTENSION_LLM=ON \ + -DEXECUTORCH_BUILD_EXTENSION_LLM_RUNNER=ON \ -DEXECUTORCH_BUILD_EXTENSION_FLAT_TENSOR=ON \ -DEXECUTORCH_BUILD_EXTENSION_RUNNER_UTIL=ON \ + -DEXECUTORCH_BUILD_KERNELS_LLM=ON \ + -DEXECUTORCH_BUILD_KERNELS_OPTIMIZED=ON \ + -DEXECUTORCH_BUILD_KERNELS_OPTIMIZED_PORTABLE=ON \ -DEXECUTORCH_BUILD_KERNELS_QUANTIZED=ON \ + -DEXECUTORCH_BUILD_XNNPACK=ON \ -DEXECUTORCH_ENABLE_LOGGING=1 \ -DEXECUTORCH_BUILD_DEVTOOLS=ON \ -DEXECUTORCH_ENABLE_EVENT_TRACER=ON \ + -DBUILD_TESTING=OFF \ + -DPYTHON_EXECUTABLE=python \ -S ${PROJECT_DIR} \ -B ${ANDROID_BUILD_DIR} - cmake --build ${ANDROID_BUILD_DIR} -j $(nproc) --target install - - LLAMA_EXAMPLE_ROOT=examples/models/llama - cmake ${PROJECT_DIR}/${LLAMA_EXAMPLE_ROOT} \ - -DCMAKE_TOOLCHAIN_FILE=$ANDROID_NDK_ROOT/build/cmake/android.toolchain.cmake \ - -DANDROID_ABI=${ANDROID_ABI} \ - -DCMAKE_INSTALL_PREFIX=${ANDROID_BUILD_DIR} \ - -DCMAKE_BUILD_TYPE=Release \ - -DANDROID_PLATFORM=${ANDROID_PLATFORM} \ - -DEXECUTORCH_ENABLE_LOGGING=ON \ - -DPYTHON_EXECUTABLE=${PYTHON_EXECUTABLE} \ - -B${ANDROID_BUILD_DIR}/${LLAMA_EXAMPLE_ROOT} - - cmake --build ${ANDROID_BUILD_DIR}/${LLAMA_EXAMPLE_ROOT} -j $(nproc) --config Release + cmake --build ${ANDROID_BUILD_DIR} -j $(nproc) --target install --config Release } # Main diff --git a/examples/samsung/executor_runner/enn_executor_llama_runner.cpp b/examples/samsung/executor_runner/enn_executor_llama_runner.cpp index a14f5f033e6..3c496412cd6 100644 --- a/examples/samsung/executor_runner/enn_executor_llama_runner.cpp +++ b/examples/samsung/executor_runner/enn_executor_llama_runner.cpp @@ -24,8 +24,15 @@ DEFINE_string( DEFINE_string(data_path, "", "Data file for the model."); +DEFINE_string( + tokenizer_type, + "tiktoken", + "Tokenizer type. One of ['bpe', 'tiktoken']."); + DEFINE_string(tokenizer_path, "tokenizer.bin", "Tokenizer stuff."); +DEFINE_string(pos_embedding_path, "", "rotary embedding stuff."); + DEFINE_string(prompt, "The answer to the ultimate question is", "Prompt."); DEFINE_double( @@ -57,6 +64,8 @@ DEFINE_bool(warmup, false, "Whether to run a warmup run."); DEFINE_bool(decode_only, false, "Whether to run with decode method only."); +DEFINE_bool(use_gemma, false, "Use Gemma model specific processing."); + int32_t main(int32_t argc, char** argv) { gflags::ParseCommandLineFlags(&argc, &argv, true); @@ -70,8 +79,14 @@ int32_t main(int32_t argc, char** argv) { data_path = FLAGS_data_path.c_str(); } + const char* tokenizer_type = FLAGS_tokenizer_type.c_str(); + const char* tokenizer_path = FLAGS_tokenizer_path.c_str(); + std::optional pos_embedding_path = + FLAGS_pos_embedding_path.empty() ? std::nullopt + : std::make_optional(FLAGS_pos_embedding_path.c_str()); + ET_LOG(Info, "prompt : %s", FLAGS_prompt.c_str()); const char* prompt = FLAGS_prompt.c_str(); @@ -85,6 +100,12 @@ int32_t main(int32_t argc, char** argv) { bool decode_only = FLAGS_decode_only; + bool use_gemma = FLAGS_use_gemma; + + int num_bos = FLAGS_num_bos; + + int num_eos = FLAGS_num_eos; + #if defined(ET_USE_THREADPOOL) uint32_t num_performant_cores = cpu_threads == -1 ? ::executorch::extension::cpuinfo::get_num_performant_cores() @@ -98,8 +119,16 @@ int32_t main(int32_t argc, char** argv) { #endif // create llama runner - std::unique_ptr runner = - example::create_llama_runner(model_path, tokenizer_path, data_path, decode_only); + example::RunnerPaths runner_paths{ + .model_path = model_path, + .tokenizer_path = tokenizer_path, + .pos_embedding_path = pos_embedding_path, + .data_path = data_path}; + example::RunnerConfig runner_config{ + .tokenizer_type = tokenizer_type, + .decode_only = decode_only, + .use_gemma = use_gemma}; + std::unique_ptr runner = example::create_llama_runner(runner_paths, runner_config); if (runner == nullptr) { ET_LOG(Error, "Failed to create llama runner"); @@ -121,7 +150,8 @@ int32_t main(int32_t argc, char** argv) { // generate executorch::extension::llm::GenerationConfig config{ - .seq_len = seq_len, .temperature = temperature}; + .seq_len = seq_len, .temperature = temperature, + .num_bos = num_bos, .num_eos = num_eos}; auto error = runner->generate(prompt, config, {}, {}); if (error != executorch::runtime::Error::Ok) { ET_LOG(Error, "Failed to generate token from llama runner"); diff --git a/examples/samsung/executor_runner/llama/runner/CMakeLists.txt b/examples/samsung/executor_runner/llama/runner/CMakeLists.txt index 016f3289142..2bf89b654f8 100644 --- a/examples/samsung/executor_runner/llama/runner/CMakeLists.txt +++ b/examples/samsung/executor_runner/llama/runner/CMakeLists.txt @@ -24,6 +24,7 @@ list( ${CMAKE_CURRENT_LIST_DIR}/decoder.cpp ${CMAKE_CURRENT_LIST_DIR}/prefiller.cpp ${CMAKE_CURRENT_LIST_DIR}/io_manager.cpp + ${CMAKE_CURRENT_LIST_DIR}/rotary_embedding.cpp ${CMAKE_CURRENT_SOURCE_DIR}/../../../../../examples/models/llama/tokenizer/llama_tiktoken.cpp ) diff --git a/examples/samsung/executor_runner/llama/runner/constants.h b/examples/samsung/executor_runner/llama/runner/constants.h index af0d3a52d2f..1ddb1c2ebb0 100644 --- a/examples/samsung/executor_runner/llama/runner/constants.h +++ b/examples/samsung/executor_runner/llama/runner/constants.h @@ -17,7 +17,7 @@ inline constexpr auto kNumOfHiddenDim = "get_hidden_dim"; inline constexpr auto kNumOfHeads = "get_n_heads"; inline constexpr auto kNumOfKVHeads = "get_n_kv_heads"; -static constexpr auto default_method_name = "forward"; -static constexpr auto default_method_name_prefill = "prefill_forward"; -static constexpr auto default_method_name_decode = "decode_forward"; -} // namespace executorch::extension::llm +inline constexpr auto default_method_name = "forward"; +inline constexpr auto default_method_name_prefill = "prefill_forward"; +inline constexpr auto default_method_name_decode = "decode_forward"; +} // namespace example diff --git a/examples/samsung/executor_runner/llama/runner/decoder.cpp b/examples/samsung/executor_runner/llama/runner/decoder.cpp index 0293c7c5abb..e25c894b47a 100644 --- a/examples/samsung/executor_runner/llama/runner/decoder.cpp +++ b/examples/samsung/executor_runner/llama/runner/decoder.cpp @@ -8,13 +8,13 @@ // Given inputs, run a text decoder and return logits. +#include #include #include #include #include - #include namespace example { @@ -25,15 +25,62 @@ class IExecutionIOBinder { virtual void afterExecution() = 0; }; +/** + * @brief Slices the k-th token's logits by calculating the memory offset. + * Corrected to return Result by dereferencing the TensorPtr. + */ +executorch::runtime::Result slice_kth_token( + const executorch::runtime::EValue& output_evalue, + int64_t k) { + ET_CHECK_MSG(output_evalue.isTensor(), "The provided EValue is not a tensor."); + executorch::aten::Tensor logits = output_evalue.toTensor(); + + constexpr int64_t kExpectedNumDim = 3; + constexpr int64_t kSequenceDim = 1; + const int64_t max_seq_len = logits.size(kSequenceDim); + const int64_t actual_dim = logits.dim(); + + const auto strides = logits.strides(); + const auto sizes = logits.sizes(); + + ET_CHECK_MSG(k >= 0 && k < max_seq_len, + "Index k (%ld) out of bounds (max: %ld)", k, max_seq_len); + + ET_CHECK_MSG(actual_dim == kExpectedNumDim, + "num dim shape of logit(%ld) is not expected(%ld)", actual_dim, kExpectedNumDim); + + // Calculate Memory Offset (zero-copy) + int64_t offset = k * strides[kSequenceDim]; + + // Define New Shape for the k-th token (1, 1, VocabSize) + // Using a fixed-size array or vector that matches SizesType + std::vector new_sizes = { + static_cast(sizes[0]), + 1, // Target single token + static_cast(sizes[2]) + }; + + // Create TensorPtr and dereference it to return Result + auto tensor_ptr = executorch::extension::from_blob( + logits.mutable_data_ptr() + offset, + new_sizes, + logits.scalar_type() + ); + + // Dereference the shared_ptr to get the Tensor object + return *tensor_ptr; +} + ::executorch::runtime::Result Decoder::step( executorch::extension::TensorPtr& tokens, int64_t start_pos) { - std::string method_stage; if (tokens->numel() > 1) { - method_stage = methods_.name[static_cast(DecoderMode::PREFILL)]; + method_stage = + methods_.name[static_cast(DecoderMode::PREFILL)]; } else { - method_stage = methods_.name[static_cast(DecoderMode::DECODE)]; + method_stage = + methods_.name[static_cast(DecoderMode::DECODE)]; } // ET_LOG(Info, "Input token %" PRIu64, input_token); @@ -45,30 +92,37 @@ ::executorch::runtime::Result Decoder::step( // Size of the second argument. This could be either input_pos or // cache_positions + const bool use_gemma = io_manager_->use_gemma(); + size_t pos_shape_size = 3; + size_t pos_idx = 3; + size_t attention_idx = 1; + // Check if we are using cache positions instead of input pos. - auto second_input_info = ET_UNWRAP(method_meta.input_tensor_meta(1)); - // For input_pos, numel is 1, for cache_positions, numel is max_seq_len - auto sizes = second_input_info.sizes(); - // Assuming 1D tensor + auto pos_input_info = ET_UNWRAP(method_meta.input_tensor_meta(pos_idx)); + auto sizes = pos_input_info.sizes(); ET_CHECK_OR_RETURN_ERROR( - sizes.size() == 1, + sizes.size() == pos_shape_size, InvalidProgram, - "The second input tensor is not 1D tensor. Got dimension (%zu)", + "The pos input tensor is not expected. Got dimension (%zu)", sizes.size()); auto numel = sizes[0]; std::vector<::executorch::aten::SizesType> sizes_vec = {numel}; auto n_inputs = method_meta.num_inputs(); - auto input_info = ET_UNWRAP(method_meta.input_tensor_meta(2)); + auto input_info = ET_UNWRAP(method_meta.input_tensor_meta(attention_idx)); sizes = input_info.sizes(); - if (sizes.size() != 2) ET_LOG(Error, "third input of %zu is invalid dimension : %zu", n_inputs, sizes.size()); - - input_info = ET_UNWRAP(method_meta.input_tensor_meta(3)); + // llama mask tensor shape : 2dim + // gemma mask tensor index : 4dim + // llama kvcache tensor index : 3~ + // gemma mask tensor index : 4~ + size_t kv_cache_idx; + if (sizes.size() != 4) ET_LOG(Error, "input %zu of %zu is invalid dimension : %zu", attention_idx, n_inputs, sizes.size()); + kv_cache_idx = attention_idx + (use_gemma ? 6 : 3); + input_info = ET_UNWRAP(method_meta.input_tensor_meta(kv_cache_idx)); sizes = input_info.sizes(); - if (sizes.size() != 4) ET_LOG(Error, "forth input of %zu is invalid dimension : %zu", n_inputs, sizes.size()); + if (sizes.size() != 4) ET_LOG(Error, "input %zu of %zu is invalid dimension : %zu", kv_cache_idx, n_inputs, sizes.size()); executorch::extension::TensorPtr start_pos_tensor; - // Assuming model is exported with input_pos, create a tensor with size 1 start_pos_tensor = executorch::extension::from_blob( &start_pos, sizes_vec, ::executorch::aten::ScalarType::Long); @@ -76,13 +130,11 @@ ::executorch::runtime::Result Decoder::step( std::vector inputs; if (method_stage == methods_.name[static_cast(DecoderMode::PREFILL)]) { - auto inputs_res = - io_manager_->prepare_prefill(tokens, start_pos_tensor); + auto inputs_res = io_manager_->prepare_prefill(tokens, start_pos_tensor); ET_CHECK_OK_OR_RETURN_ERROR(inputs_res.error()); inputs = inputs_res.get(); } else if (method_stage == methods_.name[static_cast(DecoderMode::DECODE)]) { - auto inputs_res = - io_manager_->prepare_decode(tokens, start_pos_tensor); + auto inputs_res = io_manager_->prepare_decode(tokens, start_pos_tensor); ET_CHECK_OK_OR_RETURN_ERROR(inputs_res.error()); inputs = inputs_res.get(); } @@ -106,7 +158,7 @@ ::executorch::runtime::Result Decoder::step( } // Return the logits tensor - return outputs_res.get()[0].toTensor(); + return slice_kth_token(outputs_res.get()[0], (tokens->numel() - 1)); } else { // no kv cache (void)start_pos; // unused @@ -134,9 +186,9 @@ ::executorch::runtime::Error Decoder::create_shared_planned_memory() { } auto prefill_meta = module_->method_meta( - methods_.name[static_cast(DecoderMode::PREFILL)].c_str()); + methods_.name[static_cast(DecoderMode::PREFILL)].c_str()); auto decode_meta = module_->method_meta( - methods_.name[static_cast(DecoderMode::DECODE)].c_str()); + methods_.name[static_cast(DecoderMode::DECODE)].c_str()); if (!prefill_meta.ok() || !decode_meta.ok()) { ET_LOG(Error, "Failed to get method metadata for shared memory setup"); return executorch::runtime::Error::InvalidArgument; @@ -146,26 +198,31 @@ ::executorch::runtime::Error Decoder::create_shared_planned_memory() { size_t decode_bufs = decode_meta->num_memory_planned_buffers(); size_t num_buffers = std::max(prefill_bufs, decode_bufs); - ET_LOG(Info, "Prefill_bufs(%ld), decode_bufs(%ld) num_buffers(%ld)", - prefill_bufs, decode_bufs, num_buffers); + ET_LOG( + Info, + "Prefill_bufs(%ld), decode_bufs(%ld) num_buffers(%ld)", + prefill_bufs, + decode_bufs, + num_buffers); - shared_method_memory_ = - std::make_unique(); + shared_method_memory_ = std::make_unique(); shared_method_memory_->buffers.reserve(num_buffers); shared_method_memory_->spans.reserve(num_buffers); for (size_t i = 0; i < num_buffers; ++i) { - int64_t prefill_size = - (i < prefill_bufs) + int64_t prefill_size = (i < prefill_bufs) ? prefill_meta->memory_planned_buffer_size(i).get() : 0; - int64_t decode_size = - (i < decode_bufs) + int64_t decode_size = (i < decode_bufs) ? decode_meta->memory_planned_buffer_size(i).get() : 0; size_t max_size = static_cast(std::max(prefill_size, decode_size)); - ET_LOG(Info, "Prefill_size(%ld), decode_size(%ld) max_size(%ld)", - prefill_size, decode_size, max_size); + ET_LOG( + Info, + "Prefill_size(%ld), decode_size(%ld) max_size(%ld)", + prefill_size, + decode_size, + max_size); shared_method_memory_->buffers.emplace_back(max_size); shared_method_memory_->spans.emplace_back( @@ -180,6 +237,4 @@ ::executorch::runtime::Error Decoder::create_shared_planned_memory() { return ret; } - - } // namespace example diff --git a/examples/samsung/executor_runner/llama/runner/decoder.h b/examples/samsung/executor_runner/llama/runner/decoder.h index 8d9ef611a0f..353863edc94 100644 --- a/examples/samsung/executor_runner/llama/runner/decoder.h +++ b/examples/samsung/executor_runner/llama/runner/decoder.h @@ -10,12 +10,12 @@ #pragma once +#include +#include #include #include -#include #include -#include -#include +#include namespace example { @@ -33,8 +33,7 @@ class ET_EXPERIMENTAL Decoder explicit Decoder( executorch::extension::Module* module, llm::IOManager* io_manager) - : TextDecoderRunner(module, io_manager), - shared_method_memory_(nullptr) {} + : TextDecoderRunner(module, io_manager), shared_method_memory_(nullptr) {} virtual ::executorch::runtime::Result step( executorch::extension::TensorPtr& input, @@ -43,7 +42,6 @@ class ET_EXPERIMENTAL Decoder inline ::executorch::runtime::Error load() override { executorch::runtime::Error ret = executorch::runtime::Error::Ok; - if (methods_.num > 1) { ret = create_shared_planned_memory(); if (ret != executorch::runtime::Error::Ok) { @@ -61,7 +59,8 @@ class ET_EXPERIMENTAL Decoder } if (shared_method_memory_) { - ret = module_->load_method(method, shared_method_memory_->allocator.get()); + ret = module_->load_method( + method, shared_method_memory_->allocator.get()); if (ret != executorch::runtime::Error::Ok) { ET_LOG(Error, "Failed to load method"); return ret; @@ -98,9 +97,9 @@ class ET_EXPERIMENTAL Decoder ET_LOG(Info, "method[%d] : %s", methods_.num - 1, method.c_str()); - int i = 0; for (int i = 0; i < methods_.num; i++) { - ET_LOG(Info, "Registered method[%d] : \"%s\"", i, methods_.name[i].c_str()); + ET_LOG( + Info, "Registered method[%d] : \"%s\"", i, methods_.name[i].c_str()); } } diff --git a/examples/samsung/executor_runner/llama/runner/io_manager.cpp b/examples/samsung/executor_runner/llama/runner/io_manager.cpp index c919d8d8af9..35fa5a6863e 100644 --- a/examples/samsung/executor_runner/llama/runner/io_manager.cpp +++ b/examples/samsung/executor_runner/llama/runner/io_manager.cpp @@ -6,324 +6,370 @@ * LICENSE file in the root directory of this source tree. */ +#include + #include #include namespace example { + +namespace { + +// Helper to fill a float region with a value, using memset when possible. +inline void fill_float_region( + float* begin, + float* end, + float value) { + if (begin >= end) { + return; + } + if (value == 0.0f) { + std::memset(begin, 0, static_cast(end - begin) * sizeof(float)); + } else { + std::fill(begin, end, value); + } +} + +inline void fill_float_region( + float* base, + size_t offset, + size_t len, + float value) { + fill_float_region(base + offset, base + offset + len, value); +} + +} // namespace + executorch::runtime::Error LlmIOManager::load_rotary_embedding(std::optional pos_embedding_path) { + if (pos_embedding_path.has_value()) { + return rotary_embedding_->load(pos_embedding_path.value()); + } + return executorch::runtime::Error::Ok; + } + executorch::runtime::Result> LlmIOManager::prepare_prefill( const executorch::extension::TensorPtr& input, const executorch::extension::TensorPtr& start_pos, const std::string& prefill_method) { (void)prefill_method; - std::vector inputs; - inputs.reserve(3 + config_.n_layers * 2); + inputs.reserve(get_input_reserve_size()); //re-producessing input token tensor const int64_t* prompt_tokens = input->mutable_data_ptr(); size_t prompt_token_len = input->numel(); - - input_tokens_.clear(); - input_tokens_.assign(prompt_tokens, prompt_tokens + prompt_token_len); - input_tokens_.resize(config_.prefill_seq_len, input_tokens_.back()); + input_tokens_.resize(prompt_token_len); + std::copy(prompt_tokens, prompt_tokens + prompt_token_len, input_tokens_.begin()); + input_tokens_.resize(config_.prefill_seq_len, 0); auto input_tokens_tensor = executorch::extension::from_blob( input_tokens_.data(), {1, static_cast(config_.prefill_seq_len)}, executorch::aten::ScalarType::Long); - inputs.emplace_back(input_tokens_tensor); //re-producessing start_pos tensor const int64_t* start_position = start_pos->mutable_data_ptr(); - - pos_ids_.clear(); pos_ids_.resize(config_.prefill_seq_len); + current_pos_ = *start_position; for(int64_t i = 0; i < config_.prefill_seq_len; ++i) { - if (i < prompt_token_len) pos_ids_[i] = i; - else pos_ids_[i] = prompt_token_len - 1; + pos_ids_[i] = current_pos_ + i; } - auto start_pos_tensor = executorch::extension::from_blob( - pos_ids_.data(), {static_cast(config_.prefill_seq_len)}, executorch::aten::ScalarType::Long); - - inputs.emplace_back(start_pos_tensor); - - current_pos_ = *start_position; //attention tensor update_attn_mask(current_pos_, prompt_token_len); - attn_mask_tensor_ = executorch::extension::from_blob( - attn_mask_.data(), - {static_cast(config_.attn_mask_rows), - static_cast(config_.attn_mask_cols)}); + attn_mask_.data(), + {1, 1, static_cast(config_.attn_mask_rows), + static_cast(config_.attn_mask_cols)}); inputs.emplace_back(attn_mask_tensor_); + if (config_.use_gemma) { + update_local_attn_mask(current_pos_, prompt_token_len); + local_attn_mask_tensor_ = executorch::extension::from_blob( + local_attn_mask_.data(), + {1, 1, static_cast(config_.attn_mask_rows), + static_cast(config_.attn_mask_cols)}); + inputs.emplace_back(local_attn_mask_tensor_); + } + + prepare_rotary_embedding_inputs(config_.prefill_seq_len, inputs); + prefill_k_input_tensors_.clear(); prefill_v_input_tensors_.clear(); prefill_k_output_tensors_.clear(); prefill_v_output_tensors_.clear(); - - for (size_t l = 0; l < config_.n_layers; l++) { - auto k_in = executorch::extension::from_blob( - prefill_k_input_[l].data(), - {1, - static_cast(config_.n_kv_heads), - static_cast(prefill_cache_len_), - static_cast(config_.head_dim)}, - executorch::aten::ScalarType::Half); - prefill_k_input_tensors_.push_back(k_in); - inputs.emplace_back(k_in); - } - - for (size_t l = 0; l < config_.n_layers; l++) { - auto v_in = executorch::extension::from_blob( - prefill_v_input_[l].data(), - {1, - static_cast(config_.n_kv_heads), - static_cast(prefill_cache_len_), - static_cast(config_.head_dim)}, - executorch::aten::ScalarType::Half); - prefill_v_input_tensors_.push_back(v_in); - inputs.emplace_back(v_in); - } - + append_kv_input_tensors( + prefill_k_input_, + prefill_v_input_, + prefill_k_input_tensors_, + prefill_v_input_tensors_, + executorch::aten::ScalarType::Float, + prefill_cache_len_, + inputs); dump_kv_input("prefill"); return inputs; } - executorch::runtime::Result> LlmIOManager::prepare_decode( const executorch::extension::TensorPtr& input, const executorch::extension::TensorPtr& start_pos, const std::string& decode_method) { (void)decode_method; - if (!prefill_done_) { copy_prefill_to_decode_input(); prefill_done_ = true; } - const int64_t* start_position = start_pos->mutable_data_ptr(); current_pos_ = *start_position; - - update_attn_mask(current_pos_, 1); - + update_attn_mask(current_pos_, config_.decode_seq_len); std::vector inputs; - inputs.reserve(3 + config_.n_layers * 2); + inputs.reserve(get_input_reserve_size()); inputs.emplace_back(input); - inputs.emplace_back(start_pos); attn_mask_tensor_ = executorch::extension::from_blob( - attn_mask_.data(), - {1, - static_cast(config_.attn_mask_cols)}); + attn_mask_.data(), + {1,1,1, + static_cast(config_.attn_mask_cols)}); inputs.emplace_back(attn_mask_tensor_); + if (config_.use_gemma) { + update_local_attn_mask(current_pos_, 1); + local_attn_mask_tensor_ = executorch::extension::from_blob( + local_attn_mask_.data(), + {1,1,1, + static_cast(config_.attn_mask_cols)});; + inputs.emplace_back(local_attn_mask_tensor_); + } + + pos_ids_.clear(); + pos_ids_.resize(config_.decode_seq_len); + pos_ids_[0] = current_pos_; + prepare_rotary_embedding_inputs(config_.decode_seq_len, inputs); decode_k_input_tensors_.clear(); decode_v_input_tensors_.clear(); decode_k_output_tensors_.clear(); decode_v_output_tensors_.clear(); - - for (size_t l = 0; l < config_.n_layers; l++) { - auto k_in = executorch::extension::from_blob( - decode_k_input_[l].data(), - {1, - static_cast(config_.n_kv_heads), - static_cast(decode_cache_len_), - static_cast(config_.head_dim)}, - executorch::aten::ScalarType::Half); - decode_k_input_tensors_.push_back(k_in); - inputs.emplace_back(k_in); - } - for (size_t l = 0; l < config_.n_layers; l++) { - auto v_in = executorch::extension::from_blob( - decode_v_input_[l].data(), - {1, - static_cast(config_.n_kv_heads), - static_cast(decode_cache_len_), - static_cast(config_.head_dim)}, - executorch::aten::ScalarType::Half); - decode_v_input_tensors_.push_back(v_in); - inputs.emplace_back(v_in); - } - + append_kv_input_tensors( + decode_k_input_, + decode_v_input_, + decode_k_input_tensors_, + decode_v_input_tensors_, + executorch::aten::ScalarType::Float, + decode_cache_len_, + inputs); dump_kv_input("decode"); return inputs; } - executorch::runtime::Error LlmIOManager::update_prefill( const std::vector& model_outputs, const std::string& prefill_method) { (void)prefill_method; - // model_outputs: [logits, k0, ..., k_{L-1}, v0, ..., v_{L-1}] - // k_out shape: [1, n_kv_heads, prefill_seq_len, head_dim] - // Accumulate prefill output into prefill input buffers. The copy to - // decode input is deferred until the first prepare_decode call so that - // multiple micro-batch prefill iterations work correctly. - - for (size_t l = 0; l < config_.n_layers; l++) { - const auto& k_out = model_outputs[1 + l].toTensor(); - const _Float16* k_data = k_out.const_data_ptr<_Float16>(); - copy_to_cache( - k_data, - config_.prefill_seq_len, - prefill_k_input_[l].data(), - prefill_cache_len_, - current_pos_); - } - for (size_t l = 0; l < config_.n_layers; l++) { - const auto& v_out = - model_outputs[1 + config_.n_layers + l].toTensor(); - const _Float16* v_data = v_out.const_data_ptr<_Float16>(); - copy_to_cache( - v_data, - config_.prefill_seq_len, - prefill_v_input_[l].data(), - prefill_cache_len_, - current_pos_); - } - - dump_kv_output("prefill", model_outputs); - dump_step_++; - + update_kv_cache( + model_outputs, + prefill_k_input_, + prefill_v_input_, + prefill_cache_len_, + config_.prefill_seq_len, + "prefill"); return executorch::runtime::Error::Ok; } - executorch::runtime::Error LlmIOManager::update_decode( const std::vector& model_outputs, const std::string& decode_method) { (void)decode_method; - // model_outputs: [logits, k0, ..., k_{L-1}, v0, ..., v_{L-1}] - // k_out shape: [1, n_kv_heads, 1, head_dim] - // Copy decode output into decode input cache at current_pos_. - - for (size_t l = 0; l < config_.n_layers; l++) { - const auto& k_out = model_outputs[1 + l].toTensor(); - const _Float16* k_data = k_out.const_data_ptr<_Float16>(); - copy_to_cache( - k_data, - 1, - decode_k_input_[l].data(), - decode_cache_len_, - current_pos_); - } - for (size_t l = 0; l < config_.n_layers; l++) { - const auto& v_out = - model_outputs[1 + config_.n_layers + l].toTensor(); - const _Float16* v_data = v_out.const_data_ptr<_Float16>(); - copy_to_cache( - v_data, - 1, - decode_v_input_[l].data(), - decode_cache_len_, - current_pos_); - } - - dump_kv_output("decode", model_outputs); - dump_step_++; - + update_kv_cache( + model_outputs, + decode_k_input_, + decode_v_input_, + decode_cache_len_, + 1, + "decode"); return executorch::runtime::Error::Ok; } - void LlmIOManager::allocate_buffers() { size_t prefill_in_size = config_.n_kv_heads * prefill_cache_len_ * - config_.head_dim; + config_.head_dim * sizeof(float); size_t decode_in_size = config_.n_kv_heads * decode_cache_len_ * - config_.head_dim; - + config_.head_dim * sizeof(float); prefill_k_input_.resize(config_.n_layers); prefill_v_input_.resize(config_.n_layers); decode_k_input_.resize(config_.n_layers); decode_v_input_.resize(config_.n_layers); - for (size_t l = 0; l < config_.n_layers; l++) { - prefill_k_input_[l].resize(prefill_in_size, 0.0f); - prefill_v_input_[l].resize(prefill_in_size, 0.0f); - decode_k_input_[l].resize(decode_in_size, 0.0f); - decode_v_input_[l].resize(decode_in_size, 0.0f); + prefill_k_input_[l].resize(prefill_in_size, 0); + prefill_v_input_[l].resize(prefill_in_size, 0); + decode_k_input_[l].resize(decode_in_size, 0); + decode_v_input_[l].resize(decode_in_size, 0); } - attn_mask_.resize( config_.attn_mask_rows * config_.attn_mask_cols, -1.0f); + local_attn_mask_.resize( + config_.attn_mask_rows * config_.attn_mask_cols, -1.0f); } - void LlmIOManager::copy_prefill_to_decode_input() { + size_t copy_size = prefill_cache_len_ * config_.head_dim * sizeof(float); for (size_t l = 0; l < config_.n_layers; l++) { for (size_t h = 0; h < config_.n_kv_heads; h++) { - const _Float16* src = prefill_k_input_[l].data() + - h * prefill_cache_len_ * config_.head_dim; - _Float16* dst = decode_k_input_[l].data() + - h * decode_cache_len_ * config_.head_dim; - std::memcpy( - dst, src, prefill_cache_len_ * config_.head_dim * sizeof(_Float16)); + const char* src = prefill_k_input_[l].data() + + h * prefill_cache_len_ * config_.head_dim * sizeof(float); + char* dst = decode_k_input_[l].data() + + h * decode_cache_len_ * config_.head_dim * sizeof(float); + std::memcpy(dst, src, copy_size); } for (size_t h = 0; h < config_.n_kv_heads; h++) { - const _Float16* src = prefill_v_input_[l].data() + - h * prefill_cache_len_ * config_.head_dim; - _Float16* dst = decode_v_input_[l].data() + - h * decode_cache_len_ * config_.head_dim; - std::memcpy( - dst, src, prefill_cache_len_ * config_.head_dim * sizeof(_Float16)); + const char* src = prefill_v_input_[l].data() + + h * prefill_cache_len_ * config_.head_dim * sizeof(float); + char* dst = decode_v_input_[l].data() + + h * decode_cache_len_ * config_.head_dim * sizeof(float); + std::memcpy(dst, src, copy_size); } } } - void LlmIOManager::copy_to_cache( - const _Float16* src, - size_t new_len, - _Float16* cache, + void LlmIOManager::append_kv_input_tensors( + std::vector>& k_bufs, + std::vector>& v_bufs, + std::vector& k_tensors, + std::vector& v_tensors, + executorch::aten::ScalarType kv_scalar_type, size_t cache_len, - size_t pos) { - for (size_t h = 0; h < config_.n_kv_heads; h++) { - const _Float16* src_head = src + h * new_len * config_.head_dim; - _Float16* dst_head = - cache + h * cache_len * config_.head_dim + pos * config_.head_dim; - std::memcpy( - dst_head, src_head, new_len * config_.head_dim * sizeof(_Float16)); + std::vector& inputs) const { + for (size_t l = 0; l < config_.n_layers; l++) { + auto k_in = executorch::extension::from_blob( + k_bufs[l].data(), + {1, + static_cast(config_.n_kv_heads), + static_cast(cache_len), + static_cast(config_.head_dim)}, + kv_scalar_type); + k_tensors.push_back(k_in); + inputs.emplace_back(k_in); + } + for (size_t l = 0; l < config_.n_layers; l++) { + auto v_in = executorch::extension::from_blob( + v_bufs[l].data(), + {1, + static_cast(config_.n_kv_heads), + static_cast(cache_len), + static_cast(config_.head_dim)}, + kv_scalar_type); + v_tensors.push_back(v_in); + inputs.emplace_back(v_in); } } void LlmIOManager::update_attn_mask(size_t pos, size_t seq_len) { - static int num_prompt_token = 0; // temporary variable to share num_prompt_token info from prefill to decode - std::fill(attn_mask_.begin(), attn_mask_.end(), -std::numeric_limits::infinity()); - - const size_t start_col = config_.attn_mask_cols - config_.attn_mask_rows; - if (seq_len > 1) { // prefill : start from 896 and fill same mask as last_token to all empty rows. - size_t last_valid_row = 0; - num_prompt_token = seq_len; - - for (size_t r = 0; r < seq_len && r < config_.attn_mask_rows; r++) { - size_t visible_count = pos + r + 1; - for (size_t c = 0; c < visible_count; c++) { - size_t target_col = start_col + c; - if (target_col < config_.attn_mask_cols) { - attn_mask_[r * config_.attn_mask_cols + target_col] = 0.0f; - } - } - last_valid_row = r; + const float maskFalse = -255.0f; + const float maskTrue = 0.0f; + + const size_t rows = config_.attn_mask_rows; + const size_t cols = config_.attn_mask_cols; + const size_t start_col = cols - rows; + const size_t total_size = rows * cols; + + // Initialize entire mask to maskFalse (uses memset for 0.0f) + fill_float_region(attn_mask_.data(), 0, total_size, maskFalse); + + if (seq_len > 1) { + // Prefill phase + const size_t rows = config_.attn_mask_rows; + const size_t cols = config_.attn_mask_cols; + const size_t start_col = config_.attn_mask_cols - config_.attn_mask_rows; + + for (int i = 0; i < seq_len; ++i) { + int start_idx = cols - rows; + int end_idx = start_idx + seq_len; + + for (int j = start_idx; j < end_idx && j <= start_idx + i; ++j) { + attn_mask_[i * cols + j] = maskTrue; } - for (size_t r = last_valid_row + 1; r < config_.attn_mask_rows; r++) { - float* src = &attn_mask_[last_valid_row * config_.attn_mask_cols]; - float* dst = &attn_mask_[r * config_.attn_mask_cols]; - std::memcpy(dst, src, config_.attn_mask_cols * sizeof(float)); + } + + for (int i = 0; i < rows; ++i) { + for (int j = 0; j < pos; ++j) { + attn_mask_[i * cols + j] = maskTrue; } + } + } else { + const size_t left_len = std::min(pos, cols); + if (left_len > 0) { + fill_float_region(attn_mask_.data(), 0, left_len, maskTrue); + } + + if (cols > 0) { + attn_mask_[cols - 1] = maskTrue; + } } - else { - int pos_decode = pos - num_prompt_token + 1; - for (int r = 0; r < (pos); r++) { - attn_mask_[r] = 0.0f; + } + + void LlmIOManager::update_local_attn_mask(size_t pos, size_t seq_len) { + const float maskFalse = -255.0f; + const float maskTrue = 0.0f; + + const size_t rows = config_.attn_mask_rows; + const size_t cols = config_.attn_mask_cols; + const size_t slide_window = 512; + const size_t total_size = rows * cols; + + // Start with a copy of the full attention mask + std::memcpy(local_attn_mask_.data(), attn_mask_.data(), total_size * sizeof(float)); + + if (seq_len > 1) { + // Prefill phase with sliding window + if ((pos + seq_len) >= slide_window) { + // Mask out positions before the sliding window + const size_t mask_start = slide_window - seq_len; + const size_t mask_end = pos; //512 + const size_t mask_len = mask_end - mask_start; + + if (mask_len > 0) { + for (size_t i = 0; i < rows; ++i) { + fill_float_region( + local_attn_mask_.data(), + i * cols + mask_start, + mask_len, + maskFalse); + } + } + } + } else { + // Decode phase (single token) + // Set all positions up to pos to maskTrue + const size_t end = std::min(static_cast(pos), cols); + fill_float_region(local_attn_mask_.data(), 0, end, maskTrue); + + // If there are more than 512 True values, we need to keep only the last 512 + // Count True values efficiently by scanning from the end + size_t valid_count = 1; + for (size_t i = 0; i < cols; ++i) { + if (local_attn_mask_[i] == maskTrue) { + valid_count++; + } else { + break; + } + } + + if (valid_count > slide_window) { + const size_t to_remove = valid_count - slide_window; + // Set the oldest True values to maskFalse + for (size_t i = 0; i < to_remove; ++i) { + if (local_attn_mask_[i] == maskTrue) { + local_attn_mask_[i] = maskFalse; + } + } } - attn_mask_[config_.attn_mask_cols - 1] = 0.0f; + if (cols > 0) { + attn_mask_[cols - 1] = maskTrue; + } } } void LlmIOManager::dump_buffer( - const std::string& tag, size_t layer, const _Float16* data, size_t size) { + const std::string& tag, size_t layer, const char* data, size_t size) { if (!dump_enabled_) { return; } @@ -332,14 +378,12 @@ namespace example { filename, sizeof(filename), "%s/%s_layer%zu_pos%zu_step%zu.bin", dump_dir_.c_str(), tag.c_str(), layer, current_pos_, dump_step_); - FILE* f = std::fopen(filename, "wb"); if (f) { - std::fwrite(data, sizeof(_Float16), size, f); + std::fwrite(data, 1, size, f); std::fclose(f); } } - void LlmIOManager::dump_kv_input(const std::string& phase) { if (!dump_enabled_) { return; @@ -352,7 +396,33 @@ namespace example { dump_buffer(phase + "_v_input", l, v_bufs[l].data(), v_bufs[l].size()); } } - + void LlmIOManager::update_kv_cache( + const std::vector& model_outputs, + std::vector>& k_bufs, + std::vector>& v_bufs, + size_t cache_len, + size_t new_len, + const std::string& phase) { + // model_outputs: [logits, k0..k_{L-1}, v0..v_{L-1}] + for (size_t l = 0; l < config_.n_layers; l++) { + copy_tensor_to_cache( + model_outputs[1 + l].toTensor(), + new_len, + k_bufs[l].data(), + cache_len, + current_pos_); + } + for (size_t l = 0; l < config_.n_layers; l++) { + copy_tensor_to_cache( + model_outputs[1 + config_.n_layers + l].toTensor(), + new_len, + v_bufs[l].data(), + cache_len, + current_pos_); + } + dump_kv_output(phase, model_outputs); + dump_step_++; + } void LlmIOManager::dump_kv_output( const std::string& phase, const std::vector& model_outputs) { @@ -366,13 +436,56 @@ namespace example { dump_buffer( phase + "_k_output", l, - k_out.const_data_ptr<_Float16>(), - k_out.numel()); + reinterpret_cast(k_out.const_data_ptr()), + k_out.numel() * sizeof(float)); dump_buffer( phase + "_v_output", l, - v_out.const_data_ptr<_Float16>(), - v_out.numel()); + reinterpret_cast(v_out.const_data_ptr()), + v_out.numel() * sizeof(float)); } } + + void LlmIOManager::prepare_rotary_embedding_inputs( + size_t seq_len, + std::vector& inputs) { + const size_t element_size = seq_len * config_.head_dim; + global_cos_data_.resize(element_size); + global_sin_data_.resize(element_size); + if (config_.use_gemma) { + local_cos_data_.resize(element_size); + local_sin_data_.resize(element_size); + } + + if (rotary_embedding_->update(pos_ids_, global_cos_data_.data(), global_sin_data_.data(), local_cos_data_.data(), local_sin_data_.data()) != executorch::runtime::Error::Ok) { + ET_LOG(Error, "update pos_embedding failed !"); + } + + global_cos_tensor_ = executorch::extension::from_blob( + global_cos_data_.data(), + {1, static_cast(seq_len), + static_cast(config_.head_dim)}); + global_sin_tensor_ = executorch::extension::from_blob( + global_sin_data_.data(), + {1, static_cast(seq_len), + static_cast(config_.head_dim)}); + if (config_.use_gemma) { + local_cos_tensor_ = executorch::extension::from_blob( + local_cos_data_.data(), + {1, static_cast(seq_len), + static_cast(config_.head_dim)}); + local_sin_tensor_ = executorch::extension::from_blob( + local_sin_data_.data(), + {1, static_cast(seq_len), + static_cast(config_.head_dim)}); + } + + inputs.emplace_back(global_cos_tensor_); + inputs.emplace_back(global_sin_tensor_); + if (config_.use_gemma) { + inputs.emplace_back(local_cos_tensor_); + inputs.emplace_back(local_sin_tensor_); + } + } + } // namespace example diff --git a/examples/samsung/executor_runner/llama/runner/io_manager.h b/examples/samsung/executor_runner/llama/runner/io_manager.h index 0dd78e24f25..dcaec626486 100644 --- a/examples/samsung/executor_runner/llama/runner/io_manager.h +++ b/examples/samsung/executor_runner/llama/runner/io_manager.h @@ -8,10 +8,12 @@ #pragma once +#include #include #include #include +#include namespace example { @@ -43,14 +45,17 @@ namespace example { class ET_EXPERIMENTAL LlmIOManager : public executorch::extension::llm::IOManager { public: - struct Config { + struct Config { size_t max_seq_len; size_t prefill_seq_len; + size_t decode_seq_len; size_t n_layers; size_t n_kv_heads; size_t head_dim; size_t attn_mask_rows = 128; size_t attn_mask_cols = 1024; + size_t sliding_window = 512; + bool use_gemma = false; }; /** @@ -60,11 +65,18 @@ class ET_EXPERIMENTAL LlmIOManager */ LlmIOManager(executorch::extension::Module& module, Config config) : IOManager(module), module_(module), config_(config) { - ET_LOG(Info, - "(%s) layers(%ld), kv_heads(%ld) max_sqe_len(%ld) \ - head_dim(%ld) prefill_seq_len(%ld)", - ET_FUNCTION, config_.n_layers, config_.n_kv_heads, - config_.max_seq_len, config_.head_dim, config_.prefill_seq_len); + rotary_embedding_ = std::make_unique(RotaryEmbedding::Config{config_.head_dim, config_.max_seq_len, config_.use_gemma}); + ET_LOG( + Info, + "(%s) layers(%ld), kv_heads(%ld) max_sqe_len(%ld) \ + head_dim(%ld) prefill_seq_len(%ld) use_gemma(%d)", + ET_FUNCTION, + config_.n_layers, + config_.n_kv_heads, + config_.max_seq_len, + config_.head_dim, + config_.prefill_seq_len, + config_.use_gemma); (void)module_; prefill_cache_len_ = config_.max_seq_len - config_.prefill_seq_len; @@ -73,6 +85,20 @@ class ET_EXPERIMENTAL LlmIOManager allocate_buffers(); } + /** + * @brief Load pos_embedding files.. + * @param pos_embedding_path Pos_embedding file, that include global_cos.bin, global_sin.bin, local_cos.bin, local_sin.bin. + */ + ET_NODISCARD virtual executorch::runtime::Error load_rotary_embedding(std::optional pos_embedding_path); + + /** + * @brief Get the use_gemma configuration flag. + * @return True if using Gemma model specific processing, false otherwise. + */ + bool use_gemma() const { + return config_.use_gemma; + } + /** * @brief Virtual destructor to allow proper cleanup in derived classes. */ @@ -89,10 +115,14 @@ class ET_EXPERIMENTAL LlmIOManager const std::string& prefill_method, const std::string& decode_method) override { ET_LOG(Info, - "(%s) layers(%ld), kv_heads(%ld) max_sqe_len(%ld) \ - head_dim(%ld) prefill_seq_len(%ld)", - ET_FUNCTION, config_.n_layers, config_.n_kv_heads, - config_.max_seq_len, config_.head_dim, config_.prefill_seq_len); + "(%s) layers(%ld), kv_heads(%ld) max_sqe_len(%ld) \ + head_dim(%ld) prefill_seq_len(%ld)", + ET_FUNCTION, + config_.n_layers, + config_.n_kv_heads, + config_.max_seq_len, + config_.head_dim, + config_.prefill_seq_len); return executorch::runtime::Error::Ok; } @@ -116,13 +146,13 @@ class ET_EXPERIMENTAL LlmIOManager for (size_t l = 0; l < config_.n_layers; l++) { std::fill( - prefill_k_input_[l].begin(), prefill_k_input_[l].end(), 0.0f); + prefill_k_input_[l].begin(), prefill_k_input_[l].end(), 0); std::fill( - prefill_v_input_[l].begin(), prefill_v_input_[l].end(), 0.0f); + prefill_v_input_[l].begin(), prefill_v_input_[l].end(), 0); std::fill( - decode_k_input_[l].begin(), decode_k_input_[l].end(), 0.0f); + decode_k_input_[l].begin(), decode_k_input_[l].end(), 0); std::fill( - decode_v_input_[l].begin(), decode_v_input_[l].end(), 0.0f); + decode_v_input_[l].begin(), decode_v_input_[l].end(), 0); } std::fill(attn_mask_.begin(), attn_mask_.end(), 0.0f); @@ -246,10 +276,10 @@ class ET_EXPERIMENTAL LlmIOManager } private: - void dump_buffer( + void dump_buffer( const std::string& tag, size_t layer, - const _Float16* data, + const char* data, size_t size); void dump_kv_input(const std::string& phase); void dump_kv_output( @@ -258,11 +288,27 @@ class ET_EXPERIMENTAL LlmIOManager void allocate_buffers(); + /** + * @brief Calculate the number of inputs to reserve for prefill/decode. + * + * The total depends on whether the model is Gemma: + * - Gemma: 7 + n_layers * 2 + * - Non-Gemma: 4 + n_layers * 2 + * + * @return The reserve size for the inputs vector. + */ + size_t get_input_reserve_size() const { + if (config_.use_gemma) + return 7 + config_.n_layers * 2; + else + return 4 + config_.n_layers * 2; + } + bool dump_enabled_ = false; mutable size_t dump_step_ = 0; std::string dump_dir_ = "/data/local/tmp/kv_cache_dump"; - /** + /** * Copy prefill output KV cache (layout [1,H,prefill_cache_len,D]) into * the decode input buffer (layout [1,H,decode_cache_len,D]). * @@ -271,15 +317,73 @@ class ET_EXPERIMENTAL LlmIOManager void copy_prefill_to_decode_input(); /** - * Copy new cache entries from output (layout [1,H,new_len,D]) into the - * decode input buffer (layout [1,H,cache_len,D]) at the given position. + * Copy model KV-cache outputs into the given buffers and dump them. + * + * @param model_outputs [logits, k0..k_{L-1}, v0..v_{L-1}] + * @param k_bufs per-layer K cache buffers (input/output) + * @param v_bufs per-layer V cache buffers (input/output) + * @param cache_len length of each cache buffer + * @param new_len number of new tokens to copy (prefill_seq_len or 1) + * @param phase dump label ("prefill" or "decode") + */ + void update_kv_cache( + const std::vector& model_outputs, + std::vector>& k_bufs, + std::vector>& v_bufs, + size_t cache_len, + size_t new_len, + const std::string& phase); + + /** + * Copy new cache entries from source data (layout [1,H,new_len,D]) into the + * cache buffer (layout [1,H,cache_len,D]) at the given position. * * For each head h: * cache[h * cache_len * D + pos * D .. + (pos + new_len) * D] * = src[h * new_len * D .. + new_len * D] */ - void copy_to_cache(const _Float16* src, size_t new_len, - _Float16* cache, size_t cache_len, size_t pos); + template + void copy_to_cache(const T* src, size_t new_len, + char* cache, size_t cache_len, size_t pos) { + for (size_t h = 0; h < config_.n_kv_heads; h++) { + const T* src_head = src + h * new_len * config_.head_dim; + char* dst_head = + cache + h * cache_len * config_.head_dim * sizeof(T) + + pos * config_.head_dim * sizeof(T); + std::memcpy( + dst_head, src_head, new_len * config_.head_dim * sizeof(T)); + } + } + + /** + * Copy tensor data to cache buffer, dispatching based on kv_scalar_type. + * Handles both Half (float16) and Float (float32) types. + */ + void copy_tensor_to_cache( + const executorch::aten::Tensor& tensor, + size_t new_len, + char* cache, + size_t cache_len, + size_t pos) { + const float* data = tensor.const_data_ptr(); + copy_to_cache(data, new_len, cache, cache_len, pos); + } + + /** + * Append KV-cache input tensors for all layers. + * + * For each layer l, creates k and v tensors from the given buffers + * (shape [1, n_kv_heads, cache_len, head_dim]) and appends them + * to both the per-phase tensor wrappers and the model inputs vector. + */ + void append_kv_input_tensors( + std::vector>& k_bufs, + std::vector>& v_bufs, + std::vector& k_tensors, + std::vector& v_tensors, + executorch::aten::ScalarType kv_scalar_type, + size_t cache_len, + std::vector& inputs) const; /** * Update the attention mask for a causal (lower-triangular) pattern. @@ -288,6 +392,28 @@ class ET_EXPERIMENTAL LlmIOManager */ void update_attn_mask(size_t pos, size_t seq_len); + /** + * @brief Prepare rotary embedding inputs for quantized models. + * + * This helper function computes and adds the rotary embedding tensors + * (global_cos, global_sin, local_cos, local_sin) to the inputs vector. + * It is used by both prepare_prefill and prepare_decode to avoid code + * duplication. + * + * @param seq_len The sequence length to use for the embedding computation. + * @param inputs Vector of EValue inputs to append the embedding tensors to. + */ + void prepare_rotary_embedding_inputs( + size_t seq_len, + std::vector& inputs); + + /** + * Update the local_attention mask for a causal (lower-triangular) pattern. + * For each query row r in [0, seq_len), the mask allows attending to + * positions [0, pos + r] (value 1) and masks future positions (value 0). + */ + void update_local_attn_mask(size_t pos, size_t seq_len); + /** * @brief Reference to the Module used for method metadata and execution. */ @@ -303,14 +429,29 @@ class ET_EXPERIMENTAL LlmIOManager std::vector pos_ids_; // Per-layer input buffers: [1, n_kv_heads, cache_len, head_dim] flattened - std::vector> prefill_k_input_; - std::vector> prefill_v_input_; - std::vector> decode_k_input_; - std::vector> decode_v_input_; + // Use std::vector to support both float and half (float16) data types + std::vector> prefill_k_input_; + std::vector> prefill_v_input_; + std::vector> decode_k_input_; + std::vector> decode_v_input_; + + std::unique_ptr rotary_embedding_; + + std::vector global_cos_data_; + std::vector global_sin_data_; + std::vector local_cos_data_; + std::vector local_sin_data_; + + executorch::extension::TensorPtr global_cos_tensor_; + executorch::extension::TensorPtr global_sin_tensor_; + executorch::extension::TensorPtr local_cos_tensor_; + executorch::extension::TensorPtr local_sin_tensor_; // Attention mask buffer: [attn_mask_rows, attn_mask_cols] flattened std::vector attn_mask_; + std::vector local_attn_mask_; executorch::extension::TensorPtr attn_mask_tensor_; + executorch::extension::TensorPtr local_attn_mask_tensor_; // Tensor wrappers (kept alive between prepare and model execution) std::vector prefill_k_input_tensors_; diff --git a/examples/samsung/executor_runner/llama/runner/prefiller.cpp b/examples/samsung/executor_runner/llama/runner/prefiller.cpp index ca41cd274e2..3749630fa38 100644 --- a/examples/samsung/executor_runner/llama/runner/prefiller.cpp +++ b/examples/samsung/executor_runner/llama/runner/prefiller.cpp @@ -69,9 +69,9 @@ ::executorch::runtime::Result Prefiller::prefill( ET_CHECK_OK_OR_RETURN_ERROR(chunk_result.error()); cur_token = chunk_result.get(); - if (enable_parallel_prefill_) { - start_pos += num_tokens_to_prefill_with; - } + //if (enable_parallel_prefill_) { + // start_pos += num_tokens_to_prefill_with; + //} num_tokens_to_process += num_tokens_to_prefill_with; } diff --git a/examples/samsung/executor_runner/llama/runner/rotary_embedding.cpp b/examples/samsung/executor_runner/llama/runner/rotary_embedding.cpp new file mode 100644 index 00000000000..cc833f10740 --- /dev/null +++ b/examples/samsung/executor_runner/llama/runner/rotary_embedding.cpp @@ -0,0 +1,80 @@ +/* + * Copyright (c) 2025 Samsung Electronics Co. LTD + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#include +#include +#include +#include + +namespace example { + +RotaryEmbedding::RotaryEmbedding(Config config) : config_(config) {} + +executorch::runtime::Error RotaryEmbedding::load(std::optional pos_embedding_path) { + if (!pos_embedding_path.has_value()) { + ET_LOG(Error, "RotaryEmbedding not has value"); + return executorch::runtime::Error::Ok; + } + + const size_t total_bytes = config_.max_seq_len * config_.head_dim * sizeof(float); + global_cos_ = std::make_unique(total_bytes); + global_sin_ = std::make_unique(total_bytes); + if (config_.use_gemma) { + local_cos_ = std::make_unique(total_bytes); + local_sin_ = std::make_unique(total_bytes); + } + + std::ifstream file(pos_embedding_path.value(), std::ios::binary); + if (!file) { + ET_LOG(Error, "can't open the file"); + return executorch::runtime::Error::Internal; + } + + file.seekg(0, std::ios::end); + size_t file_size = file.tellg(); + file.seekg(0, std::ios::beg); + + int bin_num = config_.use_gemma ? 4 : 2; + if (file_size != total_bytes * bin_num) { + ET_LOG(Error, "file size is not right"); + return executorch::runtime::Error::Internal; + } + + file.read(global_cos_.get(), total_bytes); + file.read(global_sin_.get(), total_bytes); + if (config_.use_gemma) { + file.read(local_cos_.get(), total_bytes); + file.read(local_sin_.get(), total_bytes); + } + + file.close(); + return executorch::runtime::Error::Ok; +} + +executorch::runtime::Error RotaryEmbedding::update(const std::vector& pos_ids, + float* global_cos, float* global_sin, + float* local_cos, float* local_sin) { + if (pos_ids.empty()) { + return executorch::runtime::Error::Internal; + } + + size_t pos = static_cast(pos_ids[0]); + size_t length = pos_ids.size(); + size_t head_dim = config_.head_dim; + + std::memcpy(global_cos, global_cos_.get() + pos * head_dim * sizeof(float), length * head_dim * sizeof(float)); + std::memcpy(global_sin, global_sin_.get() + pos * head_dim * sizeof(float), length * head_dim * sizeof(float)); + if (config_.use_gemma) { + std::memcpy(local_cos, local_cos_.get() + pos * head_dim * sizeof(float), length * head_dim * sizeof(float)); + std::memcpy(local_sin, local_sin_.get() + pos * head_dim * sizeof(float), length * head_dim * sizeof(float)); + } + + return executorch::runtime::Error::Ok; +} + +} // namespace example diff --git a/examples/samsung/executor_runner/llama/runner/rotary_embedding.h b/examples/samsung/executor_runner/llama/runner/rotary_embedding.h new file mode 100644 index 00000000000..5c3cd5d3973 --- /dev/null +++ b/examples/samsung/executor_runner/llama/runner/rotary_embedding.h @@ -0,0 +1,41 @@ +/* + * Copyright (c) 2025 Samsung Electronics Co. LTD + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +#include +#include +#include +#include + +namespace example { + +class RotaryEmbedding { + public: + struct Config { + size_t head_dim; + size_t max_seq_len; + bool use_gemma; + }; + + RotaryEmbedding(Config config); + ~RotaryEmbedding() = default; + + executorch::runtime::Error load(std::optional pos_embedding_path); + executorch::runtime::Error update(const std::vector& pos_ids, + float* global_cos, float* global_sin, + float* local_cos, float* local_sin); + private: + Config config_; + std::unique_ptr global_cos_; + std::unique_ptr global_sin_; + std::unique_ptr local_cos_; + std::unique_ptr local_sin_; +}; + +} // namespace example diff --git a/examples/samsung/executor_runner/llama/runner/runner.cpp b/examples/samsung/executor_runner/llama/runner/runner.cpp index 4a3b64a86ca..81a23b74783 100644 --- a/examples/samsung/executor_runner/llama/runner/runner.cpp +++ b/examples/samsung/executor_runner/llama/runner/runner.cpp @@ -19,12 +19,11 @@ #include #include +#include #include #include #include #include -#include - namespace example { @@ -37,22 +36,19 @@ namespace llm = ::executorch::extension::llm; /** * @brief Loads a tokenizer from the specified path * - * This function creates and initializes a tokenizer from a file, with options - * to customize special tokens and regex patterns. + * This function creates and initializes a tokenizer from a file based on the + * specified type. Supported types are: "bpe", "tiktoken", "hf", and "sp". * + * @param type Type of tokenizer to load ("bpe", "tiktoken", "hf", or "sp") * @param tokenizer_path Path to the tokenizer file - * @param special_tokens Optional list of special tokens to add to the tokenizer - * @param pattern Optional regex pattern for tokenization - * @param bos_token_index Index of the beginning-of-sequence token - * @param eos_token_index Index of the end-of-sequence token - * @return std::unique_ptr Initialized tokenizer instance + * @param version Version of the tokenizer (used for tiktoken type) + * @return std::unique_ptr Initialized tokenizer instance, + * or nullptr on failure */ ET_EXPERIMENTAL std::unique_ptr load_tokenizer( + const std::string& type, const std::string& tokenizer_path, - std::unique_ptr> special_tokens = nullptr, - std::optional pattern = std::nullopt, - size_t bos_token_index = 0, - size_t eos_token_index = 1); + Version version); /** * @brief Creates a LLMRunner instance with the specified model and @@ -62,56 +58,42 @@ ET_EXPERIMENTAL std::unique_ptr load_tokenizer( * necessary components for text generation using the specified model and * tokenizer. * - * @param model_path Path to the model file + * @param paths RunnerPaths containing model and tokenizer paths + * @param config RunnerConfig containing runtime configuration * @param tokenizer Initialized tokenizer instance - * @param data_path Optional path to additional data required by the model - * @param decode_only Optional flag to run with decode method only - * @param temperature Optional temperature parameter for controlling randomness - * (deprecated) * @return std::unique_ptr Initialized LLMRunner instance */ ET_EXPERIMENTAL std::unique_ptr create_llm_runner( - const std::string& model_path, - std::unique_ptr<::tokenizers::Tokenizer> tokenizer, - std::optional data_path = std::nullopt, - bool decode_only = false, - float temperature = -1.0f); - -std::unique_ptr<::tokenizers::Tokenizer> load_llama_tokenizer( - const std::string& tokenizer_path, - Version version) { - auto special_tokens = get_special_tokens(version); - return load_tokenizer(tokenizer_path, std::move(special_tokens)); -} + const RunnerPaths& runner_paths, + const RunnerConfig& runner_config, + std::unique_ptr<::tokenizers::Tokenizer> tokenizer); std::unique_ptr create_llama_runner( - const std::string& model_path, - const std::string& tokenizer_path, - std::optional data_path, - bool decode_only, - float temperature) { + const RunnerPaths& runner_paths, + const RunnerConfig& runner_config) { ET_LOG( Info, - "Creating LLaMa runner: model_path=%s, tokenizer_path=%s", - model_path.c_str(), - tokenizer_path.c_str()); + "Creating LLaMa runner: model_path=%s, tokenizer_path=%s, use_gemma=%d", + runner_paths.model_path.c_str(), + runner_paths.tokenizer_path.c_str(), + runner_config.use_gemma); // Create and load tokenizer std::unique_ptr<::tokenizers::Tokenizer> tokenizer = - load_llama_tokenizer(tokenizer_path, Version::Default); + load_tokenizer(runner_config.tokenizer_type, runner_paths.tokenizer_path, Version::Default); if (tokenizer == nullptr) { ET_LOG( Info, "Failed to load %s as a Tiktoken, Sentencepiece or Llama2.c tokenizer, make sure the artifact is one of these types", - tokenizer_path.c_str()); + runner_paths.tokenizer_path.c_str()); return nullptr; } ET_LOG( Info, "[Create_test_llm_runner] %s %s", - model_path.c_str(), - data_path->c_str()); - return create_llm_runner(model_path, std::move(tokenizer), data_path, decode_only); + runner_paths.model_path.c_str(), + runner_paths.data_path.has_value() ? runner_paths.data_path.value().c_str() : "null"); + return create_llm_runner(runner_paths, runner_config, std::move(tokenizer)); } static const std::unordered_map kMetadataMap{ @@ -123,11 +105,13 @@ static const std::unordered_map kMetadataMap{ {llm::kVocabSize, "get_vocab_size"}, {llm::kUseKVCache, "use_kv_cache"}, {llm::kUseSDPAWithKVCache, "use_sdpa_with_kv_cache"}, + {llm::kHeadDim, "get_head_dim"}, {example::kPrefillArLen, "get_prefill_ar_len"}, {example::kNumOfHiddenDim, "get_hidden_dim"}, {example::kNumOfHeads, "get_n_heads"}, {example::kNumOfKVHeads, "get_n_kv_heads"}, - {example::kNumOfLayers, "get_n_layers"} + {example::kNumOfLayers, "get_n_layers"}, + {"get_sliding_window", "get_sliding_window"} // TODO: find new way }; LLMRunner::LLMRunner( @@ -172,6 +156,17 @@ Error LLMRunner::load() { ET_LOG(Info, format, __VA_ARGS__); \ } +void replace_all_newlines(std::string& str) { + std::string from = "\\n"; + std::string to = "\n"; + + size_t start_pos = 0; + while ((start_pos = str.find(from, start_pos)) != std::string::npos) { + str.replace(start_pos, from.length(), to); + start_pos += to.length(); + } +} + Error LLMRunner::generate_from_pos( const std::string& prompt, int64_t start_pos, @@ -217,14 +212,25 @@ Error LLMRunner::generate_from_pos( Info, "################ input prompt ################\n %s\n", prompt.c_str()); + std::string input_prompt = prompt; + if (!io_manager_->use_gemma()) { + std::string prefix = "<|start_header_id|>system<|end_header_id|><|start_header_id|>user<|end_header_id|>\n\n"; + std::string suffix = "<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n"; + input_prompt = prefix + prompt + suffix; + } else { + std::string prefix = "user\n"; + std::string suffix = "\nmodel\n"; + input_prompt = prefix + prompt + suffix; + replace_all_newlines(input_prompt); + } - ::tokenizers::Result> encode_res = tokenizer_->encode( - prompt, + ::tokenizers::Result> encode_res = tokenizer_->encode( + input_prompt, /*bos=*/config.num_bos, /*eos=*/config.num_eos); ET_CHECK_TK_OK_OR_RETURN_ERROR( - encode_res.error(), "Failed to encode prompt %s", prompt.c_str()); + encode_res.error(), "Failed to encode prompt %s", input_prompt.c_str()); // encode the (string) prompt into tokens sequence std::vector prompt_tokens = encode_res.get(); @@ -268,7 +274,7 @@ Error LLMRunner::generate_from_pos( // print prompts if (config.echo) { - wrapped_callback(prompt); + wrapped_callback(input_prompt); } int64_t pos = start_pos; auto prefill_res = text_prefiller_->prefill(prompt_tokens, pos); @@ -372,47 +378,25 @@ void LLMRunner::stop() { } std::unique_ptr load_tokenizer( + const std::string& type, const std::string& tokenizer_path, - std::unique_ptr> special_tokens, - std::optional pattern, - size_t bos_token_index, - size_t eos_token_index) { + Version version) { ::executorch::runtime::runtime_init(); - std::unique_ptr<::tokenizers::Tiktoken> tiktoken_tokenizer; - if (special_tokens != nullptr && !pattern.has_value()) { - tiktoken_tokenizer = std::make_unique<::tokenizers::Tiktoken>( - std::move(special_tokens), bos_token_index, eos_token_index); - } else if (special_tokens != nullptr && pattern.has_value()) { - tiktoken_tokenizer = std::make_unique<::tokenizers::Tiktoken>( - pattern.value(), - std::move(special_tokens), - bos_token_index, - eos_token_index); - } else { - tiktoken_tokenizer = std::make_unique<::tokenizers::Tiktoken>(); - } - if (tiktoken_tokenizer->load(tokenizer_path) == ::tokenizers::Error::Ok) { - ET_LOG(Info, "Loaded TikToken tokenizer"); - return tiktoken_tokenizer; + std::unique_ptr<::tokenizers::Tokenizer> tokenizer; + if (type == "bpe") { + tokenizer = std::make_unique<::tokenizers::Llama2cTokenizer>(); + } else if (type == "tiktoken") { + tokenizer = example::get_tiktoken_for_llama(version); + } else if (type == "hf") { + tokenizer = std::make_unique<::tokenizers::HFTokenizer>(); + } else if (type == "sp") { + tokenizer = std::make_unique<::tokenizers::SPTokenizer>(); } - auto json_tokenizer = std::make_unique(); - if (json_tokenizer->load(tokenizer_path) == ::tokenizers::Error::Ok) { - ET_LOG(Info, "Loaded json tokenizer"); - return json_tokenizer; - } - - auto sp_tokenizer = std::make_unique<::tokenizers::SPTokenizer>(); - if (sp_tokenizer->load(tokenizer_path) == ::tokenizers::Error::Ok) { - ET_LOG(Info, "Loaded Sentencepiece tokenizer"); - return sp_tokenizer; - } - - auto bpe_tokenizer = std::make_unique<::tokenizers::Llama2cTokenizer>(); - if (bpe_tokenizer->load(tokenizer_path) == ::tokenizers::Error::Ok) { - ET_LOG(Info, "Loaded BPE tokenizer"); - return bpe_tokenizer; + if (tokenizer->load(tokenizer_path) == ::tokenizers::Error::Ok) { + ET_LOG(Info, "Loaded %s tokenizer", type.c_str()); + return tokenizer; } return nullptr; @@ -420,8 +404,17 @@ std::unique_ptr load_tokenizer( std::unordered_map get_llm_metadata( tokenizers::Tokenizer* tokenizer, - Module* module) { - // Initialize metadata with default values + Module* module, + bool use_gemma) { + // Initialize metadata with default values based on model type + // Gemma: hidden_dim=1152, n_heads=4, n_kv_heads=1, n_layers=26, head_dim=256 + // Llama: hidden_dim=2048, n_heads=32, n_kv_heads=8, n_layers=16, head_dim=64 + size_t default_hidden_dim = use_gemma ? 1152 : 2048; + size_t default_n_heads = use_gemma ? 4 : 32; + size_t default_n_kv_heads = use_gemma ? 1 : 8; + size_t default_n_layers = use_gemma ? 26 : 16; + size_t default_head_dim = use_gemma ? 256 : 64; + std::unordered_map metadata({ {kMetadataMap.at(llm::kEnableDynamicShape), false}, {kMetadataMap.at(llm::kMaxSeqLen), 128}, @@ -429,11 +422,11 @@ std::unordered_map get_llm_metadata( {kMetadataMap.at(llm::kUseKVCache), true}, {kMetadataMap.at(llm::kUseSDPAWithKVCache), false}, {kMetadataMap.at(example::kPrefillArLen), 0}, - //Params.json for llama3.2-1B - {kMetadataMap.at(example::kNumOfHiddenDim), 2048}, - {kMetadataMap.at(example::kNumOfHeads), 32}, - {kMetadataMap.at(example::kNumOfKVHeads), 8}, - {kMetadataMap.at(example::kNumOfLayers), 16}, + {kMetadataMap.at(example::kNumOfHiddenDim), static_cast(default_hidden_dim)}, + {kMetadataMap.at(example::kNumOfHeads), static_cast(default_n_heads)}, + {kMetadataMap.at(example::kNumOfKVHeads), static_cast(default_n_kv_heads)}, + {kMetadataMap.at(example::kNumOfLayers), static_cast(default_n_layers)}, + {kMetadataMap.at(llm::kHeadDim), static_cast(default_head_dim)}, }); // Read metadata from the model @@ -495,11 +488,9 @@ std::unordered_set get_eos_ids( } std::unique_ptr create_llm_runner( - const std::string& model_path, - std::unique_ptr<::tokenizers::Tokenizer> tokenizer, - std::optional data_path, - bool decode_only, - float temperature) { + const RunnerPaths& runner_paths, + const RunnerConfig& runner_config, + std::unique_ptr<::tokenizers::Tokenizer> tokenizer) { // Sanity check tokenizer if (!tokenizer || !tokenizer->is_loaded()) { ET_LOG(Error, "Tokenizer is null or not loaded"); @@ -508,11 +499,11 @@ std::unique_ptr create_llm_runner( // Create the Module std::unique_ptr module; - if (data_path.has_value()) { + if (runner_paths.data_path.has_value()) { module = std::make_unique( - model_path, data_path.value(), Module::LoadMode::File); + runner_paths.model_path, runner_paths.data_path.value(), Module::LoadMode::File); } else { - module = std::make_unique(model_path, Module::LoadMode::File); + module = std::make_unique(runner_paths.model_path, Module::LoadMode::File); } // Get methods @@ -524,7 +515,7 @@ std::unique_ptr create_llm_runner( auto methods = methods_res.get(); std::vector register_methods; - if (decode_only) { + if (runner_config.decode_only) { register_methods.emplace_back(default_method_name_decode); } else { for (auto method : methods) { @@ -535,30 +526,49 @@ std::unique_ptr create_llm_runner( ET_LOG(Info, "Register method \"%s\" to decoder", method.c_str()); register_methods.emplace_back(method); } - } + } // Get metadata from Module ET_LOG(Info, "Reading metadata from model"); - auto metadata = get_llm_metadata(tokenizer.get(), module.get()); + auto metadata = get_llm_metadata(tokenizer.get(), module.get(), runner_config.use_gemma); // Create IO Manager example::LlmIOManager::Config conf; conf.n_layers = metadata.at(example::kNumOfLayers); conf.n_kv_heads = metadata.at(example::kNumOfKVHeads); conf.max_seq_len = metadata.at(llm::kMaxSeqLen); + conf.decode_seq_len = 1; conf.prefill_seq_len = (metadata.at(example::kPrefillArLen) > 1) ? (metadata.at(example::kPrefillArLen)) : (metadata.at(llm::kMaxSeqLen)); + conf.use_gemma = runner_config.use_gemma; - if (decode_only) { + if (runner_config.decode_only) { conf.prefill_seq_len = 1; } - size_t dim = metadata.at(example::kNumOfHiddenDim); - size_t n_heads = metadata.at(example::kNumOfHeads); - conf.head_dim = dim / n_heads; + // Try to get head_dim from metadata, otherwise calculate from dim/n_heads + if (metadata.count(llm::kHeadDim)) { + conf.head_dim = metadata.at(llm::kHeadDim); + } else { + size_t dim = metadata.at(example::kNumOfHiddenDim); + size_t n_heads = metadata.at(example::kNumOfHeads); + if (runner_config.use_gemma) { + conf.head_dim = 256; + } else { + conf.head_dim = dim / n_heads; + } + } + + auto io_manager = + std::make_unique(*module.get(), conf); - auto io_manager = std::make_unique(*module.get(), conf); + if (runner_paths.pos_embedding_path.has_value()) { + if (io_manager->load_rotary_embedding(runner_paths.pos_embedding_path.value()) != + executorch::runtime::Error::Ok) { + ET_LOG(Error, "load pos_embedding file failed !"); + } + } auto eos_ids = std::make_unique>( get_eos_ids(tokenizer.get(), module.get())); @@ -575,7 +585,10 @@ std::unique_ptr create_llm_runner( // TODO(sicheon.oh) : Need to find proper flag to distinguish separaged prefill/decode opion. bool separated_prefill_decode = (conf.prefill_seq_len > 1); - ET_LOG(Info, "separated_prefill_decode(%s)", (separated_prefill_decode)?"True":"False"); + ET_LOG( + Info, + "separated_prefill_decode(%s)", + (separated_prefill_decode)?"True":"False"); // Create text_prefiller auto text_prefiller = std::make_unique( @@ -603,7 +616,7 @@ std::unique_ptr create_llm_runner( std::move(text_token_generator), std::move(stats), std::move(io_manager), - temperature); + runner_config.temperature); runner->load(); diff --git a/examples/samsung/executor_runner/llama/runner/runner.h b/examples/samsung/executor_runner/llama/runner/runner.h index a1a0d348a6d..adf444838c1 100644 --- a/examples/samsung/executor_runner/llama/runner/runner.h +++ b/examples/samsung/executor_runner/llama/runner/runner.h @@ -23,6 +23,26 @@ namespace example { namespace llm = ::executorch::extension::llm; +/** + * @brief Holds path-related configuration for creating a llama runner. + */ +struct RunnerPaths { + std::string model_path; + std::string tokenizer_path; + std::optional pos_embedding_path = std::nullopt; + std::optional data_path = std::nullopt; +}; + +/** + * @brief Holds non-path configuration for creating a llama runner. + */ +struct RunnerConfig { + std::string tokenizer_type; + bool decode_only = false; + bool use_gemma = false; + float temperature = -1.0f; +}; + class LLMRunner : public llm::IRunner { public: /** @@ -174,13 +194,7 @@ class LLMRunner : public llm::IRunner { }; std::unique_ptr create_llama_runner( - const std::string& model_path, - const std::string& tokenizer_path, - std::optional data_path = std::nullopt, - bool decode_only = false, - float temperature = -1.0f); - -std::unique_ptr load_llama_tokenizer( - const std::string& tokenizer_path); + const RunnerPaths& paths, + const RunnerConfig& config); } // namespace example diff --git a/examples/samsung/executor_runner/run_gemma3_sample.sh b/examples/samsung/executor_runner/run_gemma3_sample.sh new file mode 100644 index 00000000000..0de3a7199e4 --- /dev/null +++ b/examples/samsung/executor_runner/run_gemma3_sample.sh @@ -0,0 +1,29 @@ +# Copyright (c) 2026 Samsung Electronics Co. LTD +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +MODEL_PATH=/data/local/tmp/gemma3/gemma3.pte +WEIGHT_SHARED_MODEL_PATH=/data/local/tmp/gemma3/weights.ptd +TOKENIZER_PATH=/data/local/tmp/gemma3/tokenizer.json +POS_EMBEDDING_PATH=/data/local/tmp/gemma3/pos_embedding.bin + +SEQ_LEN=1024 +TEMPERATURE=0 +NUM_BOS=1 +INPUT_PROMPT="What is the capital of France?" + +chmod +x enn_executor_llama_runner + +./enn_executor_llama_runner \ + --model_path=$MODEL_PATH \ + --data_path=$WEIGHT_SHARED_MODEL_PATH \ + --seq_len=$SEQ_LEN \ + --temperature=$TEMPERATURE \ + --pos_embedding_path=$POS_EMBEDDING_PATH \ + --num_bos=$NUM_BOS \ + --tokenizer_type="hf" \ + --tokenizer_path=$TOKENIZER_PATH \ + --prompt="$INPUT_PROMPT" \ + --use_gemma diff --git a/examples/samsung/executor_runner/run_llama3_sample.sh b/examples/samsung/executor_runner/run_llama3_sample.sh new file mode 100644 index 00000000000..18f608e4f88 --- /dev/null +++ b/examples/samsung/executor_runner/run_llama3_sample.sh @@ -0,0 +1,27 @@ +# Copyright (c) 2026 Samsung Electronics Co. LTD +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +MODEL_PATH=/data/local/tmp/llama3/llama3.pte +WEIGHT_SHARED_MODEL_PATH=/data/local/tmp/llama3/weights.ptd +TOKENIZER_PATH=/data/local/tmp/llama3/tokenizer.model +POS_EMBEDDING_PATH=/data/local/tmp/llama3/pos_embedding.bin + +SEQ_LEN=1024 +TEMPERATURE=0 +INPUT_PROMPT="What is the capital of France?" + +chmod +x enn_executor_llama_runner + +./enn_executor_llama_runner \ + --model_path=$MODEL_PATH \ + --data_path=$WEIGHT_SHARED_MODEL_PATH \ + --seq_len=$SEQ_LEN \ + --temperature=$TEMPERATURE \ + --num_bos=1 \ + --pos_embedding_path=$POS_EMBEDDING_PATH \ + --tokenizer_type="tiktoken" \ + --tokenizer_path=$TOKENIZER_PATH \ + --prompt="$INPUT_PROMPT" diff --git a/extension/llm/runner/constants.h b/extension/llm/runner/constants.h index 60a6383a72e..b1887948cbf 100644 --- a/extension/llm/runner/constants.h +++ b/extension/llm/runner/constants.h @@ -24,6 +24,7 @@ inline constexpr auto kHeadDims = "get_head_dims"; inline constexpr auto kWindows = "get_windows"; inline constexpr auto kUseKVCache = "use_kv_cache"; inline constexpr auto kUseSDPAWithKVCache = "use_sdpa_with_kv_cache"; +inline constexpr auto kHeadDim = "get_head_dim"; // Multimodal method name conventions inline constexpr auto kVisionEncoderMethod = "vision_encoder"; diff --git a/extension/llm/runner/io_manager/io_manager.h b/extension/llm/runner/io_manager/io_manager.h index fc9a8f0641b..0a257064b77 100644 --- a/extension/llm/runner/io_manager/io_manager.h +++ b/extension/llm/runner/io_manager/io_manager.h @@ -37,6 +37,19 @@ class ET_EXPERIMENTAL IOManager { */ virtual ~IOManager() = default; + /** + * @brief Get the use_gemma configuration flag. + * @return True if using Gemma model specific processing, false by default. + * Derived classes should override this to return their specific value. + */ + virtual bool use_gemma() const { + return false; + } + + virtual bool is_quantized() const { + return false; + } + /** * @brief Load the IO manager with method metadata for prefill and * decode operations. diff --git a/extension/llm/runner/text_decoder_runner.cpp b/extension/llm/runner/text_decoder_runner.cpp index 6249719a7dd..893bbbd60bc 100644 --- a/extension/llm/runner/text_decoder_runner.cpp +++ b/extension/llm/runner/text_decoder_runner.cpp @@ -40,33 +40,18 @@ TextDecoderRunner::TextDecoderRunner( ::executorch::runtime::Result TextDecoderRunner::step( TensorPtr& tokens, int64_t start_pos) { - // ET_LOG(Info, "Input token %" PRIu64, input_token); - auto method_meta_result = module_->method_meta(method_name_); - if (!method_meta_result.ok()) { - return method_meta_result.error(); - } - auto method_meta = std::move(*method_meta_result); - + auto method_meta = ET_UNWRAP(module_->method_meta(method_name_.c_str())); // If only 1 input, we are not using kv cache bool use_kv_cache = method_meta.num_inputs() > 1; std::vector cache_positions; if (use_kv_cache) { - auto start_pos_tensor_result = populate_start_pos_or_cache_position( - module_, - start_pos, - cache_positions, - tokens->numel(), - method_name_.c_str()); - if (!start_pos_tensor_result.ok()) { - return start_pos_tensor_result.error(); - } - auto start_pos_tensor = std::move(*start_pos_tensor_result); + auto start_pos_tensor = ET_UNWRAP(populate_start_pos_or_cache_position( + module_, start_pos, cache_positions, tokens->numel(), method_name_.c_str())); std::vector inputs; - auto inputs_res = - io_manager_->prepare_decode(tokens, start_pos_tensor, method_name_); + auto inputs_res = io_manager_->prepare_decode(tokens, start_pos_tensor); ET_CHECK_OK_OR_RETURN_ERROR(inputs_res.error()); inputs = inputs_res.get(); if (stats_ != nullptr) { From ecb2752ddae3b779d360fb3d6b631c6e7d008af7 Mon Sep 17 00:00:00 2001 From: Jiseong Oh Date: Tue, 13 Oct 2026 19:37:43 +0900 Subject: [PATCH 09/17] Implement a unified LLM runner with a config parser Build the runner around a config parser and restructure Decoder and LlmIOManager around it. This also fixes a heap corruption in Decoder::register_method and lifetime issues in LlmIOManager, and renames the LLM executor source. Signed-off-by: Jiseong Oh Co-authored-by: brucekk.kim --- backends/samsung/build.sh | 1 - .../models/gemma3_1b/exynos_gemma3_1b.json | 21 + examples/models/llama/runner/CMakeLists.txt | 1 - .../models/llama3_1b/exynos_llama3_1b.json | 20 + examples/samsung/CMakeLists.txt | 25 +- examples/samsung/executor_runner/README.md | 27 + ...runner.cpp => enn_executor_llm_runner.cpp} | 93 +- .../executor_runner/llama/runner/constants.h | 20 + .../executor_runner/llama/runner/decoder.cpp | 287 ++++-- .../executor_runner/llama/runner/decoder.h | 121 +-- .../llama/runner/io_manager.cpp | 873 ++++++++++-------- .../executor_runner/llama/runner/io_manager.h | 262 +++--- .../llama/runner/prefiller.cpp | 30 +- .../executor_runner/llama/runner/prefiller.h | 30 +- .../llama/runner/rotary_embedding.cpp | 148 +-- .../llama/runner/rotary_embedding.h | 24 +- .../executor_runner/llama/runner/runner.cpp | 356 ++++--- .../executor_runner/llama/runner/runner.h | 61 +- .../llama/runner/token_generator.h | 25 +- .../executor_runner/run_gemma3_sample.sh | 10 +- .../executor_runner/run_llama3_sample.sh | 19 +- extension/llm/runner/io_manager/io_manager.h | 13 - extension/llm/runner/text_decoder_runner.cpp | 6 +- 23 files changed, 1427 insertions(+), 1046 deletions(-) create mode 100644 examples/models/gemma3_1b/exynos_gemma3_1b.json create mode 100644 examples/models/llama3_1b/exynos_llama3_1b.json create mode 100644 examples/samsung/executor_runner/README.md rename examples/samsung/executor_runner/{enn_executor_llama_runner.cpp => enn_executor_llm_runner.cpp} (60%) diff --git a/backends/samsung/build.sh b/backends/samsung/build.sh index 32b955ab765..a75c0c7bba3 100755 --- a/backends/samsung/build.sh +++ b/backends/samsung/build.sh @@ -108,7 +108,6 @@ function build_android() { -DEXECUTORCH_BUILD_DEVTOOLS=ON \ -DEXECUTORCH_ENABLE_EVENT_TRACER=ON \ -DBUILD_TESTING=OFF \ - -DPYTHON_EXECUTABLE=python \ -S ${PROJECT_DIR} \ -B ${ANDROID_BUILD_DIR} diff --git a/examples/models/gemma3_1b/exynos_gemma3_1b.json b/examples/models/gemma3_1b/exynos_gemma3_1b.json new file mode 100644 index 00000000000..554d03ec126 --- /dev/null +++ b/examples/models/gemma3_1b/exynos_gemma3_1b.json @@ -0,0 +1,21 @@ +{ + "model_id": "exynos_gemma3_1b", + "description": "Exynos Optimized Gemma3 1B Configuration", + "hidden_dim": 1152, + "n_heads": 4, + "n_kv_heads": 1, + "n_layers": 26, + "head_dim": 256, + "vocab_size": 262144, + "max_seq_len": 1024, + "max_context_len": 8192, + "embedding_strategy": "global_and_local", + "masking_strategy": "causal_local", + "model_version": "gemma3", + "tokenizer_type": "hf", + "prefill_seq_len": 128, + "decode_seq_len": 1, + "attn_mask_rows": 128, + "attn_mask_cols": 1024, + "sliding_window": 512 +} \ No newline at end of file diff --git a/examples/models/llama/runner/CMakeLists.txt b/examples/models/llama/runner/CMakeLists.txt index b459bf313d5..31bce16dd12 100644 --- a/examples/models/llama/runner/CMakeLists.txt +++ b/examples/models/llama/runner/CMakeLists.txt @@ -54,7 +54,6 @@ set(llama_runner_deps target_link_libraries(llama_runner PUBLIC ${llama_runner_deps}) target_link_libraries(llama_runner PUBLIC tokenizers::tokenizers) -target_link_libraries(llama_runner PUBLIC enn_backend) target_include_directories( llama_runner PUBLIC ${EXECUTORCH_ROOT}/extension/llm/tokenizers/include diff --git a/examples/models/llama3_1b/exynos_llama3_1b.json b/examples/models/llama3_1b/exynos_llama3_1b.json new file mode 100644 index 00000000000..c5050553a94 --- /dev/null +++ b/examples/models/llama3_1b/exynos_llama3_1b.json @@ -0,0 +1,20 @@ +{ + "model_id": "exynos_llama3_1b", + "description": "Exynos Optimized Llama3 1B Configuration", + "hidden_dim": 2048, + "n_heads": 32, + "n_kv_heads": 8, + "n_layers": 16, + "head_dim": 64, + "vocab_size": 128256, + "max_seq_len": 1024, + "max_context_len": 8192, + "embedding_strategy": "global_only", + "masking_strategy": "causal", + "model_version": "llama3", + "tokenizer_type": "tiktoken", + "prefill_seq_len": 128, + "decode_seq_len": 1, + "attn_mask_rows": 128, + "attn_mask_cols": 1024 +} \ No newline at end of file diff --git a/examples/samsung/CMakeLists.txt b/examples/samsung/CMakeLists.txt index 6182fb7bfc0..6aca9dfe853 100644 --- a/examples/samsung/CMakeLists.txt +++ b/examples/samsung/CMakeLists.txt @@ -145,26 +145,23 @@ set_target_properties( "${EXECUTORCH_ROOT}/build_samsung_android/lib/libcustom_ops.a" ) -set(__enn_executor_llama_runner__srcs ${_executor_runner__srcs}) -list(TRANSFORM __enn_executor_llama_runner__srcs - PREPEND "${EXECUTORCH_SOURCE_DIR}/" -) -list(FILTER __enn_executor_llama_runner__srcs EXCLUDE REGEX - ".*executor_runner.cpp$" -) -list(PREPEND __enn_executor_llama_runner_srcs - ${CMAKE_CURRENT_LIST_DIR}/executor_runner/enn_executor_llama_runner.cpp +set(enn_executor_llm_runner_srcs + ${CMAKE_CURRENT_LIST_DIR}/executor_runner/enn_executor_llm_runner.cpp ) -set(enn_executor_llama_runner_deps +set(enn_executor_llm_runner_deps executorch_core extension_data_loader extension_module extension_tensor extension_flat_tensor lib_extension_llm_runner ) -add_executable(enn_executor_llama_runner ${__enn_executor_llama_runner_srcs}) +add_executable(enn_executor_llm_runner ${enn_executor_llm_runner_srcs}) + +target_link_libraries( + enn_executor_llm_runner PRIVATE ${enn_executor_llm_runner_deps} +) target_include_directories( - enn_executor_llama_runner + enn_executor_llm_runner PUBLIC ${_common_include_directories} ${EXECUTORCH_ROOT}/backends/samsung/runtime ${EXECUTORCH_ROOT}/extension/llm/tokenizers/include @@ -195,13 +192,13 @@ list(APPEND _common_include_directories list(APPEND xnnpack_backend_libs quantized_kernels quantized_ops_lib) target_link_libraries( - enn_executor_llama_runner + enn_executor_llm_runner PUBLIC enn_llama_runner PRIVATE log gflags android lib_portable_ops - ${enn_executor_llama_runner_deps} + ${enn_executor_llm_runner_deps} ${xnnpack_backend_libs} sentencepiece re2 diff --git a/examples/samsung/executor_runner/README.md b/examples/samsung/executor_runner/README.md new file mode 100644 index 00000000000..8c41ae394ac --- /dev/null +++ b/examples/samsung/executor_runner/README.md @@ -0,0 +1,27 @@ +# Exynos Executor Runner Configuration + +This directory contains the executor runner for LLMs on Exynos platforms. The runner requires a JSON configuration file to correctly set up the model architecture and runtime parameters. + +## Using JSON Configuration Files + +The `enn_executor_llm_runner` binary accepts a path to a JSON configuration file via the `--config_path` argument. This file defines critical model dimensions and strategies used during execution. + +### Example Usage + +You can refer to `run_llama3_sample.sh` for a complete execution example. The key part of the command is: + +```bash +./enn_executor_llm_runner \ + --model_path="/data/local/tmp/llama3/llama3.pte" \ + --config_path="/data/local/tmp/exynos_llama3_1b.json" \ + --prompt="What is the capital of France?" \ + ... +``` + +### Where to find config files +Configuration files for supported models are located in `examples/models/`. +For example: +- Gemma 3 1B: `examples/models/gemma3_1b/exynos_gemma3_1b.json` +- Llama 3 1B: `examples/models/llama3_1b/exynos_llama3_1b.json` + +Ensure these files are pushed to the target device (e.g., to `/data/local/tmp/`) before running the script. \ No newline at end of file diff --git a/examples/samsung/executor_runner/enn_executor_llama_runner.cpp b/examples/samsung/executor_runner/enn_executor_llm_runner.cpp similarity index 60% rename from examples/samsung/executor_runner/enn_executor_llama_runner.cpp rename to examples/samsung/executor_runner/enn_executor_llm_runner.cpp index 3c496412cd6..e06cb40d5a6 100644 --- a/examples/samsung/executor_runner/enn_executor_llama_runner.cpp +++ b/examples/samsung/executor_runner/enn_executor_llm_runner.cpp @@ -1,5 +1,5 @@ /* - * Copyright (c) 2025 Samsung Electronics Co. LTD + * Copyright (c) 2026 Samsung Electronics Co. LTD * All rights reserved. * * This source code is licensed under the BSD-style license found in the @@ -24,21 +24,18 @@ DEFINE_string( DEFINE_string(data_path, "", "Data file for the model."); -DEFINE_string( - tokenizer_type, - "tiktoken", - "Tokenizer type. One of ['bpe', 'tiktoken']."); +DEFINE_string(pos_embedding_path, "", "rotary embedding stuff."); DEFINE_string(tokenizer_path, "tokenizer.bin", "Tokenizer stuff."); -DEFINE_string(pos_embedding_path, "", "rotary embedding stuff."); +DEFINE_string(config_path, "", "Path to the model configuration JSON file."); DEFINE_string(prompt, "The answer to the ultimate question is", "Prompt."); DEFINE_double( temperature, - 0.8f, - "Temperature; Default is 0.8f. 0 = greedy argmax sampling (deterministic). Lower temperature = more deterministic"); + 0.8, + "Temperature; Default is 0.8. 0 = greedy argmax sampling (deterministic). Lower temperature = more deterministic"); DEFINE_int32( seq_len, @@ -60,50 +57,41 @@ DEFINE_int32( 0, "Number of EOS tokens to append to the prompt. Defaults to 0. If > 0, the prompt will be appended with EOS tokens. This is useful for models that expect one or more EOS token at the end."); -DEFINE_bool(warmup, false, "Whether to run a warmup run."); +DEFINE_bool(is_warmup, false, "Whether to run a warmup run."); -DEFINE_bool(decode_only, false, "Whether to run with decode method only."); - -DEFINE_bool(use_gemma, false, "Use Gemma model specific processing."); +DEFINE_bool(is_decode_only, false, "Whether to run with decode method only."); int32_t main(int32_t argc, char** argv) { gflags::ParseCommandLineFlags(&argc, &argv, true); - - // Create a loader to get the data of the program file. There are other - // DataLoaders that use mmap() or point32_t to data that's already in memory, - // and users can create their own DataLoaders to load from arbitrary sources. const char* model_path = FLAGS_model_path.c_str(); - std::optional data_path = std::nullopt; if (!FLAGS_data_path.empty()) { data_path = FLAGS_data_path.c_str(); } - - const char* tokenizer_type = FLAGS_tokenizer_type.c_str(); - const char* tokenizer_path = FLAGS_tokenizer_path.c_str(); - std::optional pos_embedding_path = - FLAGS_pos_embedding_path.empty() ? std::nullopt - : std::make_optional(FLAGS_pos_embedding_path.c_str()); + FLAGS_pos_embedding_path.empty() + ? std::nullopt + : std::make_optional(FLAGS_pos_embedding_path.c_str()); + + const char* config_path = FLAGS_config_path.c_str(); - ET_LOG(Info, "prompt : %s", FLAGS_prompt.c_str()); const char* prompt = FLAGS_prompt.c_str(); - float temperature = FLAGS_temperature; + if (FLAGS_prompt.empty()) { + ET_LOG(Error, "Input prompt cannot be empty."); + return static_cast(executorch::runtime::Error::InvalidArgument); + } + ET_LOG(Info, "Input prompt message: %s", prompt); + float temperature = static_cast(FLAGS_temperature); int32_t seq_len = FLAGS_seq_len; - int32_t cpu_threads = FLAGS_cpu_threads; - bool warmup = FLAGS_warmup; - - bool decode_only = FLAGS_decode_only; - - bool use_gemma = FLAGS_use_gemma; + bool is_warmup = FLAGS_is_warmup; + bool is_decode_only = FLAGS_is_decode_only; int num_bos = FLAGS_num_bos; - int num_eos = FLAGS_num_eos; #if defined(ET_USE_THREADPOOL) @@ -118,45 +106,28 @@ int32_t main(int32_t argc, char** argv) { } #endif - // create llama runner example::RunnerPaths runner_paths{ .model_path = model_path, .tokenizer_path = tokenizer_path, .pos_embedding_path = pos_embedding_path, .data_path = data_path}; example::RunnerConfig runner_config{ - .tokenizer_type = tokenizer_type, - .decode_only = decode_only, - .use_gemma = use_gemma}; - std::unique_ptr runner = example::create_llama_runner(runner_paths, runner_config); - - if (runner == nullptr) { - ET_LOG(Error, "Failed to create llama runner"); - return 1; - } + .config_path = config_path, .is_decode_only = is_decode_only}; + std::unique_ptr runner = + example::create_llama_runner(runner_paths, runner_config); + ET_CHECK_MSG(runner != nullptr, "Failed to create llama runner"); - if (warmup) { - ET_LOG( - Info, - "Can not support warm-up (TODO : IRunner doesn't have warmup interface (pure virtual)"); -#if 0 - auto error = runner->warmup(prompt, /*max_new_tokens=*/seq_len); - if (error != executorch::runtime::Error::Ok) { - ET_LOG(Error, "Failed to warmup llama runner"); - return 1; - } -#endif - } + // TODO: support warmup interface - // generate executorch::extension::llm::GenerationConfig config{ - .seq_len = seq_len, .temperature = temperature, - .num_bos = num_bos, .num_eos = num_eos}; + .seq_len = seq_len, + .temperature = temperature, + .num_bos = num_bos, + .num_eos = num_eos}; auto error = runner->generate(prompt, config, {}, {}); - if (error != executorch::runtime::Error::Ok) { - ET_LOG(Error, "Failed to generate token from llama runner"); - return 1; - } + ET_CHECK_MSG( + error == executorch::runtime::Error::Ok, + "Failed to generate token from llama runner"); return 0; } diff --git a/examples/samsung/executor_runner/llama/runner/constants.h b/examples/samsung/executor_runner/llama/runner/constants.h index 1ddb1c2ebb0..68a180943f5 100644 --- a/examples/samsung/executor_runner/llama/runner/constants.h +++ b/examples/samsung/executor_runner/llama/runner/constants.h @@ -7,6 +7,8 @@ */ #pragma once +#include + // constants for LLM runtime namespace example { @@ -20,4 +22,22 @@ inline constexpr auto kNumOfKVHeads = "get_n_kv_heads"; inline constexpr auto default_method_name = "forward"; inline constexpr auto default_method_name_prefill = "prefill_forward"; inline constexpr auto default_method_name_decode = "decode_forward"; + +enum class DecoderModelVersion { + kLlama3, + kGemma3, + kGeneral, +}; + +enum class EmbeddingStrategy { + kGlobalOnly, + kGlobalAndLocal, +}; + +enum class MaskingStrategy { + kCausal, + kCausalLocal, + kHybridConvAttention, // LIV-conv + GQA mask constraint +}; + } // namespace example diff --git a/examples/samsung/executor_runner/llama/runner/decoder.cpp b/examples/samsung/executor_runner/llama/runner/decoder.cpp index e25c894b47a..0dc1ce2a826 100644 --- a/examples/samsung/executor_runner/llama/runner/decoder.cpp +++ b/examples/samsung/executor_runner/llama/runner/decoder.cpp @@ -8,32 +8,31 @@ // Given inputs, run a text decoder and return logits. -#include #include +#include #include +#include #include #include -#include namespace example { +namespace runtime = executorch::runtime; + class IExecutionIOBinder { virtual ~IExecutionIOBinder() = default; virtual void beforeExecution() = 0; virtual void afterExecution() = 0; }; -/** - * @brief Slices the k-th token's logits by calculating the memory offset. - * Corrected to return Result by dereferencing the TensorPtr. - */ -executorch::runtime::Result slice_kth_token( - const executorch::runtime::EValue& output_evalue, - int64_t k) { - ET_CHECK_MSG(output_evalue.isTensor(), "The provided EValue is not a tensor."); - executorch::aten::Tensor logits = output_evalue.toTensor(); +runtime::Result slice_kth_token( + const runtime::EValue& output_evalue, + int64_t k) { + ET_CHECK_MSG( + output_evalue.isTensor(), "The provided EValue is not a tensor."); + aten::Tensor logits = output_evalue.toTensor(); constexpr int64_t kExpectedNumDim = 3; constexpr int64_t kSequenceDim = 1; @@ -43,100 +42,190 @@ executorch::runtime::Result slice_kth_token( const auto strides = logits.strides(); const auto sizes = logits.sizes(); - ET_CHECK_MSG(k >= 0 && k < max_seq_len, - "Index k (%ld) out of bounds (max: %ld)", k, max_seq_len); + ET_CHECK_MSG( + k >= 0 && k < max_seq_len, + "Index k (%ld) out of bounds (max: %ld)", + k, + max_seq_len); - ET_CHECK_MSG(actual_dim == kExpectedNumDim, - "num dim shape of logit(%ld) is not expected(%ld)", actual_dim, kExpectedNumDim); + ET_CHECK_MSG( + actual_dim == kExpectedNumDim, + "num dim shape of logit(%ld) is not expected(%ld)", + actual_dim, + kExpectedNumDim); // Calculate Memory Offset (zero-copy) int64_t offset = k * strides[kSequenceDim]; // Define New Shape for the k-th token (1, 1, VocabSize) // Using a fixed-size array or vector that matches SizesType - std::vector new_sizes = { - static_cast(sizes[0]), + std::vector new_sizes = { + static_cast(sizes[0]), 1, // Target single token - static_cast(sizes[2]) - }; + static_cast(sizes[2])}; // Create TensorPtr and dereference it to return Result - auto tensor_ptr = executorch::extension::from_blob( + auto tensor_ptr = extension::from_blob( logits.mutable_data_ptr() + offset, new_sizes, - logits.scalar_type() - ); + logits.scalar_type()); // Dereference the shared_ptr to get the Tensor object return *tensor_ptr; } -::executorch::runtime::Result Decoder::step( - executorch::extension::TensorPtr& tokens, +runtime::Error Decoder::load() { + extract_llm_methods(); + + if (method_names_.size() > 1) { + auto ret_shared = create_shared_planned_memory(); + ET_CHECK_OR_RETURN_ERROR( + ret_shared == runtime::Error::Ok, + InvalidArgument, + "Failed to create share planned memory"); + } + + for (const auto& method : method_names_) { + if (method.empty() || module_->is_method_loaded(method)) { + continue; + } + + runtime::Error ret_load; + if (shared_method_memory_) { + ret_load = + module_->load_method(method, shared_method_memory_->allocator.get()); + } else { + ret_load = module_->load_method(method); + } + ET_CHECK_OR_RETURN_ERROR( + ret_load == runtime::Error::Ok, + InvalidArgument, + "Failed to load priority method %s", + method.c_str()); + ET_LOG(Info, "Loading Priority Method (%s)", method.c_str()); + } + + // Load remaining utility methods (e.g., get_genai_mode) + for (const auto& method : method_names_) { + if (module_->is_method_loaded(method)) { + continue; + } + + runtime::Error ret_load = module_->load_method(method); + ET_CHECK_OR_RETURN_ERROR( + ret_load == runtime::Error::Ok, + InvalidArgument, + "Failed to load utility method %s", + method.c_str()); + ET_LOG(Info, "Loading Utility Method (%s)", method.c_str()); + } + + auto ret = load_io_manager(); + ET_CHECK_OR_RETURN_ERROR( + ret == runtime::Error::Ok, InvalidArgument, "Failed to load io manager"); + + return ret; +} + +runtime::Result Decoder::step( + extension::TensorPtr& tokens, int64_t start_pos) { std::string method_stage; if (tokens->numel() > 1) { - method_stage = - methods_.name[static_cast(DecoderMode::PREFILL)]; + method_stage = prefill_method_; } else { - method_stage = - methods_.name[static_cast(DecoderMode::DECODE)]; + method_stage = decode_method_; } - // ET_LOG(Info, "Input token %" PRIu64, input_token); auto method_meta = ET_UNWRAP(module_->method_meta(method_stage.c_str())); - // If only 1 input, we are not using kv cache - bool use_kv_cache = method_meta.num_inputs() > 1; + const size_t total_input_count = method_meta.num_inputs(); + bool use_kv_cache = + total_input_count > 1; // kv cache is not used, when input is 1 if (use_kv_cache) { - // Size of the second argument. This could be either input_pos or + // Size of the position tracking argument. This could be either input_pos or // cache_positions + constexpr int64_t kExpectedPositionDimensions = 3; + constexpr int64_t kPositionTensorIndex = 3; + constexpr int64_t kAttentionSetupTensorIndex = 1; - const bool use_gemma = io_manager_->use_gemma(); - size_t pos_shape_size = 3; - size_t pos_idx = 3; - size_t attention_idx = 1; + constexpr int64_t kKVCacheExpectedDimensions = 4; + constexpr int64_t kHeadsDimIndex = 1; + constexpr int64_t kHeadSizeDimIndex = 3; // Check if we are using cache positions instead of input pos. - auto pos_input_info = ET_UNWRAP(method_meta.input_tensor_meta(pos_idx)); - auto sizes = pos_input_info.sizes(); + auto position_tensor_meta = + ET_UNWRAP(method_meta.input_tensor_meta(kPositionTensorIndex)); + auto position_shape = position_tensor_meta.sizes(); ET_CHECK_OR_RETURN_ERROR( - sizes.size() == pos_shape_size, + position_shape.size() == kExpectedPositionDimensions, InvalidProgram, "The pos input tensor is not expected. Got dimension (%zu)", - sizes.size()); - auto numel = sizes[0]; - std::vector<::executorch::aten::SizesType> sizes_vec = {numel}; - - auto n_inputs = method_meta.num_inputs(); - auto input_info = ET_UNWRAP(method_meta.input_tensor_meta(attention_idx)); - sizes = input_info.sizes(); - // llama mask tensor shape : 2dim - // gemma mask tensor index : 4dim - // llama kvcache tensor index : 3~ - // gemma mask tensor index : 4~ - size_t kv_cache_idx; - if (sizes.size() != 4) ET_LOG(Error, "input %zu of %zu is invalid dimension : %zu", attention_idx, n_inputs, sizes.size()); - kv_cache_idx = attention_idx + (use_gemma ? 6 : 3); - input_info = ET_UNWRAP(method_meta.input_tensor_meta(kv_cache_idx)); - sizes = input_info.sizes(); - if (sizes.size() != 4) ET_LOG(Error, "input %zu of %zu is invalid dimension : %zu", kv_cache_idx, n_inputs, sizes.size()); - - executorch::extension::TensorPtr start_pos_tensor; + position_shape.size()); + std::vector sizes_vec = {position_shape[0]}; + + auto attention_setup_tensor_meta = + ET_UNWRAP(method_meta.input_tensor_meta(kAttentionSetupTensorIndex)); + auto attention_setup_shape = attention_setup_tensor_meta.sizes(); + auto* llm_io_manager = static_cast(io_manager_); + const auto& config = llm_io_manager->get_config(); + const size_t kv_cache_tensor_index = kAttentionSetupTensorIndex + + llm_io_manager->get_attention_setup_tensor_count(); + + auto kv_cache_tensor_meta = + ET_UNWRAP(method_meta.input_tensor_meta(kv_cache_tensor_index)); + auto kv_cache_shape = kv_cache_tensor_meta.sizes(); + ET_CHECK_OR_RETURN_ERROR( + static_cast(kv_cache_shape.size()) == + kKVCacheExpectedDimensions, + InvalidArgument, + "Input tensor at index %" PRId64 + " (KV cache slot) has invalid dimensionality: %" PRId64 + " (expected %" PRId64 ")", + kv_cache_tensor_index, + static_cast(kv_cache_shape.size()), + kKVCacheExpectedDimensions); + + ET_CHECK_OR_RETURN_ERROR( + static_cast(kv_cache_shape[kHeadsDimIndex]) == + config.n_kv_heads, + InvalidArgument, + "KV cache tensor head count mismatch at index %" PRId64 + ": found %" PRId64 " (expected %" PRId64 " via configuration)", + kv_cache_tensor_index, + static_cast(kv_cache_shape[kHeadsDimIndex]), + config.n_kv_heads); + + ET_CHECK_OR_RETURN_ERROR( + static_cast(kv_cache_shape[kHeadSizeDimIndex]) == + config.head_dim, + InvalidArgument, + "KV cache tensor head dimension size mismatch at index %" PRId64 + ": found %" PRId64 " (expected %" PRId64 " via configuration)", + kv_cache_tensor_index, + static_cast(kv_cache_shape[kHeadSizeDimIndex]), + config.head_dim); + + extension::TensorPtr start_pos_tensor; // Assuming model is exported with input_pos, create a tensor with size 1 - start_pos_tensor = executorch::extension::from_blob( - &start_pos, sizes_vec, ::executorch::aten::ScalarType::Long); + start_pos_tensor = + extension::from_blob(&start_pos, sizes_vec, aten::ScalarType::Long); - std::vector inputs; + std::vector inputs; - if (method_stage == methods_.name[static_cast(DecoderMode::PREFILL)]) { + if (method_stage == prefill_method_) { auto inputs_res = io_manager_->prepare_prefill(tokens, start_pos_tensor); ET_CHECK_OK_OR_RETURN_ERROR(inputs_res.error()); inputs = inputs_res.get(); - } else if (method_stage == methods_.name[static_cast(DecoderMode::DECODE)]) { + } else if (method_stage == decode_method_) { auto inputs_res = io_manager_->prepare_decode(tokens, start_pos_tensor); ET_CHECK_OK_OR_RETURN_ERROR(inputs_res.error()); inputs = inputs_res.get(); + } else { + ET_LOG( + Debug, + "Skipping IO prepare for non-execution method stage: %s", + method_stage.c_str()); } auto outputs_res = module_->execute(method_stage.c_str(), inputs); @@ -145,18 +234,20 @@ ::executorch::runtime::Result Decoder::step( outputs_res.get()[0].isTensor(), "Non Tensor Output returned from executing LLM"); - if (method_stage == methods_.name[static_cast(DecoderMode::PREFILL)]) { - if (io_manager_->update_prefill(outputs_res.get()) != - executorch::runtime::Error::Ok) { - ET_LOG(Error, "Failed to updated prefill output by io_manager"); - } - } else if (method_stage == methods_.name[static_cast(DecoderMode::DECODE)]) { - if (io_manager_->update_decode(outputs_res.get()) != - executorch::runtime::Error::Ok) { - ET_LOG(Error, "Failed to updated decode output by io_manager"); - } + if (method_stage == prefill_method_) { + ET_CHECK_OK_OR_RETURN_ERROR( + io_manager_->update_prefill(outputs_res.get()), + "Failed to update prefill output by io_manager"); + } else if (method_stage == decode_method_) { + ET_CHECK_OK_OR_RETURN_ERROR( + io_manager_->update_decode(outputs_res.get()), + "Failed to update decode output by io_manager"); + } else { + ET_LOG( + Debug, + "Skipping IO update for non-execution method stage: %s", + method_stage.c_str()); } - // Return the logits tensor return slice_kth_token(outputs_res.get()[0], (tokens->numel() - 1)); } else { // no kv cache @@ -177,21 +268,46 @@ ::executorch::runtime::Result Decoder::step( } } -::executorch::runtime::Error Decoder::create_shared_planned_memory() { - executorch::runtime::Error ret = executorch::runtime::Error::Ok; +void Decoder::extract_llm_methods() { + if (method_names_.size() > 1) { + for (auto it = method_names_.begin(); it != method_names_.end(); ++it) { + if (it->find("prefill") != std::string::npos) { + prefill_method_ = *it; + } else if (it->find("decode") != std::string::npos) { + decode_method_ = *it; + } else { + ET_LOG( + Info, + "skipping io manager identify with %s method name.", + (*it).c_str()); + continue; + } + } + } else if (method_names_.size() == 1) { + const std::string& method = *method_names_.begin(); + prefill_method_ = method; + decode_method_ = method; + } +} + +runtime::Error Decoder::load_io_manager() { + if (method_names_.empty()) { + return runtime::Error::InvalidArgument; + } + return io_manager_->load(prefill_method_, decode_method_); +} +runtime::Error Decoder::create_shared_planned_memory() { if (shared_method_memory_ != nullptr) { ET_LOG(Info, "Shared Planned memory is already created"); - return ret; + return runtime::Error::Ok; } - auto prefill_meta = module_->method_meta( - methods_.name[static_cast(DecoderMode::PREFILL)].c_str()); - auto decode_meta = module_->method_meta( - methods_.name[static_cast(DecoderMode::DECODE)].c_str()); + auto prefill_meta = module_->method_meta(prefill_method_); + auto decode_meta = module_->method_meta(decode_method_); if (!prefill_meta.ok() || !decode_meta.ok()) { ET_LOG(Error, "Failed to get method metadata for shared memory setup"); - return executorch::runtime::Error::InvalidArgument; + return runtime::Error::InvalidArgument; } size_t prefill_bufs = prefill_meta->num_memory_planned_buffers(); @@ -230,11 +346,12 @@ ::executorch::runtime::Error Decoder::create_shared_planned_memory() { } shared_method_memory_->allocator = - std::make_unique<::executorch::runtime::HierarchicalAllocator>( - ::executorch::runtime::Span<::executorch::runtime::Span>( + std::make_unique( + runtime::Span>( shared_method_memory_->spans.data(), shared_method_memory_->spans.size())); - return ret; + + return runtime::Error::Ok; } } // namespace example diff --git a/examples/samsung/executor_runner/llama/runner/decoder.h b/examples/samsung/executor_runner/llama/runner/decoder.h index 353863edc94..688b27b3c2b 100644 --- a/examples/samsung/executor_runner/llama/runner/decoder.h +++ b/examples/samsung/executor_runner/llama/runner/decoder.h @@ -10,6 +10,8 @@ #pragma once +#include + #include #include #include @@ -20,121 +22,40 @@ namespace example { namespace llm = ::executorch::extension::llm; +namespace runtime = ::executorch::runtime; +namespace aten = ::executorch::aten; +namespace extension = ::executorch::extension; -class ET_EXPERIMENTAL Decoder - : public executorch::extension::llm::TextDecoderRunner { +class ET_EXPERIMENTAL Decoder : public llm::TextDecoderRunner { public: struct SharedMethodMemory { std::vector> buffers; - std::vector> spans; - std::unique_ptr allocator; + std::vector> spans; + std::unique_ptr allocator; }; - explicit Decoder( - executorch::extension::Module* module, - llm::IOManager* io_manager) + explicit Decoder(extension::Module* module, llm::IOManager* io_manager) : TextDecoderRunner(module, io_manager), shared_method_memory_(nullptr) {} - virtual ::executorch::runtime::Result step( - executorch::extension::TensorPtr& input, + runtime::Result step( + extension::TensorPtr& input, int64_t start_pos) override; - inline ::executorch::runtime::Error load() override { - executorch::runtime::Error ret = executorch::runtime::Error::Ok; - - if (methods_.num > 1) { - ret = create_shared_planned_memory(); - if (ret != executorch::runtime::Error::Ok) { - ET_LOG(Error, "Failed to create share planned memory"); - } - } - - /** - * Dealing with multiple method. - */ - for (auto method : methods_.name) { - if (module_->is_method_loaded(method)) { - ET_LOG(Info, "Method \"%s\" is already loaded", method.c_str()); - continue; - } - - if (shared_method_memory_) { - ret = module_->load_method( - method, shared_method_memory_->allocator.get()); - if (ret != executorch::runtime::Error::Ok) { - ET_LOG(Error, "Failed to load method"); - return ret; - } - } - - ET_LOG(Info, "Loading Method (%s)", method.c_str()); - } - - ret = load_io_manager(); - if (ret != executorch::runtime::Error::Ok) { - ET_LOG(Error, "Failed to load io manager (%d)", ret); - return ret; - } + runtime::Error load() override; - return ret; - } - - inline bool is_method_loaded() override { - bool ret = false; - for (auto method : methods_.name) { - ret = module_->is_method_loaded(method); - if (ret == false) { - ET_LOG(Info, "method name \"%s\" is not loaded", method.c_str()); - return ret; - } - } - - return ret; - } - - inline void register_method(std::string& method) { - methods_.name[methods_.num++] = method; - - ET_LOG(Info, "method[%d] : %s", methods_.num - 1, method.c_str()); - - for (int i = 0; i < methods_.num; i++) { - ET_LOG( - Info, "Registered method[%d] : \"%s\"", i, methods_.name[i].c_str()); - } + inline void set_method(std::string& method) { + method_names_.insert(method); } private: - inline ::executorch::runtime::Error load_io_manager() { - executorch::runtime::Error ret = executorch::runtime::Error::Ok; - - if (methods_.num > 1) { - ret = io_manager_->load( - methods_.name[static_cast(DecoderMode::PREFILL)], - methods_.name[static_cast(DecoderMode::DECODE)]); - } else { - ret = io_manager_->load( - methods_.name[static_cast(DecoderMode::DECODE)], - methods_.name[static_cast(DecoderMode::DECODE)]); - } - return ret; - } - - ::executorch::runtime::Error create_shared_planned_memory(); - - enum class DecoderMode : int { PREFILL = 1, DECODE = 0 }; - - struct RegisterdMethod { - uint32_t num = 0; - std::vector name = {default_method_name, default_method_name}; - std::vector objs = { - nullptr, - nullptr}; - }; - - private: - RegisterdMethod methods_; - + std::string prefill_method_; + std::string decode_method_; + std::unordered_set method_names_; std::unique_ptr shared_method_memory_; + + void extract_llm_methods(); + runtime::Error load_io_manager(); + runtime::Error create_shared_planned_memory(); }; } // namespace example diff --git a/examples/samsung/executor_runner/llama/runner/io_manager.cpp b/examples/samsung/executor_runner/llama/runner/io_manager.cpp index 35fa5a6863e..19251a79c5d 100644 --- a/examples/samsung/executor_runner/llama/runner/io_manager.cpp +++ b/examples/samsung/executor_runner/llama/runner/io_manager.cpp @@ -13,13 +13,15 @@ namespace example { +namespace llm = ::executorch::extension::llm; +namespace runtime = ::executorch::runtime; +namespace aten = ::executorch::aten; +namespace extension = ::executorch::extension; + namespace { // Helper to fill a float region with a value, using memset when possible. -inline void fill_float_region( - float* begin, - float* end, - float value) { +inline void fill_float_region(float* begin, float* end, float value) { if (begin >= end) { return; } @@ -30,462 +32,539 @@ inline void fill_float_region( } } -inline void fill_float_region( - float* base, - size_t offset, - size_t len, - float value) { +inline void +fill_float_region(float* base, size_t offset, size_t len, float value) { fill_float_region(base + offset, base + offset + len, value); } } // namespace - executorch::runtime::Error LlmIOManager::load_rotary_embedding(std::optional pos_embedding_path) { - if (pos_embedding_path.has_value()) { - return rotary_embedding_->load(pos_embedding_path.value()); - } - return executorch::runtime::Error::Ok; + +LlmIOManager::LlmIOManager(extension::Module& module, Config config) + : IOManager(module), config_(config) { + ET_LOG( + Info, + "LlmIOManager Constructor: dim(%zu), vocab(%zu), max_seq(%zu)", + config_.hidden_dim, + config_.vocab_size, + config_.max_seq_len); + + rotary_embedding_ = std::make_unique(RotaryEmbedding::Config{ + config_.head_dim, config_.max_seq_len, config_.embedding_strategy}); + + ET_CHECK_MSG( + config_.max_seq_len >= config_.prefill_seq_len, + "Inconsistent configuration: max_seq_len (%zu) must be greater than or equal to prefill_seq_len (%zu)", + config_.max_seq_len, + config_.prefill_seq_len); + + ET_CHECK_MSG( + config_.max_seq_len >= 1, + "Inconsistent configuration: max_seq_len must be at least 1, but got %zu", + config_.max_seq_len); + + prefill_cache_len_ = config_.max_seq_len - config_.prefill_seq_len; + decode_cache_len_ = config_.max_seq_len - 1; + + allocate_buffers(); +} + +executorch::runtime::Error LlmIOManager::load_rotary_embedding( + std::optional pos_embedding_path) { + if (pos_embedding_path.has_value()) { + return rotary_embedding_->load(pos_embedding_path.value()); } + return executorch::runtime::Error::Ok; +} - executorch::runtime::Result> - LlmIOManager::prepare_prefill( - const executorch::extension::TensorPtr& input, - const executorch::extension::TensorPtr& start_pos, - const std::string& prefill_method) { - (void)prefill_method; - std::vector inputs; - inputs.reserve(get_input_reserve_size()); - - //re-producessing input token tensor - const int64_t* prompt_tokens = input->mutable_data_ptr(); - size_t prompt_token_len = input->numel(); - input_tokens_.resize(prompt_token_len); - std::copy(prompt_tokens, prompt_tokens + prompt_token_len, input_tokens_.begin()); - input_tokens_.resize(config_.prefill_seq_len, 0); - - auto input_tokens_tensor = executorch::extension::from_blob( - input_tokens_.data(), {1, static_cast(config_.prefill_seq_len)}, +executorch::runtime::Result> +LlmIOManager::prepare_prefill( + const executorch::extension::TensorPtr& input, + const executorch::extension::TensorPtr& start_pos, + const std::string& prefill_method) { + (void)prefill_method; + std::vector inputs; + inputs.reserve(get_input_reserve_size()); + + // re-producessing input token tensor + const int64_t* prompt_tokens = input->mutable_data_ptr(); + size_t prompt_token_len = input->numel(); + input_tokens_.resize(prompt_token_len); + std::copy( + prompt_tokens, prompt_tokens + prompt_token_len, input_tokens_.begin()); + input_tokens_.resize(config_.prefill_seq_len, 0); + + input_tokens_tensor_ = executorch::extension::from_blob( + input_tokens_.data(), + {1, static_cast(config_.prefill_seq_len)}, executorch::aten::ScalarType::Long); - inputs.emplace_back(input_tokens_tensor); - - //re-producessing start_pos tensor - const int64_t* start_position = start_pos->mutable_data_ptr(); - pos_ids_.resize(config_.prefill_seq_len); - current_pos_ = *start_position; - for(int64_t i = 0; i < config_.prefill_seq_len; ++i) { - pos_ids_[i] = current_pos_ + i; - } + inputs.emplace_back(input_tokens_tensor_); + + // re-producessing start_pos tensor + const int64_t* start_position = start_pos->mutable_data_ptr(); + pos_ids_.resize(config_.prefill_seq_len); + current_pos_ = *start_position; + for (int64_t i = 0; i < config_.prefill_seq_len; ++i) { + pos_ids_[i] = current_pos_ + i; + } - //attention tensor - update_attn_mask(current_pos_, prompt_token_len); - attn_mask_tensor_ = executorch::extension::from_blob( + // attention tensor + update_attn_mask(current_pos_, prompt_token_len); + attn_mask_tensor_ = executorch::extension::from_blob( attn_mask_.data(), - {1, 1, static_cast(config_.attn_mask_rows), - static_cast(config_.attn_mask_cols)}); - - inputs.emplace_back(attn_mask_tensor_); - - if (config_.use_gemma) { + {1, + 1, + static_cast(config_.attn_mask_rows), + static_cast(config_.attn_mask_cols)}); + + inputs.emplace_back(attn_mask_tensor_); + + switch (config_.masking_strategy) { + case MaskingStrategy::kCausal: + break; + case MaskingStrategy::kCausalLocal: + case MaskingStrategy::kHybridConvAttention: update_local_attn_mask(current_pos_, prompt_token_len); local_attn_mask_tensor_ = executorch::extension::from_blob( local_attn_mask_.data(), - {1, 1, static_cast(config_.attn_mask_rows), - static_cast(config_.attn_mask_cols)}); + {1, + 1, + static_cast(config_.attn_mask_rows), + static_cast(config_.attn_mask_cols)}); inputs.emplace_back(local_attn_mask_tensor_); - } + break; + } - prepare_rotary_embedding_inputs(config_.prefill_seq_len, inputs); - - prefill_k_input_tensors_.clear(); - prefill_v_input_tensors_.clear(); - prefill_k_output_tensors_.clear(); - prefill_v_output_tensors_.clear(); - append_kv_input_tensors( - prefill_k_input_, - prefill_v_input_, - prefill_k_input_tensors_, - prefill_v_input_tensors_, - executorch::aten::ScalarType::Float, - prefill_cache_len_, - inputs); - dump_kv_input("prefill"); - - return inputs; + prepare_rotary_embedding_inputs(config_.prefill_seq_len, inputs); + + prefill_k_input_tensors_.clear(); + prefill_v_input_tensors_.clear(); + prefill_k_output_tensors_.clear(); + prefill_v_output_tensors_.clear(); + append_kv_input_tensors( + prefill_k_input_, + prefill_v_input_, + prefill_k_input_tensors_, + prefill_v_input_tensors_, + executorch::aten::ScalarType::Float, + prefill_cache_len_, + inputs); + dump_kv_input("prefill"); + + return inputs; +} +executorch::runtime::Result> +LlmIOManager::prepare_decode( + const executorch::extension::TensorPtr& input, + const executorch::extension::TensorPtr& start_pos, + const std::string& decode_method) { + (void)decode_method; + if (!prefill_done_) { + copy_prefill_to_decode_input(); + prefill_done_ = true; } - executorch::runtime::Result> - LlmIOManager::prepare_decode( - const executorch::extension::TensorPtr& input, - const executorch::extension::TensorPtr& start_pos, - const std::string& decode_method) { - (void)decode_method; - if (!prefill_done_) { - copy_prefill_to_decode_input(); - prefill_done_ = true; - } - const int64_t* start_position = start_pos->mutable_data_ptr(); - current_pos_ = *start_position; - update_attn_mask(current_pos_, config_.decode_seq_len); - std::vector inputs; - inputs.reserve(get_input_reserve_size()); - inputs.emplace_back(input); - - attn_mask_tensor_ = executorch::extension::from_blob( + const int64_t* start_position = start_pos->mutable_data_ptr(); + current_pos_ = *start_position; + update_attn_mask(current_pos_, config_.decode_seq_len); + std::vector inputs; + inputs.reserve(get_input_reserve_size()); + inputs.emplace_back(input); + + attn_mask_tensor_ = executorch::extension::from_blob( attn_mask_.data(), - {1,1,1, - static_cast(config_.attn_mask_cols)}); - inputs.emplace_back(attn_mask_tensor_); - if (config_.use_gemma) { + {1, + 1, + 1, + static_cast(config_.attn_mask_cols)}); + inputs.emplace_back(attn_mask_tensor_); + switch (config_.masking_strategy) { + case MaskingStrategy::kCausal: + break; + case MaskingStrategy::kCausalLocal: + case MaskingStrategy::kHybridConvAttention: update_local_attn_mask(current_pos_, 1); local_attn_mask_tensor_ = executorch::extension::from_blob( - local_attn_mask_.data(), - {1,1,1, - static_cast(config_.attn_mask_cols)});; + local_attn_mask_.data(), + {1, + 1, + 1, + static_cast(config_.attn_mask_cols)}); inputs.emplace_back(local_attn_mask_tensor_); - } - - pos_ids_.clear(); - pos_ids_.resize(config_.decode_seq_len); - pos_ids_[0] = current_pos_; - prepare_rotary_embedding_inputs(config_.decode_seq_len, inputs); - - decode_k_input_tensors_.clear(); - decode_v_input_tensors_.clear(); - decode_k_output_tensors_.clear(); - decode_v_output_tensors_.clear(); - append_kv_input_tensors( - decode_k_input_, - decode_v_input_, - decode_k_input_tensors_, - decode_v_input_tensors_, - executorch::aten::ScalarType::Float, - decode_cache_len_, - inputs); - dump_kv_input("decode"); - - return inputs; - } - executorch::runtime::Error LlmIOManager::update_prefill( - const std::vector& model_outputs, - const std::string& prefill_method) { - (void)prefill_method; - update_kv_cache( - model_outputs, - prefill_k_input_, - prefill_v_input_, - prefill_cache_len_, - config_.prefill_seq_len, - "prefill"); - return executorch::runtime::Error::Ok; + break; } - executorch::runtime::Error LlmIOManager::update_decode( - const std::vector& model_outputs, - const std::string& decode_method) { - (void)decode_method; - update_kv_cache( - model_outputs, - decode_k_input_, - decode_v_input_, - decode_cache_len_, - 1, - "decode"); - return executorch::runtime::Error::Ok; + + pos_ids_.clear(); + pos_ids_.resize(config_.decode_seq_len); + pos_ids_[0] = current_pos_; + prepare_rotary_embedding_inputs(config_.decode_seq_len, inputs); + + decode_k_input_tensors_.clear(); + decode_v_input_tensors_.clear(); + decode_k_output_tensors_.clear(); + decode_v_output_tensors_.clear(); + append_kv_input_tensors( + decode_k_input_, + decode_v_input_, + decode_k_input_tensors_, + decode_v_input_tensors_, + executorch::aten::ScalarType::Float, + decode_cache_len_, + inputs); + dump_kv_input("decode"); + + return inputs; +} +executorch::runtime::Error LlmIOManager::update_prefill( + const std::vector& model_outputs, + const std::string& prefill_method) { + (void)prefill_method; + update_kv_cache( + model_outputs, + prefill_k_input_, + prefill_v_input_, + prefill_cache_len_, + config_.prefill_seq_len, + "prefill"); + return executorch::runtime::Error::Ok; +} +executorch::runtime::Error LlmIOManager::update_decode( + const std::vector& model_outputs, + const std::string& decode_method) { + (void)decode_method; + update_kv_cache( + model_outputs, + decode_k_input_, + decode_v_input_, + decode_cache_len_, + 1, + "decode"); + return executorch::runtime::Error::Ok; +} +// TOOD: sizeof(float) dtype hard-codes the buffer dtype to be float. +void LlmIOManager::allocate_buffers() { + size_t prefill_in_size = config_.n_kv_heads * prefill_cache_len_ * + config_.head_dim * sizeof(float); + size_t decode_in_size = + config_.n_kv_heads * decode_cache_len_ * config_.head_dim * sizeof(float); + prefill_k_input_.resize(config_.n_layers); + prefill_v_input_.resize(config_.n_layers); + decode_k_input_.resize(config_.n_layers); + decode_v_input_.resize(config_.n_layers); + for (size_t l = 0; l < config_.n_layers; l++) { + prefill_k_input_[l].resize(prefill_in_size, 0); + prefill_v_input_[l].resize(prefill_in_size, 0); + decode_k_input_[l].resize(decode_in_size, 0); + decode_v_input_[l].resize(decode_in_size, 0); } - void LlmIOManager::allocate_buffers() { - size_t prefill_in_size = config_.n_kv_heads * prefill_cache_len_ * - config_.head_dim * sizeof(float); - size_t decode_in_size = config_.n_kv_heads * decode_cache_len_ * - config_.head_dim * sizeof(float); - prefill_k_input_.resize(config_.n_layers); - prefill_v_input_.resize(config_.n_layers); - decode_k_input_.resize(config_.n_layers); - decode_v_input_.resize(config_.n_layers); - for (size_t l = 0; l < config_.n_layers; l++) { - prefill_k_input_[l].resize(prefill_in_size, 0); - prefill_v_input_[l].resize(prefill_in_size, 0); - decode_k_input_[l].resize(decode_in_size, 0); - decode_v_input_[l].resize(decode_in_size, 0); + attn_mask_.resize(config_.attn_mask_rows * config_.attn_mask_cols, -1.0f); + local_attn_mask_.resize( + config_.attn_mask_rows * config_.attn_mask_cols, -1.0f); +} +void LlmIOManager::copy_prefill_to_decode_input() { + size_t copy_size = prefill_cache_len_ * config_.head_dim * sizeof(float); + for (size_t l = 0; l < config_.n_layers; l++) { + for (size_t h = 0; h < config_.n_kv_heads; h++) { + const char* src = prefill_k_input_[l].data() + + h * prefill_cache_len_ * config_.head_dim * sizeof(float); + char* dst = decode_k_input_[l].data() + + h * decode_cache_len_ * config_.head_dim * sizeof(float); + std::memcpy(dst, src, copy_size); } - attn_mask_.resize( - config_.attn_mask_rows * config_.attn_mask_cols, -1.0f); - local_attn_mask_.resize( - config_.attn_mask_rows * config_.attn_mask_cols, -1.0f); - } - void LlmIOManager::copy_prefill_to_decode_input() { - size_t copy_size = prefill_cache_len_ * config_.head_dim * sizeof(float); - for (size_t l = 0; l < config_.n_layers; l++) { - for (size_t h = 0; h < config_.n_kv_heads; h++) { - const char* src = prefill_k_input_[l].data() + - h * prefill_cache_len_ * config_.head_dim * sizeof(float); - char* dst = decode_k_input_[l].data() + - h * decode_cache_len_ * config_.head_dim * sizeof(float); - std::memcpy(dst, src, copy_size); - } - for (size_t h = 0; h < config_.n_kv_heads; h++) { - const char* src = prefill_v_input_[l].data() + - h * prefill_cache_len_ * config_.head_dim * sizeof(float); - char* dst = decode_v_input_[l].data() + - h * decode_cache_len_ * config_.head_dim * sizeof(float); - std::memcpy(dst, src, copy_size); - } + for (size_t h = 0; h < config_.n_kv_heads; h++) { + const char* src = prefill_v_input_[l].data() + + h * prefill_cache_len_ * config_.head_dim * sizeof(float); + char* dst = decode_v_input_[l].data() + + h * decode_cache_len_ * config_.head_dim * sizeof(float); + std::memcpy(dst, src, copy_size); } } +} - void LlmIOManager::append_kv_input_tensors( - std::vector>& k_bufs, - std::vector>& v_bufs, - std::vector& k_tensors, - std::vector& v_tensors, - executorch::aten::ScalarType kv_scalar_type, - size_t cache_len, - std::vector& inputs) const { - for (size_t l = 0; l < config_.n_layers; l++) { - auto k_in = executorch::extension::from_blob( - k_bufs[l].data(), - {1, - static_cast(config_.n_kv_heads), - static_cast(cache_len), - static_cast(config_.head_dim)}, - kv_scalar_type); - k_tensors.push_back(k_in); - inputs.emplace_back(k_in); - } - for (size_t l = 0; l < config_.n_layers; l++) { - auto v_in = executorch::extension::from_blob( - v_bufs[l].data(), - {1, - static_cast(config_.n_kv_heads), - static_cast(cache_len), - static_cast(config_.head_dim)}, - kv_scalar_type); - v_tensors.push_back(v_in); - inputs.emplace_back(v_in); - } +void LlmIOManager::append_kv_input_tensors( + std::vector>& k_bufs, + std::vector>& v_bufs, + std::vector& k_tensors, + std::vector& v_tensors, + executorch::aten::ScalarType kv_scalar_type, + size_t cache_len, + std::vector& inputs) const { + for (size_t l = 0; l < config_.n_layers; l++) { + auto k_in = executorch::extension::from_blob( + k_bufs[l].data(), + {1, + static_cast(config_.n_kv_heads), + static_cast(cache_len), + static_cast(config_.head_dim)}, + kv_scalar_type); + k_tensors.push_back(k_in); + inputs.emplace_back(k_in); + } + for (size_t l = 0; l < config_.n_layers; l++) { + auto v_in = executorch::extension::from_blob( + v_bufs[l].data(), + {1, + static_cast(config_.n_kv_heads), + static_cast(cache_len), + static_cast(config_.head_dim)}, + kv_scalar_type); + v_tensors.push_back(v_in); + inputs.emplace_back(v_in); } +} - void LlmIOManager::update_attn_mask(size_t pos, size_t seq_len) { - const float maskFalse = -255.0f; - const float maskTrue = 0.0f; +void LlmIOManager::update_attn_mask(size_t pos, size_t seq_len) { + const float maskFalse = -255.0f; + const float maskTrue = 0.0f; - const size_t rows = config_.attn_mask_rows; - const size_t cols = config_.attn_mask_cols; - const size_t start_col = cols - rows; - const size_t total_size = rows * cols; + const size_t rows = config_.attn_mask_rows; + const size_t cols = config_.attn_mask_cols; + const size_t start_col = cols - rows; + const size_t total_size = rows * cols; - // Initialize entire mask to maskFalse (uses memset for 0.0f) - fill_float_region(attn_mask_.data(), 0, total_size, maskFalse); + // Initialize entire mask to maskFalse (uses memset for 0.0f) + fill_float_region(attn_mask_.data(), 0, total_size, maskFalse); - if (seq_len > 1) { - // Prefill phase - const size_t rows = config_.attn_mask_rows; - const size_t cols = config_.attn_mask_cols; - const size_t start_col = config_.attn_mask_cols - config_.attn_mask_rows; + if (seq_len > 1) { + // Prefill phase + const size_t rows = config_.attn_mask_rows; + const size_t cols = config_.attn_mask_cols; + const size_t start_col = config_.attn_mask_cols - config_.attn_mask_rows; - for (int i = 0; i < seq_len; ++i) { - int start_idx = cols - rows; - int end_idx = start_idx + seq_len; + for (int i = 0; i < seq_len; ++i) { + int start_idx = cols - rows; + int end_idx = start_idx + seq_len; - for (int j = start_idx; j < end_idx && j <= start_idx + i; ++j) { - attn_mask_[i * cols + j] = maskTrue; - } + for (int j = start_idx; j < end_idx && j <= start_idx + i; ++j) { + attn_mask_[i * cols + j] = maskTrue; } + } - for (int i = 0; i < rows; ++i) { - for (int j = 0; j < pos; ++j) { - attn_mask_[i * cols + j] = maskTrue; - } - } - } else { - const size_t left_len = std::min(pos, cols); - if (left_len > 0) { - fill_float_region(attn_mask_.data(), 0, left_len, maskTrue); + for (int i = 0; i < rows; ++i) { + for (int j = 0; j < pos; ++j) { + attn_mask_[i * cols + j] = maskTrue; } + } + } else { + const size_t left_len = std::min(pos, cols); + if (left_len > 0) { + fill_float_region(attn_mask_.data(), 0, left_len, maskTrue); + } - if (cols > 0) { - attn_mask_[cols - 1] = maskTrue; - } + if (cols > 0) { + attn_mask_[cols - 1] = maskTrue; } } +} - void LlmIOManager::update_local_attn_mask(size_t pos, size_t seq_len) { - const float maskFalse = -255.0f; - const float maskTrue = 0.0f; - - const size_t rows = config_.attn_mask_rows; - const size_t cols = config_.attn_mask_cols; - const size_t slide_window = 512; - const size_t total_size = rows * cols; - - // Start with a copy of the full attention mask - std::memcpy(local_attn_mask_.data(), attn_mask_.data(), total_size * sizeof(float)); - - if (seq_len > 1) { - // Prefill phase with sliding window - if ((pos + seq_len) >= slide_window) { - // Mask out positions before the sliding window - const size_t mask_start = slide_window - seq_len; - const size_t mask_end = pos; //512 - const size_t mask_len = mask_end - mask_start; - - if (mask_len > 0) { - for (size_t i = 0; i < rows; ++i) { - fill_float_region( - local_attn_mask_.data(), - i * cols + mask_start, - mask_len, - maskFalse); - } - } +void LlmIOManager::update_local_attn_mask(size_t pos, size_t seq_len) { + const float maskFalse = -255.0f; + const float maskTrue = 0.0f; + + const size_t rows = config_.attn_mask_rows; + const size_t cols = config_.attn_mask_cols; + const size_t sliding_window = config_.sliding_window.value(); + const size_t total_size = rows * cols; + + // Start with a copy of the full attention mask + std::memcpy( + local_attn_mask_.data(), attn_mask_.data(), total_size * sizeof(float)); + + if (seq_len > 1) { + // Prefill phase with sliding window + if ((pos + seq_len) >= sliding_window) { + // Mask out positions before the sliding window + const size_t mask_start = sliding_window - seq_len; + const size_t mask_end = pos; // 512 + const size_t mask_len = mask_end - mask_start; + + if (mask_len > 0) { + for (size_t i = 0; i < rows; ++i) { + fill_float_region( + local_attn_mask_.data(), + i * cols + mask_start, + mask_len, + maskFalse); } - } else { - // Decode phase (single token) - // Set all positions up to pos to maskTrue - const size_t end = std::min(static_cast(pos), cols); - fill_float_region(local_attn_mask_.data(), 0, end, maskTrue); - - // If there are more than 512 True values, we need to keep only the last 512 - // Count True values efficiently by scanning from the end - size_t valid_count = 1; - for (size_t i = 0; i < cols; ++i) { - if (local_attn_mask_[i] == maskTrue) { - valid_count++; - } else { - break; - } } + } + } else { + // Decode phase (single token) + // Set all positions up to pos to maskTrue + const size_t end = std::min(static_cast(pos), cols); + fill_float_region(local_attn_mask_.data(), 0, end, maskTrue); + + // If there are more than 512 True values, we need to keep only the last 512 + // Count True values efficiently by scanning from the end + size_t valid_count = 1; + for (size_t i = 0; i < cols; ++i) { + if (local_attn_mask_[i] == maskTrue) { + valid_count++; + } else { + break; + } + } - if (valid_count > slide_window) { - const size_t to_remove = valid_count - slide_window; - // Set the oldest True values to maskFalse - for (size_t i = 0; i < to_remove; ++i) { - if (local_attn_mask_[i] == maskTrue) { - local_attn_mask_[i] = maskFalse; - } - } + if (valid_count > sliding_window) { + const size_t to_remove = valid_count - sliding_window; + // Set the oldest True values to maskFalse + for (size_t i = 0; i < to_remove; ++i) { + if (local_attn_mask_[i] == maskTrue) { + local_attn_mask_[i] = maskFalse; } - - if (cols > 0) { - attn_mask_[cols - 1] = maskTrue; } } - } - void LlmIOManager::dump_buffer( - const std::string& tag, size_t layer, const char* data, size_t size) { - if (!dump_enabled_) { - return; - } - char filename[512]; - std::snprintf( - filename, sizeof(filename), - "%s/%s_layer%zu_pos%zu_step%zu.bin", - dump_dir_.c_str(), tag.c_str(), layer, current_pos_, dump_step_); - FILE* f = std::fopen(filename, "wb"); - if (f) { - std::fwrite(data, 1, size, f); - std::fclose(f); + if (cols > 0) { + attn_mask_[cols - 1] = maskTrue; } } - void LlmIOManager::dump_kv_input(const std::string& phase) { - if (!dump_enabled_) { - return; - } - const bool is_prefill = (phase == "prefill"); - const auto& k_bufs = is_prefill ? prefill_k_input_ : decode_k_input_; - const auto& v_bufs = is_prefill ? prefill_v_input_ : decode_v_input_; - for (size_t l = 0; l < config_.n_layers; l++) { - dump_buffer(phase + "_k_input", l, k_bufs[l].data(), k_bufs[l].size()); - dump_buffer(phase + "_v_input", l, v_bufs[l].data(), v_bufs[l].size()); - } +} + +void LlmIOManager::dump_buffer( + const std::string& tag, + size_t layer, + const char* data, + size_t size) { + if (!dump_enabled_) { + return; } - void LlmIOManager::update_kv_cache( - const std::vector& model_outputs, - std::vector>& k_bufs, - std::vector>& v_bufs, - size_t cache_len, - size_t new_len, - const std::string& phase) { - // model_outputs: [logits, k0..k_{L-1}, v0..v_{L-1}] - for (size_t l = 0; l < config_.n_layers; l++) { - copy_tensor_to_cache( - model_outputs[1 + l].toTensor(), - new_len, - k_bufs[l].data(), - cache_len, - current_pos_); - } - for (size_t l = 0; l < config_.n_layers; l++) { - copy_tensor_to_cache( - model_outputs[1 + config_.n_layers + l].toTensor(), - new_len, - v_bufs[l].data(), - cache_len, - current_pos_); - } - dump_kv_output(phase, model_outputs); - dump_step_++; + char filename[512]; + std::snprintf( + filename, + sizeof(filename), + "%s/%s_layer%zu_pos%zu_step%zu.bin", + dump_dir_.c_str(), + tag.c_str(), + layer, + current_pos_, + dump_step_); + FILE* f = std::fopen(filename, "wb"); + if (f) { + std::fwrite(data, 1, size, f); + std::fclose(f); } - void LlmIOManager::dump_kv_output( - const std::string& phase, - const std::vector& model_outputs) { - if (!dump_enabled_) { - return; - } - for (size_t l = 0; l < config_.n_layers; l++) { - const auto& k_out = model_outputs[1 + l].toTensor(); - const auto& v_out = - model_outputs[1 + config_.n_layers + l].toTensor(); - dump_buffer( - phase + "_k_output", - l, - reinterpret_cast(k_out.const_data_ptr()), - k_out.numel() * sizeof(float)); - dump_buffer( - phase + "_v_output", - l, - reinterpret_cast(v_out.const_data_ptr()), - v_out.numel() * sizeof(float)); - } +} +void LlmIOManager::dump_kv_input(const std::string& phase) { + if (!dump_enabled_) { + return; + } + const bool is_prefill = (phase == "prefill"); + const auto& k_bufs = is_prefill ? prefill_k_input_ : decode_k_input_; + const auto& v_bufs = is_prefill ? prefill_v_input_ : decode_v_input_; + for (size_t l = 0; l < config_.n_layers; l++) { + dump_buffer(phase + "_k_input", l, k_bufs[l].data(), k_bufs[l].size()); + dump_buffer(phase + "_v_input", l, v_bufs[l].data(), v_bufs[l].size()); + } +} +void LlmIOManager::update_kv_cache( + const std::vector& model_outputs, + std::vector>& k_bufs, + std::vector>& v_bufs, + size_t cache_len, + size_t new_len, + const std::string& phase) { + // model_outputs: [logits, k0..k_{L-1}, v0..v_{L-1}] + for (size_t l = 0; l < config_.n_layers; l++) { + copy_tensor_to_cache( + model_outputs[1 + l].toTensor(), + new_len, + k_bufs[l].data(), + cache_len, + current_pos_); + } + for (size_t l = 0; l < config_.n_layers; l++) { + copy_tensor_to_cache( + model_outputs[1 + config_.n_layers + l].toTensor(), + new_len, + v_bufs[l].data(), + cache_len, + current_pos_); } + dump_kv_output(phase, model_outputs); + dump_step_++; +} +void LlmIOManager::dump_kv_output( + const std::string& phase, + const std::vector& model_outputs) { + if (!dump_enabled_) { + return; + } + for (size_t l = 0; l < config_.n_layers; l++) { + const auto& k_out = model_outputs[1 + l].toTensor(); + const auto& v_out = model_outputs[1 + config_.n_layers + l].toTensor(); + dump_buffer( + phase + "_k_output", + l, + reinterpret_cast(k_out.const_data_ptr()), + k_out.numel() * sizeof(float)); + dump_buffer( + phase + "_v_output", + l, + reinterpret_cast(v_out.const_data_ptr()), + v_out.numel() * sizeof(float)); + } +} - void LlmIOManager::prepare_rotary_embedding_inputs( - size_t seq_len, - std::vector& inputs) { - const size_t element_size = seq_len * config_.head_dim; - global_cos_data_.resize(element_size); - global_sin_data_.resize(element_size); - if (config_.use_gemma) { +void LlmIOManager::prepare_rotary_embedding_inputs( + size_t seq_len, + std::vector& inputs) { + const size_t element_size = seq_len * config_.head_dim; + global_cos_data_.resize(element_size); + global_sin_data_.resize(element_size); + switch (config_.embedding_strategy) { + case EmbeddingStrategy::kGlobalOnly: + break; + case EmbeddingStrategy::kGlobalAndLocal: local_cos_data_.resize(element_size); local_sin_data_.resize(element_size); - } + break; + } - if (rotary_embedding_->update(pos_ids_, global_cos_data_.data(), global_sin_data_.data(), local_cos_data_.data(), local_sin_data_.data()) != executorch::runtime::Error::Ok) { - ET_LOG(Error, "update pos_embedding failed !"); - } + if (rotary_embedding_->update( + pos_ids_, + global_cos_data_.data(), + global_sin_data_.data(), + local_cos_data_.data(), + local_sin_data_.data()) != executorch::runtime::Error::Ok) { + ET_LOG(Error, "update pos_embedding failed !"); + } - global_cos_tensor_ = executorch::extension::from_blob( - global_cos_data_.data(), - {1, static_cast(seq_len), - static_cast(config_.head_dim)}); - global_sin_tensor_ = executorch::extension::from_blob( - global_sin_data_.data(), - {1, static_cast(seq_len), - static_cast(config_.head_dim)}); - if (config_.use_gemma) { + global_cos_tensor_ = executorch::extension::from_blob( + global_cos_data_.data(), + {1, + static_cast(seq_len), + static_cast(config_.head_dim)}); + global_sin_tensor_ = executorch::extension::from_blob( + global_sin_data_.data(), + {1, + static_cast(seq_len), + static_cast(config_.head_dim)}); + switch (config_.embedding_strategy) { + case EmbeddingStrategy::kGlobalOnly: + break; + case EmbeddingStrategy::kGlobalAndLocal: local_cos_tensor_ = executorch::extension::from_blob( local_cos_data_.data(), - {1, static_cast(seq_len), - static_cast(config_.head_dim)}); + {1, + static_cast(seq_len), + static_cast(config_.head_dim)}); local_sin_tensor_ = executorch::extension::from_blob( local_sin_data_.data(), - {1, static_cast(seq_len), - static_cast(config_.head_dim)}); - } + {1, + static_cast(seq_len), + static_cast(config_.head_dim)}); + break; + } - inputs.emplace_back(global_cos_tensor_); - inputs.emplace_back(global_sin_tensor_); - if (config_.use_gemma) { + inputs.emplace_back(global_cos_tensor_); + inputs.emplace_back(global_sin_tensor_); + switch (config_.embedding_strategy) { + case EmbeddingStrategy::kGlobalOnly: + break; + case EmbeddingStrategy::kGlobalAndLocal: inputs.emplace_back(local_cos_tensor_); inputs.emplace_back(local_sin_tensor_); - } + break; } +} } // namespace example diff --git a/examples/samsung/executor_runner/llama/runner/io_manager.h b/examples/samsung/executor_runner/llama/runner/io_manager.h index dcaec626486..ef13283dfa8 100644 --- a/examples/samsung/executor_runner/llama/runner/io_manager.h +++ b/examples/samsung/executor_runner/llama/runner/io_manager.h @@ -11,12 +11,17 @@ #include #include -#include #include #include +#include namespace example { +namespace llm = ::executorch::extension::llm; +namespace runtime = ::executorch::runtime; +namespace aten = ::executorch::aten; +namespace extension = ::executorch::extension; + /** * IO manager for models where the runner manages KV-cache externally. * @@ -42,89 +47,120 @@ namespace example { * k_cache[1,H,max-1,D] x L, v_cache[1,H,max-1,D] x L * Output: logits, k_cache[1,H,1,D] x L, v_cache[1,H,1,D] x L */ -class ET_EXPERIMENTAL LlmIOManager - : public executorch::extension::llm::IOManager { +class ET_EXPERIMENTAL LlmIOManager : public llm::IOManager { public: struct Config { size_t max_seq_len; + size_t max_context_len; size_t prefill_seq_len; size_t decode_seq_len; size_t n_layers; + size_t n_heads; size_t n_kv_heads; size_t head_dim; + size_t hidden_dim; + size_t vocab_size; size_t attn_mask_rows = 128; size_t attn_mask_cols = 1024; - size_t sliding_window = 512; - bool use_gemma = false; + std::optional sliding_window; + EmbeddingStrategy embedding_strategy = EmbeddingStrategy::kGlobalOnly; + MaskingStrategy masking_strategy = MaskingStrategy::kCausal; }; - /** - * @brief Construct an IOManager bound to a Module. - * - * @param module The Module used for querying method metadata and execution. - */ - LlmIOManager(executorch::extension::Module& module, Config config) - : IOManager(module), module_(module), config_(config) { - rotary_embedding_ = std::make_unique(RotaryEmbedding::Config{config_.head_dim, config_.max_seq_len, config_.use_gemma}); - ET_LOG( - Info, - "(%s) layers(%ld), kv_heads(%ld) max_sqe_len(%ld) \ - head_dim(%ld) prefill_seq_len(%ld) use_gemma(%d)", - ET_FUNCTION, - config_.n_layers, - config_.n_kv_heads, - config_.max_seq_len, - config_.head_dim, - config_.prefill_seq_len, - config_.use_gemma); - (void)module_; + LlmIOManager(extension::Module& module, Config config); + + const Config& get_config() const { + return config_; + } + + ET_NODISCARD virtual runtime::Error load_rotary_embedding( + std::optional pos_embedding_path); + + size_t get_attention_setup_tensor_count() const { + constexpr size_t kGlobalMaskCount = 1; + constexpr size_t kGlobalRopeCosCount = 1; + constexpr size_t kGlobalRopeSinCount = 1; + constexpr size_t kBaselineMetadataCount = + kGlobalMaskCount + kGlobalRopeCosCount + kGlobalRopeSinCount; + + size_t setup_tensor_count = kBaselineMetadataCount; - prefill_cache_len_ = config_.max_seq_len - config_.prefill_seq_len; - decode_cache_len_ = config_.max_seq_len - 1; + const bool requires_local_mask = + (config_.masking_strategy == MaskingStrategy::kCausalLocal || + config_.masking_strategy == MaskingStrategy::kHybridConvAttention); + if (requires_local_mask) { + constexpr size_t kLocalMaskCount = 1; + setup_tensor_count += kLocalMaskCount; + } + + const bool requires_local_rope = + (config_.embedding_strategy == EmbeddingStrategy::kGlobalAndLocal); + if (requires_local_rope) { + constexpr size_t kLocalRopeCosCount = 1; + constexpr size_t kLocalRopeSinCount = 1; + setup_tensor_count += (kLocalRopeCosCount + kLocalRopeSinCount); + } - allocate_buffers(); + return setup_tensor_count; } /** - * @brief Load pos_embedding files.. - * @param pos_embedding_path Pos_embedding file, that include global_cos.bin, global_sin.bin, local_cos.bin, local_sin.bin. + * Calculates the total number of tensor slots required in the model's input + * buffer. + * + * Memory Layout Layout: + * [0] : Token/Prompt IDs + * [1 ... N] : Attention Setup Overheads (Masks, RoPE Sin/Cos + * Tensors) [N+1 ... End] : Alternating Key and Value Cache Tensors for + * all Layers */ - ET_NODISCARD virtual executorch::runtime::Error load_rotary_embedding(std::optional pos_embedding_path); + size_t get_input_reserve_size() const { + constexpr size_t kTokenInputCount = 1; + constexpr size_t kKVSlotsPerLayer = 2; - /** - * @brief Get the use_gemma configuration flag. - * @return True if using Gemma model specific processing, false otherwise. - */ - bool use_gemma() const { - return config_.use_gemma; + return kTokenInputCount + get_attention_setup_tensor_count() + + (kKVSlotsPerLayer * config_.n_layers); } - /** - * @brief Virtual destructor to allow proper cleanup in derived classes. - */ virtual ~LlmIOManager() = default; - /** - * @brief Load the IO manager with method metadata for prefill and - * decode operations. - * - * @param prefill_method The prefill method to initialize with. - * @param decode_method The decode method to initialize with. - */ - ET_NODISCARD virtual executorch::runtime::Error load( + ET_NODISCARD virtual runtime::Error load( const std::string& prefill_method, const std::string& decode_method) override { - ET_LOG(Info, - "(%s) layers(%ld), kv_heads(%ld) max_sqe_len(%ld) \ - head_dim(%ld) prefill_seq_len(%ld)", - ET_FUNCTION, - config_.n_layers, - config_.n_kv_heads, - config_.max_seq_len, - config_.head_dim, - config_.prefill_seq_len); - - return executorch::runtime::Error::Ok; + (void)prefill_method; + (void)decode_method; + + if ((config_.masking_strategy == MaskingStrategy::kCausalLocal || + config_.masking_strategy == MaskingStrategy::kHybridConvAttention) && + !config_.sliding_window.has_value()) { + ET_LOG( + Error, + "Sliding window configuration is required for the selected masking strategy"); + return runtime::Error::InvalidArgument; + } + + ET_LOG( + Info, + "(%s) Model Config: layers(%zu), kv_heads(%zu), hidden_dim(%zu), head_dim(%zu), \ + vocab(%zu), max_seq(%zu), max_ctx(%zu), prefill_seq(%zu), decode_seq(%zu), \ + mask_rows(%zu), mask_cols(%zu), slide_win(%zu), emb_strat(%d), mask_strat(%d)", + ET_FUNCTION, + config_.n_layers, + config_.n_kv_heads, + config_.hidden_dim, + config_.head_dim, + config_.vocab_size, + config_.max_seq_len, + config_.max_context_len, + config_.prefill_seq_len, + config_.decode_seq_len, + config_.attn_mask_rows, + config_.attn_mask_cols, + config_.sliding_window.value_or(0), + static_cast(config_.embedding_strategy), + static_cast(config_.masking_strategy)); + + return runtime::Error::Ok; } /** @@ -133,7 +169,7 @@ class ET_EXPERIMENTAL LlmIOManager * @param prefill_method The prefill method to reset with. * @param decode_method The decode method to reset with. */ - ET_NODISCARD virtual executorch::runtime::Error reset( + ET_NODISCARD virtual runtime::Error reset( const std::string& prefill_method, const std::string& decode_method) override { (void)prefill_method; @@ -145,18 +181,14 @@ class ET_EXPERIMENTAL LlmIOManager prefill_done_ = false; for (size_t l = 0; l < config_.n_layers; l++) { - std::fill( - prefill_k_input_[l].begin(), prefill_k_input_[l].end(), 0); - std::fill( - prefill_v_input_[l].begin(), prefill_v_input_[l].end(), 0); - std::fill( - decode_k_input_[l].begin(), decode_k_input_[l].end(), 0); - std::fill( - decode_v_input_[l].begin(), decode_v_input_[l].end(), 0); + std::fill(prefill_k_input_[l].begin(), prefill_k_input_[l].end(), 0); + std::fill(prefill_v_input_[l].begin(), prefill_v_input_[l].end(), 0); + std::fill(decode_k_input_[l].begin(), decode_k_input_[l].end(), 0); + std::fill(decode_v_input_[l].begin(), decode_v_input_[l].end(), 0); } std::fill(attn_mask_.begin(), attn_mask_.end(), 0.0f); - return executorch::runtime::Error::Ok; + return runtime::Error::Ok; } /** @@ -166,13 +198,12 @@ class ET_EXPERIMENTAL LlmIOManager * @param start_pos The tensor containing the starting position of the current * input within the context. * @param prefill_method The prefill method to prepare inputs for. - * @return std::vector Vector of prepared inputs + * @return std::vector Vector of prepared inputs * for the prefill method. */ - virtual executorch::runtime::Result> - prepare_prefill( - const executorch::extension::TensorPtr& input, - const executorch::extension::TensorPtr& start_pos, + virtual runtime::Result> prepare_prefill( + const extension::TensorPtr& input, + const extension::TensorPtr& start_pos, const std::string& prefill_method) override; /** @@ -184,11 +215,11 @@ class ET_EXPERIMENTAL LlmIOManager * @param start_pos The tensor containing the starting position. * @return Vector of prepared inputs for the prefill method. */ - executorch::runtime::Result> - prepare_prefill( - const executorch::extension::TensorPtr& input, - const executorch::extension::TensorPtr& start_pos) { - return prepare_prefill(input, start_pos, example::default_method_name_prefill); + runtime::Result> prepare_prefill( + const extension::TensorPtr& input, + const extension::TensorPtr& start_pos) { + return prepare_prefill( + input, start_pos, example::default_method_name_prefill); } /** @@ -198,13 +229,12 @@ class ET_EXPERIMENTAL LlmIOManager * @param start_pos The tensor containing the starting position of the current * input within the context. * @param decode_method The decode method to prepare inputs for. - * @return std::vector Vector of prepared inputs + * @return std::vector Vector of prepared inputs * for the decode method. */ - virtual executorch::runtime::Result> - prepare_decode( - const executorch::extension::TensorPtr& input, - const executorch::extension::TensorPtr& start_pos, + virtual runtime::Result> prepare_decode( + const extension::TensorPtr& input, + const extension::TensorPtr& start_pos, const std::string& decode_method) override; /** @@ -216,11 +246,11 @@ class ET_EXPERIMENTAL LlmIOManager * @param start_pos The tensor containing the starting position. * @return Vector of prepared inputs for the decode method. */ - executorch::runtime::Result> - prepare_decode( - const executorch::extension::TensorPtr& input, - const executorch::extension::TensorPtr& start_pos) { - return prepare_decode(input, start_pos, example::default_method_name_decode); + runtime::Result> prepare_decode( + const extension::TensorPtr& input, + const extension::TensorPtr& start_pos) { + return prepare_decode( + input, start_pos, example::default_method_name_decode); } /** @@ -230,8 +260,8 @@ class ET_EXPERIMENTAL LlmIOManager * @param prefill_method The prefill method to update with outputs. * @param model_outputs Vector of outputs from the prefill method execution. */ - ET_NODISCARD virtual executorch::runtime::Error update_prefill( - const std::vector& model_outputs, + ET_NODISCARD virtual runtime::Error update_prefill( + const std::vector& model_outputs, const std::string& prefill_method) override; /** @@ -242,8 +272,8 @@ class ET_EXPERIMENTAL LlmIOManager * @param model_outputs Vector of outputs from the prefill execution. * @return Error code. */ - ET_NODISCARD executorch::runtime::Error update_prefill( - const std::vector& model_outputs) { + ET_NODISCARD runtime::Error update_prefill( + const std::vector& model_outputs) { return update_prefill(model_outputs, example::default_method_name_prefill); } @@ -254,8 +284,8 @@ class ET_EXPERIMENTAL LlmIOManager * @param decode_method The decode method to update with outputs. * @param model_outputs Vector of outputs from the decode method execution. */ - ET_NODISCARD virtual executorch::runtime::Error update_decode( - const std::vector& model_outputs, + ET_NODISCARD virtual runtime::Error update_decode( + const std::vector& model_outputs, const std::string& decode_method) override; /** @@ -266,8 +296,8 @@ class ET_EXPERIMENTAL LlmIOManager * @param model_outputs Vector of outputs from the decode execution. * @return Error code. */ - ET_NODISCARD executorch::runtime::Error update_decode( - const std::vector& model_outputs) { + ET_NODISCARD runtime::Error update_decode( + const std::vector& model_outputs) { return update_decode(model_outputs, example::default_method_name_decode); } @@ -288,22 +318,6 @@ class ET_EXPERIMENTAL LlmIOManager void allocate_buffers(); - /** - * @brief Calculate the number of inputs to reserve for prefill/decode. - * - * The total depends on whether the model is Gemma: - * - Gemma: 7 + n_layers * 2 - * - Non-Gemma: 4 + n_layers * 2 - * - * @return The reserve size for the inputs vector. - */ - size_t get_input_reserve_size() const { - if (config_.use_gemma) - return 7 + config_.n_layers * 2; - else - return 4 + config_.n_layers * 2; - } - bool dump_enabled_ = false; mutable size_t dump_step_ = 0; std::string dump_dir_ = "/data/local/tmp/kv_cache_dump"; @@ -343,15 +357,17 @@ class ET_EXPERIMENTAL LlmIOManager * = src[h * new_len * D .. + new_len * D] */ template - void copy_to_cache(const T* src, size_t new_len, - char* cache, size_t cache_len, size_t pos) { + void copy_to_cache( + const T* src, + size_t new_len, + char* cache, + size_t cache_len, + size_t pos) { for (size_t h = 0; h < config_.n_kv_heads; h++) { const T* src_head = src + h * new_len * config_.head_dim; - char* dst_head = - cache + h * cache_len * config_.head_dim * sizeof(T) - + pos * config_.head_dim * sizeof(T); - std::memcpy( - dst_head, src_head, new_len * config_.head_dim * sizeof(T)); + char* dst_head = cache + h * cache_len * config_.head_dim * sizeof(T) + + pos * config_.head_dim * sizeof(T); + std::memcpy(dst_head, src_head, new_len * config_.head_dim * sizeof(T)); } } @@ -413,11 +429,6 @@ class ET_EXPERIMENTAL LlmIOManager * positions [0, pos + r] (value 1) and masks future positions (value 0). */ void update_local_attn_mask(size_t pos, size_t seq_len); - - /** - * @brief Reference to the Module used for method metadata and execution. - */ - executorch::extension::module::Module& module_; Config config_; size_t prefill_cache_len_; @@ -426,6 +437,7 @@ class ET_EXPERIMENTAL LlmIOManager bool prefill_done_ = false; std::vector input_tokens_; + executorch::extension::TensorPtr input_tokens_tensor_; std::vector pos_ids_; // Per-layer input buffers: [1, n_kv_heads, cache_len, head_dim] flattened diff --git a/examples/samsung/executor_runner/llama/runner/prefiller.cpp b/examples/samsung/executor_runner/llama/runner/prefiller.cpp index 3749630fa38..399afa271d8 100644 --- a/examples/samsung/executor_runner/llama/runner/prefiller.cpp +++ b/examples/samsung/executor_runner/llama/runner/prefiller.cpp @@ -14,8 +14,13 @@ namespace example { +namespace llm = ::executorch::extension::llm; +namespace runtime = ::executorch::runtime; +namespace aten = ::executorch::aten; +namespace extension = ::executorch::extension; + Prefiller::Prefiller( - executorch::extension::llm::TextDecoderRunner* text_decoder_runner, + llm::TextDecoderRunner* text_decoder_runner, bool use_kv_cache, bool enable_parallel_prefill, int64_t max_seq_len) @@ -29,7 +34,7 @@ Prefiller::Prefiller( enable_parallel_prefill_(enable_parallel_prefill), max_seq_len_(max_seq_len > 0 ? max_seq_len : 128) {} -::executorch::runtime::Result Prefiller::prefill( +runtime::Result Prefiller::prefill( std::vector& prompt_tokens, int64_t& start_pos) { ET_CHECK_MSG(!prompt_tokens.empty(), "Prompt cannot be null"); @@ -37,10 +42,8 @@ ::executorch::runtime::Result Prefiller::prefill( ET_CHECK_OK_OR_RETURN_ERROR(text_decoder_runner_->load()); } - // Check if we need to chunk the prompt tokens - int32_t num_prompt_tokens = prompt_tokens.size(); - // If prompt tokens exceed max_seq_len_, we need to chunk them + int32_t num_prompt_tokens = prompt_tokens.size(); ET_LOG( Info, "Micro Batch: num_prompt(%d), prefill::max_seq_len(%ld), start_pos(%ld)", @@ -64,14 +67,13 @@ ::executorch::runtime::Result Prefiller::prefill( num_tokens_to_prefill_with, prompt_tokens_to_process.begin()); - // Process this chunk auto chunk_result = prefill_chunk(prompt_tokens_to_process, start_pos); ET_CHECK_OK_OR_RETURN_ERROR(chunk_result.error()); cur_token = chunk_result.get(); - //if (enable_parallel_prefill_) { - // start_pos += num_tokens_to_prefill_with; - //} + if (enable_parallel_prefill_) { + start_pos += num_tokens_to_prefill_with; + } num_tokens_to_process += num_tokens_to_prefill_with; } @@ -82,7 +84,7 @@ ::executorch::runtime::Result Prefiller::prefill( } } -::executorch::runtime::Result Prefiller::prefill_chunk( +runtime::Result Prefiller::prefill_chunk( std::vector& prompt_tokens, int64_t& start_pos) { // enable_parallel_prefill_ maybe set even when not using kv cache @@ -93,10 +95,10 @@ ::executorch::runtime::Result Prefiller::prefill_chunk( uint64_t cur_token; if (enable_parallel_prefill_ || !use_kv_cache_) { // initialize tensor wrappers - auto tokens = executorch::extension::from_blob( + auto tokens = extension::from_blob( prompt_tokens.data(), {1, static_cast(prompt_tokens.size())}, - executorch::aten::ScalarType::Long); + aten::ScalarType::Long); auto outputs_res = text_decoder_runner_->step(tokens, start_pos); @@ -120,8 +122,8 @@ ::executorch::runtime::Result Prefiller::prefill_chunk( // initialize tensor wrappers // TODO(sicheon.oh) : Set tensor size by metadata::get_prefill_ar_len - auto tokens = executorch::extension::from_blob( - &cur_token, {1, 1}, executorch::aten::ScalarType::Long); + auto tokens = + extension::from_blob(&cur_token, {1, 1}, aten::ScalarType::Long); // run the first token and get back logits tensor. Assuming the first token // is bos so don't callback. diff --git a/examples/samsung/executor_runner/llama/runner/prefiller.h b/examples/samsung/executor_runner/llama/runner/prefiller.h index 42102cc7f25..fc98d284f1e 100644 --- a/examples/samsung/executor_runner/llama/runner/prefiller.h +++ b/examples/samsung/executor_runner/llama/runner/prefiller.h @@ -16,37 +16,25 @@ namespace example { -class ET_EXPERIMENTAL Prefiller - : public executorch::extension::llm::TextPrefiller { +namespace llm = ::executorch::extension::llm; +namespace runtime = ::executorch::runtime; +namespace extension = ::executorch::extension; + +class ET_EXPERIMENTAL Prefiller : public llm::TextPrefiller { public: Prefiller( - executorch::extension::llm::TextDecoderRunner* text_decoder_runner, + llm::TextDecoderRunner* text_decoder_runner, bool use_kv_cache, bool enable_parallel_prefill, int64_t max_seq_len = 128); virtual ~Prefiller() = default; - /** - * Prefill an LLM Module with the given text input. - * @param prompt_tokens The text prompt tokens to the LLM Module. Encoded by - * tokenizer. - * @param start_pos The starting position in KV cache of the input in the LLM - * Module. - * @return The next token of the LLM Module after prefill. - */ - virtual ::executorch::runtime::Result prefill( + virtual runtime::Result prefill( std::vector& prompt_tokens, int64_t& start_pos) override; - /** - * Helper method to prefill a chunk of tokens. - * @param prompt_tokens The chunk of text prompt tokens to process. - * @param start_pos The starting position in KV cache of the input in the LLM - * Module. - * @return The next token of the LLM Module after prefilling this chunk. - */ - virtual ::executorch::runtime::Result prefill_chunk( + virtual runtime::Result prefill_chunk( std::vector& prompt_tokens, int64_t& start_pos) override; @@ -57,7 +45,7 @@ class ET_EXPERIMENTAL Prefiller * lies with the outer class or entity (likely Runner) that creates * and passes the TextDecoderRunner instance to TextPrefiller. */ - executorch::extension::llm::TextDecoderRunner* text_decoder_runner_; + llm::TextDecoderRunner* text_decoder_runner_; bool use_kv_cache_; bool enable_parallel_prefill_; int64_t max_seq_len_; diff --git a/examples/samsung/executor_runner/llama/runner/rotary_embedding.cpp b/examples/samsung/executor_runner/llama/runner/rotary_embedding.cpp index cc833f10740..6884f4bf5a4 100644 --- a/examples/samsung/executor_runner/llama/runner/rotary_embedding.cpp +++ b/examples/samsung/executor_runner/llama/runner/rotary_embedding.cpp @@ -7,74 +7,104 @@ */ #include +#include #include #include -#include namespace example { +namespace runtime = ::executorch::runtime; + RotaryEmbedding::RotaryEmbedding(Config config) : config_(config) {} -executorch::runtime::Error RotaryEmbedding::load(std::optional pos_embedding_path) { - if (!pos_embedding_path.has_value()) { - ET_LOG(Error, "RotaryEmbedding not has value"); - return executorch::runtime::Error::Ok; - } - - const size_t total_bytes = config_.max_seq_len * config_.head_dim * sizeof(float); - global_cos_ = std::make_unique(total_bytes); - global_sin_ = std::make_unique(total_bytes); - if (config_.use_gemma) { - local_cos_ = std::make_unique(total_bytes); - local_sin_ = std::make_unique(total_bytes); - } - - std::ifstream file(pos_embedding_path.value(), std::ios::binary); - if (!file) { - ET_LOG(Error, "can't open the file"); - return executorch::runtime::Error::Internal; - } - - file.seekg(0, std::ios::end); - size_t file_size = file.tellg(); - file.seekg(0, std::ios::beg); - - int bin_num = config_.use_gemma ? 4 : 2; - if (file_size != total_bytes * bin_num) { - ET_LOG(Error, "file size is not right"); - return executorch::runtime::Error::Internal; - } - - file.read(global_cos_.get(), total_bytes); - file.read(global_sin_.get(), total_bytes); - if (config_.use_gemma) { - file.read(local_cos_.get(), total_bytes); - file.read(local_sin_.get(), total_bytes); - } - - file.close(); - return executorch::runtime::Error::Ok; +runtime::Error RotaryEmbedding::load( + std::optional pos_embedding_path) { + if (!pos_embedding_path.has_value()) { + ET_LOG( + Error, + "Missing position embedding path. Verify that the file path is provided in the configuration."); + return runtime::Error::InvalidArgument; + } + + const size_t total_bytes = + config_.max_seq_len * config_.head_dim * sizeof(float); + global_cos_ = std::make_unique(total_bytes); + global_sin_ = std::make_unique(total_bytes); + if (config_.strategy == EmbeddingStrategy::kGlobalAndLocal) { + local_cos_ = std::make_unique(total_bytes); + local_sin_ = std::make_unique(total_bytes); + } + + std::ifstream file(pos_embedding_path.value(), std::ios::binary); + if (!file) { + ET_LOG( + Error, + "File I/O error: Cannot open position embedding file at '%s'.", + pos_embedding_path.value().c_str()); + return runtime::Error::InvalidArgument; + } + + file.seekg(0, std::ios::end); + size_t file_size = file.tellg(); + file.seekg(0, std::ios::beg); + + int bin_num = + (config_.strategy == EmbeddingStrategy::kGlobalAndLocal) ? 4 : 2; + size_t expected_size = total_bytes * bin_num; + if (file_size != expected_size) { + ET_LOG( + Error, + "Embedding file size mismatch. Expected %zu bytes (total_bytes: %zu * bin_num: %d based on strategy), but got %zu bytes. File may be corrupted or from a different model version.", + expected_size, + total_bytes, + bin_num, + file_size); + return runtime::Error::InvalidArgument; + } + + file.read(global_cos_.get(), total_bytes); + file.read(global_sin_.get(), total_bytes); + if (config_.strategy == EmbeddingStrategy::kGlobalAndLocal) { + file.read(local_cos_.get(), total_bytes); + file.read(local_sin_.get(), total_bytes); + } + + file.close(); + return runtime::Error::Ok; } -executorch::runtime::Error RotaryEmbedding::update(const std::vector& pos_ids, - float* global_cos, float* global_sin, - float* local_cos, float* local_sin) { - if (pos_ids.empty()) { - return executorch::runtime::Error::Internal; - } - - size_t pos = static_cast(pos_ids[0]); - size_t length = pos_ids.size(); - size_t head_dim = config_.head_dim; - - std::memcpy(global_cos, global_cos_.get() + pos * head_dim * sizeof(float), length * head_dim * sizeof(float)); - std::memcpy(global_sin, global_sin_.get() + pos * head_dim * sizeof(float), length * head_dim * sizeof(float)); - if (config_.use_gemma) { - std::memcpy(local_cos, local_cos_.get() + pos * head_dim * sizeof(float), length * head_dim * sizeof(float)); - std::memcpy(local_sin, local_sin_.get() + pos * head_dim * sizeof(float), length * head_dim * sizeof(float)); - } - - return executorch::runtime::Error::Ok; +runtime::Error RotaryEmbedding::update( + const std::vector& pos_ids, + float* global_cos, + float* global_sin, + float* local_cos, + float* local_sin) { + if (pos_ids.empty()) { + return runtime::Error::InvalidArgument; + } + if (global_cos == nullptr || global_sin == nullptr) { + return runtime::Error::InvalidArgument; + } + if (config_.strategy == EmbeddingStrategy::kGlobalAndLocal && + (local_cos == nullptr || local_sin == nullptr)) { + return runtime::Error::InvalidArgument; + } + + size_t pos = static_cast(pos_ids[0]); + size_t length = pos_ids.size(); + size_t head_dim = config_.head_dim; + size_t bytes_to_copy = length * head_dim * sizeof(float); + size_t element_offset = pos * head_dim * sizeof(float); + + std::memcpy(global_cos, global_cos_.get() + element_offset, bytes_to_copy); + std::memcpy(global_sin, global_sin_.get() + element_offset, bytes_to_copy); + + if (config_.strategy == EmbeddingStrategy::kGlobalAndLocal) { + std::memcpy(local_cos, local_cos_.get() + element_offset, bytes_to_copy); + std::memcpy(local_sin, local_sin_.get() + element_offset, bytes_to_copy); + } + + return runtime::Error::Ok; } } // namespace example diff --git a/examples/samsung/executor_runner/llama/runner/rotary_embedding.h b/examples/samsung/executor_runner/llama/runner/rotary_embedding.h index 5c3cd5d3973..3a0adf84dc0 100644 --- a/examples/samsung/executor_runner/llama/runner/rotary_embedding.h +++ b/examples/samsung/executor_runner/llama/runner/rotary_embedding.h @@ -8,28 +8,36 @@ #pragma once -#include -#include -#include +#include #include +#include +#include +#include +#include namespace example { +namespace runtime = ::executorch::runtime; + class RotaryEmbedding { public: struct Config { size_t head_dim; size_t max_seq_len; - bool use_gemma; + EmbeddingStrategy strategy = EmbeddingStrategy::kGlobalOnly; }; RotaryEmbedding(Config config); ~RotaryEmbedding() = default; - executorch::runtime::Error load(std::optional pos_embedding_path); - executorch::runtime::Error update(const std::vector& pos_ids, - float* global_cos, float* global_sin, - float* local_cos, float* local_sin); + runtime::Error load(std::optional pos_embedding_path); + runtime::Error update( + const std::vector& pos_ids, + float* global_cos, + float* global_sin, + float* local_cos, + float* local_sin); + private: Config config_; std::unique_ptr global_cos_; diff --git a/examples/samsung/executor_runner/llama/runner/runner.cpp b/examples/samsung/executor_runner/llama/runner/runner.cpp index 81a23b74783..835a35290e6 100644 --- a/examples/samsung/executor_runner/llama/runner/runner.cpp +++ b/examples/samsung/executor_runner/llama/runner/runner.cpp @@ -24,14 +24,55 @@ #include #include #include +#include +#include namespace example { -using ::executorch::extension::Module; -using ::executorch::runtime::Error; -using ::executorch::runtime::Result; - namespace llm = ::executorch::extension::llm; +namespace runtime = ::executorch::runtime; +namespace aten = ::executorch::aten; +namespace extension = ::executorch::extension; + +using json = nlohmann::json; + +std::string LLMRunner::get_formatted_prompt( + const std::string& prompt, + const std::string& system_prompt, + DecoderModelVersion version) { + std::string formatted_prompt; + switch (version) { + case DecoderModelVersion::kLlama3: + if (!system_prompt.empty()) { + formatted_prompt.append( + "<|start_header_id|>system<|end_header_id|>\n\n"); + formatted_prompt.append(system_prompt); + formatted_prompt.append("<|eot_id|>"); + } + formatted_prompt.append("<|start_header_id|>user<|end_header_id|>\n\n"); + formatted_prompt.append(prompt); + formatted_prompt.append( + "<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n"); + break; + case DecoderModelVersion::kGemma3: + formatted_prompt.append("user\n"); + formatted_prompt.append(prompt); + formatted_prompt.append("\n"); + formatted_prompt.append("model\n"); + if (!system_prompt.empty()) { + formatted_prompt.append(system_prompt); + formatted_prompt.append("\n"); + } + break; + case DecoderModelVersion::kGeneral: + formatted_prompt.append(prompt); + break; + default: + formatted_prompt.append(prompt); + break; + } + return formatted_prompt; +} /** * @brief Loads a tokenizer from the specified path @@ -42,8 +83,8 @@ namespace llm = ::executorch::extension::llm; * @param type Type of tokenizer to load ("bpe", "tiktoken", "hf", or "sp") * @param tokenizer_path Path to the tokenizer file * @param version Version of the tokenizer (used for tiktoken type) - * @return std::unique_ptr Initialized tokenizer instance, - * or nullptr on failure + * @return std::unique_ptr Initialized tokenizer + * instance, or nullptr on failure */ ET_EXPERIMENTAL std::unique_ptr load_tokenizer( const std::string& type, @@ -65,26 +106,30 @@ ET_EXPERIMENTAL std::unique_ptr load_tokenizer( */ ET_EXPERIMENTAL std::unique_ptr create_llm_runner( const RunnerPaths& runner_paths, - const RunnerConfig& runner_config, + RunnerConfig runner_config, std::unique_ptr<::tokenizers::Tokenizer> tokenizer); std::unique_ptr create_llama_runner( const RunnerPaths& runner_paths, - const RunnerConfig& runner_config) { + RunnerConfig runner_config) { ET_LOG( Info, - "Creating LLaMa runner: model_path=%s, tokenizer_path=%s, use_gemma=%d", + "Creating LLaMa runner: model_path=%s, tokenizer_path=%s", runner_paths.model_path.c_str(), - runner_paths.tokenizer_path.c_str(), - runner_config.use_gemma); + runner_paths.tokenizer_path.c_str()); + + runner_config.load_from_json(); // Create and load tokenizer - std::unique_ptr<::tokenizers::Tokenizer> tokenizer = - load_tokenizer(runner_config.tokenizer_type, runner_paths.tokenizer_path, Version::Default); + std::unique_ptr<::tokenizers::Tokenizer> tokenizer = load_tokenizer( + runner_config.tokenizer_type, + runner_paths.tokenizer_path, + Version::Default); if (tokenizer == nullptr) { ET_LOG( Info, - "Failed to load %s as a Tiktoken, Sentencepiece or Llama2.c tokenizer, make sure the artifact is one of these types", + "Failed to load tokenizer. Parsed tokenizer type is %s, located at %s.", + runner_config.tokenizer_type.c_str(), runner_paths.tokenizer_path.c_str()); return nullptr; } @@ -92,8 +137,11 @@ std::unique_ptr create_llama_runner( Info, "[Create_test_llm_runner] %s %s", runner_paths.model_path.c_str(), - runner_paths.data_path.has_value() ? runner_paths.data_path.value().c_str() : "null"); - return create_llm_runner(runner_paths, runner_config, std::move(tokenizer)); + runner_paths.data_path.has_value() + ? runner_paths.data_path.value().c_str() + : "null"); + return create_llm_runner( + runner_paths, std::move(runner_config), std::move(tokenizer)); } static const std::unordered_map kMetadataMap{ @@ -111,13 +159,12 @@ static const std::unordered_map kMetadataMap{ {example::kNumOfHeads, "get_n_heads"}, {example::kNumOfKVHeads, "get_n_kv_heads"}, {example::kNumOfLayers, "get_n_layers"}, - {"get_sliding_window", "get_sliding_window"} // TODO: find new way -}; + {"get_sliding_window", "get_sliding_window"}}; LLMRunner::LLMRunner( std::unordered_map metadata, std::unique_ptr<::tokenizers::Tokenizer> tokenizer, - std::unique_ptr<::executorch::extension::Module> module, + std::unique_ptr module, std::unique_ptr text_decoder_runner, std::unique_ptr text_prefiller, std::unique_ptr text_token_generator, @@ -141,33 +188,132 @@ bool LLMRunner::is_loaded() const { return text_prefiller_->is_loaded() && text_token_generator_->is_loaded(); } -Error LLMRunner::load() { +runtime::Error LLMRunner::load() { ET_CHECK_OK_OR_RETURN_ERROR(text_decoder_runner_->load()); ET_CHECK_OK_OR_RETURN_ERROR(text_prefiller_->load()); ET_CHECK_OK_OR_RETURN_ERROR(text_token_generator_->load()); - return Error::Ok; + return runtime::Error::Ok; } // Don't print with the same priority during warmup -#define RUNNER_ET_LOG(warmup, format, ...) \ - if (warmup) { \ - ET_LOG(Debug, format, __VA_ARGS__); \ - } else { \ - ET_LOG(Info, format, __VA_ARGS__); \ +#define RUNNER_ET_LOG(is_warmup, format, ...) \ + if (is_warmup) { \ + ET_LOG(Debug, format, __VA_ARGS__); \ + } else { \ + ET_LOG(Info, format, __VA_ARGS__); \ } void replace_all_newlines(std::string& str) { - std::string from = "\\n"; - std::string to = "\n"; + std::string from = "\\n"; + std::string to = "\n"; + + size_t start_pos = 0; + while ((start_pos = str.find(from, start_pos)) != std::string::npos) { + str.replace(start_pos, from.length(), to); + start_pos += to.length(); + } +} + +void RunnerConfig::load_from_json() { + if (config_path.empty()) { + ET_LOG(Info, "Config path is empty, skipping JSON load"); + return; + } + + std::ifstream file(config_path); + if (!file.is_open()) { + ET_LOG(Error, "Could not open config file: %s", config_path.c_str()); + return; + } - size_t start_pos = 0; - while ((start_pos = str.find(from, start_pos)) != std::string::npos) { - str.replace(start_pos, from.length(), to); - start_pos += to.length(); + json j = json::parse(file, nullptr, false); + if (j.is_discarded()) { + ET_LOG( + Error, "Invalid JSON format in config file: %s", config_path.c_str()); + return; + } + + if (j.contains("hidden_dim") && j["hidden_dim"].is_number_integer()) { + hidden_dim = j["hidden_dim"].get(); + } + if (j.contains("n_heads") && j["n_heads"].is_number_integer()) { + n_heads = j["n_heads"].get(); + } + if (j.contains("n_kv_heads") && j["n_kv_heads"].is_number_integer()) { + n_kv_heads = j["n_kv_heads"].get(); + } + if (j.contains("n_layers") && j["n_layers"].is_number_integer()) { + n_layers = j["n_layers"].get(); + } + if (j.contains("head_dim") && j["head_dim"].is_number_integer()) { + head_dim = j["head_dim"].get(); + } + if (j.contains("vocab_size") && j["vocab_size"].is_number_integer()) { + vocab_size = j["vocab_size"].get(); + } + if (j.contains("max_seq_len") && j["max_seq_len"].is_number_integer()) { + max_seq_len = j["max_seq_len"].get(); + } + if (j.contains("max_context_len") && + j["max_context_len"].is_number_integer()) { + max_context_len = j["max_context_len"].get(); + } + if (j.contains("prefill_seq_len") && + j["prefill_seq_len"].is_number_integer()) { + prefill_seq_len = j["prefill_seq_len"].get(); + } + if (j.contains("decode_seq_len") && j["decode_seq_len"].is_number_integer()) { + decode_seq_len = j["decode_seq_len"].get(); + } + if (j.contains("attn_mask_rows") && j["attn_mask_rows"].is_number_integer()) { + attn_mask_rows = j["attn_mask_rows"].get(); + } + if (j.contains("attn_mask_cols") && j["attn_mask_cols"].is_number_integer()) { + attn_mask_cols = j["attn_mask_cols"].get(); + } + if (j.contains("sliding_window") && j["sliding_window"].is_number_integer()) { + sliding_window = j["sliding_window"].get(); + } else { + sliding_window = std::nullopt; + } + + if (j.contains("embedding_strategy") && j["embedding_strategy"].is_string()) { + std::string strategy = j["embedding_strategy"].get(); + if (strategy == "global_only") { + embedding_strategy = EmbeddingStrategy::kGlobalOnly; + } else if (strategy == "global_and_local") { + embedding_strategy = EmbeddingStrategy::kGlobalAndLocal; + } + } + + if (j.contains("masking_strategy") && j["masking_strategy"].is_string()) { + std::string strategy = j["masking_strategy"].get(); + if (strategy == "causal") { + masking_strategy = MaskingStrategy::kCausal; + } else if (strategy == "causal_local") { + masking_strategy = MaskingStrategy::kCausalLocal; + } else if (strategy == "hybrid_conv") { + masking_strategy = MaskingStrategy::kHybridConvAttention; } + } + + if (j.contains("model_version") && j["model_version"].is_string()) { + std::string version_str = j["model_version"].get(); + if (version_str == "llama3") { + model_version = DecoderModelVersion::kLlama3; + } else if (version_str == "gemma3") { + model_version = DecoderModelVersion::kGemma3; + } else if (version_str == "general") { + model_version = DecoderModelVersion::kGeneral; + } + } + + if (j.contains("tokenizer_type") && j["tokenizer_type"].is_string()) { + tokenizer_type = j["tokenizer_type"].get(); + } } -Error LLMRunner::generate_from_pos( +runtime::Error LLMRunner::generate_from_pos( const std::string& prompt, int64_t start_pos, const llm::GenerationConfig& config, @@ -212,19 +358,13 @@ Error LLMRunner::generate_from_pos( Info, "################ input prompt ################\n %s\n", prompt.c_str()); - std::string input_prompt = prompt; - if (!io_manager_->use_gemma()) { - std::string prefix = "<|start_header_id|>system<|end_header_id|><|start_header_id|>user<|end_header_id|>\n\n"; - std::string suffix = "<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n"; - input_prompt = prefix + prompt + suffix; - } else { - std::string prefix = "user\n"; - std::string suffix = "\nmodel\n"; - input_prompt = prefix + prompt + suffix; + + std::string input_prompt = get_formatted_prompt(prompt, "", model_version_); + if (model_version_ == DecoderModelVersion::kGemma3) { replace_all_newlines(input_prompt); } - ::tokenizers::Result> encode_res = tokenizer_->encode( + ::tokenizers::Result> encode_res = tokenizer_->encode( input_prompt, /*bos=*/config.num_bos, /*eos=*/config.num_eos); @@ -345,10 +485,10 @@ Error LLMRunner::generate_from_pos( stats_callback(*stats_); } - return Error::Ok; + return runtime::Error::Ok; } -Error LLMRunner::generate( +runtime::Error LLMRunner::generate( const std::string& prompt, const llm::GenerationConfig& config, std::function token_callback, @@ -356,13 +496,15 @@ Error LLMRunner::generate( return generate_from_pos(prompt, 0, config, token_callback, stats_callback); } -Error LLMRunner::warmup(const std::string& prompt, int32_t max_new_tokens) { +runtime::Error LLMRunner::warmup( + const std::string& prompt, + int32_t max_new_tokens) { // Create a GenerationConfig for warmup llm::GenerationConfig config{ .echo = false, .max_new_tokens = max_new_tokens, .warming = true}; // Call generate with the warmup config - Error err = generate(prompt, config); + runtime::Error err = generate(prompt, config); // Reset stats after warmup, not resetting the std::unique_ptr! stats_->reset(); @@ -392,6 +534,8 @@ std::unique_ptr load_tokenizer( tokenizer = std::make_unique<::tokenizers::HFTokenizer>(); } else if (type == "sp") { tokenizer = std::make_unique<::tokenizers::SPTokenizer>(); + } else { + return nullptr; } if (tokenizer->load(tokenizer_path) == ::tokenizers::Error::Ok) { @@ -404,34 +548,25 @@ std::unique_ptr load_tokenizer( std::unordered_map get_llm_metadata( tokenizers::Tokenizer* tokenizer, - Module* module, - bool use_gemma) { - // Initialize metadata with default values based on model type - // Gemma: hidden_dim=1152, n_heads=4, n_kv_heads=1, n_layers=26, head_dim=256 - // Llama: hidden_dim=2048, n_heads=32, n_kv_heads=8, n_layers=16, head_dim=64 - size_t default_hidden_dim = use_gemma ? 1152 : 2048; - size_t default_n_heads = use_gemma ? 4 : 32; - size_t default_n_kv_heads = use_gemma ? 1 : 8; - size_t default_n_layers = use_gemma ? 26 : 16; - size_t default_head_dim = use_gemma ? 256 : 64; - + extension::Module* module, + RunnerConfig& config) { + // Initialize metadata from JSON config values std::unordered_map metadata({ {kMetadataMap.at(llm::kEnableDynamicShape), false}, - {kMetadataMap.at(llm::kMaxSeqLen), 128}, - {kMetadataMap.at(llm::kMaxContextLen), 128}, + {kMetadataMap.at(llm::kMaxSeqLen), config.max_seq_len}, + {kMetadataMap.at(llm::kMaxContextLen), config.max_context_len}, {kMetadataMap.at(llm::kUseKVCache), true}, {kMetadataMap.at(llm::kUseSDPAWithKVCache), false}, {kMetadataMap.at(example::kPrefillArLen), 0}, - {kMetadataMap.at(example::kNumOfHiddenDim), static_cast(default_hidden_dim)}, - {kMetadataMap.at(example::kNumOfHeads), static_cast(default_n_heads)}, - {kMetadataMap.at(example::kNumOfKVHeads), static_cast(default_n_kv_heads)}, - {kMetadataMap.at(example::kNumOfLayers), static_cast(default_n_layers)}, - {kMetadataMap.at(llm::kHeadDim), static_cast(default_head_dim)}, + {kMetadataMap.at(example::kNumOfHiddenDim), config.hidden_dim}, + {kMetadataMap.at(example::kNumOfHeads), config.n_heads}, + {kMetadataMap.at(example::kNumOfKVHeads), config.n_kv_heads}, + {kMetadataMap.at(example::kNumOfLayers), config.n_layers}, + {kMetadataMap.at(llm::kHeadDim), config.head_dim}, }); - // Read metadata from the model auto method_names_result = module->method_names(); - if (method_names_result.error() != Error::Ok) { + if (method_names_result.error() != runtime::Error::Ok) { ET_LOG(Error, "Failed reading method names"); return metadata; } @@ -446,12 +581,12 @@ std::unordered_map get_llm_metadata( value = get_result.get().toScalar().to(); } else { ET_LOG( - Info, + Debug, "Method %s not found, using the default value %" PRId64, method_name.c_str(), value); } - ET_LOG(Info, "Metadata: %s = %" PRId64, method_name.c_str(), value); + ET_LOG(Debug, "Metadata: %s = %" PRId64, method_name.c_str(), value); } // Set tokenizer-related metadata metadata[llm::kBosId] = tokenizer->bos_tok(); @@ -461,11 +596,11 @@ std::unordered_map get_llm_metadata( std::unordered_set get_eos_ids( tokenizers::Tokenizer* tokenizer, - Module* module) { + extension::Module* module) { std::unordered_set eos_ids = {tokenizer->eos_tok()}; // Get EOS IDs if available auto method_names_result = module->method_names(); - if (method_names_result.error() != Error::Ok) { + if (method_names_result.error() != runtime::Error::Ok) { ET_LOG(Error, "Failed reading method names"); return eos_ids; } @@ -474,7 +609,7 @@ std::unordered_set get_eos_ids( if (method_names.count(llm::kEosIds)) { eos_ids.clear(); auto execute_result = module->execute(llm::kEosIds); - if (execute_result.error() != Error::Ok) { + if (execute_result.error() != runtime::Error::Ok) { ET_LOG(Error, "Failed to execute %s", llm::kEosIds); return eos_ids; } @@ -489,24 +624,24 @@ std::unordered_set get_eos_ids( std::unique_ptr create_llm_runner( const RunnerPaths& runner_paths, - const RunnerConfig& runner_config, + RunnerConfig runner_config, std::unique_ptr<::tokenizers::Tokenizer> tokenizer) { - // Sanity check tokenizer if (!tokenizer || !tokenizer->is_loaded()) { ET_LOG(Error, "Tokenizer is null or not loaded"); return nullptr; } - // Create the Module - std::unique_ptr module; + std::unique_ptr module; if (runner_paths.data_path.has_value()) { - module = std::make_unique( - runner_paths.model_path, runner_paths.data_path.value(), Module::LoadMode::File); + module = std::make_unique( + runner_paths.model_path, + runner_paths.data_path.value(), + extension::Module::LoadMode::File); } else { - module = std::make_unique(runner_paths.model_path, Module::LoadMode::File); + module = std::make_unique( + runner_paths.model_path, extension::Module::LoadMode::File); } - // Get methods executorch::runtime::Result> methods_res = module->method_names(); if (methods_res.error() != executorch::runtime::Error::Ok) { @@ -515,12 +650,11 @@ std::unique_ptr create_llm_runner( auto methods = methods_res.get(); std::vector register_methods; - if (runner_config.decode_only) { + if (runner_config.is_decode_only) { register_methods.emplace_back(default_method_name_decode); } else { for (auto method : methods) { if (kMetadataMap.find(method) != kMetadataMap.end()) { - ET_LOG(Info, "It is not method(%s)", method.c_str()); continue; } ET_LOG(Info, "Register method \"%s\" to decoder", method.c_str()); @@ -528,45 +662,47 @@ std::unique_ptr create_llm_runner( } } - // Get metadata from Module ET_LOG(Info, "Reading metadata from model"); - auto metadata = get_llm_metadata(tokenizer.get(), module.get(), runner_config.use_gemma); + auto metadata = + get_llm_metadata(tokenizer.get(), module.get(), runner_config); - // Create IO Manager example::LlmIOManager::Config conf; conf.n_layers = metadata.at(example::kNumOfLayers); + conf.n_heads = runner_config.n_heads; conf.n_kv_heads = metadata.at(example::kNumOfKVHeads); - conf.max_seq_len = metadata.at(llm::kMaxSeqLen); - conf.decode_seq_len = 1; - conf.prefill_seq_len = (metadata.at(example::kPrefillArLen) > 1) ? - (metadata.at(example::kPrefillArLen)) : - (metadata.at(llm::kMaxSeqLen)); - conf.use_gemma = runner_config.use_gemma; - - if (runner_config.decode_only) { + conf.hidden_dim = runner_config.hidden_dim; + conf.vocab_size = runner_config.vocab_size; + conf.max_seq_len = runner_config.max_seq_len; + conf.max_context_len = runner_config.max_context_len; + conf.decode_seq_len = runner_config.decode_seq_len; + conf.prefill_seq_len = (metadata.at(example::kPrefillArLen) > 1) + ? (metadata.at(example::kPrefillArLen)) + : (runner_config.prefill_seq_len != 0 ? runner_config.prefill_seq_len + : metadata.at(llm::kMaxSeqLen)); + conf.embedding_strategy = runner_config.embedding_strategy; + conf.masking_strategy = runner_config.masking_strategy; + conf.attn_mask_rows = runner_config.attn_mask_rows; + conf.attn_mask_cols = runner_config.attn_mask_cols; + conf.sliding_window = runner_config.sliding_window; + + if (runner_config.is_decode_only) { conf.prefill_seq_len = 1; } - // Try to get head_dim from metadata, otherwise calculate from dim/n_heads if (metadata.count(llm::kHeadDim)) { conf.head_dim = metadata.at(llm::kHeadDim); } else { - size_t dim = metadata.at(example::kNumOfHiddenDim); - size_t n_heads = metadata.at(example::kNumOfHeads); - if (runner_config.use_gemma) { - conf.head_dim = 256; - } else { - conf.head_dim = dim / n_heads; - } + conf.head_dim = runner_config.head_dim; } auto io_manager = std::make_unique(*module.get(), conf); if (runner_paths.pos_embedding_path.has_value()) { - if (io_manager->load_rotary_embedding(runner_paths.pos_embedding_path.value()) != - executorch::runtime::Error::Ok) { - ET_LOG(Error, "load pos_embedding file failed !"); + if (io_manager->load_rotary_embedding( + runner_paths.pos_embedding_path.value()) != + executorch::runtime::Error::Ok) { + ET_LOG(Error, "load pos_embedding file failed !"); } } @@ -578,26 +714,25 @@ std::unique_ptr create_llm_runner( auto text_decoder_runner = std::make_unique(module.get(), io_manager.get()); - // Register methods to cope with multi-methods for (auto method : register_methods) { - text_decoder_runner->register_method(method); + text_decoder_runner->set_method(method); } - // TODO(sicheon.oh) : Need to find proper flag to distinguish separaged prefill/decode opion. + // TODO(brucekkim): Refactor to use dynamic runtime checks (e.g., tracking the + // loop step counter). Relying solely on configuration properties will fail if + // a single-token prompt is processed. bool separated_prefill_decode = (conf.prefill_seq_len > 1); ET_LOG( Info, "separated_prefill_decode(%s)", - (separated_prefill_decode)?"True":"False"); + (separated_prefill_decode) ? "True" : "False"); - // Create text_prefiller auto text_prefiller = std::make_unique( text_decoder_runner.get(), metadata.at(llm::kUseKVCache), metadata.at(llm::kEnableDynamicShape) | (separated_prefill_decode), conf.prefill_seq_len); - // Create text_token_generator with stats auto stats = std::make_unique(); auto text_token_generator = std::make_unique( tokenizer.get(), @@ -606,7 +741,6 @@ std::unique_ptr create_llm_runner( std::move(eos_ids), stats.get()); - // Create and return the Runner instance std::unique_ptr runner = std::make_unique( std::move(metadata), std::move(tokenizer), diff --git a/examples/samsung/executor_runner/llama/runner/runner.h b/examples/samsung/executor_runner/llama/runner/runner.h index adf444838c1..ebacd7a5fb0 100644 --- a/examples/samsung/executor_runner/llama/runner/runner.h +++ b/examples/samsung/executor_runner/llama/runner/runner.h @@ -16,12 +16,15 @@ #include #include +#include #include #include namespace example { namespace llm = ::executorch::extension::llm; +namespace runtime = ::executorch::runtime; +namespace extension = ::executorch::extension; /** * @brief Holds path-related configuration for creating a llama runner. @@ -37,10 +40,32 @@ struct RunnerPaths { * @brief Holds non-path configuration for creating a llama runner. */ struct RunnerConfig { - std::string tokenizer_type; - bool decode_only = false; - bool use_gemma = false; + std::string config_path; + bool is_decode_only = false; float temperature = -1.0f; + + // Strategies + DecoderModelVersion model_version = DecoderModelVersion::kLlama3; + EmbeddingStrategy embedding_strategy = EmbeddingStrategy::kGlobalOnly; + MaskingStrategy masking_strategy = MaskingStrategy::kCausal; + + // Parsed from JSON + int64_t hidden_dim = 0; + int64_t n_heads = 0; + int64_t n_kv_heads = 0; + int64_t n_layers = 0; + int64_t head_dim = 0; + int64_t vocab_size = 0; + int64_t max_seq_len = 128; + int64_t max_context_len = 128; + int64_t prefill_seq_len = 128; + int64_t decode_seq_len = 1; + int64_t attn_mask_rows = 128; + int64_t attn_mask_cols = 1024; + std::optional sliding_window; + std::string tokenizer_type; + + void load_from_json(); }; class LLMRunner : public llm::IRunner { @@ -70,7 +95,7 @@ class LLMRunner : public llm::IRunner { explicit LLMRunner( std::unordered_map metadata, std::unique_ptr<::tokenizers::Tokenizer> tokenizer, - std::unique_ptr<::executorch::extension::Module> module, + std::unique_ptr module, std::unique_ptr text_decoder_runner, std::unique_ptr text_prefiller, std::unique_ptr text_token_generator, @@ -92,7 +117,7 @@ class LLMRunner : public llm::IRunner { * * @return ::executorch::runtime::Error Success or error status */ - ::executorch::runtime::Error load(); + runtime::Error load(); /** * @brief Generates text based on the provided prompt * @@ -109,7 +134,7 @@ class LLMRunner : public llm::IRunner { * @param stats_callback Function called with performance statistics * @return ::executorch::runtime::Error Success or error status */ - ::executorch::runtime::Error generate( + runtime::Error generate( const std::string& prompt, const llm::GenerationConfig& config, std::function token_callback = {}, @@ -132,7 +157,7 @@ class LLMRunner : public llm::IRunner { * @param stats_callback Function called with performance statistics * @return ::executorch::runtime::Error Success or error status */ - ::executorch::runtime::Error generate_from_pos( + runtime::Error generate_from_pos( const std::string& prompt, int64_t start_pos, const llm::GenerationConfig& config, @@ -149,9 +174,7 @@ class LLMRunner : public llm::IRunner { * @param max_new_tokens Maximum number of tokens to generate during warmup * @return ::executorch::runtime::Error Success or error status */ - ::executorch::runtime::Error warmup( - const std::string& prompt, - int32_t max_new_tokens); + runtime::Error warmup(const std::string& prompt, int32_t max_new_tokens); /** * @brief Stops the ongoing text generation process * @@ -160,6 +183,19 @@ class LLMRunner : public llm::IRunner { */ void stop(); + /** + * @brief Formats the prompt based on the model version. + * + * @param prompt The input text + * @param system_prompt The system prompt (optional) + * @param version The version of the decoder model + * @return Formatted prompt string + */ + static std::string get_formatted_prompt( + const std::string& prompt, + const std::string& system_prompt, + DecoderModelVersion version); + /** * @brief Force remove prefilled tokens and reset KV cache start position * @@ -174,7 +210,7 @@ class LLMRunner : public llm::IRunner { // Components std::unique_ptr<::tokenizers::Tokenizer> tokenizer_; std::unordered_map metadata_; - std::unique_ptr<::executorch::extension::Module> + std::unique_ptr module_; // Manage module's lifecycle, make sure it outlives // text_decoder_runner_llm::. std::unique_ptr @@ -191,10 +227,11 @@ class LLMRunner : public llm::IRunner { // temperature. // Deprecated, we should rely on the temperature in GenerationConfig instead. float temperature_ = -1.0f; + DecoderModelVersion model_version_ = DecoderModelVersion::kLlama3; }; std::unique_ptr create_llama_runner( const RunnerPaths& paths, - const RunnerConfig& config); + RunnerConfig config); } // namespace example diff --git a/examples/samsung/executor_runner/llama/runner/token_generator.h b/examples/samsung/executor_runner/llama/runner/token_generator.h index dcb3c8c5d75..a44e781227f 100644 --- a/examples/samsung/executor_runner/llama/runner/token_generator.h +++ b/examples/samsung/executor_runner/llama/runner/token_generator.h @@ -16,14 +16,19 @@ namespace example { +namespace llm = ::executorch::extension::llm; +namespace runtime = ::executorch::runtime; +namespace aten = ::executorch::aten; +namespace extension = ::executorch::extension; + class ET_EXPERIMENTAL TokenGenerator { public: TokenGenerator( ::tokenizers::Tokenizer* tokenizer, - executorch::extension::llm::TextDecoderRunner* text_decoder_runner, + llm::TextDecoderRunner* text_decoder_runner, bool use_kv_cache, std::unique_ptr>&& eos_ids, - executorch::extension::llm::Stats* stats) + llm::Stats* stats) : tokenizer_(tokenizer), text_decoder_runner_(text_decoder_runner), eos_ids_(std::move(eos_ids)), @@ -46,7 +51,7 @@ class ET_EXPERIMENTAL TokenGenerator { * @param token_callback what to do after a token is generated. * @return how many tokens are generated. */ - inline ::executorch::runtime::Result generate( + inline runtime::Result generate( std::vector tokens, int64_t start_pos, int32_t max_new_tokens, @@ -74,8 +79,8 @@ class ET_EXPERIMENTAL TokenGenerator { } // initialize tensor wrappers - auto tokens_managed = executorch::extension::from_blob( - token_data.data(), token_shape, executorch::aten::ScalarType::Long); + auto tokens_managed = extension::from_blob( + token_data.data(), token_shape, aten::ScalarType::Long); should_stop_ = false; @@ -103,7 +108,7 @@ class ET_EXPERIMENTAL TokenGenerator { } else { // push it to the back token_data.push_back(cur_token); - ET_CHECK_OK_OR_RETURN_ERROR(executorch::extension::resize_tensor_ptr( + ET_CHECK_OK_OR_RETURN_ERROR(extension::resize_tensor_ptr( tokens_managed, {1, static_cast(token_data.size())})); } @@ -135,9 +140,9 @@ class ET_EXPERIMENTAL TokenGenerator { * Load the necessary resources for TokenGenerator. * This method should be called before using the generate() method. */ - ::executorch::runtime::Error load() { + runtime::Error load() { // return text_decoder_runner_->load(); - return ::executorch::runtime::Error::Ok; + return runtime::Error::Ok; } /** @@ -157,7 +162,7 @@ class ET_EXPERIMENTAL TokenGenerator { * pointers remain valid for the duration of its use. */ ::tokenizers::Tokenizer* tokenizer_; - executorch::extension::llm::TextDecoderRunner* text_decoder_runner_; + llm::TextDecoderRunner* text_decoder_runner_; std::unique_ptr> eos_ids_; bool use_kv_cache_; @@ -165,7 +170,7 @@ class ET_EXPERIMENTAL TokenGenerator { bool should_stop_ = false; // stats - executorch::extension::llm::Stats* stats_; + llm::Stats* stats_; }; } // namespace example diff --git a/examples/samsung/executor_runner/run_gemma3_sample.sh b/examples/samsung/executor_runner/run_gemma3_sample.sh index 0de3a7199e4..f3fbcf3edc5 100644 --- a/examples/samsung/executor_runner/run_gemma3_sample.sh +++ b/examples/samsung/executor_runner/run_gemma3_sample.sh @@ -14,16 +14,16 @@ TEMPERATURE=0 NUM_BOS=1 INPUT_PROMPT="What is the capital of France?" -chmod +x enn_executor_llama_runner +chmod +x enn_executor_llm_runner -./enn_executor_llama_runner \ +./enn_executor_llm_runner \ --model_path=$MODEL_PATH \ --data_path=$WEIGHT_SHARED_MODEL_PATH \ --seq_len=$SEQ_LEN \ --temperature=$TEMPERATURE \ - --pos_embedding_path=$POS_EMBEDDING_PATH \ --num_bos=$NUM_BOS \ - --tokenizer_type="hf" \ + --pos_embedding_path=$POS_EMBEDDING_PATH \ + --tokenizer_type="$TOKENIZER_TYPE" \ --tokenizer_path=$TOKENIZER_PATH \ --prompt="$INPUT_PROMPT" \ - --use_gemma + --config_path "$CONFIG_PATH" diff --git a/examples/samsung/executor_runner/run_llama3_sample.sh b/examples/samsung/executor_runner/run_llama3_sample.sh index 18f608e4f88..910d9c204bb 100644 --- a/examples/samsung/executor_runner/run_llama3_sample.sh +++ b/examples/samsung/executor_runner/run_llama3_sample.sh @@ -6,22 +6,25 @@ MODEL_PATH=/data/local/tmp/llama3/llama3.pte WEIGHT_SHARED_MODEL_PATH=/data/local/tmp/llama3/weights.ptd -TOKENIZER_PATH=/data/local/tmp/llama3/tokenizer.model POS_EMBEDDING_PATH=/data/local/tmp/llama3/pos_embedding.bin +TOKENIZER_PATH=/data/local/tmp/llama3/tokenizer.model +CONFIG_PATH="/data/local/tmp/exynos_llama3_1b.json" SEQ_LEN=1024 TEMPERATURE=0 +NUM_BOS=1op INPUT_PROMPT="What is the capital of France?" -chmod +x enn_executor_llama_runner +chmod +x enn_executor_llm_runner -./enn_executor_llama_runner \ +./enn_executor_llm_runner \ --model_path=$MODEL_PATH \ --data_path=$WEIGHT_SHARED_MODEL_PATH \ - --seq_len=$SEQ_LEN \ - --temperature=$TEMPERATURE \ - --num_bos=1 \ --pos_embedding_path=$POS_EMBEDDING_PATH \ - --tokenizer_type="tiktoken" \ --tokenizer_path=$TOKENIZER_PATH \ - --prompt="$INPUT_PROMPT" + --config_path "$CONFIG_PATH" \ + --prompt="$INPUT_PROMPT" \ + --temperature=$TEMPERATURE \ + --seq_len=$SEQ_LEN \ + --num_bos=$NUM_BOS \ + diff --git a/extension/llm/runner/io_manager/io_manager.h b/extension/llm/runner/io_manager/io_manager.h index 0a257064b77..fc9a8f0641b 100644 --- a/extension/llm/runner/io_manager/io_manager.h +++ b/extension/llm/runner/io_manager/io_manager.h @@ -37,19 +37,6 @@ class ET_EXPERIMENTAL IOManager { */ virtual ~IOManager() = default; - /** - * @brief Get the use_gemma configuration flag. - * @return True if using Gemma model specific processing, false by default. - * Derived classes should override this to return their specific value. - */ - virtual bool use_gemma() const { - return false; - } - - virtual bool is_quantized() const { - return false; - } - /** * @brief Load the IO manager with method metadata for prefill and * decode operations. diff --git a/extension/llm/runner/text_decoder_runner.cpp b/extension/llm/runner/text_decoder_runner.cpp index 893bbbd60bc..25f3fc234ad 100644 --- a/extension/llm/runner/text_decoder_runner.cpp +++ b/extension/llm/runner/text_decoder_runner.cpp @@ -48,7 +48,11 @@ ::executorch::runtime::Result TextDecoderRunner::step( if (use_kv_cache) { auto start_pos_tensor = ET_UNWRAP(populate_start_pos_or_cache_position( - module_, start_pos, cache_positions, tokens->numel(), method_name_.c_str())); + module_, + start_pos, + cache_positions, + tokens->numel(), + method_name_.c_str())); std::vector inputs; auto inputs_res = io_manager_->prepare_decode(tokens, start_pos_tensor); From 6519d58bca96015ce410fd77933e07d1b1ba0763 Mon Sep 17 00:00:00 2001 From: Jiseong Oh Date: Tue, 13 Oct 2026 19:39:28 +0900 Subject: [PATCH 10/17] Fix Exynos LLM runner bugs and rename it to llm runner EnnLogImpl no longer crashes now that it uses __android_log_vprint, a failed rotary embedding load returns an explicit error, and build and prompt template bugs are fixed. The Samsung llama runner is renamed to llm runner. Signed-off-by: Jiseong Oh Co-authored-by: Hoon Choi --- backends/samsung/runtime/logging.cpp | 6 ++- examples/samsung/CMakeLists.txt | 14 +++--- .../enn_executor_llm_runner.cpp | 10 ++--- .../{llama => llm}/runner/CMakeLists.txt | 22 +++++----- .../{llama => llm}/runner/constants.h | 0 .../{llama => llm}/runner/decoder.cpp | 2 +- .../{llama => llm}/runner/decoder.h | 2 +- .../{llama => llm}/runner/io_manager.cpp | 28 +++++++----- .../{llama => llm}/runner/io_manager.h | 6 +-- .../{llama => llm}/runner/prefiller.cpp | 2 +- .../{llama => llm}/runner/prefiller.h | 0 .../runner/rotary_embedding.cpp | 7 ++- .../{llama => llm}/runner/rotary_embedding.h | 3 +- .../{llama => llm}/runner/runner.cpp | 44 ++++++++++++------- .../{llama => llm}/runner/runner.h | 16 ++++--- .../{llama => llm}/runner/token_generator.h | 0 .../executor_runner/run_gemma3_sample.sh | 2 +- .../executor_runner/run_llama3_sample.sh | 2 +- 18 files changed, 98 insertions(+), 68 deletions(-) rename examples/samsung/executor_runner/{llama => llm}/runner/CMakeLists.txt (84%) rename examples/samsung/executor_runner/{llama => llm}/runner/constants.h (100%) rename examples/samsung/executor_runner/{llama => llm}/runner/decoder.cpp (99%) rename examples/samsung/executor_runner/{llama => llm}/runner/decoder.h (95%) rename examples/samsung/executor_runner/{llama => llm}/runner/io_manager.cpp (96%) rename examples/samsung/executor_runner/{llama => llm}/runner/io_manager.h (98%) rename examples/samsung/executor_runner/{llama => llm}/runner/prefiller.cpp (98%) rename examples/samsung/executor_runner/{llama => llm}/runner/prefiller.h (100%) rename examples/samsung/executor_runner/{llama => llm}/runner/rotary_embedding.cpp (93%) rename examples/samsung/executor_runner/{llama => llm}/runner/rotary_embedding.h (91%) rename examples/samsung/executor_runner/{llama => llm}/runner/runner.cpp (94%) rename examples/samsung/executor_runner/{llama => llm}/runner/runner.h (93%) rename examples/samsung/executor_runner/{llama => llm}/runner/token_generator.h (100%) diff --git a/backends/samsung/runtime/logging.cpp b/backends/samsung/runtime/logging.cpp index 0a2a6e64282..d3f98111a9b 100644 --- a/backends/samsung/runtime/logging.cpp +++ b/backends/samsung/runtime/logging.cpp @@ -35,7 +35,11 @@ void EnnLogImpl(ENN_LOG_LEVEL log_level, const char* format, ...) { android_severity = ANDROID_LOG_UNKNOWN; break; } - __android_log_print(android_severity, "[Exynos ExecuTorch]", format, args); + va_list android_args; + va_copy(android_args, args); + __android_log_vprint( + android_severity, "[Exynos ExecuTorch]", format, android_args); + va_end(android_args); #endif const char* serverity_name; switch (log_level) { diff --git a/examples/samsung/CMakeLists.txt b/examples/samsung/CMakeLists.txt index 6aca9dfe853..c8b4d0992ae 100644 --- a/examples/samsung/CMakeLists.txt +++ b/examples/samsung/CMakeLists.txt @@ -95,17 +95,17 @@ set_target_properties( enn_executor_runner PROPERTIES CXX_VISIBILITY_PRESET hidden ) -# llama_runner library -add_subdirectory(executor_runner/llama/runner) +# LLM runner library +add_subdirectory(executor_runner/llm/runner) -target_link_directories(enn_llama_runner PRIVATE) +target_link_directories(enn_llm_runner PRIVATE) -add_library(lib_enn_llama_runner STATIC IMPORTED) +add_library(lib_enn_llm_runner STATIC IMPORTED) set_target_properties( - lib_enn_llama_runner + lib_enn_llm_runner PROPERTIES IMPORTED_LOCATION - "${EXECUTORCH_ROOT}/build_samsung_android/examples/samsung/executor_runner/llama/runner/libenn_llama_runner.a" + "${EXECUTORCH_ROOT}/build_samsung_android/examples/samsung/executor_runner/llm/runner/libenn_llm_runner.a" ) add_library(lib_tokenizer STATIC IMPORTED) @@ -193,7 +193,7 @@ list(APPEND xnnpack_backend_libs quantized_kernels quantized_ops_lib) target_link_libraries( enn_executor_llm_runner - PUBLIC enn_llama_runner + PUBLIC enn_llm_runner PRIVATE log gflags android diff --git a/examples/samsung/executor_runner/enn_executor_llm_runner.cpp b/examples/samsung/executor_runner/enn_executor_llm_runner.cpp index e06cb40d5a6..43549bd32b6 100644 --- a/examples/samsung/executor_runner/enn_executor_llm_runner.cpp +++ b/examples/samsung/executor_runner/enn_executor_llm_runner.cpp @@ -8,7 +8,7 @@ #include -#include +#include #if defined(ET_USE_THREADPOOL) #include @@ -19,7 +19,7 @@ namespace llm = ::executorch::extension::llm; DEFINE_string( model_path, - "llama2.pte", + "model.pte", "Model serialized in flatbuffer format."); DEFINE_string(data_path, "", "Data file for the model."); @@ -114,8 +114,8 @@ int32_t main(int32_t argc, char** argv) { example::RunnerConfig runner_config{ .config_path = config_path, .is_decode_only = is_decode_only}; std::unique_ptr runner = - example::create_llama_runner(runner_paths, runner_config); - ET_CHECK_MSG(runner != nullptr, "Failed to create llama runner"); + example::create_llm_runner(runner_paths, runner_config); + ET_CHECK_MSG(runner != nullptr, "Failed to create LLM runner"); // TODO: support warmup interface @@ -127,7 +127,7 @@ int32_t main(int32_t argc, char** argv) { auto error = runner->generate(prompt, config, {}, {}); ET_CHECK_MSG( error == executorch::runtime::Error::Ok, - "Failed to generate token from llama runner"); + "Failed to generate token from LLM runner"); return 0; } diff --git a/examples/samsung/executor_runner/llama/runner/CMakeLists.txt b/examples/samsung/executor_runner/llm/runner/CMakeLists.txt similarity index 84% rename from examples/samsung/executor_runner/llama/runner/CMakeLists.txt rename to examples/samsung/executor_runner/llm/runner/CMakeLists.txt index 2bf89b654f8..9cfbf858199 100644 --- a/examples/samsung/executor_runner/llama/runner/CMakeLists.txt +++ b/examples/samsung/executor_runner/llm/runner/CMakeLists.txt @@ -4,7 +4,7 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. -message("###### Build ENN llama runner ######") +message("###### Build ENN LLM runner ######") message(ROOT ${CMAKE_CURRENT_SOURCE_DIR}/../../../../..) message(EXECUTORCH_ROOT ${EXECUTORCH_ROOT}) message(CMAKE_CURRENT_SOURCE_DIR ${CMAKE_CURRENT_SOURCE_DIR}) @@ -14,12 +14,12 @@ message(CMAKE_CURRENT_LIST_DIR ${CMAKE_CURRENT_LIST_DIR}) include(${CMAKE_CURRENT_SOURCE_DIR}/../../../../../tools/cmake/Utils.cmake) include(${CMAKE_CURRENT_SOURCE_DIR}/../../../../../tools/cmake/Codegen.cmake) -list(TRANSFORM _enn_llama_runner__srcs +list(TRANSFORM _enn_llm_runner__srcs PREPEND "${CMAKE_CURRENT_SOURCE_DIR}/../../../../.." ) list( PREPEND - _enn_llama_runner__srcs + _enn_llm_runner__srcs ${CMAKE_CURRENT_LIST_DIR}/runner.cpp ${CMAKE_CURRENT_LIST_DIR}/decoder.cpp ${CMAKE_CURRENT_LIST_DIR}/prefiller.cpp @@ -54,9 +54,9 @@ if(CMAKE_TOOLCHAIN_IOS ) # Building a share library on iOS requires code signing On Android we see # duplicated registration when using shared lib - add_library(enn_llama_runner STATIC ${_enn_llama_runner__srcs}) + add_library(enn_llm_runner STATIC ${_enn_llm_runner__srcs}) else() - add_library(enn_llama_runner SHARED ${_enn_llama_runner__srcs}) + add_library(enn_llm_runner SHARED ${_enn_llm_runner__srcs}) endif() # For extension_llm_runner @@ -67,17 +67,17 @@ if(NOT TARGET extension_llm_runner) ) endif() -set(enn_llama_runner_deps +set(enn_llm_runner_deps executorch_core extension_data_loader extension_module extension_tensor extension_flat_tensor lib_extension_llm_runner ) -target_link_libraries(enn_llama_runner PUBLIC ${enn_llama_runner_deps}) -target_link_libraries(enn_llama_runner PUBLIC tokenizers::tokenizers) -target_link_libraries(enn_llama_runner PUBLIC lib_tokenizer) +target_link_libraries(enn_llm_runner PUBLIC ${enn_llm_runner_deps}) +target_link_libraries(enn_llm_runner PUBLIC tokenizers::tokenizers) +target_link_libraries(enn_llm_runner PUBLIC lib_tokenizer) target_include_directories( - enn_llama_runner + enn_llm_runner PUBLIC ${EXECUTORCH_ROOT}/.. ${CMAKE_CURRENT_SOURCE_DIR}/../../../../../extension/llm/tokenizers/include @@ -85,6 +85,6 @@ target_include_directories( ${CMAKE_CURRENT_SOURCE_DIR}/../../../../../extension/llm/tokenizers/third-party/abseil-cpp ${CMAKE_CURRENT_SOURCE_DIR}/../../../../../extension/llm/tokenizers/third-party/json/include ${CMAKE_CURRENT_SOURCE_DIR}/../../../../../examples/models/llama/tokenizer - ${CMAKE_CURRENT_SOURCE_DIR}/../../../../../examples/samsung/executor_runner/llama/runner + ${CMAKE_CURRENT_SOURCE_DIR}/../../../../../examples/samsung/executor_runner/llm/runner ${CMAKE_CURRENT_SOURCE_DIR}/../../../../../extension/llm/runner ) diff --git a/examples/samsung/executor_runner/llama/runner/constants.h b/examples/samsung/executor_runner/llm/runner/constants.h similarity index 100% rename from examples/samsung/executor_runner/llama/runner/constants.h rename to examples/samsung/executor_runner/llm/runner/constants.h diff --git a/examples/samsung/executor_runner/llama/runner/decoder.cpp b/examples/samsung/executor_runner/llm/runner/decoder.cpp similarity index 99% rename from examples/samsung/executor_runner/llama/runner/decoder.cpp rename to examples/samsung/executor_runner/llm/runner/decoder.cpp index 0dc1ce2a826..6f5a32a9dcc 100644 --- a/examples/samsung/executor_runner/llama/runner/decoder.cpp +++ b/examples/samsung/executor_runner/llm/runner/decoder.cpp @@ -13,7 +13,7 @@ #include -#include +#include #include #include diff --git a/examples/samsung/executor_runner/llama/runner/decoder.h b/examples/samsung/executor_runner/llm/runner/decoder.h similarity index 95% rename from examples/samsung/executor_runner/llama/runner/decoder.h rename to examples/samsung/executor_runner/llm/runner/decoder.h index 688b27b3c2b..6625d2fe98f 100644 --- a/examples/samsung/executor_runner/llama/runner/decoder.h +++ b/examples/samsung/executor_runner/llm/runner/decoder.h @@ -12,7 +12,7 @@ #include -#include +#include #include #include #include diff --git a/examples/samsung/executor_runner/llama/runner/io_manager.cpp b/examples/samsung/executor_runner/llm/runner/io_manager.cpp similarity index 96% rename from examples/samsung/executor_runner/llama/runner/io_manager.cpp rename to examples/samsung/executor_runner/llm/runner/io_manager.cpp index 19251a79c5d..0a165b6f8af 100644 --- a/examples/samsung/executor_runner/llama/runner/io_manager.cpp +++ b/examples/samsung/executor_runner/llm/runner/io_manager.cpp @@ -8,8 +8,8 @@ #include -#include -#include +#include +#include namespace example { @@ -134,7 +134,8 @@ LlmIOManager::prepare_prefill( break; } - prepare_rotary_embedding_inputs(config_.prefill_seq_len, inputs); + ET_CHECK_OK_OR_RETURN_ERROR( + prepare_rotary_embedding_inputs(config_.prefill_seq_len, inputs)); prefill_k_input_tensors_.clear(); prefill_v_input_tensors_.clear(); @@ -195,7 +196,8 @@ LlmIOManager::prepare_decode( pos_ids_.clear(); pos_ids_.resize(config_.decode_seq_len); pos_ids_[0] = current_pos_; - prepare_rotary_embedding_inputs(config_.decode_seq_len, inputs); + ET_CHECK_OK_OR_RETURN_ERROR( + prepare_rotary_embedding_inputs(config_.decode_seq_len, inputs)); decode_k_input_tensors_.clear(); decode_v_input_tensors_.clear(); @@ -504,7 +506,7 @@ void LlmIOManager::dump_kv_output( } } -void LlmIOManager::prepare_rotary_embedding_inputs( +executorch::runtime::Error LlmIOManager::prepare_rotary_embedding_inputs( size_t seq_len, std::vector& inputs) { const size_t element_size = seq_len * config_.head_dim; @@ -519,13 +521,15 @@ void LlmIOManager::prepare_rotary_embedding_inputs( break; } - if (rotary_embedding_->update( - pos_ids_, - global_cos_data_.data(), - global_sin_data_.data(), - local_cos_data_.data(), - local_sin_data_.data()) != executorch::runtime::Error::Ok) { + auto update_err = rotary_embedding_->update( + pos_ids_, + global_cos_data_.data(), + global_sin_data_.data(), + local_cos_data_.data(), + local_sin_data_.data()); + if (update_err != executorch::runtime::Error::Ok) { ET_LOG(Error, "update pos_embedding failed !"); + return update_err; } global_cos_tensor_ = executorch::extension::from_blob( @@ -565,6 +569,8 @@ void LlmIOManager::prepare_rotary_embedding_inputs( inputs.emplace_back(local_sin_tensor_); break; } + + return executorch::runtime::Error::Ok; } } // namespace example diff --git a/examples/samsung/executor_runner/llama/runner/io_manager.h b/examples/samsung/executor_runner/llm/runner/io_manager.h similarity index 98% rename from examples/samsung/executor_runner/llama/runner/io_manager.h rename to examples/samsung/executor_runner/llm/runner/io_manager.h index ef13283dfa8..99887f8506d 100644 --- a/examples/samsung/executor_runner/llama/runner/io_manager.h +++ b/examples/samsung/executor_runner/llm/runner/io_manager.h @@ -11,8 +11,8 @@ #include #include -#include -#include +#include +#include #include namespace example { @@ -419,7 +419,7 @@ class ET_EXPERIMENTAL LlmIOManager : public llm::IOManager { * @param seq_len The sequence length to use for the embedding computation. * @param inputs Vector of EValue inputs to append the embedding tensors to. */ - void prepare_rotary_embedding_inputs( + ET_NODISCARD runtime::Error prepare_rotary_embedding_inputs( size_t seq_len, std::vector& inputs); diff --git a/examples/samsung/executor_runner/llama/runner/prefiller.cpp b/examples/samsung/executor_runner/llm/runner/prefiller.cpp similarity index 98% rename from examples/samsung/executor_runner/llama/runner/prefiller.cpp rename to examples/samsung/executor_runner/llm/runner/prefiller.cpp index 399afa271d8..07b3e52d50a 100644 --- a/examples/samsung/executor_runner/llama/runner/prefiller.cpp +++ b/examples/samsung/executor_runner/llm/runner/prefiller.cpp @@ -9,7 +9,7 @@ // Given a text prompt, encode it using tokenizer and prefill the KV cache of a // LLM. -#include +#include #include namespace example { diff --git a/examples/samsung/executor_runner/llama/runner/prefiller.h b/examples/samsung/executor_runner/llm/runner/prefiller.h similarity index 100% rename from examples/samsung/executor_runner/llama/runner/prefiller.h rename to examples/samsung/executor_runner/llm/runner/prefiller.h diff --git a/examples/samsung/executor_runner/llama/runner/rotary_embedding.cpp b/examples/samsung/executor_runner/llm/runner/rotary_embedding.cpp similarity index 93% rename from examples/samsung/executor_runner/llama/runner/rotary_embedding.cpp rename to examples/samsung/executor_runner/llm/runner/rotary_embedding.cpp index 6884f4bf5a4..a2d833f9e37 100644 --- a/examples/samsung/executor_runner/llama/runner/rotary_embedding.cpp +++ b/examples/samsung/executor_runner/llm/runner/rotary_embedding.cpp @@ -6,7 +6,7 @@ * LICENSE file in the root directory of this source tree. */ -#include +#include #include #include #include @@ -70,6 +70,7 @@ runtime::Error RotaryEmbedding::load( } file.close(); + loaded_ = true; return runtime::Error::Ok; } @@ -79,6 +80,10 @@ runtime::Error RotaryEmbedding::update( float* global_sin, float* local_cos, float* local_sin) { + if (!loaded_) { + ET_LOG(Error, "RotaryEmbedding::update() called before load() succeeded."); + return runtime::Error::InvalidState; + } if (pos_ids.empty()) { return runtime::Error::InvalidArgument; } diff --git a/examples/samsung/executor_runner/llama/runner/rotary_embedding.h b/examples/samsung/executor_runner/llm/runner/rotary_embedding.h similarity index 91% rename from examples/samsung/executor_runner/llama/runner/rotary_embedding.h rename to examples/samsung/executor_runner/llm/runner/rotary_embedding.h index 3a0adf84dc0..5b9337863e4 100644 --- a/examples/samsung/executor_runner/llama/runner/rotary_embedding.h +++ b/examples/samsung/executor_runner/llm/runner/rotary_embedding.h @@ -8,7 +8,7 @@ #pragma once -#include +#include #include #include #include @@ -40,6 +40,7 @@ class RotaryEmbedding { private: Config config_; + bool loaded_ = false; std::unique_ptr global_cos_; std::unique_ptr global_sin_; std::unique_ptr local_cos_; diff --git a/examples/samsung/executor_runner/llama/runner/runner.cpp b/examples/samsung/executor_runner/llm/runner/runner.cpp similarity index 94% rename from examples/samsung/executor_runner/llama/runner/runner.cpp rename to examples/samsung/executor_runner/llm/runner/runner.cpp index 835a35290e6..935e0e73505 100644 --- a/examples/samsung/executor_runner/llama/runner/runner.cpp +++ b/examples/samsung/executor_runner/llm/runner/runner.cpp @@ -6,7 +6,7 @@ * LICENSE file in the root directory of this source tree. */ -// A simple llama2 runner that includes preprocessing and post processing logic. +// A simple LLM runner that includes preprocessing and post processing logic. // The module takes in a string as input and emits a string as output. #include @@ -19,11 +19,11 @@ #include #include -#include -#include -#include -#include -#include +#include +#include +#include +#include +#include #include #include @@ -109,12 +109,12 @@ ET_EXPERIMENTAL std::unique_ptr create_llm_runner( RunnerConfig runner_config, std::unique_ptr<::tokenizers::Tokenizer> tokenizer); -std::unique_ptr create_llama_runner( +std::unique_ptr create_llm_runner( const RunnerPaths& runner_paths, RunnerConfig runner_config) { ET_LOG( Info, - "Creating LLaMa runner: model_path=%s, tokenizer_path=%s", + "Creating LLM runner: model_path=%s, tokenizer_path=%s", runner_paths.model_path.c_str(), runner_paths.tokenizer_path.c_str()); @@ -135,7 +135,7 @@ std::unique_ptr create_llama_runner( } ET_LOG( Info, - "[Create_test_llm_runner] %s %s", + "[create_llm_runner] %s %s", runner_paths.model_path.c_str(), runner_paths.data_path.has_value() ? runner_paths.data_path.value().c_str() @@ -170,6 +170,7 @@ LLMRunner::LLMRunner( std::unique_ptr text_token_generator, std::unique_ptr stats, std::unique_ptr io_manager, + DecoderModelVersion model_version, float temperature) : tokenizer_(std::move(tokenizer)), metadata_(std::move(metadata)), @@ -179,7 +180,8 @@ LLMRunner::LLMRunner( text_token_generator_(std::move(text_token_generator)), io_manager_(std::move(io_manager)), stats_(std::move(stats)), - temperature_(temperature) { + temperature_(temperature), + model_version_(model_version) { // Note: This constructor assumes that text_prefiller and text_token_generator // already have references to the Module and TextDecoderRunner they need } @@ -698,12 +700,21 @@ std::unique_ptr create_llm_runner( auto io_manager = std::make_unique(*module.get(), conf); - if (runner_paths.pos_embedding_path.has_value()) { - if (io_manager->load_rotary_embedding( - runner_paths.pos_embedding_path.value()) != - executorch::runtime::Error::Ok) { - ET_LOG(Error, "load pos_embedding file failed !"); - } + // Every embedding_strategy (kGlobalOnly, kGlobalAndLocal) requires the + // global rotary embedding to be loaded, so a missing/failed load is fatal + // rather than something we can silently continue past. + if (!runner_paths.pos_embedding_path.has_value()) { + ET_LOG(Error, "pos_embedding_path is required to load rotary embedding."); + return nullptr; + } + if (io_manager->load_rotary_embedding( + runner_paths.pos_embedding_path.value()) != + executorch::runtime::Error::Ok) { + ET_LOG( + Error, + "Failed to load pos_embedding file at '%s'.", + runner_paths.pos_embedding_path.value().c_str()); + return nullptr; } auto eos_ids = std::make_unique>( @@ -750,6 +761,7 @@ std::unique_ptr create_llm_runner( std::move(text_token_generator), std::move(stats), std::move(io_manager), + runner_config.model_version, runner_config.temperature); runner->load(); diff --git a/examples/samsung/executor_runner/llama/runner/runner.h b/examples/samsung/executor_runner/llm/runner/runner.h similarity index 93% rename from examples/samsung/executor_runner/llama/runner/runner.h rename to examples/samsung/executor_runner/llm/runner/runner.h index ebacd7a5fb0..9da6e923472 100644 --- a/examples/samsung/executor_runner/llama/runner/runner.h +++ b/examples/samsung/executor_runner/llm/runner/runner.h @@ -16,9 +16,9 @@ #include #include -#include -#include -#include +#include +#include +#include namespace example { @@ -27,7 +27,7 @@ namespace runtime = ::executorch::runtime; namespace extension = ::executorch::extension; /** - * @brief Holds path-related configuration for creating a llama runner. + * @brief Holds path-related configuration for creating an LLM runner. */ struct RunnerPaths { std::string model_path; @@ -37,7 +37,7 @@ struct RunnerPaths { }; /** - * @brief Holds non-path configuration for creating a llama runner. + * @brief Holds non-path configuration for creating an LLM runner. */ struct RunnerConfig { std::string config_path; @@ -89,6 +89,7 @@ class LLMRunner : public llm::IRunner { * @param stats Statistics tracking object for performance monitoring * @param io_manager IOManager manages input/output operations for LLM * inference. + * @param model_version Decoder model family selecting the prompt template * @param temperature Temperature parameter for controlling randomness in * generation (deprecated). Please use GenerationConfig.temperature instead. */ @@ -101,6 +102,7 @@ class LLMRunner : public llm::IRunner { std::unique_ptr text_token_generator, std::unique_ptr stats, std::unique_ptr io_manager, + DecoderModelVersion model_version, float temperature = -1.0f); /** @@ -227,10 +229,10 @@ class LLMRunner : public llm::IRunner { // temperature. // Deprecated, we should rely on the temperature in GenerationConfig instead. float temperature_ = -1.0f; - DecoderModelVersion model_version_ = DecoderModelVersion::kLlama3; + DecoderModelVersion model_version_; }; -std::unique_ptr create_llama_runner( +std::unique_ptr create_llm_runner( const RunnerPaths& paths, RunnerConfig config); diff --git a/examples/samsung/executor_runner/llama/runner/token_generator.h b/examples/samsung/executor_runner/llm/runner/token_generator.h similarity index 100% rename from examples/samsung/executor_runner/llama/runner/token_generator.h rename to examples/samsung/executor_runner/llm/runner/token_generator.h diff --git a/examples/samsung/executor_runner/run_gemma3_sample.sh b/examples/samsung/executor_runner/run_gemma3_sample.sh index f3fbcf3edc5..f60ef702d12 100644 --- a/examples/samsung/executor_runner/run_gemma3_sample.sh +++ b/examples/samsung/executor_runner/run_gemma3_sample.sh @@ -8,6 +8,7 @@ MODEL_PATH=/data/local/tmp/gemma3/gemma3.pte WEIGHT_SHARED_MODEL_PATH=/data/local/tmp/gemma3/weights.ptd TOKENIZER_PATH=/data/local/tmp/gemma3/tokenizer.json POS_EMBEDDING_PATH=/data/local/tmp/gemma3/pos_embedding.bin +CONFIG_PATH="/data/local/tmp/exynos_gemma3_1b.json" SEQ_LEN=1024 TEMPERATURE=0 @@ -23,7 +24,6 @@ chmod +x enn_executor_llm_runner --temperature=$TEMPERATURE \ --num_bos=$NUM_BOS \ --pos_embedding_path=$POS_EMBEDDING_PATH \ - --tokenizer_type="$TOKENIZER_TYPE" \ --tokenizer_path=$TOKENIZER_PATH \ --prompt="$INPUT_PROMPT" \ --config_path "$CONFIG_PATH" diff --git a/examples/samsung/executor_runner/run_llama3_sample.sh b/examples/samsung/executor_runner/run_llama3_sample.sh index 910d9c204bb..da4af40d4d5 100644 --- a/examples/samsung/executor_runner/run_llama3_sample.sh +++ b/examples/samsung/executor_runner/run_llama3_sample.sh @@ -12,7 +12,7 @@ CONFIG_PATH="/data/local/tmp/exynos_llama3_1b.json" SEQ_LEN=1024 TEMPERATURE=0 -NUM_BOS=1op +NUM_BOS=1 INPUT_PROMPT="What is the capital of France?" chmod +x enn_executor_llm_runner From c3d94ad41bce8be06e51d1638f77310226a68fcb Mon Sep 17 00:00:00 2001 From: Jiseong Oh Date: Tue, 13 Oct 2026 19:39:28 +0900 Subject: [PATCH 11/17] Make RunnerConfig the single source of truth and add Gemma 4 parsing RunnerConfig now carries the model arguments, including those parsed from the Gemma 4 JSON. RoPE tables can be loaded from a binary file with precomputation as the fallback, the decoder resolves the position, mask and KV tensor indices and ranks, and the llama3 config and the prefill input resizing are fixed. Signed-off-by: Jiseong Oh Co-authored-by: brucekk.kim --- .../gemma4/config/exynos_gemma4_e2b.json | 69 +++ .../enn_executor_llm_runner.cpp | 12 +- .../executor_runner/llm/runner/CMakeLists.txt | 1 + .../executor_runner/llm/runner/config.cpp | 312 ++++++++++++ .../executor_runner/llm/runner/config.h | 456 ++++++++++++++++++ .../executor_runner/llm/runner/constants.h | 43 -- .../executor_runner/llm/runner/decoder.cpp | 241 ++++++--- .../executor_runner/llm/runner/io_manager.cpp | 317 ++++++++---- .../executor_runner/llm/runner/io_manager.h | 264 ++-------- .../llm/runner/rotary_embedding.cpp | 302 ++++++++++-- .../llm/runner/rotary_embedding.h | 40 +- .../executor_runner/llm/runner/runner.cpp | 287 +++-------- .../executor_runner/llm/runner/runner.h | 201 ++------ 13 files changed, 1690 insertions(+), 855 deletions(-) create mode 100644 examples/models/gemma4/config/exynos_gemma4_e2b.json create mode 100644 examples/samsung/executor_runner/llm/runner/config.cpp create mode 100644 examples/samsung/executor_runner/llm/runner/config.h delete mode 100644 examples/samsung/executor_runner/llm/runner/constants.h diff --git a/examples/models/gemma4/config/exynos_gemma4_e2b.json b/examples/models/gemma4/config/exynos_gemma4_e2b.json new file mode 100644 index 00000000000..62577ea953e --- /dev/null +++ b/examples/models/gemma4/config/exynos_gemma4_e2b.json @@ -0,0 +1,69 @@ +{ + "model_id": "exynos_gemma4_e2b", + "description": "Exynos Optimized Gemma4 E2B Configuration", + "embedding_strategy": "global_and_local", + "masking_strategy": "causal_local", + "model_version": "gemma4", + "tokenizer_type": "hf", + "model_type": "gemma4", + "text_config": { + "hidden_size": 1536, + "intermediate_size": 6144, + "use_double_wide_mlp": true, + "num_hidden_layers": 35, + "num_attention_heads": 8, + "head_dim": 256, + "global_head_dim": 512, + "num_key_value_heads": 1, + "vocab_size": 262144, + "vocab_size_per_layer_input": 262144, + "max_seq_len": 4096, + "max_position_embeddings": 131072, + "rms_norm_eps": 1e-06, + "rope_theta": 1000000.0, + "rope_local_base_freq": 10000.0, + "partial_rotary_factor": 0.25, + "sliding_window": 512, + "hidden_activation": "gelu_pytorch_tanh", + "final_logit_softcapping": 30.0, + "hidden_size_per_layer_input": 256, + "num_kv_shared_layers": 20, + "tie_word_embeddings": true, + "layer_types": [ + "sliding_attention","sliding_attention","sliding_attention","sliding_attention","full_attention", + "sliding_attention","sliding_attention","sliding_attention","sliding_attention","full_attention", + "sliding_attention","sliding_attention","sliding_attention","sliding_attention","full_attention", + "sliding_attention","sliding_attention","sliding_attention","sliding_attention","full_attention", + "sliding_attention","sliding_attention","sliding_attention","sliding_attention","full_attention", + "sliding_attention","sliding_attention","sliding_attention","sliding_attention","full_attention", + "sliding_attention","sliding_attention","sliding_attention","sliding_attention","full_attention" + ] + }, + "audio_config": { + "model_type": "gemma4_audio", + "hidden_size": 1024, + "output_proj_dims": 1536, + "num_attention_heads": 8, + "num_hidden_layers": 12, + "rms_norm_eps": 1e-06, + "conv_kernel_size": 5, + "residual_weight": 0.5, + "attention_chunk_size": 12, + "attention_context_left": 13, + "attention_context_right": 0, + "attention_logit_cap": 50.0 + }, + "audio_seq_length": 750, + "audio_ms_per_token": 40, + "audio_sample_rate": 16000, + "max_audio_length_sec": 30, + "special_tokens": { + "audio_token_id": 258881, + "boa_token_id": 256000, + "eoa_token_id": 258883, + "boi_token_id": 255999, + "eoi_token_id": 258882, + "image_token_id": 258880, + "video_token_id": 258884 + } + } \ No newline at end of file diff --git a/examples/samsung/executor_runner/enn_executor_llm_runner.cpp b/examples/samsung/executor_runner/enn_executor_llm_runner.cpp index 43549bd32b6..eb24071d3e6 100644 --- a/examples/samsung/executor_runner/enn_executor_llm_runner.cpp +++ b/examples/samsung/executor_runner/enn_executor_llm_runner.cpp @@ -64,15 +64,17 @@ DEFINE_bool(is_decode_only, false, "Whether to run with decode method only."); int32_t main(int32_t argc, char** argv) { gflags::ParseCommandLineFlags(&argc, &argv, true); const char* model_path = FLAGS_model_path.c_str(); + std::optional data_path = std::nullopt; if (!FLAGS_data_path.empty()) { data_path = FLAGS_data_path.c_str(); } const char* tokenizer_path = FLAGS_tokenizer_path.c_str(); - std::optional pos_embedding_path = - FLAGS_pos_embedding_path.empty() - ? std::nullopt - : std::make_optional(FLAGS_pos_embedding_path.c_str()); + + std::optional pos_embedding_path = std::nullopt; + if (!FLAGS_pos_embedding_path.empty()) { + pos_embedding_path = FLAGS_pos_embedding_path; + } const char* config_path = FLAGS_config_path.c_str(); @@ -114,7 +116,7 @@ int32_t main(int32_t argc, char** argv) { example::RunnerConfig runner_config{ .config_path = config_path, .is_decode_only = is_decode_only}; std::unique_ptr runner = - example::create_llm_runner(runner_paths, runner_config); + example::create_llm_runner(runner_paths, std::move(runner_config)); ET_CHECK_MSG(runner != nullptr, "Failed to create LLM runner"); // TODO: support warmup interface diff --git a/examples/samsung/executor_runner/llm/runner/CMakeLists.txt b/examples/samsung/executor_runner/llm/runner/CMakeLists.txt index 9cfbf858199..63db5a0a798 100644 --- a/examples/samsung/executor_runner/llm/runner/CMakeLists.txt +++ b/examples/samsung/executor_runner/llm/runner/CMakeLists.txt @@ -21,6 +21,7 @@ list( PREPEND _enn_llm_runner__srcs ${CMAKE_CURRENT_LIST_DIR}/runner.cpp + ${CMAKE_CURRENT_LIST_DIR}/config.cpp ${CMAKE_CURRENT_LIST_DIR}/decoder.cpp ${CMAKE_CURRENT_LIST_DIR}/prefiller.cpp ${CMAKE_CURRENT_LIST_DIR}/io_manager.cpp diff --git a/examples/samsung/executor_runner/llm/runner/config.cpp b/examples/samsung/executor_runner/llm/runner/config.cpp new file mode 100644 index 00000000000..df99b3bf26c --- /dev/null +++ b/examples/samsung/executor_runner/llm/runner/config.cpp @@ -0,0 +1,312 @@ +/* + * Copyright (c) 2025 Samsung Electronics Co. LTD + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#include +#include +#include +#include + +#include + +namespace example { + +using json = nlohmann::json; + +void RunnerConfig::load_from_json() { + if (config_path.empty()) { + ET_LOG(Info, "Config path is empty, skipping JSON load"); + return; + } + + std::ifstream file(config_path); + if (!file.is_open()) { + ET_LOG(Error, "Could not open config file: %s", config_path.c_str()); + return; + } + + json j = json::parse(file, nullptr, false); + if (j.is_discarded()) { + ET_LOG( + Error, "Invalid JSON format in config file: %s", config_path.c_str()); + return; + } + + if (j.contains("embedding_strategy") && j["embedding_strategy"].is_string()) { + std::string strategy = j["embedding_strategy"].get(); + if (strategy == "global_only") { + embedding_strategy = EmbeddingStrategy::kGlobalOnly; + } else if (strategy == "global_and_local") { + embedding_strategy = EmbeddingStrategy::kGlobalAndLocal; + } + } + + if (j.contains("masking_strategy") && j["masking_strategy"].is_string()) { + std::string strategy = j["masking_strategy"].get(); + if (strategy == "causal") { + masking_strategy = MaskingStrategy::kCausal; + } else if (strategy == "causal_local") { + masking_strategy = MaskingStrategy::kCausalLocal; + } else if (strategy == "hybrid_conv") { + masking_strategy = MaskingStrategy::kHybridConvAttention; + } + } + + if (j.contains("model_version") && j["model_version"].is_string()) { + std::string version_str = j["model_version"].get(); + if (version_str == "llama3") { + model_version = DecoderModelVersion::kLlama3; + } else if (version_str == "gemma3") { + model_version = DecoderModelVersion::kGemma3; + } else if (version_str == "gemma4") { + model_version = DecoderModelVersion::kGemma4; + } else if (version_str == "general") { + model_version = DecoderModelVersion::kGeneral; + } + } + + if (j.contains("tokenizer_type") && j["tokenizer_type"].is_string()) { + tokenizer_type = j["tokenizer_type"].get(); + } + + // Use text_config if present, otherwise fallback to root + const json* config_source = &j; + if (j.contains("text_config") && j["text_config"].is_object()) { + config_source = &j["text_config"]; + } + const auto& src = *config_source; + + switch (model_version) { + case DecoderModelVersion::kGemma4: { + auto args = std::make_unique(); + if (src.contains("hidden_size") && + src["hidden_size"].is_number_integer()) { + args->text_config.hidden_size = src["hidden_size"].get(); + } + if (src.contains("num_hidden_layers") && + src["num_hidden_layers"].is_number_integer()) { + args->text_config.num_hidden_layers = + src["num_hidden_layers"].get(); + } + if (src.contains("num_attention_heads") && + src["num_attention_heads"].is_number_integer()) { + args->text_config.num_attention_heads = + src["num_attention_heads"].get(); + } + if (src.contains("num_key_value_heads") && + src["num_key_value_heads"].is_number_integer()) { + args->text_config.num_key_value_heads = + src["num_key_value_heads"].get(); + } + if (src.contains("head_dim") && src["head_dim"].is_number_integer()) { + args->text_config.head_dim = src["head_dim"].get(); + } + if (src.contains("global_head_dim") && + src["global_head_dim"].is_number_integer()) { + args->text_config.global_head_dim = + src["global_head_dim"].get(); + } + if (src.contains("partial_rotary_factor") && + src["partial_rotary_factor"].is_number()) { + args->text_config.rope_percentage = + src["partial_rotary_factor"].get(); + } + if (src.contains("rope_theta") && src["rope_theta"].is_number()) { + args->text_config.base_freq_global = src["rope_theta"].get(); + } + if (src.contains("rope_local_base_freq") && + src["rope_local_base_freq"].is_number()) { + args->text_config.base_freq_local = + src["rope_local_base_freq"].get(); + } + if (src.contains("vocab_size") && src["vocab_size"].is_number_integer()) { + args->text_config.vocab_size = src["vocab_size"].get(); + } + if (src.contains("max_seq_len") && + src["max_seq_len"].is_number_integer()) { + args->text_config.max_seq_len = src["max_seq_len"].get(); + } + if (src.contains("max_position_embeddings") && + src["max_position_embeddings"].is_number_integer()) { + args->text_config.max_position_embeddings = + src["max_position_embeddings"].get(); + } + if (src.contains("sliding_window") && + src["sliding_window"].is_number_integer()) { + args->text_config.sliding_window = src["sliding_window"].get(); + } + model_args = std::move(args); + break; + } + case DecoderModelVersion::kLlama3: { + auto args = std::make_unique(); + if (src.contains("hidden_dim") && src["hidden_dim"].is_number_integer()) { + args->hidden_dim = src["hidden_dim"].get(); + } else if ( + src.contains("hidden_size") && + src["hidden_size"].is_number_integer()) { + args->hidden_dim = src["hidden_size"].get(); + } + if (src.contains("n_heads") && src["n_heads"].is_number_integer()) { + args->n_heads = src["n_heads"].get(); + } else if ( + src.contains("num_attention_heads") && + src["num_attention_heads"].is_number_integer()) { + args->n_heads = src["num_attention_heads"].get(); + } + if (src.contains("n_kv_heads") && src["n_kv_heads"].is_number_integer()) { + args->n_kv_heads = src["n_kv_heads"].get(); + } else if ( + src.contains("num_key_value_heads") && + src["num_key_value_heads"].is_number_integer()) { + args->n_kv_heads = src["num_key_value_heads"].get(); + } + if (src.contains("n_layers") && src["n_layers"].is_number_integer()) { + args->n_layers = src["n_layers"].get(); + } else if ( + src.contains("num_hidden_layers") && + src["num_hidden_layers"].is_number_integer()) { + args->n_layers = src["num_hidden_layers"].get(); + } + if (src.contains("head_dim") && src["head_dim"].is_number_integer()) { + args->head_dim = src["head_dim"].get(); + } + if (src.contains("vocab_size") && src["vocab_size"].is_number_integer()) { + args->vocab_size = src["vocab_size"].get(); + } + if (src.contains("max_seq_len") && + src["max_seq_len"].is_number_integer()) { + args->max_seq_len = src["max_seq_len"].get(); + } else if ( + src.contains("max_position_embeddings") && + src["max_position_embeddings"].is_number_integer()) { + args->max_seq_len = src["max_position_embeddings"].get(); + } + if (src.contains("max_context_len") && + src["max_context_len"].is_number_integer()) { + args->max_context_len = src["max_context_len"].get(); + } else if ( + src.contains("max_position_embeddings") && + src["max_position_embeddings"].is_number_integer()) { + args->max_context_len = src["max_position_embeddings"].get(); + } + if (src.contains("prefill_seq_len") && + src["prefill_seq_len"].is_number_integer()) { + args->prefill_seq_len = src["prefill_seq_len"].get(); + } + if (src.contains("decode_seq_len") && + src["decode_seq_len"].is_number_integer()) { + args->decode_seq_len = src["decode_seq_len"].get(); + } + if (src.contains("attn_mask_rows") && + src["attn_mask_rows"].is_number_integer()) { + args->attn_mask_rows = src["attn_mask_rows"].get(); + } + if (src.contains("attn_mask_cols") && + src["attn_mask_cols"].is_number_integer()) { + args->attn_mask_cols = src["attn_mask_cols"].get(); + } + if (src.contains("sliding_window") && + src["sliding_window"].is_number_integer()) { + args->sliding_window = src["sliding_window"].get(); + } else { + args->sliding_window = std::nullopt; + } + model_args = std::move(args); + break; + } + case DecoderModelVersion::kGemma3: { + auto args = std::make_unique(); + if (src.contains("hidden_dim") && src["hidden_dim"].is_number_integer()) { + args->hidden_dim = src["hidden_dim"].get(); + } else if ( + src.contains("hidden_size") && + src["hidden_size"].is_number_integer()) { + args->hidden_dim = src["hidden_size"].get(); + } + + if (src.contains("n_heads") && src["n_heads"].is_number_integer()) { + args->n_heads = src["n_heads"].get(); + } else if ( + src.contains("num_attention_heads") && + src["num_attention_heads"].is_number_integer()) { + args->n_heads = src["num_attention_heads"].get(); + } + + if (src.contains("n_kv_heads") && src["n_kv_heads"].is_number_integer()) { + args->n_kv_heads = src["n_kv_heads"].get(); + } else if ( + src.contains("num_key_value_heads") && + src["num_key_value_heads"].is_number_integer()) { + args->n_kv_heads = src["num_key_value_heads"].get(); + } + + if (src.contains("n_layers") && src["n_layers"].is_number_integer()) { + args->n_layers = src["n_layers"].get(); + } else if ( + src.contains("num_hidden_layers") && + src["num_hidden_layers"].is_number_integer()) { + args->n_layers = src["num_hidden_layers"].get(); + } + + if (src.contains("head_dim") && src["head_dim"].is_number_integer()) { + args->head_dim = src["head_dim"].get(); + } + + if (src.contains("vocab_size") && src["vocab_size"].is_number_integer()) { + args->vocab_size = src["vocab_size"].get(); + } + + if (src.contains("max_seq_len") && + src["max_seq_len"].is_number_integer()) { + args->max_seq_len = src["max_seq_len"].get(); + } else if ( + src.contains("max_position_embeddings") && + src["max_position_embeddings"].is_number_integer()) { + args->max_seq_len = src["max_position_embeddings"].get(); + } + + if (src.contains("max_context_len") && + src["max_context_len"].is_number_integer()) { + args->max_context_len = src["max_context_len"].get(); + } else if ( + src.contains("max_position_embeddings") && + src["max_position_embeddings"].is_number_integer()) { + args->max_context_len = src["max_position_embeddings"].get(); + } + + if (src.contains("prefill_seq_len") && + src["prefill_seq_len"].is_number_integer()) { + args->prefill_seq_len = src["prefill_seq_len"].get(); + } + if (src.contains("decode_seq_len") && + src["decode_seq_len"].is_number_integer()) { + args->decode_seq_len = src["decode_seq_len"].get(); + } + if (src.contains("attn_mask_rows") && + src["attn_mask_rows"].is_number_integer()) { + args->attn_mask_rows = src["attn_mask_rows"].get(); + } + if (src.contains("attn_mask_cols") && + src["attn_mask_cols"].is_number_integer()) { + args->attn_mask_cols = src["attn_mask_cols"].get(); + } + if (src.contains("sliding_window") && + src["sliding_window"].is_number_integer()) { + args->sliding_window = src["sliding_window"].get(); + } else { + args->sliding_window = std::nullopt; + } + model_args = std::move(args); + break; + } + case DecoderModelVersion::kGeneral: + default: + break; + } +} +} // namespace example \ No newline at end of file diff --git a/examples/samsung/executor_runner/llm/runner/config.h b/examples/samsung/executor_runner/llm/runner/config.h new file mode 100644 index 00000000000..0246a6bc0ec --- /dev/null +++ b/examples/samsung/executor_runner/llm/runner/config.h @@ -0,0 +1,456 @@ +/* + * Copyright (c) 2026 Samsung Electronics Co. LTD + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +#include + +#include +#include +#include +#include +#include + +namespace example { + +// Exynos Runtime metadata key constants +inline constexpr auto kPrefillArLen = "get_prefill_ar_len"; +inline constexpr auto kNumOfLayers = "get_n_layers"; +inline constexpr auto kNumOfHiddenDim = "get_hidden_dim"; +inline constexpr auto kNumOfHeads = "get_n_heads"; +inline constexpr auto kNumOfKVHeads = "get_n_kv_heads"; + +inline constexpr auto default_method_name = "forward"; +inline constexpr auto default_method_name_prefill = "prefill_forward"; +inline constexpr auto default_method_name_decode = "decode_forward"; + +enum class DecoderModelVersion { + kLlama3, + kGemma3, + kGemma4, + kGeneral, +}; + +struct ModelVersionConfig { + size_t position_rank; + size_t attention_mask_rank; + size_t kv_cache_rank; + + size_t attention_mask_tensor_index; + size_t attention_mask_tensor_count; + + size_t position_tensor_index; + size_t position_tensor_count; + + size_t k_cache_tensor_index; + size_t v_cache_tensor_index; + size_t kv_cache_tensor_count; + + size_t heads_dim_index; + size_t head_size_dim_index; + + bool has_local_global_rope; + bool has_partial_rope; + bool has_shared_kv; + + size_t sliding_window; +}; + +inline ModelVersionConfig getModelVersionConfig(DecoderModelVersion version) { + switch (version) { + case DecoderModelVersion::kGemma3: + return { + .position_rank = 3, + .attention_mask_rank = 3, + .kv_cache_rank = 4, + + // Flat PTE: + // 0 input_ids + // 1-2 attention_setup + // 3-4 masks + // 5-6 position + // 7-32 k_cache_0..25 + // 33-58 v_cache_0..25 + .attention_mask_tensor_index = 3, + .attention_mask_tensor_count = 2, + + .position_tensor_index = 5, + .position_tensor_count = 2, + + .k_cache_tensor_index = 7, + .v_cache_tensor_index = 33, + .kv_cache_tensor_count = 26, + + .heads_dim_index = 1, + .head_size_dim_index = 3, + + .has_local_global_rope = true, + .has_partial_rope = false, + .has_shared_kv = false, + .sliding_window = 512, + }; + + case DecoderModelVersion::kGemma4: + return { + .position_rank = 3, + .attention_mask_rank = 4, + .kv_cache_rank = 4, + + // Flat PTE: + // 0 input_ids + // 1 inputs_embeds + // 2 mm_mask + // 3 attn_mask_full + // 4 attn_mask_sliding + // 5 pos_emb_global_cos + // 6 pos_emb_global_sin + // 7 pos_emb_local_cos + // 8 pos_emb_local_sin + // 9-43 k_cache_0..34 + // 44-78 v_cache_0..34 + .attention_mask_tensor_index = 3, + .attention_mask_tensor_count = 2, + + .position_tensor_index = 5, + .position_tensor_count = 4, + + .k_cache_tensor_index = 9, + .v_cache_tensor_index = 44, + .kv_cache_tensor_count = 35, + + .heads_dim_index = 1, + .head_size_dim_index = 3, + + .has_local_global_rope = true, + .has_partial_rope = true, + .has_shared_kv = true, + .sliding_window = 512, + }; + + case DecoderModelVersion::kLlama3: + return { + .position_rank = 3, + .attention_mask_rank = 4, + .kv_cache_rank = 4, + + .attention_mask_tensor_index = 1, + .attention_mask_tensor_count = 1, + + .position_tensor_index = 2, + .position_tensor_count = 2, + + .k_cache_tensor_index = 4, + .v_cache_tensor_index = 20, + .kv_cache_tensor_count = 16, + + .heads_dim_index = 1, + .head_size_dim_index = 3, + + .has_local_global_rope = false, + .has_partial_rope = false, + .has_shared_kv = false, + .sliding_window = 0, + }; + default: + return { + .position_rank = 3, + .kv_cache_rank = 4, + .has_local_global_rope = false, + .has_partial_rope = false, + .has_shared_kv = false, + .sliding_window = 0, + }; + } +} + +enum class EmbeddingStrategy { + kGlobalOnly, + kGlobalAndLocal, +}; + +enum class MaskingStrategy { + kCausal, + kCausalLocal, + kHybridConvAttention, // LIV-conv + GQA mask constraint +}; + +class RunnerPaths { + public: + std::string model_path; + std::string tokenizer_path; + std::optional pos_embedding_path = std::nullopt; + std::optional data_path = std::nullopt; +}; + +class BaseModelArgs { + public: + virtual ~BaseModelArgs() = default; + virtual DecoderModelVersion get_version() const = 0; + virtual int64_t get_hidden_dim() const { + return 0; + } + virtual int64_t get_n_heads() const { + return 0; + } + virtual int64_t get_n_kv_heads() const { + return 0; + } + virtual int64_t get_n_layers() const { + return 0; + } + virtual int64_t get_head_dim() const { + return 0; + } + virtual int64_t get_global_head_dim() const { + return 0; + } + virtual int64_t get_vocab_size() const { + return 0; + } + virtual int64_t get_max_seq_len() const { + return 128; + } + virtual int64_t get_max_context_len() const { + return 128; + } + virtual int64_t get_prefill_seq_len() const { + return 128; + } + virtual int64_t get_decode_seq_len() const { + return 1; + } + virtual int64_t get_attn_mask_rows() const { + return 128; + } + virtual int64_t get_attn_mask_cols() const { + return 1024; + } + virtual std::optional get_sliding_window() const { + return std::nullopt; + } +}; + +class Llama3Args : public BaseModelArgs { + public: + DecoderModelVersion get_version() const override { + return DecoderModelVersion::kLlama3; + } + int64_t hidden_dim = 0; + int64_t n_heads = 0; + int64_t n_kv_heads = 0; + int64_t n_layers = 0; + int64_t head_dim = 0; + int64_t vocab_size = 0; + int64_t max_seq_len = 128; + int64_t max_context_len = 128; + int64_t prefill_seq_len = 128; + int64_t decode_seq_len = 1; + int64_t attn_mask_rows = 128; + int64_t attn_mask_cols = 1024; + std::optional sliding_window; + + int64_t get_hidden_dim() const override { + return hidden_dim; + } + int64_t get_n_heads() const override { + return n_heads; + } + int64_t get_n_kv_heads() const override { + return n_kv_heads; + } + int64_t get_n_layers() const override { + return n_layers; + } + int64_t get_head_dim() const override { + return head_dim; + } + int64_t get_global_head_dim() const override { + return head_dim; + } + int64_t get_vocab_size() const override { + return vocab_size; + } + int64_t get_max_seq_len() const override { + return max_seq_len; + } + int64_t get_max_context_len() const override { + return max_context_len; + } + int64_t get_prefill_seq_len() const override { + return prefill_seq_len; + } + int64_t get_decode_seq_len() const override { + return decode_seq_len; + } + int64_t get_attn_mask_rows() const override { + return attn_mask_rows; + } + int64_t get_attn_mask_cols() const override { + return attn_mask_cols; + } + std::optional get_sliding_window() const override { + return sliding_window; + } +}; + +class Gemma3Args : public BaseModelArgs { + public: + DecoderModelVersion get_version() const override { + return DecoderModelVersion::kGemma3; + } + int64_t hidden_dim = 0; + int64_t n_heads = 0; + int64_t n_kv_heads = 0; + int64_t n_layers = 0; + int64_t head_dim = 0; + int64_t vocab_size = 0; + int64_t max_seq_len = 128; + int64_t max_context_len = 128; + int64_t prefill_seq_len = 128; + int64_t decode_seq_len = 1; + int64_t attn_mask_rows = 128; + int64_t attn_mask_cols = 1024; + std::optional sliding_window; + + int64_t get_hidden_dim() const override { + return hidden_dim; + } + int64_t get_n_heads() const override { + return n_heads; + } + int64_t get_n_kv_heads() const override { + return n_kv_heads; + } + int64_t get_n_layers() const override { + return n_layers; + } + int64_t get_head_dim() const override { + return head_dim; + } + int64_t get_global_head_dim() const override { + return head_dim; + } + int64_t get_vocab_size() const override { + return vocab_size; + } + int64_t get_max_seq_len() const override { + return max_seq_len; + } + int64_t get_max_context_len() const override { + return max_context_len; + } + int64_t get_prefill_seq_len() const override { + return prefill_seq_len; + } + int64_t get_decode_seq_len() const override { + return decode_seq_len; + } + int64_t get_attn_mask_rows() const override { + return attn_mask_rows; + } + int64_t get_attn_mask_cols() const override { + return attn_mask_cols; + } + std::optional get_sliding_window() const override { + return sliding_window; + } +}; + +class Gemma4Args : public BaseModelArgs { + public: + DecoderModelVersion get_version() const override { + return DecoderModelVersion::kGemma4; + } + + struct TextConfig { + int64_t hidden_size = 0; + int64_t intermediate_size = 0; + bool use_double_wide_mlp = false; + int64_t num_hidden_layers = 0; + int64_t num_attention_heads = 0; + int64_t head_dim = 0; + int64_t global_head_dim = 0; + int64_t num_key_value_heads = 0; + int64_t vocab_size = 0; + int64_t max_seq_len = 0; + int64_t max_position_embeddings = 0; + std::optional sliding_window; + + // RoPE specific configuration for Gemma 4 + float rope_percentage = 0.25f; + float base_freq_local = 10000.0f; + float base_freq_global = 1000000.0f; + } text_config; + + struct AudioConfig { + int64_t hidden_size = 0; + int64_t num_hidden_layers = 0; + int64_t num_attention_heads = 0; + } audio_config; + + int64_t audio_seq_length = 0; + int64_t attn_mask_rows = 128; + int64_t attn_mask_cols = 1024; + + int64_t get_hidden_dim() const override { + return text_config.hidden_size; + } + int64_t get_n_heads() const override { + return text_config.num_attention_heads; + } + int64_t get_n_kv_heads() const override { + return text_config.num_key_value_heads; + } + int64_t get_n_layers() const override { + return text_config.num_hidden_layers; + } + int64_t get_head_dim() const override { + return text_config.head_dim; + } + int64_t get_global_head_dim() const override { + return text_config.global_head_dim; + } + int64_t get_vocab_size() const override { + return text_config.vocab_size; + } + int64_t get_max_seq_len() const override { + return text_config.max_seq_len; + } + int64_t get_max_context_len() const override { + return text_config.max_position_embeddings; + } + std::optional get_sliding_window() const override { + return text_config.sliding_window; + } + int64_t get_attn_mask_rows() const override { + return attn_mask_rows; + } + int64_t get_attn_mask_cols() const override { + return attn_mask_cols; + } +}; + +class RunnerConfig { + public: + std::string config_path; + bool is_decode_only = false; + float temperature = -1.0f; + + // Strategies + EmbeddingStrategy embedding_strategy = EmbeddingStrategy::kGlobalOnly; + MaskingStrategy masking_strategy = MaskingStrategy::kCausal; + + std::string tokenizer_type; + + DecoderModelVersion model_version; + std::unique_ptr model_args; + + void load_from_json(); +}; + +} // namespace example \ No newline at end of file diff --git a/examples/samsung/executor_runner/llm/runner/constants.h b/examples/samsung/executor_runner/llm/runner/constants.h deleted file mode 100644 index 68a180943f5..00000000000 --- a/examples/samsung/executor_runner/llm/runner/constants.h +++ /dev/null @@ -1,43 +0,0 @@ -/* - * Copyright (c) Samsung Electronics Co. LTD - * All rights reserved. - * - * This source code is licensed under the BSD-style license found in the - * LICENSE file in the root directory of this source tree. - */ - -#pragma once -#include - -// constants for LLM runtime -namespace example { - -// Exynos Runtime metadata key constants -inline constexpr auto kPrefillArLen = "get_prefill_ar_len"; -inline constexpr auto kNumOfLayers = "get_n_layers"; -inline constexpr auto kNumOfHiddenDim = "get_hidden_dim"; -inline constexpr auto kNumOfHeads = "get_n_heads"; -inline constexpr auto kNumOfKVHeads = "get_n_kv_heads"; - -inline constexpr auto default_method_name = "forward"; -inline constexpr auto default_method_name_prefill = "prefill_forward"; -inline constexpr auto default_method_name_decode = "decode_forward"; - -enum class DecoderModelVersion { - kLlama3, - kGemma3, - kGeneral, -}; - -enum class EmbeddingStrategy { - kGlobalOnly, - kGlobalAndLocal, -}; - -enum class MaskingStrategy { - kCausal, - kCausalLocal, - kHybridConvAttention, // LIV-conv + GQA mask constraint -}; - -} // namespace example diff --git a/examples/samsung/executor_runner/llm/runner/decoder.cpp b/examples/samsung/executor_runner/llm/runner/decoder.cpp index 6f5a32a9dcc..5d6c17edfc2 100644 --- a/examples/samsung/executor_runner/llm/runner/decoder.cpp +++ b/examples/samsung/executor_runner/llm/runner/decoder.cpp @@ -8,14 +8,13 @@ // Given inputs, run a text decoder and return logits. -#include -#include - -#include - #include #include #include +#include + +#include +#include namespace example { @@ -138,85 +137,190 @@ runtime::Result Decoder::step( } auto method_meta = ET_UNWRAP(module_->method_meta(method_stage.c_str())); - const size_t total_input_count = method_meta.num_inputs(); - bool use_kv_cache = - total_input_count > 1; // kv cache is not used, when input is 1 + + bool use_kv_cache = total_input_count > 1; if (use_kv_cache) { - // Size of the position tracking argument. This could be either input_pos or - // cache_positions - constexpr int64_t kExpectedPositionDimensions = 3; - constexpr int64_t kPositionTensorIndex = 3; - constexpr int64_t kAttentionSetupTensorIndex = 1; - - constexpr int64_t kKVCacheExpectedDimensions = 4; - constexpr int64_t kHeadsDimIndex = 1; - constexpr int64_t kHeadSizeDimIndex = 3; - - // Check if we are using cache positions instead of input pos. - auto position_tensor_meta = - ET_UNWRAP(method_meta.input_tensor_meta(kPositionTensorIndex)); - auto position_shape = position_tensor_meta.sizes(); + auto* llm_io_manager = static_cast(io_manager_); + const auto& config = llm_io_manager->get_config(); + auto model_version_config = getModelVersionConfig(config.model_version); + + // + // Validate attention-mask tensors. + // + for (size_t i = 0; i < model_version_config.attention_mask_tensor_count; + ++i) { + const size_t tensor_index = + model_version_config.attention_mask_tensor_index + i; + + auto tensor_meta = ET_UNWRAP(method_meta.input_tensor_meta(tensor_index)); + + const auto tensor_shape = tensor_meta.sizes(); + + ET_CHECK_OR_RETURN_ERROR( + tensor_shape.size() == model_version_config.attention_mask_rank, + InvalidProgram, + "Unexpected attention mask rank at input %zu. " + "Expected %zu, got %zu.", + tensor_index, + model_version_config.attention_mask_rank, + tensor_shape.size()); + } + + // + // Validate position / RoPE tensors. + // + // Gemma3: + // position[0], position[1] + // + // Gemma4: + // global_cos, global_sin, local_cos, local_sin + // + for (size_t i = 0; i < model_version_config.position_tensor_count; ++i) { + const size_t tensor_index = + model_version_config.position_tensor_index + i; + + auto tensor_meta = ET_UNWRAP(method_meta.input_tensor_meta(tensor_index)); + + const auto tensor_shape = tensor_meta.sizes(); + + ET_CHECK_OR_RETURN_ERROR( + tensor_shape.size() == model_version_config.position_rank, + InvalidProgram, + "Unexpected position tensor rank at input %zu. " + "Expected %zu, got %zu.", + tensor_index, + model_version_config.position_rank, + tensor_shape.size()); + } + + // + // Validate KV-cache tensors. + // + // We only need one cache tensor to determine/validate the common + // KV-cache layout [B, H, S, D], assuming all cache leaves use + // the same rank/layout. + // + auto k_cache_tensor_meta = ET_UNWRAP(method_meta.input_tensor_meta( + model_version_config.k_cache_tensor_index)); + + const auto k_cache_shape = k_cache_tensor_meta.sizes(); + ET_CHECK_OR_RETURN_ERROR( - position_shape.size() == kExpectedPositionDimensions, + k_cache_shape.size() == model_version_config.kv_cache_rank, InvalidProgram, - "The pos input tensor is not expected. Got dimension (%zu)", - position_shape.size()); - std::vector sizes_vec = {position_shape[0]}; + "Unexpected K-cache rank. Expected %zu, got %zu.", + model_version_config.kv_cache_rank, + k_cache_shape.size()); - auto attention_setup_tensor_meta = - ET_UNWRAP(method_meta.input_tensor_meta(kAttentionSetupTensorIndex)); - auto attention_setup_shape = attention_setup_tensor_meta.sizes(); - auto* llm_io_manager = static_cast(io_manager_); - const auto& config = llm_io_manager->get_config(); - const size_t kv_cache_tensor_index = kAttentionSetupTensorIndex + - llm_io_manager->get_attention_setup_tensor_count(); + auto v_cache_tensor_meta = ET_UNWRAP(method_meta.input_tensor_meta( + model_version_config.v_cache_tensor_index)); + + const auto v_cache_shape = v_cache_tensor_meta.sizes(); - auto kv_cache_tensor_meta = - ET_UNWRAP(method_meta.input_tensor_meta(kv_cache_tensor_index)); - auto kv_cache_shape = kv_cache_tensor_meta.sizes(); ET_CHECK_OR_RETURN_ERROR( - static_cast(kv_cache_shape.size()) == - kKVCacheExpectedDimensions, - InvalidArgument, - "Input tensor at index %" PRId64 - " (KV cache slot) has invalid dimensionality: %" PRId64 - " (expected %" PRId64 ")", - kv_cache_tensor_index, - static_cast(kv_cache_shape.size()), - kKVCacheExpectedDimensions); + v_cache_shape.size() == model_version_config.kv_cache_rank, + InvalidProgram, + "Unexpected V-cache rank. Expected %zu, got %zu.", + model_version_config.kv_cache_rank, + v_cache_shape.size()); + // + // Validate the configured KV-cache axes before indexing them. + // ET_CHECK_OR_RETURN_ERROR( - static_cast(kv_cache_shape[kHeadsDimIndex]) == - config.n_kv_heads, - InvalidArgument, - "KV cache tensor head count mismatch at index %" PRId64 - ": found %" PRId64 " (expected %" PRId64 " via configuration)", - kv_cache_tensor_index, - static_cast(kv_cache_shape[kHeadsDimIndex]), - config.n_kv_heads); + model_version_config.heads_dim_index < k_cache_shape.size(), + InvalidProgram, + "KV heads dimension index %zu exceeds cache rank %zu.", + model_version_config.heads_dim_index, + k_cache_shape.size()); ET_CHECK_OR_RETURN_ERROR( - static_cast(kv_cache_shape[kHeadSizeDimIndex]) == - config.head_dim, - InvalidArgument, - "KV cache tensor head dimension size mismatch at index %" PRId64 - ": found %" PRId64 " (expected %" PRId64 " via configuration)", - kv_cache_tensor_index, - static_cast(kv_cache_shape[kHeadSizeDimIndex]), - config.head_dim); + model_version_config.head_size_dim_index < k_cache_shape.size(), + InvalidProgram, + "KV head-size dimension index %zu exceeds cache rank %zu.", + model_version_config.head_size_dim_index, + k_cache_shape.size()); - extension::TensorPtr start_pos_tensor; - // Assuming model is exported with input_pos, create a tensor with size 1 - start_pos_tensor = - extension::from_blob(&start_pos, sizes_vec, aten::ScalarType::Long); + const auto num_kv_heads = + k_cache_shape[model_version_config.heads_dim_index]; - std::vector inputs; + const auto kv_head_size = + k_cache_shape[model_version_config.head_size_dim_index]; + + ET_LOG( + Debug, + "KV cache: K base=%zu, V base=%zu, count=%zu, " + "heads=%d, head_size=%d", + model_version_config.k_cache_tensor_index, + model_version_config.v_cache_tensor_index, + model_version_config.kv_cache_tensor_count, + num_kv_heads, + kv_head_size); + + // + // Optional sanity check that the final V-cache input actually + // exists in this PTE method. + // + const size_t last_v_cache_index = + model_version_config.v_cache_tensor_index + + model_version_config.kv_cache_tensor_count - 1; + ET_CHECK_OR_RETURN_ERROR( + last_v_cache_index < method_meta.num_inputs(), + InvalidProgram, + "KV cache layout exceeds PTE input count. " + "Last V-cache input=%zu, num_inputs=%zu.", + last_v_cache_index, + method_meta.num_inputs()); + + // Assuming model is exported with input_pos, create a tensor with size 1 + extension::TensorPtr start_pos_tensor; + std::vector start_pos_sizes = {1}; + start_pos_tensor = extension::from_blob( + &start_pos, start_pos_sizes, aten::ScalarType::Long); + std::vector inputs; if (method_stage == prefill_method_) { auto inputs_res = io_manager_->prepare_prefill(tokens, start_pos_tensor); ET_CHECK_OK_OR_RETURN_ERROR(inputs_res.error()); inputs = inputs_res.get(); + + ET_LOG( + Info, + "Prepared %zu inputs, PTE expects %zu inputs", + inputs.size(), + method_meta.num_inputs()); + + const size_t count = std::min(inputs.size(), method_meta.num_inputs()); + + for (size_t i = 0; i < count; ++i) { + if (!inputs[i].isTensor()) { + ET_LOG(Info, "input[%zu]: prepared EValue is not Tensor", i); + continue; + } + + const auto prepared_tensor = inputs[i].toTensor(); + const auto prepared_sizes = prepared_tensor.sizes(); + + auto expected_meta = ET_UNWRAP(method_meta.input_tensor_meta(i)); + const auto expected_sizes = expected_meta.sizes(); + + ET_LOG( + Info, + "input[%zu]: prepared_rank=%zu expected_rank=%zu", + i, + prepared_sizes.size(), + expected_sizes.size()); + + for (size_t d = 0; d < prepared_sizes.size(); ++d) { + ET_LOG(Info, " prepared[%zu]=%d", d, prepared_sizes[d]); + } + + for (size_t d = 0; d < expected_sizes.size(); ++d) { + ET_LOG(Info, " expected[%zu]=%d", d, expected_sizes[d]); + } + } + } else if (method_stage == decode_method_) { auto inputs_res = io_manager_->prepare_decode(tokens, start_pos_tensor); ET_CHECK_OK_OR_RETURN_ERROR(inputs_res.error()); @@ -250,9 +354,8 @@ runtime::Result Decoder::step( } // Return the logits tensor return slice_kth_token(outputs_res.get()[0], (tokens->numel() - 1)); - } else { // no kv cache - (void)start_pos; // unused - + } else { // kv cache is not used, when input is 1 + (void)start_pos; auto outputs_res = module_->execute(method_stage.c_str(), tokens); ET_LOG(Info, "After module forward"); ET_CHECK_OK_OR_RETURN_ERROR(outputs_res.error()); diff --git a/examples/samsung/executor_runner/llm/runner/io_manager.cpp b/examples/samsung/executor_runner/llm/runner/io_manager.cpp index 0a165b6f8af..529b1567859 100644 --- a/examples/samsung/executor_runner/llm/runner/io_manager.cpp +++ b/examples/samsung/executor_runner/llm/runner/io_manager.cpp @@ -8,7 +8,7 @@ #include -#include +#include #include namespace example { @@ -39,41 +39,122 @@ fill_float_region(float* base, size_t offset, size_t len, float value) { } // namespace -LlmIOManager::LlmIOManager(extension::Module& module, Config config) +LlmIOManager::LlmIOManager( + extension::Module& module, + const RunnerConfig& config) : IOManager(module), config_(config) { - ET_LOG( - Info, - "LlmIOManager Constructor: dim(%zu), vocab(%zu), max_seq(%zu)", - config_.hidden_dim, - config_.vocab_size, - config_.max_seq_len); - - rotary_embedding_ = std::make_unique(RotaryEmbedding::Config{ - config_.head_dim, config_.max_seq_len, config_.embedding_strategy}); + rotary_embedding_ = std::make_unique(config); ET_CHECK_MSG( - config_.max_seq_len >= config_.prefill_seq_len, + config_.model_args->get_max_seq_len() >= + config_.model_args->get_prefill_seq_len(), "Inconsistent configuration: max_seq_len (%zu) must be greater than or equal to prefill_seq_len (%zu)", - config_.max_seq_len, - config_.prefill_seq_len); + config_.model_args->get_max_seq_len(), + config_.model_args->get_prefill_seq_len()); ET_CHECK_MSG( - config_.max_seq_len >= 1, + config_.model_args->get_max_seq_len() >= 1, "Inconsistent configuration: max_seq_len must be at least 1, but got %zu", - config_.max_seq_len); + config_.model_args->get_max_seq_len()); - prefill_cache_len_ = config_.max_seq_len - config_.prefill_seq_len; - decode_cache_len_ = config_.max_seq_len - 1; + prefill_cache_len_ = config_.model_args->get_max_seq_len() - + config_.model_args->get_prefill_seq_len(); + decode_cache_len_ = config_.model_args->get_max_seq_len() - 1; allocate_buffers(); } +size_t LlmIOManager::get_attention_setup_tensor_count() const { + constexpr size_t kGlobalMaskCount = 1; + constexpr size_t kGlobalRopeCosCount = 1; + constexpr size_t kGlobalRopeSinCount = 1; + constexpr size_t kBaselineMetadataCount = + kGlobalMaskCount + kGlobalRopeCosCount + kGlobalRopeSinCount; + + size_t setup_tensor_count = kBaselineMetadataCount; + + const bool requires_local_mask = + (config_.masking_strategy == MaskingStrategy::kCausalLocal || + config_.masking_strategy == MaskingStrategy::kHybridConvAttention); + if (requires_local_mask) { + constexpr size_t kLocalMaskCount = 1; + setup_tensor_count += kLocalMaskCount; + } + + const bool requires_local_rope = + (config_.embedding_strategy == EmbeddingStrategy::kGlobalAndLocal); + if (requires_local_rope) { + constexpr size_t kLocalRopeCosCount = 1; + constexpr size_t kLocalRopeSinCount = 1; + setup_tensor_count += (kLocalRopeCosCount + kLocalRopeSinCount); + } + + return setup_tensor_count; +} + +ET_NODISCARD runtime::Error LlmIOManager::load( + const std::string& prefill_method, + const std::string& decode_method) { + (void)prefill_method; + (void)decode_method; + + if ((config_.masking_strategy == MaskingStrategy::kCausalLocal || + config_.masking_strategy == MaskingStrategy::kHybridConvAttention) && + !config_.model_args->get_sliding_window().has_value()) { + ET_LOG( + Error, + "Sliding window configuration is required for the selected masking strategy"); + return runtime::Error::InvalidArgument; + } + + ET_LOG( + Info, + "(%s) Model Config: layers(%zu), kv_heads(%zu), hidden_dim(%zu), head_dim(%zu), \ + vocab(%zu), max_seq(%zu), max_ctx(%zu), prefill_seq(%zu), decode_seq(%zu), \ + mask_rows(%zu), mask_cols(%zu), slide_win(%zu), emb_strat(%d), mask_strat(%d)", + ET_FUNCTION, + config_.model_args->get_n_layers(), + config_.model_args->get_n_kv_heads(), + config_.model_args->get_hidden_dim(), + config_.model_args->get_head_dim(), + config_.model_args->get_vocab_size(), + config_.model_args->get_max_seq_len(), + config_.model_args->get_max_context_len(), + config_.model_args->get_prefill_seq_len(), + config_.model_args->get_decode_seq_len(), + config_.model_args->get_attn_mask_rows(), + config_.model_args->get_attn_mask_cols(), + config_.model_args->get_sliding_window().value_or(0), + static_cast(config_.embedding_strategy), + static_cast(config_.masking_strategy)); + + return runtime::Error::Ok; +} + +ET_NODISCARD runtime::Error LlmIOManager::reset( + const std::string& prefill_method, + const std::string& decode_method) { + (void)prefill_method; + (void)decode_method; + + dump_step_ = 0; + current_pos_ = 0; + prefill_done_ = false; + + for (size_t l = 0; l < config_.model_args->get_n_layers(); l++) { + std::fill(prefill_k_input_[l].begin(), prefill_k_input_[l].end(), 0); + std::fill(prefill_v_input_[l].begin(), prefill_v_input_[l].end(), 0); + std::fill(decode_k_input_[l].begin(), decode_k_input_[l].end(), 0); + std::fill(decode_v_input_[l].begin(), decode_v_input_[l].end(), 0); + } + std::fill(attn_mask_.begin(), attn_mask_.end(), 0.0f); + + return runtime::Error::Ok; +} + executorch::runtime::Error LlmIOManager::load_rotary_embedding( std::optional pos_embedding_path) { - if (pos_embedding_path.has_value()) { - return rotary_embedding_->load(pos_embedding_path.value()); - } - return executorch::runtime::Error::Ok; + return rotary_embedding_->load(pos_embedding_path); } executorch::runtime::Result> @@ -84,58 +165,78 @@ LlmIOManager::prepare_prefill( (void)prefill_method; std::vector inputs; inputs.reserve(get_input_reserve_size()); + const size_t prefill_seq_len = config_.model_args->get_prefill_seq_len(); + const size_t prompt_token_len = input->numel(); + + ET_CHECK_OR_RETURN_ERROR( + prompt_token_len <= prefill_seq_len, + InvalidArgument, + "Prompt chunk length %zu exceeds prefill sequence length %zu.", + prompt_token_len, + prefill_seq_len); + + const int64_t* prompt_tokens = input->const_data_ptr(); + input_tokens_.assign(prefill_seq_len, 0); - // re-producessing input token tensor - const int64_t* prompt_tokens = input->mutable_data_ptr(); - size_t prompt_token_len = input->numel(); - input_tokens_.resize(prompt_token_len); std::copy( prompt_tokens, prompt_tokens + prompt_token_len, input_tokens_.begin()); - input_tokens_.resize(config_.prefill_seq_len, 0); input_tokens_tensor_ = executorch::extension::from_blob( input_tokens_.data(), - {1, static_cast(config_.prefill_seq_len)}, + {1, static_cast(prefill_seq_len)}, executorch::aten::ScalarType::Long); inputs.emplace_back(input_tokens_tensor_); - // re-producessing start_pos tensor - const int64_t* start_position = start_pos->mutable_data_ptr(); - pos_ids_.resize(config_.prefill_seq_len); + const int64_t* start_position = start_pos->const_data_ptr(); current_pos_ = *start_position; - for (int64_t i = 0; i < config_.prefill_seq_len; ++i) { - pos_ids_[i] = current_pos_ + i; + pos_ids_.resize(prefill_seq_len); + for (size_t i = 0; i < prefill_seq_len; ++i) { + pos_ids_[i] = current_pos_ + static_cast(i); } + // TODO (brucekk.kim): + // 1. Add two missing inputs + // inputs[1]: inputs_embeds + // inputs[2]: mm_mask + + // TODO (brucekk.kim) + // 2. Fix/refactor update_attn_mask so that dim changes from [1,1,128,1024] to + // [1,1,128,4096] + // attention tensor update_attn_mask(current_pos_, prompt_token_len); attn_mask_tensor_ = executorch::extension::from_blob( attn_mask_.data(), {1, 1, - static_cast(config_.attn_mask_rows), - static_cast(config_.attn_mask_cols)}); + static_cast( + config_.model_args->get_attn_mask_rows()), + static_cast( + config_.model_args->get_attn_mask_cols())}); inputs.emplace_back(attn_mask_tensor_); - switch (config_.masking_strategy) { + const auto masking_strategy = config_.masking_strategy; + switch (masking_strategy) { case MaskingStrategy::kCausal: break; case MaskingStrategy::kCausalLocal: case MaskingStrategy::kHybridConvAttention: update_local_attn_mask(current_pos_, prompt_token_len); + const size_t mask_rows = config_.model_args->get_attn_mask_rows(); + const size_t mask_cols = config_.model_args->get_attn_mask_cols(); local_attn_mask_tensor_ = executorch::extension::from_blob( local_attn_mask_.data(), {1, 1, - static_cast(config_.attn_mask_rows), - static_cast(config_.attn_mask_cols)}); + static_cast(mask_rows), + static_cast(mask_cols)}); inputs.emplace_back(local_attn_mask_tensor_); break; } ET_CHECK_OK_OR_RETURN_ERROR( - prepare_rotary_embedding_inputs(config_.prefill_seq_len, inputs)); + prepare_rotary_embedding_inputs(prefill_seq_len, inputs)); prefill_k_input_tensors_.clear(); prefill_v_input_tensors_.clear(); @@ -165,7 +266,7 @@ LlmIOManager::prepare_decode( } const int64_t* start_position = start_pos->mutable_data_ptr(); current_pos_ = *start_position; - update_attn_mask(current_pos_, config_.decode_seq_len); + update_attn_mask(current_pos_, config_.model_args->get_decode_seq_len()); std::vector inputs; inputs.reserve(get_input_reserve_size()); inputs.emplace_back(input); @@ -175,7 +276,8 @@ LlmIOManager::prepare_decode( {1, 1, 1, - static_cast(config_.attn_mask_cols)}); + static_cast( + config_.model_args->get_attn_mask_cols())}); inputs.emplace_back(attn_mask_tensor_); switch (config_.masking_strategy) { case MaskingStrategy::kCausal: @@ -188,16 +290,17 @@ LlmIOManager::prepare_decode( {1, 1, 1, - static_cast(config_.attn_mask_cols)}); + static_cast( + config_.model_args->get_attn_mask_cols())}); inputs.emplace_back(local_attn_mask_tensor_); break; } pos_ids_.clear(); - pos_ids_.resize(config_.decode_seq_len); + pos_ids_.resize(config_.model_args->get_decode_seq_len()); pos_ids_[0] = current_pos_; - ET_CHECK_OK_OR_RETURN_ERROR( - prepare_rotary_embedding_inputs(config_.decode_seq_len, inputs)); + ET_CHECK_OK_OR_RETURN_ERROR(prepare_rotary_embedding_inputs( + config_.model_args->get_decode_seq_len(), inputs)); decode_k_input_tensors_.clear(); decode_v_input_tensors_.clear(); @@ -224,7 +327,7 @@ executorch::runtime::Error LlmIOManager::update_prefill( prefill_k_input_, prefill_v_input_, prefill_cache_len_, - config_.prefill_seq_len, + config_.model_args->get_prefill_seq_len(), "prefill"); return executorch::runtime::Error::Ok; } @@ -243,39 +346,49 @@ executorch::runtime::Error LlmIOManager::update_decode( } // TOOD: sizeof(float) dtype hard-codes the buffer dtype to be float. void LlmIOManager::allocate_buffers() { - size_t prefill_in_size = config_.n_kv_heads * prefill_cache_len_ * - config_.head_dim * sizeof(float); - size_t decode_in_size = - config_.n_kv_heads * decode_cache_len_ * config_.head_dim * sizeof(float); - prefill_k_input_.resize(config_.n_layers); - prefill_v_input_.resize(config_.n_layers); - decode_k_input_.resize(config_.n_layers); - decode_v_input_.resize(config_.n_layers); - for (size_t l = 0; l < config_.n_layers; l++) { + size_t prefill_in_size = config_.model_args->get_n_kv_heads() * + prefill_cache_len_ * config_.model_args->get_head_dim() * sizeof(float); + size_t decode_in_size = config_.model_args->get_n_kv_heads() * + decode_cache_len_ * config_.model_args->get_head_dim() * sizeof(float); + prefill_k_input_.resize(config_.model_args->get_n_layers()); + prefill_v_input_.resize(config_.model_args->get_n_layers()); + decode_k_input_.resize(config_.model_args->get_n_layers()); + decode_v_input_.resize(config_.model_args->get_n_layers()); + for (size_t l = 0; l < config_.model_args->get_n_layers(); l++) { prefill_k_input_[l].resize(prefill_in_size, 0); prefill_v_input_[l].resize(prefill_in_size, 0); decode_k_input_[l].resize(decode_in_size, 0); decode_v_input_[l].resize(decode_in_size, 0); } - attn_mask_.resize(config_.attn_mask_rows * config_.attn_mask_cols, -1.0f); + attn_mask_.resize( + config_.model_args->get_attn_mask_rows() * + config_.model_args->get_attn_mask_cols(), + -1.0f); local_attn_mask_.resize( - config_.attn_mask_rows * config_.attn_mask_cols, -1.0f); + config_.model_args->get_attn_mask_rows() * + config_.model_args->get_attn_mask_cols(), + -1.0f); } void LlmIOManager::copy_prefill_to_decode_input() { - size_t copy_size = prefill_cache_len_ * config_.head_dim * sizeof(float); - for (size_t l = 0; l < config_.n_layers; l++) { - for (size_t h = 0; h < config_.n_kv_heads; h++) { + size_t copy_size = + prefill_cache_len_ * config_.model_args->get_head_dim() * sizeof(float); + for (size_t l = 0; l < config_.model_args->get_n_layers(); l++) { + for (size_t h = 0; h < config_.model_args->get_n_kv_heads(); h++) { const char* src = prefill_k_input_[l].data() + - h * prefill_cache_len_ * config_.head_dim * sizeof(float); + h * prefill_cache_len_ * config_.model_args->get_head_dim() * + sizeof(float); char* dst = decode_k_input_[l].data() + - h * decode_cache_len_ * config_.head_dim * sizeof(float); + h * decode_cache_len_ * config_.model_args->get_head_dim() * + sizeof(float); std::memcpy(dst, src, copy_size); } - for (size_t h = 0; h < config_.n_kv_heads; h++) { + for (size_t h = 0; h < config_.model_args->get_n_kv_heads(); h++) { const char* src = prefill_v_input_[l].data() + - h * prefill_cache_len_ * config_.head_dim * sizeof(float); + h * prefill_cache_len_ * config_.model_args->get_head_dim() * + sizeof(float); char* dst = decode_v_input_[l].data() + - h * decode_cache_len_ * config_.head_dim * sizeof(float); + h * decode_cache_len_ * config_.model_args->get_head_dim() * + sizeof(float); std::memcpy(dst, src, copy_size); } } @@ -289,24 +402,28 @@ void LlmIOManager::append_kv_input_tensors( executorch::aten::ScalarType kv_scalar_type, size_t cache_len, std::vector& inputs) const { - for (size_t l = 0; l < config_.n_layers; l++) { + for (size_t l = 0; l < config_.model_args->get_n_layers(); l++) { auto k_in = executorch::extension::from_blob( k_bufs[l].data(), {1, - static_cast(config_.n_kv_heads), + static_cast( + config_.model_args->get_n_kv_heads()), static_cast(cache_len), - static_cast(config_.head_dim)}, + static_cast( + config_.model_args->get_head_dim())}, kv_scalar_type); k_tensors.push_back(k_in); inputs.emplace_back(k_in); } - for (size_t l = 0; l < config_.n_layers; l++) { + for (size_t l = 0; l < config_.model_args->get_n_layers(); l++) { auto v_in = executorch::extension::from_blob( v_bufs[l].data(), {1, - static_cast(config_.n_kv_heads), + static_cast( + config_.model_args->get_n_kv_heads()), static_cast(cache_len), - static_cast(config_.head_dim)}, + static_cast( + config_.model_args->get_head_dim())}, kv_scalar_type); v_tensors.push_back(v_in); inputs.emplace_back(v_in); @@ -317,8 +434,8 @@ void LlmIOManager::update_attn_mask(size_t pos, size_t seq_len) { const float maskFalse = -255.0f; const float maskTrue = 0.0f; - const size_t rows = config_.attn_mask_rows; - const size_t cols = config_.attn_mask_cols; + const size_t rows = config_.model_args->get_attn_mask_rows(); + const size_t cols = config_.model_args->get_attn_mask_cols(); const size_t start_col = cols - rows; const size_t total_size = rows * cols; @@ -327,9 +444,10 @@ void LlmIOManager::update_attn_mask(size_t pos, size_t seq_len) { if (seq_len > 1) { // Prefill phase - const size_t rows = config_.attn_mask_rows; - const size_t cols = config_.attn_mask_cols; - const size_t start_col = config_.attn_mask_cols - config_.attn_mask_rows; + const size_t rows = config_.model_args->get_attn_mask_rows(); + const size_t cols = config_.model_args->get_attn_mask_cols(); + const size_t start_col = config_.model_args->get_attn_mask_cols() - + config_.model_args->get_attn_mask_rows(); for (int i = 0; i < seq_len; ++i) { int start_idx = cols - rows; @@ -361,9 +479,10 @@ void LlmIOManager::update_local_attn_mask(size_t pos, size_t seq_len) { const float maskFalse = -255.0f; const float maskTrue = 0.0f; - const size_t rows = config_.attn_mask_rows; - const size_t cols = config_.attn_mask_cols; - const size_t sliding_window = config_.sliding_window.value(); + const size_t rows = config_.model_args->get_attn_mask_rows(); + const size_t cols = config_.model_args->get_attn_mask_cols(); + const size_t sliding_window = + config_.model_args->get_sliding_window().value(); const size_t total_size = rows * cols; // Start with a copy of the full attention mask @@ -452,7 +571,8 @@ void LlmIOManager::dump_kv_input(const std::string& phase) { const bool is_prefill = (phase == "prefill"); const auto& k_bufs = is_prefill ? prefill_k_input_ : decode_k_input_; const auto& v_bufs = is_prefill ? prefill_v_input_ : decode_v_input_; - for (size_t l = 0; l < config_.n_layers; l++) { + const size_t n_layers = config_.model_args->get_n_layers(); + for (size_t l = 0; l < n_layers; l++) { dump_buffer(phase + "_k_input", l, k_bufs[l].data(), k_bufs[l].size()); dump_buffer(phase + "_v_input", l, v_bufs[l].data(), v_bufs[l].size()); } @@ -465,7 +585,7 @@ void LlmIOManager::update_kv_cache( size_t new_len, const std::string& phase) { // model_outputs: [logits, k0..k_{L-1}, v0..v_{L-1}] - for (size_t l = 0; l < config_.n_layers; l++) { + for (size_t l = 0; l < config_.model_args->get_n_layers(); l++) { copy_tensor_to_cache( model_outputs[1 + l].toTensor(), new_len, @@ -473,9 +593,9 @@ void LlmIOManager::update_kv_cache( cache_len, current_pos_); } - for (size_t l = 0; l < config_.n_layers; l++) { + for (size_t l = 0; l < config_.model_args->get_n_layers(); l++) { copy_tensor_to_cache( - model_outputs[1 + config_.n_layers + l].toTensor(), + model_outputs[1 + config_.model_args->get_n_layers() + l].toTensor(), new_len, v_bufs[l].data(), cache_len, @@ -490,9 +610,10 @@ void LlmIOManager::dump_kv_output( if (!dump_enabled_) { return; } - for (size_t l = 0; l < config_.n_layers; l++) { + for (size_t l = 0; l < config_.model_args->get_n_layers(); l++) { const auto& k_out = model_outputs[1 + l].toTensor(); - const auto& v_out = model_outputs[1 + config_.n_layers + l].toTensor(); + const auto& v_out = + model_outputs[1 + config_.model_args->get_n_layers() + l].toTensor(); dump_buffer( phase + "_k_output", l, @@ -509,16 +630,22 @@ void LlmIOManager::dump_kv_output( executorch::runtime::Error LlmIOManager::prepare_rotary_embedding_inputs( size_t seq_len, std::vector& inputs) { - const size_t element_size = seq_len * config_.head_dim; - global_cos_data_.resize(element_size); - global_sin_data_.resize(element_size); + const int64_t head_dim = config_.model_args->get_head_dim(); + const int64_t global_head_dim = config_.model_args->get_global_head_dim(); + + const size_t global_element_size = seq_len * global_head_dim; + global_cos_data_.resize(global_element_size); + global_sin_data_.resize(global_element_size); + switch (config_.embedding_strategy) { case EmbeddingStrategy::kGlobalOnly: break; - case EmbeddingStrategy::kGlobalAndLocal: - local_cos_data_.resize(element_size); - local_sin_data_.resize(element_size); + case EmbeddingStrategy::kGlobalAndLocal: { + const size_t local_element_size = seq_len * head_dim; + local_cos_data_.resize(local_element_size); + local_sin_data_.resize(local_element_size); break; + } } auto update_err = rotary_embedding_->update( @@ -536,12 +663,12 @@ executorch::runtime::Error LlmIOManager::prepare_rotary_embedding_inputs( global_cos_data_.data(), {1, static_cast(seq_len), - static_cast(config_.head_dim)}); + static_cast(global_head_dim)}); global_sin_tensor_ = executorch::extension::from_blob( global_sin_data_.data(), {1, static_cast(seq_len), - static_cast(config_.head_dim)}); + static_cast(global_head_dim)}); switch (config_.embedding_strategy) { case EmbeddingStrategy::kGlobalOnly: break; @@ -550,12 +677,12 @@ executorch::runtime::Error LlmIOManager::prepare_rotary_embedding_inputs( local_cos_data_.data(), {1, static_cast(seq_len), - static_cast(config_.head_dim)}); + static_cast(head_dim)}); local_sin_tensor_ = executorch::extension::from_blob( local_sin_data_.data(), {1, static_cast(seq_len), - static_cast(config_.head_dim)}); + static_cast(head_dim)}); break; } diff --git a/examples/samsung/executor_runner/llm/runner/io_manager.h b/examples/samsung/executor_runner/llm/runner/io_manager.h index 99887f8506d..72f55363298 100644 --- a/examples/samsung/executor_runner/llm/runner/io_manager.h +++ b/examples/samsung/executor_runner/llm/runner/io_manager.h @@ -11,7 +11,7 @@ #include #include -#include +#include #include #include @@ -49,172 +49,29 @@ namespace extension = ::executorch::extension; */ class ET_EXPERIMENTAL LlmIOManager : public llm::IOManager { public: - struct Config { - size_t max_seq_len; - size_t max_context_len; - size_t prefill_seq_len; - size_t decode_seq_len; - size_t n_layers; - size_t n_heads; - size_t n_kv_heads; - size_t head_dim; - size_t hidden_dim; - size_t vocab_size; - size_t attn_mask_rows = 128; - size_t attn_mask_cols = 1024; - std::optional sliding_window; - EmbeddingStrategy embedding_strategy = EmbeddingStrategy::kGlobalOnly; - MaskingStrategy masking_strategy = MaskingStrategy::kCausal; - }; - - LlmIOManager(extension::Module& module, Config config); - - const Config& get_config() const { - return config_; - } - - ET_NODISCARD virtual runtime::Error load_rotary_embedding( - std::optional pos_embedding_path); - - size_t get_attention_setup_tensor_count() const { - constexpr size_t kGlobalMaskCount = 1; - constexpr size_t kGlobalRopeCosCount = 1; - constexpr size_t kGlobalRopeSinCount = 1; - constexpr size_t kBaselineMetadataCount = - kGlobalMaskCount + kGlobalRopeCosCount + kGlobalRopeSinCount; - - size_t setup_tensor_count = kBaselineMetadataCount; - - const bool requires_local_mask = - (config_.masking_strategy == MaskingStrategy::kCausalLocal || - config_.masking_strategy == MaskingStrategy::kHybridConvAttention); - if (requires_local_mask) { - constexpr size_t kLocalMaskCount = 1; - setup_tensor_count += kLocalMaskCount; - } - - const bool requires_local_rope = - (config_.embedding_strategy == EmbeddingStrategy::kGlobalAndLocal); - if (requires_local_rope) { - constexpr size_t kLocalRopeCosCount = 1; - constexpr size_t kLocalRopeSinCount = 1; - setup_tensor_count += (kLocalRopeCosCount + kLocalRopeSinCount); - } - - return setup_tensor_count; - } - - /** - * Calculates the total number of tensor slots required in the model's input - * buffer. - * - * Memory Layout Layout: - * [0] : Token/Prompt IDs - * [1 ... N] : Attention Setup Overheads (Masks, RoPE Sin/Cos - * Tensors) [N+1 ... End] : Alternating Key and Value Cache Tensors for - * all Layers - */ - size_t get_input_reserve_size() const { - constexpr size_t kTokenInputCount = 1; - constexpr size_t kKVSlotsPerLayer = 2; - - return kTokenInputCount + get_attention_setup_tensor_count() + - (kKVSlotsPerLayer * config_.n_layers); - } - - virtual ~LlmIOManager() = default; - - ET_NODISCARD virtual runtime::Error load( + LlmIOManager(extension::Module& module, const RunnerConfig& config); + ~LlmIOManager() override = default; + LlmIOManager(const LlmIOManager& other) = delete; + LlmIOManager& operator=(const LlmIOManager& rhs) = delete; + LlmIOManager(LlmIOManager&& other) = delete; + LlmIOManager& operator=(const LlmIOManager&& rhs) = delete; + + ET_NODISCARD runtime::Error load( const std::string& prefill_method, - const std::string& decode_method) override { - (void)prefill_method; - (void)decode_method; - - if ((config_.masking_strategy == MaskingStrategy::kCausalLocal || - config_.masking_strategy == MaskingStrategy::kHybridConvAttention) && - !config_.sliding_window.has_value()) { - ET_LOG( - Error, - "Sliding window configuration is required for the selected masking strategy"); - return runtime::Error::InvalidArgument; - } - - ET_LOG( - Info, - "(%s) Model Config: layers(%zu), kv_heads(%zu), hidden_dim(%zu), head_dim(%zu), \ - vocab(%zu), max_seq(%zu), max_ctx(%zu), prefill_seq(%zu), decode_seq(%zu), \ - mask_rows(%zu), mask_cols(%zu), slide_win(%zu), emb_strat(%d), mask_strat(%d)", - ET_FUNCTION, - config_.n_layers, - config_.n_kv_heads, - config_.hidden_dim, - config_.head_dim, - config_.vocab_size, - config_.max_seq_len, - config_.max_context_len, - config_.prefill_seq_len, - config_.decode_seq_len, - config_.attn_mask_rows, - config_.attn_mask_cols, - config_.sliding_window.value_or(0), - static_cast(config_.embedding_strategy), - static_cast(config_.masking_strategy)); - - return runtime::Error::Ok; - } + const std::string& decode_method) override; - /** - * @brief Reset the IO manager state. - * - * @param prefill_method The prefill method to reset with. - * @param decode_method The decode method to reset with. - */ - ET_NODISCARD virtual runtime::Error reset( + ET_NODISCARD runtime::Error reset( const std::string& prefill_method, - const std::string& decode_method) override { - (void)prefill_method; - (void)decode_method; - - dump_step_ = 0; - - current_pos_ = 0; - prefill_done_ = false; - - for (size_t l = 0; l < config_.n_layers; l++) { - std::fill(prefill_k_input_[l].begin(), prefill_k_input_[l].end(), 0); - std::fill(prefill_v_input_[l].begin(), prefill_v_input_[l].end(), 0); - std::fill(decode_k_input_[l].begin(), decode_k_input_[l].end(), 0); - std::fill(decode_v_input_[l].begin(), decode_v_input_[l].end(), 0); - } - std::fill(attn_mask_.begin(), attn_mask_.end(), 0.0f); + const std::string& decode_method) override; - return runtime::Error::Ok; - } + ET_NODISCARD virtual runtime::Error load_rotary_embedding( + std::optional pos_embedding_path); - /** - * @brief Prepare inputs for the prefill phase of LLM inference. - * - * @param input The input tensor containing token IDs. - * @param start_pos The tensor containing the starting position of the current - * input within the context. - * @param prefill_method The prefill method to prepare inputs for. - * @return std::vector Vector of prepared inputs - * for the prefill method. - */ virtual runtime::Result> prepare_prefill( const extension::TensorPtr& input, const extension::TensorPtr& start_pos, const std::string& prefill_method) override; - /** - * @brief Prepare inputs for the prefill phase using the default method name. - * - * Uses "forward" as the prefill method. - * - * @param input The input tensor containing token IDs. - * @param start_pos The tensor containing the starting position. - * @return Vector of prepared inputs for the prefill method. - */ runtime::Result> prepare_prefill( const extension::TensorPtr& input, const extension::TensorPtr& start_pos) { @@ -222,30 +79,11 @@ class ET_EXPERIMENTAL LlmIOManager : public llm::IOManager { input, start_pos, example::default_method_name_prefill); } - /** - * @brief Prepare inputs for the decode phase of LLM inference. - * - * @param input The input tensor containing token IDs. - * @param start_pos The tensor containing the starting position of the current - * input within the context. - * @param decode_method The decode method to prepare inputs for. - * @return std::vector Vector of prepared inputs - * for the decode method. - */ - virtual runtime::Result> prepare_decode( + runtime::Result> prepare_decode( const extension::TensorPtr& input, const extension::TensorPtr& start_pos, const std::string& decode_method) override; - /** - * @brief Prepare inputs for the decode phase using the default method name. - * - * Uses "forward" as the decode method. - * - * @param input The input tensor containing token IDs. - * @param start_pos The tensor containing the starting position. - * @return Vector of prepared inputs for the decode method. - */ runtime::Result> prepare_decode( const extension::TensorPtr& input, const extension::TensorPtr& start_pos) { @@ -253,54 +91,45 @@ class ET_EXPERIMENTAL LlmIOManager : public llm::IOManager { input, start_pos, example::default_method_name_decode); } - /** - * @brief Process and update internal state with outputs from the prefill - * phase. - * - * @param prefill_method The prefill method to update with outputs. - * @param model_outputs Vector of outputs from the prefill method execution. - */ - ET_NODISCARD virtual runtime::Error update_prefill( + ET_NODISCARD runtime::Error update_prefill( const std::vector& model_outputs, const std::string& prefill_method) override; - /** - * @brief Process outputs from the prefill phase using the default method. - * - * Uses "forward" as the prefill method. - * - * @param model_outputs Vector of outputs from the prefill execution. - * @return Error code. - */ ET_NODISCARD runtime::Error update_prefill( const std::vector& model_outputs) { return update_prefill(model_outputs, example::default_method_name_prefill); } - /** - * @brief Process and update internal state with outputs from the decode - * phase. - * - * @param decode_method The decode method to update with outputs. - * @param model_outputs Vector of outputs from the decode method execution. - */ - ET_NODISCARD virtual runtime::Error update_decode( + ET_NODISCARD runtime::Error update_decode( const std::vector& model_outputs, const std::string& decode_method) override; - /** - * @brief Process outputs from the decode phase using the default method. - * - * Uses "forward" as the decode method. - * - * @param model_outputs Vector of outputs from the decode execution. - * @return Error code. - */ ET_NODISCARD runtime::Error update_decode( const std::vector& model_outputs) { return update_decode(model_outputs, example::default_method_name_decode); } + const RunnerConfig& get_config() const { + return config_; + } + + size_t get_attention_setup_tensor_count() const; + /** + * Calculates the total number of tensor slots required in the model's input + * buffer. + * + * Memory Layout Layout: + * [0] : Token/Prompt IDs + * [1 ... N] : Attention Setup Overheads (Masks, RoPE Sin/Cos + * Tensors) [N+1 ... End] : Alternating Key and Value Cache Tensors for + * all Layers + */ + size_t get_input_reserve_size() const { + auto model_config = getModelVersionConfig(config_.model_version); + return model_config.v_cache_tensor_index + + model_config.kv_cache_tensor_count; + } + size_t current_pos() const { return current_pos_; } @@ -363,11 +192,16 @@ class ET_EXPERIMENTAL LlmIOManager : public llm::IOManager { char* cache, size_t cache_len, size_t pos) { - for (size_t h = 0; h < config_.n_kv_heads; h++) { - const T* src_head = src + h * new_len * config_.head_dim; - char* dst_head = cache + h * cache_len * config_.head_dim * sizeof(T) + - pos * config_.head_dim * sizeof(T); - std::memcpy(dst_head, src_head, new_len * config_.head_dim * sizeof(T)); + for (size_t h = 0; h < config_.model_args->get_n_kv_heads(); h++) { + const T* src_head = + src + h * new_len * config_.model_args->get_head_dim(); + char* dst_head = cache + + h * cache_len * config_.model_args->get_head_dim() * sizeof(T) + + pos * config_.model_args->get_head_dim() * sizeof(T); + std::memcpy( + dst_head, + src_head, + new_len * config_.model_args->get_head_dim() * sizeof(T)); } } @@ -429,7 +263,7 @@ class ET_EXPERIMENTAL LlmIOManager : public llm::IOManager { * positions [0, pos + r] (value 1) and masks future positions (value 0). */ void update_local_attn_mask(size_t pos, size_t seq_len); - Config config_; + const RunnerConfig& config_; size_t prefill_cache_len_; size_t decode_cache_len_; diff --git a/examples/samsung/executor_runner/llm/runner/rotary_embedding.cpp b/examples/samsung/executor_runner/llm/runner/rotary_embedding.cpp index a2d833f9e37..dc8ab5cc0ef 100644 --- a/examples/samsung/executor_runner/llm/runner/rotary_embedding.cpp +++ b/examples/samsung/executor_runner/llm/runner/rotary_embedding.cpp @@ -1,79 +1,270 @@ /* - * Copyright (c) 2025 Samsung Electronics Co. LTD + * Copyright (c) 2026 Samsung Electronics Co. LTD * All rights reserved. * * This source code is licensed under the BSD-style license found in the * LICENSE file in the root directory of this source tree. */ -#include +#include +#include #include #include #include +#include + namespace example { namespace runtime = ::executorch::runtime; -RotaryEmbedding::RotaryEmbedding(Config config) : config_(config) {} +RotaryEmbedding::RotaryEmbedding(const RunnerConfig& config) + : config_(config), + loaded_(false), + seq_len_(config.model_args->get_max_seq_len()), + base_freq_global_(1000000.0f), + base_freq_local_(1000000.0f), + rope_percentage_(1.0f) { + switch (config_.model_version) { + case DecoderModelVersion::kGemma4: { + auto* args = static_cast(config_.model_args.get()); + normalized_global_head_dim_ = args->text_config.global_head_dim; + normalized_local_head_dim_ = args->text_config.head_dim; + base_freq_global_ = args->text_config.base_freq_global; + base_freq_local_ = args->text_config.base_freq_local; + rope_percentage_ = args->text_config.rope_percentage; + break; + } + case DecoderModelVersion::kLlama3: { + auto* args = static_cast(config_.model_args.get()); + normalized_global_head_dim_ = args->get_head_dim(); + normalized_local_head_dim_ = args->get_head_dim(); + base_freq_global_ = 500000.0f; + rope_percentage_ = 1.0f; + break; + } + case DecoderModelVersion::kGemma3: { + auto* args = static_cast(config_.model_args.get()); + normalized_global_head_dim_ = args->get_head_dim(); + normalized_local_head_dim_ = args->get_head_dim(); + base_freq_global_ = 1000000.0f; + base_freq_local_ = 10000.0f; + rope_percentage_ = 1.0f; + break; + } + default: { + break; + } + } + + // TODO(brucekk.kim): size should be configurable + global_bytes_ = config_.model_args->get_max_seq_len() * + normalized_global_head_dim_ * sizeof(float); + local_bytes_ = config_.model_args->get_max_seq_len() * + normalized_local_head_dim_ * sizeof(float); +} runtime::Error RotaryEmbedding::load( std::optional pos_embedding_path) { - if (!pos_embedding_path.has_value()) { + allocate_buffers(); + + if (pos_embedding_path.has_value() && !pos_embedding_path->empty()) { + runtime::Error err = load_from_file(pos_embedding_path.value()); + if (err == runtime::Error::Ok) { + return runtime::Error::Ok; + } ET_LOG( - Error, - "Missing position embedding path. Verify that the file path is provided in the configuration."); - return runtime::Error::InvalidArgument; + Info, + "Failed to load embeddings from file, falling back to precompute."); + } else { + ET_LOG( + Info, + "No valid position embedding path provided, falling back to precompute."); } - const size_t total_bytes = - config_.max_seq_len * config_.head_dim * sizeof(float); - global_cos_ = std::make_unique(total_bytes); - global_sin_ = std::make_unique(total_bytes); - if (config_.strategy == EmbeddingStrategy::kGlobalAndLocal) { - local_cos_ = std::make_unique(total_bytes); - local_sin_ = std::make_unique(total_bytes); - } + return load_by_precompute_rope_tables(); +} - std::ifstream file(pos_embedding_path.value(), std::ios::binary); - if (!file) { - ET_LOG( - Error, - "File I/O error: Cannot open position embedding file at '%s'.", - pos_embedding_path.value().c_str()); - return runtime::Error::InvalidArgument; - } +runtime::Error RotaryEmbedding::load_from_file(const std::string& path) { + std::ifstream file(path, std::ios::binary); + ET_CHECK_OR_RETURN_ERROR( + file.is_open(), + InvalidArgument, + "File I/O error: Cannot open file at '%s'.", + path.c_str()); file.seekg(0, std::ios::end); size_t file_size = file.tellg(); file.seekg(0, std::ios::beg); - int bin_num = - (config_.strategy == EmbeddingStrategy::kGlobalAndLocal) ? 4 : 2; - size_t expected_size = total_bytes * bin_num; - if (file_size != expected_size) { - ET_LOG( - Error, - "Embedding file size mismatch. Expected %zu bytes (total_bytes: %zu * bin_num: %d based on strategy), but got %zu bytes. File may be corrupted or from a different model version.", - expected_size, - total_bytes, - bin_num, - file_size); - return runtime::Error::InvalidArgument; + size_t expected_size = global_bytes_ * 2; + if (config_.embedding_strategy == EmbeddingStrategy::kGlobalAndLocal) { + expected_size += local_bytes_ * 2; } - file.read(global_cos_.get(), total_bytes); - file.read(global_sin_.get(), total_bytes); - if (config_.strategy == EmbeddingStrategy::kGlobalAndLocal) { - file.read(local_cos_.get(), total_bytes); - file.read(local_sin_.get(), total_bytes); + ET_CHECK_OR_RETURN_ERROR( + file_size == expected_size, + InvalidArgument, + "Embedding file size mismatch. Expected %zu, got %zu.", + expected_size, + file_size); + + // ========================================================================= + // BINARY ASSET LAYOUT ASSUMPTION WARNING: + // The raw .bin file contains NO metadata, headers, or structural markers. + // This code relies entirely on a strict, sequential file-pointer progression. + // Any discrepancy in the packing sequence or byte padding during model + // export will cause silent memory corruption or misaligned frequencies. + // + // Expected File Byte Ordering Layout: + // +-----------------------------+-----------------------------+ + // | global_cos (global_bytes_) | global_sin (global_bytes_) | + // +-----------------------------+-----------------------------+ + // | local_cos (local_bytes_)* | local_sin (local_bytes_)* | + // +-----------------------------+-----------------------------+ + // (*Optional: Only if embedding_strategy == kGlobalAndLocal) + // ========================================================================= + + // Stream out the mandatory global attention layer matrices + file.read(global_cos_.get(), global_bytes_); + file.read(global_sin_.get(), global_bytes_); + if (config_.embedding_strategy == EmbeddingStrategy::kGlobalAndLocal) { + file.read(local_cos_.get(), local_bytes_); + file.read(local_sin_.get(), local_bytes_); } - file.close(); loaded_ = true; return runtime::Error::Ok; } +runtime::Error RotaryEmbedding::load_by_precompute_rope_tables() { + float* global_cos_ptr = reinterpret_cast(global_cos_.get()); + float* global_sin_ptr = reinterpret_cast(global_sin_.get()); + + assert( + reinterpret_cast(global_cos_ptr) % alignof(float) == 0 && + "global_cos_ is misaligned for float access!"); + assert( + reinterpret_cast(global_sin_ptr) % alignof(float) == 0 && + "global_sin_ is misaligned for float access!"); + + compute_rope_table( + global_cos_ptr, + global_sin_ptr, + seq_len_, + normalized_global_head_dim_, + base_freq_global_, + rope_percentage_); + + if (config_.embedding_strategy == EmbeddingStrategy::kGlobalAndLocal) { + float* local_cos_ptr = reinterpret_cast(local_cos_.get()); + float* local_sin_ptr = reinterpret_cast(local_sin_.get()); + + assert( + reinterpret_cast(local_cos_ptr) % alignof(float) == 0 && + "local_cos_ is misaligned for float access!"); + assert( + reinterpret_cast(local_sin_ptr) % alignof(float) == 0 && + "local_sin_ is misaligned for float access!"); + + compute_rope_table( + local_cos_ptr, + local_sin_ptr, + seq_len_, + normalized_local_head_dim_, + base_freq_local_, + rope_percentage_); + } + loaded_ = true; + ET_LOG(Info, "Rotary Embedding tables precomputed successfully."); + return runtime::Error::Ok; +} + +void RotaryEmbedding::allocate_buffers() { + global_cos_ = std::make_unique(global_bytes_); + global_sin_ = std::make_unique(global_bytes_); + + if (config_.embedding_strategy == EmbeddingStrategy::kGlobalAndLocal) { + local_cos_ = std::make_unique(local_bytes_); + local_sin_ = std::make_unique(local_bytes_); + } +} + +runtime::Error RotaryEmbedding::compute_rope_table( + float* cos_ptr, + float* sin_ptr, + size_t seq_len, + size_t dim, + float base, + float rope_pct) { + ET_CHECK_OR_RETURN_ERROR( + cos_ptr != nullptr, InvalidArgument, "cos_ptr cannot be nullptr."); + ET_CHECK_OR_RETURN_ERROR( + sin_ptr != nullptr, InvalidArgument, "sin_ptr cannot be nullptr."); + ET_CHECK_OR_RETURN_ERROR( + dim % 2 == 0, + InvalidArgument, + "The head dimension (dim) must be an even number for rotary pairs."); + ET_CHECK_OR_RETURN_ERROR( + base > 0.0f, + InvalidArgument, + "The RoPE base frequency must be greater than zero."); + ET_CHECK_OR_RETURN_ERROR( + rope_pct >= 0.0f && rope_pct <= 1.0f, + InvalidArgument, + "The rope_pct parameter must be a percentage value between 0.0 and 1.0 inclusive."); + + const size_t half_dim = dim / 2; + size_t rotary_dim = static_cast(static_cast(dim) * rope_pct); + rotary_dim = std::min(rotary_dim, dim); + rotary_dim &= ~size_t{1}; + + ET_CHECK_OR_RETURN_ERROR( + rotary_dim > 0, + InvalidArgument, + "rope_pct results in zero rotary dimensions."); + + const size_t rotary_half = rotary_dim / 2; + + for (size_t pos = 0; pos < seq_len; ++pos) { + float* cos_row = cos_ptr + pos * dim; + float* sin_row = sin_ptr + pos * dim; + + for (size_t i = 0; i < dim; ++i) { + cos_ptr[i] = 1.0f; + sin_ptr[i] = 0.0f; + } + + for (size_t i = 0; i < rotary_half; ++i) { + // inv_freq_i = base^(-2i / rotary_dim) + const float exponent = + -2.0f * static_cast(i) / static_cast(rotary_dim); + const float inv_freq = std::pow(base, exponent); + const float angle = static_cast(pos) * inv_freq; + + const float c = std::cos(angle); + const float s = std::sin(angle); + + // Split-half convention: + // + // first half: i + // second half: half_dim + i + // + // x = [x_0 ... x_{d/2-1} | x_{d/2} ... x_{d-1}] + // where each rotary pair is: (x_i, x_{i + d/2}) + // Both coordinates of the pair use the same cos/sin angle. + + cos_row[i] = c; + sin_row[i] = s; + + cos_row[half_dim + i] = c; + sin_row[half_dim + i] = s; + } + } + + return runtime::Error::Ok; +} + runtime::Error RotaryEmbedding::update( const std::vector& pos_ids, float* global_cos, @@ -90,26 +281,35 @@ runtime::Error RotaryEmbedding::update( if (global_cos == nullptr || global_sin == nullptr) { return runtime::Error::InvalidArgument; } - if (config_.strategy == EmbeddingStrategy::kGlobalAndLocal && + if (config_.embedding_strategy == EmbeddingStrategy::kGlobalAndLocal && (local_cos == nullptr || local_sin == nullptr)) { return runtime::Error::InvalidArgument; } size_t pos = static_cast(pos_ids[0]); size_t length = pos_ids.size(); - size_t head_dim = config_.head_dim; - size_t bytes_to_copy = length * head_dim * sizeof(float); - size_t element_offset = pos * head_dim * sizeof(float); - std::memcpy(global_cos, global_cos_.get() + element_offset, bytes_to_copy); - std::memcpy(global_sin, global_sin_.get() + element_offset, bytes_to_copy); + // Global Copy + size_t g_head_dim = normalized_global_head_dim_; + size_t g_bytes_to_copy = length * g_head_dim * sizeof(float); + size_t g_element_offset = pos * g_head_dim * sizeof(float); + std::memcpy( + global_cos, global_cos_.get() + g_element_offset, g_bytes_to_copy); + std::memcpy( + global_sin, global_sin_.get() + g_element_offset, g_bytes_to_copy); - if (config_.strategy == EmbeddingStrategy::kGlobalAndLocal) { - std::memcpy(local_cos, local_cos_.get() + element_offset, bytes_to_copy); - std::memcpy(local_sin, local_sin_.get() + element_offset, bytes_to_copy); + // Local Copy + if (config_.embedding_strategy == EmbeddingStrategy::kGlobalAndLocal) { + size_t l_head_dim = config_.model_args->get_head_dim(); + size_t l_bytes_to_copy = length * l_head_dim * sizeof(float); + size_t l_element_offset = pos * l_head_dim * sizeof(float); + std::memcpy( + local_cos, local_cos_.get() + l_element_offset, l_bytes_to_copy); + std::memcpy( + local_sin, local_sin_.get() + l_element_offset, l_bytes_to_copy); } return runtime::Error::Ok; } -} // namespace example +} // namespace example \ No newline at end of file diff --git a/examples/samsung/executor_runner/llm/runner/rotary_embedding.h b/examples/samsung/executor_runner/llm/runner/rotary_embedding.h index 5b9337863e4..3f9b0882c49 100644 --- a/examples/samsung/executor_runner/llm/runner/rotary_embedding.h +++ b/examples/samsung/executor_runner/llm/runner/rotary_embedding.h @@ -1,5 +1,5 @@ /* - * Copyright (c) 2025 Samsung Electronics Co. LTD + * Copyright (c) 2026 Samsung Electronics Co. LTD * All rights reserved. * * This source code is licensed under the BSD-style license found in the @@ -8,26 +8,21 @@ #pragma once -#include -#include #include #include #include #include +#include +#include + namespace example { namespace runtime = ::executorch::runtime; class RotaryEmbedding { public: - struct Config { - size_t head_dim; - size_t max_seq_len; - EmbeddingStrategy strategy = EmbeddingStrategy::kGlobalOnly; - }; - - RotaryEmbedding(Config config); + RotaryEmbedding(const RunnerConfig& config); ~RotaryEmbedding() = default; runtime::Error load(std::optional pos_embedding_path); @@ -39,8 +34,29 @@ class RotaryEmbedding { float* local_sin); private: - Config config_; - bool loaded_ = false; + runtime::Error load_from_file(const std::string& path); + runtime::Error load_by_precompute_rope_tables(); + void allocate_buffers(); + runtime::Error compute_rope_table( + float* cos_ptr, + float* sin_ptr, + size_t seq_len, + size_t dim, + float base, + float rope_pct); + + const RunnerConfig& config_; + bool loaded_; + + size_t seq_len_; + size_t normalized_global_head_dim_; + size_t normalized_local_head_dim_; + size_t global_bytes_; + size_t local_bytes_; + float base_freq_global_; + float base_freq_local_; + float rope_percentage_; + std::unique_ptr global_cos_; std::unique_ptr global_sin_; std::unique_ptr local_cos_; diff --git a/examples/samsung/executor_runner/llm/runner/runner.cpp b/examples/samsung/executor_runner/llm/runner/runner.cpp index 935e0e73505..467d9049f0b 100644 --- a/examples/samsung/executor_runner/llm/runner/runner.cpp +++ b/examples/samsung/executor_runner/llm/runner/runner.cpp @@ -9,7 +9,7 @@ // A simple LLM runner that includes preprocessing and post processing logic. // The module takes in a string as input and emits a string as output. -#include +#include #include #include #include @@ -19,7 +19,7 @@ #include #include -#include +#include #include #include #include @@ -64,6 +64,17 @@ std::string LLMRunner::get_formatted_prompt( formatted_prompt.append("\n"); } break; + case DecoderModelVersion::kGemma4: + if (!system_prompt.empty()) { + formatted_prompt.append("<|turn>system "); + formatted_prompt.append(system_prompt); + formatted_prompt.append("\n"); + } + formatted_prompt.append("<|turn>user "); + formatted_prompt.append(prompt); + formatted_prompt.append("\n"); + formatted_prompt.append("<|turn>model\n"); + break; case DecoderModelVersion::kGeneral: formatted_prompt.append(prompt); break; @@ -75,40 +86,14 @@ std::string LLMRunner::get_formatted_prompt( } /** - * @brief Loads a tokenizer from the specified path - * * This function creates and initializes a tokenizer from a file based on the * specified type. Supported types are: "bpe", "tiktoken", "hf", and "sp". - * - * @param type Type of tokenizer to load ("bpe", "tiktoken", "hf", or "sp") - * @param tokenizer_path Path to the tokenizer file - * @param version Version of the tokenizer (used for tiktoken type) - * @return std::unique_ptr Initialized tokenizer - * instance, or nullptr on failure */ ET_EXPERIMENTAL std::unique_ptr load_tokenizer( const std::string& type, const std::string& tokenizer_path, Version version); -/** - * @brief Creates a LLMRunner instance with the specified model and - * tokenizer - * - * This factory function creates and initializes a LLMRunner with all - * necessary components for text generation using the specified model and - * tokenizer. - * - * @param paths RunnerPaths containing model and tokenizer paths - * @param config RunnerConfig containing runtime configuration - * @param tokenizer Initialized tokenizer instance - * @return std::unique_ptr Initialized LLMRunner instance - */ -ET_EXPERIMENTAL std::unique_ptr create_llm_runner( - const RunnerPaths& runner_paths, - RunnerConfig runner_config, - std::unique_ptr<::tokenizers::Tokenizer> tokenizer); - std::unique_ptr create_llm_runner( const RunnerPaths& runner_paths, RunnerConfig runner_config) { @@ -120,7 +105,6 @@ std::unique_ptr create_llm_runner( runner_config.load_from_json(); - // Create and load tokenizer std::unique_ptr<::tokenizers::Tokenizer> tokenizer = load_tokenizer( runner_config.tokenizer_type, runner_paths.tokenizer_path, @@ -170,9 +154,11 @@ LLMRunner::LLMRunner( std::unique_ptr text_token_generator, std::unique_ptr stats, std::unique_ptr io_manager, + RunnerConfig config, DecoderModelVersion model_version, float temperature) - : tokenizer_(std::move(tokenizer)), + : config_(std::move(config)), + tokenizer_(std::move(tokenizer)), metadata_(std::move(metadata)), module_(std::move(module)), text_decoder_runner_(std::move(text_decoder_runner)), @@ -216,105 +202,6 @@ void replace_all_newlines(std::string& str) { } } -void RunnerConfig::load_from_json() { - if (config_path.empty()) { - ET_LOG(Info, "Config path is empty, skipping JSON load"); - return; - } - - std::ifstream file(config_path); - if (!file.is_open()) { - ET_LOG(Error, "Could not open config file: %s", config_path.c_str()); - return; - } - - json j = json::parse(file, nullptr, false); - if (j.is_discarded()) { - ET_LOG( - Error, "Invalid JSON format in config file: %s", config_path.c_str()); - return; - } - - if (j.contains("hidden_dim") && j["hidden_dim"].is_number_integer()) { - hidden_dim = j["hidden_dim"].get(); - } - if (j.contains("n_heads") && j["n_heads"].is_number_integer()) { - n_heads = j["n_heads"].get(); - } - if (j.contains("n_kv_heads") && j["n_kv_heads"].is_number_integer()) { - n_kv_heads = j["n_kv_heads"].get(); - } - if (j.contains("n_layers") && j["n_layers"].is_number_integer()) { - n_layers = j["n_layers"].get(); - } - if (j.contains("head_dim") && j["head_dim"].is_number_integer()) { - head_dim = j["head_dim"].get(); - } - if (j.contains("vocab_size") && j["vocab_size"].is_number_integer()) { - vocab_size = j["vocab_size"].get(); - } - if (j.contains("max_seq_len") && j["max_seq_len"].is_number_integer()) { - max_seq_len = j["max_seq_len"].get(); - } - if (j.contains("max_context_len") && - j["max_context_len"].is_number_integer()) { - max_context_len = j["max_context_len"].get(); - } - if (j.contains("prefill_seq_len") && - j["prefill_seq_len"].is_number_integer()) { - prefill_seq_len = j["prefill_seq_len"].get(); - } - if (j.contains("decode_seq_len") && j["decode_seq_len"].is_number_integer()) { - decode_seq_len = j["decode_seq_len"].get(); - } - if (j.contains("attn_mask_rows") && j["attn_mask_rows"].is_number_integer()) { - attn_mask_rows = j["attn_mask_rows"].get(); - } - if (j.contains("attn_mask_cols") && j["attn_mask_cols"].is_number_integer()) { - attn_mask_cols = j["attn_mask_cols"].get(); - } - if (j.contains("sliding_window") && j["sliding_window"].is_number_integer()) { - sliding_window = j["sliding_window"].get(); - } else { - sliding_window = std::nullopt; - } - - if (j.contains("embedding_strategy") && j["embedding_strategy"].is_string()) { - std::string strategy = j["embedding_strategy"].get(); - if (strategy == "global_only") { - embedding_strategy = EmbeddingStrategy::kGlobalOnly; - } else if (strategy == "global_and_local") { - embedding_strategy = EmbeddingStrategy::kGlobalAndLocal; - } - } - - if (j.contains("masking_strategy") && j["masking_strategy"].is_string()) { - std::string strategy = j["masking_strategy"].get(); - if (strategy == "causal") { - masking_strategy = MaskingStrategy::kCausal; - } else if (strategy == "causal_local") { - masking_strategy = MaskingStrategy::kCausalLocal; - } else if (strategy == "hybrid_conv") { - masking_strategy = MaskingStrategy::kHybridConvAttention; - } - } - - if (j.contains("model_version") && j["model_version"].is_string()) { - std::string version_str = j["model_version"].get(); - if (version_str == "llama3") { - model_version = DecoderModelVersion::kLlama3; - } else if (version_str == "gemma3") { - model_version = DecoderModelVersion::kGemma3; - } else if (version_str == "general") { - model_version = DecoderModelVersion::kGeneral; - } - } - - if (j.contains("tokenizer_type") && j["tokenizer_type"].is_string()) { - tokenizer_type = j["tokenizer_type"].get(); - } -} - runtime::Error LLMRunner::generate_from_pos( const std::string& prompt, int64_t start_pos, @@ -350,9 +237,9 @@ runtime::Error LLMRunner::generate_from_pos( token_callback(piece); } }; + // First token time only measures the time it takes to encode the prompt and // return a response token. - stats_->inference_start_ms = llm::time_in_ms(); shouldStop_ = false; @@ -379,6 +266,7 @@ runtime::Error LLMRunner::generate_from_pos( int num_prompt_tokens = prompt_tokens.size(); // Reduce max_context_len by start_pos + // TODO(brucekk.kim): Check where to fetch from: metatdata or runner_config int64_t max_context_len = metadata_.at(llm::kMaxContextLen) - start_pos; ET_CHECK_OR_RETURN_ERROR( num_prompt_tokens >= 1, @@ -555,16 +443,23 @@ std::unordered_map get_llm_metadata( // Initialize metadata from JSON config values std::unordered_map metadata({ {kMetadataMap.at(llm::kEnableDynamicShape), false}, - {kMetadataMap.at(llm::kMaxSeqLen), config.max_seq_len}, - {kMetadataMap.at(llm::kMaxContextLen), config.max_context_len}, + {kMetadataMap.at(llm::kMaxSeqLen), + config.model_args ? config.model_args->get_max_seq_len() : 128}, + {kMetadataMap.at(llm::kMaxContextLen), + config.model_args ? config.model_args->get_max_context_len() : 128}, {kMetadataMap.at(llm::kUseKVCache), true}, {kMetadataMap.at(llm::kUseSDPAWithKVCache), false}, {kMetadataMap.at(example::kPrefillArLen), 0}, - {kMetadataMap.at(example::kNumOfHiddenDim), config.hidden_dim}, - {kMetadataMap.at(example::kNumOfHeads), config.n_heads}, - {kMetadataMap.at(example::kNumOfKVHeads), config.n_kv_heads}, - {kMetadataMap.at(example::kNumOfLayers), config.n_layers}, - {kMetadataMap.at(llm::kHeadDim), config.head_dim}, + {kMetadataMap.at(example::kNumOfHiddenDim), + config.model_args ? config.model_args->get_hidden_dim() : 0}, + {kMetadataMap.at(example::kNumOfHeads), + config.model_args ? config.model_args->get_n_heads() : 0}, + {kMetadataMap.at(example::kNumOfKVHeads), + config.model_args ? config.model_args->get_n_kv_heads() : 0}, + {kMetadataMap.at(example::kNumOfLayers), + config.model_args ? config.model_args->get_n_layers() : 0}, + {kMetadataMap.at(llm::kHeadDim), + config.model_args ? config.model_args->get_head_dim() : 0}, }); auto method_names_result = module->method_names(); @@ -588,7 +483,7 @@ std::unordered_map get_llm_metadata( method_name.c_str(), value); } - ET_LOG(Debug, "Metadata: %s = %" PRId64, method_name.c_str(), value); + ET_LOG(Info, "Metadata: %s = %" PRId64, method_name.c_str(), value); } // Set tokenizer-related metadata metadata[llm::kBosId] = tokenizer->bos_tok(); @@ -624,7 +519,7 @@ std::unordered_set get_eos_ids( return eos_ids; } -std::unique_ptr create_llm_runner( +std::unique_ptr create_llm_runner( const RunnerPaths& runner_paths, RunnerConfig runner_config, std::unique_ptr<::tokenizers::Tokenizer> tokenizer) { @@ -633,6 +528,10 @@ std::unique_ptr create_llm_runner( return nullptr; } + DecoderModelVersion model_version = runner_config.model_version; + float temperature = runner_config.temperature; + bool is_decode_only = runner_config.is_decode_only; + std::unique_ptr module; if (runner_paths.data_path.has_value()) { module = std::make_unique( @@ -652,14 +551,14 @@ std::unique_ptr create_llm_runner( auto methods = methods_res.get(); std::vector register_methods; - if (runner_config.is_decode_only) { + if (is_decode_only) { register_methods.emplace_back(default_method_name_decode); } else { for (auto method : methods) { if (kMetadataMap.find(method) != kMetadataMap.end()) { continue; } - ET_LOG(Info, "Register method \"%s\" to decoder", method.c_str()); + ET_LOG(Info, "Register method: \"%s\"", method.c_str()); register_methods.emplace_back(method); } } @@ -668,71 +567,48 @@ std::unique_ptr create_llm_runner( auto metadata = get_llm_metadata(tokenizer.get(), module.get(), runner_config); - example::LlmIOManager::Config conf; - conf.n_layers = metadata.at(example::kNumOfLayers); - conf.n_heads = runner_config.n_heads; - conf.n_kv_heads = metadata.at(example::kNumOfKVHeads); - conf.hidden_dim = runner_config.hidden_dim; - conf.vocab_size = runner_config.vocab_size; - conf.max_seq_len = runner_config.max_seq_len; - conf.max_context_len = runner_config.max_context_len; - conf.decode_seq_len = runner_config.decode_seq_len; - conf.prefill_seq_len = (metadata.at(example::kPrefillArLen) > 1) - ? (metadata.at(example::kPrefillArLen)) - : (runner_config.prefill_seq_len != 0 ? runner_config.prefill_seq_len - : metadata.at(llm::kMaxSeqLen)); - conf.embedding_strategy = runner_config.embedding_strategy; - conf.masking_strategy = runner_config.masking_strategy; - conf.attn_mask_rows = runner_config.attn_mask_rows; - conf.attn_mask_cols = runner_config.attn_mask_cols; - conf.sliding_window = runner_config.sliding_window; - - if (runner_config.is_decode_only) { - conf.prefill_seq_len = 1; - } + auto stats = std::make_unique(); + llm::Stats* stats_ptr = stats.get(); - if (metadata.count(llm::kHeadDim)) { - conf.head_dim = metadata.at(llm::kHeadDim); - } else { - conf.head_dim = runner_config.head_dim; - } + ET_LOG(Info, "Instantiate LLMRunner with null placeholders"); + std::unique_ptr runner = std::make_unique( + std::move(metadata), + std::move(tokenizer), + std::move(module), + nullptr, // text_decoder_runner placeholder + nullptr, // text_prefiller placeholder + nullptr, // text_token_generator placeholder + std::move(stats), + nullptr, // io_manager placeholder + std::move(runner_config), + model_version, + temperature); + const auto& config = runner->get_config(); + + ET_LOG(Info, "Instantiate IO Manager"); auto io_manager = - std::make_unique(*module.get(), conf); + std::make_unique(*runner->get_module(), config); - // Every embedding_strategy (kGlobalOnly, kGlobalAndLocal) requires the - // global rotary embedding to be loaded, so a missing/failed load is fatal - // rather than something we can silently continue past. - if (!runner_paths.pos_embedding_path.has_value()) { - ET_LOG(Error, "pos_embedding_path is required to load rotary embedding."); - return nullptr; - } - if (io_manager->load_rotary_embedding( - runner_paths.pos_embedding_path.value()) != + ET_LOG(Info, "Start loading Rotary Embedding"); + if (io_manager->load_rotary_embedding(runner_paths.pos_embedding_path) != executorch::runtime::Error::Ok) { - ET_LOG( - Error, - "Failed to load pos_embedding file at '%s'.", - runner_paths.pos_embedding_path.value().c_str()); + ET_LOG(Error, "Failed to load or precompute pos_embedding."); return nullptr; } auto eos_ids = std::make_unique>( - get_eos_ids(tokenizer.get(), module.get())); + get_eos_ids(runner->get_tokenizer(), runner->get_module())); - // Create text_decoder_runner. Use a shared_ptr so that it can be shared with - // TextPrefiller and TextTokenGenerator - auto text_decoder_runner = - std::make_unique(module.get(), io_manager.get()); + auto text_decoder_runner = std::make_unique( + runner->get_module(), io_manager.get()); for (auto method : register_methods) { text_decoder_runner->set_method(method); } - // TODO(brucekkim): Refactor to use dynamic runtime checks (e.g., tracking the - // loop step counter). Relying solely on configuration properties will fail if - // a single-token prompt is processed. - bool separated_prefill_decode = (conf.prefill_seq_len > 1); + bool separated_prefill_decode = + (config.model_args->get_prefill_seq_len() > 1); ET_LOG( Info, "separated_prefill_decode(%s)", @@ -740,29 +616,22 @@ std::unique_ptr create_llm_runner( auto text_prefiller = std::make_unique( text_decoder_runner.get(), - metadata.at(llm::kUseKVCache), - metadata.at(llm::kEnableDynamicShape) | (separated_prefill_decode), - conf.prefill_seq_len); + runner->get_metadata().at(llm::kUseKVCache), + runner->get_metadata().at(llm::kEnableDynamicShape) | + (separated_prefill_decode), + config.model_args->get_prefill_seq_len()); - auto stats = std::make_unique(); auto text_token_generator = std::make_unique( - tokenizer.get(), + runner->get_tokenizer(), text_decoder_runner.get(), - metadata.at(llm::kUseKVCache), + runner->get_metadata().at(llm::kUseKVCache), std::move(eos_ids), - stats.get()); + stats_ptr); - std::unique_ptr runner = std::make_unique( - std::move(metadata), - std::move(tokenizer), - std::move(module), - std::move(text_decoder_runner), - std::move(text_prefiller), - std::move(text_token_generator), - std::move(stats), - std::move(io_manager), - runner_config.model_version, - runner_config.temperature); + runner->set_io_manager(std::move(io_manager)); + runner->set_text_decoder_runner(std::move(text_decoder_runner)); + runner->set_text_prefiller(std::move(text_prefiller)); + runner->set_text_token_generator(std::move(text_token_generator)); runner->load(); diff --git a/examples/samsung/executor_runner/llm/runner/runner.h b/examples/samsung/executor_runner/llm/runner/runner.h index 9da6e923472..986d0b9850f 100644 --- a/examples/samsung/executor_runner/llm/runner/runner.h +++ b/examples/samsung/executor_runner/llm/runner/runner.h @@ -1,5 +1,5 @@ /* - * Copyright (c) 2025 Samsung Electronics Co. LTD + * Copyright (c) 2026 Samsung Electronics Co. LTD * All rights reserved. * * This source code is licensed under the BSD-style license found in the @@ -16,7 +16,7 @@ #include #include -#include +#include #include #include @@ -26,73 +26,8 @@ namespace llm = ::executorch::extension::llm; namespace runtime = ::executorch::runtime; namespace extension = ::executorch::extension; -/** - * @brief Holds path-related configuration for creating an LLM runner. - */ -struct RunnerPaths { - std::string model_path; - std::string tokenizer_path; - std::optional pos_embedding_path = std::nullopt; - std::optional data_path = std::nullopt; -}; - -/** - * @brief Holds non-path configuration for creating an LLM runner. - */ -struct RunnerConfig { - std::string config_path; - bool is_decode_only = false; - float temperature = -1.0f; - - // Strategies - DecoderModelVersion model_version = DecoderModelVersion::kLlama3; - EmbeddingStrategy embedding_strategy = EmbeddingStrategy::kGlobalOnly; - MaskingStrategy masking_strategy = MaskingStrategy::kCausal; - - // Parsed from JSON - int64_t hidden_dim = 0; - int64_t n_heads = 0; - int64_t n_kv_heads = 0; - int64_t n_layers = 0; - int64_t head_dim = 0; - int64_t vocab_size = 0; - int64_t max_seq_len = 128; - int64_t max_context_len = 128; - int64_t prefill_seq_len = 128; - int64_t decode_seq_len = 1; - int64_t attn_mask_rows = 128; - int64_t attn_mask_cols = 1024; - std::optional sliding_window; - std::string tokenizer_type; - - void load_from_json(); -}; - class LLMRunner : public llm::IRunner { public: - /** - * @brief Constructor for LLMRunner with dependency injection - * - * Creates a LLMRunner instance with all required components for text - * generation. - * - * @param metadata Key-value pairs containing model metadata (e.g., - * vocab_size, context_length) - * @param tokenizer Tokenizer for converting between text and token IDs - * @param module The underlying model module that performs inference - * @param text_decoder_runner Component responsible for running the decoder - * part of the model - * @param text_prefiller Component for handling the prefill phase of text - * generation - * @param text_token_generator Component for generating tokens during the - * decode phase - * @param stats Statistics tracking object for performance monitoring - * @param io_manager IOManager manages input/output operations for LLM - * inference. - * @param model_version Decoder model family selecting the prompt template - * @param temperature Temperature parameter for controlling randomness in - * generation (deprecated). Please use GenerationConfig.temperature instead. - */ explicit LLMRunner( std::unordered_map metadata, std::unique_ptr<::tokenizers::Tokenizer> tokenizer, @@ -102,63 +37,21 @@ class LLMRunner : public llm::IRunner { std::unique_ptr text_token_generator, std::unique_ptr stats, std::unique_ptr io_manager, + RunnerConfig config, DecoderModelVersion model_version, float temperature = -1.0f); - /** - * @brief Checks if the model is loaded and ready for inference - * - * @return bool True if the model is loaded, false otherwise - */ bool is_loaded() const; - /** - * @brief Loads the model and prepares it for inference - * - * This method initializes all components and prepares the model for text - * generation. - * - * @return ::executorch::runtime::Error Success or error status - */ + runtime::Error load(); - /** - * @brief Generates text based on the provided prompt - * - * This method performs text generation using the loaded model. It processes - * the input prompt, runs the model in prefill and decode phases until max - * tokens to generate is reached or eos token is generated, then returns - * generated text and perf stats through callbacks. - * - * @param prompt The input text to generate from - * @param config Configuration parameters for text generation (e.g., - * max_new_tokens, temperature) - * @param token_callback Function called for each generated token with the - * decoded text - * @param stats_callback Function called with performance statistics - * @return ::executorch::runtime::Error Success or error status - */ + runtime::Error generate( const std::string& prompt, const llm::GenerationConfig& config, std::function token_callback = {}, std::function stats_callback = {}); - /** - * @brief Generates text based on the provided prompt and start position - * - * This method performs text generation using the loaded model. It processes - * the input prompt, runs the model in prefill and decode phases using the - * start position until max tokens to generate is reached or eos token is - * generated, then returns generated text and perf stats through callbacks. - * - * @param prompt The input text to generate from - * @param start_pos The starting position in KV cache of the input - * @param config Configuration parameters for text generation (e.g., - * max_new_tokens, temperature) - * @param token_callback Function called for each generated token with the - * decoded text - * @param stats_callback Function called with performance statistics - * @return ::executorch::runtime::Error Success or error status - */ + // TODO (brucekk.kim): consider removing unncessary args runtime::Error generate_from_pos( const std::string& prompt, int64_t start_pos, @@ -166,74 +59,70 @@ class LLMRunner : public llm::IRunner { std::function token_callback = {}, std::function stats_callback = {}); - /** - * @brief Warms up the model with a sample prompt - * - * This method runs a complete generation cycle without returning results, - * which helps initialize internal caches and optimize subsequent inferences. - * - * @param prompt The sample prompt to use for warmup - * @param max_new_tokens Maximum number of tokens to generate during warmup - * @return ::executorch::runtime::Error Success or error status - */ runtime::Error warmup(const std::string& prompt, int32_t max_new_tokens); - /** - * @brief Stops the ongoing text generation process - * - * This method signals the generator to stop producing new tokens and - * terminate the current generation process. - */ + void stop(); - /** - * @brief Formats the prompt based on the model version. - * - * @param prompt The input text - * @param system_prompt The system prompt (optional) - * @param version The version of the decoder model - * @return Formatted prompt string - */ + const RunnerConfig& get_config() const { + return config_; + } + extension::Module* get_module() const { + return module_.get(); + } + ::tokenizers::Tokenizer* get_tokenizer() const { + return tokenizer_.get(); + } + const std::unordered_map& get_metadata() const { + return metadata_; + } + + void set_io_manager(std::unique_ptr io_manager) { + io_manager_ = std::move(io_manager); + } + void set_text_decoder_runner(std::unique_ptr runner) { + text_decoder_runner_ = std::move(runner); + } + void set_text_prefiller(std::unique_ptr prefiller) { + text_prefiller_ = std::move(prefiller); + } + void set_text_token_generator( + std::unique_ptr generator) { + text_token_generator_ = std::move(generator); + } + + // TODO (brucekk.kim): should it be static? static std::string get_formatted_prompt( const std::string& prompt, const std::string& system_prompt, DecoderModelVersion version); - /** - * @brief Force remove prefilled tokens and reset KV cache start position - * - * This mehtod removes the prefilled tokens from the KV cache and resets the - * start position to 0. - */ void reset() {}; private: bool shouldStop_{false}; - // Components + RunnerConfig config_; std::unique_ptr<::tokenizers::Tokenizer> tokenizer_; std::unordered_map metadata_; - std::unique_ptr - module_; // Manage module's lifecycle, make sure it outlives - // text_decoder_runner_llm::. - std::unique_ptr - text_decoder_runner_; // Manage text_decoder_runner_'s lifecycle, make - // sure it outlives text_prefiller_ & - // text_token_generator_. + std::unique_ptr module_; + std::unique_ptr text_decoder_runner_; std::unique_ptr text_prefiller_; std::unique_ptr text_token_generator_; std::unique_ptr io_manager_; - // Stats std::unique_ptr stats_; - // temperature. - // Deprecated, we should rely on the temperature in GenerationConfig instead. float temperature_ = -1.0f; DecoderModelVersion model_version_; }; std::unique_ptr create_llm_runner( - const RunnerPaths& paths, - RunnerConfig config); + const RunnerPaths& runner_paths, + RunnerConfig runner_config); + +std::unique_ptr create_llm_runner( + const RunnerPaths& runner_paths, + RunnerConfig runner_config, + std::unique_ptr<::tokenizers::Tokenizer> tokenizer); } // namespace example From 8d12eb46e2e14611fd81ebe1d0e9264096ef6bdd Mon Sep 17 00:00:00 2001 From: Jiseong Oh Date: Tue, 13 Oct 2026 19:39:29 +0900 Subject: [PATCH 12/17] Add Gemma 4 text and vision support to the LLM runner Fix the Gemma 4 text decoder I/O, the partial-RoPE table math and the chat template separator, precompute the RoPE position embeddings, and fix a double start_pos increment in multi-chunk prefill. A vision encoder adds multimodal prefill. The run script gains --quiet, the runner reprints the full answer at the end behind a verbose toggle, and the runner is exported as a library other sample apps can link. Signed-off-by: Jiseong Oh Co-authored-by: MJ Kim --- .../generate_pos_embedding_bin.py | 84 + examples/samsung/CMakeLists.txt | 35 + .../enn_executor_llm_runner.cpp | 22 +- examples/samsung/executor_runner/ennllm.cpp | 131 + examples/samsung/executor_runner/ennllm.h | 110 + .../executor_runner/llm/runner/CMakeLists.txt | 3 + .../executor_runner/llm/runner/config.cpp | 17 + .../executor_runner/llm/runner/config.h | 61 +- .../executor_runner/llm/runner/image_utils.h | 218 + .../executor_runner/llm/runner/io_manager.cpp | 233 +- .../executor_runner/llm/runner/io_manager.h | 72 +- .../executor_runner/llm/runner/prefiller.cpp | 3 - .../llm/runner/rotary_embedding.cpp | 19 +- .../executor_runner/llm/runner/runner.cpp | 136 +- .../executor_runner/llm/runner/runner.h | 14 +- .../llm/runner/stb_image_impl.cpp | 15 + .../llm/runner/vision_encoder.cpp | 221 + .../llm/runner/vision_encoder.h | 73 + third-party/stb/stb_image.h | 7988 +++++++++++++++++ third-party/stb/stb_image_resize.h | 2634 ++++++ 20 files changed, 11990 insertions(+), 99 deletions(-) create mode 100644 examples/models/gemma4/text_decoder/generate_pos_embedding_bin.py create mode 100644 examples/samsung/executor_runner/ennllm.cpp create mode 100644 examples/samsung/executor_runner/ennllm.h create mode 100644 examples/samsung/executor_runner/llm/runner/image_utils.h create mode 100644 examples/samsung/executor_runner/llm/runner/stb_image_impl.cpp create mode 100644 examples/samsung/executor_runner/llm/runner/vision_encoder.cpp create mode 100644 examples/samsung/executor_runner/llm/runner/vision_encoder.h create mode 100644 third-party/stb/stb_image.h create mode 100644 third-party/stb/stb_image_resize.h diff --git a/examples/models/gemma4/text_decoder/generate_pos_embedding_bin.py b/examples/models/gemma4/text_decoder/generate_pos_embedding_bin.py new file mode 100644 index 00000000000..2138c80a30d --- /dev/null +++ b/examples/models/gemma4/text_decoder/generate_pos_embedding_bin.py @@ -0,0 +1,84 @@ +#!/usr/bin/env python3 +"""Precompute the Gemma4 RoPE position-embedding .bin for the Exynos LLM runner. + +Reproduces the RoPE convention from gemma4_attention.py: inv_freq denominator is +head_dim, full-attention (global) layers use partial rotary while sliding (local) +layers use full rotary, non-rotary dims zero-padded so split-half rotate_half +pairs correctly. Emits the layout RotaryEmbedding::load_from_file expects: + + [global_cos | global_sin | local_cos | local_sin] + +each table row-major [max_seq_len, head_dim] float32 (global head_dim 512, local +head_dim 256). +""" + +import argparse + +import torch +from gemma4_config import Gemma4Config + + +def rope_table( + max_seq_len: int, + head_dim: int, + base: float, + partial_rotary_factor: float, + full_rope: bool, +): + rotary_dim = head_dim if full_rope else int(head_dim * partial_rotary_factor) + rope_angles = rotary_dim // 2 + inv_freq_rotated = 1.0 / ( + base ** (torch.arange(0, rotary_dim, 2).float() / head_dim) + ) + nope_angles = head_dim // 2 - rope_angles + if nope_angles > 0: + inv_freq = torch.cat([inv_freq_rotated, torch.zeros(nope_angles)]) + else: + inv_freq = inv_freq_rotated + + pos = torch.arange(max_seq_len).float() + freqs = torch.outer(pos, inv_freq) + emb = torch.cat((freqs, freqs), dim=-1) + return torch.cos(emb), torch.sin(emb) + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--config", required=True, help="HF config.json path") + ap.add_argument("--max-seq-len", type=int, default=4096) + ap.add_argument("--out", required=True, help="output pos_embedding.bin path") + args = ap.parse_args() + + cfg = Gemma4Config.from_json(args.config) + + g_cos, g_sin = rope_table( + args.max_seq_len, + cfg.global_head_dim, + cfg.rope_theta, + cfg.partial_rotary_factor, + full_rope=False, + ) + l_cos, l_sin = rope_table( + args.max_seq_len, + cfg.head_dim, + cfg.rope_local_base_freq, + cfg.partial_rotary_factor, + full_rope=True, + ) + + with open(args.out, "wb") as f: + for t in (g_cos, g_sin, l_cos, l_sin): + f.write(t.contiguous().to(torch.float32).numpy().tobytes()) + + nbytes = args.max_seq_len * (cfg.global_head_dim + cfg.head_dim) * 2 * 4 + print( + f"wrote {args.out}: {nbytes} bytes " + f"(max_seq_len={args.max_seq_len}, global_head_dim={cfg.global_head_dim}, " + f"head_dim={cfg.head_dim}, rope_theta={cfg.rope_theta}, " + f"rope_local_base_freq={cfg.rope_local_base_freq}, " + f"partial_rotary_factor={cfg.partial_rotary_factor})" + ) + + +if __name__ == "__main__": + main() diff --git a/examples/samsung/CMakeLists.txt b/examples/samsung/CMakeLists.txt index c8b4d0992ae..1138327e2e6 100644 --- a/examples/samsung/CMakeLists.txt +++ b/examples/samsung/CMakeLists.txt @@ -210,3 +210,38 @@ target_link_libraries( enn_logging -Wl,--no-whole-archive ) + +# Public shared library exposing the C API in ennllm.h. Bundles all +# dependencies so only the .so + header need to be shipped to third parties. +add_library(ennllm SHARED ${CMAKE_CURRENT_LIST_DIR}/executor_runner/ennllm.cpp) + +target_compile_definitions(ennllm PRIVATE EXECUTORCH_VERSION=2) +target_include_directories( + ennllm + PUBLIC ${CMAKE_CURRENT_LIST_DIR}/executor_runner +) + +target_link_libraries( + ennllm + PUBLIC enn_llm_runner + PRIVATE log + gflags + android + lib_portable_ops + ${enn_executor_llm_runner_deps} + ${xnnpack_backend_libs} + sentencepiece + re2 + lib_tokenizer + -Wl,--whole-archive + lib_executorch + lib_custom_op + lib_enn_backend + enn_logging + -Wl,--no-whole-archive + -Wl,--exclude-libs,ALL +) + +set_target_properties( + ennllm PROPERTIES CXX_VISIBILITY_PRESET hidden VISIBILITY_INLINES_HIDDEN ON +) diff --git a/examples/samsung/executor_runner/enn_executor_llm_runner.cpp b/examples/samsung/executor_runner/enn_executor_llm_runner.cpp index eb24071d3e6..e527ace4eab 100644 --- a/examples/samsung/executor_runner/enn_executor_llm_runner.cpp +++ b/examples/samsung/executor_runner/enn_executor_llm_runner.cpp @@ -30,6 +30,15 @@ DEFINE_string(tokenizer_path, "tokenizer.bin", "Tokenizer stuff."); DEFINE_string(config_path, "", "Path to the model configuration JSON file."); +DEFINE_string(vision_model_path, "", "Path to the Gemma 4 vision encoder PTE."); + +DEFINE_string(image_path, "", "Path to an input image (Gemma 4 multimodal)."); + +DEFINE_string( + prompt_before_image, + "", + "Text placed before the image in the user turn (Gemma 4 multimodal)."); + DEFINE_string(prompt, "The answer to the ultimate question is", "Prompt."); DEFINE_double( @@ -77,6 +86,10 @@ int32_t main(int32_t argc, char** argv) { } const char* config_path = FLAGS_config_path.c_str(); + std::optional vision_model_path = std::nullopt; + if (!FLAGS_vision_model_path.empty()) { + vision_model_path = FLAGS_vision_model_path; + } const char* prompt = FLAGS_prompt.c_str(); @@ -112,9 +125,14 @@ int32_t main(int32_t argc, char** argv) { .model_path = model_path, .tokenizer_path = tokenizer_path, .pos_embedding_path = pos_embedding_path, - .data_path = data_path}; + .data_path = data_path, + .vision_model_path = vision_model_path}; example::RunnerConfig runner_config{ - .config_path = config_path, .is_decode_only = is_decode_only}; + .config_path = config_path, + .is_decode_only = is_decode_only, + .verbose = true, + .image_path = FLAGS_image_path, + .prompt_before_image = FLAGS_prompt_before_image}; std::unique_ptr runner = example::create_llm_runner(runner_paths, std::move(runner_config)); ET_CHECK_MSG(runner != nullptr, "Failed to create LLM runner"); diff --git a/examples/samsung/executor_runner/ennllm.cpp b/examples/samsung/executor_runner/ennllm.cpp new file mode 100644 index 00000000000..426fd2be14e --- /dev/null +++ b/examples/samsung/executor_runner/ennllm.cpp @@ -0,0 +1,131 @@ +/* + * Copyright (c) 2026 Samsung Electronics Co. LTD + * All rights reserved. + * + * C API implementation over the Exynos LLM runner. Keeps the model(s) loaded + * across generate() calls so consecutive inferences do not re-prepare. + */ + +#include + +#include + +#include + +#include +#include +#include +#include + +namespace llm = ::executorch::extension::llm; + +struct enn_runner { + std::unique_ptr impl; +}; + +namespace { + +std::optional opt_or_null(const char* s) { + if (s == nullptr || *s == '\0') { + return std::nullopt; + } + return std::string(s); +} + +std::string str_or_empty(const char* s) { + return s == nullptr ? std::string() : std::string(s); +} + +} // namespace + +extern "C" { + +enn_runner* enn_llm_create( + const char* model_path, + const char* data_path, + const char* tokenizer_path, + const char* pos_embedding_path, + const char* config_path, + const char* vision_model_path) { + if (model_path == nullptr || tokenizer_path == nullptr || + config_path == nullptr) { + ET_LOG(Error, "enn_llm_create: model/tokenizer/config path required"); + return nullptr; + } + + example::RunnerPaths paths{ + .model_path = model_path, + .tokenizer_path = tokenizer_path, + .pos_embedding_path = opt_or_null(pos_embedding_path), + .data_path = opt_or_null(data_path), + .vision_model_path = opt_or_null(vision_model_path)}; + + example::RunnerConfig cfg{.config_path = config_path}; + + auto runner = example::create_llm_runner(paths, std::move(cfg)); + if (runner == nullptr) { + ET_LOG(Error, "enn_llm_create: failed to create LLM runner"); + return nullptr; + } + + return new enn_runner{std::move(runner)}; +} + +void enn_llm_destroy(enn_runner* runner) { + delete runner; +} + +int32_t enn_llm_generate( + enn_runner* runner, + const char* prompt, + const char* image_path, + const char* prompt_before_image, + int32_t seq_len, + float temperature, + int32_t num_bos, + int32_t num_eos, + int32_t verbose, + enn_token_cb on_token, + void* userdata) { + if (runner == nullptr || prompt == nullptr || *prompt == '\0') { + ET_LOG(Error, "enn_llm_generate: invalid argument"); + return ENN_ERR_INVALID_ARG; + } + + auto* lrunner = dynamic_cast(runner->impl.get()); + if (lrunner == nullptr) { + ET_LOG(Error, "enn_llm_generate: unexpected runner type"); + return ENN_ERR_INVALID_ARG; + } + + // Start each generation from a clean KV-cache. + lrunner->reset(); + + // Image inputs are threaded through the runner's config (read by generate). + lrunner->get_config().image_path = str_or_empty(image_path); + lrunner->get_config().prompt_before_image = str_or_empty(prompt_before_image); + lrunner->get_config().verbose = verbose != 0; + + llm::GenerationConfig gen_cfg{ + .seq_len = seq_len >= 0 ? seq_len : -1, + .temperature = temperature, + .num_bos = num_bos, + .num_eos = num_eos}; + + auto err = runner->impl->generate( + prompt, + gen_cfg, + [on_token, userdata](const std::string& piece) { + if (on_token != nullptr) { + on_token(piece.data(), static_cast(piece.size()), userdata); + } + }, + {}); + if (err != executorch::runtime::Error::Ok) { + ET_LOG(Error, "enn_llm_generate: generation failed"); + return ENN_ERR_GENERATE; + } + return ENN_OK; +} + +} // extern "C" diff --git a/examples/samsung/executor_runner/ennllm.h b/examples/samsung/executor_runner/ennllm.h new file mode 100644 index 00000000000..d62b3671b05 --- /dev/null +++ b/examples/samsung/executor_runner/ennllm.h @@ -0,0 +1,110 @@ +/* + * Copyright (c) 2026 Samsung Electronics Co. LTD + * All rights reserved. + * + * Public C API for the Exynos LLM runner. + * + * Only these symbols are exported from libennllm.so. The underlying + * implementation (model loading, ENN backend, tokenizer, KV-cache) is kept + * internal to the library. + * + * Model preparation happens once in enn_llm_create. enn_llm_generate can then + * be called repeatedly (with or without an image) on the same handle without + * reloading or re-preparing the models. + */ + +#ifndef ENNLLM_H_ +#define ENNLLM_H_ + +#include + +#ifdef __cplusplus +extern "C" { +#endif + +#if defined(__GNUC__) || defined(__clang__) +#define ENN_API __attribute__((visibility("default"))) +#else +#define ENN_API +#endif + +/** Values returned by the API functions. 0 = OK, non-zero = error. */ +enum { + ENN_OK = 0, + ENN_ERR_INVALID_ARG = 1, + ENN_ERR_LOAD = 2, + ENN_ERR_GENERATE = 3, +}; + +/** + * Callback invoked with each decoded text piece as tokens are generated. + * text points to len bytes of UTF-8 text (may not be NUL-terminated); copy it + * if you keep it beyond the call. Called on the same thread that calls + * enn_llm_generate; must not block. + */ +typedef void (*enn_token_cb)(const char* text, int32_t len, void* userdata); + +typedef struct enn_runner enn_runner; + +/** + * Load and prepare the model(s). Does all one-time work: decoder .pte/.ptd, + * tokenizer, positional embedding, optional vision encoder, and KV-cache + * allocation. Repeated calls are not allowed on the same handle. + * + * @param model_path Decoder .pte file (required). + * @param data_path Decoder .ptd weights, or NULL if none. + * @param tokenizer_path Tokenizer model file (required). + * @param pos_embedding_path Precomputed rotary embeddings, or NULL. + * @param config_path Runner config file (required). + * @param vision_model_path Vision encoder .pte, or NULL => text-only runner. + * @return A handle on success, NULL on failure. + */ +ENN_API enn_runner* enn_llm_create( + const char* model_path, + const char* data_path, + const char* tokenizer_path, + const char* pos_embedding_path, + const char* config_path, + const char* vision_model_path); + +/** Destroy the runner and release all resources. Safe to call with NULL. */ +ENN_API void enn_llm_destroy(enn_runner* runner); + +/** + * Generate a response for prompt (optionally with an image). The models stay + * prepared across calls; KV-cache is reset before each generation so every + * prompt starts fresh. + * + * @param runner Handle from enn_llm_create. + * @param prompt User prompt text (required, non-empty). + * @param image_path Image file for this turn, or NULL => text-only. + * @param prompt_before_image Text placed before the image in the user turn + * (Gemma 4 multimodal); ignored when image_path is NULL. + * @param seq_len Total token budget (prompt + generated), capped by the model's + * max context length; -1 => use the model's max context length. + * @param temperature Sampling temperature. + * @param num_bos BOS tokens to prepend to the prompt. + * @param num_eos EOS tokens to append to the prompt. + * @param verbose Non-zero also streams tokens to stdout. + * @param on_token Per-piece callback, or NULL to discard output. + * @param userdata Opaque pointer passed through to on_token. + * @return ENN_OK on success, an ENN_ERR_* code otherwise. + */ +ENN_API int32_t enn_llm_generate( + enn_runner* runner, + const char* prompt, + const char* image_path, + const char* prompt_before_image, + int32_t seq_len, + float temperature, + int32_t num_bos, + int32_t num_eos, + int32_t verbose, + enn_token_cb on_token, + void* userdata); + +#ifdef __cplusplus +} +#endif + +#endif // ENNLLM_H_ diff --git a/examples/samsung/executor_runner/llm/runner/CMakeLists.txt b/examples/samsung/executor_runner/llm/runner/CMakeLists.txt index 63db5a0a798..87c02561dfc 100644 --- a/examples/samsung/executor_runner/llm/runner/CMakeLists.txt +++ b/examples/samsung/executor_runner/llm/runner/CMakeLists.txt @@ -26,6 +26,8 @@ list( ${CMAKE_CURRENT_LIST_DIR}/prefiller.cpp ${CMAKE_CURRENT_LIST_DIR}/io_manager.cpp ${CMAKE_CURRENT_LIST_DIR}/rotary_embedding.cpp + ${CMAKE_CURRENT_LIST_DIR}/vision_encoder.cpp + ${CMAKE_CURRENT_LIST_DIR}/stb_image_impl.cpp ${CMAKE_CURRENT_SOURCE_DIR}/../../../../../examples/models/llama/tokenizer/llama_tiktoken.cpp ) @@ -87,5 +89,6 @@ target_include_directories( ${CMAKE_CURRENT_SOURCE_DIR}/../../../../../extension/llm/tokenizers/third-party/json/include ${CMAKE_CURRENT_SOURCE_DIR}/../../../../../examples/models/llama/tokenizer ${CMAKE_CURRENT_SOURCE_DIR}/../../../../../examples/samsung/executor_runner/llm/runner + ${CMAKE_CURRENT_SOURCE_DIR}/../../../../../third-party/stb ${CMAKE_CURRENT_SOURCE_DIR}/../../../../../extension/llm/runner ) diff --git a/examples/samsung/executor_runner/llm/runner/config.cpp b/examples/samsung/executor_runner/llm/runner/config.cpp index df99b3bf26c..7f51d1d1e43 100644 --- a/examples/samsung/executor_runner/llm/runner/config.cpp +++ b/examples/samsung/executor_runner/llm/runner/config.cpp @@ -139,6 +139,23 @@ void RunnerConfig::load_from_json() { src["sliding_window"].is_number_integer()) { args->text_config.sliding_window = src["sliding_window"].get(); } + if (src.contains("num_kv_shared_layers") && + src["num_kv_shared_layers"].is_number_integer()) { + args->num_kv_shared_layers = src["num_kv_shared_layers"].get(); + } + if (src.contains("sliding_window_pattern") && + src["sliding_window_pattern"].is_number_integer()) { + args->sliding_window_pattern = + src["sliding_window_pattern"].get(); + } + if (src.contains("layer_types") && src["layer_types"].is_array()) { + args->is_full_attention.clear(); + for (const auto& layer_type : src["layer_types"]) { + args->is_full_attention.push_back( + layer_type.is_string() && + layer_type.get() == "full_attention"); + } + } model_args = std::move(args); break; } diff --git a/examples/samsung/executor_runner/llm/runner/config.h b/examples/samsung/executor_runner/llm/runner/config.h index 0246a6bc0ec..98be7b42e0b 100644 --- a/examples/samsung/executor_runner/llm/runner/config.h +++ b/examples/samsung/executor_runner/llm/runner/config.h @@ -15,6 +15,7 @@ #include #include #include +#include namespace example { @@ -29,6 +30,15 @@ inline constexpr auto default_method_name = "forward"; inline constexpr auto default_method_name_prefill = "prefill_forward"; inline constexpr auto default_method_name_decode = "decode_forward"; +// Gemma 4 multimodal special tokens (matches the model's tokenizer). +inline constexpr int64_t kGemma4PadId = 0; +inline constexpr int64_t kGemma4BosId = 2; +inline constexpr int64_t kGemma4TurnStartId = 105; // <|turn| +inline constexpr int64_t kGemma4TurnEndId = 106; // +inline constexpr int64_t kGemma4ImageTokenId = 258880; +inline constexpr int64_t kGemma4BoiTokenId = 255999; // <|boi|> +inline constexpr int64_t kGemma4EoiTokenId = 258882; // <|eoi|> + enum class DecoderModelVersion { kLlama3, kGemma3, @@ -185,6 +195,7 @@ class RunnerPaths { std::string tokenizer_path; std::optional pos_embedding_path = std::nullopt; std::optional data_path = std::nullopt; + std::optional vision_model_path = std::nullopt; }; class BaseModelArgs { @@ -233,6 +244,23 @@ class BaseModelArgs { virtual std::optional get_sliding_window() const { return std::nullopt; } + // Number of trailing layers that reuse (share) KV cache from earlier layers + // and therefore do not maintain their own externally-managed cache. + virtual int64_t get_num_kv_shared_layers() const { + return 0; + } + // Per-layer head dimension. Defaults to a uniform head_dim; models with + // mixed local/global attention (e.g. Gemma 4) override this. + virtual int64_t get_head_dim_for_layer(int64_t layer_idx) const { + (void)layer_idx; + return get_head_dim(); + } + // Layers that actually produce/consume an externally-managed KV cache + // (the leading self-decoder layers). Shared layers are handled inside the + // delegate and are not cached by the runner. + int64_t get_num_cache_layers() const { + return get_n_layers() - get_num_kv_shared_layers(); + } }; class Llama3Args : public BaseModelArgs { @@ -397,6 +425,14 @@ class Gemma4Args : public BaseModelArgs { int64_t attn_mask_rows = 128; int64_t attn_mask_cols = 1024; + // YOCO KV sharing: the last num_kv_shared_layers reuse earlier layers' KV. + int64_t num_kv_shared_layers = 0; + // Attention pattern period (every pattern-th layer is full_attention). + int64_t sliding_window_pattern = 5; + // Parsed from "layer_types"; true where the layer uses full (global) + // attention. Empty falls back to the sliding_window_pattern derivation. + std::vector is_full_attention; + int64_t get_hidden_dim() const override { return text_config.hidden_size; } @@ -427,11 +463,25 @@ class Gemma4Args : public BaseModelArgs { std::optional get_sliding_window() const override { return text_config.sliding_window; } + // Attention mask spans [prefill_seq_len, max_seq_len] (rows x cols). int64_t get_attn_mask_rows() const override { - return attn_mask_rows; + return get_prefill_seq_len(); } int64_t get_attn_mask_cols() const override { - return attn_mask_cols; + return text_config.max_seq_len; + } + int64_t get_num_kv_shared_layers() const override { + return num_kv_shared_layers; + } + int64_t get_head_dim_for_layer(int64_t layer_idx) const override { + bool full; + if (layer_idx >= 0 && + layer_idx < static_cast(is_full_attention.size())) { + full = is_full_attention[layer_idx]; + } else { + full = ((layer_idx + 1) % sliding_window_pattern == 0); + } + return full ? text_config.global_head_dim : text_config.head_dim; } }; @@ -441,6 +491,13 @@ class RunnerConfig { bool is_decode_only = false; float temperature = -1.0f; + // Stream generated tokens to stdout (in addition to the token callback). + bool verbose = false; + + // Vision inputs (Gemma 4 multimodal). Empty image_path => text-only. + std::string image_path; + std::string prompt_before_image; + // Strategies EmbeddingStrategy embedding_strategy = EmbeddingStrategy::kGlobalOnly; MaskingStrategy masking_strategy = MaskingStrategy::kCausal; diff --git a/examples/samsung/executor_runner/llm/runner/image_utils.h b/examples/samsung/executor_runner/llm/runner/image_utils.h new file mode 100644 index 00000000000..849df80b8a2 --- /dev/null +++ b/examples/samsung/executor_runner/llm/runner/image_utils.h @@ -0,0 +1,218 @@ +/* + * Copyright (c) 2026 Samsung Electronics Co. LTD + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +#include +#include + +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +namespace example { + +using ::executorch::extension::TensorPtr; + +// Vision encoder constants (Gemma 4). Matches HF's image processor. +inline constexpr int32_t kVisionPatchSize = 16; +inline constexpr int32_t kVisionPatchDim = + 3 * kVisionPatchSize * kVisionPatchSize; // 768 +inline constexpr int32_t kVisionPoolingKernel = 3; +inline constexpr int32_t kVisionCellSize = + kVisionPatchSize * kVisionPoolingKernel; // 48 + +// The shipped gemma4_vision.pte accepts [1, 1260, 768] pixel_values and +// [1, 1260, 2] pixel_position_ids (STATIC shapes). 1260 = 140 pooled +// tokens * 9 (pooling kernel squared). +inline constexpr int32_t kVisionMaxSoftTokens = 140; +inline constexpr int32_t kVisionMaxPatches = + kVisionMaxSoftTokens * kVisionPoolingKernel * kVisionPoolingKernel; // 1260 +inline constexpr int32_t kVisionPooledTokens = + kVisionMaxPatches / (kVisionPoolingKernel * kVisionPoolingKernel); // 140 + +struct ImageData { + TensorPtr pixel_values; // (1, num_patches, 768) float32 + TensorPtr pixel_position_ids; // (1, num_patches, 2) int64, padding = -1 + int64_t num_valid_patches; + int32_t original_width; + int32_t original_height; +}; + +/// Load an image file (JPEG/PNG/etc.) into HWC interleaved RGB bytes. +inline std::vector +load_image_rgb(const std::string& path, int32_t* width, int32_t* height) { + int w = 0; + int h = 0; + int channels = 0; + unsigned char* data = stbi_load(path.c_str(), &w, &h, &channels, 3); + if (data == nullptr) { + throw std::runtime_error( + "Failed to load image: " + path + " (" + stbi_failure_reason() + ")"); + } + std::vector rgb(data, data + static_cast(w) * h * 3); + stbi_image_free(data); + *width = w; + *height = h; + return rgb; +} + +/// Compute target dimensions preserving aspect ratio. +/// Both dimensions will be divisible by kVisionCellSize (48). +inline std::pair get_aspect_ratio_preserving_size( + int32_t height, + int32_t width, + int32_t max_patches) { + double target_px = + static_cast(max_patches) * kVisionPatchSize * kVisionPatchSize; + double factor = std::sqrt(target_px / (static_cast(height) * width)); + double ideal_h = factor * height; + double ideal_w = factor * width; + + auto target_h = static_cast(std::floor(ideal_h / kVisionCellSize)) * + kVisionCellSize; + auto target_w = static_cast(std::floor(ideal_w / kVisionCellSize)) * + kVisionCellSize; + + int32_t max_side = + (max_patches / (kVisionPoolingKernel * kVisionPoolingKernel)) * + kVisionCellSize; + if (target_h == 0 && target_w == 0) { + throw std::runtime_error("Image too small for patch size"); + } + if (target_h == 0) { + target_h = kVisionCellSize; + target_w = std::min( + static_cast(std::floor(static_cast(width) / height)) * + kVisionCellSize, + max_side); + } else if (target_w == 0) { + target_w = kVisionCellSize; + target_h = std::min( + static_cast(std::floor(static_cast(height) / width)) * + kVisionCellSize, + max_side); + } + return {target_h, target_w}; +} + +/// Compute number of soft tokens after pooling. +inline int64_t compute_vision_num_tokens(int64_t num_patches) { + return num_patches / (kVisionPoolingKernel * kVisionPoolingKernel); +} + +/// Resize RGB image (HWC uint8) to target dimensions using bicubic +/// (Catmull-Rom) interpolation, matching Python's PIL.Image.BICUBIC. +/// PIL filters the stored sRGB values directly, so stb must not linearize +/// them (STBIR_COLORSPACE_SRGB would gamma-decode before filtering). +inline std::vector resize_rgb( + const uint8_t* src, + int32_t src_w, + int32_t src_h, + int32_t dst_w, + int32_t dst_h) { + std::vector dst(static_cast(dst_w) * dst_h * 3); + stbir_resize_uint8_generic( + src, + src_w, + src_h, + 0, + dst.data(), + dst_w, + dst_h, + 0, + 3, + STBIR_ALPHA_CHANNEL_NONE, + 0, + STBIR_EDGE_CLAMP, + STBIR_FILTER_CATMULLROM, + STBIR_COLORSPACE_LINEAR, + nullptr); + return dst; +} + +/// Patchify an RGB image (HWC uint8, already resized) into pixel_values +/// and position_ids tensors, padded to max_patches. +inline ImageData patchify_rgb_image( + const uint8_t* rgb, + int32_t width, + int32_t height, + int32_t max_soft_tokens) { + int32_t max_patches = + max_soft_tokens * kVisionPoolingKernel * kVisionPoolingKernel; + + // Resize if needed + auto [target_h, target_w] = + get_aspect_ratio_preserving_size(height, width, max_patches); + + std::vector resized; + const uint8_t* img_data = rgb; + if (target_h != height || target_w != width) { + resized = resize_rgb(rgb, width, height, target_w, target_h); + img_data = resized.data(); + height = target_h; + width = target_w; + } + + int32_t h_patches = height / kVisionPatchSize; + int32_t w_patches = width / kVisionPatchSize; + int32_t num_patches = h_patches * w_patches; + + // Allocate padded tensors + auto pixel_values = executorch::extension::zeros( + {1, max_patches, kVisionPatchDim}, executorch::aten::ScalarType::Float); + auto pixel_position_ids = executorch::extension::full( + {1, max_patches, 2}, -1, executorch::aten::ScalarType::Long); + + float* pv_data = pixel_values->mutable_data_ptr(); + int64_t* pos_data = pixel_position_ids->mutable_data_ptr(); + + // Extract patches in row-major order (y then x), matching HF's layout: + // patches[py * w_patches + px] = image[py*16:(py+1)*16, px*16:(px+1)*16] + for (int32_t py = 0; py < h_patches; ++py) { + for (int32_t px = 0; px < w_patches; ++px) { + int32_t patch_idx = py * w_patches + px; + float* patch_out = + pv_data + static_cast(patch_idx) * kVisionPatchDim; + + for (int32_t dy = 0; dy < kVisionPatchSize; ++dy) { + for (int32_t dx = 0; dx < kVisionPatchSize; ++dx) { + int32_t img_y = py * kVisionPatchSize + dy; + int32_t img_x = px * kVisionPatchSize + dx; + int32_t src_idx = (img_y * width + img_x) * 3; + int32_t dst_idx = (dy * kVisionPatchSize + dx) * 3; + for (int32_t c = 0; c < 3; ++c) { + patch_out[dst_idx + c] = img_data[src_idx + c] / 255.0f; + } + } + } + + // Position IDs: (x, y) patch coordinates + pos_data[patch_idx * 2] = px; + pos_data[patch_idx * 2 + 1] = py; + } + } + + return ImageData{ + std::move(pixel_values), + std::move(pixel_position_ids), + num_patches, + width, + height, + }; +} + +} // namespace example diff --git a/examples/samsung/executor_runner/llm/runner/io_manager.cpp b/examples/samsung/executor_runner/llm/runner/io_manager.cpp index 529b1567859..fda72beeff5 100644 --- a/examples/samsung/executor_runner/llm/runner/io_manager.cpp +++ b/examples/samsung/executor_runner/llm/runner/io_manager.cpp @@ -6,6 +6,7 @@ * LICENSE file in the root directory of this source tree. */ +#include #include #include @@ -141,7 +142,7 @@ ET_NODISCARD runtime::Error LlmIOManager::reset( current_pos_ = 0; prefill_done_ = false; - for (size_t l = 0; l < config_.model_args->get_n_layers(); l++) { + for (size_t l = 0; l < config_.model_args->get_num_cache_layers(); l++) { std::fill(prefill_k_input_[l].begin(), prefill_k_input_[l].end(), 0); std::fill(prefill_v_input_[l].begin(), prefill_v_input_[l].end(), 0); std::fill(decode_k_input_[l].begin(), decode_k_input_[l].end(), 0); @@ -187,6 +188,8 @@ LlmIOManager::prepare_prefill( executorch::aten::ScalarType::Long); inputs.emplace_back(input_tokens_tensor_); + append_prelude_inputs(prefill_seq_len, inputs); + const int64_t* start_position = start_pos->const_data_ptr(); current_pos_ = *start_position; pos_ids_.resize(prefill_seq_len); @@ -194,15 +197,6 @@ LlmIOManager::prepare_prefill( pos_ids_[i] = current_pos_ + static_cast(i); } - // TODO (brucekk.kim): - // 1. Add two missing inputs - // inputs[1]: inputs_embeds - // inputs[2]: mm_mask - - // TODO (brucekk.kim) - // 2. Fix/refactor update_attn_mask so that dim changes from [1,1,128,1024] to - // [1,1,128,4096] - // attention tensor update_attn_mask(current_pos_, prompt_token_len); attn_mask_tensor_ = executorch::extension::from_blob( @@ -242,6 +236,7 @@ LlmIOManager::prepare_prefill( prefill_v_input_tensors_.clear(); prefill_k_output_tensors_.clear(); prefill_v_output_tensors_.clear(); + dummy_kv_tensors_.clear(); append_kv_input_tensors( prefill_k_input_, prefill_v_input_, @@ -249,6 +244,7 @@ LlmIOManager::prepare_prefill( prefill_v_input_tensors_, executorch::aten::ScalarType::Float, prefill_cache_len_, + prefill_dummy_kv_, inputs); dump_kv_input("prefill"); @@ -271,6 +267,8 @@ LlmIOManager::prepare_decode( inputs.reserve(get_input_reserve_size()); inputs.emplace_back(input); + append_prelude_inputs(config_.model_args->get_decode_seq_len(), inputs); + attn_mask_tensor_ = executorch::extension::from_blob( attn_mask_.data(), {1, @@ -306,6 +304,7 @@ LlmIOManager::prepare_decode( decode_v_input_tensors_.clear(); decode_k_output_tensors_.clear(); decode_v_output_tensors_.clear(); + dummy_kv_tensors_.clear(); append_kv_input_tensors( decode_k_input_, decode_v_input_, @@ -313,6 +312,7 @@ LlmIOManager::prepare_decode( decode_v_input_tensors_, executorch::aten::ScalarType::Float, decode_cache_len_, + decode_dummy_kv_, inputs); dump_kv_input("decode"); @@ -346,20 +346,32 @@ executorch::runtime::Error LlmIOManager::update_decode( } // TOOD: sizeof(float) dtype hard-codes the buffer dtype to be float. void LlmIOManager::allocate_buffers() { - size_t prefill_in_size = config_.model_args->get_n_kv_heads() * - prefill_cache_len_ * config_.model_args->get_head_dim() * sizeof(float); - size_t decode_in_size = config_.model_args->get_n_kv_heads() * - decode_cache_len_ * config_.model_args->get_head_dim() * sizeof(float); - prefill_k_input_.resize(config_.model_args->get_n_layers()); - prefill_v_input_.resize(config_.model_args->get_n_layers()); - decode_k_input_.resize(config_.model_args->get_n_layers()); - decode_v_input_.resize(config_.model_args->get_n_layers()); - for (size_t l = 0; l < config_.model_args->get_n_layers(); l++) { + const size_t n_kv_heads = config_.model_args->get_n_kv_heads(); + const size_t num_cache_layers = config_.model_args->get_num_cache_layers(); + prefill_k_input_.resize(num_cache_layers); + prefill_v_input_.resize(num_cache_layers); + decode_k_input_.resize(num_cache_layers); + decode_v_input_.resize(num_cache_layers); + for (size_t l = 0; l < num_cache_layers; l++) { + const size_t head_dim = config_.model_args->get_head_dim_for_layer(l); + const size_t prefill_in_size = + n_kv_heads * prefill_cache_len_ * head_dim * sizeof(float); + const size_t decode_in_size = + n_kv_heads * decode_cache_len_ * head_dim * sizeof(float); prefill_k_input_[l].resize(prefill_in_size, 0); prefill_v_input_[l].resize(prefill_in_size, 0); decode_k_input_[l].resize(decode_in_size, 0); decode_v_input_[l].resize(decode_in_size, 0); } + + // Dummy scratch for KV-shared layer slots, sized for the widest head_dim. + const size_t max_head_dim = std::max( + config_.model_args->get_head_dim(), + config_.model_args->get_global_head_dim()); + prefill_dummy_kv_.assign( + n_kv_heads * prefill_cache_len_ * max_head_dim * sizeof(float), 0); + decode_dummy_kv_.assign( + n_kv_heads * decode_cache_len_ * max_head_dim * sizeof(float), 0); attn_mask_.resize( config_.model_args->get_attn_mask_rows() * config_.model_args->get_attn_mask_cols(), @@ -370,26 +382,21 @@ void LlmIOManager::allocate_buffers() { -1.0f); } void LlmIOManager::copy_prefill_to_decode_input() { - size_t copy_size = - prefill_cache_len_ * config_.model_args->get_head_dim() * sizeof(float); - for (size_t l = 0; l < config_.model_args->get_n_layers(); l++) { - for (size_t h = 0; h < config_.model_args->get_n_kv_heads(); h++) { - const char* src = prefill_k_input_[l].data() + - h * prefill_cache_len_ * config_.model_args->get_head_dim() * - sizeof(float); - char* dst = decode_k_input_[l].data() + - h * decode_cache_len_ * config_.model_args->get_head_dim() * - sizeof(float); - std::memcpy(dst, src, copy_size); - } - for (size_t h = 0; h < config_.model_args->get_n_kv_heads(); h++) { - const char* src = prefill_v_input_[l].data() + - h * prefill_cache_len_ * config_.model_args->get_head_dim() * - sizeof(float); - char* dst = decode_v_input_[l].data() + - h * decode_cache_len_ * config_.model_args->get_head_dim() * - sizeof(float); - std::memcpy(dst, src, copy_size); + const size_t n_kv_heads = config_.model_args->get_n_kv_heads(); + for (size_t l = 0; l < config_.model_args->get_num_cache_layers(); l++) { + const size_t head_dim = config_.model_args->get_head_dim_for_layer(l); + const size_t copy_size = prefill_cache_len_ * head_dim * sizeof(float); + for (size_t h = 0; h < n_kv_heads; h++) { + const char* k_src = prefill_k_input_[l].data() + + h * prefill_cache_len_ * head_dim * sizeof(float); + char* k_dst = decode_k_input_[l].data() + + h * decode_cache_len_ * head_dim * sizeof(float); + std::memcpy(k_dst, k_src, copy_size); + const char* v_src = prefill_v_input_[l].data() + + h * prefill_cache_len_ * head_dim * sizeof(float); + char* v_dst = decode_v_input_[l].data() + + h * decode_cache_len_ * head_dim * sizeof(float); + std::memcpy(v_dst, v_src, copy_size); } } } @@ -401,31 +408,47 @@ void LlmIOManager::append_kv_input_tensors( std::vector& v_tensors, executorch::aten::ScalarType kv_scalar_type, size_t cache_len, + std::vector& dummy_buf, std::vector& inputs) const { - for (size_t l = 0; l < config_.model_args->get_n_layers(); l++) { - auto k_in = executorch::extension::from_blob( - k_bufs[l].data(), + const size_t n_layers = config_.model_args->get_n_layers(); + const size_t num_cache_layers = config_.model_args->get_num_cache_layers(); + const auto n_kv_heads = static_cast( + config_.model_args->get_n_kv_heads()); + + // The model declares a KV input slot for every layer, in K[0..L-1] then + // V[0..L-1] order. The leading self-decoder layers carry the real cache; + // the trailing KV-shared (cross-decoder) layers are resolved inside the + // delegate and their slots are never read, so point them at zeroed scratch. + auto make_tensor = [&](std::vector>& bufs, + size_t layer) -> executorch::extension::TensorPtr { + char* data = + (layer < num_cache_layers) ? bufs[layer].data() : dummy_buf.data(); + return executorch::extension::from_blob( + data, {1, - static_cast( - config_.model_args->get_n_kv_heads()), + n_kv_heads, static_cast(cache_len), static_cast( - config_.model_args->get_head_dim())}, + config_.model_args->get_head_dim_for_layer(layer))}, kv_scalar_type); - k_tensors.push_back(k_in); + }; + + for (size_t l = 0; l < n_layers; l++) { + auto k_in = make_tensor(k_bufs, l); + if (l < num_cache_layers) { + k_tensors.push_back(k_in); + } else { + dummy_kv_tensors_.push_back(k_in); + } inputs.emplace_back(k_in); } - for (size_t l = 0; l < config_.model_args->get_n_layers(); l++) { - auto v_in = executorch::extension::from_blob( - v_bufs[l].data(), - {1, - static_cast( - config_.model_args->get_n_kv_heads()), - static_cast(cache_len), - static_cast( - config_.model_args->get_head_dim())}, - kv_scalar_type); - v_tensors.push_back(v_in); + for (size_t l = 0; l < n_layers; l++) { + auto v_in = make_tensor(v_bufs, l); + if (l < num_cache_layers) { + v_tensors.push_back(v_in); + } else { + dummy_kv_tensors_.push_back(v_in); + } inputs.emplace_back(v_in); } } @@ -571,8 +594,8 @@ void LlmIOManager::dump_kv_input(const std::string& phase) { const bool is_prefill = (phase == "prefill"); const auto& k_bufs = is_prefill ? prefill_k_input_ : decode_k_input_; const auto& v_bufs = is_prefill ? prefill_v_input_ : decode_v_input_; - const size_t n_layers = config_.model_args->get_n_layers(); - for (size_t l = 0; l < n_layers; l++) { + const size_t num_cache_layers = config_.model_args->get_num_cache_layers(); + for (size_t l = 0; l < num_cache_layers; l++) { dump_buffer(phase + "_k_input", l, k_bufs[l].data(), k_bufs[l].size()); dump_buffer(phase + "_v_input", l, v_bufs[l].data(), v_bufs[l].size()); } @@ -584,22 +607,26 @@ void LlmIOManager::update_kv_cache( size_t cache_len, size_t new_len, const std::string& phase) { - // model_outputs: [logits, k0..k_{L-1}, v0..v_{L-1}] - for (size_t l = 0; l < config_.model_args->get_n_layers(); l++) { + // model_outputs: [logits, k0..k_{C-1}, v0..v_{C-1}] where C is the number of + // cached (self-decoder) layers; KV-shared layers emit no new cache. + const size_t num_cache_layers = config_.model_args->get_num_cache_layers(); + for (size_t l = 0; l < num_cache_layers; l++) { copy_tensor_to_cache( model_outputs[1 + l].toTensor(), new_len, k_bufs[l].data(), cache_len, - current_pos_); + current_pos_, + config_.model_args->get_head_dim_for_layer(l)); } - for (size_t l = 0; l < config_.model_args->get_n_layers(); l++) { + for (size_t l = 0; l < num_cache_layers; l++) { copy_tensor_to_cache( - model_outputs[1 + config_.model_args->get_n_layers() + l].toTensor(), + model_outputs[1 + num_cache_layers + l].toTensor(), new_len, v_bufs[l].data(), cache_len, - current_pos_); + current_pos_, + config_.model_args->get_head_dim_for_layer(l)); } dump_kv_output(phase, model_outputs); dump_step_++; @@ -610,10 +637,10 @@ void LlmIOManager::dump_kv_output( if (!dump_enabled_) { return; } - for (size_t l = 0; l < config_.model_args->get_n_layers(); l++) { + const size_t num_cache_layers = config_.model_args->get_num_cache_layers(); + for (size_t l = 0; l < num_cache_layers; l++) { const auto& k_out = model_outputs[1 + l].toTensor(); - const auto& v_out = - model_outputs[1 + config_.model_args->get_n_layers() + l].toTensor(); + const auto& v_out = model_outputs[1 + num_cache_layers + l].toTensor(); dump_buffer( phase + "_k_output", l, @@ -700,4 +727,74 @@ executorch::runtime::Error LlmIOManager::prepare_rotary_embedding_inputs( return executorch::runtime::Error::Ok; } +void LlmIOManager::set_vision_embeddings( + std::vector embeds, + int64_t num_tokens) { + vision_embeds_ = std::move(embeds); + num_vision_tokens_ = num_tokens; + vision_embed_idx_ = 0; + ET_LOG( + Info, + "Vision embeddings set: %lld tokens, %zu floats", + static_cast(num_vision_tokens_), + vision_embeds_.size()); +} + +void LlmIOManager::append_prelude_inputs( + size_t seq_len, + std::vector& inputs) { + if (config_.model_version != DecoderModelVersion::kGemma4) { + return; + } + const int64_t hidden_dim = config_.model_args->get_hidden_dim(); + + // inputs_embeds [1, seq_len, hidden_dim]: zeros for text-only inference. The + // decoder embeds token ids internally; this slot only matters at multimodal + // placeholder positions, which text-only prompts never contain. + inputs_embeds_data_.assign(seq_len * hidden_dim, 0.0f); + + // mm_mask [1, seq_len]: zeros select the internal text embeddings everywhere. + mm_mask_data_.assign(seq_len, 0.0f); + + // Inject vision embeddings at image-token positions (Gemma 4 multimodal). + // Only during prefill (seq_len > 1); decode tokens are never image tokens. + // Each position whose token is the image-token id consumes the next packed + // vision embedding and flips its mm_mask bit so the decoder picks + // inputs_embeds at that position. vision_embed_idx_ persists across + // micro-batch prefill chunks. + // The token id is also replaced with pad: the exported text decoder looks up + // per-layer embeddings (PLE) from input_ids directly, and HF / the export + // reference (Gemma4Model_ENN.forward) feed pad ids at multimodal positions. + // input_tokens_tensor_ aliases input_tokens_, so the in-place edit is seen + // by the model. + size_t vision_idx = vision_embed_idx_; + if (!vision_embeds_.empty() && seq_len > 1) { + for (size_t i = 0; i < seq_len; ++i) { + if (static_cast(input_tokens_[i]) == kGemma4ImageTokenId && + vision_idx < static_cast(num_vision_tokens_)) { + std::copy_n( + vision_embeds_.data() + vision_idx * hidden_dim, + hidden_dim, + inputs_embeds_data_.data() + i * hidden_dim); + mm_mask_data_[i] = 1.0f; + input_tokens_[i] = kGemma4PadId; + ++vision_idx; + } + } + } + vision_embed_idx_ = vision_idx; + + inputs_embeds_tensor_ = executorch::extension::from_blob( + inputs_embeds_data_.data(), + {1, + static_cast(seq_len), + static_cast(hidden_dim)}); + inputs.emplace_back(inputs_embeds_tensor_); + + mm_mask_tensor_ = executorch::extension::from_blob( + mm_mask_data_.data(), + {1, static_cast(seq_len)}); + inputs.emplace_back(mm_mask_tensor_); +} + } // namespace example diff --git a/examples/samsung/executor_runner/llm/runner/io_manager.h b/examples/samsung/executor_runner/llm/runner/io_manager.h index 72f55363298..2b410746752 100644 --- a/examples/samsung/executor_runner/llm/runner/io_manager.h +++ b/examples/samsung/executor_runner/llm/runner/io_manager.h @@ -109,6 +109,19 @@ class ET_EXPERIMENTAL LlmIOManager : public llm::IOManager { return update_decode(model_outputs, example::default_method_name_decode); } + /** + * Supply pre-computed vision embeddings for Gemma 4 multimodal inference. + * + * @param embeds Packed valid vision embeddings (num_vision_tokens * hidden, + * contiguous float). The caller must have dropped padding. + * @param num_tokens Number of valid vision tokens (= embeds.size()/hidden). + */ + void set_vision_embeddings(std::vector embeds, int64_t num_tokens); + + size_t num_vision_embeds_consumed() const { + return vision_embed_idx_; + } + const RunnerConfig& get_config() const { return config_; } @@ -191,17 +204,13 @@ class ET_EXPERIMENTAL LlmIOManager : public llm::IOManager { size_t new_len, char* cache, size_t cache_len, - size_t pos) { + size_t pos, + size_t head_dim) { for (size_t h = 0; h < config_.model_args->get_n_kv_heads(); h++) { - const T* src_head = - src + h * new_len * config_.model_args->get_head_dim(); - char* dst_head = cache + - h * cache_len * config_.model_args->get_head_dim() * sizeof(T) + - pos * config_.model_args->get_head_dim() * sizeof(T); - std::memcpy( - dst_head, - src_head, - new_len * config_.model_args->get_head_dim() * sizeof(T)); + const T* src_head = src + h * new_len * head_dim; + char* dst_head = cache + h * cache_len * head_dim * sizeof(T) + + pos * head_dim * sizeof(T); + std::memcpy(dst_head, src_head, new_len * head_dim * sizeof(T)); } } @@ -214,9 +223,10 @@ class ET_EXPERIMENTAL LlmIOManager : public llm::IOManager { size_t new_len, char* cache, size_t cache_len, - size_t pos) { + size_t pos, + size_t head_dim) { const float* data = tensor.const_data_ptr(); - copy_to_cache(data, new_len, cache, cache_len, pos); + copy_to_cache(data, new_len, cache, cache_len, pos, head_dim); } /** @@ -233,8 +243,21 @@ class ET_EXPERIMENTAL LlmIOManager : public llm::IOManager { std::vector& v_tensors, executorch::aten::ScalarType kv_scalar_type, size_t cache_len, + std::vector& dummy_buf, std::vector& inputs) const; + /** + * Append the Gemma 4 text-decoder prelude inputs that follow the token ids: + * inputs_embeds [1, seq_len, hidden_dim] and mm_mask [1, seq_len]. + * + * For text-only inference both are zero-filled: the decoder embeds the token + * ids internally and mm_mask == 0 selects those text embeddings everywhere, + * so inputs_embeds (the multimodal injection path) is ignored. + */ + void append_prelude_inputs( + size_t seq_len, + std::vector& inputs); + /** * Update the attention mask for a causal (lower-triangular) pattern. * For each query row r in [0, seq_len), the mask allows attending to @@ -274,6 +297,31 @@ class ET_EXPERIMENTAL LlmIOManager : public llm::IOManager { executorch::extension::TensorPtr input_tokens_tensor_; std::vector pos_ids_; + // Gemma 4 prelude inputs (zero-filled for text-only inference). + std::vector inputs_embeds_data_; + executorch::extension::TensorPtr inputs_embeds_tensor_; + std::vector mm_mask_data_; + executorch::extension::TensorPtr mm_mask_tensor_; + + // Gemma 4 vision embeddings injected at image-token positions. + // vision_embeds_ holds only the valid (non-padding) embeddings packed + // contiguously (num_vision_tokens * hidden). append_prelude_inputs copies + // them into inputs_embeds_data_ at the image-token slots and marks those + // slots in mm_mask_data_. + std::vector vision_embeds_; + int64_t num_vision_tokens_ = 0; + // Cursor into vision_embeds_ consumed by append_prelude_inputs. Persists + // across micro-batch prefill chunks so image tokens split across chunks map + // to the correct embeddings. Reset on set_vision_embeddings. + size_t vision_embed_idx_ = 0; + + // Zero-filled scratch feeding the KV input slots of KV-shared (cross-decoder) + // layers. These slots are declared model inputs but never read by the graph, + // so a single shared buffer (sized for the widest head_dim) backs them all. + std::vector prefill_dummy_kv_; + std::vector decode_dummy_kv_; + mutable std::vector dummy_kv_tensors_; + // Per-layer input buffers: [1, n_kv_heads, cache_len, head_dim] flattened // Use std::vector to support both float and half (float16) data types std::vector> prefill_k_input_; diff --git a/examples/samsung/executor_runner/llm/runner/prefiller.cpp b/examples/samsung/executor_runner/llm/runner/prefiller.cpp index 07b3e52d50a..57eefb96ef9 100644 --- a/examples/samsung/executor_runner/llm/runner/prefiller.cpp +++ b/examples/samsung/executor_runner/llm/runner/prefiller.cpp @@ -71,9 +71,6 @@ runtime::Result Prefiller::prefill( ET_CHECK_OK_OR_RETURN_ERROR(chunk_result.error()); cur_token = chunk_result.get(); - if (enable_parallel_prefill_) { - start_pos += num_tokens_to_prefill_with; - } num_tokens_to_process += num_tokens_to_prefill_with; } diff --git a/examples/samsung/executor_runner/llm/runner/rotary_embedding.cpp b/examples/samsung/executor_runner/llm/runner/rotary_embedding.cpp index dc8ab5cc0ef..a085e4552ca 100644 --- a/examples/samsung/executor_runner/llm/runner/rotary_embedding.cpp +++ b/examples/samsung/executor_runner/llm/runner/rotary_embedding.cpp @@ -167,13 +167,15 @@ runtime::Error RotaryEmbedding::load_by_precompute_rope_tables() { reinterpret_cast(local_sin_ptr) % alignof(float) == 0 && "local_sin_ is misaligned for float access!"); + // Sliding/local layers use full RoPE; the partial rotary factor applies + // only to the global (full-attention) layers. compute_rope_table( local_cos_ptr, local_sin_ptr, seq_len_, normalized_local_head_dim_, base_freq_local_, - rope_percentage_); + /*rope_pct=*/1.0f); } loaded_ = true; ET_LOG(Info, "Rotary Embedding tables precomputed successfully."); @@ -226,19 +228,24 @@ runtime::Error RotaryEmbedding::compute_rope_table( const size_t rotary_half = rotary_dim / 2; + // Partial RoPE uses the full head_dim as the inv_freq denominator (Gemma + // convention), so for full RoPE this reduces to the standard formula. + const float freq_denom = static_cast(dim); + for (size_t pos = 0; pos < seq_len; ++pos) { float* cos_row = cos_ptr + pos * dim; float* sin_row = sin_ptr + pos * dim; + // Non-rotary dimensions (partial RoPE) pass through unchanged: cos=1, + // sin=0. for (size_t i = 0; i < dim; ++i) { - cos_ptr[i] = 1.0f; - sin_ptr[i] = 0.0f; + cos_row[i] = 1.0f; + sin_row[i] = 0.0f; } for (size_t i = 0; i < rotary_half; ++i) { - // inv_freq_i = base^(-2i / rotary_dim) - const float exponent = - -2.0f * static_cast(i) / static_cast(rotary_dim); + // inv_freq_i = base^(-2i / head_dim) + const float exponent = -2.0f * static_cast(i) / freq_denom; const float inv_freq = std::pow(base, exponent); const float angle = static_cast(pos) * inv_freq; diff --git a/examples/samsung/executor_runner/llm/runner/runner.cpp b/examples/samsung/executor_runner/llm/runner/runner.cpp index 467d9049f0b..3014bea7d64 100644 --- a/examples/samsung/executor_runner/llm/runner/runner.cpp +++ b/examples/samsung/executor_runner/llm/runner/runner.cpp @@ -66,11 +66,11 @@ std::string LLMRunner::get_formatted_prompt( break; case DecoderModelVersion::kGemma4: if (!system_prompt.empty()) { - formatted_prompt.append("<|turn>system "); + formatted_prompt.append("<|turn>system\n"); formatted_prompt.append(system_prompt); formatted_prompt.append("\n"); } - formatted_prompt.append("<|turn>user "); + formatted_prompt.append("<|turn>user\n"); formatted_prompt.append(prompt); formatted_prompt.append("\n"); formatted_prompt.append("<|turn>model\n"); @@ -228,8 +228,8 @@ runtime::Error LLMRunner::generate_from_pos( // Wrap the token_callback with print function std::function wrapped_callback = - [token_callback, config](const std::string& piece) { - if (!config.warming) { + [this, token_callback, config](const std::string& piece) { + if (!config.warming && config_.verbose) { llm::safe_printf(piece.c_str()); fflush(stdout); } @@ -263,6 +263,71 @@ runtime::Error LLMRunner::generate_from_pos( // encode the (string) prompt into tokens sequence std::vector prompt_tokens = encode_res.get(); + + // Gemma 4 multimodal: run the vision encoder and inject the image token + // block. The user turn becomes: + // <|turn>user\n xN + // \nmodel\n + // holds the text AFTER the image; + // (config_.prompt_before_image) is the text placed BEFORE the image. + if (model_version_ == DecoderModelVersion::kGemma4 && + vision_encoder_ != nullptr && !config_.image_path.empty()) { + auto* llm_io = static_cast(io_manager_.get()); + ET_CHECK_MSG(llm_io != nullptr, "IOManager is null"); + + // Vision embeddings are injected only by the multi-token prefill method; + // single-token steps go through decode and would silently drop them. + ET_CHECK_OR_RETURN_ERROR( + !config_.is_decode_only && + config_.model_args->get_prefill_seq_len() > 1, + InvalidArgument, + "Image input requires a prefill method with seq_len > 1"); + + auto encode_err = vision_encoder_->encode(config_.image_path); + ET_CHECK_OK_OR_RETURN_ERROR(encode_err); + int64_t num_vision_tokens = vision_encoder_->num_valid_tokens(); + ET_CHECK_OR_RETURN_ERROR( + vision_encoder_->hidden_size() == + static_cast(config_.model_args->get_hidden_dim()), + InvalidProgram, + "Vision hidden size %lld != decoder hidden dim %lld", + static_cast(vision_encoder_->hidden_size()), + static_cast(config_.model_args->get_hidden_dim())); + + auto enc = [this]( + const std::string& s, std::vector& out) -> bool { + auto r = tokenizer_->encode(s, /*bos=*/0, /*eos=*/0); + if (r.error() != ::tokenizers::Error::Ok) { + ET_LOG(Error, "Failed to encode text segment"); + return false; + } + auto toks = std::move(r.get()); + out.insert(out.end(), toks.begin(), toks.end()); + return true; + }; + + std::vector vision_tokens; + vision_tokens.push_back(kGemma4BosId); + if (!enc("<|turn>user\n", vision_tokens) || + (!config_.prompt_before_image.empty() && + !enc(config_.prompt_before_image, vision_tokens))) { + return runtime::Error::InvalidArgument; + } + vision_tokens.push_back(kGemma4BoiTokenId); + for (int64_t i = 0; i < num_vision_tokens; ++i) { + vision_tokens.push_back(kGemma4ImageTokenId); + } + vision_tokens.push_back(kGemma4EoiTokenId); + if (!enc(prompt, vision_tokens) || + !enc("\n<|turn>model\n", vision_tokens)) { + return runtime::Error::InvalidArgument; + } + prompt_tokens = std::move(vision_tokens); + + llm_io->set_vision_embeddings( + vision_encoder_->vision_embeds(), num_vision_tokens); + } + int num_prompt_tokens = prompt_tokens.size(); // Reduce max_context_len by start_pos @@ -310,6 +375,17 @@ runtime::Error LLMRunner::generate_from_pos( auto prefill_res = text_prefiller_->prefill(prompt_tokens, pos); ET_CHECK_OK_OR_RETURN_ERROR(prefill_res.error()); uint64_t cur_token = prefill_res.get(); + if (model_version_ == DecoderModelVersion::kGemma4 && + vision_encoder_ != nullptr && !config_.image_path.empty()) { + const auto consumed = static_cast(io_manager_.get()) + ->num_vision_embeds_consumed(); + ET_CHECK_OR_RETURN_ERROR( + consumed == static_cast(vision_encoder_->num_valid_tokens()), + InvalidState, + "Only %zu of %lld vision embeddings were injected during prefill", + consumed, + static_cast(vision_encoder_->num_valid_tokens())); + } stats_->first_token_ms = llm::time_in_ms(); stats_->prompt_eval_end_ms = llm::time_in_ms(); @@ -327,6 +403,15 @@ runtime::Error LLMRunner::generate_from_pos( Info, "################ The start of token generation ################\n"); + // Accumulate the full model answer (first token from prefill plus all + // decoded tokens) so it can be reprinted cleanly at the very end, free of + // interleaved logs. Warmed-up runs are skipped (they only measure speed). + std::string full_answer; + if (!config.warming) { + full_answer += + ET_UNWRAP_TOKENIZER(tokenizer_->decode(cur_token, cur_token)); + } + // print the first token from prefill. No prev_token so use cur_token for it. wrapped_callback( ET_UNWRAP_TOKENIZER(tokenizer_->decode(cur_token, cur_token))); @@ -335,12 +420,19 @@ runtime::Error LLMRunner::generate_from_pos( prompt_tokens.push_back(cur_token); // Generate max_new_tokens - 1 because prefill already generated 1 token. + // Forward every decoded piece both to the stream and to the accumulator so + // the complete answer can be reprinted verbatim at the end. + std::function collect_callback = + [&full_answer, &wrapped_callback](const std::string& piece) { + full_answer += piece; + wrapped_callback(piece); + }; int64_t num_generated_tokens = ET_UNWRAP(text_token_generator_->generate( prompt_tokens, num_prompt_tokens, max_new_tokens - 1, temperature_ == -1.0f ? config.temperature : temperature_, - wrapped_callback)); + collect_callback)); // A temporary new line to separate stdout mixed with results. wrapped_callback(ET_UNWRAP_TOKENIZER( @@ -375,6 +467,15 @@ runtime::Error LLMRunner::generate_from_pos( stats_callback(*stats_); } + if (!config.warming && config_.verbose) { + // Reprint the complete answer on its own lines, after all logs and stats, + // so it is easy to identify and never overlaps other output. + printf("======== FULL ANSWER ========\n"); + printf("%s\n", full_answer.c_str()); + printf("=============================\n"); + fflush(stdout); + } + return runtime::Error::Ok; } @@ -409,6 +510,17 @@ void LLMRunner::stop() { } } +void LLMRunner::reset() { + auto* llm_io = static_cast(io_manager_.get()); + if (llm_io == nullptr) { + return; + } + auto err = llm_io->reset( + example::default_method_name_prefill, + example::default_method_name_decode); + (void)err; +} + std::unique_ptr load_tokenizer( const std::string& type, const std::string& tokenizer_path, @@ -633,6 +745,20 @@ std::unique_ptr create_llm_runner( runner->set_text_prefiller(std::move(text_prefiller)); runner->set_text_token_generator(std::move(text_token_generator)); + // Optional Gemma 4 vision encoder (separate PTE). If no vision model path is + // configured, the runner stays text-only. + if (runner_paths.vision_model_path.has_value() && + !runner_paths.vision_model_path.value().empty()) { + auto vision_encoder = std::make_unique( + runner_paths.vision_model_path.value()); + auto verr = vision_encoder->load(); + if (verr != executorch::runtime::Error::Ok) { + ET_LOG(Error, "Failed to load vision encoder."); + return nullptr; + } + runner->set_vision_encoder(std::move(vision_encoder)); + } + runner->load(); return runner; diff --git a/examples/samsung/executor_runner/llm/runner/runner.h b/examples/samsung/executor_runner/llm/runner/runner.h index 986d0b9850f..4bf26ee76fb 100644 --- a/examples/samsung/executor_runner/llm/runner/runner.h +++ b/examples/samsung/executor_runner/llm/runner/runner.h @@ -19,6 +19,7 @@ #include #include #include +#include namespace example { @@ -66,6 +67,9 @@ class LLMRunner : public llm::IRunner { const RunnerConfig& get_config() const { return config_; } + RunnerConfig& get_config() { + return config_; + } extension::Module* get_module() const { return module_.get(); } @@ -89,6 +93,12 @@ class LLMRunner : public llm::IRunner { std::unique_ptr generator) { text_token_generator_ = std::move(generator); } + void set_vision_encoder(std::unique_ptr encoder) { + vision_encoder_ = std::move(encoder); + } + VisionEncoder* get_vision_encoder() const { + return vision_encoder_.get(); + } // TODO (brucekk.kim): should it be static? static std::string get_formatted_prompt( @@ -96,7 +106,7 @@ class LLMRunner : public llm::IRunner { const std::string& system_prompt, DecoderModelVersion version); - void reset() {}; + void reset(); private: bool shouldStop_{false}; @@ -110,6 +120,8 @@ class LLMRunner : public llm::IRunner { std::unique_ptr text_token_generator_; std::unique_ptr io_manager_; + std::unique_ptr vision_encoder_; + std::unique_ptr stats_; float temperature_ = -1.0f; diff --git a/examples/samsung/executor_runner/llm/runner/stb_image_impl.cpp b/examples/samsung/executor_runner/llm/runner/stb_image_impl.cpp new file mode 100644 index 00000000000..02b01d4f407 --- /dev/null +++ b/examples/samsung/executor_runner/llm/runner/stb_image_impl.cpp @@ -0,0 +1,15 @@ +/* + * Copyright (c) 2026 Samsung Electronics Co. LTD + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +// Single translation unit that instantiates the stb_image implementation. +// This must be included from exactly one .cpp in the runner. +#define STB_IMAGE_IMPLEMENTATION +#include + +#define STB_IMAGE_RESIZE_IMPLEMENTATION +#include diff --git a/examples/samsung/executor_runner/llm/runner/vision_encoder.cpp b/examples/samsung/executor_runner/llm/runner/vision_encoder.cpp new file mode 100644 index 00000000000..7cb0af6de2a --- /dev/null +++ b/examples/samsung/executor_runner/llm/runner/vision_encoder.cpp @@ -0,0 +1,221 @@ +/* + * Copyright (c) 2026 Samsung Electronics Co. LTD + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#include + +#include + +#include +#include +#include +#include + +namespace example { + +namespace runtime = ::executorch::runtime; +namespace extension = ::executorch::extension; + +VisionEncoder::VisionEncoder(const std::string& model_path) + : model_path_(model_path) {} + +runtime::Error VisionEncoder::load() { + if (loaded_) { + return runtime::Error::Ok; + } + if (model_path_.empty()) { + ET_LOG(Error, "Vision model path is empty"); + return runtime::Error::InvalidArgument; + } + + module_ = std::make_unique( + model_path_, extension::Module::LoadMode::File); + + auto err = module_->load_method("vision_encoder"); + if (err != runtime::Error::Ok) { + ET_LOG( + Error, + "Failed to load vision_encoder method from %s", + model_path_.c_str()); + return err; + } + + auto method_meta = module_->method_meta("vision_encoder"); + ET_CHECK_OK_OR_RETURN_ERROR(method_meta.error()); + ET_CHECK_OR_RETURN_ERROR( + method_meta->num_inputs() >= 2 && method_meta->num_outputs() >= 2, + InvalidProgram, + "vision_encoder expects 2 inputs / 2 outputs, got %zu / %zu", + method_meta->num_inputs(), + method_meta->num_outputs()); + + auto pixel_meta = method_meta->input_tensor_meta(0); + ET_CHECK_OK_OR_RETURN_ERROR(pixel_meta.error()); + ET_CHECK_OR_RETURN_ERROR( + pixel_meta->scalar_type() == executorch::aten::ScalarType::Float, + InvalidProgram, + "vision_encoder pixel_values must be float32, got dtype %d", + static_cast(pixel_meta->scalar_type())); + const auto pixel_sizes = pixel_meta->sizes(); + ET_CHECK_OR_RETURN_ERROR( + pixel_sizes.size() == 3 && pixel_sizes[2] == kVisionPatchDim && + pixel_sizes[1] % (kVisionPoolingKernel * kVisionPoolingKernel) == 0, + InvalidProgram, + "Unexpected vision_encoder pixel_values shape"); + max_soft_tokens_ = + pixel_sizes[1] / (kVisionPoolingKernel * kVisionPoolingKernel); + + auto pos_meta = method_meta->input_tensor_meta(1); + ET_CHECK_OK_OR_RETURN_ERROR(pos_meta.error()); + position_ids_dtype_ = pos_meta->scalar_type(); + ET_CHECK_OR_RETURN_ERROR( + position_ids_dtype_ == executorch::aten::ScalarType::Long || + position_ids_dtype_ == executorch::aten::ScalarType::Int, + InvalidProgram, + "vision_encoder pixel_position_ids must be int64/int32, got dtype %d", + static_cast(position_ids_dtype_)); + + auto out_meta = method_meta->output_tensor_meta(0); + if (out_meta.ok() && out_meta->sizes().size() >= 3) { + hidden_size_ = out_meta->sizes()[2]; + } + ET_LOG( + Info, + "Vision encoder: max_soft_tokens %d, position_ids dtype %d, hidden %lld", + max_soft_tokens_, + static_cast(position_ids_dtype_), + static_cast(hidden_size_)); + + loaded_ = true; + ET_LOG(Info, "Vision encoder loaded: %s", model_path_.c_str()); + return runtime::Error::Ok; +} + +runtime::Error VisionEncoder::encode(const std::string& image_path) { + if (!loaded_) { + ET_LOG(Error, "Vision encoder is not loaded"); + return runtime::Error::InvalidState; + } + + // Load and patchify the image. + int32_t width = 0; + int32_t height = 0; + std::vector rgb; + try { + rgb = load_image_rgb(image_path, &width, &height); + } catch (const std::exception& e) { + ET_LOG(Error, "Failed to load image: %s", e.what()); + return runtime::Error::InvalidArgument; + } + + ImageData image = + patchify_rgb_image(rgb.data(), width, height, max_soft_tokens_); + + extension::TensorPtr position_ids = image.pixel_position_ids; + if (position_ids_dtype_ == executorch::aten::ScalarType::Int) { + position_ids = extension::empty( + {1, static_cast(position_ids->size(1)), 2}, + executorch::aten::ScalarType::Int); + std::copy_n( + image.pixel_position_ids->const_data_ptr(), + position_ids->numel(), + position_ids->mutable_data_ptr()); + } + + std::vector input_vec = {image.pixel_values, position_ids}; + + auto results = module_->execute("vision_encoder", input_vec); + ET_CHECK_OK_OR_RETURN_ERROR(results.error()); + + auto& outputs = results.get(); + if (outputs.size() < 2 || !outputs[0].isTensor() || !outputs[1].isTensor()) { + ET_LOG(Error, "Unexpected vision encoder outputs (%zu)", outputs.size()); + return runtime::Error::InvalidProgram; + } + + const auto& embeds = outputs[0].toTensor(); + const auto& mask = outputs[1].toTensor(); + + if (embeds.scalar_type() != executorch::aten::ScalarType::Float) { + ET_LOG(Error, "Vision embeddings are not float32"); + return runtime::Error::InvalidProgram; + } + + const float* embeds_data = embeds.const_data_ptr(); + const auto pooled = embeds.size(1); + const auto hidden = embeds.size(2); + hidden_size_ = hidden; + + ET_CHECK_OR_RETURN_ERROR( + mask.numel() == pooled, + InvalidProgram, + "Vision mask numel %lld != pooled tokens %lld", + static_cast(mask.numel()), + static_cast(pooled)); + + // The delegate may not preserve the BOOL dtype of the exported mask. + std::vector valid(pooled); + switch (mask.scalar_type()) { + case executorch::aten::ScalarType::Bool: + std::copy_n(mask.const_data_ptr(), pooled, valid.begin()); + break; + case executorch::aten::ScalarType::Byte: + std::copy_n(mask.const_data_ptr(), pooled, valid.begin()); + break; + case executorch::aten::ScalarType::Int: + std::copy_n(mask.const_data_ptr(), pooled, valid.begin()); + break; + case executorch::aten::ScalarType::Long: + std::copy_n(mask.const_data_ptr(), pooled, valid.begin()); + break; + case executorch::aten::ScalarType::Float: + std::copy_n(mask.const_data_ptr(), pooled, valid.begin()); + break; + default: + ET_LOG( + Error, + "Unsupported vision mask dtype %d", + static_cast(mask.scalar_type())); + return runtime::Error::InvalidProgram; + } + num_valid_tokens_ = std::count(valid.begin(), valid.end(), true); + + const int64_t expected_tokens = + compute_vision_num_tokens(image.num_valid_patches); + if (num_valid_tokens_ != expected_tokens) { + ET_LOG( + Error, + "Vision mask reports %lld valid tokens, expected %lld for %dx%d resized image", + static_cast(num_valid_tokens_), + static_cast(expected_tokens), + image.original_width, + image.original_height); + return runtime::Error::InvalidProgram; + } + + // Keep only valid tokens' embeddings, packed contiguously. + vision_embeds_.clear(); + vision_embeds_.reserve(static_cast(num_valid_tokens_) * hidden); + for (int64_t i = 0; i < pooled; ++i) { + if (valid[i]) { + vision_embeds_.insert( + vision_embeds_.end(), + embeds_data + i * hidden, + embeds_data + (i + 1) * hidden); + } + } + + ET_LOG( + Info, + "Vision encoded: %lld valid tokens out of %lld", + static_cast(num_valid_tokens_), + static_cast(pooled)); + + return runtime::Error::Ok; +} + +} // namespace example diff --git a/examples/samsung/executor_runner/llm/runner/vision_encoder.h b/examples/samsung/executor_runner/llm/runner/vision_encoder.h new file mode 100644 index 00000000000..7f2954320e9 --- /dev/null +++ b/examples/samsung/executor_runner/llm/runner/vision_encoder.h @@ -0,0 +1,73 @@ +/* + * Copyright (c) 2026 Samsung Electronics Co. LTD + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +#include +#include +#include + +#include +#include +#include + +namespace example { + +namespace runtime = ::executorch::runtime; +namespace extension = ::executorch::extension; + +/** + * Runs the Gemma 4 vision encoder (gemma4_vision.pte) on an input image and + * produces LM-ready vision embeddings plus the number of valid (non-padding) + * vision tokens. + * + * The vision encoder is a separate module from the text decoder. Its PTE has a + * single "vision_encoder" method accepting pre-patchified pixel values + * [1, max_patches, 768] and 2D patch positions [1, max_patches, 2], and + * returning embeddings [1, pooled, hidden] plus a BOOL validity mask [1, + * pooled]. + */ +class VisionEncoder { + public: + explicit VisionEncoder(const std::string& model_path); + + runtime::Error load(); + + /// Encode an image file and produce vision embeddings. + /// On success fills vision_embeds_ ([1, pooled, hidden], float) and sets + /// num_valid_tokens_ to the number of real (non-padding) vision tokens. + runtime::Error encode(const std::string& image_path); + + const std::vector& vision_embeds() const { + return vision_embeds_; + } + int64_t num_valid_tokens() const { + return num_valid_tokens_; + } + int64_t hidden_size() const { + return hidden_size_; + } + bool loaded() const { + return loaded_; + } + + private: + std::string model_path_; + std::unique_ptr module_; + bool loaded_ = false; + // Patch budget and position-id dtype the PTE was exported with; read from + // the method metadata in load(). + int32_t max_soft_tokens_ = kVisionMaxSoftTokens; + executorch::aten::ScalarType position_ids_dtype_ = + executorch::aten::ScalarType::Long; + int64_t hidden_size_ = 0; + int64_t num_valid_tokens_ = 0; + std::vector vision_embeds_; +}; + +} // namespace example diff --git a/third-party/stb/stb_image.h b/third-party/stb/stb_image.h new file mode 100644 index 00000000000..9eedabedc45 --- /dev/null +++ b/third-party/stb/stb_image.h @@ -0,0 +1,7988 @@ +/* stb_image - v2.30 - public domain image loader - http://nothings.org/stb + no warranty implied; use at your own risk + + Do this: + #define STB_IMAGE_IMPLEMENTATION + before you include this file in *one* C or C++ file to create the implementation. + + // i.e. it should look like this: + #include ... + #include ... + #include ... + #define STB_IMAGE_IMPLEMENTATION + #include "stb_image.h" + + You can #define STBI_ASSERT(x) before the #include to avoid using assert.h. + And #define STBI_MALLOC, STBI_REALLOC, and STBI_FREE to avoid using malloc,realloc,free + + + QUICK NOTES: + Primarily of interest to game developers and other people who can + avoid problematic images and only need the trivial interface + + JPEG baseline & progressive (12 bpc/arithmetic not supported, same as stock IJG lib) + PNG 1/2/4/8/16-bit-per-channel + + TGA (not sure what subset, if a subset) + BMP non-1bpp, non-RLE + PSD (composited view only, no extra channels, 8/16 bit-per-channel) + + GIF (*comp always reports as 4-channel) + HDR (radiance rgbE format) + PIC (Softimage PIC) + PNM (PPM and PGM binary only) + + Animated GIF still needs a proper API, but here's one way to do it: + http://gist.github.com/urraka/685d9a6340b26b830d49 + + - decode from memory or through FILE (define STBI_NO_STDIO to remove code) + - decode from arbitrary I/O callbacks + - SIMD acceleration on x86/x64 (SSE2) and ARM (NEON) + + Full documentation under "DOCUMENTATION" below. + + +LICENSE + + See end of file for license information. + +RECENT REVISION HISTORY: + + 2.30 (2024-05-31) avoid erroneous gcc warning + 2.29 (2023-05-xx) optimizations + 2.28 (2023-01-29) many error fixes, security errors, just tons of stuff + 2.27 (2021-07-11) document stbi_info better, 16-bit PNM support, bug fixes + 2.26 (2020-07-13) many minor fixes + 2.25 (2020-02-02) fix warnings + 2.24 (2020-02-02) fix warnings; thread-local failure_reason and flip_vertically + 2.23 (2019-08-11) fix clang static analysis warning + 2.22 (2019-03-04) gif fixes, fix warnings + 2.21 (2019-02-25) fix typo in comment + 2.20 (2019-02-07) support utf8 filenames in Windows; fix warnings and platform ifdefs + 2.19 (2018-02-11) fix warning + 2.18 (2018-01-30) fix warnings + 2.17 (2018-01-29) bugfix, 1-bit BMP, 16-bitness query, fix warnings + 2.16 (2017-07-23) all functions have 16-bit variants; optimizations; bugfixes + 2.15 (2017-03-18) fix png-1,2,4; all Imagenet JPGs; no runtime SSE detection on GCC + 2.14 (2017-03-03) remove deprecated STBI_JPEG_OLD; fixes for Imagenet JPGs + 2.13 (2016-12-04) experimental 16-bit API, only for PNG so far; fixes + 2.12 (2016-04-02) fix typo in 2.11 PSD fix that caused crashes + 2.11 (2016-04-02) 16-bit PNGS; enable SSE2 in non-gcc x64 + RGB-format JPEG; remove white matting in PSD; + allocate large structures on the stack; + correct channel count for PNG & BMP + 2.10 (2016-01-22) avoid warning introduced in 2.09 + 2.09 (2016-01-16) 16-bit TGA; comments in PNM files; STBI_REALLOC_SIZED + + See end of file for full revision history. + + + ============================ Contributors ========================= + + Image formats Extensions, features + Sean Barrett (jpeg, png, bmp) Jetro Lauha (stbi_info) + Nicolas Schulz (hdr, psd) Martin "SpartanJ" Golini (stbi_info) + Jonathan Dummer (tga) James "moose2000" Brown (iPhone PNG) + Jean-Marc Lienher (gif) Ben "Disch" Wenger (io callbacks) + Tom Seddon (pic) Omar Cornut (1/2/4-bit PNG) + Thatcher Ulrich (psd) Nicolas Guillemot (vertical flip) + Ken Miller (pgm, ppm) Richard Mitton (16-bit PSD) + github:urraka (animated gif) Junggon Kim (PNM comments) + Christopher Forseth (animated gif) Daniel Gibson (16-bit TGA) + socks-the-fox (16-bit PNG) + Jeremy Sawicki (handle all ImageNet JPGs) + Optimizations & bugfixes Mikhail Morozov (1-bit BMP) + Fabian "ryg" Giesen Anael Seghezzi (is-16-bit query) + Arseny Kapoulkine Simon Breuss (16-bit PNM) + John-Mark Allen + Carmelo J Fdez-Aguera + + Bug & warning fixes + Marc LeBlanc David Woo Guillaume George Martins Mozeiko + Christpher Lloyd Jerry Jansson Joseph Thomson Blazej Dariusz Roszkowski + Phil Jordan Dave Moore Roy Eltham + Hayaki Saito Nathan Reed Won Chun + Luke Graham Johan Duparc Nick Verigakis the Horde3D community + Thomas Ruf Ronny Chevalier github:rlyeh + Janez Zemva John Bartholomew Michal Cichon github:romigrou + Jonathan Blow Ken Hamada Tero Hanninen github:svdijk + Eugene Golushkov Laurent Gomila Cort Stratton github:snagar + Aruelien Pocheville Sergio Gonzalez Thibault Reuille github:Zelex + Cass Everitt Ryamond Barbiero github:grim210 + Paul Du Bois Engin Manap Aldo Culquicondor github:sammyhw + Philipp Wiesemann Dale Weiler Oriol Ferrer Mesia github:phprus + Josh Tobin Neil Bickford Matthew Gregan github:poppolopoppo + Julian Raschke Gregory Mullen Christian Floisand github:darealshinji + Baldur Karlsson Kevin Schmidt JR Smith github:Michaelangel007 + Brad Weinberger Matvey Cherevko github:mosra + Luca Sas Alexander Veselov Zack Middleton [reserved] + Ryan C. Gordon [reserved] [reserved] + DO NOT ADD YOUR NAME HERE + + Jacko Dirks + + To add your name to the credits, pick a random blank space in the middle and fill it. + 80% of merge conflicts on stb PRs are due to people adding their name at the end + of the credits. +*/ + +#ifndef STBI_INCLUDE_STB_IMAGE_H +#define STBI_INCLUDE_STB_IMAGE_H + +// DOCUMENTATION +// +// Limitations: +// - no 12-bit-per-channel JPEG +// - no JPEGs with arithmetic coding +// - GIF always returns *comp=4 +// +// Basic usage (see HDR discussion below for HDR usage): +// int x,y,n; +// unsigned char *data = stbi_load(filename, &x, &y, &n, 0); +// // ... process data if not NULL ... +// // ... x = width, y = height, n = # 8-bit components per pixel ... +// // ... replace '0' with '1'..'4' to force that many components per pixel +// // ... but 'n' will always be the number that it would have been if you said 0 +// stbi_image_free(data); +// +// Standard parameters: +// int *x -- outputs image width in pixels +// int *y -- outputs image height in pixels +// int *channels_in_file -- outputs # of image components in image file +// int desired_channels -- if non-zero, # of image components requested in result +// +// The return value from an image loader is an 'unsigned char *' which points +// to the pixel data, or NULL on an allocation failure or if the image is +// corrupt or invalid. The pixel data consists of *y scanlines of *x pixels, +// with each pixel consisting of N interleaved 8-bit components; the first +// pixel pointed to is top-left-most in the image. There is no padding between +// image scanlines or between pixels, regardless of format. The number of +// components N is 'desired_channels' if desired_channels is non-zero, or +// *channels_in_file otherwise. If desired_channels is non-zero, +// *channels_in_file has the number of components that _would_ have been +// output otherwise. E.g. if you set desired_channels to 4, you will always +// get RGBA output, but you can check *channels_in_file to see if it's trivially +// opaque because e.g. there were only 3 channels in the source image. +// +// An output image with N components has the following components interleaved +// in this order in each pixel: +// +// N=#comp components +// 1 grey +// 2 grey, alpha +// 3 red, green, blue +// 4 red, green, blue, alpha +// +// If image loading fails for any reason, the return value will be NULL, +// and *x, *y, *channels_in_file will be unchanged. The function +// stbi_failure_reason() can be queried for an extremely brief, end-user +// unfriendly explanation of why the load failed. Define STBI_NO_FAILURE_STRINGS +// to avoid compiling these strings at all, and STBI_FAILURE_USERMSG to get slightly +// more user-friendly ones. +// +// Paletted PNG, BMP, GIF, and PIC images are automatically depalettized. +// +// To query the width, height and component count of an image without having to +// decode the full file, you can use the stbi_info family of functions: +// +// int x,y,n,ok; +// ok = stbi_info(filename, &x, &y, &n); +// // returns ok=1 and sets x, y, n if image is a supported format, +// // 0 otherwise. +// +// Note that stb_image pervasively uses ints in its public API for sizes, +// including sizes of memory buffers. This is now part of the API and thus +// hard to change without causing breakage. As a result, the various image +// loaders all have certain limits on image size; these differ somewhat +// by format but generally boil down to either just under 2GB or just under +// 1GB. When the decoded image would be larger than this, stb_image decoding +// will fail. +// +// Additionally, stb_image will reject image files that have any of their +// dimensions set to a larger value than the configurable STBI_MAX_DIMENSIONS, +// which defaults to 2**24 = 16777216 pixels. Due to the above memory limit, +// the only way to have an image with such dimensions load correctly +// is for it to have a rather extreme aspect ratio. Either way, the +// assumption here is that such larger images are likely to be malformed +// or malicious. If you do need to load an image with individual dimensions +// larger than that, and it still fits in the overall size limit, you can +// #define STBI_MAX_DIMENSIONS on your own to be something larger. +// +// =========================================================================== +// +// UNICODE: +// +// If compiling for Windows and you wish to use Unicode filenames, compile +// with +// #define STBI_WINDOWS_UTF8 +// and pass utf8-encoded filenames. Call stbi_convert_wchar_to_utf8 to convert +// Windows wchar_t filenames to utf8. +// +// =========================================================================== +// +// Philosophy +// +// stb libraries are designed with the following priorities: +// +// 1. easy to use +// 2. easy to maintain +// 3. good performance +// +// Sometimes I let "good performance" creep up in priority over "easy to maintain", +// and for best performance I may provide less-easy-to-use APIs that give higher +// performance, in addition to the easy-to-use ones. Nevertheless, it's important +// to keep in mind that from the standpoint of you, a client of this library, +// all you care about is #1 and #3, and stb libraries DO NOT emphasize #3 above all. +// +// Some secondary priorities arise directly from the first two, some of which +// provide more explicit reasons why performance can't be emphasized. +// +// - Portable ("ease of use") +// - Small source code footprint ("easy to maintain") +// - No dependencies ("ease of use") +// +// =========================================================================== +// +// I/O callbacks +// +// I/O callbacks allow you to read from arbitrary sources, like packaged +// files or some other source. Data read from callbacks are processed +// through a small internal buffer (currently 128 bytes) to try to reduce +// overhead. +// +// The three functions you must define are "read" (reads some bytes of data), +// "skip" (skips some bytes of data), "eof" (reports if the stream is at the end). +// +// =========================================================================== +// +// SIMD support +// +// The JPEG decoder will try to automatically use SIMD kernels on x86 when +// supported by the compiler. For ARM Neon support, you must explicitly +// request it. +// +// (The old do-it-yourself SIMD API is no longer supported in the current +// code.) +// +// On x86, SSE2 will automatically be used when available based on a run-time +// test; if not, the generic C versions are used as a fall-back. On ARM targets, +// the typical path is to have separate builds for NEON and non-NEON devices +// (at least this is true for iOS and Android). Therefore, the NEON support is +// toggled by a build flag: define STBI_NEON to get NEON loops. +// +// If for some reason you do not want to use any of SIMD code, or if +// you have issues compiling it, you can disable it entirely by +// defining STBI_NO_SIMD. +// +// =========================================================================== +// +// HDR image support (disable by defining STBI_NO_HDR) +// +// stb_image supports loading HDR images in general, and currently the Radiance +// .HDR file format specifically. You can still load any file through the existing +// interface; if you attempt to load an HDR file, it will be automatically remapped +// to LDR, assuming gamma 2.2 and an arbitrary scale factor defaulting to 1; +// both of these constants can be reconfigured through this interface: +// +// stbi_hdr_to_ldr_gamma(2.2f); +// stbi_hdr_to_ldr_scale(1.0f); +// +// (note, do not use _inverse_ constants; stbi_image will invert them +// appropriately). +// +// Additionally, there is a new, parallel interface for loading files as +// (linear) floats to preserve the full dynamic range: +// +// float *data = stbi_loadf(filename, &x, &y, &n, 0); +// +// If you load LDR images through this interface, those images will +// be promoted to floating point values, run through the inverse of +// constants corresponding to the above: +// +// stbi_ldr_to_hdr_scale(1.0f); +// stbi_ldr_to_hdr_gamma(2.2f); +// +// Finally, given a filename (or an open file or memory block--see header +// file for details) containing image data, you can query for the "most +// appropriate" interface to use (that is, whether the image is HDR or +// not), using: +// +// stbi_is_hdr(char *filename); +// +// =========================================================================== +// +// iPhone PNG support: +// +// We optionally support converting iPhone-formatted PNGs (which store +// premultiplied BGRA) back to RGB, even though they're internally encoded +// differently. To enable this conversion, call +// stbi_convert_iphone_png_to_rgb(1). +// +// Call stbi_set_unpremultiply_on_load(1) as well to force a divide per +// pixel to remove any premultiplied alpha *only* if the image file explicitly +// says there's premultiplied data (currently only happens in iPhone images, +// and only if iPhone convert-to-rgb processing is on). +// +// =========================================================================== +// +// ADDITIONAL CONFIGURATION +// +// - You can suppress implementation of any of the decoders to reduce +// your code footprint by #defining one or more of the following +// symbols before creating the implementation. +// +// STBI_NO_JPEG +// STBI_NO_PNG +// STBI_NO_BMP +// STBI_NO_PSD +// STBI_NO_TGA +// STBI_NO_GIF +// STBI_NO_HDR +// STBI_NO_PIC +// STBI_NO_PNM (.ppm and .pgm) +// +// - You can request *only* certain decoders and suppress all other ones +// (this will be more forward-compatible, as addition of new decoders +// doesn't require you to disable them explicitly): +// +// STBI_ONLY_JPEG +// STBI_ONLY_PNG +// STBI_ONLY_BMP +// STBI_ONLY_PSD +// STBI_ONLY_TGA +// STBI_ONLY_GIF +// STBI_ONLY_HDR +// STBI_ONLY_PIC +// STBI_ONLY_PNM (.ppm and .pgm) +// +// - If you use STBI_NO_PNG (or _ONLY_ without PNG), and you still +// want the zlib decoder to be available, #define STBI_SUPPORT_ZLIB +// +// - If you define STBI_MAX_DIMENSIONS, stb_image will reject images greater +// than that size (in either width or height) without further processing. +// This is to let programs in the wild set an upper bound to prevent +// denial-of-service attacks on untrusted data, as one could generate a +// valid image of gigantic dimensions and force stb_image to allocate a +// huge block of memory and spend disproportionate time decoding it. By +// default this is set to (1 << 24), which is 16777216, but that's still +// very big. + +#ifndef STBI_NO_STDIO +#include +#endif // STBI_NO_STDIO + +#define STBI_VERSION 1 + +enum +{ + STBI_default = 0, // only used for desired_channels + + STBI_grey = 1, + STBI_grey_alpha = 2, + STBI_rgb = 3, + STBI_rgb_alpha = 4 +}; + +#include +typedef unsigned char stbi_uc; +typedef unsigned short stbi_us; + +#ifdef __cplusplus +extern "C" { +#endif + +#ifndef STBIDEF +#ifdef STB_IMAGE_STATIC +#define STBIDEF static +#else +#define STBIDEF extern +#endif +#endif + +////////////////////////////////////////////////////////////////////////////// +// +// PRIMARY API - works on images of any type +// + +// +// load image by filename, open file, or memory buffer +// + +typedef struct +{ + int (*read) (void *user,char *data,int size); // fill 'data' with 'size' bytes. return number of bytes actually read + void (*skip) (void *user,int n); // skip the next 'n' bytes, or 'unget' the last -n bytes if negative + int (*eof) (void *user); // returns nonzero if we are at end of file/data +} stbi_io_callbacks; + +//////////////////////////////////// +// +// 8-bits-per-channel interface +// + +STBIDEF stbi_uc *stbi_load_from_memory (stbi_uc const *buffer, int len , int *x, int *y, int *channels_in_file, int desired_channels); +STBIDEF stbi_uc *stbi_load_from_callbacks(stbi_io_callbacks const *clbk , void *user, int *x, int *y, int *channels_in_file, int desired_channels); + +#ifndef STBI_NO_STDIO +STBIDEF stbi_uc *stbi_load (char const *filename, int *x, int *y, int *channels_in_file, int desired_channels); +STBIDEF stbi_uc *stbi_load_from_file (FILE *f, int *x, int *y, int *channels_in_file, int desired_channels); +// for stbi_load_from_file, file pointer is left pointing immediately after image +#endif + +#ifndef STBI_NO_GIF +STBIDEF stbi_uc *stbi_load_gif_from_memory(stbi_uc const *buffer, int len, int **delays, int *x, int *y, int *z, int *comp, int req_comp); +#endif + +#ifdef STBI_WINDOWS_UTF8 +STBIDEF int stbi_convert_wchar_to_utf8(char *buffer, size_t bufferlen, const wchar_t* input); +#endif + +//////////////////////////////////// +// +// 16-bits-per-channel interface +// + +STBIDEF stbi_us *stbi_load_16_from_memory (stbi_uc const *buffer, int len, int *x, int *y, int *channels_in_file, int desired_channels); +STBIDEF stbi_us *stbi_load_16_from_callbacks(stbi_io_callbacks const *clbk, void *user, int *x, int *y, int *channels_in_file, int desired_channels); + +#ifndef STBI_NO_STDIO +STBIDEF stbi_us *stbi_load_16 (char const *filename, int *x, int *y, int *channels_in_file, int desired_channels); +STBIDEF stbi_us *stbi_load_from_file_16(FILE *f, int *x, int *y, int *channels_in_file, int desired_channels); +#endif + +//////////////////////////////////// +// +// float-per-channel interface +// +#ifndef STBI_NO_LINEAR + STBIDEF float *stbi_loadf_from_memory (stbi_uc const *buffer, int len, int *x, int *y, int *channels_in_file, int desired_channels); + STBIDEF float *stbi_loadf_from_callbacks (stbi_io_callbacks const *clbk, void *user, int *x, int *y, int *channels_in_file, int desired_channels); + + #ifndef STBI_NO_STDIO + STBIDEF float *stbi_loadf (char const *filename, int *x, int *y, int *channels_in_file, int desired_channels); + STBIDEF float *stbi_loadf_from_file (FILE *f, int *x, int *y, int *channels_in_file, int desired_channels); + #endif +#endif + +#ifndef STBI_NO_HDR + STBIDEF void stbi_hdr_to_ldr_gamma(float gamma); + STBIDEF void stbi_hdr_to_ldr_scale(float scale); +#endif // STBI_NO_HDR + +#ifndef STBI_NO_LINEAR + STBIDEF void stbi_ldr_to_hdr_gamma(float gamma); + STBIDEF void stbi_ldr_to_hdr_scale(float scale); +#endif // STBI_NO_LINEAR + +// stbi_is_hdr is always defined, but always returns false if STBI_NO_HDR +STBIDEF int stbi_is_hdr_from_callbacks(stbi_io_callbacks const *clbk, void *user); +STBIDEF int stbi_is_hdr_from_memory(stbi_uc const *buffer, int len); +#ifndef STBI_NO_STDIO +STBIDEF int stbi_is_hdr (char const *filename); +STBIDEF int stbi_is_hdr_from_file(FILE *f); +#endif // STBI_NO_STDIO + + +// get a VERY brief reason for failure +// on most compilers (and ALL modern mainstream compilers) this is threadsafe +STBIDEF const char *stbi_failure_reason (void); + +// free the loaded image -- this is just free() +STBIDEF void stbi_image_free (void *retval_from_stbi_load); + +// get image dimensions & components without fully decoding +STBIDEF int stbi_info_from_memory(stbi_uc const *buffer, int len, int *x, int *y, int *comp); +STBIDEF int stbi_info_from_callbacks(stbi_io_callbacks const *clbk, void *user, int *x, int *y, int *comp); +STBIDEF int stbi_is_16_bit_from_memory(stbi_uc const *buffer, int len); +STBIDEF int stbi_is_16_bit_from_callbacks(stbi_io_callbacks const *clbk, void *user); + +#ifndef STBI_NO_STDIO +STBIDEF int stbi_info (char const *filename, int *x, int *y, int *comp); +STBIDEF int stbi_info_from_file (FILE *f, int *x, int *y, int *comp); +STBIDEF int stbi_is_16_bit (char const *filename); +STBIDEF int stbi_is_16_bit_from_file(FILE *f); +#endif + + + +// for image formats that explicitly notate that they have premultiplied alpha, +// we just return the colors as stored in the file. set this flag to force +// unpremultiplication. results are undefined if the unpremultiply overflow. +STBIDEF void stbi_set_unpremultiply_on_load(int flag_true_if_should_unpremultiply); + +// indicate whether we should process iphone images back to canonical format, +// or just pass them through "as-is" +STBIDEF void stbi_convert_iphone_png_to_rgb(int flag_true_if_should_convert); + +// flip the image vertically, so the first pixel in the output array is the bottom left +STBIDEF void stbi_set_flip_vertically_on_load(int flag_true_if_should_flip); + +// as above, but only applies to images loaded on the thread that calls the function +// this function is only available if your compiler supports thread-local variables; +// calling it will fail to link if your compiler doesn't +STBIDEF void stbi_set_unpremultiply_on_load_thread(int flag_true_if_should_unpremultiply); +STBIDEF void stbi_convert_iphone_png_to_rgb_thread(int flag_true_if_should_convert); +STBIDEF void stbi_set_flip_vertically_on_load_thread(int flag_true_if_should_flip); + +// ZLIB client - used by PNG, available for other purposes + +STBIDEF char *stbi_zlib_decode_malloc_guesssize(const char *buffer, int len, int initial_size, int *outlen); +STBIDEF char *stbi_zlib_decode_malloc_guesssize_headerflag(const char *buffer, int len, int initial_size, int *outlen, int parse_header); +STBIDEF char *stbi_zlib_decode_malloc(const char *buffer, int len, int *outlen); +STBIDEF int stbi_zlib_decode_buffer(char *obuffer, int olen, const char *ibuffer, int ilen); + +STBIDEF char *stbi_zlib_decode_noheader_malloc(const char *buffer, int len, int *outlen); +STBIDEF int stbi_zlib_decode_noheader_buffer(char *obuffer, int olen, const char *ibuffer, int ilen); + + +#ifdef __cplusplus +} +#endif + +// +// +//// end header file ///////////////////////////////////////////////////// +#endif // STBI_INCLUDE_STB_IMAGE_H + +#ifdef STB_IMAGE_IMPLEMENTATION + +#if defined(STBI_ONLY_JPEG) || defined(STBI_ONLY_PNG) || defined(STBI_ONLY_BMP) \ + || defined(STBI_ONLY_TGA) || defined(STBI_ONLY_GIF) || defined(STBI_ONLY_PSD) \ + || defined(STBI_ONLY_HDR) || defined(STBI_ONLY_PIC) || defined(STBI_ONLY_PNM) \ + || defined(STBI_ONLY_ZLIB) + #ifndef STBI_ONLY_JPEG + #define STBI_NO_JPEG + #endif + #ifndef STBI_ONLY_PNG + #define STBI_NO_PNG + #endif + #ifndef STBI_ONLY_BMP + #define STBI_NO_BMP + #endif + #ifndef STBI_ONLY_PSD + #define STBI_NO_PSD + #endif + #ifndef STBI_ONLY_TGA + #define STBI_NO_TGA + #endif + #ifndef STBI_ONLY_GIF + #define STBI_NO_GIF + #endif + #ifndef STBI_ONLY_HDR + #define STBI_NO_HDR + #endif + #ifndef STBI_ONLY_PIC + #define STBI_NO_PIC + #endif + #ifndef STBI_ONLY_PNM + #define STBI_NO_PNM + #endif +#endif + +#if defined(STBI_NO_PNG) && !defined(STBI_SUPPORT_ZLIB) && !defined(STBI_NO_ZLIB) +#define STBI_NO_ZLIB +#endif + + +#include +#include // ptrdiff_t on osx +#include +#include +#include + +#if !defined(STBI_NO_LINEAR) || !defined(STBI_NO_HDR) +#include // ldexp, pow +#endif + +#ifndef STBI_NO_STDIO +#include +#endif + +#ifndef STBI_ASSERT +#include +#define STBI_ASSERT(x) assert(x) +#endif + +#ifdef __cplusplus +#define STBI_EXTERN extern "C" +#else +#define STBI_EXTERN extern +#endif + + +#ifndef _MSC_VER + #ifdef __cplusplus + #define stbi_inline inline + #else + #define stbi_inline + #endif +#else + #define stbi_inline __forceinline +#endif + +#ifndef STBI_NO_THREAD_LOCALS + #if defined(__cplusplus) && __cplusplus >= 201103L + #define STBI_THREAD_LOCAL thread_local + #elif defined(__GNUC__) && __GNUC__ < 5 + #define STBI_THREAD_LOCAL __thread + #elif defined(_MSC_VER) + #define STBI_THREAD_LOCAL __declspec(thread) + #elif defined (__STDC_VERSION__) && __STDC_VERSION__ >= 201112L && !defined(__STDC_NO_THREADS__) + #define STBI_THREAD_LOCAL _Thread_local + #endif + + #ifndef STBI_THREAD_LOCAL + #if defined(__GNUC__) + #define STBI_THREAD_LOCAL __thread + #endif + #endif +#endif + +#if defined(_MSC_VER) || defined(__SYMBIAN32__) +typedef unsigned short stbi__uint16; +typedef signed short stbi__int16; +typedef unsigned int stbi__uint32; +typedef signed int stbi__int32; +#else +#include +typedef uint16_t stbi__uint16; +typedef int16_t stbi__int16; +typedef uint32_t stbi__uint32; +typedef int32_t stbi__int32; +#endif + +// should produce compiler error if size is wrong +typedef unsigned char validate_uint32[sizeof(stbi__uint32)==4 ? 1 : -1]; + +#ifdef _MSC_VER +#define STBI_NOTUSED(v) (void)(v) +#else +#define STBI_NOTUSED(v) (void)sizeof(v) +#endif + +#ifdef _MSC_VER +#define STBI_HAS_LROTL +#endif + +#ifdef STBI_HAS_LROTL + #define stbi_lrot(x,y) _lrotl(x,y) +#else + #define stbi_lrot(x,y) (((x) << (y)) | ((x) >> (-(y) & 31))) +#endif + +#if defined(STBI_MALLOC) && defined(STBI_FREE) && (defined(STBI_REALLOC) || defined(STBI_REALLOC_SIZED)) +// ok +#elif !defined(STBI_MALLOC) && !defined(STBI_FREE) && !defined(STBI_REALLOC) && !defined(STBI_REALLOC_SIZED) +// ok +#else +#error "Must define all or none of STBI_MALLOC, STBI_FREE, and STBI_REALLOC (or STBI_REALLOC_SIZED)." +#endif + +#ifndef STBI_MALLOC +#define STBI_MALLOC(sz) malloc(sz) +#define STBI_REALLOC(p,newsz) realloc(p,newsz) +#define STBI_FREE(p) free(p) +#endif + +#ifndef STBI_REALLOC_SIZED +#define STBI_REALLOC_SIZED(p,oldsz,newsz) STBI_REALLOC(p,newsz) +#endif + +// x86/x64 detection +#if defined(__x86_64__) || defined(_M_X64) +#define STBI__X64_TARGET +#elif defined(__i386) || defined(_M_IX86) +#define STBI__X86_TARGET +#endif + +#if defined(__GNUC__) && defined(STBI__X86_TARGET) && !defined(__SSE2__) && !defined(STBI_NO_SIMD) +// gcc doesn't support sse2 intrinsics unless you compile with -msse2, +// which in turn means it gets to use SSE2 everywhere. This is unfortunate, +// but previous attempts to provide the SSE2 functions with runtime +// detection caused numerous issues. The way architecture extensions are +// exposed in GCC/Clang is, sadly, not really suited for one-file libs. +// New behavior: if compiled with -msse2, we use SSE2 without any +// detection; if not, we don't use it at all. +#define STBI_NO_SIMD +#endif + +#if defined(__MINGW32__) && defined(STBI__X86_TARGET) && !defined(STBI_MINGW_ENABLE_SSE2) && !defined(STBI_NO_SIMD) +// Note that __MINGW32__ doesn't actually mean 32-bit, so we have to avoid STBI__X64_TARGET +// +// 32-bit MinGW wants ESP to be 16-byte aligned, but this is not in the +// Windows ABI and VC++ as well as Windows DLLs don't maintain that invariant. +// As a result, enabling SSE2 on 32-bit MinGW is dangerous when not +// simultaneously enabling "-mstackrealign". +// +// See https://github.com/nothings/stb/issues/81 for more information. +// +// So default to no SSE2 on 32-bit MinGW. If you've read this far and added +// -mstackrealign to your build settings, feel free to #define STBI_MINGW_ENABLE_SSE2. +#define STBI_NO_SIMD +#endif + +#if !defined(STBI_NO_SIMD) && (defined(STBI__X86_TARGET) || defined(STBI__X64_TARGET)) +#define STBI_SSE2 +#include + +#ifdef _MSC_VER + +#if _MSC_VER >= 1400 // not VC6 +#include // __cpuid +static int stbi__cpuid3(void) +{ + int info[4]; + __cpuid(info,1); + return info[3]; +} +#else +static int stbi__cpuid3(void) +{ + int res; + __asm { + mov eax,1 + cpuid + mov res,edx + } + return res; +} +#endif + +#define STBI_SIMD_ALIGN(type, name) __declspec(align(16)) type name + +#if !defined(STBI_NO_JPEG) && defined(STBI_SSE2) +static int stbi__sse2_available(void) +{ + int info3 = stbi__cpuid3(); + return ((info3 >> 26) & 1) != 0; +} +#endif + +#else // assume GCC-style if not VC++ +#define STBI_SIMD_ALIGN(type, name) type name __attribute__((aligned(16))) + +#if !defined(STBI_NO_JPEG) && defined(STBI_SSE2) +static int stbi__sse2_available(void) +{ + // If we're even attempting to compile this on GCC/Clang, that means + // -msse2 is on, which means the compiler is allowed to use SSE2 + // instructions at will, and so are we. + return 1; +} +#endif + +#endif +#endif + +// ARM NEON +#if defined(STBI_NO_SIMD) && defined(STBI_NEON) +#undef STBI_NEON +#endif + +#ifdef STBI_NEON +#include +#ifdef _MSC_VER +#define STBI_SIMD_ALIGN(type, name) __declspec(align(16)) type name +#else +#define STBI_SIMD_ALIGN(type, name) type name __attribute__((aligned(16))) +#endif +#endif + +#ifndef STBI_SIMD_ALIGN +#define STBI_SIMD_ALIGN(type, name) type name +#endif + +#ifndef STBI_MAX_DIMENSIONS +#define STBI_MAX_DIMENSIONS (1 << 24) +#endif + +/////////////////////////////////////////////// +// +// stbi__context struct and start_xxx functions + +// stbi__context structure is our basic context used by all images, so it +// contains all the IO context, plus some basic image information +typedef struct +{ + stbi__uint32 img_x, img_y; + int img_n, img_out_n; + + stbi_io_callbacks io; + void *io_user_data; + + int read_from_callbacks; + int buflen; + stbi_uc buffer_start[128]; + int callback_already_read; + + stbi_uc *img_buffer, *img_buffer_end; + stbi_uc *img_buffer_original, *img_buffer_original_end; +} stbi__context; + + +static void stbi__refill_buffer(stbi__context *s); + +// initialize a memory-decode context +static void stbi__start_mem(stbi__context *s, stbi_uc const *buffer, int len) +{ + s->io.read = NULL; + s->read_from_callbacks = 0; + s->callback_already_read = 0; + s->img_buffer = s->img_buffer_original = (stbi_uc *) buffer; + s->img_buffer_end = s->img_buffer_original_end = (stbi_uc *) buffer+len; +} + +// initialize a callback-based context +static void stbi__start_callbacks(stbi__context *s, stbi_io_callbacks *c, void *user) +{ + s->io = *c; + s->io_user_data = user; + s->buflen = sizeof(s->buffer_start); + s->read_from_callbacks = 1; + s->callback_already_read = 0; + s->img_buffer = s->img_buffer_original = s->buffer_start; + stbi__refill_buffer(s); + s->img_buffer_original_end = s->img_buffer_end; +} + +#ifndef STBI_NO_STDIO + +static int stbi__stdio_read(void *user, char *data, int size) +{ + return (int) fread(data,1,size,(FILE*) user); +} + +static void stbi__stdio_skip(void *user, int n) +{ + int ch; + fseek((FILE*) user, n, SEEK_CUR); + ch = fgetc((FILE*) user); /* have to read a byte to reset feof()'s flag */ + if (ch != EOF) { + ungetc(ch, (FILE *) user); /* push byte back onto stream if valid. */ + } +} + +static int stbi__stdio_eof(void *user) +{ + return feof((FILE*) user) || ferror((FILE *) user); +} + +static stbi_io_callbacks stbi__stdio_callbacks = +{ + stbi__stdio_read, + stbi__stdio_skip, + stbi__stdio_eof, +}; + +static void stbi__start_file(stbi__context *s, FILE *f) +{ + stbi__start_callbacks(s, &stbi__stdio_callbacks, (void *) f); +} + +//static void stop_file(stbi__context *s) { } + +#endif // !STBI_NO_STDIO + +static void stbi__rewind(stbi__context *s) +{ + // conceptually rewind SHOULD rewind to the beginning of the stream, + // but we just rewind to the beginning of the initial buffer, because + // we only use it after doing 'test', which only ever looks at at most 92 bytes + s->img_buffer = s->img_buffer_original; + s->img_buffer_end = s->img_buffer_original_end; +} + +enum +{ + STBI_ORDER_RGB, + STBI_ORDER_BGR +}; + +typedef struct +{ + int bits_per_channel; + int num_channels; + int channel_order; +} stbi__result_info; + +#ifndef STBI_NO_JPEG +static int stbi__jpeg_test(stbi__context *s); +static void *stbi__jpeg_load(stbi__context *s, int *x, int *y, int *comp, int req_comp, stbi__result_info *ri); +static int stbi__jpeg_info(stbi__context *s, int *x, int *y, int *comp); +#endif + +#ifndef STBI_NO_PNG +static int stbi__png_test(stbi__context *s); +static void *stbi__png_load(stbi__context *s, int *x, int *y, int *comp, int req_comp, stbi__result_info *ri); +static int stbi__png_info(stbi__context *s, int *x, int *y, int *comp); +static int stbi__png_is16(stbi__context *s); +#endif + +#ifndef STBI_NO_BMP +static int stbi__bmp_test(stbi__context *s); +static void *stbi__bmp_load(stbi__context *s, int *x, int *y, int *comp, int req_comp, stbi__result_info *ri); +static int stbi__bmp_info(stbi__context *s, int *x, int *y, int *comp); +#endif + +#ifndef STBI_NO_TGA +static int stbi__tga_test(stbi__context *s); +static void *stbi__tga_load(stbi__context *s, int *x, int *y, int *comp, int req_comp, stbi__result_info *ri); +static int stbi__tga_info(stbi__context *s, int *x, int *y, int *comp); +#endif + +#ifndef STBI_NO_PSD +static int stbi__psd_test(stbi__context *s); +static void *stbi__psd_load(stbi__context *s, int *x, int *y, int *comp, int req_comp, stbi__result_info *ri, int bpc); +static int stbi__psd_info(stbi__context *s, int *x, int *y, int *comp); +static int stbi__psd_is16(stbi__context *s); +#endif + +#ifndef STBI_NO_HDR +static int stbi__hdr_test(stbi__context *s); +static float *stbi__hdr_load(stbi__context *s, int *x, int *y, int *comp, int req_comp, stbi__result_info *ri); +static int stbi__hdr_info(stbi__context *s, int *x, int *y, int *comp); +#endif + +#ifndef STBI_NO_PIC +static int stbi__pic_test(stbi__context *s); +static void *stbi__pic_load(stbi__context *s, int *x, int *y, int *comp, int req_comp, stbi__result_info *ri); +static int stbi__pic_info(stbi__context *s, int *x, int *y, int *comp); +#endif + +#ifndef STBI_NO_GIF +static int stbi__gif_test(stbi__context *s); +static void *stbi__gif_load(stbi__context *s, int *x, int *y, int *comp, int req_comp, stbi__result_info *ri); +static void *stbi__load_gif_main(stbi__context *s, int **delays, int *x, int *y, int *z, int *comp, int req_comp); +static int stbi__gif_info(stbi__context *s, int *x, int *y, int *comp); +#endif + +#ifndef STBI_NO_PNM +static int stbi__pnm_test(stbi__context *s); +static void *stbi__pnm_load(stbi__context *s, int *x, int *y, int *comp, int req_comp, stbi__result_info *ri); +static int stbi__pnm_info(stbi__context *s, int *x, int *y, int *comp); +static int stbi__pnm_is16(stbi__context *s); +#endif + +static +#ifdef STBI_THREAD_LOCAL +STBI_THREAD_LOCAL +#endif +const char *stbi__g_failure_reason; + +STBIDEF const char *stbi_failure_reason(void) +{ + return stbi__g_failure_reason; +} + +#ifndef STBI_NO_FAILURE_STRINGS +static int stbi__err(const char *str) +{ + stbi__g_failure_reason = str; + return 0; +} +#endif + +static void *stbi__malloc(size_t size) +{ + return STBI_MALLOC(size); +} + +// stb_image uses ints pervasively, including for offset calculations. +// therefore the largest decoded image size we can support with the +// current code, even on 64-bit targets, is INT_MAX. this is not a +// significant limitation for the intended use case. +// +// we do, however, need to make sure our size calculations don't +// overflow. hence a few helper functions for size calculations that +// multiply integers together, making sure that they're non-negative +// and no overflow occurs. + +// return 1 if the sum is valid, 0 on overflow. +// negative terms are considered invalid. +static int stbi__addsizes_valid(int a, int b) +{ + if (b < 0) return 0; + // now 0 <= b <= INT_MAX, hence also + // 0 <= INT_MAX - b <= INTMAX. + // And "a + b <= INT_MAX" (which might overflow) is the + // same as a <= INT_MAX - b (no overflow) + return a <= INT_MAX - b; +} + +// returns 1 if the product is valid, 0 on overflow. +// negative factors are considered invalid. +static int stbi__mul2sizes_valid(int a, int b) +{ + if (a < 0 || b < 0) return 0; + if (b == 0) return 1; // mul-by-0 is always safe + // portable way to check for no overflows in a*b + return a <= INT_MAX/b; +} + +#if !defined(STBI_NO_JPEG) || !defined(STBI_NO_PNG) || !defined(STBI_NO_TGA) || !defined(STBI_NO_HDR) +// returns 1 if "a*b + add" has no negative terms/factors and doesn't overflow +static int stbi__mad2sizes_valid(int a, int b, int add) +{ + return stbi__mul2sizes_valid(a, b) && stbi__addsizes_valid(a*b, add); +} +#endif + +// returns 1 if "a*b*c + add" has no negative terms/factors and doesn't overflow +static int stbi__mad3sizes_valid(int a, int b, int c, int add) +{ + return stbi__mul2sizes_valid(a, b) && stbi__mul2sizes_valid(a*b, c) && + stbi__addsizes_valid(a*b*c, add); +} + +// returns 1 if "a*b*c*d + add" has no negative terms/factors and doesn't overflow +#if !defined(STBI_NO_LINEAR) || !defined(STBI_NO_HDR) || !defined(STBI_NO_PNM) +static int stbi__mad4sizes_valid(int a, int b, int c, int d, int add) +{ + return stbi__mul2sizes_valid(a, b) && stbi__mul2sizes_valid(a*b, c) && + stbi__mul2sizes_valid(a*b*c, d) && stbi__addsizes_valid(a*b*c*d, add); +} +#endif + +#if !defined(STBI_NO_JPEG) || !defined(STBI_NO_PNG) || !defined(STBI_NO_TGA) || !defined(STBI_NO_HDR) +// mallocs with size overflow checking +static void *stbi__malloc_mad2(int a, int b, int add) +{ + if (!stbi__mad2sizes_valid(a, b, add)) return NULL; + return stbi__malloc(a*b + add); +} +#endif + +static void *stbi__malloc_mad3(int a, int b, int c, int add) +{ + if (!stbi__mad3sizes_valid(a, b, c, add)) return NULL; + return stbi__malloc(a*b*c + add); +} + +#if !defined(STBI_NO_LINEAR) || !defined(STBI_NO_HDR) || !defined(STBI_NO_PNM) +static void *stbi__malloc_mad4(int a, int b, int c, int d, int add) +{ + if (!stbi__mad4sizes_valid(a, b, c, d, add)) return NULL; + return stbi__malloc(a*b*c*d + add); +} +#endif + +// returns 1 if the sum of two signed ints is valid (between -2^31 and 2^31-1 inclusive), 0 on overflow. +static int stbi__addints_valid(int a, int b) +{ + if ((a >= 0) != (b >= 0)) return 1; // a and b have different signs, so no overflow + if (a < 0 && b < 0) return a >= INT_MIN - b; // same as a + b >= INT_MIN; INT_MIN - b cannot overflow since b < 0. + return a <= INT_MAX - b; +} + +// returns 1 if the product of two ints fits in a signed short, 0 on overflow. +static int stbi__mul2shorts_valid(int a, int b) +{ + if (b == 0 || b == -1) return 1; // multiplication by 0 is always 0; check for -1 so SHRT_MIN/b doesn't overflow + if ((a >= 0) == (b >= 0)) return a <= SHRT_MAX/b; // product is positive, so similar to mul2sizes_valid + if (b < 0) return a <= SHRT_MIN / b; // same as a * b >= SHRT_MIN + return a >= SHRT_MIN / b; +} + +// stbi__err - error +// stbi__errpf - error returning pointer to float +// stbi__errpuc - error returning pointer to unsigned char + +#ifdef STBI_NO_FAILURE_STRINGS + #define stbi__err(x,y) 0 +#elif defined(STBI_FAILURE_USERMSG) + #define stbi__err(x,y) stbi__err(y) +#else + #define stbi__err(x,y) stbi__err(x) +#endif + +#define stbi__errpf(x,y) ((float *)(size_t) (stbi__err(x,y)?NULL:NULL)) +#define stbi__errpuc(x,y) ((unsigned char *)(size_t) (stbi__err(x,y)?NULL:NULL)) + +STBIDEF void stbi_image_free(void *retval_from_stbi_load) +{ + STBI_FREE(retval_from_stbi_load); +} + +#ifndef STBI_NO_LINEAR +static float *stbi__ldr_to_hdr(stbi_uc *data, int x, int y, int comp); +#endif + +#ifndef STBI_NO_HDR +static stbi_uc *stbi__hdr_to_ldr(float *data, int x, int y, int comp); +#endif + +static int stbi__vertically_flip_on_load_global = 0; + +STBIDEF void stbi_set_flip_vertically_on_load(int flag_true_if_should_flip) +{ + stbi__vertically_flip_on_load_global = flag_true_if_should_flip; +} + +#ifndef STBI_THREAD_LOCAL +#define stbi__vertically_flip_on_load stbi__vertically_flip_on_load_global +#else +static STBI_THREAD_LOCAL int stbi__vertically_flip_on_load_local, stbi__vertically_flip_on_load_set; + +STBIDEF void stbi_set_flip_vertically_on_load_thread(int flag_true_if_should_flip) +{ + stbi__vertically_flip_on_load_local = flag_true_if_should_flip; + stbi__vertically_flip_on_load_set = 1; +} + +#define stbi__vertically_flip_on_load (stbi__vertically_flip_on_load_set \ + ? stbi__vertically_flip_on_load_local \ + : stbi__vertically_flip_on_load_global) +#endif // STBI_THREAD_LOCAL + +static void *stbi__load_main(stbi__context *s, int *x, int *y, int *comp, int req_comp, stbi__result_info *ri, int bpc) +{ + memset(ri, 0, sizeof(*ri)); // make sure it's initialized if we add new fields + ri->bits_per_channel = 8; // default is 8 so most paths don't have to be changed + ri->channel_order = STBI_ORDER_RGB; // all current input & output are this, but this is here so we can add BGR order + ri->num_channels = 0; + + // test the formats with a very explicit header first (at least a FOURCC + // or distinctive magic number first) + #ifndef STBI_NO_PNG + if (stbi__png_test(s)) return stbi__png_load(s,x,y,comp,req_comp, ri); + #endif + #ifndef STBI_NO_BMP + if (stbi__bmp_test(s)) return stbi__bmp_load(s,x,y,comp,req_comp, ri); + #endif + #ifndef STBI_NO_GIF + if (stbi__gif_test(s)) return stbi__gif_load(s,x,y,comp,req_comp, ri); + #endif + #ifndef STBI_NO_PSD + if (stbi__psd_test(s)) return stbi__psd_load(s,x,y,comp,req_comp, ri, bpc); + #else + STBI_NOTUSED(bpc); + #endif + #ifndef STBI_NO_PIC + if (stbi__pic_test(s)) return stbi__pic_load(s,x,y,comp,req_comp, ri); + #endif + + // then the formats that can end up attempting to load with just 1 or 2 + // bytes matching expectations; these are prone to false positives, so + // try them later + #ifndef STBI_NO_JPEG + if (stbi__jpeg_test(s)) return stbi__jpeg_load(s,x,y,comp,req_comp, ri); + #endif + #ifndef STBI_NO_PNM + if (stbi__pnm_test(s)) return stbi__pnm_load(s,x,y,comp,req_comp, ri); + #endif + + #ifndef STBI_NO_HDR + if (stbi__hdr_test(s)) { + float *hdr = stbi__hdr_load(s, x,y,comp,req_comp, ri); + return stbi__hdr_to_ldr(hdr, *x, *y, req_comp ? req_comp : *comp); + } + #endif + + #ifndef STBI_NO_TGA + // test tga last because it's a crappy test! + if (stbi__tga_test(s)) + return stbi__tga_load(s,x,y,comp,req_comp, ri); + #endif + + return stbi__errpuc("unknown image type", "Image not of any known type, or corrupt"); +} + +static stbi_uc *stbi__convert_16_to_8(stbi__uint16 *orig, int w, int h, int channels) +{ + int i; + int img_len = w * h * channels; + stbi_uc *reduced; + + reduced = (stbi_uc *) stbi__malloc(img_len); + if (reduced == NULL) return stbi__errpuc("outofmem", "Out of memory"); + + for (i = 0; i < img_len; ++i) + reduced[i] = (stbi_uc)((orig[i] >> 8) & 0xFF); // top half of each byte is sufficient approx of 16->8 bit scaling + + STBI_FREE(orig); + return reduced; +} + +static stbi__uint16 *stbi__convert_8_to_16(stbi_uc *orig, int w, int h, int channels) +{ + int i; + int img_len = w * h * channels; + stbi__uint16 *enlarged; + + enlarged = (stbi__uint16 *) stbi__malloc(img_len*2); + if (enlarged == NULL) return (stbi__uint16 *) stbi__errpuc("outofmem", "Out of memory"); + + for (i = 0; i < img_len; ++i) + enlarged[i] = (stbi__uint16)((orig[i] << 8) + orig[i]); // replicate to high and low byte, maps 0->0, 255->0xffff + + STBI_FREE(orig); + return enlarged; +} + +static void stbi__vertical_flip(void *image, int w, int h, int bytes_per_pixel) +{ + int row; + size_t bytes_per_row = (size_t)w * bytes_per_pixel; + stbi_uc temp[2048]; + stbi_uc *bytes = (stbi_uc *)image; + + for (row = 0; row < (h>>1); row++) { + stbi_uc *row0 = bytes + row*bytes_per_row; + stbi_uc *row1 = bytes + (h - row - 1)*bytes_per_row; + // swap row0 with row1 + size_t bytes_left = bytes_per_row; + while (bytes_left) { + size_t bytes_copy = (bytes_left < sizeof(temp)) ? bytes_left : sizeof(temp); + memcpy(temp, row0, bytes_copy); + memcpy(row0, row1, bytes_copy); + memcpy(row1, temp, bytes_copy); + row0 += bytes_copy; + row1 += bytes_copy; + bytes_left -= bytes_copy; + } + } +} + +#ifndef STBI_NO_GIF +static void stbi__vertical_flip_slices(void *image, int w, int h, int z, int bytes_per_pixel) +{ + int slice; + int slice_size = w * h * bytes_per_pixel; + + stbi_uc *bytes = (stbi_uc *)image; + for (slice = 0; slice < z; ++slice) { + stbi__vertical_flip(bytes, w, h, bytes_per_pixel); + bytes += slice_size; + } +} +#endif + +static unsigned char *stbi__load_and_postprocess_8bit(stbi__context *s, int *x, int *y, int *comp, int req_comp) +{ + stbi__result_info ri; + void *result = stbi__load_main(s, x, y, comp, req_comp, &ri, 8); + + if (result == NULL) + return NULL; + + // it is the responsibility of the loaders to make sure we get either 8 or 16 bit. + STBI_ASSERT(ri.bits_per_channel == 8 || ri.bits_per_channel == 16); + + if (ri.bits_per_channel != 8) { + result = stbi__convert_16_to_8((stbi__uint16 *) result, *x, *y, req_comp == 0 ? *comp : req_comp); + ri.bits_per_channel = 8; + } + + // @TODO: move stbi__convert_format to here + + if (stbi__vertically_flip_on_load) { + int channels = req_comp ? req_comp : *comp; + stbi__vertical_flip(result, *x, *y, channels * sizeof(stbi_uc)); + } + + return (unsigned char *) result; +} + +static stbi__uint16 *stbi__load_and_postprocess_16bit(stbi__context *s, int *x, int *y, int *comp, int req_comp) +{ + stbi__result_info ri; + void *result = stbi__load_main(s, x, y, comp, req_comp, &ri, 16); + + if (result == NULL) + return NULL; + + // it is the responsibility of the loaders to make sure we get either 8 or 16 bit. + STBI_ASSERT(ri.bits_per_channel == 8 || ri.bits_per_channel == 16); + + if (ri.bits_per_channel != 16) { + result = stbi__convert_8_to_16((stbi_uc *) result, *x, *y, req_comp == 0 ? *comp : req_comp); + ri.bits_per_channel = 16; + } + + // @TODO: move stbi__convert_format16 to here + // @TODO: special case RGB-to-Y (and RGBA-to-YA) for 8-bit-to-16-bit case to keep more precision + + if (stbi__vertically_flip_on_load) { + int channels = req_comp ? req_comp : *comp; + stbi__vertical_flip(result, *x, *y, channels * sizeof(stbi__uint16)); + } + + return (stbi__uint16 *) result; +} + +#if !defined(STBI_NO_HDR) && !defined(STBI_NO_LINEAR) +static void stbi__float_postprocess(float *result, int *x, int *y, int *comp, int req_comp) +{ + if (stbi__vertically_flip_on_load && result != NULL) { + int channels = req_comp ? req_comp : *comp; + stbi__vertical_flip(result, *x, *y, channels * sizeof(float)); + } +} +#endif + +#ifndef STBI_NO_STDIO + +#if defined(_WIN32) && defined(STBI_WINDOWS_UTF8) +STBI_EXTERN __declspec(dllimport) int __stdcall MultiByteToWideChar(unsigned int cp, unsigned long flags, const char *str, int cbmb, wchar_t *widestr, int cchwide); +STBI_EXTERN __declspec(dllimport) int __stdcall WideCharToMultiByte(unsigned int cp, unsigned long flags, const wchar_t *widestr, int cchwide, char *str, int cbmb, const char *defchar, int *used_default); +#endif + +#if defined(_WIN32) && defined(STBI_WINDOWS_UTF8) +STBIDEF int stbi_convert_wchar_to_utf8(char *buffer, size_t bufferlen, const wchar_t* input) +{ + return WideCharToMultiByte(65001 /* UTF8 */, 0, input, -1, buffer, (int) bufferlen, NULL, NULL); +} +#endif + +static FILE *stbi__fopen(char const *filename, char const *mode) +{ + FILE *f; +#if defined(_WIN32) && defined(STBI_WINDOWS_UTF8) + wchar_t wMode[64]; + wchar_t wFilename[1024]; + if (0 == MultiByteToWideChar(65001 /* UTF8 */, 0, filename, -1, wFilename, sizeof(wFilename)/sizeof(*wFilename))) + return 0; + + if (0 == MultiByteToWideChar(65001 /* UTF8 */, 0, mode, -1, wMode, sizeof(wMode)/sizeof(*wMode))) + return 0; + +#if defined(_MSC_VER) && _MSC_VER >= 1400 + if (0 != _wfopen_s(&f, wFilename, wMode)) + f = 0; +#else + f = _wfopen(wFilename, wMode); +#endif + +#elif defined(_MSC_VER) && _MSC_VER >= 1400 + if (0 != fopen_s(&f, filename, mode)) + f=0; +#else + f = fopen(filename, mode); +#endif + return f; +} + + +STBIDEF stbi_uc *stbi_load(char const *filename, int *x, int *y, int *comp, int req_comp) +{ + FILE *f = stbi__fopen(filename, "rb"); + unsigned char *result; + if (!f) return stbi__errpuc("can't fopen", "Unable to open file"); + result = stbi_load_from_file(f,x,y,comp,req_comp); + fclose(f); + return result; +} + +STBIDEF stbi_uc *stbi_load_from_file(FILE *f, int *x, int *y, int *comp, int req_comp) +{ + unsigned char *result; + stbi__context s; + stbi__start_file(&s,f); + result = stbi__load_and_postprocess_8bit(&s,x,y,comp,req_comp); + if (result) { + // need to 'unget' all the characters in the IO buffer + fseek(f, - (int) (s.img_buffer_end - s.img_buffer), SEEK_CUR); + } + return result; +} + +STBIDEF stbi__uint16 *stbi_load_from_file_16(FILE *f, int *x, int *y, int *comp, int req_comp) +{ + stbi__uint16 *result; + stbi__context s; + stbi__start_file(&s,f); + result = stbi__load_and_postprocess_16bit(&s,x,y,comp,req_comp); + if (result) { + // need to 'unget' all the characters in the IO buffer + fseek(f, - (int) (s.img_buffer_end - s.img_buffer), SEEK_CUR); + } + return result; +} + +STBIDEF stbi_us *stbi_load_16(char const *filename, int *x, int *y, int *comp, int req_comp) +{ + FILE *f = stbi__fopen(filename, "rb"); + stbi__uint16 *result; + if (!f) return (stbi_us *) stbi__errpuc("can't fopen", "Unable to open file"); + result = stbi_load_from_file_16(f,x,y,comp,req_comp); + fclose(f); + return result; +} + + +#endif //!STBI_NO_STDIO + +STBIDEF stbi_us *stbi_load_16_from_memory(stbi_uc const *buffer, int len, int *x, int *y, int *channels_in_file, int desired_channels) +{ + stbi__context s; + stbi__start_mem(&s,buffer,len); + return stbi__load_and_postprocess_16bit(&s,x,y,channels_in_file,desired_channels); +} + +STBIDEF stbi_us *stbi_load_16_from_callbacks(stbi_io_callbacks const *clbk, void *user, int *x, int *y, int *channels_in_file, int desired_channels) +{ + stbi__context s; + stbi__start_callbacks(&s, (stbi_io_callbacks *)clbk, user); + return stbi__load_and_postprocess_16bit(&s,x,y,channels_in_file,desired_channels); +} + +STBIDEF stbi_uc *stbi_load_from_memory(stbi_uc const *buffer, int len, int *x, int *y, int *comp, int req_comp) +{ + stbi__context s; + stbi__start_mem(&s,buffer,len); + return stbi__load_and_postprocess_8bit(&s,x,y,comp,req_comp); +} + +STBIDEF stbi_uc *stbi_load_from_callbacks(stbi_io_callbacks const *clbk, void *user, int *x, int *y, int *comp, int req_comp) +{ + stbi__context s; + stbi__start_callbacks(&s, (stbi_io_callbacks *) clbk, user); + return stbi__load_and_postprocess_8bit(&s,x,y,comp,req_comp); +} + +#ifndef STBI_NO_GIF +STBIDEF stbi_uc *stbi_load_gif_from_memory(stbi_uc const *buffer, int len, int **delays, int *x, int *y, int *z, int *comp, int req_comp) +{ + unsigned char *result; + stbi__context s; + stbi__start_mem(&s,buffer,len); + + result = (unsigned char*) stbi__load_gif_main(&s, delays, x, y, z, comp, req_comp); + if (stbi__vertically_flip_on_load) { + stbi__vertical_flip_slices( result, *x, *y, *z, *comp ); + } + + return result; +} +#endif + +#ifndef STBI_NO_LINEAR +static float *stbi__loadf_main(stbi__context *s, int *x, int *y, int *comp, int req_comp) +{ + unsigned char *data; + #ifndef STBI_NO_HDR + if (stbi__hdr_test(s)) { + stbi__result_info ri; + float *hdr_data = stbi__hdr_load(s,x,y,comp,req_comp, &ri); + if (hdr_data) + stbi__float_postprocess(hdr_data,x,y,comp,req_comp); + return hdr_data; + } + #endif + data = stbi__load_and_postprocess_8bit(s, x, y, comp, req_comp); + if (data) + return stbi__ldr_to_hdr(data, *x, *y, req_comp ? req_comp : *comp); + return stbi__errpf("unknown image type", "Image not of any known type, or corrupt"); +} + +STBIDEF float *stbi_loadf_from_memory(stbi_uc const *buffer, int len, int *x, int *y, int *comp, int req_comp) +{ + stbi__context s; + stbi__start_mem(&s,buffer,len); + return stbi__loadf_main(&s,x,y,comp,req_comp); +} + +STBIDEF float *stbi_loadf_from_callbacks(stbi_io_callbacks const *clbk, void *user, int *x, int *y, int *comp, int req_comp) +{ + stbi__context s; + stbi__start_callbacks(&s, (stbi_io_callbacks *) clbk, user); + return stbi__loadf_main(&s,x,y,comp,req_comp); +} + +#ifndef STBI_NO_STDIO +STBIDEF float *stbi_loadf(char const *filename, int *x, int *y, int *comp, int req_comp) +{ + float *result; + FILE *f = stbi__fopen(filename, "rb"); + if (!f) return stbi__errpf("can't fopen", "Unable to open file"); + result = stbi_loadf_from_file(f,x,y,comp,req_comp); + fclose(f); + return result; +} + +STBIDEF float *stbi_loadf_from_file(FILE *f, int *x, int *y, int *comp, int req_comp) +{ + stbi__context s; + stbi__start_file(&s,f); + return stbi__loadf_main(&s,x,y,comp,req_comp); +} +#endif // !STBI_NO_STDIO + +#endif // !STBI_NO_LINEAR + +// these is-hdr-or-not is defined independent of whether STBI_NO_LINEAR is +// defined, for API simplicity; if STBI_NO_LINEAR is defined, it always +// reports false! + +STBIDEF int stbi_is_hdr_from_memory(stbi_uc const *buffer, int len) +{ + #ifndef STBI_NO_HDR + stbi__context s; + stbi__start_mem(&s,buffer,len); + return stbi__hdr_test(&s); + #else + STBI_NOTUSED(buffer); + STBI_NOTUSED(len); + return 0; + #endif +} + +#ifndef STBI_NO_STDIO +STBIDEF int stbi_is_hdr (char const *filename) +{ + FILE *f = stbi__fopen(filename, "rb"); + int result=0; + if (f) { + result = stbi_is_hdr_from_file(f); + fclose(f); + } + return result; +} + +STBIDEF int stbi_is_hdr_from_file(FILE *f) +{ + #ifndef STBI_NO_HDR + long pos = ftell(f); + int res; + stbi__context s; + stbi__start_file(&s,f); + res = stbi__hdr_test(&s); + fseek(f, pos, SEEK_SET); + return res; + #else + STBI_NOTUSED(f); + return 0; + #endif +} +#endif // !STBI_NO_STDIO + +STBIDEF int stbi_is_hdr_from_callbacks(stbi_io_callbacks const *clbk, void *user) +{ + #ifndef STBI_NO_HDR + stbi__context s; + stbi__start_callbacks(&s, (stbi_io_callbacks *) clbk, user); + return stbi__hdr_test(&s); + #else + STBI_NOTUSED(clbk); + STBI_NOTUSED(user); + return 0; + #endif +} + +#ifndef STBI_NO_LINEAR +static float stbi__l2h_gamma=2.2f, stbi__l2h_scale=1.0f; + +STBIDEF void stbi_ldr_to_hdr_gamma(float gamma) { stbi__l2h_gamma = gamma; } +STBIDEF void stbi_ldr_to_hdr_scale(float scale) { stbi__l2h_scale = scale; } +#endif + +static float stbi__h2l_gamma_i=1.0f/2.2f, stbi__h2l_scale_i=1.0f; + +STBIDEF void stbi_hdr_to_ldr_gamma(float gamma) { stbi__h2l_gamma_i = 1/gamma; } +STBIDEF void stbi_hdr_to_ldr_scale(float scale) { stbi__h2l_scale_i = 1/scale; } + + +////////////////////////////////////////////////////////////////////////////// +// +// Common code used by all image loaders +// + +enum +{ + STBI__SCAN_load=0, + STBI__SCAN_type, + STBI__SCAN_header +}; + +static void stbi__refill_buffer(stbi__context *s) +{ + int n = (s->io.read)(s->io_user_data,(char*)s->buffer_start,s->buflen); + s->callback_already_read += (int) (s->img_buffer - s->img_buffer_original); + if (n == 0) { + // at end of file, treat same as if from memory, but need to handle case + // where s->img_buffer isn't pointing to safe memory, e.g. 0-byte file + s->read_from_callbacks = 0; + s->img_buffer = s->buffer_start; + s->img_buffer_end = s->buffer_start+1; + *s->img_buffer = 0; + } else { + s->img_buffer = s->buffer_start; + s->img_buffer_end = s->buffer_start + n; + } +} + +stbi_inline static stbi_uc stbi__get8(stbi__context *s) +{ + if (s->img_buffer < s->img_buffer_end) + return *s->img_buffer++; + if (s->read_from_callbacks) { + stbi__refill_buffer(s); + return *s->img_buffer++; + } + return 0; +} + +#if defined(STBI_NO_JPEG) && defined(STBI_NO_HDR) && defined(STBI_NO_PIC) && defined(STBI_NO_PNM) +// nothing +#else +stbi_inline static int stbi__at_eof(stbi__context *s) +{ + if (s->io.read) { + if (!(s->io.eof)(s->io_user_data)) return 0; + // if feof() is true, check if buffer = end + // special case: we've only got the special 0 character at the end + if (s->read_from_callbacks == 0) return 1; + } + + return s->img_buffer >= s->img_buffer_end; +} +#endif + +#if defined(STBI_NO_JPEG) && defined(STBI_NO_PNG) && defined(STBI_NO_BMP) && defined(STBI_NO_PSD) && defined(STBI_NO_TGA) && defined(STBI_NO_GIF) && defined(STBI_NO_PIC) +// nothing +#else +static void stbi__skip(stbi__context *s, int n) +{ + if (n == 0) return; // already there! + if (n < 0) { + s->img_buffer = s->img_buffer_end; + return; + } + if (s->io.read) { + int blen = (int) (s->img_buffer_end - s->img_buffer); + if (blen < n) { + s->img_buffer = s->img_buffer_end; + (s->io.skip)(s->io_user_data, n - blen); + return; + } + } + s->img_buffer += n; +} +#endif + +#if defined(STBI_NO_PNG) && defined(STBI_NO_TGA) && defined(STBI_NO_HDR) && defined(STBI_NO_PNM) +// nothing +#else +static int stbi__getn(stbi__context *s, stbi_uc *buffer, int n) +{ + if (s->io.read) { + int blen = (int) (s->img_buffer_end - s->img_buffer); + if (blen < n) { + int res, count; + + memcpy(buffer, s->img_buffer, blen); + + count = (s->io.read)(s->io_user_data, (char*) buffer + blen, n - blen); + res = (count == (n-blen)); + s->img_buffer = s->img_buffer_end; + return res; + } + } + + if (s->img_buffer+n <= s->img_buffer_end) { + memcpy(buffer, s->img_buffer, n); + s->img_buffer += n; + return 1; + } else + return 0; +} +#endif + +#if defined(STBI_NO_JPEG) && defined(STBI_NO_PNG) && defined(STBI_NO_PSD) && defined(STBI_NO_PIC) +// nothing +#else +static int stbi__get16be(stbi__context *s) +{ + int z = stbi__get8(s); + return (z << 8) + stbi__get8(s); +} +#endif + +#if defined(STBI_NO_PNG) && defined(STBI_NO_PSD) && defined(STBI_NO_PIC) +// nothing +#else +static stbi__uint32 stbi__get32be(stbi__context *s) +{ + stbi__uint32 z = stbi__get16be(s); + return (z << 16) + stbi__get16be(s); +} +#endif + +#if defined(STBI_NO_BMP) && defined(STBI_NO_TGA) && defined(STBI_NO_GIF) +// nothing +#else +static int stbi__get16le(stbi__context *s) +{ + int z = stbi__get8(s); + return z + (stbi__get8(s) << 8); +} +#endif + +#ifndef STBI_NO_BMP +static stbi__uint32 stbi__get32le(stbi__context *s) +{ + stbi__uint32 z = stbi__get16le(s); + z += (stbi__uint32)stbi__get16le(s) << 16; + return z; +} +#endif + +#define STBI__BYTECAST(x) ((stbi_uc) ((x) & 255)) // truncate int to byte without warnings + +#if defined(STBI_NO_JPEG) && defined(STBI_NO_PNG) && defined(STBI_NO_BMP) && defined(STBI_NO_PSD) && defined(STBI_NO_TGA) && defined(STBI_NO_GIF) && defined(STBI_NO_PIC) && defined(STBI_NO_PNM) +// nothing +#else +////////////////////////////////////////////////////////////////////////////// +// +// generic converter from built-in img_n to req_comp +// individual types do this automatically as much as possible (e.g. jpeg +// does all cases internally since it needs to colorspace convert anyway, +// and it never has alpha, so very few cases ). png can automatically +// interleave an alpha=255 channel, but falls back to this for other cases +// +// assume data buffer is malloced, so malloc a new one and free that one +// only failure mode is malloc failing + +static stbi_uc stbi__compute_y(int r, int g, int b) +{ + return (stbi_uc) (((r*77) + (g*150) + (29*b)) >> 8); +} +#endif + +#if defined(STBI_NO_PNG) && defined(STBI_NO_BMP) && defined(STBI_NO_PSD) && defined(STBI_NO_TGA) && defined(STBI_NO_GIF) && defined(STBI_NO_PIC) && defined(STBI_NO_PNM) +// nothing +#else +static unsigned char *stbi__convert_format(unsigned char *data, int img_n, int req_comp, unsigned int x, unsigned int y) +{ + int i,j; + unsigned char *good; + + if (req_comp == img_n) return data; + STBI_ASSERT(req_comp >= 1 && req_comp <= 4); + + good = (unsigned char *) stbi__malloc_mad3(req_comp, x, y, 0); + if (good == NULL) { + STBI_FREE(data); + return stbi__errpuc("outofmem", "Out of memory"); + } + + for (j=0; j < (int) y; ++j) { + unsigned char *src = data + j * x * img_n ; + unsigned char *dest = good + j * x * req_comp; + + #define STBI__COMBO(a,b) ((a)*8+(b)) + #define STBI__CASE(a,b) case STBI__COMBO(a,b): for(i=x-1; i >= 0; --i, src += a, dest += b) + // convert source image with img_n components to one with req_comp components; + // avoid switch per pixel, so use switch per scanline and massive macros + switch (STBI__COMBO(img_n, req_comp)) { + STBI__CASE(1,2) { dest[0]=src[0]; dest[1]=255; } break; + STBI__CASE(1,3) { dest[0]=dest[1]=dest[2]=src[0]; } break; + STBI__CASE(1,4) { dest[0]=dest[1]=dest[2]=src[0]; dest[3]=255; } break; + STBI__CASE(2,1) { dest[0]=src[0]; } break; + STBI__CASE(2,3) { dest[0]=dest[1]=dest[2]=src[0]; } break; + STBI__CASE(2,4) { dest[0]=dest[1]=dest[2]=src[0]; dest[3]=src[1]; } break; + STBI__CASE(3,4) { dest[0]=src[0];dest[1]=src[1];dest[2]=src[2];dest[3]=255; } break; + STBI__CASE(3,1) { dest[0]=stbi__compute_y(src[0],src[1],src[2]); } break; + STBI__CASE(3,2) { dest[0]=stbi__compute_y(src[0],src[1],src[2]); dest[1] = 255; } break; + STBI__CASE(4,1) { dest[0]=stbi__compute_y(src[0],src[1],src[2]); } break; + STBI__CASE(4,2) { dest[0]=stbi__compute_y(src[0],src[1],src[2]); dest[1] = src[3]; } break; + STBI__CASE(4,3) { dest[0]=src[0];dest[1]=src[1];dest[2]=src[2]; } break; + default: STBI_ASSERT(0); STBI_FREE(data); STBI_FREE(good); return stbi__errpuc("unsupported", "Unsupported format conversion"); + } + #undef STBI__CASE + } + + STBI_FREE(data); + return good; +} +#endif + +#if defined(STBI_NO_PNG) && defined(STBI_NO_PSD) +// nothing +#else +static stbi__uint16 stbi__compute_y_16(int r, int g, int b) +{ + return (stbi__uint16) (((r*77) + (g*150) + (29*b)) >> 8); +} +#endif + +#if defined(STBI_NO_PNG) && defined(STBI_NO_PSD) +// nothing +#else +static stbi__uint16 *stbi__convert_format16(stbi__uint16 *data, int img_n, int req_comp, unsigned int x, unsigned int y) +{ + int i,j; + stbi__uint16 *good; + + if (req_comp == img_n) return data; + STBI_ASSERT(req_comp >= 1 && req_comp <= 4); + + good = (stbi__uint16 *) stbi__malloc(req_comp * x * y * 2); + if (good == NULL) { + STBI_FREE(data); + return (stbi__uint16 *) stbi__errpuc("outofmem", "Out of memory"); + } + + for (j=0; j < (int) y; ++j) { + stbi__uint16 *src = data + j * x * img_n ; + stbi__uint16 *dest = good + j * x * req_comp; + + #define STBI__COMBO(a,b) ((a)*8+(b)) + #define STBI__CASE(a,b) case STBI__COMBO(a,b): for(i=x-1; i >= 0; --i, src += a, dest += b) + // convert source image with img_n components to one with req_comp components; + // avoid switch per pixel, so use switch per scanline and massive macros + switch (STBI__COMBO(img_n, req_comp)) { + STBI__CASE(1,2) { dest[0]=src[0]; dest[1]=0xffff; } break; + STBI__CASE(1,3) { dest[0]=dest[1]=dest[2]=src[0]; } break; + STBI__CASE(1,4) { dest[0]=dest[1]=dest[2]=src[0]; dest[3]=0xffff; } break; + STBI__CASE(2,1) { dest[0]=src[0]; } break; + STBI__CASE(2,3) { dest[0]=dest[1]=dest[2]=src[0]; } break; + STBI__CASE(2,4) { dest[0]=dest[1]=dest[2]=src[0]; dest[3]=src[1]; } break; + STBI__CASE(3,4) { dest[0]=src[0];dest[1]=src[1];dest[2]=src[2];dest[3]=0xffff; } break; + STBI__CASE(3,1) { dest[0]=stbi__compute_y_16(src[0],src[1],src[2]); } break; + STBI__CASE(3,2) { dest[0]=stbi__compute_y_16(src[0],src[1],src[2]); dest[1] = 0xffff; } break; + STBI__CASE(4,1) { dest[0]=stbi__compute_y_16(src[0],src[1],src[2]); } break; + STBI__CASE(4,2) { dest[0]=stbi__compute_y_16(src[0],src[1],src[2]); dest[1] = src[3]; } break; + STBI__CASE(4,3) { dest[0]=src[0];dest[1]=src[1];dest[2]=src[2]; } break; + default: STBI_ASSERT(0); STBI_FREE(data); STBI_FREE(good); return (stbi__uint16*) stbi__errpuc("unsupported", "Unsupported format conversion"); + } + #undef STBI__CASE + } + + STBI_FREE(data); + return good; +} +#endif + +#ifndef STBI_NO_LINEAR +static float *stbi__ldr_to_hdr(stbi_uc *data, int x, int y, int comp) +{ + int i,k,n; + float *output; + if (!data) return NULL; + output = (float *) stbi__malloc_mad4(x, y, comp, sizeof(float), 0); + if (output == NULL) { STBI_FREE(data); return stbi__errpf("outofmem", "Out of memory"); } + // compute number of non-alpha components + if (comp & 1) n = comp; else n = comp-1; + for (i=0; i < x*y; ++i) { + for (k=0; k < n; ++k) { + output[i*comp + k] = (float) (pow(data[i*comp+k]/255.0f, stbi__l2h_gamma) * stbi__l2h_scale); + } + } + if (n < comp) { + for (i=0; i < x*y; ++i) { + output[i*comp + n] = data[i*comp + n]/255.0f; + } + } + STBI_FREE(data); + return output; +} +#endif + +#ifndef STBI_NO_HDR +#define stbi__float2int(x) ((int) (x)) +static stbi_uc *stbi__hdr_to_ldr(float *data, int x, int y, int comp) +{ + int i,k,n; + stbi_uc *output; + if (!data) return NULL; + output = (stbi_uc *) stbi__malloc_mad3(x, y, comp, 0); + if (output == NULL) { STBI_FREE(data); return stbi__errpuc("outofmem", "Out of memory"); } + // compute number of non-alpha components + if (comp & 1) n = comp; else n = comp-1; + for (i=0; i < x*y; ++i) { + for (k=0; k < n; ++k) { + float z = (float) pow(data[i*comp+k]*stbi__h2l_scale_i, stbi__h2l_gamma_i) * 255 + 0.5f; + if (z < 0) z = 0; + if (z > 255) z = 255; + output[i*comp + k] = (stbi_uc) stbi__float2int(z); + } + if (k < comp) { + float z = data[i*comp+k] * 255 + 0.5f; + if (z < 0) z = 0; + if (z > 255) z = 255; + output[i*comp + k] = (stbi_uc) stbi__float2int(z); + } + } + STBI_FREE(data); + return output; +} +#endif + +////////////////////////////////////////////////////////////////////////////// +// +// "baseline" JPEG/JFIF decoder +// +// simple implementation +// - doesn't support delayed output of y-dimension +// - simple interface (only one output format: 8-bit interleaved RGB) +// - doesn't try to recover corrupt jpegs +// - doesn't allow partial loading, loading multiple at once +// - still fast on x86 (copying globals into locals doesn't help x86) +// - allocates lots of intermediate memory (full size of all components) +// - non-interleaved case requires this anyway +// - allows good upsampling (see next) +// high-quality +// - upsampled channels are bilinearly interpolated, even across blocks +// - quality integer IDCT derived from IJG's 'slow' +// performance +// - fast huffman; reasonable integer IDCT +// - some SIMD kernels for common paths on targets with SSE2/NEON +// - uses a lot of intermediate memory, could cache poorly + +#ifndef STBI_NO_JPEG + +// huffman decoding acceleration +#define FAST_BITS 9 // larger handles more cases; smaller stomps less cache + +typedef struct +{ + stbi_uc fast[1 << FAST_BITS]; + // weirdly, repacking this into AoS is a 10% speed loss, instead of a win + stbi__uint16 code[256]; + stbi_uc values[256]; + stbi_uc size[257]; + unsigned int maxcode[18]; + int delta[17]; // old 'firstsymbol' - old 'firstcode' +} stbi__huffman; + +typedef struct +{ + stbi__context *s; + stbi__huffman huff_dc[4]; + stbi__huffman huff_ac[4]; + stbi__uint16 dequant[4][64]; + stbi__int16 fast_ac[4][1 << FAST_BITS]; + +// sizes for components, interleaved MCUs + int img_h_max, img_v_max; + int img_mcu_x, img_mcu_y; + int img_mcu_w, img_mcu_h; + +// definition of jpeg image component + struct + { + int id; + int h,v; + int tq; + int hd,ha; + int dc_pred; + + int x,y,w2,h2; + stbi_uc *data; + void *raw_data, *raw_coeff; + stbi_uc *linebuf; + short *coeff; // progressive only + int coeff_w, coeff_h; // number of 8x8 coefficient blocks + } img_comp[4]; + + stbi__uint32 code_buffer; // jpeg entropy-coded buffer + int code_bits; // number of valid bits + unsigned char marker; // marker seen while filling entropy buffer + int nomore; // flag if we saw a marker so must stop + + int progressive; + int spec_start; + int spec_end; + int succ_high; + int succ_low; + int eob_run; + int jfif; + int app14_color_transform; // Adobe APP14 tag + int rgb; + + int scan_n, order[4]; + int restart_interval, todo; + +// kernels + void (*idct_block_kernel)(stbi_uc *out, int out_stride, short data[64]); + void (*YCbCr_to_RGB_kernel)(stbi_uc *out, const stbi_uc *y, const stbi_uc *pcb, const stbi_uc *pcr, int count, int step); + stbi_uc *(*resample_row_hv_2_kernel)(stbi_uc *out, stbi_uc *in_near, stbi_uc *in_far, int w, int hs); +} stbi__jpeg; + +static int stbi__build_huffman(stbi__huffman *h, int *count) +{ + int i,j,k=0; + unsigned int code; + // build size list for each symbol (from JPEG spec) + for (i=0; i < 16; ++i) { + for (j=0; j < count[i]; ++j) { + h->size[k++] = (stbi_uc) (i+1); + if(k >= 257) return stbi__err("bad size list","Corrupt JPEG"); + } + } + h->size[k] = 0; + + // compute actual symbols (from jpeg spec) + code = 0; + k = 0; + for(j=1; j <= 16; ++j) { + // compute delta to add to code to compute symbol id + h->delta[j] = k - code; + if (h->size[k] == j) { + while (h->size[k] == j) + h->code[k++] = (stbi__uint16) (code++); + if (code-1 >= (1u << j)) return stbi__err("bad code lengths","Corrupt JPEG"); + } + // compute largest code + 1 for this size, preshifted as needed later + h->maxcode[j] = code << (16-j); + code <<= 1; + } + h->maxcode[j] = 0xffffffff; + + // build non-spec acceleration table; 255 is flag for not-accelerated + memset(h->fast, 255, 1 << FAST_BITS); + for (i=0; i < k; ++i) { + int s = h->size[i]; + if (s <= FAST_BITS) { + int c = h->code[i] << (FAST_BITS-s); + int m = 1 << (FAST_BITS-s); + for (j=0; j < m; ++j) { + h->fast[c+j] = (stbi_uc) i; + } + } + } + return 1; +} + +// build a table that decodes both magnitude and value of small ACs in +// one go. +static void stbi__build_fast_ac(stbi__int16 *fast_ac, stbi__huffman *h) +{ + int i; + for (i=0; i < (1 << FAST_BITS); ++i) { + stbi_uc fast = h->fast[i]; + fast_ac[i] = 0; + if (fast < 255) { + int rs = h->values[fast]; + int run = (rs >> 4) & 15; + int magbits = rs & 15; + int len = h->size[fast]; + + if (magbits && len + magbits <= FAST_BITS) { + // magnitude code followed by receive_extend code + int k = ((i << len) & ((1 << FAST_BITS) - 1)) >> (FAST_BITS - magbits); + int m = 1 << (magbits - 1); + if (k < m) k += (~0U << magbits) + 1; + // if the result is small enough, we can fit it in fast_ac table + if (k >= -128 && k <= 127) + fast_ac[i] = (stbi__int16) ((k * 256) + (run * 16) + (len + magbits)); + } + } + } +} + +static void stbi__grow_buffer_unsafe(stbi__jpeg *j) +{ + do { + unsigned int b = j->nomore ? 0 : stbi__get8(j->s); + if (b == 0xff) { + int c = stbi__get8(j->s); + while (c == 0xff) c = stbi__get8(j->s); // consume fill bytes + if (c != 0) { + j->marker = (unsigned char) c; + j->nomore = 1; + return; + } + } + j->code_buffer |= b << (24 - j->code_bits); + j->code_bits += 8; + } while (j->code_bits <= 24); +} + +// (1 << n) - 1 +static const stbi__uint32 stbi__bmask[17]={0,1,3,7,15,31,63,127,255,511,1023,2047,4095,8191,16383,32767,65535}; + +// decode a jpeg huffman value from the bitstream +stbi_inline static int stbi__jpeg_huff_decode(stbi__jpeg *j, stbi__huffman *h) +{ + unsigned int temp; + int c,k; + + if (j->code_bits < 16) stbi__grow_buffer_unsafe(j); + + // look at the top FAST_BITS and determine what symbol ID it is, + // if the code is <= FAST_BITS + c = (j->code_buffer >> (32 - FAST_BITS)) & ((1 << FAST_BITS)-1); + k = h->fast[c]; + if (k < 255) { + int s = h->size[k]; + if (s > j->code_bits) + return -1; + j->code_buffer <<= s; + j->code_bits -= s; + return h->values[k]; + } + + // naive test is to shift the code_buffer down so k bits are + // valid, then test against maxcode. To speed this up, we've + // preshifted maxcode left so that it has (16-k) 0s at the + // end; in other words, regardless of the number of bits, it + // wants to be compared against something shifted to have 16; + // that way we don't need to shift inside the loop. + temp = j->code_buffer >> 16; + for (k=FAST_BITS+1 ; ; ++k) + if (temp < h->maxcode[k]) + break; + if (k == 17) { + // error! code not found + j->code_bits -= 16; + return -1; + } + + if (k > j->code_bits) + return -1; + + // convert the huffman code to the symbol id + c = ((j->code_buffer >> (32 - k)) & stbi__bmask[k]) + h->delta[k]; + if(c < 0 || c >= 256) // symbol id out of bounds! + return -1; + STBI_ASSERT((((j->code_buffer) >> (32 - h->size[c])) & stbi__bmask[h->size[c]]) == h->code[c]); + + // convert the id to a symbol + j->code_bits -= k; + j->code_buffer <<= k; + return h->values[c]; +} + +// bias[n] = (-1<code_bits < n) stbi__grow_buffer_unsafe(j); + if (j->code_bits < n) return 0; // ran out of bits from stream, return 0s intead of continuing + + sgn = j->code_buffer >> 31; // sign bit always in MSB; 0 if MSB clear (positive), 1 if MSB set (negative) + k = stbi_lrot(j->code_buffer, n); + j->code_buffer = k & ~stbi__bmask[n]; + k &= stbi__bmask[n]; + j->code_bits -= n; + return k + (stbi__jbias[n] & (sgn - 1)); +} + +// get some unsigned bits +stbi_inline static int stbi__jpeg_get_bits(stbi__jpeg *j, int n) +{ + unsigned int k; + if (j->code_bits < n) stbi__grow_buffer_unsafe(j); + if (j->code_bits < n) return 0; // ran out of bits from stream, return 0s intead of continuing + k = stbi_lrot(j->code_buffer, n); + j->code_buffer = k & ~stbi__bmask[n]; + k &= stbi__bmask[n]; + j->code_bits -= n; + return k; +} + +stbi_inline static int stbi__jpeg_get_bit(stbi__jpeg *j) +{ + unsigned int k; + if (j->code_bits < 1) stbi__grow_buffer_unsafe(j); + if (j->code_bits < 1) return 0; // ran out of bits from stream, return 0s intead of continuing + k = j->code_buffer; + j->code_buffer <<= 1; + --j->code_bits; + return k & 0x80000000; +} + +// given a value that's at position X in the zigzag stream, +// where does it appear in the 8x8 matrix coded as row-major? +static const stbi_uc stbi__jpeg_dezigzag[64+15] = +{ + 0, 1, 8, 16, 9, 2, 3, 10, + 17, 24, 32, 25, 18, 11, 4, 5, + 12, 19, 26, 33, 40, 48, 41, 34, + 27, 20, 13, 6, 7, 14, 21, 28, + 35, 42, 49, 56, 57, 50, 43, 36, + 29, 22, 15, 23, 30, 37, 44, 51, + 58, 59, 52, 45, 38, 31, 39, 46, + 53, 60, 61, 54, 47, 55, 62, 63, + // let corrupt input sample past end + 63, 63, 63, 63, 63, 63, 63, 63, + 63, 63, 63, 63, 63, 63, 63 +}; + +// decode one 64-entry block-- +static int stbi__jpeg_decode_block(stbi__jpeg *j, short data[64], stbi__huffman *hdc, stbi__huffman *hac, stbi__int16 *fac, int b, stbi__uint16 *dequant) +{ + int diff,dc,k; + int t; + + if (j->code_bits < 16) stbi__grow_buffer_unsafe(j); + t = stbi__jpeg_huff_decode(j, hdc); + if (t < 0 || t > 15) return stbi__err("bad huffman code","Corrupt JPEG"); + + // 0 all the ac values now so we can do it 32-bits at a time + memset(data,0,64*sizeof(data[0])); + + diff = t ? stbi__extend_receive(j, t) : 0; + if (!stbi__addints_valid(j->img_comp[b].dc_pred, diff)) return stbi__err("bad delta","Corrupt JPEG"); + dc = j->img_comp[b].dc_pred + diff; + j->img_comp[b].dc_pred = dc; + if (!stbi__mul2shorts_valid(dc, dequant[0])) return stbi__err("can't merge dc and ac", "Corrupt JPEG"); + data[0] = (short) (dc * dequant[0]); + + // decode AC components, see JPEG spec + k = 1; + do { + unsigned int zig; + int c,r,s; + if (j->code_bits < 16) stbi__grow_buffer_unsafe(j); + c = (j->code_buffer >> (32 - FAST_BITS)) & ((1 << FAST_BITS)-1); + r = fac[c]; + if (r) { // fast-AC path + k += (r >> 4) & 15; // run + s = r & 15; // combined length + if (s > j->code_bits) return stbi__err("bad huffman code", "Combined length longer than code bits available"); + j->code_buffer <<= s; + j->code_bits -= s; + // decode into unzigzag'd location + zig = stbi__jpeg_dezigzag[k++]; + data[zig] = (short) ((r >> 8) * dequant[zig]); + } else { + int rs = stbi__jpeg_huff_decode(j, hac); + if (rs < 0) return stbi__err("bad huffman code","Corrupt JPEG"); + s = rs & 15; + r = rs >> 4; + if (s == 0) { + if (rs != 0xf0) break; // end block + k += 16; + } else { + k += r; + // decode into unzigzag'd location + zig = stbi__jpeg_dezigzag[k++]; + data[zig] = (short) (stbi__extend_receive(j,s) * dequant[zig]); + } + } + } while (k < 64); + return 1; +} + +static int stbi__jpeg_decode_block_prog_dc(stbi__jpeg *j, short data[64], stbi__huffman *hdc, int b) +{ + int diff,dc; + int t; + if (j->spec_end != 0) return stbi__err("can't merge dc and ac", "Corrupt JPEG"); + + if (j->code_bits < 16) stbi__grow_buffer_unsafe(j); + + if (j->succ_high == 0) { + // first scan for DC coefficient, must be first + memset(data,0,64*sizeof(data[0])); // 0 all the ac values now + t = stbi__jpeg_huff_decode(j, hdc); + if (t < 0 || t > 15) return stbi__err("can't merge dc and ac", "Corrupt JPEG"); + diff = t ? stbi__extend_receive(j, t) : 0; + + if (!stbi__addints_valid(j->img_comp[b].dc_pred, diff)) return stbi__err("bad delta", "Corrupt JPEG"); + dc = j->img_comp[b].dc_pred + diff; + j->img_comp[b].dc_pred = dc; + if (!stbi__mul2shorts_valid(dc, 1 << j->succ_low)) return stbi__err("can't merge dc and ac", "Corrupt JPEG"); + data[0] = (short) (dc * (1 << j->succ_low)); + } else { + // refinement scan for DC coefficient + if (stbi__jpeg_get_bit(j)) + data[0] += (short) (1 << j->succ_low); + } + return 1; +} + +// @OPTIMIZE: store non-zigzagged during the decode passes, +// and only de-zigzag when dequantizing +static int stbi__jpeg_decode_block_prog_ac(stbi__jpeg *j, short data[64], stbi__huffman *hac, stbi__int16 *fac) +{ + int k; + if (j->spec_start == 0) return stbi__err("can't merge dc and ac", "Corrupt JPEG"); + + if (j->succ_high == 0) { + int shift = j->succ_low; + + if (j->eob_run) { + --j->eob_run; + return 1; + } + + k = j->spec_start; + do { + unsigned int zig; + int c,r,s; + if (j->code_bits < 16) stbi__grow_buffer_unsafe(j); + c = (j->code_buffer >> (32 - FAST_BITS)) & ((1 << FAST_BITS)-1); + r = fac[c]; + if (r) { // fast-AC path + k += (r >> 4) & 15; // run + s = r & 15; // combined length + if (s > j->code_bits) return stbi__err("bad huffman code", "Combined length longer than code bits available"); + j->code_buffer <<= s; + j->code_bits -= s; + zig = stbi__jpeg_dezigzag[k++]; + data[zig] = (short) ((r >> 8) * (1 << shift)); + } else { + int rs = stbi__jpeg_huff_decode(j, hac); + if (rs < 0) return stbi__err("bad huffman code","Corrupt JPEG"); + s = rs & 15; + r = rs >> 4; + if (s == 0) { + if (r < 15) { + j->eob_run = (1 << r); + if (r) + j->eob_run += stbi__jpeg_get_bits(j, r); + --j->eob_run; + break; + } + k += 16; + } else { + k += r; + zig = stbi__jpeg_dezigzag[k++]; + data[zig] = (short) (stbi__extend_receive(j,s) * (1 << shift)); + } + } + } while (k <= j->spec_end); + } else { + // refinement scan for these AC coefficients + + short bit = (short) (1 << j->succ_low); + + if (j->eob_run) { + --j->eob_run; + for (k = j->spec_start; k <= j->spec_end; ++k) { + short *p = &data[stbi__jpeg_dezigzag[k]]; + if (*p != 0) + if (stbi__jpeg_get_bit(j)) + if ((*p & bit)==0) { + if (*p > 0) + *p += bit; + else + *p -= bit; + } + } + } else { + k = j->spec_start; + do { + int r,s; + int rs = stbi__jpeg_huff_decode(j, hac); // @OPTIMIZE see if we can use the fast path here, advance-by-r is so slow, eh + if (rs < 0) return stbi__err("bad huffman code","Corrupt JPEG"); + s = rs & 15; + r = rs >> 4; + if (s == 0) { + if (r < 15) { + j->eob_run = (1 << r) - 1; + if (r) + j->eob_run += stbi__jpeg_get_bits(j, r); + r = 64; // force end of block + } else { + // r=15 s=0 should write 16 0s, so we just do + // a run of 15 0s and then write s (which is 0), + // so we don't have to do anything special here + } + } else { + if (s != 1) return stbi__err("bad huffman code", "Corrupt JPEG"); + // sign bit + if (stbi__jpeg_get_bit(j)) + s = bit; + else + s = -bit; + } + + // advance by r + while (k <= j->spec_end) { + short *p = &data[stbi__jpeg_dezigzag[k++]]; + if (*p != 0) { + if (stbi__jpeg_get_bit(j)) + if ((*p & bit)==0) { + if (*p > 0) + *p += bit; + else + *p -= bit; + } + } else { + if (r == 0) { + *p = (short) s; + break; + } + --r; + } + } + } while (k <= j->spec_end); + } + } + return 1; +} + +// take a -128..127 value and stbi__clamp it and convert to 0..255 +stbi_inline static stbi_uc stbi__clamp(int x) +{ + // trick to use a single test to catch both cases + if ((unsigned int) x > 255) { + if (x < 0) return 0; + if (x > 255) return 255; + } + return (stbi_uc) x; +} + +#define stbi__f2f(x) ((int) (((x) * 4096 + 0.5))) +#define stbi__fsh(x) ((x) * 4096) + +// derived from jidctint -- DCT_ISLOW +#define STBI__IDCT_1D(s0,s1,s2,s3,s4,s5,s6,s7) \ + int t0,t1,t2,t3,p1,p2,p3,p4,p5,x0,x1,x2,x3; \ + p2 = s2; \ + p3 = s6; \ + p1 = (p2+p3) * stbi__f2f(0.5411961f); \ + t2 = p1 + p3*stbi__f2f(-1.847759065f); \ + t3 = p1 + p2*stbi__f2f( 0.765366865f); \ + p2 = s0; \ + p3 = s4; \ + t0 = stbi__fsh(p2+p3); \ + t1 = stbi__fsh(p2-p3); \ + x0 = t0+t3; \ + x3 = t0-t3; \ + x1 = t1+t2; \ + x2 = t1-t2; \ + t0 = s7; \ + t1 = s5; \ + t2 = s3; \ + t3 = s1; \ + p3 = t0+t2; \ + p4 = t1+t3; \ + p1 = t0+t3; \ + p2 = t1+t2; \ + p5 = (p3+p4)*stbi__f2f( 1.175875602f); \ + t0 = t0*stbi__f2f( 0.298631336f); \ + t1 = t1*stbi__f2f( 2.053119869f); \ + t2 = t2*stbi__f2f( 3.072711026f); \ + t3 = t3*stbi__f2f( 1.501321110f); \ + p1 = p5 + p1*stbi__f2f(-0.899976223f); \ + p2 = p5 + p2*stbi__f2f(-2.562915447f); \ + p3 = p3*stbi__f2f(-1.961570560f); \ + p4 = p4*stbi__f2f(-0.390180644f); \ + t3 += p1+p4; \ + t2 += p2+p3; \ + t1 += p2+p4; \ + t0 += p1+p3; + +static void stbi__idct_block(stbi_uc *out, int out_stride, short data[64]) +{ + int i,val[64],*v=val; + stbi_uc *o; + short *d = data; + + // columns + for (i=0; i < 8; ++i,++d, ++v) { + // if all zeroes, shortcut -- this avoids dequantizing 0s and IDCTing + if (d[ 8]==0 && d[16]==0 && d[24]==0 && d[32]==0 + && d[40]==0 && d[48]==0 && d[56]==0) { + // no shortcut 0 seconds + // (1|2|3|4|5|6|7)==0 0 seconds + // all separate -0.047 seconds + // 1 && 2|3 && 4|5 && 6|7: -0.047 seconds + int dcterm = d[0]*4; + v[0] = v[8] = v[16] = v[24] = v[32] = v[40] = v[48] = v[56] = dcterm; + } else { + STBI__IDCT_1D(d[ 0],d[ 8],d[16],d[24],d[32],d[40],d[48],d[56]) + // constants scaled things up by 1<<12; let's bring them back + // down, but keep 2 extra bits of precision + x0 += 512; x1 += 512; x2 += 512; x3 += 512; + v[ 0] = (x0+t3) >> 10; + v[56] = (x0-t3) >> 10; + v[ 8] = (x1+t2) >> 10; + v[48] = (x1-t2) >> 10; + v[16] = (x2+t1) >> 10; + v[40] = (x2-t1) >> 10; + v[24] = (x3+t0) >> 10; + v[32] = (x3-t0) >> 10; + } + } + + for (i=0, v=val, o=out; i < 8; ++i,v+=8,o+=out_stride) { + // no fast case since the first 1D IDCT spread components out + STBI__IDCT_1D(v[0],v[1],v[2],v[3],v[4],v[5],v[6],v[7]) + // constants scaled things up by 1<<12, plus we had 1<<2 from first + // loop, plus horizontal and vertical each scale by sqrt(8) so together + // we've got an extra 1<<3, so 1<<17 total we need to remove. + // so we want to round that, which means adding 0.5 * 1<<17, + // aka 65536. Also, we'll end up with -128 to 127 that we want + // to encode as 0..255 by adding 128, so we'll add that before the shift + x0 += 65536 + (128<<17); + x1 += 65536 + (128<<17); + x2 += 65536 + (128<<17); + x3 += 65536 + (128<<17); + // tried computing the shifts into temps, or'ing the temps to see + // if any were out of range, but that was slower + o[0] = stbi__clamp((x0+t3) >> 17); + o[7] = stbi__clamp((x0-t3) >> 17); + o[1] = stbi__clamp((x1+t2) >> 17); + o[6] = stbi__clamp((x1-t2) >> 17); + o[2] = stbi__clamp((x2+t1) >> 17); + o[5] = stbi__clamp((x2-t1) >> 17); + o[3] = stbi__clamp((x3+t0) >> 17); + o[4] = stbi__clamp((x3-t0) >> 17); + } +} + +#ifdef STBI_SSE2 +// sse2 integer IDCT. not the fastest possible implementation but it +// produces bit-identical results to the generic C version so it's +// fully "transparent". +static void stbi__idct_simd(stbi_uc *out, int out_stride, short data[64]) +{ + // This is constructed to match our regular (generic) integer IDCT exactly. + __m128i row0, row1, row2, row3, row4, row5, row6, row7; + __m128i tmp; + + // dot product constant: even elems=x, odd elems=y + #define dct_const(x,y) _mm_setr_epi16((x),(y),(x),(y),(x),(y),(x),(y)) + + // out(0) = c0[even]*x + c0[odd]*y (c0, x, y 16-bit, out 32-bit) + // out(1) = c1[even]*x + c1[odd]*y + #define dct_rot(out0,out1, x,y,c0,c1) \ + __m128i c0##lo = _mm_unpacklo_epi16((x),(y)); \ + __m128i c0##hi = _mm_unpackhi_epi16((x),(y)); \ + __m128i out0##_l = _mm_madd_epi16(c0##lo, c0); \ + __m128i out0##_h = _mm_madd_epi16(c0##hi, c0); \ + __m128i out1##_l = _mm_madd_epi16(c0##lo, c1); \ + __m128i out1##_h = _mm_madd_epi16(c0##hi, c1) + + // out = in << 12 (in 16-bit, out 32-bit) + #define dct_widen(out, in) \ + __m128i out##_l = _mm_srai_epi32(_mm_unpacklo_epi16(_mm_setzero_si128(), (in)), 4); \ + __m128i out##_h = _mm_srai_epi32(_mm_unpackhi_epi16(_mm_setzero_si128(), (in)), 4) + + // wide add + #define dct_wadd(out, a, b) \ + __m128i out##_l = _mm_add_epi32(a##_l, b##_l); \ + __m128i out##_h = _mm_add_epi32(a##_h, b##_h) + + // wide sub + #define dct_wsub(out, a, b) \ + __m128i out##_l = _mm_sub_epi32(a##_l, b##_l); \ + __m128i out##_h = _mm_sub_epi32(a##_h, b##_h) + + // butterfly a/b, add bias, then shift by "s" and pack + #define dct_bfly32o(out0, out1, a,b,bias,s) \ + { \ + __m128i abiased_l = _mm_add_epi32(a##_l, bias); \ + __m128i abiased_h = _mm_add_epi32(a##_h, bias); \ + dct_wadd(sum, abiased, b); \ + dct_wsub(dif, abiased, b); \ + out0 = _mm_packs_epi32(_mm_srai_epi32(sum_l, s), _mm_srai_epi32(sum_h, s)); \ + out1 = _mm_packs_epi32(_mm_srai_epi32(dif_l, s), _mm_srai_epi32(dif_h, s)); \ + } + + // 8-bit interleave step (for transposes) + #define dct_interleave8(a, b) \ + tmp = a; \ + a = _mm_unpacklo_epi8(a, b); \ + b = _mm_unpackhi_epi8(tmp, b) + + // 16-bit interleave step (for transposes) + #define dct_interleave16(a, b) \ + tmp = a; \ + a = _mm_unpacklo_epi16(a, b); \ + b = _mm_unpackhi_epi16(tmp, b) + + #define dct_pass(bias,shift) \ + { \ + /* even part */ \ + dct_rot(t2e,t3e, row2,row6, rot0_0,rot0_1); \ + __m128i sum04 = _mm_add_epi16(row0, row4); \ + __m128i dif04 = _mm_sub_epi16(row0, row4); \ + dct_widen(t0e, sum04); \ + dct_widen(t1e, dif04); \ + dct_wadd(x0, t0e, t3e); \ + dct_wsub(x3, t0e, t3e); \ + dct_wadd(x1, t1e, t2e); \ + dct_wsub(x2, t1e, t2e); \ + /* odd part */ \ + dct_rot(y0o,y2o, row7,row3, rot2_0,rot2_1); \ + dct_rot(y1o,y3o, row5,row1, rot3_0,rot3_1); \ + __m128i sum17 = _mm_add_epi16(row1, row7); \ + __m128i sum35 = _mm_add_epi16(row3, row5); \ + dct_rot(y4o,y5o, sum17,sum35, rot1_0,rot1_1); \ + dct_wadd(x4, y0o, y4o); \ + dct_wadd(x5, y1o, y5o); \ + dct_wadd(x6, y2o, y5o); \ + dct_wadd(x7, y3o, y4o); \ + dct_bfly32o(row0,row7, x0,x7,bias,shift); \ + dct_bfly32o(row1,row6, x1,x6,bias,shift); \ + dct_bfly32o(row2,row5, x2,x5,bias,shift); \ + dct_bfly32o(row3,row4, x3,x4,bias,shift); \ + } + + __m128i rot0_0 = dct_const(stbi__f2f(0.5411961f), stbi__f2f(0.5411961f) + stbi__f2f(-1.847759065f)); + __m128i rot0_1 = dct_const(stbi__f2f(0.5411961f) + stbi__f2f( 0.765366865f), stbi__f2f(0.5411961f)); + __m128i rot1_0 = dct_const(stbi__f2f(1.175875602f) + stbi__f2f(-0.899976223f), stbi__f2f(1.175875602f)); + __m128i rot1_1 = dct_const(stbi__f2f(1.175875602f), stbi__f2f(1.175875602f) + stbi__f2f(-2.562915447f)); + __m128i rot2_0 = dct_const(stbi__f2f(-1.961570560f) + stbi__f2f( 0.298631336f), stbi__f2f(-1.961570560f)); + __m128i rot2_1 = dct_const(stbi__f2f(-1.961570560f), stbi__f2f(-1.961570560f) + stbi__f2f( 3.072711026f)); + __m128i rot3_0 = dct_const(stbi__f2f(-0.390180644f) + stbi__f2f( 2.053119869f), stbi__f2f(-0.390180644f)); + __m128i rot3_1 = dct_const(stbi__f2f(-0.390180644f), stbi__f2f(-0.390180644f) + stbi__f2f( 1.501321110f)); + + // rounding biases in column/row passes, see stbi__idct_block for explanation. + __m128i bias_0 = _mm_set1_epi32(512); + __m128i bias_1 = _mm_set1_epi32(65536 + (128<<17)); + + // load + row0 = _mm_load_si128((const __m128i *) (data + 0*8)); + row1 = _mm_load_si128((const __m128i *) (data + 1*8)); + row2 = _mm_load_si128((const __m128i *) (data + 2*8)); + row3 = _mm_load_si128((const __m128i *) (data + 3*8)); + row4 = _mm_load_si128((const __m128i *) (data + 4*8)); + row5 = _mm_load_si128((const __m128i *) (data + 5*8)); + row6 = _mm_load_si128((const __m128i *) (data + 6*8)); + row7 = _mm_load_si128((const __m128i *) (data + 7*8)); + + // column pass + dct_pass(bias_0, 10); + + { + // 16bit 8x8 transpose pass 1 + dct_interleave16(row0, row4); + dct_interleave16(row1, row5); + dct_interleave16(row2, row6); + dct_interleave16(row3, row7); + + // transpose pass 2 + dct_interleave16(row0, row2); + dct_interleave16(row1, row3); + dct_interleave16(row4, row6); + dct_interleave16(row5, row7); + + // transpose pass 3 + dct_interleave16(row0, row1); + dct_interleave16(row2, row3); + dct_interleave16(row4, row5); + dct_interleave16(row6, row7); + } + + // row pass + dct_pass(bias_1, 17); + + { + // pack + __m128i p0 = _mm_packus_epi16(row0, row1); // a0a1a2a3...a7b0b1b2b3...b7 + __m128i p1 = _mm_packus_epi16(row2, row3); + __m128i p2 = _mm_packus_epi16(row4, row5); + __m128i p3 = _mm_packus_epi16(row6, row7); + + // 8bit 8x8 transpose pass 1 + dct_interleave8(p0, p2); // a0e0a1e1... + dct_interleave8(p1, p3); // c0g0c1g1... + + // transpose pass 2 + dct_interleave8(p0, p1); // a0c0e0g0... + dct_interleave8(p2, p3); // b0d0f0h0... + + // transpose pass 3 + dct_interleave8(p0, p2); // a0b0c0d0... + dct_interleave8(p1, p3); // a4b4c4d4... + + // store + _mm_storel_epi64((__m128i *) out, p0); out += out_stride; + _mm_storel_epi64((__m128i *) out, _mm_shuffle_epi32(p0, 0x4e)); out += out_stride; + _mm_storel_epi64((__m128i *) out, p2); out += out_stride; + _mm_storel_epi64((__m128i *) out, _mm_shuffle_epi32(p2, 0x4e)); out += out_stride; + _mm_storel_epi64((__m128i *) out, p1); out += out_stride; + _mm_storel_epi64((__m128i *) out, _mm_shuffle_epi32(p1, 0x4e)); out += out_stride; + _mm_storel_epi64((__m128i *) out, p3); out += out_stride; + _mm_storel_epi64((__m128i *) out, _mm_shuffle_epi32(p3, 0x4e)); + } + +#undef dct_const +#undef dct_rot +#undef dct_widen +#undef dct_wadd +#undef dct_wsub +#undef dct_bfly32o +#undef dct_interleave8 +#undef dct_interleave16 +#undef dct_pass +} + +#endif // STBI_SSE2 + +#ifdef STBI_NEON + +// NEON integer IDCT. should produce bit-identical +// results to the generic C version. +static void stbi__idct_simd(stbi_uc *out, int out_stride, short data[64]) +{ + int16x8_t row0, row1, row2, row3, row4, row5, row6, row7; + + int16x4_t rot0_0 = vdup_n_s16(stbi__f2f(0.5411961f)); + int16x4_t rot0_1 = vdup_n_s16(stbi__f2f(-1.847759065f)); + int16x4_t rot0_2 = vdup_n_s16(stbi__f2f( 0.765366865f)); + int16x4_t rot1_0 = vdup_n_s16(stbi__f2f( 1.175875602f)); + int16x4_t rot1_1 = vdup_n_s16(stbi__f2f(-0.899976223f)); + int16x4_t rot1_2 = vdup_n_s16(stbi__f2f(-2.562915447f)); + int16x4_t rot2_0 = vdup_n_s16(stbi__f2f(-1.961570560f)); + int16x4_t rot2_1 = vdup_n_s16(stbi__f2f(-0.390180644f)); + int16x4_t rot3_0 = vdup_n_s16(stbi__f2f( 0.298631336f)); + int16x4_t rot3_1 = vdup_n_s16(stbi__f2f( 2.053119869f)); + int16x4_t rot3_2 = vdup_n_s16(stbi__f2f( 3.072711026f)); + int16x4_t rot3_3 = vdup_n_s16(stbi__f2f( 1.501321110f)); + +#define dct_long_mul(out, inq, coeff) \ + int32x4_t out##_l = vmull_s16(vget_low_s16(inq), coeff); \ + int32x4_t out##_h = vmull_s16(vget_high_s16(inq), coeff) + +#define dct_long_mac(out, acc, inq, coeff) \ + int32x4_t out##_l = vmlal_s16(acc##_l, vget_low_s16(inq), coeff); \ + int32x4_t out##_h = vmlal_s16(acc##_h, vget_high_s16(inq), coeff) + +#define dct_widen(out, inq) \ + int32x4_t out##_l = vshll_n_s16(vget_low_s16(inq), 12); \ + int32x4_t out##_h = vshll_n_s16(vget_high_s16(inq), 12) + +// wide add +#define dct_wadd(out, a, b) \ + int32x4_t out##_l = vaddq_s32(a##_l, b##_l); \ + int32x4_t out##_h = vaddq_s32(a##_h, b##_h) + +// wide sub +#define dct_wsub(out, a, b) \ + int32x4_t out##_l = vsubq_s32(a##_l, b##_l); \ + int32x4_t out##_h = vsubq_s32(a##_h, b##_h) + +// butterfly a/b, then shift using "shiftop" by "s" and pack +#define dct_bfly32o(out0,out1, a,b,shiftop,s) \ + { \ + dct_wadd(sum, a, b); \ + dct_wsub(dif, a, b); \ + out0 = vcombine_s16(shiftop(sum_l, s), shiftop(sum_h, s)); \ + out1 = vcombine_s16(shiftop(dif_l, s), shiftop(dif_h, s)); \ + } + +#define dct_pass(shiftop, shift) \ + { \ + /* even part */ \ + int16x8_t sum26 = vaddq_s16(row2, row6); \ + dct_long_mul(p1e, sum26, rot0_0); \ + dct_long_mac(t2e, p1e, row6, rot0_1); \ + dct_long_mac(t3e, p1e, row2, rot0_2); \ + int16x8_t sum04 = vaddq_s16(row0, row4); \ + int16x8_t dif04 = vsubq_s16(row0, row4); \ + dct_widen(t0e, sum04); \ + dct_widen(t1e, dif04); \ + dct_wadd(x0, t0e, t3e); \ + dct_wsub(x3, t0e, t3e); \ + dct_wadd(x1, t1e, t2e); \ + dct_wsub(x2, t1e, t2e); \ + /* odd part */ \ + int16x8_t sum15 = vaddq_s16(row1, row5); \ + int16x8_t sum17 = vaddq_s16(row1, row7); \ + int16x8_t sum35 = vaddq_s16(row3, row5); \ + int16x8_t sum37 = vaddq_s16(row3, row7); \ + int16x8_t sumodd = vaddq_s16(sum17, sum35); \ + dct_long_mul(p5o, sumodd, rot1_0); \ + dct_long_mac(p1o, p5o, sum17, rot1_1); \ + dct_long_mac(p2o, p5o, sum35, rot1_2); \ + dct_long_mul(p3o, sum37, rot2_0); \ + dct_long_mul(p4o, sum15, rot2_1); \ + dct_wadd(sump13o, p1o, p3o); \ + dct_wadd(sump24o, p2o, p4o); \ + dct_wadd(sump23o, p2o, p3o); \ + dct_wadd(sump14o, p1o, p4o); \ + dct_long_mac(x4, sump13o, row7, rot3_0); \ + dct_long_mac(x5, sump24o, row5, rot3_1); \ + dct_long_mac(x6, sump23o, row3, rot3_2); \ + dct_long_mac(x7, sump14o, row1, rot3_3); \ + dct_bfly32o(row0,row7, x0,x7,shiftop,shift); \ + dct_bfly32o(row1,row6, x1,x6,shiftop,shift); \ + dct_bfly32o(row2,row5, x2,x5,shiftop,shift); \ + dct_bfly32o(row3,row4, x3,x4,shiftop,shift); \ + } + + // load + row0 = vld1q_s16(data + 0*8); + row1 = vld1q_s16(data + 1*8); + row2 = vld1q_s16(data + 2*8); + row3 = vld1q_s16(data + 3*8); + row4 = vld1q_s16(data + 4*8); + row5 = vld1q_s16(data + 5*8); + row6 = vld1q_s16(data + 6*8); + row7 = vld1q_s16(data + 7*8); + + // add DC bias + row0 = vaddq_s16(row0, vsetq_lane_s16(1024, vdupq_n_s16(0), 0)); + + // column pass + dct_pass(vrshrn_n_s32, 10); + + // 16bit 8x8 transpose + { +// these three map to a single VTRN.16, VTRN.32, and VSWP, respectively. +// whether compilers actually get this is another story, sadly. +#define dct_trn16(x, y) { int16x8x2_t t = vtrnq_s16(x, y); x = t.val[0]; y = t.val[1]; } +#define dct_trn32(x, y) { int32x4x2_t t = vtrnq_s32(vreinterpretq_s32_s16(x), vreinterpretq_s32_s16(y)); x = vreinterpretq_s16_s32(t.val[0]); y = vreinterpretq_s16_s32(t.val[1]); } +#define dct_trn64(x, y) { int16x8_t x0 = x; int16x8_t y0 = y; x = vcombine_s16(vget_low_s16(x0), vget_low_s16(y0)); y = vcombine_s16(vget_high_s16(x0), vget_high_s16(y0)); } + + // pass 1 + dct_trn16(row0, row1); // a0b0a2b2a4b4a6b6 + dct_trn16(row2, row3); + dct_trn16(row4, row5); + dct_trn16(row6, row7); + + // pass 2 + dct_trn32(row0, row2); // a0b0c0d0a4b4c4d4 + dct_trn32(row1, row3); + dct_trn32(row4, row6); + dct_trn32(row5, row7); + + // pass 3 + dct_trn64(row0, row4); // a0b0c0d0e0f0g0h0 + dct_trn64(row1, row5); + dct_trn64(row2, row6); + dct_trn64(row3, row7); + +#undef dct_trn16 +#undef dct_trn32 +#undef dct_trn64 + } + + // row pass + // vrshrn_n_s32 only supports shifts up to 16, we need + // 17. so do a non-rounding shift of 16 first then follow + // up with a rounding shift by 1. + dct_pass(vshrn_n_s32, 16); + + { + // pack and round + uint8x8_t p0 = vqrshrun_n_s16(row0, 1); + uint8x8_t p1 = vqrshrun_n_s16(row1, 1); + uint8x8_t p2 = vqrshrun_n_s16(row2, 1); + uint8x8_t p3 = vqrshrun_n_s16(row3, 1); + uint8x8_t p4 = vqrshrun_n_s16(row4, 1); + uint8x8_t p5 = vqrshrun_n_s16(row5, 1); + uint8x8_t p6 = vqrshrun_n_s16(row6, 1); + uint8x8_t p7 = vqrshrun_n_s16(row7, 1); + + // again, these can translate into one instruction, but often don't. +#define dct_trn8_8(x, y) { uint8x8x2_t t = vtrn_u8(x, y); x = t.val[0]; y = t.val[1]; } +#define dct_trn8_16(x, y) { uint16x4x2_t t = vtrn_u16(vreinterpret_u16_u8(x), vreinterpret_u16_u8(y)); x = vreinterpret_u8_u16(t.val[0]); y = vreinterpret_u8_u16(t.val[1]); } +#define dct_trn8_32(x, y) { uint32x2x2_t t = vtrn_u32(vreinterpret_u32_u8(x), vreinterpret_u32_u8(y)); x = vreinterpret_u8_u32(t.val[0]); y = vreinterpret_u8_u32(t.val[1]); } + + // sadly can't use interleaved stores here since we only write + // 8 bytes to each scan line! + + // 8x8 8-bit transpose pass 1 + dct_trn8_8(p0, p1); + dct_trn8_8(p2, p3); + dct_trn8_8(p4, p5); + dct_trn8_8(p6, p7); + + // pass 2 + dct_trn8_16(p0, p2); + dct_trn8_16(p1, p3); + dct_trn8_16(p4, p6); + dct_trn8_16(p5, p7); + + // pass 3 + dct_trn8_32(p0, p4); + dct_trn8_32(p1, p5); + dct_trn8_32(p2, p6); + dct_trn8_32(p3, p7); + + // store + vst1_u8(out, p0); out += out_stride; + vst1_u8(out, p1); out += out_stride; + vst1_u8(out, p2); out += out_stride; + vst1_u8(out, p3); out += out_stride; + vst1_u8(out, p4); out += out_stride; + vst1_u8(out, p5); out += out_stride; + vst1_u8(out, p6); out += out_stride; + vst1_u8(out, p7); + +#undef dct_trn8_8 +#undef dct_trn8_16 +#undef dct_trn8_32 + } + +#undef dct_long_mul +#undef dct_long_mac +#undef dct_widen +#undef dct_wadd +#undef dct_wsub +#undef dct_bfly32o +#undef dct_pass +} + +#endif // STBI_NEON + +#define STBI__MARKER_none 0xff +// if there's a pending marker from the entropy stream, return that +// otherwise, fetch from the stream and get a marker. if there's no +// marker, return 0xff, which is never a valid marker value +static stbi_uc stbi__get_marker(stbi__jpeg *j) +{ + stbi_uc x; + if (j->marker != STBI__MARKER_none) { x = j->marker; j->marker = STBI__MARKER_none; return x; } + x = stbi__get8(j->s); + if (x != 0xff) return STBI__MARKER_none; + while (x == 0xff) + x = stbi__get8(j->s); // consume repeated 0xff fill bytes + return x; +} + +// in each scan, we'll have scan_n components, and the order +// of the components is specified by order[] +#define STBI__RESTART(x) ((x) >= 0xd0 && (x) <= 0xd7) + +// after a restart interval, stbi__jpeg_reset the entropy decoder and +// the dc prediction +static void stbi__jpeg_reset(stbi__jpeg *j) +{ + j->code_bits = 0; + j->code_buffer = 0; + j->nomore = 0; + j->img_comp[0].dc_pred = j->img_comp[1].dc_pred = j->img_comp[2].dc_pred = j->img_comp[3].dc_pred = 0; + j->marker = STBI__MARKER_none; + j->todo = j->restart_interval ? j->restart_interval : 0x7fffffff; + j->eob_run = 0; + // no more than 1<<31 MCUs if no restart_interal? that's plenty safe, + // since we don't even allow 1<<30 pixels +} + +static int stbi__parse_entropy_coded_data(stbi__jpeg *z) +{ + stbi__jpeg_reset(z); + if (!z->progressive) { + if (z->scan_n == 1) { + int i,j; + STBI_SIMD_ALIGN(short, data[64]); + int n = z->order[0]; + // non-interleaved data, we just need to process one block at a time, + // in trivial scanline order + // number of blocks to do just depends on how many actual "pixels" this + // component has, independent of interleaved MCU blocking and such + int w = (z->img_comp[n].x+7) >> 3; + int h = (z->img_comp[n].y+7) >> 3; + for (j=0; j < h; ++j) { + for (i=0; i < w; ++i) { + int ha = z->img_comp[n].ha; + if (!stbi__jpeg_decode_block(z, data, z->huff_dc+z->img_comp[n].hd, z->huff_ac+ha, z->fast_ac[ha], n, z->dequant[z->img_comp[n].tq])) return 0; + z->idct_block_kernel(z->img_comp[n].data+z->img_comp[n].w2*j*8+i*8, z->img_comp[n].w2, data); + // every data block is an MCU, so countdown the restart interval + if (--z->todo <= 0) { + if (z->code_bits < 24) stbi__grow_buffer_unsafe(z); + // if it's NOT a restart, then just bail, so we get corrupt data + // rather than no data + if (!STBI__RESTART(z->marker)) return 1; + stbi__jpeg_reset(z); + } + } + } + return 1; + } else { // interleaved + int i,j,k,x,y; + STBI_SIMD_ALIGN(short, data[64]); + for (j=0; j < z->img_mcu_y; ++j) { + for (i=0; i < z->img_mcu_x; ++i) { + // scan an interleaved mcu... process scan_n components in order + for (k=0; k < z->scan_n; ++k) { + int n = z->order[k]; + // scan out an mcu's worth of this component; that's just determined + // by the basic H and V specified for the component + for (y=0; y < z->img_comp[n].v; ++y) { + for (x=0; x < z->img_comp[n].h; ++x) { + int x2 = (i*z->img_comp[n].h + x)*8; + int y2 = (j*z->img_comp[n].v + y)*8; + int ha = z->img_comp[n].ha; + if (!stbi__jpeg_decode_block(z, data, z->huff_dc+z->img_comp[n].hd, z->huff_ac+ha, z->fast_ac[ha], n, z->dequant[z->img_comp[n].tq])) return 0; + z->idct_block_kernel(z->img_comp[n].data+z->img_comp[n].w2*y2+x2, z->img_comp[n].w2, data); + } + } + } + // after all interleaved components, that's an interleaved MCU, + // so now count down the restart interval + if (--z->todo <= 0) { + if (z->code_bits < 24) stbi__grow_buffer_unsafe(z); + if (!STBI__RESTART(z->marker)) return 1; + stbi__jpeg_reset(z); + } + } + } + return 1; + } + } else { + if (z->scan_n == 1) { + int i,j; + int n = z->order[0]; + // non-interleaved data, we just need to process one block at a time, + // in trivial scanline order + // number of blocks to do just depends on how many actual "pixels" this + // component has, independent of interleaved MCU blocking and such + int w = (z->img_comp[n].x+7) >> 3; + int h = (z->img_comp[n].y+7) >> 3; + for (j=0; j < h; ++j) { + for (i=0; i < w; ++i) { + short *data = z->img_comp[n].coeff + 64 * (i + j * z->img_comp[n].coeff_w); + if (z->spec_start == 0) { + if (!stbi__jpeg_decode_block_prog_dc(z, data, &z->huff_dc[z->img_comp[n].hd], n)) + return 0; + } else { + int ha = z->img_comp[n].ha; + if (!stbi__jpeg_decode_block_prog_ac(z, data, &z->huff_ac[ha], z->fast_ac[ha])) + return 0; + } + // every data block is an MCU, so countdown the restart interval + if (--z->todo <= 0) { + if (z->code_bits < 24) stbi__grow_buffer_unsafe(z); + if (!STBI__RESTART(z->marker)) return 1; + stbi__jpeg_reset(z); + } + } + } + return 1; + } else { // interleaved + int i,j,k,x,y; + for (j=0; j < z->img_mcu_y; ++j) { + for (i=0; i < z->img_mcu_x; ++i) { + // scan an interleaved mcu... process scan_n components in order + for (k=0; k < z->scan_n; ++k) { + int n = z->order[k]; + // scan out an mcu's worth of this component; that's just determined + // by the basic H and V specified for the component + for (y=0; y < z->img_comp[n].v; ++y) { + for (x=0; x < z->img_comp[n].h; ++x) { + int x2 = (i*z->img_comp[n].h + x); + int y2 = (j*z->img_comp[n].v + y); + short *data = z->img_comp[n].coeff + 64 * (x2 + y2 * z->img_comp[n].coeff_w); + if (!stbi__jpeg_decode_block_prog_dc(z, data, &z->huff_dc[z->img_comp[n].hd], n)) + return 0; + } + } + } + // after all interleaved components, that's an interleaved MCU, + // so now count down the restart interval + if (--z->todo <= 0) { + if (z->code_bits < 24) stbi__grow_buffer_unsafe(z); + if (!STBI__RESTART(z->marker)) return 1; + stbi__jpeg_reset(z); + } + } + } + return 1; + } + } +} + +static void stbi__jpeg_dequantize(short *data, stbi__uint16 *dequant) +{ + int i; + for (i=0; i < 64; ++i) + data[i] *= dequant[i]; +} + +static void stbi__jpeg_finish(stbi__jpeg *z) +{ + if (z->progressive) { + // dequantize and idct the data + int i,j,n; + for (n=0; n < z->s->img_n; ++n) { + int w = (z->img_comp[n].x+7) >> 3; + int h = (z->img_comp[n].y+7) >> 3; + for (j=0; j < h; ++j) { + for (i=0; i < w; ++i) { + short *data = z->img_comp[n].coeff + 64 * (i + j * z->img_comp[n].coeff_w); + stbi__jpeg_dequantize(data, z->dequant[z->img_comp[n].tq]); + z->idct_block_kernel(z->img_comp[n].data+z->img_comp[n].w2*j*8+i*8, z->img_comp[n].w2, data); + } + } + } + } +} + +static int stbi__process_marker(stbi__jpeg *z, int m) +{ + int L; + switch (m) { + case STBI__MARKER_none: // no marker found + return stbi__err("expected marker","Corrupt JPEG"); + + case 0xDD: // DRI - specify restart interval + if (stbi__get16be(z->s) != 4) return stbi__err("bad DRI len","Corrupt JPEG"); + z->restart_interval = stbi__get16be(z->s); + return 1; + + case 0xDB: // DQT - define quantization table + L = stbi__get16be(z->s)-2; + while (L > 0) { + int q = stbi__get8(z->s); + int p = q >> 4, sixteen = (p != 0); + int t = q & 15,i; + if (p != 0 && p != 1) return stbi__err("bad DQT type","Corrupt JPEG"); + if (t > 3) return stbi__err("bad DQT table","Corrupt JPEG"); + + for (i=0; i < 64; ++i) + z->dequant[t][stbi__jpeg_dezigzag[i]] = (stbi__uint16)(sixteen ? stbi__get16be(z->s) : stbi__get8(z->s)); + L -= (sixteen ? 129 : 65); + } + return L==0; + + case 0xC4: // DHT - define huffman table + L = stbi__get16be(z->s)-2; + while (L > 0) { + stbi_uc *v; + int sizes[16],i,n=0; + int q = stbi__get8(z->s); + int tc = q >> 4; + int th = q & 15; + if (tc > 1 || th > 3) return stbi__err("bad DHT header","Corrupt JPEG"); + for (i=0; i < 16; ++i) { + sizes[i] = stbi__get8(z->s); + n += sizes[i]; + } + if(n > 256) return stbi__err("bad DHT header","Corrupt JPEG"); // Loop over i < n would write past end of values! + L -= 17; + if (tc == 0) { + if (!stbi__build_huffman(z->huff_dc+th, sizes)) return 0; + v = z->huff_dc[th].values; + } else { + if (!stbi__build_huffman(z->huff_ac+th, sizes)) return 0; + v = z->huff_ac[th].values; + } + for (i=0; i < n; ++i) + v[i] = stbi__get8(z->s); + if (tc != 0) + stbi__build_fast_ac(z->fast_ac[th], z->huff_ac + th); + L -= n; + } + return L==0; + } + + // check for comment block or APP blocks + if ((m >= 0xE0 && m <= 0xEF) || m == 0xFE) { + L = stbi__get16be(z->s); + if (L < 2) { + if (m == 0xFE) + return stbi__err("bad COM len","Corrupt JPEG"); + else + return stbi__err("bad APP len","Corrupt JPEG"); + } + L -= 2; + + if (m == 0xE0 && L >= 5) { // JFIF APP0 segment + static const unsigned char tag[5] = {'J','F','I','F','\0'}; + int ok = 1; + int i; + for (i=0; i < 5; ++i) + if (stbi__get8(z->s) != tag[i]) + ok = 0; + L -= 5; + if (ok) + z->jfif = 1; + } else if (m == 0xEE && L >= 12) { // Adobe APP14 segment + static const unsigned char tag[6] = {'A','d','o','b','e','\0'}; + int ok = 1; + int i; + for (i=0; i < 6; ++i) + if (stbi__get8(z->s) != tag[i]) + ok = 0; + L -= 6; + if (ok) { + stbi__get8(z->s); // version + stbi__get16be(z->s); // flags0 + stbi__get16be(z->s); // flags1 + z->app14_color_transform = stbi__get8(z->s); // color transform + L -= 6; + } + } + + stbi__skip(z->s, L); + return 1; + } + + return stbi__err("unknown marker","Corrupt JPEG"); +} + +// after we see SOS +static int stbi__process_scan_header(stbi__jpeg *z) +{ + int i; + int Ls = stbi__get16be(z->s); + z->scan_n = stbi__get8(z->s); + if (z->scan_n < 1 || z->scan_n > 4 || z->scan_n > (int) z->s->img_n) return stbi__err("bad SOS component count","Corrupt JPEG"); + if (Ls != 6+2*z->scan_n) return stbi__err("bad SOS len","Corrupt JPEG"); + for (i=0; i < z->scan_n; ++i) { + int id = stbi__get8(z->s), which; + int q = stbi__get8(z->s); + for (which = 0; which < z->s->img_n; ++which) + if (z->img_comp[which].id == id) + break; + if (which == z->s->img_n) return 0; // no match + z->img_comp[which].hd = q >> 4; if (z->img_comp[which].hd > 3) return stbi__err("bad DC huff","Corrupt JPEG"); + z->img_comp[which].ha = q & 15; if (z->img_comp[which].ha > 3) return stbi__err("bad AC huff","Corrupt JPEG"); + z->order[i] = which; + } + + { + int aa; + z->spec_start = stbi__get8(z->s); + z->spec_end = stbi__get8(z->s); // should be 63, but might be 0 + aa = stbi__get8(z->s); + z->succ_high = (aa >> 4); + z->succ_low = (aa & 15); + if (z->progressive) { + if (z->spec_start > 63 || z->spec_end > 63 || z->spec_start > z->spec_end || z->succ_high > 13 || z->succ_low > 13) + return stbi__err("bad SOS", "Corrupt JPEG"); + } else { + if (z->spec_start != 0) return stbi__err("bad SOS","Corrupt JPEG"); + if (z->succ_high != 0 || z->succ_low != 0) return stbi__err("bad SOS","Corrupt JPEG"); + z->spec_end = 63; + } + } + + return 1; +} + +static int stbi__free_jpeg_components(stbi__jpeg *z, int ncomp, int why) +{ + int i; + for (i=0; i < ncomp; ++i) { + if (z->img_comp[i].raw_data) { + STBI_FREE(z->img_comp[i].raw_data); + z->img_comp[i].raw_data = NULL; + z->img_comp[i].data = NULL; + } + if (z->img_comp[i].raw_coeff) { + STBI_FREE(z->img_comp[i].raw_coeff); + z->img_comp[i].raw_coeff = 0; + z->img_comp[i].coeff = 0; + } + if (z->img_comp[i].linebuf) { + STBI_FREE(z->img_comp[i].linebuf); + z->img_comp[i].linebuf = NULL; + } + } + return why; +} + +static int stbi__process_frame_header(stbi__jpeg *z, int scan) +{ + stbi__context *s = z->s; + int Lf,p,i,q, h_max=1,v_max=1,c; + Lf = stbi__get16be(s); if (Lf < 11) return stbi__err("bad SOF len","Corrupt JPEG"); // JPEG + p = stbi__get8(s); if (p != 8) return stbi__err("only 8-bit","JPEG format not supported: 8-bit only"); // JPEG baseline + s->img_y = stbi__get16be(s); if (s->img_y == 0) return stbi__err("no header height", "JPEG format not supported: delayed height"); // Legal, but we don't handle it--but neither does IJG + s->img_x = stbi__get16be(s); if (s->img_x == 0) return stbi__err("0 width","Corrupt JPEG"); // JPEG requires + if (s->img_y > STBI_MAX_DIMENSIONS) return stbi__err("too large","Very large image (corrupt?)"); + if (s->img_x > STBI_MAX_DIMENSIONS) return stbi__err("too large","Very large image (corrupt?)"); + c = stbi__get8(s); + if (c != 3 && c != 1 && c != 4) return stbi__err("bad component count","Corrupt JPEG"); + s->img_n = c; + for (i=0; i < c; ++i) { + z->img_comp[i].data = NULL; + z->img_comp[i].linebuf = NULL; + } + + if (Lf != 8+3*s->img_n) return stbi__err("bad SOF len","Corrupt JPEG"); + + z->rgb = 0; + for (i=0; i < s->img_n; ++i) { + static const unsigned char rgb[3] = { 'R', 'G', 'B' }; + z->img_comp[i].id = stbi__get8(s); + if (s->img_n == 3 && z->img_comp[i].id == rgb[i]) + ++z->rgb; + q = stbi__get8(s); + z->img_comp[i].h = (q >> 4); if (!z->img_comp[i].h || z->img_comp[i].h > 4) return stbi__err("bad H","Corrupt JPEG"); + z->img_comp[i].v = q & 15; if (!z->img_comp[i].v || z->img_comp[i].v > 4) return stbi__err("bad V","Corrupt JPEG"); + z->img_comp[i].tq = stbi__get8(s); if (z->img_comp[i].tq > 3) return stbi__err("bad TQ","Corrupt JPEG"); + } + + if (scan != STBI__SCAN_load) return 1; + + if (!stbi__mad3sizes_valid(s->img_x, s->img_y, s->img_n, 0)) return stbi__err("too large", "Image too large to decode"); + + for (i=0; i < s->img_n; ++i) { + if (z->img_comp[i].h > h_max) h_max = z->img_comp[i].h; + if (z->img_comp[i].v > v_max) v_max = z->img_comp[i].v; + } + + // check that plane subsampling factors are integer ratios; our resamplers can't deal with fractional ratios + // and I've never seen a non-corrupted JPEG file actually use them + for (i=0; i < s->img_n; ++i) { + if (h_max % z->img_comp[i].h != 0) return stbi__err("bad H","Corrupt JPEG"); + if (v_max % z->img_comp[i].v != 0) return stbi__err("bad V","Corrupt JPEG"); + } + + // compute interleaved mcu info + z->img_h_max = h_max; + z->img_v_max = v_max; + z->img_mcu_w = h_max * 8; + z->img_mcu_h = v_max * 8; + // these sizes can't be more than 17 bits + z->img_mcu_x = (s->img_x + z->img_mcu_w-1) / z->img_mcu_w; + z->img_mcu_y = (s->img_y + z->img_mcu_h-1) / z->img_mcu_h; + + for (i=0; i < s->img_n; ++i) { + // number of effective pixels (e.g. for non-interleaved MCU) + z->img_comp[i].x = (s->img_x * z->img_comp[i].h + h_max-1) / h_max; + z->img_comp[i].y = (s->img_y * z->img_comp[i].v + v_max-1) / v_max; + // to simplify generation, we'll allocate enough memory to decode + // the bogus oversized data from using interleaved MCUs and their + // big blocks (e.g. a 16x16 iMCU on an image of width 33); we won't + // discard the extra data until colorspace conversion + // + // img_mcu_x, img_mcu_y: <=17 bits; comp[i].h and .v are <=4 (checked earlier) + // so these muls can't overflow with 32-bit ints (which we require) + z->img_comp[i].w2 = z->img_mcu_x * z->img_comp[i].h * 8; + z->img_comp[i].h2 = z->img_mcu_y * z->img_comp[i].v * 8; + z->img_comp[i].coeff = 0; + z->img_comp[i].raw_coeff = 0; + z->img_comp[i].linebuf = NULL; + z->img_comp[i].raw_data = stbi__malloc_mad2(z->img_comp[i].w2, z->img_comp[i].h2, 15); + if (z->img_comp[i].raw_data == NULL) + return stbi__free_jpeg_components(z, i+1, stbi__err("outofmem", "Out of memory")); + // align blocks for idct using mmx/sse + z->img_comp[i].data = (stbi_uc*) (((size_t) z->img_comp[i].raw_data + 15) & ~15); + if (z->progressive) { + // w2, h2 are multiples of 8 (see above) + z->img_comp[i].coeff_w = z->img_comp[i].w2 / 8; + z->img_comp[i].coeff_h = z->img_comp[i].h2 / 8; + z->img_comp[i].raw_coeff = stbi__malloc_mad3(z->img_comp[i].w2, z->img_comp[i].h2, sizeof(short), 15); + if (z->img_comp[i].raw_coeff == NULL) + return stbi__free_jpeg_components(z, i+1, stbi__err("outofmem", "Out of memory")); + z->img_comp[i].coeff = (short*) (((size_t) z->img_comp[i].raw_coeff + 15) & ~15); + } + } + + return 1; +} + +// use comparisons since in some cases we handle more than one case (e.g. SOF) +#define stbi__DNL(x) ((x) == 0xdc) +#define stbi__SOI(x) ((x) == 0xd8) +#define stbi__EOI(x) ((x) == 0xd9) +#define stbi__SOF(x) ((x) == 0xc0 || (x) == 0xc1 || (x) == 0xc2) +#define stbi__SOS(x) ((x) == 0xda) + +#define stbi__SOF_progressive(x) ((x) == 0xc2) + +static int stbi__decode_jpeg_header(stbi__jpeg *z, int scan) +{ + int m; + z->jfif = 0; + z->app14_color_transform = -1; // valid values are 0,1,2 + z->marker = STBI__MARKER_none; // initialize cached marker to empty + m = stbi__get_marker(z); + if (!stbi__SOI(m)) return stbi__err("no SOI","Corrupt JPEG"); + if (scan == STBI__SCAN_type) return 1; + m = stbi__get_marker(z); + while (!stbi__SOF(m)) { + if (!stbi__process_marker(z,m)) return 0; + m = stbi__get_marker(z); + while (m == STBI__MARKER_none) { + // some files have extra padding after their blocks, so ok, we'll scan + if (stbi__at_eof(z->s)) return stbi__err("no SOF", "Corrupt JPEG"); + m = stbi__get_marker(z); + } + } + z->progressive = stbi__SOF_progressive(m); + if (!stbi__process_frame_header(z, scan)) return 0; + return 1; +} + +static stbi_uc stbi__skip_jpeg_junk_at_end(stbi__jpeg *j) +{ + // some JPEGs have junk at end, skip over it but if we find what looks + // like a valid marker, resume there + while (!stbi__at_eof(j->s)) { + stbi_uc x = stbi__get8(j->s); + while (x == 0xff) { // might be a marker + if (stbi__at_eof(j->s)) return STBI__MARKER_none; + x = stbi__get8(j->s); + if (x != 0x00 && x != 0xff) { + // not a stuffed zero or lead-in to another marker, looks + // like an actual marker, return it + return x; + } + // stuffed zero has x=0 now which ends the loop, meaning we go + // back to regular scan loop. + // repeated 0xff keeps trying to read the next byte of the marker. + } + } + return STBI__MARKER_none; +} + +// decode image to YCbCr format +static int stbi__decode_jpeg_image(stbi__jpeg *j) +{ + int m; + for (m = 0; m < 4; m++) { + j->img_comp[m].raw_data = NULL; + j->img_comp[m].raw_coeff = NULL; + } + j->restart_interval = 0; + if (!stbi__decode_jpeg_header(j, STBI__SCAN_load)) return 0; + m = stbi__get_marker(j); + while (!stbi__EOI(m)) { + if (stbi__SOS(m)) { + if (!stbi__process_scan_header(j)) return 0; + if (!stbi__parse_entropy_coded_data(j)) return 0; + if (j->marker == STBI__MARKER_none ) { + j->marker = stbi__skip_jpeg_junk_at_end(j); + // if we reach eof without hitting a marker, stbi__get_marker() below will fail and we'll eventually return 0 + } + m = stbi__get_marker(j); + if (STBI__RESTART(m)) + m = stbi__get_marker(j); + } else if (stbi__DNL(m)) { + int Ld = stbi__get16be(j->s); + stbi__uint32 NL = stbi__get16be(j->s); + if (Ld != 4) return stbi__err("bad DNL len", "Corrupt JPEG"); + if (NL != j->s->img_y) return stbi__err("bad DNL height", "Corrupt JPEG"); + m = stbi__get_marker(j); + } else { + if (!stbi__process_marker(j, m)) return 1; + m = stbi__get_marker(j); + } + } + if (j->progressive) + stbi__jpeg_finish(j); + return 1; +} + +// static jfif-centered resampling (across block boundaries) + +typedef stbi_uc *(*resample_row_func)(stbi_uc *out, stbi_uc *in0, stbi_uc *in1, + int w, int hs); + +#define stbi__div4(x) ((stbi_uc) ((x) >> 2)) + +static stbi_uc *resample_row_1(stbi_uc *out, stbi_uc *in_near, stbi_uc *in_far, int w, int hs) +{ + STBI_NOTUSED(out); + STBI_NOTUSED(in_far); + STBI_NOTUSED(w); + STBI_NOTUSED(hs); + return in_near; +} + +static stbi_uc* stbi__resample_row_v_2(stbi_uc *out, stbi_uc *in_near, stbi_uc *in_far, int w, int hs) +{ + // need to generate two samples vertically for every one in input + int i; + STBI_NOTUSED(hs); + for (i=0; i < w; ++i) + out[i] = stbi__div4(3*in_near[i] + in_far[i] + 2); + return out; +} + +static stbi_uc* stbi__resample_row_h_2(stbi_uc *out, stbi_uc *in_near, stbi_uc *in_far, int w, int hs) +{ + // need to generate two samples horizontally for every one in input + int i; + stbi_uc *input = in_near; + + if (w == 1) { + // if only one sample, can't do any interpolation + out[0] = out[1] = input[0]; + return out; + } + + out[0] = input[0]; + out[1] = stbi__div4(input[0]*3 + input[1] + 2); + for (i=1; i < w-1; ++i) { + int n = 3*input[i]+2; + out[i*2+0] = stbi__div4(n+input[i-1]); + out[i*2+1] = stbi__div4(n+input[i+1]); + } + out[i*2+0] = stbi__div4(input[w-2]*3 + input[w-1] + 2); + out[i*2+1] = input[w-1]; + + STBI_NOTUSED(in_far); + STBI_NOTUSED(hs); + + return out; +} + +#define stbi__div16(x) ((stbi_uc) ((x) >> 4)) + +static stbi_uc *stbi__resample_row_hv_2(stbi_uc *out, stbi_uc *in_near, stbi_uc *in_far, int w, int hs) +{ + // need to generate 2x2 samples for every one in input + int i,t0,t1; + if (w == 1) { + out[0] = out[1] = stbi__div4(3*in_near[0] + in_far[0] + 2); + return out; + } + + t1 = 3*in_near[0] + in_far[0]; + out[0] = stbi__div4(t1+2); + for (i=1; i < w; ++i) { + t0 = t1; + t1 = 3*in_near[i]+in_far[i]; + out[i*2-1] = stbi__div16(3*t0 + t1 + 8); + out[i*2 ] = stbi__div16(3*t1 + t0 + 8); + } + out[w*2-1] = stbi__div4(t1+2); + + STBI_NOTUSED(hs); + + return out; +} + +#if defined(STBI_SSE2) || defined(STBI_NEON) +static stbi_uc *stbi__resample_row_hv_2_simd(stbi_uc *out, stbi_uc *in_near, stbi_uc *in_far, int w, int hs) +{ + // need to generate 2x2 samples for every one in input + int i=0,t0,t1; + + if (w == 1) { + out[0] = out[1] = stbi__div4(3*in_near[0] + in_far[0] + 2); + return out; + } + + t1 = 3*in_near[0] + in_far[0]; + // process groups of 8 pixels for as long as we can. + // note we can't handle the last pixel in a row in this loop + // because we need to handle the filter boundary conditions. + for (; i < ((w-1) & ~7); i += 8) { +#if defined(STBI_SSE2) + // load and perform the vertical filtering pass + // this uses 3*x + y = 4*x + (y - x) + __m128i zero = _mm_setzero_si128(); + __m128i farb = _mm_loadl_epi64((__m128i *) (in_far + i)); + __m128i nearb = _mm_loadl_epi64((__m128i *) (in_near + i)); + __m128i farw = _mm_unpacklo_epi8(farb, zero); + __m128i nearw = _mm_unpacklo_epi8(nearb, zero); + __m128i diff = _mm_sub_epi16(farw, nearw); + __m128i nears = _mm_slli_epi16(nearw, 2); + __m128i curr = _mm_add_epi16(nears, diff); // current row + + // horizontal filter works the same based on shifted vers of current + // row. "prev" is current row shifted right by 1 pixel; we need to + // insert the previous pixel value (from t1). + // "next" is current row shifted left by 1 pixel, with first pixel + // of next block of 8 pixels added in. + __m128i prv0 = _mm_slli_si128(curr, 2); + __m128i nxt0 = _mm_srli_si128(curr, 2); + __m128i prev = _mm_insert_epi16(prv0, t1, 0); + __m128i next = _mm_insert_epi16(nxt0, 3*in_near[i+8] + in_far[i+8], 7); + + // horizontal filter, polyphase implementation since it's convenient: + // even pixels = 3*cur + prev = cur*4 + (prev - cur) + // odd pixels = 3*cur + next = cur*4 + (next - cur) + // note the shared term. + __m128i bias = _mm_set1_epi16(8); + __m128i curs = _mm_slli_epi16(curr, 2); + __m128i prvd = _mm_sub_epi16(prev, curr); + __m128i nxtd = _mm_sub_epi16(next, curr); + __m128i curb = _mm_add_epi16(curs, bias); + __m128i even = _mm_add_epi16(prvd, curb); + __m128i odd = _mm_add_epi16(nxtd, curb); + + // interleave even and odd pixels, then undo scaling. + __m128i int0 = _mm_unpacklo_epi16(even, odd); + __m128i int1 = _mm_unpackhi_epi16(even, odd); + __m128i de0 = _mm_srli_epi16(int0, 4); + __m128i de1 = _mm_srli_epi16(int1, 4); + + // pack and write output + __m128i outv = _mm_packus_epi16(de0, de1); + _mm_storeu_si128((__m128i *) (out + i*2), outv); +#elif defined(STBI_NEON) + // load and perform the vertical filtering pass + // this uses 3*x + y = 4*x + (y - x) + uint8x8_t farb = vld1_u8(in_far + i); + uint8x8_t nearb = vld1_u8(in_near + i); + int16x8_t diff = vreinterpretq_s16_u16(vsubl_u8(farb, nearb)); + int16x8_t nears = vreinterpretq_s16_u16(vshll_n_u8(nearb, 2)); + int16x8_t curr = vaddq_s16(nears, diff); // current row + + // horizontal filter works the same based on shifted vers of current + // row. "prev" is current row shifted right by 1 pixel; we need to + // insert the previous pixel value (from t1). + // "next" is current row shifted left by 1 pixel, with first pixel + // of next block of 8 pixels added in. + int16x8_t prv0 = vextq_s16(curr, curr, 7); + int16x8_t nxt0 = vextq_s16(curr, curr, 1); + int16x8_t prev = vsetq_lane_s16(t1, prv0, 0); + int16x8_t next = vsetq_lane_s16(3*in_near[i+8] + in_far[i+8], nxt0, 7); + + // horizontal filter, polyphase implementation since it's convenient: + // even pixels = 3*cur + prev = cur*4 + (prev - cur) + // odd pixels = 3*cur + next = cur*4 + (next - cur) + // note the shared term. + int16x8_t curs = vshlq_n_s16(curr, 2); + int16x8_t prvd = vsubq_s16(prev, curr); + int16x8_t nxtd = vsubq_s16(next, curr); + int16x8_t even = vaddq_s16(curs, prvd); + int16x8_t odd = vaddq_s16(curs, nxtd); + + // undo scaling and round, then store with even/odd phases interleaved + uint8x8x2_t o; + o.val[0] = vqrshrun_n_s16(even, 4); + o.val[1] = vqrshrun_n_s16(odd, 4); + vst2_u8(out + i*2, o); +#endif + + // "previous" value for next iter + t1 = 3*in_near[i+7] + in_far[i+7]; + } + + t0 = t1; + t1 = 3*in_near[i] + in_far[i]; + out[i*2] = stbi__div16(3*t1 + t0 + 8); + + for (++i; i < w; ++i) { + t0 = t1; + t1 = 3*in_near[i]+in_far[i]; + out[i*2-1] = stbi__div16(3*t0 + t1 + 8); + out[i*2 ] = stbi__div16(3*t1 + t0 + 8); + } + out[w*2-1] = stbi__div4(t1+2); + + STBI_NOTUSED(hs); + + return out; +} +#endif + +static stbi_uc *stbi__resample_row_generic(stbi_uc *out, stbi_uc *in_near, stbi_uc *in_far, int w, int hs) +{ + // resample with nearest-neighbor + int i,j; + STBI_NOTUSED(in_far); + for (i=0; i < w; ++i) + for (j=0; j < hs; ++j) + out[i*hs+j] = in_near[i]; + return out; +} + +// this is a reduced-precision calculation of YCbCr-to-RGB introduced +// to make sure the code produces the same results in both SIMD and scalar +#define stbi__float2fixed(x) (((int) ((x) * 4096.0f + 0.5f)) << 8) +static void stbi__YCbCr_to_RGB_row(stbi_uc *out, const stbi_uc *y, const stbi_uc *pcb, const stbi_uc *pcr, int count, int step) +{ + int i; + for (i=0; i < count; ++i) { + int y_fixed = (y[i] << 20) + (1<<19); // rounding + int r,g,b; + int cr = pcr[i] - 128; + int cb = pcb[i] - 128; + r = y_fixed + cr* stbi__float2fixed(1.40200f); + g = y_fixed + (cr*-stbi__float2fixed(0.71414f)) + ((cb*-stbi__float2fixed(0.34414f)) & 0xffff0000); + b = y_fixed + cb* stbi__float2fixed(1.77200f); + r >>= 20; + g >>= 20; + b >>= 20; + if ((unsigned) r > 255) { if (r < 0) r = 0; else r = 255; } + if ((unsigned) g > 255) { if (g < 0) g = 0; else g = 255; } + if ((unsigned) b > 255) { if (b < 0) b = 0; else b = 255; } + out[0] = (stbi_uc)r; + out[1] = (stbi_uc)g; + out[2] = (stbi_uc)b; + out[3] = 255; + out += step; + } +} + +#if defined(STBI_SSE2) || defined(STBI_NEON) +static void stbi__YCbCr_to_RGB_simd(stbi_uc *out, stbi_uc const *y, stbi_uc const *pcb, stbi_uc const *pcr, int count, int step) +{ + int i = 0; + +#ifdef STBI_SSE2 + // step == 3 is pretty ugly on the final interleave, and i'm not convinced + // it's useful in practice (you wouldn't use it for textures, for example). + // so just accelerate step == 4 case. + if (step == 4) { + // this is a fairly straightforward implementation and not super-optimized. + __m128i signflip = _mm_set1_epi8(-0x80); + __m128i cr_const0 = _mm_set1_epi16( (short) ( 1.40200f*4096.0f+0.5f)); + __m128i cr_const1 = _mm_set1_epi16( - (short) ( 0.71414f*4096.0f+0.5f)); + __m128i cb_const0 = _mm_set1_epi16( - (short) ( 0.34414f*4096.0f+0.5f)); + __m128i cb_const1 = _mm_set1_epi16( (short) ( 1.77200f*4096.0f+0.5f)); + __m128i y_bias = _mm_set1_epi8((char) (unsigned char) 128); + __m128i xw = _mm_set1_epi16(255); // alpha channel + + for (; i+7 < count; i += 8) { + // load + __m128i y_bytes = _mm_loadl_epi64((__m128i *) (y+i)); + __m128i cr_bytes = _mm_loadl_epi64((__m128i *) (pcr+i)); + __m128i cb_bytes = _mm_loadl_epi64((__m128i *) (pcb+i)); + __m128i cr_biased = _mm_xor_si128(cr_bytes, signflip); // -128 + __m128i cb_biased = _mm_xor_si128(cb_bytes, signflip); // -128 + + // unpack to short (and left-shift cr, cb by 8) + __m128i yw = _mm_unpacklo_epi8(y_bias, y_bytes); + __m128i crw = _mm_unpacklo_epi8(_mm_setzero_si128(), cr_biased); + __m128i cbw = _mm_unpacklo_epi8(_mm_setzero_si128(), cb_biased); + + // color transform + __m128i yws = _mm_srli_epi16(yw, 4); + __m128i cr0 = _mm_mulhi_epi16(cr_const0, crw); + __m128i cb0 = _mm_mulhi_epi16(cb_const0, cbw); + __m128i cb1 = _mm_mulhi_epi16(cbw, cb_const1); + __m128i cr1 = _mm_mulhi_epi16(crw, cr_const1); + __m128i rws = _mm_add_epi16(cr0, yws); + __m128i gwt = _mm_add_epi16(cb0, yws); + __m128i bws = _mm_add_epi16(yws, cb1); + __m128i gws = _mm_add_epi16(gwt, cr1); + + // descale + __m128i rw = _mm_srai_epi16(rws, 4); + __m128i bw = _mm_srai_epi16(bws, 4); + __m128i gw = _mm_srai_epi16(gws, 4); + + // back to byte, set up for transpose + __m128i brb = _mm_packus_epi16(rw, bw); + __m128i gxb = _mm_packus_epi16(gw, xw); + + // transpose to interleave channels + __m128i t0 = _mm_unpacklo_epi8(brb, gxb); + __m128i t1 = _mm_unpackhi_epi8(brb, gxb); + __m128i o0 = _mm_unpacklo_epi16(t0, t1); + __m128i o1 = _mm_unpackhi_epi16(t0, t1); + + // store + _mm_storeu_si128((__m128i *) (out + 0), o0); + _mm_storeu_si128((__m128i *) (out + 16), o1); + out += 32; + } + } +#endif + +#ifdef STBI_NEON + // in this version, step=3 support would be easy to add. but is there demand? + if (step == 4) { + // this is a fairly straightforward implementation and not super-optimized. + uint8x8_t signflip = vdup_n_u8(0x80); + int16x8_t cr_const0 = vdupq_n_s16( (short) ( 1.40200f*4096.0f+0.5f)); + int16x8_t cr_const1 = vdupq_n_s16( - (short) ( 0.71414f*4096.0f+0.5f)); + int16x8_t cb_const0 = vdupq_n_s16( - (short) ( 0.34414f*4096.0f+0.5f)); + int16x8_t cb_const1 = vdupq_n_s16( (short) ( 1.77200f*4096.0f+0.5f)); + + for (; i+7 < count; i += 8) { + // load + uint8x8_t y_bytes = vld1_u8(y + i); + uint8x8_t cr_bytes = vld1_u8(pcr + i); + uint8x8_t cb_bytes = vld1_u8(pcb + i); + int8x8_t cr_biased = vreinterpret_s8_u8(vsub_u8(cr_bytes, signflip)); + int8x8_t cb_biased = vreinterpret_s8_u8(vsub_u8(cb_bytes, signflip)); + + // expand to s16 + int16x8_t yws = vreinterpretq_s16_u16(vshll_n_u8(y_bytes, 4)); + int16x8_t crw = vshll_n_s8(cr_biased, 7); + int16x8_t cbw = vshll_n_s8(cb_biased, 7); + + // color transform + int16x8_t cr0 = vqdmulhq_s16(crw, cr_const0); + int16x8_t cb0 = vqdmulhq_s16(cbw, cb_const0); + int16x8_t cr1 = vqdmulhq_s16(crw, cr_const1); + int16x8_t cb1 = vqdmulhq_s16(cbw, cb_const1); + int16x8_t rws = vaddq_s16(yws, cr0); + int16x8_t gws = vaddq_s16(vaddq_s16(yws, cb0), cr1); + int16x8_t bws = vaddq_s16(yws, cb1); + + // undo scaling, round, convert to byte + uint8x8x4_t o; + o.val[0] = vqrshrun_n_s16(rws, 4); + o.val[1] = vqrshrun_n_s16(gws, 4); + o.val[2] = vqrshrun_n_s16(bws, 4); + o.val[3] = vdup_n_u8(255); + + // store, interleaving r/g/b/a + vst4_u8(out, o); + out += 8*4; + } + } +#endif + + for (; i < count; ++i) { + int y_fixed = (y[i] << 20) + (1<<19); // rounding + int r,g,b; + int cr = pcr[i] - 128; + int cb = pcb[i] - 128; + r = y_fixed + cr* stbi__float2fixed(1.40200f); + g = y_fixed + cr*-stbi__float2fixed(0.71414f) + ((cb*-stbi__float2fixed(0.34414f)) & 0xffff0000); + b = y_fixed + cb* stbi__float2fixed(1.77200f); + r >>= 20; + g >>= 20; + b >>= 20; + if ((unsigned) r > 255) { if (r < 0) r = 0; else r = 255; } + if ((unsigned) g > 255) { if (g < 0) g = 0; else g = 255; } + if ((unsigned) b > 255) { if (b < 0) b = 0; else b = 255; } + out[0] = (stbi_uc)r; + out[1] = (stbi_uc)g; + out[2] = (stbi_uc)b; + out[3] = 255; + out += step; + } +} +#endif + +// set up the kernels +static void stbi__setup_jpeg(stbi__jpeg *j) +{ + j->idct_block_kernel = stbi__idct_block; + j->YCbCr_to_RGB_kernel = stbi__YCbCr_to_RGB_row; + j->resample_row_hv_2_kernel = stbi__resample_row_hv_2; + +#ifdef STBI_SSE2 + if (stbi__sse2_available()) { + j->idct_block_kernel = stbi__idct_simd; + j->YCbCr_to_RGB_kernel = stbi__YCbCr_to_RGB_simd; + j->resample_row_hv_2_kernel = stbi__resample_row_hv_2_simd; + } +#endif + +#ifdef STBI_NEON + j->idct_block_kernel = stbi__idct_simd; + j->YCbCr_to_RGB_kernel = stbi__YCbCr_to_RGB_simd; + j->resample_row_hv_2_kernel = stbi__resample_row_hv_2_simd; +#endif +} + +// clean up the temporary component buffers +static void stbi__cleanup_jpeg(stbi__jpeg *j) +{ + stbi__free_jpeg_components(j, j->s->img_n, 0); +} + +typedef struct +{ + resample_row_func resample; + stbi_uc *line0,*line1; + int hs,vs; // expansion factor in each axis + int w_lores; // horizontal pixels pre-expansion + int ystep; // how far through vertical expansion we are + int ypos; // which pre-expansion row we're on +} stbi__resample; + +// fast 0..255 * 0..255 => 0..255 rounded multiplication +static stbi_uc stbi__blinn_8x8(stbi_uc x, stbi_uc y) +{ + unsigned int t = x*y + 128; + return (stbi_uc) ((t + (t >>8)) >> 8); +} + +static stbi_uc *load_jpeg_image(stbi__jpeg *z, int *out_x, int *out_y, int *comp, int req_comp) +{ + int n, decode_n, is_rgb; + z->s->img_n = 0; // make stbi__cleanup_jpeg safe + + // validate req_comp + if (req_comp < 0 || req_comp > 4) return stbi__errpuc("bad req_comp", "Internal error"); + + // load a jpeg image from whichever source, but leave in YCbCr format + if (!stbi__decode_jpeg_image(z)) { stbi__cleanup_jpeg(z); return NULL; } + + // determine actual number of components to generate + n = req_comp ? req_comp : z->s->img_n >= 3 ? 3 : 1; + + is_rgb = z->s->img_n == 3 && (z->rgb == 3 || (z->app14_color_transform == 0 && !z->jfif)); + + if (z->s->img_n == 3 && n < 3 && !is_rgb) + decode_n = 1; + else + decode_n = z->s->img_n; + + // nothing to do if no components requested; check this now to avoid + // accessing uninitialized coutput[0] later + if (decode_n <= 0) { stbi__cleanup_jpeg(z); return NULL; } + + // resample and color-convert + { + int k; + unsigned int i,j; + stbi_uc *output; + stbi_uc *coutput[4] = { NULL, NULL, NULL, NULL }; + + stbi__resample res_comp[4]; + + for (k=0; k < decode_n; ++k) { + stbi__resample *r = &res_comp[k]; + + // allocate line buffer big enough for upsampling off the edges + // with upsample factor of 4 + z->img_comp[k].linebuf = (stbi_uc *) stbi__malloc(z->s->img_x + 3); + if (!z->img_comp[k].linebuf) { stbi__cleanup_jpeg(z); return stbi__errpuc("outofmem", "Out of memory"); } + + r->hs = z->img_h_max / z->img_comp[k].h; + r->vs = z->img_v_max / z->img_comp[k].v; + r->ystep = r->vs >> 1; + r->w_lores = (z->s->img_x + r->hs-1) / r->hs; + r->ypos = 0; + r->line0 = r->line1 = z->img_comp[k].data; + + if (r->hs == 1 && r->vs == 1) r->resample = resample_row_1; + else if (r->hs == 1 && r->vs == 2) r->resample = stbi__resample_row_v_2; + else if (r->hs == 2 && r->vs == 1) r->resample = stbi__resample_row_h_2; + else if (r->hs == 2 && r->vs == 2) r->resample = z->resample_row_hv_2_kernel; + else r->resample = stbi__resample_row_generic; + } + + // can't error after this so, this is safe + output = (stbi_uc *) stbi__malloc_mad3(n, z->s->img_x, z->s->img_y, 1); + if (!output) { stbi__cleanup_jpeg(z); return stbi__errpuc("outofmem", "Out of memory"); } + + // now go ahead and resample + for (j=0; j < z->s->img_y; ++j) { + stbi_uc *out = output + n * z->s->img_x * j; + for (k=0; k < decode_n; ++k) { + stbi__resample *r = &res_comp[k]; + int y_bot = r->ystep >= (r->vs >> 1); + coutput[k] = r->resample(z->img_comp[k].linebuf, + y_bot ? r->line1 : r->line0, + y_bot ? r->line0 : r->line1, + r->w_lores, r->hs); + if (++r->ystep >= r->vs) { + r->ystep = 0; + r->line0 = r->line1; + if (++r->ypos < z->img_comp[k].y) + r->line1 += z->img_comp[k].w2; + } + } + if (n >= 3) { + stbi_uc *y = coutput[0]; + if (z->s->img_n == 3) { + if (is_rgb) { + for (i=0; i < z->s->img_x; ++i) { + out[0] = y[i]; + out[1] = coutput[1][i]; + out[2] = coutput[2][i]; + out[3] = 255; + out += n; + } + } else { + z->YCbCr_to_RGB_kernel(out, y, coutput[1], coutput[2], z->s->img_x, n); + } + } else if (z->s->img_n == 4) { + if (z->app14_color_transform == 0) { // CMYK + for (i=0; i < z->s->img_x; ++i) { + stbi_uc m = coutput[3][i]; + out[0] = stbi__blinn_8x8(coutput[0][i], m); + out[1] = stbi__blinn_8x8(coutput[1][i], m); + out[2] = stbi__blinn_8x8(coutput[2][i], m); + out[3] = 255; + out += n; + } + } else if (z->app14_color_transform == 2) { // YCCK + z->YCbCr_to_RGB_kernel(out, y, coutput[1], coutput[2], z->s->img_x, n); + for (i=0; i < z->s->img_x; ++i) { + stbi_uc m = coutput[3][i]; + out[0] = stbi__blinn_8x8(255 - out[0], m); + out[1] = stbi__blinn_8x8(255 - out[1], m); + out[2] = stbi__blinn_8x8(255 - out[2], m); + out += n; + } + } else { // YCbCr + alpha? Ignore the fourth channel for now + z->YCbCr_to_RGB_kernel(out, y, coutput[1], coutput[2], z->s->img_x, n); + } + } else + for (i=0; i < z->s->img_x; ++i) { + out[0] = out[1] = out[2] = y[i]; + out[3] = 255; // not used if n==3 + out += n; + } + } else { + if (is_rgb) { + if (n == 1) + for (i=0; i < z->s->img_x; ++i) + *out++ = stbi__compute_y(coutput[0][i], coutput[1][i], coutput[2][i]); + else { + for (i=0; i < z->s->img_x; ++i, out += 2) { + out[0] = stbi__compute_y(coutput[0][i], coutput[1][i], coutput[2][i]); + out[1] = 255; + } + } + } else if (z->s->img_n == 4 && z->app14_color_transform == 0) { + for (i=0; i < z->s->img_x; ++i) { + stbi_uc m = coutput[3][i]; + stbi_uc r = stbi__blinn_8x8(coutput[0][i], m); + stbi_uc g = stbi__blinn_8x8(coutput[1][i], m); + stbi_uc b = stbi__blinn_8x8(coutput[2][i], m); + out[0] = stbi__compute_y(r, g, b); + out[1] = 255; + out += n; + } + } else if (z->s->img_n == 4 && z->app14_color_transform == 2) { + for (i=0; i < z->s->img_x; ++i) { + out[0] = stbi__blinn_8x8(255 - coutput[0][i], coutput[3][i]); + out[1] = 255; + out += n; + } + } else { + stbi_uc *y = coutput[0]; + if (n == 1) + for (i=0; i < z->s->img_x; ++i) out[i] = y[i]; + else + for (i=0; i < z->s->img_x; ++i) { *out++ = y[i]; *out++ = 255; } + } + } + } + stbi__cleanup_jpeg(z); + *out_x = z->s->img_x; + *out_y = z->s->img_y; + if (comp) *comp = z->s->img_n >= 3 ? 3 : 1; // report original components, not output + return output; + } +} + +static void *stbi__jpeg_load(stbi__context *s, int *x, int *y, int *comp, int req_comp, stbi__result_info *ri) +{ + unsigned char* result; + stbi__jpeg* j = (stbi__jpeg*) stbi__malloc(sizeof(stbi__jpeg)); + if (!j) return stbi__errpuc("outofmem", "Out of memory"); + memset(j, 0, sizeof(stbi__jpeg)); + STBI_NOTUSED(ri); + j->s = s; + stbi__setup_jpeg(j); + result = load_jpeg_image(j, x,y,comp,req_comp); + STBI_FREE(j); + return result; +} + +static int stbi__jpeg_test(stbi__context *s) +{ + int r; + stbi__jpeg* j = (stbi__jpeg*)stbi__malloc(sizeof(stbi__jpeg)); + if (!j) return stbi__err("outofmem", "Out of memory"); + memset(j, 0, sizeof(stbi__jpeg)); + j->s = s; + stbi__setup_jpeg(j); + r = stbi__decode_jpeg_header(j, STBI__SCAN_type); + stbi__rewind(s); + STBI_FREE(j); + return r; +} + +static int stbi__jpeg_info_raw(stbi__jpeg *j, int *x, int *y, int *comp) +{ + if (!stbi__decode_jpeg_header(j, STBI__SCAN_header)) { + stbi__rewind( j->s ); + return 0; + } + if (x) *x = j->s->img_x; + if (y) *y = j->s->img_y; + if (comp) *comp = j->s->img_n >= 3 ? 3 : 1; + return 1; +} + +static int stbi__jpeg_info(stbi__context *s, int *x, int *y, int *comp) +{ + int result; + stbi__jpeg* j = (stbi__jpeg*) (stbi__malloc(sizeof(stbi__jpeg))); + if (!j) return stbi__err("outofmem", "Out of memory"); + memset(j, 0, sizeof(stbi__jpeg)); + j->s = s; + result = stbi__jpeg_info_raw(j, x, y, comp); + STBI_FREE(j); + return result; +} +#endif + +// public domain zlib decode v0.2 Sean Barrett 2006-11-18 +// simple implementation +// - all input must be provided in an upfront buffer +// - all output is written to a single output buffer (can malloc/realloc) +// performance +// - fast huffman + +#ifndef STBI_NO_ZLIB + +// fast-way is faster to check than jpeg huffman, but slow way is slower +#define STBI__ZFAST_BITS 9 // accelerate all cases in default tables +#define STBI__ZFAST_MASK ((1 << STBI__ZFAST_BITS) - 1) +#define STBI__ZNSYMS 288 // number of symbols in literal/length alphabet + +// zlib-style huffman encoding +// (jpegs packs from left, zlib from right, so can't share code) +typedef struct +{ + stbi__uint16 fast[1 << STBI__ZFAST_BITS]; + stbi__uint16 firstcode[16]; + int maxcode[17]; + stbi__uint16 firstsymbol[16]; + stbi_uc size[STBI__ZNSYMS]; + stbi__uint16 value[STBI__ZNSYMS]; +} stbi__zhuffman; + +stbi_inline static int stbi__bitreverse16(int n) +{ + n = ((n & 0xAAAA) >> 1) | ((n & 0x5555) << 1); + n = ((n & 0xCCCC) >> 2) | ((n & 0x3333) << 2); + n = ((n & 0xF0F0) >> 4) | ((n & 0x0F0F) << 4); + n = ((n & 0xFF00) >> 8) | ((n & 0x00FF) << 8); + return n; +} + +stbi_inline static int stbi__bit_reverse(int v, int bits) +{ + STBI_ASSERT(bits <= 16); + // to bit reverse n bits, reverse 16 and shift + // e.g. 11 bits, bit reverse and shift away 5 + return stbi__bitreverse16(v) >> (16-bits); +} + +static int stbi__zbuild_huffman(stbi__zhuffman *z, const stbi_uc *sizelist, int num) +{ + int i,k=0; + int code, next_code[16], sizes[17]; + + // DEFLATE spec for generating codes + memset(sizes, 0, sizeof(sizes)); + memset(z->fast, 0, sizeof(z->fast)); + for (i=0; i < num; ++i) + ++sizes[sizelist[i]]; + sizes[0] = 0; + for (i=1; i < 16; ++i) + if (sizes[i] > (1 << i)) + return stbi__err("bad sizes", "Corrupt PNG"); + code = 0; + for (i=1; i < 16; ++i) { + next_code[i] = code; + z->firstcode[i] = (stbi__uint16) code; + z->firstsymbol[i] = (stbi__uint16) k; + code = (code + sizes[i]); + if (sizes[i]) + if (code-1 >= (1 << i)) return stbi__err("bad codelengths","Corrupt PNG"); + z->maxcode[i] = code << (16-i); // preshift for inner loop + code <<= 1; + k += sizes[i]; + } + z->maxcode[16] = 0x10000; // sentinel + for (i=0; i < num; ++i) { + int s = sizelist[i]; + if (s) { + int c = next_code[s] - z->firstcode[s] + z->firstsymbol[s]; + stbi__uint16 fastv = (stbi__uint16) ((s << 9) | i); + z->size [c] = (stbi_uc ) s; + z->value[c] = (stbi__uint16) i; + if (s <= STBI__ZFAST_BITS) { + int j = stbi__bit_reverse(next_code[s],s); + while (j < (1 << STBI__ZFAST_BITS)) { + z->fast[j] = fastv; + j += (1 << s); + } + } + ++next_code[s]; + } + } + return 1; +} + +// zlib-from-memory implementation for PNG reading +// because PNG allows splitting the zlib stream arbitrarily, +// and it's annoying structurally to have PNG call ZLIB call PNG, +// we require PNG read all the IDATs and combine them into a single +// memory buffer + +typedef struct +{ + stbi_uc *zbuffer, *zbuffer_end; + int num_bits; + int hit_zeof_once; + stbi__uint32 code_buffer; + + char *zout; + char *zout_start; + char *zout_end; + int z_expandable; + + stbi__zhuffman z_length, z_distance; +} stbi__zbuf; + +stbi_inline static int stbi__zeof(stbi__zbuf *z) +{ + return (z->zbuffer >= z->zbuffer_end); +} + +stbi_inline static stbi_uc stbi__zget8(stbi__zbuf *z) +{ + return stbi__zeof(z) ? 0 : *z->zbuffer++; +} + +static void stbi__fill_bits(stbi__zbuf *z) +{ + do { + if (z->code_buffer >= (1U << z->num_bits)) { + z->zbuffer = z->zbuffer_end; /* treat this as EOF so we fail. */ + return; + } + z->code_buffer |= (unsigned int) stbi__zget8(z) << z->num_bits; + z->num_bits += 8; + } while (z->num_bits <= 24); +} + +stbi_inline static unsigned int stbi__zreceive(stbi__zbuf *z, int n) +{ + unsigned int k; + if (z->num_bits < n) stbi__fill_bits(z); + k = z->code_buffer & ((1 << n) - 1); + z->code_buffer >>= n; + z->num_bits -= n; + return k; +} + +static int stbi__zhuffman_decode_slowpath(stbi__zbuf *a, stbi__zhuffman *z) +{ + int b,s,k; + // not resolved by fast table, so compute it the slow way + // use jpeg approach, which requires MSbits at top + k = stbi__bit_reverse(a->code_buffer, 16); + for (s=STBI__ZFAST_BITS+1; ; ++s) + if (k < z->maxcode[s]) + break; + if (s >= 16) return -1; // invalid code! + // code size is s, so: + b = (k >> (16-s)) - z->firstcode[s] + z->firstsymbol[s]; + if (b >= STBI__ZNSYMS) return -1; // some data was corrupt somewhere! + if (z->size[b] != s) return -1; // was originally an assert, but report failure instead. + a->code_buffer >>= s; + a->num_bits -= s; + return z->value[b]; +} + +stbi_inline static int stbi__zhuffman_decode(stbi__zbuf *a, stbi__zhuffman *z) +{ + int b,s; + if (a->num_bits < 16) { + if (stbi__zeof(a)) { + if (!a->hit_zeof_once) { + // This is the first time we hit eof, insert 16 extra padding btis + // to allow us to keep going; if we actually consume any of them + // though, that is invalid data. This is caught later. + a->hit_zeof_once = 1; + a->num_bits += 16; // add 16 implicit zero bits + } else { + // We already inserted our extra 16 padding bits and are again + // out, this stream is actually prematurely terminated. + return -1; + } + } else { + stbi__fill_bits(a); + } + } + b = z->fast[a->code_buffer & STBI__ZFAST_MASK]; + if (b) { + s = b >> 9; + a->code_buffer >>= s; + a->num_bits -= s; + return b & 511; + } + return stbi__zhuffman_decode_slowpath(a, z); +} + +static int stbi__zexpand(stbi__zbuf *z, char *zout, int n) // need to make room for n bytes +{ + char *q; + unsigned int cur, limit, old_limit; + z->zout = zout; + if (!z->z_expandable) return stbi__err("output buffer limit","Corrupt PNG"); + cur = (unsigned int) (z->zout - z->zout_start); + limit = old_limit = (unsigned) (z->zout_end - z->zout_start); + if (UINT_MAX - cur < (unsigned) n) return stbi__err("outofmem", "Out of memory"); + while (cur + n > limit) { + if(limit > UINT_MAX / 2) return stbi__err("outofmem", "Out of memory"); + limit *= 2; + } + q = (char *) STBI_REALLOC_SIZED(z->zout_start, old_limit, limit); + STBI_NOTUSED(old_limit); + if (q == NULL) return stbi__err("outofmem", "Out of memory"); + z->zout_start = q; + z->zout = q + cur; + z->zout_end = q + limit; + return 1; +} + +static const int stbi__zlength_base[31] = { + 3,4,5,6,7,8,9,10,11,13, + 15,17,19,23,27,31,35,43,51,59, + 67,83,99,115,131,163,195,227,258,0,0 }; + +static const int stbi__zlength_extra[31]= +{ 0,0,0,0,0,0,0,0,1,1,1,1,2,2,2,2,3,3,3,3,4,4,4,4,5,5,5,5,0,0,0 }; + +static const int stbi__zdist_base[32] = { 1,2,3,4,5,7,9,13,17,25,33,49,65,97,129,193, +257,385,513,769,1025,1537,2049,3073,4097,6145,8193,12289,16385,24577,0,0}; + +static const int stbi__zdist_extra[32] = +{ 0,0,0,0,1,1,2,2,3,3,4,4,5,5,6,6,7,7,8,8,9,9,10,10,11,11,12,12,13,13}; + +static int stbi__parse_huffman_block(stbi__zbuf *a) +{ + char *zout = a->zout; + for(;;) { + int z = stbi__zhuffman_decode(a, &a->z_length); + if (z < 256) { + if (z < 0) return stbi__err("bad huffman code","Corrupt PNG"); // error in huffman codes + if (zout >= a->zout_end) { + if (!stbi__zexpand(a, zout, 1)) return 0; + zout = a->zout; + } + *zout++ = (char) z; + } else { + stbi_uc *p; + int len,dist; + if (z == 256) { + a->zout = zout; + if (a->hit_zeof_once && a->num_bits < 16) { + // The first time we hit zeof, we inserted 16 extra zero bits into our bit + // buffer so the decoder can just do its speculative decoding. But if we + // actually consumed any of those bits (which is the case when num_bits < 16), + // the stream actually read past the end so it is malformed. + return stbi__err("unexpected end","Corrupt PNG"); + } + return 1; + } + if (z >= 286) return stbi__err("bad huffman code","Corrupt PNG"); // per DEFLATE, length codes 286 and 287 must not appear in compressed data + z -= 257; + len = stbi__zlength_base[z]; + if (stbi__zlength_extra[z]) len += stbi__zreceive(a, stbi__zlength_extra[z]); + z = stbi__zhuffman_decode(a, &a->z_distance); + if (z < 0 || z >= 30) return stbi__err("bad huffman code","Corrupt PNG"); // per DEFLATE, distance codes 30 and 31 must not appear in compressed data + dist = stbi__zdist_base[z]; + if (stbi__zdist_extra[z]) dist += stbi__zreceive(a, stbi__zdist_extra[z]); + if (zout - a->zout_start < dist) return stbi__err("bad dist","Corrupt PNG"); + if (len > a->zout_end - zout) { + if (!stbi__zexpand(a, zout, len)) return 0; + zout = a->zout; + } + p = (stbi_uc *) (zout - dist); + if (dist == 1) { // run of one byte; common in images. + stbi_uc v = *p; + if (len) { do *zout++ = v; while (--len); } + } else { + if (len) { do *zout++ = *p++; while (--len); } + } + } + } +} + +static int stbi__compute_huffman_codes(stbi__zbuf *a) +{ + static const stbi_uc length_dezigzag[19] = { 16,17,18,0,8,7,9,6,10,5,11,4,12,3,13,2,14,1,15 }; + stbi__zhuffman z_codelength; + stbi_uc lencodes[286+32+137];//padding for maximum single op + stbi_uc codelength_sizes[19]; + int i,n; + + int hlit = stbi__zreceive(a,5) + 257; + int hdist = stbi__zreceive(a,5) + 1; + int hclen = stbi__zreceive(a,4) + 4; + int ntot = hlit + hdist; + + memset(codelength_sizes, 0, sizeof(codelength_sizes)); + for (i=0; i < hclen; ++i) { + int s = stbi__zreceive(a,3); + codelength_sizes[length_dezigzag[i]] = (stbi_uc) s; + } + if (!stbi__zbuild_huffman(&z_codelength, codelength_sizes, 19)) return 0; + + n = 0; + while (n < ntot) { + int c = stbi__zhuffman_decode(a, &z_codelength); + if (c < 0 || c >= 19) return stbi__err("bad codelengths", "Corrupt PNG"); + if (c < 16) + lencodes[n++] = (stbi_uc) c; + else { + stbi_uc fill = 0; + if (c == 16) { + c = stbi__zreceive(a,2)+3; + if (n == 0) return stbi__err("bad codelengths", "Corrupt PNG"); + fill = lencodes[n-1]; + } else if (c == 17) { + c = stbi__zreceive(a,3)+3; + } else if (c == 18) { + c = stbi__zreceive(a,7)+11; + } else { + return stbi__err("bad codelengths", "Corrupt PNG"); + } + if (ntot - n < c) return stbi__err("bad codelengths", "Corrupt PNG"); + memset(lencodes+n, fill, c); + n += c; + } + } + if (n != ntot) return stbi__err("bad codelengths","Corrupt PNG"); + if (!stbi__zbuild_huffman(&a->z_length, lencodes, hlit)) return 0; + if (!stbi__zbuild_huffman(&a->z_distance, lencodes+hlit, hdist)) return 0; + return 1; +} + +static int stbi__parse_uncompressed_block(stbi__zbuf *a) +{ + stbi_uc header[4]; + int len,nlen,k; + if (a->num_bits & 7) + stbi__zreceive(a, a->num_bits & 7); // discard + // drain the bit-packed data into header + k = 0; + while (a->num_bits > 0) { + header[k++] = (stbi_uc) (a->code_buffer & 255); // suppress MSVC run-time check + a->code_buffer >>= 8; + a->num_bits -= 8; + } + if (a->num_bits < 0) return stbi__err("zlib corrupt","Corrupt PNG"); + // now fill header the normal way + while (k < 4) + header[k++] = stbi__zget8(a); + len = header[1] * 256 + header[0]; + nlen = header[3] * 256 + header[2]; + if (nlen != (len ^ 0xffff)) return stbi__err("zlib corrupt","Corrupt PNG"); + if (a->zbuffer + len > a->zbuffer_end) return stbi__err("read past buffer","Corrupt PNG"); + if (a->zout + len > a->zout_end) + if (!stbi__zexpand(a, a->zout, len)) return 0; + memcpy(a->zout, a->zbuffer, len); + a->zbuffer += len; + a->zout += len; + return 1; +} + +static int stbi__parse_zlib_header(stbi__zbuf *a) +{ + int cmf = stbi__zget8(a); + int cm = cmf & 15; + /* int cinfo = cmf >> 4; */ + int flg = stbi__zget8(a); + if (stbi__zeof(a)) return stbi__err("bad zlib header","Corrupt PNG"); // zlib spec + if ((cmf*256+flg) % 31 != 0) return stbi__err("bad zlib header","Corrupt PNG"); // zlib spec + if (flg & 32) return stbi__err("no preset dict","Corrupt PNG"); // preset dictionary not allowed in png + if (cm != 8) return stbi__err("bad compression","Corrupt PNG"); // DEFLATE required for png + // window = 1 << (8 + cinfo)... but who cares, we fully buffer output + return 1; +} + +static const stbi_uc stbi__zdefault_length[STBI__ZNSYMS] = +{ + 8,8,8,8,8,8,8,8,8,8,8,8,8,8,8,8, 8,8,8,8,8,8,8,8,8,8,8,8,8,8,8,8, + 8,8,8,8,8,8,8,8,8,8,8,8,8,8,8,8, 8,8,8,8,8,8,8,8,8,8,8,8,8,8,8,8, + 8,8,8,8,8,8,8,8,8,8,8,8,8,8,8,8, 8,8,8,8,8,8,8,8,8,8,8,8,8,8,8,8, + 8,8,8,8,8,8,8,8,8,8,8,8,8,8,8,8, 8,8,8,8,8,8,8,8,8,8,8,8,8,8,8,8, + 8,8,8,8,8,8,8,8,8,8,8,8,8,8,8,8, 9,9,9,9,9,9,9,9,9,9,9,9,9,9,9,9, + 9,9,9,9,9,9,9,9,9,9,9,9,9,9,9,9, 9,9,9,9,9,9,9,9,9,9,9,9,9,9,9,9, + 9,9,9,9,9,9,9,9,9,9,9,9,9,9,9,9, 9,9,9,9,9,9,9,9,9,9,9,9,9,9,9,9, + 9,9,9,9,9,9,9,9,9,9,9,9,9,9,9,9, 9,9,9,9,9,9,9,9,9,9,9,9,9,9,9,9, + 7,7,7,7,7,7,7,7,7,7,7,7,7,7,7,7, 7,7,7,7,7,7,7,7,8,8,8,8,8,8,8,8 +}; +static const stbi_uc stbi__zdefault_distance[32] = +{ + 5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5 +}; +/* +Init algorithm: +{ + int i; // use <= to match clearly with spec + for (i=0; i <= 143; ++i) stbi__zdefault_length[i] = 8; + for ( ; i <= 255; ++i) stbi__zdefault_length[i] = 9; + for ( ; i <= 279; ++i) stbi__zdefault_length[i] = 7; + for ( ; i <= 287; ++i) stbi__zdefault_length[i] = 8; + + for (i=0; i <= 31; ++i) stbi__zdefault_distance[i] = 5; +} +*/ + +static int stbi__parse_zlib(stbi__zbuf *a, int parse_header) +{ + int final, type; + if (parse_header) + if (!stbi__parse_zlib_header(a)) return 0; + a->num_bits = 0; + a->code_buffer = 0; + a->hit_zeof_once = 0; + do { + final = stbi__zreceive(a,1); + type = stbi__zreceive(a,2); + if (type == 0) { + if (!stbi__parse_uncompressed_block(a)) return 0; + } else if (type == 3) { + return 0; + } else { + if (type == 1) { + // use fixed code lengths + if (!stbi__zbuild_huffman(&a->z_length , stbi__zdefault_length , STBI__ZNSYMS)) return 0; + if (!stbi__zbuild_huffman(&a->z_distance, stbi__zdefault_distance, 32)) return 0; + } else { + if (!stbi__compute_huffman_codes(a)) return 0; + } + if (!stbi__parse_huffman_block(a)) return 0; + } + } while (!final); + return 1; +} + +static int stbi__do_zlib(stbi__zbuf *a, char *obuf, int olen, int exp, int parse_header) +{ + a->zout_start = obuf; + a->zout = obuf; + a->zout_end = obuf + olen; + a->z_expandable = exp; + + return stbi__parse_zlib(a, parse_header); +} + +STBIDEF char *stbi_zlib_decode_malloc_guesssize(const char *buffer, int len, int initial_size, int *outlen) +{ + stbi__zbuf a; + char *p = (char *) stbi__malloc(initial_size); + if (p == NULL) return NULL; + a.zbuffer = (stbi_uc *) buffer; + a.zbuffer_end = (stbi_uc *) buffer + len; + if (stbi__do_zlib(&a, p, initial_size, 1, 1)) { + if (outlen) *outlen = (int) (a.zout - a.zout_start); + return a.zout_start; + } else { + STBI_FREE(a.zout_start); + return NULL; + } +} + +STBIDEF char *stbi_zlib_decode_malloc(char const *buffer, int len, int *outlen) +{ + return stbi_zlib_decode_malloc_guesssize(buffer, len, 16384, outlen); +} + +STBIDEF char *stbi_zlib_decode_malloc_guesssize_headerflag(const char *buffer, int len, int initial_size, int *outlen, int parse_header) +{ + stbi__zbuf a; + char *p = (char *) stbi__malloc(initial_size); + if (p == NULL) return NULL; + a.zbuffer = (stbi_uc *) buffer; + a.zbuffer_end = (stbi_uc *) buffer + len; + if (stbi__do_zlib(&a, p, initial_size, 1, parse_header)) { + if (outlen) *outlen = (int) (a.zout - a.zout_start); + return a.zout_start; + } else { + STBI_FREE(a.zout_start); + return NULL; + } +} + +STBIDEF int stbi_zlib_decode_buffer(char *obuffer, int olen, char const *ibuffer, int ilen) +{ + stbi__zbuf a; + a.zbuffer = (stbi_uc *) ibuffer; + a.zbuffer_end = (stbi_uc *) ibuffer + ilen; + if (stbi__do_zlib(&a, obuffer, olen, 0, 1)) + return (int) (a.zout - a.zout_start); + else + return -1; +} + +STBIDEF char *stbi_zlib_decode_noheader_malloc(char const *buffer, int len, int *outlen) +{ + stbi__zbuf a; + char *p = (char *) stbi__malloc(16384); + if (p == NULL) return NULL; + a.zbuffer = (stbi_uc *) buffer; + a.zbuffer_end = (stbi_uc *) buffer+len; + if (stbi__do_zlib(&a, p, 16384, 1, 0)) { + if (outlen) *outlen = (int) (a.zout - a.zout_start); + return a.zout_start; + } else { + STBI_FREE(a.zout_start); + return NULL; + } +} + +STBIDEF int stbi_zlib_decode_noheader_buffer(char *obuffer, int olen, const char *ibuffer, int ilen) +{ + stbi__zbuf a; + a.zbuffer = (stbi_uc *) ibuffer; + a.zbuffer_end = (stbi_uc *) ibuffer + ilen; + if (stbi__do_zlib(&a, obuffer, olen, 0, 0)) + return (int) (a.zout - a.zout_start); + else + return -1; +} +#endif + +// public domain "baseline" PNG decoder v0.10 Sean Barrett 2006-11-18 +// simple implementation +// - only 8-bit samples +// - no CRC checking +// - allocates lots of intermediate memory +// - avoids problem of streaming data between subsystems +// - avoids explicit window management +// performance +// - uses stb_zlib, a PD zlib implementation with fast huffman decoding + +#ifndef STBI_NO_PNG +typedef struct +{ + stbi__uint32 length; + stbi__uint32 type; +} stbi__pngchunk; + +static stbi__pngchunk stbi__get_chunk_header(stbi__context *s) +{ + stbi__pngchunk c; + c.length = stbi__get32be(s); + c.type = stbi__get32be(s); + return c; +} + +static int stbi__check_png_header(stbi__context *s) +{ + static const stbi_uc png_sig[8] = { 137,80,78,71,13,10,26,10 }; + int i; + for (i=0; i < 8; ++i) + if (stbi__get8(s) != png_sig[i]) return stbi__err("bad png sig","Not a PNG"); + return 1; +} + +typedef struct +{ + stbi__context *s; + stbi_uc *idata, *expanded, *out; + int depth; +} stbi__png; + + +enum { + STBI__F_none=0, + STBI__F_sub=1, + STBI__F_up=2, + STBI__F_avg=3, + STBI__F_paeth=4, + // synthetic filter used for first scanline to avoid needing a dummy row of 0s + STBI__F_avg_first +}; + +static stbi_uc first_row_filter[5] = +{ + STBI__F_none, + STBI__F_sub, + STBI__F_none, + STBI__F_avg_first, + STBI__F_sub // Paeth with b=c=0 turns out to be equivalent to sub +}; + +static int stbi__paeth(int a, int b, int c) +{ + // This formulation looks very different from the reference in the PNG spec, but is + // actually equivalent and has favorable data dependencies and admits straightforward + // generation of branch-free code, which helps performance significantly. + int thresh = c*3 - (a + b); + int lo = a < b ? a : b; + int hi = a < b ? b : a; + int t0 = (hi <= thresh) ? lo : c; + int t1 = (thresh <= lo) ? hi : t0; + return t1; +} + +static const stbi_uc stbi__depth_scale_table[9] = { 0, 0xff, 0x55, 0, 0x11, 0,0,0, 0x01 }; + +// adds an extra all-255 alpha channel +// dest == src is legal +// img_n must be 1 or 3 +static void stbi__create_png_alpha_expand8(stbi_uc *dest, stbi_uc *src, stbi__uint32 x, int img_n) +{ + int i; + // must process data backwards since we allow dest==src + if (img_n == 1) { + for (i=x-1; i >= 0; --i) { + dest[i*2+1] = 255; + dest[i*2+0] = src[i]; + } + } else { + STBI_ASSERT(img_n == 3); + for (i=x-1; i >= 0; --i) { + dest[i*4+3] = 255; + dest[i*4+2] = src[i*3+2]; + dest[i*4+1] = src[i*3+1]; + dest[i*4+0] = src[i*3+0]; + } + } +} + +// create the png data from post-deflated data +static int stbi__create_png_image_raw(stbi__png *a, stbi_uc *raw, stbi__uint32 raw_len, int out_n, stbi__uint32 x, stbi__uint32 y, int depth, int color) +{ + int bytes = (depth == 16 ? 2 : 1); + stbi__context *s = a->s; + stbi__uint32 i,j,stride = x*out_n*bytes; + stbi__uint32 img_len, img_width_bytes; + stbi_uc *filter_buf; + int all_ok = 1; + int k; + int img_n = s->img_n; // copy it into a local for later + + int output_bytes = out_n*bytes; + int filter_bytes = img_n*bytes; + int width = x; + + STBI_ASSERT(out_n == s->img_n || out_n == s->img_n+1); + a->out = (stbi_uc *) stbi__malloc_mad3(x, y, output_bytes, 0); // extra bytes to write off the end into + if (!a->out) return stbi__err("outofmem", "Out of memory"); + + // note: error exits here don't need to clean up a->out individually, + // stbi__do_png always does on error. + if (!stbi__mad3sizes_valid(img_n, x, depth, 7)) return stbi__err("too large", "Corrupt PNG"); + img_width_bytes = (((img_n * x * depth) + 7) >> 3); + if (!stbi__mad2sizes_valid(img_width_bytes, y, img_width_bytes)) return stbi__err("too large", "Corrupt PNG"); + img_len = (img_width_bytes + 1) * y; + + // we used to check for exact match between raw_len and img_len on non-interlaced PNGs, + // but issue #276 reported a PNG in the wild that had extra data at the end (all zeros), + // so just check for raw_len < img_len always. + if (raw_len < img_len) return stbi__err("not enough pixels","Corrupt PNG"); + + // Allocate two scan lines worth of filter workspace buffer. + filter_buf = (stbi_uc *) stbi__malloc_mad2(img_width_bytes, 2, 0); + if (!filter_buf) return stbi__err("outofmem", "Out of memory"); + + // Filtering for low-bit-depth images + if (depth < 8) { + filter_bytes = 1; + width = img_width_bytes; + } + + for (j=0; j < y; ++j) { + // cur/prior filter buffers alternate + stbi_uc *cur = filter_buf + (j & 1)*img_width_bytes; + stbi_uc *prior = filter_buf + (~j & 1)*img_width_bytes; + stbi_uc *dest = a->out + stride*j; + int nk = width * filter_bytes; + int filter = *raw++; + + // check filter type + if (filter > 4) { + all_ok = stbi__err("invalid filter","Corrupt PNG"); + break; + } + + // if first row, use special filter that doesn't sample previous row + if (j == 0) filter = first_row_filter[filter]; + + // perform actual filtering + switch (filter) { + case STBI__F_none: + memcpy(cur, raw, nk); + break; + case STBI__F_sub: + memcpy(cur, raw, filter_bytes); + for (k = filter_bytes; k < nk; ++k) + cur[k] = STBI__BYTECAST(raw[k] + cur[k-filter_bytes]); + break; + case STBI__F_up: + for (k = 0; k < nk; ++k) + cur[k] = STBI__BYTECAST(raw[k] + prior[k]); + break; + case STBI__F_avg: + for (k = 0; k < filter_bytes; ++k) + cur[k] = STBI__BYTECAST(raw[k] + (prior[k]>>1)); + for (k = filter_bytes; k < nk; ++k) + cur[k] = STBI__BYTECAST(raw[k] + ((prior[k] + cur[k-filter_bytes])>>1)); + break; + case STBI__F_paeth: + for (k = 0; k < filter_bytes; ++k) + cur[k] = STBI__BYTECAST(raw[k] + prior[k]); // prior[k] == stbi__paeth(0,prior[k],0) + for (k = filter_bytes; k < nk; ++k) + cur[k] = STBI__BYTECAST(raw[k] + stbi__paeth(cur[k-filter_bytes], prior[k], prior[k-filter_bytes])); + break; + case STBI__F_avg_first: + memcpy(cur, raw, filter_bytes); + for (k = filter_bytes; k < nk; ++k) + cur[k] = STBI__BYTECAST(raw[k] + (cur[k-filter_bytes] >> 1)); + break; + } + + raw += nk; + + // expand decoded bits in cur to dest, also adding an extra alpha channel if desired + if (depth < 8) { + stbi_uc scale = (color == 0) ? stbi__depth_scale_table[depth] : 1; // scale grayscale values to 0..255 range + stbi_uc *in = cur; + stbi_uc *out = dest; + stbi_uc inb = 0; + stbi__uint32 nsmp = x*img_n; + + // expand bits to bytes first + if (depth == 4) { + for (i=0; i < nsmp; ++i) { + if ((i & 1) == 0) inb = *in++; + *out++ = scale * (inb >> 4); + inb <<= 4; + } + } else if (depth == 2) { + for (i=0; i < nsmp; ++i) { + if ((i & 3) == 0) inb = *in++; + *out++ = scale * (inb >> 6); + inb <<= 2; + } + } else { + STBI_ASSERT(depth == 1); + for (i=0; i < nsmp; ++i) { + if ((i & 7) == 0) inb = *in++; + *out++ = scale * (inb >> 7); + inb <<= 1; + } + } + + // insert alpha=255 values if desired + if (img_n != out_n) + stbi__create_png_alpha_expand8(dest, dest, x, img_n); + } else if (depth == 8) { + if (img_n == out_n) + memcpy(dest, cur, x*img_n); + else + stbi__create_png_alpha_expand8(dest, cur, x, img_n); + } else if (depth == 16) { + // convert the image data from big-endian to platform-native + stbi__uint16 *dest16 = (stbi__uint16*)dest; + stbi__uint32 nsmp = x*img_n; + + if (img_n == out_n) { + for (i = 0; i < nsmp; ++i, ++dest16, cur += 2) + *dest16 = (cur[0] << 8) | cur[1]; + } else { + STBI_ASSERT(img_n+1 == out_n); + if (img_n == 1) { + for (i = 0; i < x; ++i, dest16 += 2, cur += 2) { + dest16[0] = (cur[0] << 8) | cur[1]; + dest16[1] = 0xffff; + } + } else { + STBI_ASSERT(img_n == 3); + for (i = 0; i < x; ++i, dest16 += 4, cur += 6) { + dest16[0] = (cur[0] << 8) | cur[1]; + dest16[1] = (cur[2] << 8) | cur[3]; + dest16[2] = (cur[4] << 8) | cur[5]; + dest16[3] = 0xffff; + } + } + } + } + } + + STBI_FREE(filter_buf); + if (!all_ok) return 0; + + return 1; +} + +static int stbi__create_png_image(stbi__png *a, stbi_uc *image_data, stbi__uint32 image_data_len, int out_n, int depth, int color, int interlaced) +{ + int bytes = (depth == 16 ? 2 : 1); + int out_bytes = out_n * bytes; + stbi_uc *final; + int p; + if (!interlaced) + return stbi__create_png_image_raw(a, image_data, image_data_len, out_n, a->s->img_x, a->s->img_y, depth, color); + + // de-interlacing + final = (stbi_uc *) stbi__malloc_mad3(a->s->img_x, a->s->img_y, out_bytes, 0); + if (!final) return stbi__err("outofmem", "Out of memory"); + for (p=0; p < 7; ++p) { + int xorig[] = { 0,4,0,2,0,1,0 }; + int yorig[] = { 0,0,4,0,2,0,1 }; + int xspc[] = { 8,8,4,4,2,2,1 }; + int yspc[] = { 8,8,8,4,4,2,2 }; + int i,j,x,y; + // pass1_x[4] = 0, pass1_x[5] = 1, pass1_x[12] = 1 + x = (a->s->img_x - xorig[p] + xspc[p]-1) / xspc[p]; + y = (a->s->img_y - yorig[p] + yspc[p]-1) / yspc[p]; + if (x && y) { + stbi__uint32 img_len = ((((a->s->img_n * x * depth) + 7) >> 3) + 1) * y; + if (!stbi__create_png_image_raw(a, image_data, image_data_len, out_n, x, y, depth, color)) { + STBI_FREE(final); + return 0; + } + for (j=0; j < y; ++j) { + for (i=0; i < x; ++i) { + int out_y = j*yspc[p]+yorig[p]; + int out_x = i*xspc[p]+xorig[p]; + memcpy(final + out_y*a->s->img_x*out_bytes + out_x*out_bytes, + a->out + (j*x+i)*out_bytes, out_bytes); + } + } + STBI_FREE(a->out); + image_data += img_len; + image_data_len -= img_len; + } + } + a->out = final; + + return 1; +} + +static int stbi__compute_transparency(stbi__png *z, stbi_uc tc[3], int out_n) +{ + stbi__context *s = z->s; + stbi__uint32 i, pixel_count = s->img_x * s->img_y; + stbi_uc *p = z->out; + + // compute color-based transparency, assuming we've + // already got 255 as the alpha value in the output + STBI_ASSERT(out_n == 2 || out_n == 4); + + if (out_n == 2) { + for (i=0; i < pixel_count; ++i) { + p[1] = (p[0] == tc[0] ? 0 : 255); + p += 2; + } + } else { + for (i=0; i < pixel_count; ++i) { + if (p[0] == tc[0] && p[1] == tc[1] && p[2] == tc[2]) + p[3] = 0; + p += 4; + } + } + return 1; +} + +static int stbi__compute_transparency16(stbi__png *z, stbi__uint16 tc[3], int out_n) +{ + stbi__context *s = z->s; + stbi__uint32 i, pixel_count = s->img_x * s->img_y; + stbi__uint16 *p = (stbi__uint16*) z->out; + + // compute color-based transparency, assuming we've + // already got 65535 as the alpha value in the output + STBI_ASSERT(out_n == 2 || out_n == 4); + + if (out_n == 2) { + for (i = 0; i < pixel_count; ++i) { + p[1] = (p[0] == tc[0] ? 0 : 65535); + p += 2; + } + } else { + for (i = 0; i < pixel_count; ++i) { + if (p[0] == tc[0] && p[1] == tc[1] && p[2] == tc[2]) + p[3] = 0; + p += 4; + } + } + return 1; +} + +static int stbi__expand_png_palette(stbi__png *a, stbi_uc *palette, int len, int pal_img_n) +{ + stbi__uint32 i, pixel_count = a->s->img_x * a->s->img_y; + stbi_uc *p, *temp_out, *orig = a->out; + + p = (stbi_uc *) stbi__malloc_mad2(pixel_count, pal_img_n, 0); + if (p == NULL) return stbi__err("outofmem", "Out of memory"); + + // between here and free(out) below, exitting would leak + temp_out = p; + + if (pal_img_n == 3) { + for (i=0; i < pixel_count; ++i) { + int n = orig[i]*4; + p[0] = palette[n ]; + p[1] = palette[n+1]; + p[2] = palette[n+2]; + p += 3; + } + } else { + for (i=0; i < pixel_count; ++i) { + int n = orig[i]*4; + p[0] = palette[n ]; + p[1] = palette[n+1]; + p[2] = palette[n+2]; + p[3] = palette[n+3]; + p += 4; + } + } + STBI_FREE(a->out); + a->out = temp_out; + + STBI_NOTUSED(len); + + return 1; +} + +static int stbi__unpremultiply_on_load_global = 0; +static int stbi__de_iphone_flag_global = 0; + +STBIDEF void stbi_set_unpremultiply_on_load(int flag_true_if_should_unpremultiply) +{ + stbi__unpremultiply_on_load_global = flag_true_if_should_unpremultiply; +} + +STBIDEF void stbi_convert_iphone_png_to_rgb(int flag_true_if_should_convert) +{ + stbi__de_iphone_flag_global = flag_true_if_should_convert; +} + +#ifndef STBI_THREAD_LOCAL +#define stbi__unpremultiply_on_load stbi__unpremultiply_on_load_global +#define stbi__de_iphone_flag stbi__de_iphone_flag_global +#else +static STBI_THREAD_LOCAL int stbi__unpremultiply_on_load_local, stbi__unpremultiply_on_load_set; +static STBI_THREAD_LOCAL int stbi__de_iphone_flag_local, stbi__de_iphone_flag_set; + +STBIDEF void stbi_set_unpremultiply_on_load_thread(int flag_true_if_should_unpremultiply) +{ + stbi__unpremultiply_on_load_local = flag_true_if_should_unpremultiply; + stbi__unpremultiply_on_load_set = 1; +} + +STBIDEF void stbi_convert_iphone_png_to_rgb_thread(int flag_true_if_should_convert) +{ + stbi__de_iphone_flag_local = flag_true_if_should_convert; + stbi__de_iphone_flag_set = 1; +} + +#define stbi__unpremultiply_on_load (stbi__unpremultiply_on_load_set \ + ? stbi__unpremultiply_on_load_local \ + : stbi__unpremultiply_on_load_global) +#define stbi__de_iphone_flag (stbi__de_iphone_flag_set \ + ? stbi__de_iphone_flag_local \ + : stbi__de_iphone_flag_global) +#endif // STBI_THREAD_LOCAL + +static void stbi__de_iphone(stbi__png *z) +{ + stbi__context *s = z->s; + stbi__uint32 i, pixel_count = s->img_x * s->img_y; + stbi_uc *p = z->out; + + if (s->img_out_n == 3) { // convert bgr to rgb + for (i=0; i < pixel_count; ++i) { + stbi_uc t = p[0]; + p[0] = p[2]; + p[2] = t; + p += 3; + } + } else { + STBI_ASSERT(s->img_out_n == 4); + if (stbi__unpremultiply_on_load) { + // convert bgr to rgb and unpremultiply + for (i=0; i < pixel_count; ++i) { + stbi_uc a = p[3]; + stbi_uc t = p[0]; + if (a) { + stbi_uc half = a / 2; + p[0] = (p[2] * 255 + half) / a; + p[1] = (p[1] * 255 + half) / a; + p[2] = ( t * 255 + half) / a; + } else { + p[0] = p[2]; + p[2] = t; + } + p += 4; + } + } else { + // convert bgr to rgb + for (i=0; i < pixel_count; ++i) { + stbi_uc t = p[0]; + p[0] = p[2]; + p[2] = t; + p += 4; + } + } + } +} + +#define STBI__PNG_TYPE(a,b,c,d) (((unsigned) (a) << 24) + ((unsigned) (b) << 16) + ((unsigned) (c) << 8) + (unsigned) (d)) + +static int stbi__parse_png_file(stbi__png *z, int scan, int req_comp) +{ + stbi_uc palette[1024], pal_img_n=0; + stbi_uc has_trans=0, tc[3]={0}; + stbi__uint16 tc16[3]; + stbi__uint32 ioff=0, idata_limit=0, i, pal_len=0; + int first=1,k,interlace=0, color=0, is_iphone=0; + stbi__context *s = z->s; + + z->expanded = NULL; + z->idata = NULL; + z->out = NULL; + + if (!stbi__check_png_header(s)) return 0; + + if (scan == STBI__SCAN_type) return 1; + + for (;;) { + stbi__pngchunk c = stbi__get_chunk_header(s); + switch (c.type) { + case STBI__PNG_TYPE('C','g','B','I'): + is_iphone = 1; + stbi__skip(s, c.length); + break; + case STBI__PNG_TYPE('I','H','D','R'): { + int comp,filter; + if (!first) return stbi__err("multiple IHDR","Corrupt PNG"); + first = 0; + if (c.length != 13) return stbi__err("bad IHDR len","Corrupt PNG"); + s->img_x = stbi__get32be(s); + s->img_y = stbi__get32be(s); + if (s->img_y > STBI_MAX_DIMENSIONS) return stbi__err("too large","Very large image (corrupt?)"); + if (s->img_x > STBI_MAX_DIMENSIONS) return stbi__err("too large","Very large image (corrupt?)"); + z->depth = stbi__get8(s); if (z->depth != 1 && z->depth != 2 && z->depth != 4 && z->depth != 8 && z->depth != 16) return stbi__err("1/2/4/8/16-bit only","PNG not supported: 1/2/4/8/16-bit only"); + color = stbi__get8(s); if (color > 6) return stbi__err("bad ctype","Corrupt PNG"); + if (color == 3 && z->depth == 16) return stbi__err("bad ctype","Corrupt PNG"); + if (color == 3) pal_img_n = 3; else if (color & 1) return stbi__err("bad ctype","Corrupt PNG"); + comp = stbi__get8(s); if (comp) return stbi__err("bad comp method","Corrupt PNG"); + filter= stbi__get8(s); if (filter) return stbi__err("bad filter method","Corrupt PNG"); + interlace = stbi__get8(s); if (interlace>1) return stbi__err("bad interlace method","Corrupt PNG"); + if (!s->img_x || !s->img_y) return stbi__err("0-pixel image","Corrupt PNG"); + if (!pal_img_n) { + s->img_n = (color & 2 ? 3 : 1) + (color & 4 ? 1 : 0); + if ((1 << 30) / s->img_x / s->img_n < s->img_y) return stbi__err("too large", "Image too large to decode"); + } else { + // if paletted, then pal_n is our final components, and + // img_n is # components to decompress/filter. + s->img_n = 1; + if ((1 << 30) / s->img_x / 4 < s->img_y) return stbi__err("too large","Corrupt PNG"); + } + // even with SCAN_header, have to scan to see if we have a tRNS + break; + } + + case STBI__PNG_TYPE('P','L','T','E'): { + if (first) return stbi__err("first not IHDR", "Corrupt PNG"); + if (c.length > 256*3) return stbi__err("invalid PLTE","Corrupt PNG"); + pal_len = c.length / 3; + if (pal_len * 3 != c.length) return stbi__err("invalid PLTE","Corrupt PNG"); + for (i=0; i < pal_len; ++i) { + palette[i*4+0] = stbi__get8(s); + palette[i*4+1] = stbi__get8(s); + palette[i*4+2] = stbi__get8(s); + palette[i*4+3] = 255; + } + break; + } + + case STBI__PNG_TYPE('t','R','N','S'): { + if (first) return stbi__err("first not IHDR", "Corrupt PNG"); + if (z->idata) return stbi__err("tRNS after IDAT","Corrupt PNG"); + if (pal_img_n) { + if (scan == STBI__SCAN_header) { s->img_n = 4; return 1; } + if (pal_len == 0) return stbi__err("tRNS before PLTE","Corrupt PNG"); + if (c.length > pal_len) return stbi__err("bad tRNS len","Corrupt PNG"); + pal_img_n = 4; + for (i=0; i < c.length; ++i) + palette[i*4+3] = stbi__get8(s); + } else { + if (!(s->img_n & 1)) return stbi__err("tRNS with alpha","Corrupt PNG"); + if (c.length != (stbi__uint32) s->img_n*2) return stbi__err("bad tRNS len","Corrupt PNG"); + has_trans = 1; + // non-paletted with tRNS = constant alpha. if header-scanning, we can stop now. + if (scan == STBI__SCAN_header) { ++s->img_n; return 1; } + if (z->depth == 16) { + for (k = 0; k < s->img_n && k < 3; ++k) // extra loop test to suppress false GCC warning + tc16[k] = (stbi__uint16)stbi__get16be(s); // copy the values as-is + } else { + for (k = 0; k < s->img_n && k < 3; ++k) + tc[k] = (stbi_uc)(stbi__get16be(s) & 255) * stbi__depth_scale_table[z->depth]; // non 8-bit images will be larger + } + } + break; + } + + case STBI__PNG_TYPE('I','D','A','T'): { + if (first) return stbi__err("first not IHDR", "Corrupt PNG"); + if (pal_img_n && !pal_len) return stbi__err("no PLTE","Corrupt PNG"); + if (scan == STBI__SCAN_header) { + // header scan definitely stops at first IDAT + if (pal_img_n) + s->img_n = pal_img_n; + return 1; + } + if (c.length > (1u << 30)) return stbi__err("IDAT size limit", "IDAT section larger than 2^30 bytes"); + if ((int)(ioff + c.length) < (int)ioff) return 0; + if (ioff + c.length > idata_limit) { + stbi__uint32 idata_limit_old = idata_limit; + stbi_uc *p; + if (idata_limit == 0) idata_limit = c.length > 4096 ? c.length : 4096; + while (ioff + c.length > idata_limit) + idata_limit *= 2; + STBI_NOTUSED(idata_limit_old); + p = (stbi_uc *) STBI_REALLOC_SIZED(z->idata, idata_limit_old, idata_limit); if (p == NULL) return stbi__err("outofmem", "Out of memory"); + z->idata = p; + } + if (!stbi__getn(s, z->idata+ioff,c.length)) return stbi__err("outofdata","Corrupt PNG"); + ioff += c.length; + break; + } + + case STBI__PNG_TYPE('I','E','N','D'): { + stbi__uint32 raw_len, bpl; + if (first) return stbi__err("first not IHDR", "Corrupt PNG"); + if (scan != STBI__SCAN_load) return 1; + if (z->idata == NULL) return stbi__err("no IDAT","Corrupt PNG"); + // initial guess for decoded data size to avoid unnecessary reallocs + bpl = (s->img_x * z->depth + 7) / 8; // bytes per line, per component + raw_len = bpl * s->img_y * s->img_n /* pixels */ + s->img_y /* filter mode per row */; + z->expanded = (stbi_uc *) stbi_zlib_decode_malloc_guesssize_headerflag((char *) z->idata, ioff, raw_len, (int *) &raw_len, !is_iphone); + if (z->expanded == NULL) return 0; // zlib should set error + STBI_FREE(z->idata); z->idata = NULL; + if ((req_comp == s->img_n+1 && req_comp != 3 && !pal_img_n) || has_trans) + s->img_out_n = s->img_n+1; + else + s->img_out_n = s->img_n; + if (!stbi__create_png_image(z, z->expanded, raw_len, s->img_out_n, z->depth, color, interlace)) return 0; + if (has_trans) { + if (z->depth == 16) { + if (!stbi__compute_transparency16(z, tc16, s->img_out_n)) return 0; + } else { + if (!stbi__compute_transparency(z, tc, s->img_out_n)) return 0; + } + } + if (is_iphone && stbi__de_iphone_flag && s->img_out_n > 2) + stbi__de_iphone(z); + if (pal_img_n) { + // pal_img_n == 3 or 4 + s->img_n = pal_img_n; // record the actual colors we had + s->img_out_n = pal_img_n; + if (req_comp >= 3) s->img_out_n = req_comp; + if (!stbi__expand_png_palette(z, palette, pal_len, s->img_out_n)) + return 0; + } else if (has_trans) { + // non-paletted image with tRNS -> source image has (constant) alpha + ++s->img_n; + } + STBI_FREE(z->expanded); z->expanded = NULL; + // end of PNG chunk, read and skip CRC + stbi__get32be(s); + return 1; + } + + default: + // if critical, fail + if (first) return stbi__err("first not IHDR", "Corrupt PNG"); + if ((c.type & (1 << 29)) == 0) { + #ifndef STBI_NO_FAILURE_STRINGS + // not threadsafe + static char invalid_chunk[] = "XXXX PNG chunk not known"; + invalid_chunk[0] = STBI__BYTECAST(c.type >> 24); + invalid_chunk[1] = STBI__BYTECAST(c.type >> 16); + invalid_chunk[2] = STBI__BYTECAST(c.type >> 8); + invalid_chunk[3] = STBI__BYTECAST(c.type >> 0); + #endif + return stbi__err(invalid_chunk, "PNG not supported: unknown PNG chunk type"); + } + stbi__skip(s, c.length); + break; + } + // end of PNG chunk, read and skip CRC + stbi__get32be(s); + } +} + +static void *stbi__do_png(stbi__png *p, int *x, int *y, int *n, int req_comp, stbi__result_info *ri) +{ + void *result=NULL; + if (req_comp < 0 || req_comp > 4) return stbi__errpuc("bad req_comp", "Internal error"); + if (stbi__parse_png_file(p, STBI__SCAN_load, req_comp)) { + if (p->depth <= 8) + ri->bits_per_channel = 8; + else if (p->depth == 16) + ri->bits_per_channel = 16; + else + return stbi__errpuc("bad bits_per_channel", "PNG not supported: unsupported color depth"); + result = p->out; + p->out = NULL; + if (req_comp && req_comp != p->s->img_out_n) { + if (ri->bits_per_channel == 8) + result = stbi__convert_format((unsigned char *) result, p->s->img_out_n, req_comp, p->s->img_x, p->s->img_y); + else + result = stbi__convert_format16((stbi__uint16 *) result, p->s->img_out_n, req_comp, p->s->img_x, p->s->img_y); + p->s->img_out_n = req_comp; + if (result == NULL) return result; + } + *x = p->s->img_x; + *y = p->s->img_y; + if (n) *n = p->s->img_n; + } + STBI_FREE(p->out); p->out = NULL; + STBI_FREE(p->expanded); p->expanded = NULL; + STBI_FREE(p->idata); p->idata = NULL; + + return result; +} + +static void *stbi__png_load(stbi__context *s, int *x, int *y, int *comp, int req_comp, stbi__result_info *ri) +{ + stbi__png p; + p.s = s; + return stbi__do_png(&p, x,y,comp,req_comp, ri); +} + +static int stbi__png_test(stbi__context *s) +{ + int r; + r = stbi__check_png_header(s); + stbi__rewind(s); + return r; +} + +static int stbi__png_info_raw(stbi__png *p, int *x, int *y, int *comp) +{ + if (!stbi__parse_png_file(p, STBI__SCAN_header, 0)) { + stbi__rewind( p->s ); + return 0; + } + if (x) *x = p->s->img_x; + if (y) *y = p->s->img_y; + if (comp) *comp = p->s->img_n; + return 1; +} + +static int stbi__png_info(stbi__context *s, int *x, int *y, int *comp) +{ + stbi__png p; + p.s = s; + return stbi__png_info_raw(&p, x, y, comp); +} + +static int stbi__png_is16(stbi__context *s) +{ + stbi__png p; + p.s = s; + if (!stbi__png_info_raw(&p, NULL, NULL, NULL)) + return 0; + if (p.depth != 16) { + stbi__rewind(p.s); + return 0; + } + return 1; +} +#endif + +// Microsoft/Windows BMP image + +#ifndef STBI_NO_BMP +static int stbi__bmp_test_raw(stbi__context *s) +{ + int r; + int sz; + if (stbi__get8(s) != 'B') return 0; + if (stbi__get8(s) != 'M') return 0; + stbi__get32le(s); // discard filesize + stbi__get16le(s); // discard reserved + stbi__get16le(s); // discard reserved + stbi__get32le(s); // discard data offset + sz = stbi__get32le(s); + r = (sz == 12 || sz == 40 || sz == 56 || sz == 108 || sz == 124); + return r; +} + +static int stbi__bmp_test(stbi__context *s) +{ + int r = stbi__bmp_test_raw(s); + stbi__rewind(s); + return r; +} + + +// returns 0..31 for the highest set bit +static int stbi__high_bit(unsigned int z) +{ + int n=0; + if (z == 0) return -1; + if (z >= 0x10000) { n += 16; z >>= 16; } + if (z >= 0x00100) { n += 8; z >>= 8; } + if (z >= 0x00010) { n += 4; z >>= 4; } + if (z >= 0x00004) { n += 2; z >>= 2; } + if (z >= 0x00002) { n += 1;/* >>= 1;*/ } + return n; +} + +static int stbi__bitcount(unsigned int a) +{ + a = (a & 0x55555555) + ((a >> 1) & 0x55555555); // max 2 + a = (a & 0x33333333) + ((a >> 2) & 0x33333333); // max 4 + a = (a + (a >> 4)) & 0x0f0f0f0f; // max 8 per 4, now 8 bits + a = (a + (a >> 8)); // max 16 per 8 bits + a = (a + (a >> 16)); // max 32 per 8 bits + return a & 0xff; +} + +// extract an arbitrarily-aligned N-bit value (N=bits) +// from v, and then make it 8-bits long and fractionally +// extend it to full full range. +static int stbi__shiftsigned(unsigned int v, int shift, int bits) +{ + static unsigned int mul_table[9] = { + 0, + 0xff/*0b11111111*/, 0x55/*0b01010101*/, 0x49/*0b01001001*/, 0x11/*0b00010001*/, + 0x21/*0b00100001*/, 0x41/*0b01000001*/, 0x81/*0b10000001*/, 0x01/*0b00000001*/, + }; + static unsigned int shift_table[9] = { + 0, 0,0,1,0,2,4,6,0, + }; + if (shift < 0) + v <<= -shift; + else + v >>= shift; + STBI_ASSERT(v < 256); + v >>= (8-bits); + STBI_ASSERT(bits >= 0 && bits <= 8); + return (int) ((unsigned) v * mul_table[bits]) >> shift_table[bits]; +} + +typedef struct +{ + int bpp, offset, hsz; + unsigned int mr,mg,mb,ma, all_a; + int extra_read; +} stbi__bmp_data; + +static int stbi__bmp_set_mask_defaults(stbi__bmp_data *info, int compress) +{ + // BI_BITFIELDS specifies masks explicitly, don't override + if (compress == 3) + return 1; + + if (compress == 0) { + if (info->bpp == 16) { + info->mr = 31u << 10; + info->mg = 31u << 5; + info->mb = 31u << 0; + } else if (info->bpp == 32) { + info->mr = 0xffu << 16; + info->mg = 0xffu << 8; + info->mb = 0xffu << 0; + info->ma = 0xffu << 24; + info->all_a = 0; // if all_a is 0 at end, then we loaded alpha channel but it was all 0 + } else { + // otherwise, use defaults, which is all-0 + info->mr = info->mg = info->mb = info->ma = 0; + } + return 1; + } + return 0; // error +} + +static void *stbi__bmp_parse_header(stbi__context *s, stbi__bmp_data *info) +{ + int hsz; + if (stbi__get8(s) != 'B' || stbi__get8(s) != 'M') return stbi__errpuc("not BMP", "Corrupt BMP"); + stbi__get32le(s); // discard filesize + stbi__get16le(s); // discard reserved + stbi__get16le(s); // discard reserved + info->offset = stbi__get32le(s); + info->hsz = hsz = stbi__get32le(s); + info->mr = info->mg = info->mb = info->ma = 0; + info->extra_read = 14; + + if (info->offset < 0) return stbi__errpuc("bad BMP", "bad BMP"); + + if (hsz != 12 && hsz != 40 && hsz != 56 && hsz != 108 && hsz != 124) return stbi__errpuc("unknown BMP", "BMP type not supported: unknown"); + if (hsz == 12) { + s->img_x = stbi__get16le(s); + s->img_y = stbi__get16le(s); + } else { + s->img_x = stbi__get32le(s); + s->img_y = stbi__get32le(s); + } + if (stbi__get16le(s) != 1) return stbi__errpuc("bad BMP", "bad BMP"); + info->bpp = stbi__get16le(s); + if (hsz != 12) { + int compress = stbi__get32le(s); + if (compress == 1 || compress == 2) return stbi__errpuc("BMP RLE", "BMP type not supported: RLE"); + if (compress >= 4) return stbi__errpuc("BMP JPEG/PNG", "BMP type not supported: unsupported compression"); // this includes PNG/JPEG modes + if (compress == 3 && info->bpp != 16 && info->bpp != 32) return stbi__errpuc("bad BMP", "bad BMP"); // bitfields requires 16 or 32 bits/pixel + stbi__get32le(s); // discard sizeof + stbi__get32le(s); // discard hres + stbi__get32le(s); // discard vres + stbi__get32le(s); // discard colorsused + stbi__get32le(s); // discard max important + if (hsz == 40 || hsz == 56) { + if (hsz == 56) { + stbi__get32le(s); + stbi__get32le(s); + stbi__get32le(s); + stbi__get32le(s); + } + if (info->bpp == 16 || info->bpp == 32) { + if (compress == 0) { + stbi__bmp_set_mask_defaults(info, compress); + } else if (compress == 3) { + info->mr = stbi__get32le(s); + info->mg = stbi__get32le(s); + info->mb = stbi__get32le(s); + info->extra_read += 12; + // not documented, but generated by photoshop and handled by mspaint + if (info->mr == info->mg && info->mg == info->mb) { + // ?!?!? + return stbi__errpuc("bad BMP", "bad BMP"); + } + } else + return stbi__errpuc("bad BMP", "bad BMP"); + } + } else { + // V4/V5 header + int i; + if (hsz != 108 && hsz != 124) + return stbi__errpuc("bad BMP", "bad BMP"); + info->mr = stbi__get32le(s); + info->mg = stbi__get32le(s); + info->mb = stbi__get32le(s); + info->ma = stbi__get32le(s); + if (compress != 3) // override mr/mg/mb unless in BI_BITFIELDS mode, as per docs + stbi__bmp_set_mask_defaults(info, compress); + stbi__get32le(s); // discard color space + for (i=0; i < 12; ++i) + stbi__get32le(s); // discard color space parameters + if (hsz == 124) { + stbi__get32le(s); // discard rendering intent + stbi__get32le(s); // discard offset of profile data + stbi__get32le(s); // discard size of profile data + stbi__get32le(s); // discard reserved + } + } + } + return (void *) 1; +} + + +static void *stbi__bmp_load(stbi__context *s, int *x, int *y, int *comp, int req_comp, stbi__result_info *ri) +{ + stbi_uc *out; + unsigned int mr=0,mg=0,mb=0,ma=0, all_a; + stbi_uc pal[256][4]; + int psize=0,i,j,width; + int flip_vertically, pad, target; + stbi__bmp_data info; + STBI_NOTUSED(ri); + + info.all_a = 255; + if (stbi__bmp_parse_header(s, &info) == NULL) + return NULL; // error code already set + + flip_vertically = ((int) s->img_y) > 0; + s->img_y = abs((int) s->img_y); + + if (s->img_y > STBI_MAX_DIMENSIONS) return stbi__errpuc("too large","Very large image (corrupt?)"); + if (s->img_x > STBI_MAX_DIMENSIONS) return stbi__errpuc("too large","Very large image (corrupt?)"); + + mr = info.mr; + mg = info.mg; + mb = info.mb; + ma = info.ma; + all_a = info.all_a; + + if (info.hsz == 12) { + if (info.bpp < 24) + psize = (info.offset - info.extra_read - 24) / 3; + } else { + if (info.bpp < 16) + psize = (info.offset - info.extra_read - info.hsz) >> 2; + } + if (psize == 0) { + // accept some number of extra bytes after the header, but if the offset points either to before + // the header ends or implies a large amount of extra data, reject the file as malformed + int bytes_read_so_far = s->callback_already_read + (int)(s->img_buffer - s->img_buffer_original); + int header_limit = 1024; // max we actually read is below 256 bytes currently. + int extra_data_limit = 256*4; // what ordinarily goes here is a palette; 256 entries*4 bytes is its max size. + if (bytes_read_so_far <= 0 || bytes_read_so_far > header_limit) { + return stbi__errpuc("bad header", "Corrupt BMP"); + } + // we established that bytes_read_so_far is positive and sensible. + // the first half of this test rejects offsets that are either too small positives, or + // negative, and guarantees that info.offset >= bytes_read_so_far > 0. this in turn + // ensures the number computed in the second half of the test can't overflow. + if (info.offset < bytes_read_so_far || info.offset - bytes_read_so_far > extra_data_limit) { + return stbi__errpuc("bad offset", "Corrupt BMP"); + } else { + stbi__skip(s, info.offset - bytes_read_so_far); + } + } + + if (info.bpp == 24 && ma == 0xff000000) + s->img_n = 3; + else + s->img_n = ma ? 4 : 3; + if (req_comp && req_comp >= 3) // we can directly decode 3 or 4 + target = req_comp; + else + target = s->img_n; // if they want monochrome, we'll post-convert + + // sanity-check size + if (!stbi__mad3sizes_valid(target, s->img_x, s->img_y, 0)) + return stbi__errpuc("too large", "Corrupt BMP"); + + out = (stbi_uc *) stbi__malloc_mad3(target, s->img_x, s->img_y, 0); + if (!out) return stbi__errpuc("outofmem", "Out of memory"); + if (info.bpp < 16) { + int z=0; + if (psize == 0 || psize > 256) { STBI_FREE(out); return stbi__errpuc("invalid", "Corrupt BMP"); } + for (i=0; i < psize; ++i) { + pal[i][2] = stbi__get8(s); + pal[i][1] = stbi__get8(s); + pal[i][0] = stbi__get8(s); + if (info.hsz != 12) stbi__get8(s); + pal[i][3] = 255; + } + stbi__skip(s, info.offset - info.extra_read - info.hsz - psize * (info.hsz == 12 ? 3 : 4)); + if (info.bpp == 1) width = (s->img_x + 7) >> 3; + else if (info.bpp == 4) width = (s->img_x + 1) >> 1; + else if (info.bpp == 8) width = s->img_x; + else { STBI_FREE(out); return stbi__errpuc("bad bpp", "Corrupt BMP"); } + pad = (-width)&3; + if (info.bpp == 1) { + for (j=0; j < (int) s->img_y; ++j) { + int bit_offset = 7, v = stbi__get8(s); + for (i=0; i < (int) s->img_x; ++i) { + int color = (v>>bit_offset)&0x1; + out[z++] = pal[color][0]; + out[z++] = pal[color][1]; + out[z++] = pal[color][2]; + if (target == 4) out[z++] = 255; + if (i+1 == (int) s->img_x) break; + if((--bit_offset) < 0) { + bit_offset = 7; + v = stbi__get8(s); + } + } + stbi__skip(s, pad); + } + } else { + for (j=0; j < (int) s->img_y; ++j) { + for (i=0; i < (int) s->img_x; i += 2) { + int v=stbi__get8(s),v2=0; + if (info.bpp == 4) { + v2 = v & 15; + v >>= 4; + } + out[z++] = pal[v][0]; + out[z++] = pal[v][1]; + out[z++] = pal[v][2]; + if (target == 4) out[z++] = 255; + if (i+1 == (int) s->img_x) break; + v = (info.bpp == 8) ? stbi__get8(s) : v2; + out[z++] = pal[v][0]; + out[z++] = pal[v][1]; + out[z++] = pal[v][2]; + if (target == 4) out[z++] = 255; + } + stbi__skip(s, pad); + } + } + } else { + int rshift=0,gshift=0,bshift=0,ashift=0,rcount=0,gcount=0,bcount=0,acount=0; + int z = 0; + int easy=0; + stbi__skip(s, info.offset - info.extra_read - info.hsz); + if (info.bpp == 24) width = 3 * s->img_x; + else if (info.bpp == 16) width = 2*s->img_x; + else /* bpp = 32 and pad = 0 */ width=0; + pad = (-width) & 3; + if (info.bpp == 24) { + easy = 1; + } else if (info.bpp == 32) { + if (mb == 0xff && mg == 0xff00 && mr == 0x00ff0000 && ma == 0xff000000) + easy = 2; + } + if (!easy) { + if (!mr || !mg || !mb) { STBI_FREE(out); return stbi__errpuc("bad masks", "Corrupt BMP"); } + // right shift amt to put high bit in position #7 + rshift = stbi__high_bit(mr)-7; rcount = stbi__bitcount(mr); + gshift = stbi__high_bit(mg)-7; gcount = stbi__bitcount(mg); + bshift = stbi__high_bit(mb)-7; bcount = stbi__bitcount(mb); + ashift = stbi__high_bit(ma)-7; acount = stbi__bitcount(ma); + if (rcount > 8 || gcount > 8 || bcount > 8 || acount > 8) { STBI_FREE(out); return stbi__errpuc("bad masks", "Corrupt BMP"); } + } + for (j=0; j < (int) s->img_y; ++j) { + if (easy) { + for (i=0; i < (int) s->img_x; ++i) { + unsigned char a; + out[z+2] = stbi__get8(s); + out[z+1] = stbi__get8(s); + out[z+0] = stbi__get8(s); + z += 3; + a = (easy == 2 ? stbi__get8(s) : 255); + all_a |= a; + if (target == 4) out[z++] = a; + } + } else { + int bpp = info.bpp; + for (i=0; i < (int) s->img_x; ++i) { + stbi__uint32 v = (bpp == 16 ? (stbi__uint32) stbi__get16le(s) : stbi__get32le(s)); + unsigned int a; + out[z++] = STBI__BYTECAST(stbi__shiftsigned(v & mr, rshift, rcount)); + out[z++] = STBI__BYTECAST(stbi__shiftsigned(v & mg, gshift, gcount)); + out[z++] = STBI__BYTECAST(stbi__shiftsigned(v & mb, bshift, bcount)); + a = (ma ? stbi__shiftsigned(v & ma, ashift, acount) : 255); + all_a |= a; + if (target == 4) out[z++] = STBI__BYTECAST(a); + } + } + stbi__skip(s, pad); + } + } + + // if alpha channel is all 0s, replace with all 255s + if (target == 4 && all_a == 0) + for (i=4*s->img_x*s->img_y-1; i >= 0; i -= 4) + out[i] = 255; + + if (flip_vertically) { + stbi_uc t; + for (j=0; j < (int) s->img_y>>1; ++j) { + stbi_uc *p1 = out + j *s->img_x*target; + stbi_uc *p2 = out + (s->img_y-1-j)*s->img_x*target; + for (i=0; i < (int) s->img_x*target; ++i) { + t = p1[i]; p1[i] = p2[i]; p2[i] = t; + } + } + } + + if (req_comp && req_comp != target) { + out = stbi__convert_format(out, target, req_comp, s->img_x, s->img_y); + if (out == NULL) return out; // stbi__convert_format frees input on failure + } + + *x = s->img_x; + *y = s->img_y; + if (comp) *comp = s->img_n; + return out; +} +#endif + +// Targa Truevision - TGA +// by Jonathan Dummer +#ifndef STBI_NO_TGA +// returns STBI_rgb or whatever, 0 on error +static int stbi__tga_get_comp(int bits_per_pixel, int is_grey, int* is_rgb16) +{ + // only RGB or RGBA (incl. 16bit) or grey allowed + if (is_rgb16) *is_rgb16 = 0; + switch(bits_per_pixel) { + case 8: return STBI_grey; + case 16: if(is_grey) return STBI_grey_alpha; + // fallthrough + case 15: if(is_rgb16) *is_rgb16 = 1; + return STBI_rgb; + case 24: // fallthrough + case 32: return bits_per_pixel/8; + default: return 0; + } +} + +static int stbi__tga_info(stbi__context *s, int *x, int *y, int *comp) +{ + int tga_w, tga_h, tga_comp, tga_image_type, tga_bits_per_pixel, tga_colormap_bpp; + int sz, tga_colormap_type; + stbi__get8(s); // discard Offset + tga_colormap_type = stbi__get8(s); // colormap type + if( tga_colormap_type > 1 ) { + stbi__rewind(s); + return 0; // only RGB or indexed allowed + } + tga_image_type = stbi__get8(s); // image type + if ( tga_colormap_type == 1 ) { // colormapped (paletted) image + if (tga_image_type != 1 && tga_image_type != 9) { + stbi__rewind(s); + return 0; + } + stbi__skip(s,4); // skip index of first colormap entry and number of entries + sz = stbi__get8(s); // check bits per palette color entry + if ( (sz != 8) && (sz != 15) && (sz != 16) && (sz != 24) && (sz != 32) ) { + stbi__rewind(s); + return 0; + } + stbi__skip(s,4); // skip image x and y origin + tga_colormap_bpp = sz; + } else { // "normal" image w/o colormap - only RGB or grey allowed, +/- RLE + if ( (tga_image_type != 2) && (tga_image_type != 3) && (tga_image_type != 10) && (tga_image_type != 11) ) { + stbi__rewind(s); + return 0; // only RGB or grey allowed, +/- RLE + } + stbi__skip(s,9); // skip colormap specification and image x/y origin + tga_colormap_bpp = 0; + } + tga_w = stbi__get16le(s); + if( tga_w < 1 ) { + stbi__rewind(s); + return 0; // test width + } + tga_h = stbi__get16le(s); + if( tga_h < 1 ) { + stbi__rewind(s); + return 0; // test height + } + tga_bits_per_pixel = stbi__get8(s); // bits per pixel + stbi__get8(s); // ignore alpha bits + if (tga_colormap_bpp != 0) { + if((tga_bits_per_pixel != 8) && (tga_bits_per_pixel != 16)) { + // when using a colormap, tga_bits_per_pixel is the size of the indexes + // I don't think anything but 8 or 16bit indexes makes sense + stbi__rewind(s); + return 0; + } + tga_comp = stbi__tga_get_comp(tga_colormap_bpp, 0, NULL); + } else { + tga_comp = stbi__tga_get_comp(tga_bits_per_pixel, (tga_image_type == 3) || (tga_image_type == 11), NULL); + } + if(!tga_comp) { + stbi__rewind(s); + return 0; + } + if (x) *x = tga_w; + if (y) *y = tga_h; + if (comp) *comp = tga_comp; + return 1; // seems to have passed everything +} + +static int stbi__tga_test(stbi__context *s) +{ + int res = 0; + int sz, tga_color_type; + stbi__get8(s); // discard Offset + tga_color_type = stbi__get8(s); // color type + if ( tga_color_type > 1 ) goto errorEnd; // only RGB or indexed allowed + sz = stbi__get8(s); // image type + if ( tga_color_type == 1 ) { // colormapped (paletted) image + if (sz != 1 && sz != 9) goto errorEnd; // colortype 1 demands image type 1 or 9 + stbi__skip(s,4); // skip index of first colormap entry and number of entries + sz = stbi__get8(s); // check bits per palette color entry + if ( (sz != 8) && (sz != 15) && (sz != 16) && (sz != 24) && (sz != 32) ) goto errorEnd; + stbi__skip(s,4); // skip image x and y origin + } else { // "normal" image w/o colormap + if ( (sz != 2) && (sz != 3) && (sz != 10) && (sz != 11) ) goto errorEnd; // only RGB or grey allowed, +/- RLE + stbi__skip(s,9); // skip colormap specification and image x/y origin + } + if ( stbi__get16le(s) < 1 ) goto errorEnd; // test width + if ( stbi__get16le(s) < 1 ) goto errorEnd; // test height + sz = stbi__get8(s); // bits per pixel + if ( (tga_color_type == 1) && (sz != 8) && (sz != 16) ) goto errorEnd; // for colormapped images, bpp is size of an index + if ( (sz != 8) && (sz != 15) && (sz != 16) && (sz != 24) && (sz != 32) ) goto errorEnd; + + res = 1; // if we got this far, everything's good and we can return 1 instead of 0 + +errorEnd: + stbi__rewind(s); + return res; +} + +// read 16bit value and convert to 24bit RGB +static void stbi__tga_read_rgb16(stbi__context *s, stbi_uc* out) +{ + stbi__uint16 px = (stbi__uint16)stbi__get16le(s); + stbi__uint16 fiveBitMask = 31; + // we have 3 channels with 5bits each + int r = (px >> 10) & fiveBitMask; + int g = (px >> 5) & fiveBitMask; + int b = px & fiveBitMask; + // Note that this saves the data in RGB(A) order, so it doesn't need to be swapped later + out[0] = (stbi_uc)((r * 255)/31); + out[1] = (stbi_uc)((g * 255)/31); + out[2] = (stbi_uc)((b * 255)/31); + + // some people claim that the most significant bit might be used for alpha + // (possibly if an alpha-bit is set in the "image descriptor byte") + // but that only made 16bit test images completely translucent.. + // so let's treat all 15 and 16bit TGAs as RGB with no alpha. +} + +static void *stbi__tga_load(stbi__context *s, int *x, int *y, int *comp, int req_comp, stbi__result_info *ri) +{ + // read in the TGA header stuff + int tga_offset = stbi__get8(s); + int tga_indexed = stbi__get8(s); + int tga_image_type = stbi__get8(s); + int tga_is_RLE = 0; + int tga_palette_start = stbi__get16le(s); + int tga_palette_len = stbi__get16le(s); + int tga_palette_bits = stbi__get8(s); + int tga_x_origin = stbi__get16le(s); + int tga_y_origin = stbi__get16le(s); + int tga_width = stbi__get16le(s); + int tga_height = stbi__get16le(s); + int tga_bits_per_pixel = stbi__get8(s); + int tga_comp, tga_rgb16=0; + int tga_inverted = stbi__get8(s); + // int tga_alpha_bits = tga_inverted & 15; // the 4 lowest bits - unused (useless?) + // image data + unsigned char *tga_data; + unsigned char *tga_palette = NULL; + int i, j; + unsigned char raw_data[4] = {0}; + int RLE_count = 0; + int RLE_repeating = 0; + int read_next_pixel = 1; + STBI_NOTUSED(ri); + STBI_NOTUSED(tga_x_origin); // @TODO + STBI_NOTUSED(tga_y_origin); // @TODO + + if (tga_height > STBI_MAX_DIMENSIONS) return stbi__errpuc("too large","Very large image (corrupt?)"); + if (tga_width > STBI_MAX_DIMENSIONS) return stbi__errpuc("too large","Very large image (corrupt?)"); + + // do a tiny bit of precessing + if ( tga_image_type >= 8 ) + { + tga_image_type -= 8; + tga_is_RLE = 1; + } + tga_inverted = 1 - ((tga_inverted >> 5) & 1); + + // If I'm paletted, then I'll use the number of bits from the palette + if ( tga_indexed ) tga_comp = stbi__tga_get_comp(tga_palette_bits, 0, &tga_rgb16); + else tga_comp = stbi__tga_get_comp(tga_bits_per_pixel, (tga_image_type == 3), &tga_rgb16); + + if(!tga_comp) // shouldn't really happen, stbi__tga_test() should have ensured basic consistency + return stbi__errpuc("bad format", "Can't find out TGA pixelformat"); + + // tga info + *x = tga_width; + *y = tga_height; + if (comp) *comp = tga_comp; + + if (!stbi__mad3sizes_valid(tga_width, tga_height, tga_comp, 0)) + return stbi__errpuc("too large", "Corrupt TGA"); + + tga_data = (unsigned char*)stbi__malloc_mad3(tga_width, tga_height, tga_comp, 0); + if (!tga_data) return stbi__errpuc("outofmem", "Out of memory"); + + // skip to the data's starting position (offset usually = 0) + stbi__skip(s, tga_offset ); + + if ( !tga_indexed && !tga_is_RLE && !tga_rgb16 ) { + for (i=0; i < tga_height; ++i) { + int row = tga_inverted ? tga_height -i - 1 : i; + stbi_uc *tga_row = tga_data + row*tga_width*tga_comp; + stbi__getn(s, tga_row, tga_width * tga_comp); + } + } else { + // do I need to load a palette? + if ( tga_indexed) + { + if (tga_palette_len == 0) { /* you have to have at least one entry! */ + STBI_FREE(tga_data); + return stbi__errpuc("bad palette", "Corrupt TGA"); + } + + // any data to skip? (offset usually = 0) + stbi__skip(s, tga_palette_start ); + // load the palette + tga_palette = (unsigned char*)stbi__malloc_mad2(tga_palette_len, tga_comp, 0); + if (!tga_palette) { + STBI_FREE(tga_data); + return stbi__errpuc("outofmem", "Out of memory"); + } + if (tga_rgb16) { + stbi_uc *pal_entry = tga_palette; + STBI_ASSERT(tga_comp == STBI_rgb); + for (i=0; i < tga_palette_len; ++i) { + stbi__tga_read_rgb16(s, pal_entry); + pal_entry += tga_comp; + } + } else if (!stbi__getn(s, tga_palette, tga_palette_len * tga_comp)) { + STBI_FREE(tga_data); + STBI_FREE(tga_palette); + return stbi__errpuc("bad palette", "Corrupt TGA"); + } + } + // load the data + for (i=0; i < tga_width * tga_height; ++i) + { + // if I'm in RLE mode, do I need to get a RLE stbi__pngchunk? + if ( tga_is_RLE ) + { + if ( RLE_count == 0 ) + { + // yep, get the next byte as a RLE command + int RLE_cmd = stbi__get8(s); + RLE_count = 1 + (RLE_cmd & 127); + RLE_repeating = RLE_cmd >> 7; + read_next_pixel = 1; + } else if ( !RLE_repeating ) + { + read_next_pixel = 1; + } + } else + { + read_next_pixel = 1; + } + // OK, if I need to read a pixel, do it now + if ( read_next_pixel ) + { + // load however much data we did have + if ( tga_indexed ) + { + // read in index, then perform the lookup + int pal_idx = (tga_bits_per_pixel == 8) ? stbi__get8(s) : stbi__get16le(s); + if ( pal_idx >= tga_palette_len ) { + // invalid index + pal_idx = 0; + } + pal_idx *= tga_comp; + for (j = 0; j < tga_comp; ++j) { + raw_data[j] = tga_palette[pal_idx+j]; + } + } else if(tga_rgb16) { + STBI_ASSERT(tga_comp == STBI_rgb); + stbi__tga_read_rgb16(s, raw_data); + } else { + // read in the data raw + for (j = 0; j < tga_comp; ++j) { + raw_data[j] = stbi__get8(s); + } + } + // clear the reading flag for the next pixel + read_next_pixel = 0; + } // end of reading a pixel + + // copy data + for (j = 0; j < tga_comp; ++j) + tga_data[i*tga_comp+j] = raw_data[j]; + + // in case we're in RLE mode, keep counting down + --RLE_count; + } + // do I need to invert the image? + if ( tga_inverted ) + { + for (j = 0; j*2 < tga_height; ++j) + { + int index1 = j * tga_width * tga_comp; + int index2 = (tga_height - 1 - j) * tga_width * tga_comp; + for (i = tga_width * tga_comp; i > 0; --i) + { + unsigned char temp = tga_data[index1]; + tga_data[index1] = tga_data[index2]; + tga_data[index2] = temp; + ++index1; + ++index2; + } + } + } + // clear my palette, if I had one + if ( tga_palette != NULL ) + { + STBI_FREE( tga_palette ); + } + } + + // swap RGB - if the source data was RGB16, it already is in the right order + if (tga_comp >= 3 && !tga_rgb16) + { + unsigned char* tga_pixel = tga_data; + for (i=0; i < tga_width * tga_height; ++i) + { + unsigned char temp = tga_pixel[0]; + tga_pixel[0] = tga_pixel[2]; + tga_pixel[2] = temp; + tga_pixel += tga_comp; + } + } + + // convert to target component count + if (req_comp && req_comp != tga_comp) + tga_data = stbi__convert_format(tga_data, tga_comp, req_comp, tga_width, tga_height); + + // the things I do to get rid of an error message, and yet keep + // Microsoft's C compilers happy... [8^( + tga_palette_start = tga_palette_len = tga_palette_bits = + tga_x_origin = tga_y_origin = 0; + STBI_NOTUSED(tga_palette_start); + // OK, done + return tga_data; +} +#endif + +// ************************************************************************************************* +// Photoshop PSD loader -- PD by Thatcher Ulrich, integration by Nicolas Schulz, tweaked by STB + +#ifndef STBI_NO_PSD +static int stbi__psd_test(stbi__context *s) +{ + int r = (stbi__get32be(s) == 0x38425053); + stbi__rewind(s); + return r; +} + +static int stbi__psd_decode_rle(stbi__context *s, stbi_uc *p, int pixelCount) +{ + int count, nleft, len; + + count = 0; + while ((nleft = pixelCount - count) > 0) { + len = stbi__get8(s); + if (len == 128) { + // No-op. + } else if (len < 128) { + // Copy next len+1 bytes literally. + len++; + if (len > nleft) return 0; // corrupt data + count += len; + while (len) { + *p = stbi__get8(s); + p += 4; + len--; + } + } else if (len > 128) { + stbi_uc val; + // Next -len+1 bytes in the dest are replicated from next source byte. + // (Interpret len as a negative 8-bit int.) + len = 257 - len; + if (len > nleft) return 0; // corrupt data + val = stbi__get8(s); + count += len; + while (len) { + *p = val; + p += 4; + len--; + } + } + } + + return 1; +} + +static void *stbi__psd_load(stbi__context *s, int *x, int *y, int *comp, int req_comp, stbi__result_info *ri, int bpc) +{ + int pixelCount; + int channelCount, compression; + int channel, i; + int bitdepth; + int w,h; + stbi_uc *out; + STBI_NOTUSED(ri); + + // Check identifier + if (stbi__get32be(s) != 0x38425053) // "8BPS" + return stbi__errpuc("not PSD", "Corrupt PSD image"); + + // Check file type version. + if (stbi__get16be(s) != 1) + return stbi__errpuc("wrong version", "Unsupported version of PSD image"); + + // Skip 6 reserved bytes. + stbi__skip(s, 6 ); + + // Read the number of channels (R, G, B, A, etc). + channelCount = stbi__get16be(s); + if (channelCount < 0 || channelCount > 16) + return stbi__errpuc("wrong channel count", "Unsupported number of channels in PSD image"); + + // Read the rows and columns of the image. + h = stbi__get32be(s); + w = stbi__get32be(s); + + if (h > STBI_MAX_DIMENSIONS) return stbi__errpuc("too large","Very large image (corrupt?)"); + if (w > STBI_MAX_DIMENSIONS) return stbi__errpuc("too large","Very large image (corrupt?)"); + + // Make sure the depth is 8 bits. + bitdepth = stbi__get16be(s); + if (bitdepth != 8 && bitdepth != 16) + return stbi__errpuc("unsupported bit depth", "PSD bit depth is not 8 or 16 bit"); + + // Make sure the color mode is RGB. + // Valid options are: + // 0: Bitmap + // 1: Grayscale + // 2: Indexed color + // 3: RGB color + // 4: CMYK color + // 7: Multichannel + // 8: Duotone + // 9: Lab color + if (stbi__get16be(s) != 3) + return stbi__errpuc("wrong color format", "PSD is not in RGB color format"); + + // Skip the Mode Data. (It's the palette for indexed color; other info for other modes.) + stbi__skip(s,stbi__get32be(s) ); + + // Skip the image resources. (resolution, pen tool paths, etc) + stbi__skip(s, stbi__get32be(s) ); + + // Skip the reserved data. + stbi__skip(s, stbi__get32be(s) ); + + // Find out if the data is compressed. + // Known values: + // 0: no compression + // 1: RLE compressed + compression = stbi__get16be(s); + if (compression > 1) + return stbi__errpuc("bad compression", "PSD has an unknown compression format"); + + // Check size + if (!stbi__mad3sizes_valid(4, w, h, 0)) + return stbi__errpuc("too large", "Corrupt PSD"); + + // Create the destination image. + + if (!compression && bitdepth == 16 && bpc == 16) { + out = (stbi_uc *) stbi__malloc_mad3(8, w, h, 0); + ri->bits_per_channel = 16; + } else + out = (stbi_uc *) stbi__malloc(4 * w*h); + + if (!out) return stbi__errpuc("outofmem", "Out of memory"); + pixelCount = w*h; + + // Initialize the data to zero. + //memset( out, 0, pixelCount * 4 ); + + // Finally, the image data. + if (compression) { + // RLE as used by .PSD and .TIFF + // Loop until you get the number of unpacked bytes you are expecting: + // Read the next source byte into n. + // If n is between 0 and 127 inclusive, copy the next n+1 bytes literally. + // Else if n is between -127 and -1 inclusive, copy the next byte -n+1 times. + // Else if n is 128, noop. + // Endloop + + // The RLE-compressed data is preceded by a 2-byte data count for each row in the data, + // which we're going to just skip. + stbi__skip(s, h * channelCount * 2 ); + + // Read the RLE data by channel. + for (channel = 0; channel < 4; channel++) { + stbi_uc *p; + + p = out+channel; + if (channel >= channelCount) { + // Fill this channel with default data. + for (i = 0; i < pixelCount; i++, p += 4) + *p = (channel == 3 ? 255 : 0); + } else { + // Read the RLE data. + if (!stbi__psd_decode_rle(s, p, pixelCount)) { + STBI_FREE(out); + return stbi__errpuc("corrupt", "bad RLE data"); + } + } + } + + } else { + // We're at the raw image data. It's each channel in order (Red, Green, Blue, Alpha, ...) + // where each channel consists of an 8-bit (or 16-bit) value for each pixel in the image. + + // Read the data by channel. + for (channel = 0; channel < 4; channel++) { + if (channel >= channelCount) { + // Fill this channel with default data. + if (bitdepth == 16 && bpc == 16) { + stbi__uint16 *q = ((stbi__uint16 *) out) + channel; + stbi__uint16 val = channel == 3 ? 65535 : 0; + for (i = 0; i < pixelCount; i++, q += 4) + *q = val; + } else { + stbi_uc *p = out+channel; + stbi_uc val = channel == 3 ? 255 : 0; + for (i = 0; i < pixelCount; i++, p += 4) + *p = val; + } + } else { + if (ri->bits_per_channel == 16) { // output bpc + stbi__uint16 *q = ((stbi__uint16 *) out) + channel; + for (i = 0; i < pixelCount; i++, q += 4) + *q = (stbi__uint16) stbi__get16be(s); + } else { + stbi_uc *p = out+channel; + if (bitdepth == 16) { // input bpc + for (i = 0; i < pixelCount; i++, p += 4) + *p = (stbi_uc) (stbi__get16be(s) >> 8); + } else { + for (i = 0; i < pixelCount; i++, p += 4) + *p = stbi__get8(s); + } + } + } + } + } + + // remove weird white matte from PSD + if (channelCount >= 4) { + if (ri->bits_per_channel == 16) { + for (i=0; i < w*h; ++i) { + stbi__uint16 *pixel = (stbi__uint16 *) out + 4*i; + if (pixel[3] != 0 && pixel[3] != 65535) { + float a = pixel[3] / 65535.0f; + float ra = 1.0f / a; + float inv_a = 65535.0f * (1 - ra); + pixel[0] = (stbi__uint16) (pixel[0]*ra + inv_a); + pixel[1] = (stbi__uint16) (pixel[1]*ra + inv_a); + pixel[2] = (stbi__uint16) (pixel[2]*ra + inv_a); + } + } + } else { + for (i=0; i < w*h; ++i) { + unsigned char *pixel = out + 4*i; + if (pixel[3] != 0 && pixel[3] != 255) { + float a = pixel[3] / 255.0f; + float ra = 1.0f / a; + float inv_a = 255.0f * (1 - ra); + pixel[0] = (unsigned char) (pixel[0]*ra + inv_a); + pixel[1] = (unsigned char) (pixel[1]*ra + inv_a); + pixel[2] = (unsigned char) (pixel[2]*ra + inv_a); + } + } + } + } + + // convert to desired output format + if (req_comp && req_comp != 4) { + if (ri->bits_per_channel == 16) + out = (stbi_uc *) stbi__convert_format16((stbi__uint16 *) out, 4, req_comp, w, h); + else + out = stbi__convert_format(out, 4, req_comp, w, h); + if (out == NULL) return out; // stbi__convert_format frees input on failure + } + + if (comp) *comp = 4; + *y = h; + *x = w; + + return out; +} +#endif + +// ************************************************************************************************* +// Softimage PIC loader +// by Tom Seddon +// +// See http://softimage.wiki.softimage.com/index.php/INFO:_PIC_file_format +// See http://ozviz.wasp.uwa.edu.au/~pbourke/dataformats/softimagepic/ + +#ifndef STBI_NO_PIC +static int stbi__pic_is4(stbi__context *s,const char *str) +{ + int i; + for (i=0; i<4; ++i) + if (stbi__get8(s) != (stbi_uc)str[i]) + return 0; + + return 1; +} + +static int stbi__pic_test_core(stbi__context *s) +{ + int i; + + if (!stbi__pic_is4(s,"\x53\x80\xF6\x34")) + return 0; + + for(i=0;i<84;++i) + stbi__get8(s); + + if (!stbi__pic_is4(s,"PICT")) + return 0; + + return 1; +} + +typedef struct +{ + stbi_uc size,type,channel; +} stbi__pic_packet; + +static stbi_uc *stbi__readval(stbi__context *s, int channel, stbi_uc *dest) +{ + int mask=0x80, i; + + for (i=0; i<4; ++i, mask>>=1) { + if (channel & mask) { + if (stbi__at_eof(s)) return stbi__errpuc("bad file","PIC file too short"); + dest[i]=stbi__get8(s); + } + } + + return dest; +} + +static void stbi__copyval(int channel,stbi_uc *dest,const stbi_uc *src) +{ + int mask=0x80,i; + + for (i=0;i<4; ++i, mask>>=1) + if (channel&mask) + dest[i]=src[i]; +} + +static stbi_uc *stbi__pic_load_core(stbi__context *s,int width,int height,int *comp, stbi_uc *result) +{ + int act_comp=0,num_packets=0,y,chained; + stbi__pic_packet packets[10]; + + // this will (should...) cater for even some bizarre stuff like having data + // for the same channel in multiple packets. + do { + stbi__pic_packet *packet; + + if (num_packets==sizeof(packets)/sizeof(packets[0])) + return stbi__errpuc("bad format","too many packets"); + + packet = &packets[num_packets++]; + + chained = stbi__get8(s); + packet->size = stbi__get8(s); + packet->type = stbi__get8(s); + packet->channel = stbi__get8(s); + + act_comp |= packet->channel; + + if (stbi__at_eof(s)) return stbi__errpuc("bad file","file too short (reading packets)"); + if (packet->size != 8) return stbi__errpuc("bad format","packet isn't 8bpp"); + } while (chained); + + *comp = (act_comp & 0x10 ? 4 : 3); // has alpha channel? + + for(y=0; ytype) { + default: + return stbi__errpuc("bad format","packet has bad compression type"); + + case 0: {//uncompressed + int x; + + for(x=0;xchannel,dest)) + return 0; + break; + } + + case 1://Pure RLE + { + int left=width, i; + + while (left>0) { + stbi_uc count,value[4]; + + count=stbi__get8(s); + if (stbi__at_eof(s)) return stbi__errpuc("bad file","file too short (pure read count)"); + + if (count > left) + count = (stbi_uc) left; + + if (!stbi__readval(s,packet->channel,value)) return 0; + + for(i=0; ichannel,dest,value); + left -= count; + } + } + break; + + case 2: {//Mixed RLE + int left=width; + while (left>0) { + int count = stbi__get8(s), i; + if (stbi__at_eof(s)) return stbi__errpuc("bad file","file too short (mixed read count)"); + + if (count >= 128) { // Repeated + stbi_uc value[4]; + + if (count==128) + count = stbi__get16be(s); + else + count -= 127; + if (count > left) + return stbi__errpuc("bad file","scanline overrun"); + + if (!stbi__readval(s,packet->channel,value)) + return 0; + + for(i=0;ichannel,dest,value); + } else { // Raw + ++count; + if (count>left) return stbi__errpuc("bad file","scanline overrun"); + + for(i=0;ichannel,dest)) + return 0; + } + left-=count; + } + break; + } + } + } + } + + return result; +} + +static void *stbi__pic_load(stbi__context *s,int *px,int *py,int *comp,int req_comp, stbi__result_info *ri) +{ + stbi_uc *result; + int i, x,y, internal_comp; + STBI_NOTUSED(ri); + + if (!comp) comp = &internal_comp; + + for (i=0; i<92; ++i) + stbi__get8(s); + + x = stbi__get16be(s); + y = stbi__get16be(s); + + if (y > STBI_MAX_DIMENSIONS) return stbi__errpuc("too large","Very large image (corrupt?)"); + if (x > STBI_MAX_DIMENSIONS) return stbi__errpuc("too large","Very large image (corrupt?)"); + + if (stbi__at_eof(s)) return stbi__errpuc("bad file","file too short (pic header)"); + if (!stbi__mad3sizes_valid(x, y, 4, 0)) return stbi__errpuc("too large", "PIC image too large to decode"); + + stbi__get32be(s); //skip `ratio' + stbi__get16be(s); //skip `fields' + stbi__get16be(s); //skip `pad' + + // intermediate buffer is RGBA + result = (stbi_uc *) stbi__malloc_mad3(x, y, 4, 0); + if (!result) return stbi__errpuc("outofmem", "Out of memory"); + memset(result, 0xff, x*y*4); + + if (!stbi__pic_load_core(s,x,y,comp, result)) { + STBI_FREE(result); + result=0; + } + *px = x; + *py = y; + if (req_comp == 0) req_comp = *comp; + result=stbi__convert_format(result,4,req_comp,x,y); + + return result; +} + +static int stbi__pic_test(stbi__context *s) +{ + int r = stbi__pic_test_core(s); + stbi__rewind(s); + return r; +} +#endif + +// ************************************************************************************************* +// GIF loader -- public domain by Jean-Marc Lienher -- simplified/shrunk by stb + +#ifndef STBI_NO_GIF +typedef struct +{ + stbi__int16 prefix; + stbi_uc first; + stbi_uc suffix; +} stbi__gif_lzw; + +typedef struct +{ + int w,h; + stbi_uc *out; // output buffer (always 4 components) + stbi_uc *background; // The current "background" as far as a gif is concerned + stbi_uc *history; + int flags, bgindex, ratio, transparent, eflags; + stbi_uc pal[256][4]; + stbi_uc lpal[256][4]; + stbi__gif_lzw codes[8192]; + stbi_uc *color_table; + int parse, step; + int lflags; + int start_x, start_y; + int max_x, max_y; + int cur_x, cur_y; + int line_size; + int delay; +} stbi__gif; + +static int stbi__gif_test_raw(stbi__context *s) +{ + int sz; + if (stbi__get8(s) != 'G' || stbi__get8(s) != 'I' || stbi__get8(s) != 'F' || stbi__get8(s) != '8') return 0; + sz = stbi__get8(s); + if (sz != '9' && sz != '7') return 0; + if (stbi__get8(s) != 'a') return 0; + return 1; +} + +static int stbi__gif_test(stbi__context *s) +{ + int r = stbi__gif_test_raw(s); + stbi__rewind(s); + return r; +} + +static void stbi__gif_parse_colortable(stbi__context *s, stbi_uc pal[256][4], int num_entries, int transp) +{ + int i; + for (i=0; i < num_entries; ++i) { + pal[i][2] = stbi__get8(s); + pal[i][1] = stbi__get8(s); + pal[i][0] = stbi__get8(s); + pal[i][3] = transp == i ? 0 : 255; + } +} + +static int stbi__gif_header(stbi__context *s, stbi__gif *g, int *comp, int is_info) +{ + stbi_uc version; + if (stbi__get8(s) != 'G' || stbi__get8(s) != 'I' || stbi__get8(s) != 'F' || stbi__get8(s) != '8') + return stbi__err("not GIF", "Corrupt GIF"); + + version = stbi__get8(s); + if (version != '7' && version != '9') return stbi__err("not GIF", "Corrupt GIF"); + if (stbi__get8(s) != 'a') return stbi__err("not GIF", "Corrupt GIF"); + + stbi__g_failure_reason = ""; + g->w = stbi__get16le(s); + g->h = stbi__get16le(s); + g->flags = stbi__get8(s); + g->bgindex = stbi__get8(s); + g->ratio = stbi__get8(s); + g->transparent = -1; + + if (g->w > STBI_MAX_DIMENSIONS) return stbi__err("too large","Very large image (corrupt?)"); + if (g->h > STBI_MAX_DIMENSIONS) return stbi__err("too large","Very large image (corrupt?)"); + + if (comp != 0) *comp = 4; // can't actually tell whether it's 3 or 4 until we parse the comments + + if (is_info) return 1; + + if (g->flags & 0x80) + stbi__gif_parse_colortable(s,g->pal, 2 << (g->flags & 7), -1); + + return 1; +} + +static int stbi__gif_info_raw(stbi__context *s, int *x, int *y, int *comp) +{ + stbi__gif* g = (stbi__gif*) stbi__malloc(sizeof(stbi__gif)); + if (!g) return stbi__err("outofmem", "Out of memory"); + if (!stbi__gif_header(s, g, comp, 1)) { + STBI_FREE(g); + stbi__rewind( s ); + return 0; + } + if (x) *x = g->w; + if (y) *y = g->h; + STBI_FREE(g); + return 1; +} + +static void stbi__out_gif_code(stbi__gif *g, stbi__uint16 code) +{ + stbi_uc *p, *c; + int idx; + + // recurse to decode the prefixes, since the linked-list is backwards, + // and working backwards through an interleaved image would be nasty + if (g->codes[code].prefix >= 0) + stbi__out_gif_code(g, g->codes[code].prefix); + + if (g->cur_y >= g->max_y) return; + + idx = g->cur_x + g->cur_y; + p = &g->out[idx]; + g->history[idx / 4] = 1; + + c = &g->color_table[g->codes[code].suffix * 4]; + if (c[3] > 128) { // don't render transparent pixels; + p[0] = c[2]; + p[1] = c[1]; + p[2] = c[0]; + p[3] = c[3]; + } + g->cur_x += 4; + + if (g->cur_x >= g->max_x) { + g->cur_x = g->start_x; + g->cur_y += g->step; + + while (g->cur_y >= g->max_y && g->parse > 0) { + g->step = (1 << g->parse) * g->line_size; + g->cur_y = g->start_y + (g->step >> 1); + --g->parse; + } + } +} + +static stbi_uc *stbi__process_gif_raster(stbi__context *s, stbi__gif *g) +{ + stbi_uc lzw_cs; + stbi__int32 len, init_code; + stbi__uint32 first; + stbi__int32 codesize, codemask, avail, oldcode, bits, valid_bits, clear; + stbi__gif_lzw *p; + + lzw_cs = stbi__get8(s); + if (lzw_cs > 12) return NULL; + clear = 1 << lzw_cs; + first = 1; + codesize = lzw_cs + 1; + codemask = (1 << codesize) - 1; + bits = 0; + valid_bits = 0; + for (init_code = 0; init_code < clear; init_code++) { + g->codes[init_code].prefix = -1; + g->codes[init_code].first = (stbi_uc) init_code; + g->codes[init_code].suffix = (stbi_uc) init_code; + } + + // support no starting clear code + avail = clear+2; + oldcode = -1; + + len = 0; + for(;;) { + if (valid_bits < codesize) { + if (len == 0) { + len = stbi__get8(s); // start new block + if (len == 0) + return g->out; + } + --len; + bits |= (stbi__int32) stbi__get8(s) << valid_bits; + valid_bits += 8; + } else { + stbi__int32 code = bits & codemask; + bits >>= codesize; + valid_bits -= codesize; + // @OPTIMIZE: is there some way we can accelerate the non-clear path? + if (code == clear) { // clear code + codesize = lzw_cs + 1; + codemask = (1 << codesize) - 1; + avail = clear + 2; + oldcode = -1; + first = 0; + } else if (code == clear + 1) { // end of stream code + stbi__skip(s, len); + while ((len = stbi__get8(s)) > 0) + stbi__skip(s,len); + return g->out; + } else if (code <= avail) { + if (first) { + return stbi__errpuc("no clear code", "Corrupt GIF"); + } + + if (oldcode >= 0) { + p = &g->codes[avail++]; + if (avail > 8192) { + return stbi__errpuc("too many codes", "Corrupt GIF"); + } + + p->prefix = (stbi__int16) oldcode; + p->first = g->codes[oldcode].first; + p->suffix = (code == avail) ? p->first : g->codes[code].first; + } else if (code == avail) + return stbi__errpuc("illegal code in raster", "Corrupt GIF"); + + stbi__out_gif_code(g, (stbi__uint16) code); + + if ((avail & codemask) == 0 && avail <= 0x0FFF) { + codesize++; + codemask = (1 << codesize) - 1; + } + + oldcode = code; + } else { + return stbi__errpuc("illegal code in raster", "Corrupt GIF"); + } + } + } +} + +// this function is designed to support animated gifs, although stb_image doesn't support it +// two back is the image from two frames ago, used for a very specific disposal format +static stbi_uc *stbi__gif_load_next(stbi__context *s, stbi__gif *g, int *comp, int req_comp, stbi_uc *two_back) +{ + int dispose; + int first_frame; + int pi; + int pcount; + STBI_NOTUSED(req_comp); + + // on first frame, any non-written pixels get the background colour (non-transparent) + first_frame = 0; + if (g->out == 0) { + if (!stbi__gif_header(s, g, comp,0)) return 0; // stbi__g_failure_reason set by stbi__gif_header + if (!stbi__mad3sizes_valid(4, g->w, g->h, 0)) + return stbi__errpuc("too large", "GIF image is too large"); + pcount = g->w * g->h; + g->out = (stbi_uc *) stbi__malloc(4 * pcount); + g->background = (stbi_uc *) stbi__malloc(4 * pcount); + g->history = (stbi_uc *) stbi__malloc(pcount); + if (!g->out || !g->background || !g->history) + return stbi__errpuc("outofmem", "Out of memory"); + + // image is treated as "transparent" at the start - ie, nothing overwrites the current background; + // background colour is only used for pixels that are not rendered first frame, after that "background" + // color refers to the color that was there the previous frame. + memset(g->out, 0x00, 4 * pcount); + memset(g->background, 0x00, 4 * pcount); // state of the background (starts transparent) + memset(g->history, 0x00, pcount); // pixels that were affected previous frame + first_frame = 1; + } else { + // second frame - how do we dispose of the previous one? + dispose = (g->eflags & 0x1C) >> 2; + pcount = g->w * g->h; + + if ((dispose == 3) && (two_back == 0)) { + dispose = 2; // if I don't have an image to revert back to, default to the old background + } + + if (dispose == 3) { // use previous graphic + for (pi = 0; pi < pcount; ++pi) { + if (g->history[pi]) { + memcpy( &g->out[pi * 4], &two_back[pi * 4], 4 ); + } + } + } else if (dispose == 2) { + // restore what was changed last frame to background before that frame; + for (pi = 0; pi < pcount; ++pi) { + if (g->history[pi]) { + memcpy( &g->out[pi * 4], &g->background[pi * 4], 4 ); + } + } + } else { + // This is a non-disposal case eithe way, so just + // leave the pixels as is, and they will become the new background + // 1: do not dispose + // 0: not specified. + } + + // background is what out is after the undoing of the previou frame; + memcpy( g->background, g->out, 4 * g->w * g->h ); + } + + // clear my history; + memset( g->history, 0x00, g->w * g->h ); // pixels that were affected previous frame + + for (;;) { + int tag = stbi__get8(s); + switch (tag) { + case 0x2C: /* Image Descriptor */ + { + stbi__int32 x, y, w, h; + stbi_uc *o; + + x = stbi__get16le(s); + y = stbi__get16le(s); + w = stbi__get16le(s); + h = stbi__get16le(s); + if (((x + w) > (g->w)) || ((y + h) > (g->h))) + return stbi__errpuc("bad Image Descriptor", "Corrupt GIF"); + + g->line_size = g->w * 4; + g->start_x = x * 4; + g->start_y = y * g->line_size; + g->max_x = g->start_x + w * 4; + g->max_y = g->start_y + h * g->line_size; + g->cur_x = g->start_x; + g->cur_y = g->start_y; + + // if the width of the specified rectangle is 0, that means + // we may not see *any* pixels or the image is malformed; + // to make sure this is caught, move the current y down to + // max_y (which is what out_gif_code checks). + if (w == 0) + g->cur_y = g->max_y; + + g->lflags = stbi__get8(s); + + if (g->lflags & 0x40) { + g->step = 8 * g->line_size; // first interlaced spacing + g->parse = 3; + } else { + g->step = g->line_size; + g->parse = 0; + } + + if (g->lflags & 0x80) { + stbi__gif_parse_colortable(s,g->lpal, 2 << (g->lflags & 7), g->eflags & 0x01 ? g->transparent : -1); + g->color_table = (stbi_uc *) g->lpal; + } else if (g->flags & 0x80) { + g->color_table = (stbi_uc *) g->pal; + } else + return stbi__errpuc("missing color table", "Corrupt GIF"); + + o = stbi__process_gif_raster(s, g); + if (!o) return NULL; + + // if this was the first frame, + pcount = g->w * g->h; + if (first_frame && (g->bgindex > 0)) { + // if first frame, any pixel not drawn to gets the background color + for (pi = 0; pi < pcount; ++pi) { + if (g->history[pi] == 0) { + g->pal[g->bgindex][3] = 255; // just in case it was made transparent, undo that; It will be reset next frame if need be; + memcpy( &g->out[pi * 4], &g->pal[g->bgindex], 4 ); + } + } + } + + return o; + } + + case 0x21: // Comment Extension. + { + int len; + int ext = stbi__get8(s); + if (ext == 0xF9) { // Graphic Control Extension. + len = stbi__get8(s); + if (len == 4) { + g->eflags = stbi__get8(s); + g->delay = 10 * stbi__get16le(s); // delay - 1/100th of a second, saving as 1/1000ths. + + // unset old transparent + if (g->transparent >= 0) { + g->pal[g->transparent][3] = 255; + } + if (g->eflags & 0x01) { + g->transparent = stbi__get8(s); + if (g->transparent >= 0) { + g->pal[g->transparent][3] = 0; + } + } else { + // don't need transparent + stbi__skip(s, 1); + g->transparent = -1; + } + } else { + stbi__skip(s, len); + break; + } + } + while ((len = stbi__get8(s)) != 0) { + stbi__skip(s, len); + } + break; + } + + case 0x3B: // gif stream termination code + return (stbi_uc *) s; // using '1' causes warning on some compilers + + default: + return stbi__errpuc("unknown code", "Corrupt GIF"); + } + } +} + +static void *stbi__load_gif_main_outofmem(stbi__gif *g, stbi_uc *out, int **delays) +{ + STBI_FREE(g->out); + STBI_FREE(g->history); + STBI_FREE(g->background); + + if (out) STBI_FREE(out); + if (delays && *delays) STBI_FREE(*delays); + return stbi__errpuc("outofmem", "Out of memory"); +} + +static void *stbi__load_gif_main(stbi__context *s, int **delays, int *x, int *y, int *z, int *comp, int req_comp) +{ + if (stbi__gif_test(s)) { + int layers = 0; + stbi_uc *u = 0; + stbi_uc *out = 0; + stbi_uc *two_back = 0; + stbi__gif g; + int stride; + int out_size = 0; + int delays_size = 0; + + STBI_NOTUSED(out_size); + STBI_NOTUSED(delays_size); + + memset(&g, 0, sizeof(g)); + if (delays) { + *delays = 0; + } + + do { + u = stbi__gif_load_next(s, &g, comp, req_comp, two_back); + if (u == (stbi_uc *) s) u = 0; // end of animated gif marker + + if (u) { + *x = g.w; + *y = g.h; + ++layers; + stride = g.w * g.h * 4; + + if (out) { + void *tmp = (stbi_uc*) STBI_REALLOC_SIZED( out, out_size, layers * stride ); + if (!tmp) + return stbi__load_gif_main_outofmem(&g, out, delays); + else { + out = (stbi_uc*) tmp; + out_size = layers * stride; + } + + if (delays) { + int *new_delays = (int*) STBI_REALLOC_SIZED( *delays, delays_size, sizeof(int) * layers ); + if (!new_delays) + return stbi__load_gif_main_outofmem(&g, out, delays); + *delays = new_delays; + delays_size = layers * sizeof(int); + } + } else { + out = (stbi_uc*)stbi__malloc( layers * stride ); + if (!out) + return stbi__load_gif_main_outofmem(&g, out, delays); + out_size = layers * stride; + if (delays) { + *delays = (int*) stbi__malloc( layers * sizeof(int) ); + if (!*delays) + return stbi__load_gif_main_outofmem(&g, out, delays); + delays_size = layers * sizeof(int); + } + } + memcpy( out + ((layers - 1) * stride), u, stride ); + if (layers >= 2) { + two_back = out - 2 * stride; + } + + if (delays) { + (*delays)[layers - 1U] = g.delay; + } + } + } while (u != 0); + + // free temp buffer; + STBI_FREE(g.out); + STBI_FREE(g.history); + STBI_FREE(g.background); + + // do the final conversion after loading everything; + if (req_comp && req_comp != 4) + out = stbi__convert_format(out, 4, req_comp, layers * g.w, g.h); + + *z = layers; + return out; + } else { + return stbi__errpuc("not GIF", "Image was not as a gif type."); + } +} + +static void *stbi__gif_load(stbi__context *s, int *x, int *y, int *comp, int req_comp, stbi__result_info *ri) +{ + stbi_uc *u = 0; + stbi__gif g; + memset(&g, 0, sizeof(g)); + STBI_NOTUSED(ri); + + u = stbi__gif_load_next(s, &g, comp, req_comp, 0); + if (u == (stbi_uc *) s) u = 0; // end of animated gif marker + if (u) { + *x = g.w; + *y = g.h; + + // moved conversion to after successful load so that the same + // can be done for multiple frames. + if (req_comp && req_comp != 4) + u = stbi__convert_format(u, 4, req_comp, g.w, g.h); + } else if (g.out) { + // if there was an error and we allocated an image buffer, free it! + STBI_FREE(g.out); + } + + // free buffers needed for multiple frame loading; + STBI_FREE(g.history); + STBI_FREE(g.background); + + return u; +} + +static int stbi__gif_info(stbi__context *s, int *x, int *y, int *comp) +{ + return stbi__gif_info_raw(s,x,y,comp); +} +#endif + +// ************************************************************************************************* +// Radiance RGBE HDR loader +// originally by Nicolas Schulz +#ifndef STBI_NO_HDR +static int stbi__hdr_test_core(stbi__context *s, const char *signature) +{ + int i; + for (i=0; signature[i]; ++i) + if (stbi__get8(s) != signature[i]) + return 0; + stbi__rewind(s); + return 1; +} + +static int stbi__hdr_test(stbi__context* s) +{ + int r = stbi__hdr_test_core(s, "#?RADIANCE\n"); + stbi__rewind(s); + if(!r) { + r = stbi__hdr_test_core(s, "#?RGBE\n"); + stbi__rewind(s); + } + return r; +} + +#define STBI__HDR_BUFLEN 1024 +static char *stbi__hdr_gettoken(stbi__context *z, char *buffer) +{ + int len=0; + char c = '\0'; + + c = (char) stbi__get8(z); + + while (!stbi__at_eof(z) && c != '\n') { + buffer[len++] = c; + if (len == STBI__HDR_BUFLEN-1) { + // flush to end of line + while (!stbi__at_eof(z) && stbi__get8(z) != '\n') + ; + break; + } + c = (char) stbi__get8(z); + } + + buffer[len] = 0; + return buffer; +} + +static void stbi__hdr_convert(float *output, stbi_uc *input, int req_comp) +{ + if ( input[3] != 0 ) { + float f1; + // Exponent + f1 = (float) ldexp(1.0f, input[3] - (int)(128 + 8)); + if (req_comp <= 2) + output[0] = (input[0] + input[1] + input[2]) * f1 / 3; + else { + output[0] = input[0] * f1; + output[1] = input[1] * f1; + output[2] = input[2] * f1; + } + if (req_comp == 2) output[1] = 1; + if (req_comp == 4) output[3] = 1; + } else { + switch (req_comp) { + case 4: output[3] = 1; /* fallthrough */ + case 3: output[0] = output[1] = output[2] = 0; + break; + case 2: output[1] = 1; /* fallthrough */ + case 1: output[0] = 0; + break; + } + } +} + +static float *stbi__hdr_load(stbi__context *s, int *x, int *y, int *comp, int req_comp, stbi__result_info *ri) +{ + char buffer[STBI__HDR_BUFLEN]; + char *token; + int valid = 0; + int width, height; + stbi_uc *scanline; + float *hdr_data; + int len; + unsigned char count, value; + int i, j, k, c1,c2, z; + const char *headerToken; + STBI_NOTUSED(ri); + + // Check identifier + headerToken = stbi__hdr_gettoken(s,buffer); + if (strcmp(headerToken, "#?RADIANCE") != 0 && strcmp(headerToken, "#?RGBE") != 0) + return stbi__errpf("not HDR", "Corrupt HDR image"); + + // Parse header + for(;;) { + token = stbi__hdr_gettoken(s,buffer); + if (token[0] == 0) break; + if (strcmp(token, "FORMAT=32-bit_rle_rgbe") == 0) valid = 1; + } + + if (!valid) return stbi__errpf("unsupported format", "Unsupported HDR format"); + + // Parse width and height + // can't use sscanf() if we're not using stdio! + token = stbi__hdr_gettoken(s,buffer); + if (strncmp(token, "-Y ", 3)) return stbi__errpf("unsupported data layout", "Unsupported HDR format"); + token += 3; + height = (int) strtol(token, &token, 10); + while (*token == ' ') ++token; + if (strncmp(token, "+X ", 3)) return stbi__errpf("unsupported data layout", "Unsupported HDR format"); + token += 3; + width = (int) strtol(token, NULL, 10); + + if (height > STBI_MAX_DIMENSIONS) return stbi__errpf("too large","Very large image (corrupt?)"); + if (width > STBI_MAX_DIMENSIONS) return stbi__errpf("too large","Very large image (corrupt?)"); + + *x = width; + *y = height; + + if (comp) *comp = 3; + if (req_comp == 0) req_comp = 3; + + if (!stbi__mad4sizes_valid(width, height, req_comp, sizeof(float), 0)) + return stbi__errpf("too large", "HDR image is too large"); + + // Read data + hdr_data = (float *) stbi__malloc_mad4(width, height, req_comp, sizeof(float), 0); + if (!hdr_data) + return stbi__errpf("outofmem", "Out of memory"); + + // Load image data + // image data is stored as some number of sca + if ( width < 8 || width >= 32768) { + // Read flat data + for (j=0; j < height; ++j) { + for (i=0; i < width; ++i) { + stbi_uc rgbe[4]; + main_decode_loop: + stbi__getn(s, rgbe, 4); + stbi__hdr_convert(hdr_data + j * width * req_comp + i * req_comp, rgbe, req_comp); + } + } + } else { + // Read RLE-encoded data + scanline = NULL; + + for (j = 0; j < height; ++j) { + c1 = stbi__get8(s); + c2 = stbi__get8(s); + len = stbi__get8(s); + if (c1 != 2 || c2 != 2 || (len & 0x80)) { + // not run-length encoded, so we have to actually use THIS data as a decoded + // pixel (note this can't be a valid pixel--one of RGB must be >= 128) + stbi_uc rgbe[4]; + rgbe[0] = (stbi_uc) c1; + rgbe[1] = (stbi_uc) c2; + rgbe[2] = (stbi_uc) len; + rgbe[3] = (stbi_uc) stbi__get8(s); + stbi__hdr_convert(hdr_data, rgbe, req_comp); + i = 1; + j = 0; + STBI_FREE(scanline); + goto main_decode_loop; // yes, this makes no sense + } + len <<= 8; + len |= stbi__get8(s); + if (len != width) { STBI_FREE(hdr_data); STBI_FREE(scanline); return stbi__errpf("invalid decoded scanline length", "corrupt HDR"); } + if (scanline == NULL) { + scanline = (stbi_uc *) stbi__malloc_mad2(width, 4, 0); + if (!scanline) { + STBI_FREE(hdr_data); + return stbi__errpf("outofmem", "Out of memory"); + } + } + + for (k = 0; k < 4; ++k) { + int nleft; + i = 0; + while ((nleft = width - i) > 0) { + count = stbi__get8(s); + if (count > 128) { + // Run + value = stbi__get8(s); + count -= 128; + if ((count == 0) || (count > nleft)) { STBI_FREE(hdr_data); STBI_FREE(scanline); return stbi__errpf("corrupt", "bad RLE data in HDR"); } + for (z = 0; z < count; ++z) + scanline[i++ * 4 + k] = value; + } else { + // Dump + if ((count == 0) || (count > nleft)) { STBI_FREE(hdr_data); STBI_FREE(scanline); return stbi__errpf("corrupt", "bad RLE data in HDR"); } + for (z = 0; z < count; ++z) + scanline[i++ * 4 + k] = stbi__get8(s); + } + } + } + for (i=0; i < width; ++i) + stbi__hdr_convert(hdr_data+(j*width + i)*req_comp, scanline + i*4, req_comp); + } + if (scanline) + STBI_FREE(scanline); + } + + return hdr_data; +} + +static int stbi__hdr_info(stbi__context *s, int *x, int *y, int *comp) +{ + char buffer[STBI__HDR_BUFLEN]; + char *token; + int valid = 0; + int dummy; + + if (!x) x = &dummy; + if (!y) y = &dummy; + if (!comp) comp = &dummy; + + if (stbi__hdr_test(s) == 0) { + stbi__rewind( s ); + return 0; + } + + for(;;) { + token = stbi__hdr_gettoken(s,buffer); + if (token[0] == 0) break; + if (strcmp(token, "FORMAT=32-bit_rle_rgbe") == 0) valid = 1; + } + + if (!valid) { + stbi__rewind( s ); + return 0; + } + token = stbi__hdr_gettoken(s,buffer); + if (strncmp(token, "-Y ", 3)) { + stbi__rewind( s ); + return 0; + } + token += 3; + *y = (int) strtol(token, &token, 10); + while (*token == ' ') ++token; + if (strncmp(token, "+X ", 3)) { + stbi__rewind( s ); + return 0; + } + token += 3; + *x = (int) strtol(token, NULL, 10); + *comp = 3; + return 1; +} +#endif // STBI_NO_HDR + +#ifndef STBI_NO_BMP +static int stbi__bmp_info(stbi__context *s, int *x, int *y, int *comp) +{ + void *p; + stbi__bmp_data info; + + info.all_a = 255; + p = stbi__bmp_parse_header(s, &info); + if (p == NULL) { + stbi__rewind( s ); + return 0; + } + if (x) *x = s->img_x; + if (y) *y = s->img_y; + if (comp) { + if (info.bpp == 24 && info.ma == 0xff000000) + *comp = 3; + else + *comp = info.ma ? 4 : 3; + } + return 1; +} +#endif + +#ifndef STBI_NO_PSD +static int stbi__psd_info(stbi__context *s, int *x, int *y, int *comp) +{ + int channelCount, dummy, depth; + if (!x) x = &dummy; + if (!y) y = &dummy; + if (!comp) comp = &dummy; + if (stbi__get32be(s) != 0x38425053) { + stbi__rewind( s ); + return 0; + } + if (stbi__get16be(s) != 1) { + stbi__rewind( s ); + return 0; + } + stbi__skip(s, 6); + channelCount = stbi__get16be(s); + if (channelCount < 0 || channelCount > 16) { + stbi__rewind( s ); + return 0; + } + *y = stbi__get32be(s); + *x = stbi__get32be(s); + depth = stbi__get16be(s); + if (depth != 8 && depth != 16) { + stbi__rewind( s ); + return 0; + } + if (stbi__get16be(s) != 3) { + stbi__rewind( s ); + return 0; + } + *comp = 4; + return 1; +} + +static int stbi__psd_is16(stbi__context *s) +{ + int channelCount, depth; + if (stbi__get32be(s) != 0x38425053) { + stbi__rewind( s ); + return 0; + } + if (stbi__get16be(s) != 1) { + stbi__rewind( s ); + return 0; + } + stbi__skip(s, 6); + channelCount = stbi__get16be(s); + if (channelCount < 0 || channelCount > 16) { + stbi__rewind( s ); + return 0; + } + STBI_NOTUSED(stbi__get32be(s)); + STBI_NOTUSED(stbi__get32be(s)); + depth = stbi__get16be(s); + if (depth != 16) { + stbi__rewind( s ); + return 0; + } + return 1; +} +#endif + +#ifndef STBI_NO_PIC +static int stbi__pic_info(stbi__context *s, int *x, int *y, int *comp) +{ + int act_comp=0,num_packets=0,chained,dummy; + stbi__pic_packet packets[10]; + + if (!x) x = &dummy; + if (!y) y = &dummy; + if (!comp) comp = &dummy; + + if (!stbi__pic_is4(s,"\x53\x80\xF6\x34")) { + stbi__rewind(s); + return 0; + } + + stbi__skip(s, 88); + + *x = stbi__get16be(s); + *y = stbi__get16be(s); + if (stbi__at_eof(s)) { + stbi__rewind( s); + return 0; + } + if ( (*x) != 0 && (1 << 28) / (*x) < (*y)) { + stbi__rewind( s ); + return 0; + } + + stbi__skip(s, 8); + + do { + stbi__pic_packet *packet; + + if (num_packets==sizeof(packets)/sizeof(packets[0])) + return 0; + + packet = &packets[num_packets++]; + chained = stbi__get8(s); + packet->size = stbi__get8(s); + packet->type = stbi__get8(s); + packet->channel = stbi__get8(s); + act_comp |= packet->channel; + + if (stbi__at_eof(s)) { + stbi__rewind( s ); + return 0; + } + if (packet->size != 8) { + stbi__rewind( s ); + return 0; + } + } while (chained); + + *comp = (act_comp & 0x10 ? 4 : 3); + + return 1; +} +#endif + +// ************************************************************************************************* +// Portable Gray Map and Portable Pixel Map loader +// by Ken Miller +// +// PGM: http://netpbm.sourceforge.net/doc/pgm.html +// PPM: http://netpbm.sourceforge.net/doc/ppm.html +// +// Known limitations: +// Does not support comments in the header section +// Does not support ASCII image data (formats P2 and P3) + +#ifndef STBI_NO_PNM + +static int stbi__pnm_test(stbi__context *s) +{ + char p, t; + p = (char) stbi__get8(s); + t = (char) stbi__get8(s); + if (p != 'P' || (t != '5' && t != '6')) { + stbi__rewind( s ); + return 0; + } + return 1; +} + +static void *stbi__pnm_load(stbi__context *s, int *x, int *y, int *comp, int req_comp, stbi__result_info *ri) +{ + stbi_uc *out; + STBI_NOTUSED(ri); + + ri->bits_per_channel = stbi__pnm_info(s, (int *)&s->img_x, (int *)&s->img_y, (int *)&s->img_n); + if (ri->bits_per_channel == 0) + return 0; + + if (s->img_y > STBI_MAX_DIMENSIONS) return stbi__errpuc("too large","Very large image (corrupt?)"); + if (s->img_x > STBI_MAX_DIMENSIONS) return stbi__errpuc("too large","Very large image (corrupt?)"); + + *x = s->img_x; + *y = s->img_y; + if (comp) *comp = s->img_n; + + if (!stbi__mad4sizes_valid(s->img_n, s->img_x, s->img_y, ri->bits_per_channel / 8, 0)) + return stbi__errpuc("too large", "PNM too large"); + + out = (stbi_uc *) stbi__malloc_mad4(s->img_n, s->img_x, s->img_y, ri->bits_per_channel / 8, 0); + if (!out) return stbi__errpuc("outofmem", "Out of memory"); + if (!stbi__getn(s, out, s->img_n * s->img_x * s->img_y * (ri->bits_per_channel / 8))) { + STBI_FREE(out); + return stbi__errpuc("bad PNM", "PNM file truncated"); + } + + if (req_comp && req_comp != s->img_n) { + if (ri->bits_per_channel == 16) { + out = (stbi_uc *) stbi__convert_format16((stbi__uint16 *) out, s->img_n, req_comp, s->img_x, s->img_y); + } else { + out = stbi__convert_format(out, s->img_n, req_comp, s->img_x, s->img_y); + } + if (out == NULL) return out; // stbi__convert_format frees input on failure + } + return out; +} + +static int stbi__pnm_isspace(char c) +{ + return c == ' ' || c == '\t' || c == '\n' || c == '\v' || c == '\f' || c == '\r'; +} + +static void stbi__pnm_skip_whitespace(stbi__context *s, char *c) +{ + for (;;) { + while (!stbi__at_eof(s) && stbi__pnm_isspace(*c)) + *c = (char) stbi__get8(s); + + if (stbi__at_eof(s) || *c != '#') + break; + + while (!stbi__at_eof(s) && *c != '\n' && *c != '\r' ) + *c = (char) stbi__get8(s); + } +} + +static int stbi__pnm_isdigit(char c) +{ + return c >= '0' && c <= '9'; +} + +static int stbi__pnm_getinteger(stbi__context *s, char *c) +{ + int value = 0; + + while (!stbi__at_eof(s) && stbi__pnm_isdigit(*c)) { + value = value*10 + (*c - '0'); + *c = (char) stbi__get8(s); + if((value > 214748364) || (value == 214748364 && *c > '7')) + return stbi__err("integer parse overflow", "Parsing an integer in the PPM header overflowed a 32-bit int"); + } + + return value; +} + +static int stbi__pnm_info(stbi__context *s, int *x, int *y, int *comp) +{ + int maxv, dummy; + char c, p, t; + + if (!x) x = &dummy; + if (!y) y = &dummy; + if (!comp) comp = &dummy; + + stbi__rewind(s); + + // Get identifier + p = (char) stbi__get8(s); + t = (char) stbi__get8(s); + if (p != 'P' || (t != '5' && t != '6')) { + stbi__rewind(s); + return 0; + } + + *comp = (t == '6') ? 3 : 1; // '5' is 1-component .pgm; '6' is 3-component .ppm + + c = (char) stbi__get8(s); + stbi__pnm_skip_whitespace(s, &c); + + *x = stbi__pnm_getinteger(s, &c); // read width + if(*x == 0) + return stbi__err("invalid width", "PPM image header had zero or overflowing width"); + stbi__pnm_skip_whitespace(s, &c); + + *y = stbi__pnm_getinteger(s, &c); // read height + if (*y == 0) + return stbi__err("invalid width", "PPM image header had zero or overflowing width"); + stbi__pnm_skip_whitespace(s, &c); + + maxv = stbi__pnm_getinteger(s, &c); // read max value + if (maxv > 65535) + return stbi__err("max value > 65535", "PPM image supports only 8-bit and 16-bit images"); + else if (maxv > 255) + return 16; + else + return 8; +} + +static int stbi__pnm_is16(stbi__context *s) +{ + if (stbi__pnm_info(s, NULL, NULL, NULL) == 16) + return 1; + return 0; +} +#endif + +static int stbi__info_main(stbi__context *s, int *x, int *y, int *comp) +{ + #ifndef STBI_NO_JPEG + if (stbi__jpeg_info(s, x, y, comp)) return 1; + #endif + + #ifndef STBI_NO_PNG + if (stbi__png_info(s, x, y, comp)) return 1; + #endif + + #ifndef STBI_NO_GIF + if (stbi__gif_info(s, x, y, comp)) return 1; + #endif + + #ifndef STBI_NO_BMP + if (stbi__bmp_info(s, x, y, comp)) return 1; + #endif + + #ifndef STBI_NO_PSD + if (stbi__psd_info(s, x, y, comp)) return 1; + #endif + + #ifndef STBI_NO_PIC + if (stbi__pic_info(s, x, y, comp)) return 1; + #endif + + #ifndef STBI_NO_PNM + if (stbi__pnm_info(s, x, y, comp)) return 1; + #endif + + #ifndef STBI_NO_HDR + if (stbi__hdr_info(s, x, y, comp)) return 1; + #endif + + // test tga last because it's a crappy test! + #ifndef STBI_NO_TGA + if (stbi__tga_info(s, x, y, comp)) + return 1; + #endif + return stbi__err("unknown image type", "Image not of any known type, or corrupt"); +} + +static int stbi__is_16_main(stbi__context *s) +{ + #ifndef STBI_NO_PNG + if (stbi__png_is16(s)) return 1; + #endif + + #ifndef STBI_NO_PSD + if (stbi__psd_is16(s)) return 1; + #endif + + #ifndef STBI_NO_PNM + if (stbi__pnm_is16(s)) return 1; + #endif + return 0; +} + +#ifndef STBI_NO_STDIO +STBIDEF int stbi_info(char const *filename, int *x, int *y, int *comp) +{ + FILE *f = stbi__fopen(filename, "rb"); + int result; + if (!f) return stbi__err("can't fopen", "Unable to open file"); + result = stbi_info_from_file(f, x, y, comp); + fclose(f); + return result; +} + +STBIDEF int stbi_info_from_file(FILE *f, int *x, int *y, int *comp) +{ + int r; + stbi__context s; + long pos = ftell(f); + stbi__start_file(&s, f); + r = stbi__info_main(&s,x,y,comp); + fseek(f,pos,SEEK_SET); + return r; +} + +STBIDEF int stbi_is_16_bit(char const *filename) +{ + FILE *f = stbi__fopen(filename, "rb"); + int result; + if (!f) return stbi__err("can't fopen", "Unable to open file"); + result = stbi_is_16_bit_from_file(f); + fclose(f); + return result; +} + +STBIDEF int stbi_is_16_bit_from_file(FILE *f) +{ + int r; + stbi__context s; + long pos = ftell(f); + stbi__start_file(&s, f); + r = stbi__is_16_main(&s); + fseek(f,pos,SEEK_SET); + return r; +} +#endif // !STBI_NO_STDIO + +STBIDEF int stbi_info_from_memory(stbi_uc const *buffer, int len, int *x, int *y, int *comp) +{ + stbi__context s; + stbi__start_mem(&s,buffer,len); + return stbi__info_main(&s,x,y,comp); +} + +STBIDEF int stbi_info_from_callbacks(stbi_io_callbacks const *c, void *user, int *x, int *y, int *comp) +{ + stbi__context s; + stbi__start_callbacks(&s, (stbi_io_callbacks *) c, user); + return stbi__info_main(&s,x,y,comp); +} + +STBIDEF int stbi_is_16_bit_from_memory(stbi_uc const *buffer, int len) +{ + stbi__context s; + stbi__start_mem(&s,buffer,len); + return stbi__is_16_main(&s); +} + +STBIDEF int stbi_is_16_bit_from_callbacks(stbi_io_callbacks const *c, void *user) +{ + stbi__context s; + stbi__start_callbacks(&s, (stbi_io_callbacks *) c, user); + return stbi__is_16_main(&s); +} + +#endif // STB_IMAGE_IMPLEMENTATION + +/* + revision history: + 2.20 (2019-02-07) support utf8 filenames in Windows; fix warnings and platform ifdefs + 2.19 (2018-02-11) fix warning + 2.18 (2018-01-30) fix warnings + 2.17 (2018-01-29) change sbti__shiftsigned to avoid clang -O2 bug + 1-bit BMP + *_is_16_bit api + avoid warnings + 2.16 (2017-07-23) all functions have 16-bit variants; + STBI_NO_STDIO works again; + compilation fixes; + fix rounding in unpremultiply; + optimize vertical flip; + disable raw_len validation; + documentation fixes + 2.15 (2017-03-18) fix png-1,2,4 bug; now all Imagenet JPGs decode; + warning fixes; disable run-time SSE detection on gcc; + uniform handling of optional "return" values; + thread-safe initialization of zlib tables + 2.14 (2017-03-03) remove deprecated STBI_JPEG_OLD; fixes for Imagenet JPGs + 2.13 (2016-11-29) add 16-bit API, only supported for PNG right now + 2.12 (2016-04-02) fix typo in 2.11 PSD fix that caused crashes + 2.11 (2016-04-02) allocate large structures on the stack + remove white matting for transparent PSD + fix reported channel count for PNG & BMP + re-enable SSE2 in non-gcc 64-bit + support RGB-formatted JPEG + read 16-bit PNGs (only as 8-bit) + 2.10 (2016-01-22) avoid warning introduced in 2.09 by STBI_REALLOC_SIZED + 2.09 (2016-01-16) allow comments in PNM files + 16-bit-per-pixel TGA (not bit-per-component) + info() for TGA could break due to .hdr handling + info() for BMP to shares code instead of sloppy parse + can use STBI_REALLOC_SIZED if allocator doesn't support realloc + code cleanup + 2.08 (2015-09-13) fix to 2.07 cleanup, reading RGB PSD as RGBA + 2.07 (2015-09-13) fix compiler warnings + partial animated GIF support + limited 16-bpc PSD support + #ifdef unused functions + bug with < 92 byte PIC,PNM,HDR,TGA + 2.06 (2015-04-19) fix bug where PSD returns wrong '*comp' value + 2.05 (2015-04-19) fix bug in progressive JPEG handling, fix warning + 2.04 (2015-04-15) try to re-enable SIMD on MinGW 64-bit + 2.03 (2015-04-12) extra corruption checking (mmozeiko) + stbi_set_flip_vertically_on_load (nguillemot) + fix NEON support; fix mingw support + 2.02 (2015-01-19) fix incorrect assert, fix warning + 2.01 (2015-01-17) fix various warnings; suppress SIMD on gcc 32-bit without -msse2 + 2.00b (2014-12-25) fix STBI_MALLOC in progressive JPEG + 2.00 (2014-12-25) optimize JPG, including x86 SSE2 & NEON SIMD (ryg) + progressive JPEG (stb) + PGM/PPM support (Ken Miller) + STBI_MALLOC,STBI_REALLOC,STBI_FREE + GIF bugfix -- seemingly never worked + STBI_NO_*, STBI_ONLY_* + 1.48 (2014-12-14) fix incorrectly-named assert() + 1.47 (2014-12-14) 1/2/4-bit PNG support, both direct and paletted (Omar Cornut & stb) + optimize PNG (ryg) + fix bug in interlaced PNG with user-specified channel count (stb) + 1.46 (2014-08-26) + fix broken tRNS chunk (colorkey-style transparency) in non-paletted PNG + 1.45 (2014-08-16) + fix MSVC-ARM internal compiler error by wrapping malloc + 1.44 (2014-08-07) + various warning fixes from Ronny Chevalier + 1.43 (2014-07-15) + fix MSVC-only compiler problem in code changed in 1.42 + 1.42 (2014-07-09) + don't define _CRT_SECURE_NO_WARNINGS (affects user code) + fixes to stbi__cleanup_jpeg path + added STBI_ASSERT to avoid requiring assert.h + 1.41 (2014-06-25) + fix search&replace from 1.36 that messed up comments/error messages + 1.40 (2014-06-22) + fix gcc struct-initialization warning + 1.39 (2014-06-15) + fix to TGA optimization when req_comp != number of components in TGA; + fix to GIF loading because BMP wasn't rewinding (whoops, no GIFs in my test suite) + add support for BMP version 5 (more ignored fields) + 1.38 (2014-06-06) + suppress MSVC warnings on integer casts truncating values + fix accidental rename of 'skip' field of I/O + 1.37 (2014-06-04) + remove duplicate typedef + 1.36 (2014-06-03) + convert to header file single-file library + if de-iphone isn't set, load iphone images color-swapped instead of returning NULL + 1.35 (2014-05-27) + various warnings + fix broken STBI_SIMD path + fix bug where stbi_load_from_file no longer left file pointer in correct place + fix broken non-easy path for 32-bit BMP (possibly never used) + TGA optimization by Arseny Kapoulkine + 1.34 (unknown) + use STBI_NOTUSED in stbi__resample_row_generic(), fix one more leak in tga failure case + 1.33 (2011-07-14) + make stbi_is_hdr work in STBI_NO_HDR (as specified), minor compiler-friendly improvements + 1.32 (2011-07-13) + support for "info" function for all supported filetypes (SpartanJ) + 1.31 (2011-06-20) + a few more leak fixes, bug in PNG handling (SpartanJ) + 1.30 (2011-06-11) + added ability to load files via callbacks to accomidate custom input streams (Ben Wenger) + removed deprecated format-specific test/load functions + removed support for installable file formats (stbi_loader) -- would have been broken for IO callbacks anyway + error cases in bmp and tga give messages and don't leak (Raymond Barbiero, grisha) + fix inefficiency in decoding 32-bit BMP (David Woo) + 1.29 (2010-08-16) + various warning fixes from Aurelien Pocheville + 1.28 (2010-08-01) + fix bug in GIF palette transparency (SpartanJ) + 1.27 (2010-08-01) + cast-to-stbi_uc to fix warnings + 1.26 (2010-07-24) + fix bug in file buffering for PNG reported by SpartanJ + 1.25 (2010-07-17) + refix trans_data warning (Won Chun) + 1.24 (2010-07-12) + perf improvements reading from files on platforms with lock-heavy fgetc() + minor perf improvements for jpeg + deprecated type-specific functions so we'll get feedback if they're needed + attempt to fix trans_data warning (Won Chun) + 1.23 fixed bug in iPhone support + 1.22 (2010-07-10) + removed image *writing* support + stbi_info support from Jetro Lauha + GIF support from Jean-Marc Lienher + iPhone PNG-extensions from James Brown + warning-fixes from Nicolas Schulz and Janez Zemva (i.stbi__err. Janez (U+017D)emva) + 1.21 fix use of 'stbi_uc' in header (reported by jon blow) + 1.20 added support for Softimage PIC, by Tom Seddon + 1.19 bug in interlaced PNG corruption check (found by ryg) + 1.18 (2008-08-02) + fix a threading bug (local mutable static) + 1.17 support interlaced PNG + 1.16 major bugfix - stbi__convert_format converted one too many pixels + 1.15 initialize some fields for thread safety + 1.14 fix threadsafe conversion bug + header-file-only version (#define STBI_HEADER_FILE_ONLY before including) + 1.13 threadsafe + 1.12 const qualifiers in the API + 1.11 Support installable IDCT, colorspace conversion routines + 1.10 Fixes for 64-bit (don't use "unsigned long") + optimized upsampling by Fabian "ryg" Giesen + 1.09 Fix format-conversion for PSD code (bad global variables!) + 1.08 Thatcher Ulrich's PSD code integrated by Nicolas Schulz + 1.07 attempt to fix C++ warning/errors again + 1.06 attempt to fix C++ warning/errors again + 1.05 fix TGA loading to return correct *comp and use good luminance calc + 1.04 default float alpha is 1, not 255; use 'void *' for stbi_image_free + 1.03 bugfixes to STBI_NO_STDIO, STBI_NO_HDR + 1.02 support for (subset of) HDR files, float interface for preferred access to them + 1.01 fix bug: possible bug in handling right-side up bmps... not sure + fix bug: the stbi__bmp_load() and stbi__tga_load() functions didn't work at all + 1.00 interface to zlib that skips zlib header + 0.99 correct handling of alpha in palette + 0.98 TGA loader by lonesock; dynamically add loaders (untested) + 0.97 jpeg errors on too large a file; also catch another malloc failure + 0.96 fix detection of invalid v value - particleman@mollyrocket forum + 0.95 during header scan, seek to markers in case of padding + 0.94 STBI_NO_STDIO to disable stdio usage; rename all #defines the same + 0.93 handle jpegtran output; verbose errors + 0.92 read 4,8,16,24,32-bit BMP files of several formats + 0.91 output 24-bit Windows 3.0 BMP files + 0.90 fix a few more warnings; bump version number to approach 1.0 + 0.61 bugfixes due to Marc LeBlanc, Christopher Lloyd + 0.60 fix compiling as c++ + 0.59 fix warnings: merge Dave Moore's -Wall fixes + 0.58 fix bug: zlib uncompressed mode len/nlen was wrong endian + 0.57 fix bug: jpg last huffman symbol before marker was >9 bits but less than 16 available + 0.56 fix bug: zlib uncompressed mode len vs. nlen + 0.55 fix bug: restart_interval not initialized to 0 + 0.54 allow NULL for 'int *comp' + 0.53 fix bug in png 3->4; speedup png decoding + 0.52 png handles req_comp=3,4 directly; minor cleanup; jpeg comments + 0.51 obey req_comp requests, 1-component jpegs return as 1-component, + on 'test' only check type, not whether we support this variant + 0.50 (2006-11-19) + first released version +*/ + + +/* +------------------------------------------------------------------------------ +This software is available under 2 licenses -- choose whichever you prefer. +------------------------------------------------------------------------------ +ALTERNATIVE A - MIT License +Copyright (c) 2017 Sean Barrett +Permission is hereby granted, free of charge, to any person obtaining a copy of +this software and associated documentation files (the "Software"), to deal in +the Software without restriction, including without limitation the rights to +use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies +of the Software, and to permit persons to whom the Software is furnished to do +so, subject to the following conditions: +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. +------------------------------------------------------------------------------ +ALTERNATIVE B - Public Domain (www.unlicense.org) +This is free and unencumbered software released into the public domain. +Anyone is free to copy, modify, publish, use, compile, sell, or distribute this +software, either in source code form or as a compiled binary, for any purpose, +commercial or non-commercial, and by any means. +In jurisdictions that recognize copyright laws, the author or authors of this +software dedicate any and all copyright interest in the software to the public +domain. We make this dedication for the benefit of the public at large and to +the detriment of our heirs and successors. We intend this dedication to be an +overt act of relinquishment in perpetuity of all present and future rights to +this software under copyright law. +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN +ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION +WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. +------------------------------------------------------------------------------ +*/ diff --git a/third-party/stb/stb_image_resize.h b/third-party/stb/stb_image_resize.h new file mode 100644 index 00000000000..ef9e6fe8704 --- /dev/null +++ b/third-party/stb/stb_image_resize.h @@ -0,0 +1,2634 @@ +/* stb_image_resize - v0.97 - public domain image resizing + by Jorge L Rodriguez (@VinoBS) - 2014 + http://github.com/nothings/stb + + Written with emphasis on usability, portability, and efficiency. (No + SIMD or threads, so it be easily outperformed by libs that use those.) + Only scaling and translation is supported, no rotations or shears. + Easy API downsamples w/Mitchell filter, upsamples w/cubic interpolation. + + COMPILING & LINKING + In one C/C++ file that #includes this file, do this: + #define STB_IMAGE_RESIZE_IMPLEMENTATION + before the #include. That will create the implementation in that file. + + QUICKSTART + stbir_resize_uint8( input_pixels , in_w , in_h , 0, + output_pixels, out_w, out_h, 0, num_channels) + stbir_resize_float(...) + stbir_resize_uint8_srgb( input_pixels , in_w , in_h , 0, + output_pixels, out_w, out_h, 0, + num_channels , alpha_chan , 0) + stbir_resize_uint8_srgb_edgemode( + input_pixels , in_w , in_h , 0, + output_pixels, out_w, out_h, 0, + num_channels , alpha_chan , 0, STBIR_EDGE_CLAMP) + // WRAP/REFLECT/ZERO + + FULL API + See the "header file" section of the source for API documentation. + + ADDITIONAL DOCUMENTATION + + SRGB & FLOATING POINT REPRESENTATION + The sRGB functions presume IEEE floating point. If you do not have + IEEE floating point, define STBIR_NON_IEEE_FLOAT. This will use + a slower implementation. + + MEMORY ALLOCATION + The resize functions here perform a single memory allocation using + malloc. To control the memory allocation, before the #include that + triggers the implementation, do: + + #define STBIR_MALLOC(size,context) ... + #define STBIR_FREE(ptr,context) ... + + Each resize function makes exactly one call to malloc/free, so to use + temp memory, store the temp memory in the context and return that. + + ASSERT + Define STBIR_ASSERT(boolval) to override assert() and not use assert.h + + OPTIMIZATION + Define STBIR_SATURATE_INT to compute clamp values in-range using + integer operations instead of float operations. This may be faster + on some platforms. + + DEFAULT FILTERS + For functions which don't provide explicit control over what filters + to use, you can change the compile-time defaults with + + #define STBIR_DEFAULT_FILTER_UPSAMPLE STBIR_FILTER_something + #define STBIR_DEFAULT_FILTER_DOWNSAMPLE STBIR_FILTER_something + + See stbir_filter in the header-file section for the list of filters. + + NEW FILTERS + A number of 1D filter kernels are used. For a list of + supported filters see the stbir_filter enum. To add a new filter, + write a filter function and add it to stbir__filter_info_table. + + PROGRESS + For interactive use with slow resize operations, you can install + a progress-report callback: + + #define STBIR_PROGRESS_REPORT(val) some_func(val) + + The parameter val is a float which goes from 0 to 1 as progress is made. + + For example: + + static void my_progress_report(float progress); + #define STBIR_PROGRESS_REPORT(val) my_progress_report(val) + + #define STB_IMAGE_RESIZE_IMPLEMENTATION + #include "stb_image_resize.h" + + static void my_progress_report(float progress) + { + printf("Progress: %f%%\n", progress*100); + } + + MAX CHANNELS + If your image has more than 64 channels, define STBIR_MAX_CHANNELS + to the max you'll have. + + ALPHA CHANNEL + Most of the resizing functions provide the ability to control how + the alpha channel of an image is processed. The important things + to know about this: + + 1. The best mathematically-behaved version of alpha to use is + called "premultiplied alpha", in which the other color channels + have had the alpha value multiplied in. If you use premultiplied + alpha, linear filtering (such as image resampling done by this + library, or performed in texture units on GPUs) does the "right + thing". While premultiplied alpha is standard in the movie CGI + industry, it is still uncommon in the videogame/real-time world. + + If you linearly filter non-premultiplied alpha, strange effects + occur. (For example, the 50/50 average of 99% transparent bright green + and 1% transparent black produces 50% transparent dark green when + non-premultiplied, whereas premultiplied it produces 50% + transparent near-black. The former introduces green energy + that doesn't exist in the source image.) + + 2. Artists should not edit premultiplied-alpha images; artists + want non-premultiplied alpha images. Thus, art tools generally output + non-premultiplied alpha images. + + 3. You will get best results in most cases by converting images + to premultiplied alpha before processing them mathematically. + + 4. If you pass the flag STBIR_FLAG_ALPHA_PREMULTIPLIED, the + resizer does not do anything special for the alpha channel; + it is resampled identically to other channels. This produces + the correct results for premultiplied-alpha images, but produces + less-than-ideal results for non-premultiplied-alpha images. + + 5. If you do not pass the flag STBIR_FLAG_ALPHA_PREMULTIPLIED, + then the resizer weights the contribution of input pixels + based on their alpha values, or, equivalently, it multiplies + the alpha value into the color channels, resamples, then divides + by the resultant alpha value. Input pixels which have alpha=0 do + not contribute at all to output pixels unless _all_ of the input + pixels affecting that output pixel have alpha=0, in which case + the result for that pixel is the same as it would be without + STBIR_FLAG_ALPHA_PREMULTIPLIED. However, this is only true for + input images in integer formats. For input images in float format, + input pixels with alpha=0 have no effect, and output pixels + which have alpha=0 will be 0 in all channels. (For float images, + you can manually achieve the same result by adding a tiny epsilon + value to the alpha channel of every image, and then subtracting + or clamping it at the end.) + + 6. You can suppress the behavior described in #5 and make + all-0-alpha pixels have 0 in all channels by #defining + STBIR_NO_ALPHA_EPSILON. + + 7. You can separately control whether the alpha channel is + interpreted as linear or affected by the colorspace. By default + it is linear; you almost never want to apply the colorspace. + (For example, graphics hardware does not apply sRGB conversion + to the alpha channel.) + + CONTRIBUTORS + Jorge L Rodriguez: Implementation + Sean Barrett: API design, optimizations + Aras Pranckevicius: bugfix + Nathan Reed: warning fixes + + REVISIONS + 0.97 (2020-02-02) fixed warning + 0.96 (2019-03-04) fixed warnings + 0.95 (2017-07-23) fixed warnings + 0.94 (2017-03-18) fixed warnings + 0.93 (2017-03-03) fixed bug with certain combinations of heights + 0.92 (2017-01-02) fix integer overflow on large (>2GB) images + 0.91 (2016-04-02) fix warnings; fix handling of subpixel regions + 0.90 (2014-09-17) first released version + + LICENSE + See end of file for license information. + + TODO + Don't decode all of the image data when only processing a partial tile + Don't use full-width decode buffers when only processing a partial tile + When processing wide images, break processing into tiles so data fits in L1 cache + Installable filters? + Resize that respects alpha test coverage + (Reference code: FloatImage::alphaTestCoverage and FloatImage::scaleAlphaToCoverage: + https://code.google.com/p/nvidia-texture-tools/source/browse/trunk/src/nvimage/FloatImage.cpp ) +*/ + +#ifndef STBIR_INCLUDE_STB_IMAGE_RESIZE_H +#define STBIR_INCLUDE_STB_IMAGE_RESIZE_H + +#ifdef _MSC_VER +typedef unsigned char stbir_uint8; +typedef unsigned short stbir_uint16; +typedef unsigned int stbir_uint32; +#else +#include +typedef uint8_t stbir_uint8; +typedef uint16_t stbir_uint16; +typedef uint32_t stbir_uint32; +#endif + +#ifndef STBIRDEF +#ifdef STB_IMAGE_RESIZE_STATIC +#define STBIRDEF static +#else +#ifdef __cplusplus +#define STBIRDEF extern "C" +#else +#define STBIRDEF extern +#endif +#endif +#endif + +////////////////////////////////////////////////////////////////////////////// +// +// Easy-to-use API: +// +// * "input pixels" points to an array of image data with 'num_channels' channels (e.g. RGB=3, RGBA=4) +// * input_w is input image width (x-axis), input_h is input image height (y-axis) +// * stride is the offset between successive rows of image data in memory, in bytes. you can +// specify 0 to mean packed continuously in memory +// * alpha channel is treated identically to other channels. +// * colorspace is linear or sRGB as specified by function name +// * returned result is 1 for success or 0 in case of an error. +// #define STBIR_ASSERT() to trigger an assert on parameter validation errors. +// * Memory required grows approximately linearly with input and output size, but with +// discontinuities at input_w == output_w and input_h == output_h. +// * These functions use a "default" resampling filter defined at compile time. To change the filter, +// you can change the compile-time defaults by #defining STBIR_DEFAULT_FILTER_UPSAMPLE +// and STBIR_DEFAULT_FILTER_DOWNSAMPLE, or you can use the medium-complexity API. + +STBIRDEF int stbir_resize_uint8( const unsigned char *input_pixels , int input_w , int input_h , int input_stride_in_bytes, + unsigned char *output_pixels, int output_w, int output_h, int output_stride_in_bytes, + int num_channels); + +STBIRDEF int stbir_resize_float( const float *input_pixels , int input_w , int input_h , int input_stride_in_bytes, + float *output_pixels, int output_w, int output_h, int output_stride_in_bytes, + int num_channels); + + +// The following functions interpret image data as gamma-corrected sRGB. +// Specify STBIR_ALPHA_CHANNEL_NONE if you have no alpha channel, +// or otherwise provide the index of the alpha channel. Flags value +// of 0 will probably do the right thing if you're not sure what +// the flags mean. + +#define STBIR_ALPHA_CHANNEL_NONE -1 + +// Set this flag if your texture has premultiplied alpha. Otherwise, stbir will +// use alpha-weighted resampling (effectively premultiplying, resampling, +// then unpremultiplying). +#define STBIR_FLAG_ALPHA_PREMULTIPLIED (1 << 0) +// The specified alpha channel should be handled as gamma-corrected value even +// when doing sRGB operations. +#define STBIR_FLAG_ALPHA_USES_COLORSPACE (1 << 1) + +STBIRDEF int stbir_resize_uint8_srgb(const unsigned char *input_pixels , int input_w , int input_h , int input_stride_in_bytes, + unsigned char *output_pixels, int output_w, int output_h, int output_stride_in_bytes, + int num_channels, int alpha_channel, int flags); + + +typedef enum +{ + STBIR_EDGE_CLAMP = 1, + STBIR_EDGE_REFLECT = 2, + STBIR_EDGE_WRAP = 3, + STBIR_EDGE_ZERO = 4, +} stbir_edge; + +// This function adds the ability to specify how requests to sample off the edge of the image are handled. +STBIRDEF int stbir_resize_uint8_srgb_edgemode(const unsigned char *input_pixels , int input_w , int input_h , int input_stride_in_bytes, + unsigned char *output_pixels, int output_w, int output_h, int output_stride_in_bytes, + int num_channels, int alpha_channel, int flags, + stbir_edge edge_wrap_mode); + +////////////////////////////////////////////////////////////////////////////// +// +// Medium-complexity API +// +// This extends the easy-to-use API as follows: +// +// * Alpha-channel can be processed separately +// * If alpha_channel is not STBIR_ALPHA_CHANNEL_NONE +// * Alpha channel will not be gamma corrected (unless flags&STBIR_FLAG_GAMMA_CORRECT) +// * Filters will be weighted by alpha channel (unless flags&STBIR_FLAG_ALPHA_PREMULTIPLIED) +// * Filter can be selected explicitly +// * uint16 image type +// * sRGB colorspace available for all types +// * context parameter for passing to STBIR_MALLOC + +typedef enum +{ + STBIR_FILTER_DEFAULT = 0, // use same filter type that easy-to-use API chooses + STBIR_FILTER_BOX = 1, // A trapezoid w/1-pixel wide ramps, same result as box for integer scale ratios + STBIR_FILTER_TRIANGLE = 2, // On upsampling, produces same results as bilinear texture filtering + STBIR_FILTER_CUBICBSPLINE = 3, // The cubic b-spline (aka Mitchell-Netrevalli with B=1,C=0), gaussian-esque + STBIR_FILTER_CATMULLROM = 4, // An interpolating cubic spline + STBIR_FILTER_MITCHELL = 5, // Mitchell-Netrevalli filter with B=1/3, C=1/3 +} stbir_filter; + +typedef enum +{ + STBIR_COLORSPACE_LINEAR, + STBIR_COLORSPACE_SRGB, + + STBIR_MAX_COLORSPACES, +} stbir_colorspace; + +// The following functions are all identical except for the type of the image data + +STBIRDEF int stbir_resize_uint8_generic( const unsigned char *input_pixels , int input_w , int input_h , int input_stride_in_bytes, + unsigned char *output_pixels, int output_w, int output_h, int output_stride_in_bytes, + int num_channels, int alpha_channel, int flags, + stbir_edge edge_wrap_mode, stbir_filter filter, stbir_colorspace space, + void *alloc_context); + +STBIRDEF int stbir_resize_uint16_generic(const stbir_uint16 *input_pixels , int input_w , int input_h , int input_stride_in_bytes, + stbir_uint16 *output_pixels , int output_w, int output_h, int output_stride_in_bytes, + int num_channels, int alpha_channel, int flags, + stbir_edge edge_wrap_mode, stbir_filter filter, stbir_colorspace space, + void *alloc_context); + +STBIRDEF int stbir_resize_float_generic( const float *input_pixels , int input_w , int input_h , int input_stride_in_bytes, + float *output_pixels , int output_w, int output_h, int output_stride_in_bytes, + int num_channels, int alpha_channel, int flags, + stbir_edge edge_wrap_mode, stbir_filter filter, stbir_colorspace space, + void *alloc_context); + + + +////////////////////////////////////////////////////////////////////////////// +// +// Full-complexity API +// +// This extends the medium API as follows: +// +// * uint32 image type +// * not typesafe +// * separate filter types for each axis +// * separate edge modes for each axis +// * can specify scale explicitly for subpixel correctness +// * can specify image source tile using texture coordinates + +typedef enum +{ + STBIR_TYPE_UINT8 , + STBIR_TYPE_UINT16, + STBIR_TYPE_UINT32, + STBIR_TYPE_FLOAT , + + STBIR_MAX_TYPES +} stbir_datatype; + +STBIRDEF int stbir_resize( const void *input_pixels , int input_w , int input_h , int input_stride_in_bytes, + void *output_pixels, int output_w, int output_h, int output_stride_in_bytes, + stbir_datatype datatype, + int num_channels, int alpha_channel, int flags, + stbir_edge edge_mode_horizontal, stbir_edge edge_mode_vertical, + stbir_filter filter_horizontal, stbir_filter filter_vertical, + stbir_colorspace space, void *alloc_context); + +STBIRDEF int stbir_resize_subpixel(const void *input_pixels , int input_w , int input_h , int input_stride_in_bytes, + void *output_pixels, int output_w, int output_h, int output_stride_in_bytes, + stbir_datatype datatype, + int num_channels, int alpha_channel, int flags, + stbir_edge edge_mode_horizontal, stbir_edge edge_mode_vertical, + stbir_filter filter_horizontal, stbir_filter filter_vertical, + stbir_colorspace space, void *alloc_context, + float x_scale, float y_scale, + float x_offset, float y_offset); + +STBIRDEF int stbir_resize_region( const void *input_pixels , int input_w , int input_h , int input_stride_in_bytes, + void *output_pixels, int output_w, int output_h, int output_stride_in_bytes, + stbir_datatype datatype, + int num_channels, int alpha_channel, int flags, + stbir_edge edge_mode_horizontal, stbir_edge edge_mode_vertical, + stbir_filter filter_horizontal, stbir_filter filter_vertical, + stbir_colorspace space, void *alloc_context, + float s0, float t0, float s1, float t1); +// (s0, t0) & (s1, t1) are the top-left and bottom right corner (uv addressing style: [0, 1]x[0, 1]) of a region of the input image to use. + +// +// +//// end header file ///////////////////////////////////////////////////// +#endif // STBIR_INCLUDE_STB_IMAGE_RESIZE_H + + + + + +#ifdef STB_IMAGE_RESIZE_IMPLEMENTATION + +#ifndef STBIR_ASSERT +#include +#define STBIR_ASSERT(x) assert(x) +#endif + +// For memset +#include + +#include + +#ifndef STBIR_MALLOC +#include +// use comma operator to evaluate c, to avoid "unused parameter" warnings +#define STBIR_MALLOC(size,c) ((void)(c), malloc(size)) +#define STBIR_FREE(ptr,c) ((void)(c), free(ptr)) +#endif + +#ifndef _MSC_VER +#ifdef __cplusplus +#define stbir__inline inline +#else +#define stbir__inline +#endif +#else +#define stbir__inline __forceinline +#endif + + +// should produce compiler error if size is wrong +typedef unsigned char stbir__validate_uint32[sizeof(stbir_uint32) == 4 ? 1 : -1]; + +#ifdef _MSC_VER +#define STBIR__NOTUSED(v) (void)(v) +#else +#define STBIR__NOTUSED(v) (void)sizeof(v) +#endif + +#define STBIR__ARRAY_SIZE(a) (sizeof((a))/sizeof((a)[0])) + +#ifndef STBIR_DEFAULT_FILTER_UPSAMPLE +#define STBIR_DEFAULT_FILTER_UPSAMPLE STBIR_FILTER_CATMULLROM +#endif + +#ifndef STBIR_DEFAULT_FILTER_DOWNSAMPLE +#define STBIR_DEFAULT_FILTER_DOWNSAMPLE STBIR_FILTER_MITCHELL +#endif + +#ifndef STBIR_PROGRESS_REPORT +#define STBIR_PROGRESS_REPORT(float_0_to_1) +#endif + +#ifndef STBIR_MAX_CHANNELS +#define STBIR_MAX_CHANNELS 64 +#endif + +#if STBIR_MAX_CHANNELS > 65536 +#error "Too many channels; STBIR_MAX_CHANNELS must be no more than 65536." +// because we store the indices in 16-bit variables +#endif + +// This value is added to alpha just before premultiplication to avoid +// zeroing out color values. It is equivalent to 2^-80. If you don't want +// that behavior (it may interfere if you have floating point images with +// very small alpha values) then you can define STBIR_NO_ALPHA_EPSILON to +// disable it. +#ifndef STBIR_ALPHA_EPSILON +#define STBIR_ALPHA_EPSILON ((float)1 / (1 << 20) / (1 << 20) / (1 << 20) / (1 << 20)) +#endif + + + +#ifdef _MSC_VER +#define STBIR__UNUSED_PARAM(v) (void)(v) +#else +#define STBIR__UNUSED_PARAM(v) (void)sizeof(v) +#endif + +// must match stbir_datatype +static unsigned char stbir__type_size[] = { + 1, // STBIR_TYPE_UINT8 + 2, // STBIR_TYPE_UINT16 + 4, // STBIR_TYPE_UINT32 + 4, // STBIR_TYPE_FLOAT +}; + +// Kernel function centered at 0 +typedef float (stbir__kernel_fn)(float x, float scale); +typedef float (stbir__support_fn)(float scale); + +typedef struct +{ + stbir__kernel_fn* kernel; + stbir__support_fn* support; +} stbir__filter_info; + +// When upsampling, the contributors are which source pixels contribute. +// When downsampling, the contributors are which destination pixels are contributed to. +typedef struct +{ + int n0; // First contributing pixel + int n1; // Last contributing pixel +} stbir__contributors; + +typedef struct +{ + const void* input_data; + int input_w; + int input_h; + int input_stride_bytes; + + void* output_data; + int output_w; + int output_h; + int output_stride_bytes; + + float s0, t0, s1, t1; + + float horizontal_shift; // Units: output pixels + float vertical_shift; // Units: output pixels + float horizontal_scale; + float vertical_scale; + + int channels; + int alpha_channel; + stbir_uint32 flags; + stbir_datatype type; + stbir_filter horizontal_filter; + stbir_filter vertical_filter; + stbir_edge edge_horizontal; + stbir_edge edge_vertical; + stbir_colorspace colorspace; + + stbir__contributors* horizontal_contributors; + float* horizontal_coefficients; + + stbir__contributors* vertical_contributors; + float* vertical_coefficients; + + int decode_buffer_pixels; + float* decode_buffer; + + float* horizontal_buffer; + + // cache these because ceil/floor are inexplicably showing up in profile + int horizontal_coefficient_width; + int vertical_coefficient_width; + int horizontal_filter_pixel_width; + int vertical_filter_pixel_width; + int horizontal_filter_pixel_margin; + int vertical_filter_pixel_margin; + int horizontal_num_contributors; + int vertical_num_contributors; + + int ring_buffer_length_bytes; // The length of an individual entry in the ring buffer. The total number of ring buffers is stbir__get_filter_pixel_width(filter) + int ring_buffer_num_entries; // Total number of entries in the ring buffer. + int ring_buffer_first_scanline; + int ring_buffer_last_scanline; + int ring_buffer_begin_index; // first_scanline is at this index in the ring buffer + float* ring_buffer; + + float* encode_buffer; // A temporary buffer to store floats so we don't lose precision while we do multiply-adds. + + int horizontal_contributors_size; + int horizontal_coefficients_size; + int vertical_contributors_size; + int vertical_coefficients_size; + int decode_buffer_size; + int horizontal_buffer_size; + int ring_buffer_size; + int encode_buffer_size; +} stbir__info; + + +static const float stbir__max_uint8_as_float = 255.0f; +static const float stbir__max_uint16_as_float = 65535.0f; +static const double stbir__max_uint32_as_float = 4294967295.0; + + +static stbir__inline int stbir__min(int a, int b) +{ + return a < b ? a : b; +} + +static stbir__inline float stbir__saturate(float x) +{ + if (x < 0) + return 0; + + if (x > 1) + return 1; + + return x; +} + +#ifdef STBIR_SATURATE_INT +static stbir__inline stbir_uint8 stbir__saturate8(int x) +{ + if ((unsigned int) x <= 255) + return x; + + if (x < 0) + return 0; + + return 255; +} + +static stbir__inline stbir_uint16 stbir__saturate16(int x) +{ + if ((unsigned int) x <= 65535) + return x; + + if (x < 0) + return 0; + + return 65535; +} +#endif + +static float stbir__srgb_uchar_to_linear_float[256] = { + 0.000000f, 0.000304f, 0.000607f, 0.000911f, 0.001214f, 0.001518f, 0.001821f, 0.002125f, 0.002428f, 0.002732f, 0.003035f, + 0.003347f, 0.003677f, 0.004025f, 0.004391f, 0.004777f, 0.005182f, 0.005605f, 0.006049f, 0.006512f, 0.006995f, 0.007499f, + 0.008023f, 0.008568f, 0.009134f, 0.009721f, 0.010330f, 0.010960f, 0.011612f, 0.012286f, 0.012983f, 0.013702f, 0.014444f, + 0.015209f, 0.015996f, 0.016807f, 0.017642f, 0.018500f, 0.019382f, 0.020289f, 0.021219f, 0.022174f, 0.023153f, 0.024158f, + 0.025187f, 0.026241f, 0.027321f, 0.028426f, 0.029557f, 0.030713f, 0.031896f, 0.033105f, 0.034340f, 0.035601f, 0.036889f, + 0.038204f, 0.039546f, 0.040915f, 0.042311f, 0.043735f, 0.045186f, 0.046665f, 0.048172f, 0.049707f, 0.051269f, 0.052861f, + 0.054480f, 0.056128f, 0.057805f, 0.059511f, 0.061246f, 0.063010f, 0.064803f, 0.066626f, 0.068478f, 0.070360f, 0.072272f, + 0.074214f, 0.076185f, 0.078187f, 0.080220f, 0.082283f, 0.084376f, 0.086500f, 0.088656f, 0.090842f, 0.093059f, 0.095307f, + 0.097587f, 0.099899f, 0.102242f, 0.104616f, 0.107023f, 0.109462f, 0.111932f, 0.114435f, 0.116971f, 0.119538f, 0.122139f, + 0.124772f, 0.127438f, 0.130136f, 0.132868f, 0.135633f, 0.138432f, 0.141263f, 0.144128f, 0.147027f, 0.149960f, 0.152926f, + 0.155926f, 0.158961f, 0.162029f, 0.165132f, 0.168269f, 0.171441f, 0.174647f, 0.177888f, 0.181164f, 0.184475f, 0.187821f, + 0.191202f, 0.194618f, 0.198069f, 0.201556f, 0.205079f, 0.208637f, 0.212231f, 0.215861f, 0.219526f, 0.223228f, 0.226966f, + 0.230740f, 0.234551f, 0.238398f, 0.242281f, 0.246201f, 0.250158f, 0.254152f, 0.258183f, 0.262251f, 0.266356f, 0.270498f, + 0.274677f, 0.278894f, 0.283149f, 0.287441f, 0.291771f, 0.296138f, 0.300544f, 0.304987f, 0.309469f, 0.313989f, 0.318547f, + 0.323143f, 0.327778f, 0.332452f, 0.337164f, 0.341914f, 0.346704f, 0.351533f, 0.356400f, 0.361307f, 0.366253f, 0.371238f, + 0.376262f, 0.381326f, 0.386430f, 0.391573f, 0.396755f, 0.401978f, 0.407240f, 0.412543f, 0.417885f, 0.423268f, 0.428691f, + 0.434154f, 0.439657f, 0.445201f, 0.450786f, 0.456411f, 0.462077f, 0.467784f, 0.473532f, 0.479320f, 0.485150f, 0.491021f, + 0.496933f, 0.502887f, 0.508881f, 0.514918f, 0.520996f, 0.527115f, 0.533276f, 0.539480f, 0.545725f, 0.552011f, 0.558340f, + 0.564712f, 0.571125f, 0.577581f, 0.584078f, 0.590619f, 0.597202f, 0.603827f, 0.610496f, 0.617207f, 0.623960f, 0.630757f, + 0.637597f, 0.644480f, 0.651406f, 0.658375f, 0.665387f, 0.672443f, 0.679543f, 0.686685f, 0.693872f, 0.701102f, 0.708376f, + 0.715694f, 0.723055f, 0.730461f, 0.737911f, 0.745404f, 0.752942f, 0.760525f, 0.768151f, 0.775822f, 0.783538f, 0.791298f, + 0.799103f, 0.806952f, 0.814847f, 0.822786f, 0.830770f, 0.838799f, 0.846873f, 0.854993f, 0.863157f, 0.871367f, 0.879622f, + 0.887923f, 0.896269f, 0.904661f, 0.913099f, 0.921582f, 0.930111f, 0.938686f, 0.947307f, 0.955974f, 0.964686f, 0.973445f, + 0.982251f, 0.991102f, 1.0f +}; + +static float stbir__srgb_to_linear(float f) +{ + if (f <= 0.04045f) + return f / 12.92f; + else + return (float)pow((f + 0.055f) / 1.055f, 2.4f); +} + +static float stbir__linear_to_srgb(float f) +{ + if (f <= 0.0031308f) + return f * 12.92f; + else + return 1.055f * (float)pow(f, 1 / 2.4f) - 0.055f; +} + +#ifndef STBIR_NON_IEEE_FLOAT +// From https://gist.github.com/rygorous/2203834 + +typedef union +{ + stbir_uint32 u; + float f; +} stbir__FP32; + +static const stbir_uint32 fp32_to_srgb8_tab4[104] = { + 0x0073000d, 0x007a000d, 0x0080000d, 0x0087000d, 0x008d000d, 0x0094000d, 0x009a000d, 0x00a1000d, + 0x00a7001a, 0x00b4001a, 0x00c1001a, 0x00ce001a, 0x00da001a, 0x00e7001a, 0x00f4001a, 0x0101001a, + 0x010e0033, 0x01280033, 0x01410033, 0x015b0033, 0x01750033, 0x018f0033, 0x01a80033, 0x01c20033, + 0x01dc0067, 0x020f0067, 0x02430067, 0x02760067, 0x02aa0067, 0x02dd0067, 0x03110067, 0x03440067, + 0x037800ce, 0x03df00ce, 0x044600ce, 0x04ad00ce, 0x051400ce, 0x057b00c5, 0x05dd00bc, 0x063b00b5, + 0x06970158, 0x07420142, 0x07e30130, 0x087b0120, 0x090b0112, 0x09940106, 0x0a1700fc, 0x0a9500f2, + 0x0b0f01cb, 0x0bf401ae, 0x0ccb0195, 0x0d950180, 0x0e56016e, 0x0f0d015e, 0x0fbc0150, 0x10630143, + 0x11070264, 0x1238023e, 0x1357021d, 0x14660201, 0x156601e9, 0x165a01d3, 0x174401c0, 0x182401af, + 0x18fe0331, 0x1a9602fe, 0x1c1502d2, 0x1d7e02ad, 0x1ed4028d, 0x201a0270, 0x21520256, 0x227d0240, + 0x239f0443, 0x25c003fe, 0x27bf03c4, 0x29a10392, 0x2b6a0367, 0x2d1d0341, 0x2ebe031f, 0x304d0300, + 0x31d105b0, 0x34a80555, 0x37520507, 0x39d504c5, 0x3c37048b, 0x3e7c0458, 0x40a8042a, 0x42bd0401, + 0x44c20798, 0x488e071e, 0x4c1c06b6, 0x4f76065d, 0x52a50610, 0x55ac05cc, 0x5892058f, 0x5b590559, + 0x5e0c0a23, 0x631c0980, 0x67db08f6, 0x6c55087f, 0x70940818, 0x74a007bd, 0x787d076c, 0x7c330723, +}; + +static stbir_uint8 stbir__linear_to_srgb_uchar(float in) +{ + static const stbir__FP32 almostone = { 0x3f7fffff }; // 1-eps + static const stbir__FP32 minval = { (127-13) << 23 }; + stbir_uint32 tab,bias,scale,t; + stbir__FP32 f; + + // Clamp to [2^(-13), 1-eps]; these two values map to 0 and 1, respectively. + // The tests are carefully written so that NaNs map to 0, same as in the reference + // implementation. + if (!(in > minval.f)) // written this way to catch NaNs + in = minval.f; + if (in > almostone.f) + in = almostone.f; + + // Do the table lookup and unpack bias, scale + f.f = in; + tab = fp32_to_srgb8_tab4[(f.u - minval.u) >> 20]; + bias = (tab >> 16) << 9; + scale = tab & 0xffff; + + // Grab next-highest mantissa bits and perform linear interpolation + t = (f.u >> 12) & 0xff; + return (unsigned char) ((bias + scale*t) >> 16); +} + +#else +// sRGB transition values, scaled by 1<<28 +static int stbir__srgb_offset_to_linear_scaled[256] = +{ + 0, 40738, 122216, 203693, 285170, 366648, 448125, 529603, + 611080, 692557, 774035, 855852, 942009, 1033024, 1128971, 1229926, + 1335959, 1447142, 1563542, 1685229, 1812268, 1944725, 2082664, 2226148, + 2375238, 2529996, 2690481, 2856753, 3028870, 3206888, 3390865, 3580856, + 3776916, 3979100, 4187460, 4402049, 4622919, 4850123, 5083710, 5323731, + 5570236, 5823273, 6082892, 6349140, 6622065, 6901714, 7188133, 7481369, + 7781466, 8088471, 8402427, 8723380, 9051372, 9386448, 9728650, 10078021, + 10434603, 10798439, 11169569, 11548036, 11933879, 12327139, 12727857, 13136073, + 13551826, 13975156, 14406100, 14844697, 15290987, 15745007, 16206795, 16676389, + 17153826, 17639142, 18132374, 18633560, 19142734, 19659934, 20185196, 20718552, + 21260042, 21809696, 22367554, 22933648, 23508010, 24090680, 24681686, 25281066, + 25888850, 26505076, 27129772, 27762974, 28404716, 29055026, 29713942, 30381490, + 31057708, 31742624, 32436272, 33138682, 33849884, 34569912, 35298800, 36036568, + 36783260, 37538896, 38303512, 39077136, 39859796, 40651528, 41452360, 42262316, + 43081432, 43909732, 44747252, 45594016, 46450052, 47315392, 48190064, 49074096, + 49967516, 50870356, 51782636, 52704392, 53635648, 54576432, 55526772, 56486700, + 57456236, 58435408, 59424248, 60422780, 61431036, 62449032, 63476804, 64514376, + 65561776, 66619028, 67686160, 68763192, 69850160, 70947088, 72053992, 73170912, + 74297864, 75434880, 76581976, 77739184, 78906536, 80084040, 81271736, 82469648, + 83677792, 84896192, 86124888, 87363888, 88613232, 89872928, 91143016, 92423512, + 93714432, 95015816, 96327688, 97650056, 98982952, 100326408, 101680440, 103045072, + 104420320, 105806224, 107202800, 108610064, 110028048, 111456776, 112896264, 114346544, + 115807632, 117279552, 118762328, 120255976, 121760536, 123276016, 124802440, 126339832, + 127888216, 129447616, 131018048, 132599544, 134192112, 135795792, 137410592, 139036528, + 140673648, 142321952, 143981456, 145652208, 147334208, 149027488, 150732064, 152447968, + 154175200, 155913792, 157663776, 159425168, 161197984, 162982240, 164777968, 166585184, + 168403904, 170234160, 172075968, 173929344, 175794320, 177670896, 179559120, 181458992, + 183370528, 185293776, 187228736, 189175424, 191133888, 193104112, 195086128, 197079968, + 199085648, 201103184, 203132592, 205173888, 207227120, 209292272, 211369392, 213458480, + 215559568, 217672656, 219797792, 221934976, 224084240, 226245600, 228419056, 230604656, + 232802400, 235012320, 237234432, 239468736, 241715280, 243974080, 246245120, 248528464, + 250824112, 253132064, 255452368, 257785040, 260130080, 262487520, 264857376, 267239664, +}; + +static stbir_uint8 stbir__linear_to_srgb_uchar(float f) +{ + int x = (int) (f * (1 << 28)); // has headroom so you don't need to clamp + int v = 0; + int i; + + // Refine the guess with a short binary search. + i = v + 128; if (x >= stbir__srgb_offset_to_linear_scaled[i]) v = i; + i = v + 64; if (x >= stbir__srgb_offset_to_linear_scaled[i]) v = i; + i = v + 32; if (x >= stbir__srgb_offset_to_linear_scaled[i]) v = i; + i = v + 16; if (x >= stbir__srgb_offset_to_linear_scaled[i]) v = i; + i = v + 8; if (x >= stbir__srgb_offset_to_linear_scaled[i]) v = i; + i = v + 4; if (x >= stbir__srgb_offset_to_linear_scaled[i]) v = i; + i = v + 2; if (x >= stbir__srgb_offset_to_linear_scaled[i]) v = i; + i = v + 1; if (x >= stbir__srgb_offset_to_linear_scaled[i]) v = i; + + return (stbir_uint8) v; +} +#endif + +static float stbir__filter_trapezoid(float x, float scale) +{ + float halfscale = scale / 2; + float t = 0.5f + halfscale; + STBIR_ASSERT(scale <= 1); + + x = (float)fabs(x); + + if (x >= t) + return 0; + else + { + float r = 0.5f - halfscale; + if (x <= r) + return 1; + else + return (t - x) / scale; + } +} + +static float stbir__support_trapezoid(float scale) +{ + STBIR_ASSERT(scale <= 1); + return 0.5f + scale / 2; +} + +static float stbir__filter_triangle(float x, float s) +{ + STBIR__UNUSED_PARAM(s); + + x = (float)fabs(x); + + if (x <= 1.0f) + return 1 - x; + else + return 0; +} + +static float stbir__filter_cubic(float x, float s) +{ + STBIR__UNUSED_PARAM(s); + + x = (float)fabs(x); + + if (x < 1.0f) + return (4 + x*x*(3*x - 6))/6; + else if (x < 2.0f) + return (8 + x*(-12 + x*(6 - x)))/6; + + return (0.0f); +} + +static float stbir__filter_catmullrom(float x, float s) +{ + STBIR__UNUSED_PARAM(s); + + x = (float)fabs(x); + + if (x < 1.0f) + return 1 - x*x*(2.5f - 1.5f*x); + else if (x < 2.0f) + return 2 - x*(4 + x*(0.5f*x - 2.5f)); + + return (0.0f); +} + +static float stbir__filter_mitchell(float x, float s) +{ + STBIR__UNUSED_PARAM(s); + + x = (float)fabs(x); + + if (x < 1.0f) + return (16 + x*x*(21 * x - 36))/18; + else if (x < 2.0f) + return (32 + x*(-60 + x*(36 - 7*x)))/18; + + return (0.0f); +} + +static float stbir__support_zero(float s) +{ + STBIR__UNUSED_PARAM(s); + return 0; +} + +static float stbir__support_one(float s) +{ + STBIR__UNUSED_PARAM(s); + return 1; +} + +static float stbir__support_two(float s) +{ + STBIR__UNUSED_PARAM(s); + return 2; +} + +static stbir__filter_info stbir__filter_info_table[] = { + { NULL, stbir__support_zero }, + { stbir__filter_trapezoid, stbir__support_trapezoid }, + { stbir__filter_triangle, stbir__support_one }, + { stbir__filter_cubic, stbir__support_two }, + { stbir__filter_catmullrom, stbir__support_two }, + { stbir__filter_mitchell, stbir__support_two }, +}; + +stbir__inline static int stbir__use_upsampling(float ratio) +{ + return ratio > 1; +} + +stbir__inline static int stbir__use_width_upsampling(stbir__info* stbir_info) +{ + return stbir__use_upsampling(stbir_info->horizontal_scale); +} + +stbir__inline static int stbir__use_height_upsampling(stbir__info* stbir_info) +{ + return stbir__use_upsampling(stbir_info->vertical_scale); +} + +// This is the maximum number of input samples that can affect an output sample +// with the given filter +static int stbir__get_filter_pixel_width(stbir_filter filter, float scale) +{ + STBIR_ASSERT(filter != 0); + STBIR_ASSERT(filter < STBIR__ARRAY_SIZE(stbir__filter_info_table)); + + if (stbir__use_upsampling(scale)) + return (int)ceil(stbir__filter_info_table[filter].support(1/scale) * 2); + else + return (int)ceil(stbir__filter_info_table[filter].support(scale) * 2 / scale); +} + +// This is how much to expand buffers to account for filters seeking outside +// the image boundaries. +static int stbir__get_filter_pixel_margin(stbir_filter filter, float scale) +{ + return stbir__get_filter_pixel_width(filter, scale) / 2; +} + +static int stbir__get_coefficient_width(stbir_filter filter, float scale) +{ + if (stbir__use_upsampling(scale)) + return (int)ceil(stbir__filter_info_table[filter].support(1 / scale) * 2); + else + return (int)ceil(stbir__filter_info_table[filter].support(scale) * 2); +} + +static int stbir__get_contributors(float scale, stbir_filter filter, int input_size, int output_size) +{ + if (stbir__use_upsampling(scale)) + return output_size; + else + return (input_size + stbir__get_filter_pixel_margin(filter, scale) * 2); +} + +static int stbir__get_total_horizontal_coefficients(stbir__info* info) +{ + return info->horizontal_num_contributors + * stbir__get_coefficient_width (info->horizontal_filter, info->horizontal_scale); +} + +static int stbir__get_total_vertical_coefficients(stbir__info* info) +{ + return info->vertical_num_contributors + * stbir__get_coefficient_width (info->vertical_filter, info->vertical_scale); +} + +static stbir__contributors* stbir__get_contributor(stbir__contributors* contributors, int n) +{ + return &contributors[n]; +} + +// For perf reasons this code is duplicated in stbir__resample_horizontal_upsample/downsample, +// if you change it here change it there too. +static float* stbir__get_coefficient(float* coefficients, stbir_filter filter, float scale, int n, int c) +{ + int width = stbir__get_coefficient_width(filter, scale); + return &coefficients[width*n + c]; +} + +static int stbir__edge_wrap_slow(stbir_edge edge, int n, int max) +{ + switch (edge) + { + case STBIR_EDGE_ZERO: + return 0; // we'll decode the wrong pixel here, and then overwrite with 0s later + + case STBIR_EDGE_CLAMP: + if (n < 0) + return 0; + + if (n >= max) + return max - 1; + + return n; // NOTREACHED + + case STBIR_EDGE_REFLECT: + { + if (n < 0) + { + if (n < max) + return -n; + else + return max - 1; + } + + if (n >= max) + { + int max2 = max * 2; + if (n >= max2) + return 0; + else + return max2 - n - 1; + } + + return n; // NOTREACHED + } + + case STBIR_EDGE_WRAP: + if (n >= 0) + return (n % max); + else + { + int m = (-n) % max; + + if (m != 0) + m = max - m; + + return (m); + } + // NOTREACHED + + default: + STBIR_ASSERT(!"Unimplemented edge type"); + return 0; + } +} + +stbir__inline static int stbir__edge_wrap(stbir_edge edge, int n, int max) +{ + // avoid per-pixel switch + if (n >= 0 && n < max) + return n; + return stbir__edge_wrap_slow(edge, n, max); +} + +// What input pixels contribute to this output pixel? +static void stbir__calculate_sample_range_upsample(int n, float out_filter_radius, float scale_ratio, float out_shift, int* in_first_pixel, int* in_last_pixel, float* in_center_of_out) +{ + float out_pixel_center = (float)n + 0.5f; + float out_pixel_influence_lowerbound = out_pixel_center - out_filter_radius; + float out_pixel_influence_upperbound = out_pixel_center + out_filter_radius; + + float in_pixel_influence_lowerbound = (out_pixel_influence_lowerbound + out_shift) / scale_ratio; + float in_pixel_influence_upperbound = (out_pixel_influence_upperbound + out_shift) / scale_ratio; + + *in_center_of_out = (out_pixel_center + out_shift) / scale_ratio; + *in_first_pixel = (int)(floor(in_pixel_influence_lowerbound + 0.5)); + *in_last_pixel = (int)(floor(in_pixel_influence_upperbound - 0.5)); +} + +// What output pixels does this input pixel contribute to? +static void stbir__calculate_sample_range_downsample(int n, float in_pixels_radius, float scale_ratio, float out_shift, int* out_first_pixel, int* out_last_pixel, float* out_center_of_in) +{ + float in_pixel_center = (float)n + 0.5f; + float in_pixel_influence_lowerbound = in_pixel_center - in_pixels_radius; + float in_pixel_influence_upperbound = in_pixel_center + in_pixels_radius; + + float out_pixel_influence_lowerbound = in_pixel_influence_lowerbound * scale_ratio - out_shift; + float out_pixel_influence_upperbound = in_pixel_influence_upperbound * scale_ratio - out_shift; + + *out_center_of_in = in_pixel_center * scale_ratio - out_shift; + *out_first_pixel = (int)(floor(out_pixel_influence_lowerbound + 0.5)); + *out_last_pixel = (int)(floor(out_pixel_influence_upperbound - 0.5)); +} + +static void stbir__calculate_coefficients_upsample(stbir_filter filter, float scale, int in_first_pixel, int in_last_pixel, float in_center_of_out, stbir__contributors* contributor, float* coefficient_group) +{ + int i; + float total_filter = 0; + float filter_scale; + + STBIR_ASSERT(in_last_pixel - in_first_pixel <= (int)ceil(stbir__filter_info_table[filter].support(1/scale) * 2)); // Taken directly from stbir__get_coefficient_width() which we can't call because we don't know if we're horizontal or vertical. + + contributor->n0 = in_first_pixel; + contributor->n1 = in_last_pixel; + + STBIR_ASSERT(contributor->n1 >= contributor->n0); + + for (i = 0; i <= in_last_pixel - in_first_pixel; i++) + { + float in_pixel_center = (float)(i + in_first_pixel) + 0.5f; + coefficient_group[i] = stbir__filter_info_table[filter].kernel(in_center_of_out - in_pixel_center, 1 / scale); + + // If the coefficient is zero, skip it. (Don't do the <0 check here, we want the influence of those outside pixels.) + if (i == 0 && !coefficient_group[i]) + { + contributor->n0 = ++in_first_pixel; + i--; + continue; + } + + total_filter += coefficient_group[i]; + } + + // NOTE(fg): Not actually true in general, nor is there any reason to expect it should be. + // It would be true in exact math but is at best approximately true in floating-point math, + // and it would not make sense to try and put actual bounds on this here because it depends + // on the image aspect ratio which can get pretty extreme. + //STBIR_ASSERT(stbir__filter_info_table[filter].kernel((float)(in_last_pixel + 1) + 0.5f - in_center_of_out, 1/scale) == 0); + + STBIR_ASSERT(total_filter > 0.9); + STBIR_ASSERT(total_filter < 1.1f); // Make sure it's not way off. + + // Make sure the sum of all coefficients is 1. + filter_scale = 1 / total_filter; + + for (i = 0; i <= in_last_pixel - in_first_pixel; i++) + coefficient_group[i] *= filter_scale; + + for (i = in_last_pixel - in_first_pixel; i >= 0; i--) + { + if (coefficient_group[i]) + break; + + // This line has no weight. We can skip it. + contributor->n1 = contributor->n0 + i - 1; + } +} + +static void stbir__calculate_coefficients_downsample(stbir_filter filter, float scale_ratio, int out_first_pixel, int out_last_pixel, float out_center_of_in, stbir__contributors* contributor, float* coefficient_group) +{ + int i; + + STBIR_ASSERT(out_last_pixel - out_first_pixel <= (int)ceil(stbir__filter_info_table[filter].support(scale_ratio) * 2)); // Taken directly from stbir__get_coefficient_width() which we can't call because we don't know if we're horizontal or vertical. + + contributor->n0 = out_first_pixel; + contributor->n1 = out_last_pixel; + + STBIR_ASSERT(contributor->n1 >= contributor->n0); + + for (i = 0; i <= out_last_pixel - out_first_pixel; i++) + { + float out_pixel_center = (float)(i + out_first_pixel) + 0.5f; + float x = out_pixel_center - out_center_of_in; + coefficient_group[i] = stbir__filter_info_table[filter].kernel(x, scale_ratio) * scale_ratio; + } + + // NOTE(fg): Not actually true in general, nor is there any reason to expect it should be. + // It would be true in exact math but is at best approximately true in floating-point math, + // and it would not make sense to try and put actual bounds on this here because it depends + // on the image aspect ratio which can get pretty extreme. + //STBIR_ASSERT(stbir__filter_info_table[filter].kernel((float)(out_last_pixel + 1) + 0.5f - out_center_of_in, scale_ratio) == 0); + + for (i = out_last_pixel - out_first_pixel; i >= 0; i--) + { + if (coefficient_group[i]) + break; + + // This line has no weight. We can skip it. + contributor->n1 = contributor->n0 + i - 1; + } +} + +static void stbir__normalize_downsample_coefficients(stbir__contributors* contributors, float* coefficients, stbir_filter filter, float scale_ratio, int input_size, int output_size) +{ + int num_contributors = stbir__get_contributors(scale_ratio, filter, input_size, output_size); + int num_coefficients = stbir__get_coefficient_width(filter, scale_ratio); + int i, j; + int skip; + + for (i = 0; i < output_size; i++) + { + float scale; + float total = 0; + + for (j = 0; j < num_contributors; j++) + { + if (i >= contributors[j].n0 && i <= contributors[j].n1) + { + float coefficient = *stbir__get_coefficient(coefficients, filter, scale_ratio, j, i - contributors[j].n0); + total += coefficient; + } + else if (i < contributors[j].n0) + break; + } + + STBIR_ASSERT(total > 0.9f); + STBIR_ASSERT(total < 1.1f); + + scale = 1 / total; + + for (j = 0; j < num_contributors; j++) + { + if (i >= contributors[j].n0 && i <= contributors[j].n1) + *stbir__get_coefficient(coefficients, filter, scale_ratio, j, i - contributors[j].n0) *= scale; + else if (i < contributors[j].n0) + break; + } + } + + // Optimize: Skip zero coefficients and contributions outside of image bounds. + // Do this after normalizing because normalization depends on the n0/n1 values. + for (j = 0; j < num_contributors; j++) + { + int range, max, width; + + skip = 0; + while (*stbir__get_coefficient(coefficients, filter, scale_ratio, j, skip) == 0) + skip++; + + contributors[j].n0 += skip; + + while (contributors[j].n0 < 0) + { + contributors[j].n0++; + skip++; + } + + range = contributors[j].n1 - contributors[j].n0 + 1; + max = stbir__min(num_coefficients, range); + + width = stbir__get_coefficient_width(filter, scale_ratio); + for (i = 0; i < max; i++) + { + if (i + skip >= width) + break; + + *stbir__get_coefficient(coefficients, filter, scale_ratio, j, i) = *stbir__get_coefficient(coefficients, filter, scale_ratio, j, i + skip); + } + + continue; + } + + // Using min to avoid writing into invalid pixels. + for (i = 0; i < num_contributors; i++) + contributors[i].n1 = stbir__min(contributors[i].n1, output_size - 1); +} + +// Each scan line uses the same kernel values so we should calculate the kernel +// values once and then we can use them for every scan line. +static void stbir__calculate_filters(stbir__contributors* contributors, float* coefficients, stbir_filter filter, float scale_ratio, float shift, int input_size, int output_size) +{ + int n; + int total_contributors = stbir__get_contributors(scale_ratio, filter, input_size, output_size); + + if (stbir__use_upsampling(scale_ratio)) + { + float out_pixels_radius = stbir__filter_info_table[filter].support(1 / scale_ratio) * scale_ratio; + + // Looping through out pixels + for (n = 0; n < total_contributors; n++) + { + float in_center_of_out; // Center of the current out pixel in the in pixel space + int in_first_pixel, in_last_pixel; + + stbir__calculate_sample_range_upsample(n, out_pixels_radius, scale_ratio, shift, &in_first_pixel, &in_last_pixel, &in_center_of_out); + + stbir__calculate_coefficients_upsample(filter, scale_ratio, in_first_pixel, in_last_pixel, in_center_of_out, stbir__get_contributor(contributors, n), stbir__get_coefficient(coefficients, filter, scale_ratio, n, 0)); + } + } + else + { + float in_pixels_radius = stbir__filter_info_table[filter].support(scale_ratio) / scale_ratio; + + // Looping through in pixels + for (n = 0; n < total_contributors; n++) + { + float out_center_of_in; // Center of the current out pixel in the in pixel space + int out_first_pixel, out_last_pixel; + int n_adjusted = n - stbir__get_filter_pixel_margin(filter, scale_ratio); + + stbir__calculate_sample_range_downsample(n_adjusted, in_pixels_radius, scale_ratio, shift, &out_first_pixel, &out_last_pixel, &out_center_of_in); + + stbir__calculate_coefficients_downsample(filter, scale_ratio, out_first_pixel, out_last_pixel, out_center_of_in, stbir__get_contributor(contributors, n), stbir__get_coefficient(coefficients, filter, scale_ratio, n, 0)); + } + + stbir__normalize_downsample_coefficients(contributors, coefficients, filter, scale_ratio, input_size, output_size); + } +} + +static float* stbir__get_decode_buffer(stbir__info* stbir_info) +{ + // The 0 index of the decode buffer starts after the margin. This makes + // it okay to use negative indexes on the decode buffer. + return &stbir_info->decode_buffer[stbir_info->horizontal_filter_pixel_margin * stbir_info->channels]; +} + +#define STBIR__DECODE(type, colorspace) ((int)(type) * (STBIR_MAX_COLORSPACES) + (int)(colorspace)) + +static void stbir__decode_scanline(stbir__info* stbir_info, int n) +{ + int c; + int channels = stbir_info->channels; + int alpha_channel = stbir_info->alpha_channel; + int type = stbir_info->type; + int colorspace = stbir_info->colorspace; + int input_w = stbir_info->input_w; + size_t input_stride_bytes = stbir_info->input_stride_bytes; + float* decode_buffer = stbir__get_decode_buffer(stbir_info); + stbir_edge edge_horizontal = stbir_info->edge_horizontal; + stbir_edge edge_vertical = stbir_info->edge_vertical; + size_t in_buffer_row_offset = stbir__edge_wrap(edge_vertical, n, stbir_info->input_h) * input_stride_bytes; + const void* input_data = (char *) stbir_info->input_data + in_buffer_row_offset; + int max_x = input_w + stbir_info->horizontal_filter_pixel_margin; + int decode = STBIR__DECODE(type, colorspace); + + int x = -stbir_info->horizontal_filter_pixel_margin; + + // special handling for STBIR_EDGE_ZERO because it needs to return an item that doesn't appear in the input, + // and we want to avoid paying overhead on every pixel if not STBIR_EDGE_ZERO + if (edge_vertical == STBIR_EDGE_ZERO && (n < 0 || n >= stbir_info->input_h)) + { + for (; x < max_x; x++) + for (c = 0; c < channels; c++) + decode_buffer[x*channels + c] = 0; + return; + } + + switch (decode) + { + case STBIR__DECODE(STBIR_TYPE_UINT8, STBIR_COLORSPACE_LINEAR): + for (; x < max_x; x++) + { + int decode_pixel_index = x * channels; + int input_pixel_index = stbir__edge_wrap(edge_horizontal, x, input_w) * channels; + for (c = 0; c < channels; c++) + decode_buffer[decode_pixel_index + c] = ((float)((const unsigned char*)input_data)[input_pixel_index + c]) / stbir__max_uint8_as_float; + } + break; + + case STBIR__DECODE(STBIR_TYPE_UINT8, STBIR_COLORSPACE_SRGB): + for (; x < max_x; x++) + { + int decode_pixel_index = x * channels; + int input_pixel_index = stbir__edge_wrap(edge_horizontal, x, input_w) * channels; + for (c = 0; c < channels; c++) + decode_buffer[decode_pixel_index + c] = stbir__srgb_uchar_to_linear_float[((const unsigned char*)input_data)[input_pixel_index + c]]; + + if (!(stbir_info->flags&STBIR_FLAG_ALPHA_USES_COLORSPACE)) + decode_buffer[decode_pixel_index + alpha_channel] = ((float)((const unsigned char*)input_data)[input_pixel_index + alpha_channel]) / stbir__max_uint8_as_float; + } + break; + + case STBIR__DECODE(STBIR_TYPE_UINT16, STBIR_COLORSPACE_LINEAR): + for (; x < max_x; x++) + { + int decode_pixel_index = x * channels; + int input_pixel_index = stbir__edge_wrap(edge_horizontal, x, input_w) * channels; + for (c = 0; c < channels; c++) + decode_buffer[decode_pixel_index + c] = ((float)((const unsigned short*)input_data)[input_pixel_index + c]) / stbir__max_uint16_as_float; + } + break; + + case STBIR__DECODE(STBIR_TYPE_UINT16, STBIR_COLORSPACE_SRGB): + for (; x < max_x; x++) + { + int decode_pixel_index = x * channels; + int input_pixel_index = stbir__edge_wrap(edge_horizontal, x, input_w) * channels; + for (c = 0; c < channels; c++) + decode_buffer[decode_pixel_index + c] = stbir__srgb_to_linear(((float)((const unsigned short*)input_data)[input_pixel_index + c]) / stbir__max_uint16_as_float); + + if (!(stbir_info->flags&STBIR_FLAG_ALPHA_USES_COLORSPACE)) + decode_buffer[decode_pixel_index + alpha_channel] = ((float)((const unsigned short*)input_data)[input_pixel_index + alpha_channel]) / stbir__max_uint16_as_float; + } + break; + + case STBIR__DECODE(STBIR_TYPE_UINT32, STBIR_COLORSPACE_LINEAR): + for (; x < max_x; x++) + { + int decode_pixel_index = x * channels; + int input_pixel_index = stbir__edge_wrap(edge_horizontal, x, input_w) * channels; + for (c = 0; c < channels; c++) + decode_buffer[decode_pixel_index + c] = (float)(((double)((const unsigned int*)input_data)[input_pixel_index + c]) / stbir__max_uint32_as_float); + } + break; + + case STBIR__DECODE(STBIR_TYPE_UINT32, STBIR_COLORSPACE_SRGB): + for (; x < max_x; x++) + { + int decode_pixel_index = x * channels; + int input_pixel_index = stbir__edge_wrap(edge_horizontal, x, input_w) * channels; + for (c = 0; c < channels; c++) + decode_buffer[decode_pixel_index + c] = stbir__srgb_to_linear((float)(((double)((const unsigned int*)input_data)[input_pixel_index + c]) / stbir__max_uint32_as_float)); + + if (!(stbir_info->flags&STBIR_FLAG_ALPHA_USES_COLORSPACE)) + decode_buffer[decode_pixel_index + alpha_channel] = (float)(((double)((const unsigned int*)input_data)[input_pixel_index + alpha_channel]) / stbir__max_uint32_as_float); + } + break; + + case STBIR__DECODE(STBIR_TYPE_FLOAT, STBIR_COLORSPACE_LINEAR): + for (; x < max_x; x++) + { + int decode_pixel_index = x * channels; + int input_pixel_index = stbir__edge_wrap(edge_horizontal, x, input_w) * channels; + for (c = 0; c < channels; c++) + decode_buffer[decode_pixel_index + c] = ((const float*)input_data)[input_pixel_index + c]; + } + break; + + case STBIR__DECODE(STBIR_TYPE_FLOAT, STBIR_COLORSPACE_SRGB): + for (; x < max_x; x++) + { + int decode_pixel_index = x * channels; + int input_pixel_index = stbir__edge_wrap(edge_horizontal, x, input_w) * channels; + for (c = 0; c < channels; c++) + decode_buffer[decode_pixel_index + c] = stbir__srgb_to_linear(((const float*)input_data)[input_pixel_index + c]); + + if (!(stbir_info->flags&STBIR_FLAG_ALPHA_USES_COLORSPACE)) + decode_buffer[decode_pixel_index + alpha_channel] = ((const float*)input_data)[input_pixel_index + alpha_channel]; + } + + break; + + default: + STBIR_ASSERT(!"Unknown type/colorspace/channels combination."); + break; + } + + if (!(stbir_info->flags & STBIR_FLAG_ALPHA_PREMULTIPLIED)) + { + for (x = -stbir_info->horizontal_filter_pixel_margin; x < max_x; x++) + { + int decode_pixel_index = x * channels; + + // If the alpha value is 0 it will clobber the color values. Make sure it's not. + float alpha = decode_buffer[decode_pixel_index + alpha_channel]; +#ifndef STBIR_NO_ALPHA_EPSILON + if (stbir_info->type != STBIR_TYPE_FLOAT) { + alpha += STBIR_ALPHA_EPSILON; + decode_buffer[decode_pixel_index + alpha_channel] = alpha; + } +#endif + for (c = 0; c < channels; c++) + { + if (c == alpha_channel) + continue; + + decode_buffer[decode_pixel_index + c] *= alpha; + } + } + } + + if (edge_horizontal == STBIR_EDGE_ZERO) + { + for (x = -stbir_info->horizontal_filter_pixel_margin; x < 0; x++) + { + for (c = 0; c < channels; c++) + decode_buffer[x*channels + c] = 0; + } + for (x = input_w; x < max_x; x++) + { + for (c = 0; c < channels; c++) + decode_buffer[x*channels + c] = 0; + } + } +} + +static float* stbir__get_ring_buffer_entry(float* ring_buffer, int index, int ring_buffer_length) +{ + return &ring_buffer[index * ring_buffer_length]; +} + +static float* stbir__add_empty_ring_buffer_entry(stbir__info* stbir_info, int n) +{ + int ring_buffer_index; + float* ring_buffer; + + stbir_info->ring_buffer_last_scanline = n; + + if (stbir_info->ring_buffer_begin_index < 0) + { + ring_buffer_index = stbir_info->ring_buffer_begin_index = 0; + stbir_info->ring_buffer_first_scanline = n; + } + else + { + ring_buffer_index = (stbir_info->ring_buffer_begin_index + (stbir_info->ring_buffer_last_scanline - stbir_info->ring_buffer_first_scanline)) % stbir_info->ring_buffer_num_entries; + STBIR_ASSERT(ring_buffer_index != stbir_info->ring_buffer_begin_index); + } + + ring_buffer = stbir__get_ring_buffer_entry(stbir_info->ring_buffer, ring_buffer_index, stbir_info->ring_buffer_length_bytes / sizeof(float)); + memset(ring_buffer, 0, stbir_info->ring_buffer_length_bytes); + + return ring_buffer; +} + + +static void stbir__resample_horizontal_upsample(stbir__info* stbir_info, float* output_buffer) +{ + int x, k; + int output_w = stbir_info->output_w; + int channels = stbir_info->channels; + float* decode_buffer = stbir__get_decode_buffer(stbir_info); + stbir__contributors* horizontal_contributors = stbir_info->horizontal_contributors; + float* horizontal_coefficients = stbir_info->horizontal_coefficients; + int coefficient_width = stbir_info->horizontal_coefficient_width; + + for (x = 0; x < output_w; x++) + { + int n0 = horizontal_contributors[x].n0; + int n1 = horizontal_contributors[x].n1; + + int out_pixel_index = x * channels; + int coefficient_group = coefficient_width * x; + int coefficient_counter = 0; + + STBIR_ASSERT(n1 >= n0); + STBIR_ASSERT(n0 >= -stbir_info->horizontal_filter_pixel_margin); + STBIR_ASSERT(n1 >= -stbir_info->horizontal_filter_pixel_margin); + STBIR_ASSERT(n0 < stbir_info->input_w + stbir_info->horizontal_filter_pixel_margin); + STBIR_ASSERT(n1 < stbir_info->input_w + stbir_info->horizontal_filter_pixel_margin); + + switch (channels) { + case 1: + for (k = n0; k <= n1; k++) + { + int in_pixel_index = k * 1; + float coefficient = horizontal_coefficients[coefficient_group + coefficient_counter++]; + STBIR_ASSERT(coefficient != 0); + output_buffer[out_pixel_index + 0] += decode_buffer[in_pixel_index + 0] * coefficient; + } + break; + case 2: + for (k = n0; k <= n1; k++) + { + int in_pixel_index = k * 2; + float coefficient = horizontal_coefficients[coefficient_group + coefficient_counter++]; + STBIR_ASSERT(coefficient != 0); + output_buffer[out_pixel_index + 0] += decode_buffer[in_pixel_index + 0] * coefficient; + output_buffer[out_pixel_index + 1] += decode_buffer[in_pixel_index + 1] * coefficient; + } + break; + case 3: + for (k = n0; k <= n1; k++) + { + int in_pixel_index = k * 3; + float coefficient = horizontal_coefficients[coefficient_group + coefficient_counter++]; + STBIR_ASSERT(coefficient != 0); + output_buffer[out_pixel_index + 0] += decode_buffer[in_pixel_index + 0] * coefficient; + output_buffer[out_pixel_index + 1] += decode_buffer[in_pixel_index + 1] * coefficient; + output_buffer[out_pixel_index + 2] += decode_buffer[in_pixel_index + 2] * coefficient; + } + break; + case 4: + for (k = n0; k <= n1; k++) + { + int in_pixel_index = k * 4; + float coefficient = horizontal_coefficients[coefficient_group + coefficient_counter++]; + STBIR_ASSERT(coefficient != 0); + output_buffer[out_pixel_index + 0] += decode_buffer[in_pixel_index + 0] * coefficient; + output_buffer[out_pixel_index + 1] += decode_buffer[in_pixel_index + 1] * coefficient; + output_buffer[out_pixel_index + 2] += decode_buffer[in_pixel_index + 2] * coefficient; + output_buffer[out_pixel_index + 3] += decode_buffer[in_pixel_index + 3] * coefficient; + } + break; + default: + for (k = n0; k <= n1; k++) + { + int in_pixel_index = k * channels; + float coefficient = horizontal_coefficients[coefficient_group + coefficient_counter++]; + int c; + STBIR_ASSERT(coefficient != 0); + for (c = 0; c < channels; c++) + output_buffer[out_pixel_index + c] += decode_buffer[in_pixel_index + c] * coefficient; + } + break; + } + } +} + +static void stbir__resample_horizontal_downsample(stbir__info* stbir_info, float* output_buffer) +{ + int x, k; + int input_w = stbir_info->input_w; + int channels = stbir_info->channels; + float* decode_buffer = stbir__get_decode_buffer(stbir_info); + stbir__contributors* horizontal_contributors = stbir_info->horizontal_contributors; + float* horizontal_coefficients = stbir_info->horizontal_coefficients; + int coefficient_width = stbir_info->horizontal_coefficient_width; + int filter_pixel_margin = stbir_info->horizontal_filter_pixel_margin; + int max_x = input_w + filter_pixel_margin * 2; + + STBIR_ASSERT(!stbir__use_width_upsampling(stbir_info)); + + switch (channels) { + case 1: + for (x = 0; x < max_x; x++) + { + int n0 = horizontal_contributors[x].n0; + int n1 = horizontal_contributors[x].n1; + + int in_x = x - filter_pixel_margin; + int in_pixel_index = in_x * 1; + int max_n = n1; + int coefficient_group = coefficient_width * x; + + for (k = n0; k <= max_n; k++) + { + int out_pixel_index = k * 1; + float coefficient = horizontal_coefficients[coefficient_group + k - n0]; + output_buffer[out_pixel_index + 0] += decode_buffer[in_pixel_index + 0] * coefficient; + } + } + break; + + case 2: + for (x = 0; x < max_x; x++) + { + int n0 = horizontal_contributors[x].n0; + int n1 = horizontal_contributors[x].n1; + + int in_x = x - filter_pixel_margin; + int in_pixel_index = in_x * 2; + int max_n = n1; + int coefficient_group = coefficient_width * x; + + for (k = n0; k <= max_n; k++) + { + int out_pixel_index = k * 2; + float coefficient = horizontal_coefficients[coefficient_group + k - n0]; + output_buffer[out_pixel_index + 0] += decode_buffer[in_pixel_index + 0] * coefficient; + output_buffer[out_pixel_index + 1] += decode_buffer[in_pixel_index + 1] * coefficient; + } + } + break; + + case 3: + for (x = 0; x < max_x; x++) + { + int n0 = horizontal_contributors[x].n0; + int n1 = horizontal_contributors[x].n1; + + int in_x = x - filter_pixel_margin; + int in_pixel_index = in_x * 3; + int max_n = n1; + int coefficient_group = coefficient_width * x; + + for (k = n0; k <= max_n; k++) + { + int out_pixel_index = k * 3; + float coefficient = horizontal_coefficients[coefficient_group + k - n0]; + output_buffer[out_pixel_index + 0] += decode_buffer[in_pixel_index + 0] * coefficient; + output_buffer[out_pixel_index + 1] += decode_buffer[in_pixel_index + 1] * coefficient; + output_buffer[out_pixel_index + 2] += decode_buffer[in_pixel_index + 2] * coefficient; + } + } + break; + + case 4: + for (x = 0; x < max_x; x++) + { + int n0 = horizontal_contributors[x].n0; + int n1 = horizontal_contributors[x].n1; + + int in_x = x - filter_pixel_margin; + int in_pixel_index = in_x * 4; + int max_n = n1; + int coefficient_group = coefficient_width * x; + + for (k = n0; k <= max_n; k++) + { + int out_pixel_index = k * 4; + float coefficient = horizontal_coefficients[coefficient_group + k - n0]; + output_buffer[out_pixel_index + 0] += decode_buffer[in_pixel_index + 0] * coefficient; + output_buffer[out_pixel_index + 1] += decode_buffer[in_pixel_index + 1] * coefficient; + output_buffer[out_pixel_index + 2] += decode_buffer[in_pixel_index + 2] * coefficient; + output_buffer[out_pixel_index + 3] += decode_buffer[in_pixel_index + 3] * coefficient; + } + } + break; + + default: + for (x = 0; x < max_x; x++) + { + int n0 = horizontal_contributors[x].n0; + int n1 = horizontal_contributors[x].n1; + + int in_x = x - filter_pixel_margin; + int in_pixel_index = in_x * channels; + int max_n = n1; + int coefficient_group = coefficient_width * x; + + for (k = n0; k <= max_n; k++) + { + int c; + int out_pixel_index = k * channels; + float coefficient = horizontal_coefficients[coefficient_group + k - n0]; + for (c = 0; c < channels; c++) + output_buffer[out_pixel_index + c] += decode_buffer[in_pixel_index + c] * coefficient; + } + } + break; + } +} + +static void stbir__decode_and_resample_upsample(stbir__info* stbir_info, int n) +{ + // Decode the nth scanline from the source image into the decode buffer. + stbir__decode_scanline(stbir_info, n); + + // Now resample it into the ring buffer. + if (stbir__use_width_upsampling(stbir_info)) + stbir__resample_horizontal_upsample(stbir_info, stbir__add_empty_ring_buffer_entry(stbir_info, n)); + else + stbir__resample_horizontal_downsample(stbir_info, stbir__add_empty_ring_buffer_entry(stbir_info, n)); + + // Now it's sitting in the ring buffer ready to be used as source for the vertical sampling. +} + +static void stbir__decode_and_resample_downsample(stbir__info* stbir_info, int n) +{ + // Decode the nth scanline from the source image into the decode buffer. + stbir__decode_scanline(stbir_info, n); + + memset(stbir_info->horizontal_buffer, 0, stbir_info->output_w * stbir_info->channels * sizeof(float)); + + // Now resample it into the horizontal buffer. + if (stbir__use_width_upsampling(stbir_info)) + stbir__resample_horizontal_upsample(stbir_info, stbir_info->horizontal_buffer); + else + stbir__resample_horizontal_downsample(stbir_info, stbir_info->horizontal_buffer); + + // Now it's sitting in the horizontal buffer ready to be distributed into the ring buffers. +} + +// Get the specified scan line from the ring buffer. +static float* stbir__get_ring_buffer_scanline(int get_scanline, float* ring_buffer, int begin_index, int first_scanline, int ring_buffer_num_entries, int ring_buffer_length) +{ + int ring_buffer_index = (begin_index + (get_scanline - first_scanline)) % ring_buffer_num_entries; + return stbir__get_ring_buffer_entry(ring_buffer, ring_buffer_index, ring_buffer_length); +} + + +static void stbir__encode_scanline(stbir__info* stbir_info, int num_pixels, void *output_buffer, float *encode_buffer, int channels, int alpha_channel, int decode) +{ + int x; + int n; + int num_nonalpha; + stbir_uint16 nonalpha[STBIR_MAX_CHANNELS]; + + if (!(stbir_info->flags&STBIR_FLAG_ALPHA_PREMULTIPLIED)) + { + for (x=0; x < num_pixels; ++x) + { + int pixel_index = x*channels; + + float alpha = encode_buffer[pixel_index + alpha_channel]; + float reciprocal_alpha = alpha ? 1.0f / alpha : 0; + + // unrolling this produced a 1% slowdown upscaling a large RGBA linear-space image on my machine - stb + for (n = 0; n < channels; n++) + if (n != alpha_channel) + encode_buffer[pixel_index + n] *= reciprocal_alpha; + + // We added in a small epsilon to prevent the color channel from being deleted with zero alpha. + // Because we only add it for integer types, it will automatically be discarded on integer + // conversion, so we don't need to subtract it back out (which would be problematic for + // numeric precision reasons). + } + } + + // build a table of all channels that need colorspace correction, so + // we don't perform colorspace correction on channels that don't need it. + for (x = 0, num_nonalpha = 0; x < channels; ++x) + { + if (x != alpha_channel || (stbir_info->flags & STBIR_FLAG_ALPHA_USES_COLORSPACE)) + { + nonalpha[num_nonalpha++] = (stbir_uint16)x; + } + } + + #define STBIR__ROUND_INT(f) ((int) ((f)+0.5)) + #define STBIR__ROUND_UINT(f) ((stbir_uint32) ((f)+0.5)) + + #ifdef STBIR__SATURATE_INT + #define STBIR__ENCODE_LINEAR8(f) stbir__saturate8 (STBIR__ROUND_INT((f) * stbir__max_uint8_as_float )) + #define STBIR__ENCODE_LINEAR16(f) stbir__saturate16(STBIR__ROUND_INT((f) * stbir__max_uint16_as_float)) + #else + #define STBIR__ENCODE_LINEAR8(f) (unsigned char ) STBIR__ROUND_INT(stbir__saturate(f) * stbir__max_uint8_as_float ) + #define STBIR__ENCODE_LINEAR16(f) (unsigned short) STBIR__ROUND_INT(stbir__saturate(f) * stbir__max_uint16_as_float) + #endif + + switch (decode) + { + case STBIR__DECODE(STBIR_TYPE_UINT8, STBIR_COLORSPACE_LINEAR): + for (x=0; x < num_pixels; ++x) + { + int pixel_index = x*channels; + + for (n = 0; n < channels; n++) + { + int index = pixel_index + n; + ((unsigned char*)output_buffer)[index] = STBIR__ENCODE_LINEAR8(encode_buffer[index]); + } + } + break; + + case STBIR__DECODE(STBIR_TYPE_UINT8, STBIR_COLORSPACE_SRGB): + for (x=0; x < num_pixels; ++x) + { + int pixel_index = x*channels; + + for (n = 0; n < num_nonalpha; n++) + { + int index = pixel_index + nonalpha[n]; + ((unsigned char*)output_buffer)[index] = stbir__linear_to_srgb_uchar(encode_buffer[index]); + } + + if (!(stbir_info->flags & STBIR_FLAG_ALPHA_USES_COLORSPACE)) + ((unsigned char *)output_buffer)[pixel_index + alpha_channel] = STBIR__ENCODE_LINEAR8(encode_buffer[pixel_index+alpha_channel]); + } + break; + + case STBIR__DECODE(STBIR_TYPE_UINT16, STBIR_COLORSPACE_LINEAR): + for (x=0; x < num_pixels; ++x) + { + int pixel_index = x*channels; + + for (n = 0; n < channels; n++) + { + int index = pixel_index + n; + ((unsigned short*)output_buffer)[index] = STBIR__ENCODE_LINEAR16(encode_buffer[index]); + } + } + break; + + case STBIR__DECODE(STBIR_TYPE_UINT16, STBIR_COLORSPACE_SRGB): + for (x=0; x < num_pixels; ++x) + { + int pixel_index = x*channels; + + for (n = 0; n < num_nonalpha; n++) + { + int index = pixel_index + nonalpha[n]; + ((unsigned short*)output_buffer)[index] = (unsigned short)STBIR__ROUND_INT(stbir__linear_to_srgb(stbir__saturate(encode_buffer[index])) * stbir__max_uint16_as_float); + } + + if (!(stbir_info->flags&STBIR_FLAG_ALPHA_USES_COLORSPACE)) + ((unsigned short*)output_buffer)[pixel_index + alpha_channel] = STBIR__ENCODE_LINEAR16(encode_buffer[pixel_index + alpha_channel]); + } + + break; + + case STBIR__DECODE(STBIR_TYPE_UINT32, STBIR_COLORSPACE_LINEAR): + for (x=0; x < num_pixels; ++x) + { + int pixel_index = x*channels; + + for (n = 0; n < channels; n++) + { + int index = pixel_index + n; + ((unsigned int*)output_buffer)[index] = (unsigned int)STBIR__ROUND_UINT(((double)stbir__saturate(encode_buffer[index])) * stbir__max_uint32_as_float); + } + } + break; + + case STBIR__DECODE(STBIR_TYPE_UINT32, STBIR_COLORSPACE_SRGB): + for (x=0; x < num_pixels; ++x) + { + int pixel_index = x*channels; + + for (n = 0; n < num_nonalpha; n++) + { + int index = pixel_index + nonalpha[n]; + ((unsigned int*)output_buffer)[index] = (unsigned int)STBIR__ROUND_UINT(((double)stbir__linear_to_srgb(stbir__saturate(encode_buffer[index]))) * stbir__max_uint32_as_float); + } + + if (!(stbir_info->flags&STBIR_FLAG_ALPHA_USES_COLORSPACE)) + ((unsigned int*)output_buffer)[pixel_index + alpha_channel] = (unsigned int)STBIR__ROUND_INT(((double)stbir__saturate(encode_buffer[pixel_index + alpha_channel])) * stbir__max_uint32_as_float); + } + break; + + case STBIR__DECODE(STBIR_TYPE_FLOAT, STBIR_COLORSPACE_LINEAR): + for (x=0; x < num_pixels; ++x) + { + int pixel_index = x*channels; + + for (n = 0; n < channels; n++) + { + int index = pixel_index + n; + ((float*)output_buffer)[index] = encode_buffer[index]; + } + } + break; + + case STBIR__DECODE(STBIR_TYPE_FLOAT, STBIR_COLORSPACE_SRGB): + for (x=0; x < num_pixels; ++x) + { + int pixel_index = x*channels; + + for (n = 0; n < num_nonalpha; n++) + { + int index = pixel_index + nonalpha[n]; + ((float*)output_buffer)[index] = stbir__linear_to_srgb(encode_buffer[index]); + } + + if (!(stbir_info->flags&STBIR_FLAG_ALPHA_USES_COLORSPACE)) + ((float*)output_buffer)[pixel_index + alpha_channel] = encode_buffer[pixel_index + alpha_channel]; + } + break; + + default: + STBIR_ASSERT(!"Unknown type/colorspace/channels combination."); + break; + } +} + +static void stbir__resample_vertical_upsample(stbir__info* stbir_info, int n) +{ + int x, k; + int output_w = stbir_info->output_w; + stbir__contributors* vertical_contributors = stbir_info->vertical_contributors; + float* vertical_coefficients = stbir_info->vertical_coefficients; + int channels = stbir_info->channels; + int alpha_channel = stbir_info->alpha_channel; + int type = stbir_info->type; + int colorspace = stbir_info->colorspace; + int ring_buffer_entries = stbir_info->ring_buffer_num_entries; + void* output_data = stbir_info->output_data; + float* encode_buffer = stbir_info->encode_buffer; + int decode = STBIR__DECODE(type, colorspace); + int coefficient_width = stbir_info->vertical_coefficient_width; + int coefficient_counter; + int contributor = n; + + float* ring_buffer = stbir_info->ring_buffer; + int ring_buffer_begin_index = stbir_info->ring_buffer_begin_index; + int ring_buffer_first_scanline = stbir_info->ring_buffer_first_scanline; + int ring_buffer_length = stbir_info->ring_buffer_length_bytes/sizeof(float); + + int n0,n1, output_row_start; + int coefficient_group = coefficient_width * contributor; + + n0 = vertical_contributors[contributor].n0; + n1 = vertical_contributors[contributor].n1; + + output_row_start = n * stbir_info->output_stride_bytes; + + STBIR_ASSERT(stbir__use_height_upsampling(stbir_info)); + + memset(encode_buffer, 0, output_w * sizeof(float) * channels); + + // I tried reblocking this for better cache usage of encode_buffer + // (using x_outer, k, x_inner), but it lost speed. -- stb + + coefficient_counter = 0; + switch (channels) { + case 1: + for (k = n0; k <= n1; k++) + { + int coefficient_index = coefficient_counter++; + float* ring_buffer_entry = stbir__get_ring_buffer_scanline(k, ring_buffer, ring_buffer_begin_index, ring_buffer_first_scanline, ring_buffer_entries, ring_buffer_length); + float coefficient = vertical_coefficients[coefficient_group + coefficient_index]; + for (x = 0; x < output_w; ++x) + { + int in_pixel_index = x * 1; + encode_buffer[in_pixel_index + 0] += ring_buffer_entry[in_pixel_index + 0] * coefficient; + } + } + break; + case 2: + for (k = n0; k <= n1; k++) + { + int coefficient_index = coefficient_counter++; + float* ring_buffer_entry = stbir__get_ring_buffer_scanline(k, ring_buffer, ring_buffer_begin_index, ring_buffer_first_scanline, ring_buffer_entries, ring_buffer_length); + float coefficient = vertical_coefficients[coefficient_group + coefficient_index]; + for (x = 0; x < output_w; ++x) + { + int in_pixel_index = x * 2; + encode_buffer[in_pixel_index + 0] += ring_buffer_entry[in_pixel_index + 0] * coefficient; + encode_buffer[in_pixel_index + 1] += ring_buffer_entry[in_pixel_index + 1] * coefficient; + } + } + break; + case 3: + for (k = n0; k <= n1; k++) + { + int coefficient_index = coefficient_counter++; + float* ring_buffer_entry = stbir__get_ring_buffer_scanline(k, ring_buffer, ring_buffer_begin_index, ring_buffer_first_scanline, ring_buffer_entries, ring_buffer_length); + float coefficient = vertical_coefficients[coefficient_group + coefficient_index]; + for (x = 0; x < output_w; ++x) + { + int in_pixel_index = x * 3; + encode_buffer[in_pixel_index + 0] += ring_buffer_entry[in_pixel_index + 0] * coefficient; + encode_buffer[in_pixel_index + 1] += ring_buffer_entry[in_pixel_index + 1] * coefficient; + encode_buffer[in_pixel_index + 2] += ring_buffer_entry[in_pixel_index + 2] * coefficient; + } + } + break; + case 4: + for (k = n0; k <= n1; k++) + { + int coefficient_index = coefficient_counter++; + float* ring_buffer_entry = stbir__get_ring_buffer_scanline(k, ring_buffer, ring_buffer_begin_index, ring_buffer_first_scanline, ring_buffer_entries, ring_buffer_length); + float coefficient = vertical_coefficients[coefficient_group + coefficient_index]; + for (x = 0; x < output_w; ++x) + { + int in_pixel_index = x * 4; + encode_buffer[in_pixel_index + 0] += ring_buffer_entry[in_pixel_index + 0] * coefficient; + encode_buffer[in_pixel_index + 1] += ring_buffer_entry[in_pixel_index + 1] * coefficient; + encode_buffer[in_pixel_index + 2] += ring_buffer_entry[in_pixel_index + 2] * coefficient; + encode_buffer[in_pixel_index + 3] += ring_buffer_entry[in_pixel_index + 3] * coefficient; + } + } + break; + default: + for (k = n0; k <= n1; k++) + { + int coefficient_index = coefficient_counter++; + float* ring_buffer_entry = stbir__get_ring_buffer_scanline(k, ring_buffer, ring_buffer_begin_index, ring_buffer_first_scanline, ring_buffer_entries, ring_buffer_length); + float coefficient = vertical_coefficients[coefficient_group + coefficient_index]; + for (x = 0; x < output_w; ++x) + { + int in_pixel_index = x * channels; + int c; + for (c = 0; c < channels; c++) + encode_buffer[in_pixel_index + c] += ring_buffer_entry[in_pixel_index + c] * coefficient; + } + } + break; + } + stbir__encode_scanline(stbir_info, output_w, (char *) output_data + output_row_start, encode_buffer, channels, alpha_channel, decode); +} + +static void stbir__resample_vertical_downsample(stbir__info* stbir_info, int n) +{ + int x, k; + int output_w = stbir_info->output_w; + stbir__contributors* vertical_contributors = stbir_info->vertical_contributors; + float* vertical_coefficients = stbir_info->vertical_coefficients; + int channels = stbir_info->channels; + int ring_buffer_entries = stbir_info->ring_buffer_num_entries; + float* horizontal_buffer = stbir_info->horizontal_buffer; + int coefficient_width = stbir_info->vertical_coefficient_width; + int contributor = n + stbir_info->vertical_filter_pixel_margin; + + float* ring_buffer = stbir_info->ring_buffer; + int ring_buffer_begin_index = stbir_info->ring_buffer_begin_index; + int ring_buffer_first_scanline = stbir_info->ring_buffer_first_scanline; + int ring_buffer_length = stbir_info->ring_buffer_length_bytes/sizeof(float); + int n0,n1; + + n0 = vertical_contributors[contributor].n0; + n1 = vertical_contributors[contributor].n1; + + STBIR_ASSERT(!stbir__use_height_upsampling(stbir_info)); + + for (k = n0; k <= n1; k++) + { + int coefficient_index = k - n0; + int coefficient_group = coefficient_width * contributor; + float coefficient = vertical_coefficients[coefficient_group + coefficient_index]; + + float* ring_buffer_entry = stbir__get_ring_buffer_scanline(k, ring_buffer, ring_buffer_begin_index, ring_buffer_first_scanline, ring_buffer_entries, ring_buffer_length); + + switch (channels) { + case 1: + for (x = 0; x < output_w; x++) + { + int in_pixel_index = x * 1; + ring_buffer_entry[in_pixel_index + 0] += horizontal_buffer[in_pixel_index + 0] * coefficient; + } + break; + case 2: + for (x = 0; x < output_w; x++) + { + int in_pixel_index = x * 2; + ring_buffer_entry[in_pixel_index + 0] += horizontal_buffer[in_pixel_index + 0] * coefficient; + ring_buffer_entry[in_pixel_index + 1] += horizontal_buffer[in_pixel_index + 1] * coefficient; + } + break; + case 3: + for (x = 0; x < output_w; x++) + { + int in_pixel_index = x * 3; + ring_buffer_entry[in_pixel_index + 0] += horizontal_buffer[in_pixel_index + 0] * coefficient; + ring_buffer_entry[in_pixel_index + 1] += horizontal_buffer[in_pixel_index + 1] * coefficient; + ring_buffer_entry[in_pixel_index + 2] += horizontal_buffer[in_pixel_index + 2] * coefficient; + } + break; + case 4: + for (x = 0; x < output_w; x++) + { + int in_pixel_index = x * 4; + ring_buffer_entry[in_pixel_index + 0] += horizontal_buffer[in_pixel_index + 0] * coefficient; + ring_buffer_entry[in_pixel_index + 1] += horizontal_buffer[in_pixel_index + 1] * coefficient; + ring_buffer_entry[in_pixel_index + 2] += horizontal_buffer[in_pixel_index + 2] * coefficient; + ring_buffer_entry[in_pixel_index + 3] += horizontal_buffer[in_pixel_index + 3] * coefficient; + } + break; + default: + for (x = 0; x < output_w; x++) + { + int in_pixel_index = x * channels; + + int c; + for (c = 0; c < channels; c++) + ring_buffer_entry[in_pixel_index + c] += horizontal_buffer[in_pixel_index + c] * coefficient; + } + break; + } + } +} + +static void stbir__buffer_loop_upsample(stbir__info* stbir_info) +{ + int y; + float scale_ratio = stbir_info->vertical_scale; + float out_scanlines_radius = stbir__filter_info_table[stbir_info->vertical_filter].support(1/scale_ratio) * scale_ratio; + + STBIR_ASSERT(stbir__use_height_upsampling(stbir_info)); + + for (y = 0; y < stbir_info->output_h; y++) + { + float in_center_of_out = 0; // Center of the current out scanline in the in scanline space + int in_first_scanline = 0, in_last_scanline = 0; + + stbir__calculate_sample_range_upsample(y, out_scanlines_radius, scale_ratio, stbir_info->vertical_shift, &in_first_scanline, &in_last_scanline, &in_center_of_out); + + STBIR_ASSERT(in_last_scanline - in_first_scanline + 1 <= stbir_info->ring_buffer_num_entries); + + if (stbir_info->ring_buffer_begin_index >= 0) + { + // Get rid of whatever we don't need anymore. + while (in_first_scanline > stbir_info->ring_buffer_first_scanline) + { + if (stbir_info->ring_buffer_first_scanline == stbir_info->ring_buffer_last_scanline) + { + // We just popped the last scanline off the ring buffer. + // Reset it to the empty state. + stbir_info->ring_buffer_begin_index = -1; + stbir_info->ring_buffer_first_scanline = 0; + stbir_info->ring_buffer_last_scanline = 0; + break; + } + else + { + stbir_info->ring_buffer_first_scanline++; + stbir_info->ring_buffer_begin_index = (stbir_info->ring_buffer_begin_index + 1) % stbir_info->ring_buffer_num_entries; + } + } + } + + // Load in new ones. + if (stbir_info->ring_buffer_begin_index < 0) + stbir__decode_and_resample_upsample(stbir_info, in_first_scanline); + + while (in_last_scanline > stbir_info->ring_buffer_last_scanline) + stbir__decode_and_resample_upsample(stbir_info, stbir_info->ring_buffer_last_scanline + 1); + + // Now all buffers should be ready to write a row of vertical sampling. + stbir__resample_vertical_upsample(stbir_info, y); + + STBIR_PROGRESS_REPORT((float)y / stbir_info->output_h); + } +} + +static void stbir__empty_ring_buffer(stbir__info* stbir_info, int first_necessary_scanline) +{ + int output_stride_bytes = stbir_info->output_stride_bytes; + int channels = stbir_info->channels; + int alpha_channel = stbir_info->alpha_channel; + int type = stbir_info->type; + int colorspace = stbir_info->colorspace; + int output_w = stbir_info->output_w; + void* output_data = stbir_info->output_data; + int decode = STBIR__DECODE(type, colorspace); + + float* ring_buffer = stbir_info->ring_buffer; + int ring_buffer_length = stbir_info->ring_buffer_length_bytes/sizeof(float); + + if (stbir_info->ring_buffer_begin_index >= 0) + { + // Get rid of whatever we don't need anymore. + while (first_necessary_scanline > stbir_info->ring_buffer_first_scanline) + { + if (stbir_info->ring_buffer_first_scanline >= 0 && stbir_info->ring_buffer_first_scanline < stbir_info->output_h) + { + int output_row_start = stbir_info->ring_buffer_first_scanline * output_stride_bytes; + float* ring_buffer_entry = stbir__get_ring_buffer_entry(ring_buffer, stbir_info->ring_buffer_begin_index, ring_buffer_length); + stbir__encode_scanline(stbir_info, output_w, (char *) output_data + output_row_start, ring_buffer_entry, channels, alpha_channel, decode); + STBIR_PROGRESS_REPORT((float)stbir_info->ring_buffer_first_scanline / stbir_info->output_h); + } + + if (stbir_info->ring_buffer_first_scanline == stbir_info->ring_buffer_last_scanline) + { + // We just popped the last scanline off the ring buffer. + // Reset it to the empty state. + stbir_info->ring_buffer_begin_index = -1; + stbir_info->ring_buffer_first_scanline = 0; + stbir_info->ring_buffer_last_scanline = 0; + break; + } + else + { + stbir_info->ring_buffer_first_scanline++; + stbir_info->ring_buffer_begin_index = (stbir_info->ring_buffer_begin_index + 1) % stbir_info->ring_buffer_num_entries; + } + } + } +} + +static void stbir__buffer_loop_downsample(stbir__info* stbir_info) +{ + int y; + float scale_ratio = stbir_info->vertical_scale; + int output_h = stbir_info->output_h; + float in_pixels_radius = stbir__filter_info_table[stbir_info->vertical_filter].support(scale_ratio) / scale_ratio; + int pixel_margin = stbir_info->vertical_filter_pixel_margin; + int max_y = stbir_info->input_h + pixel_margin; + + STBIR_ASSERT(!stbir__use_height_upsampling(stbir_info)); + + for (y = -pixel_margin; y < max_y; y++) + { + float out_center_of_in; // Center of the current out scanline in the in scanline space + int out_first_scanline, out_last_scanline; + + stbir__calculate_sample_range_downsample(y, in_pixels_radius, scale_ratio, stbir_info->vertical_shift, &out_first_scanline, &out_last_scanline, &out_center_of_in); + + STBIR_ASSERT(out_last_scanline - out_first_scanline + 1 <= stbir_info->ring_buffer_num_entries); + + if (out_last_scanline < 0 || out_first_scanline >= output_h) + continue; + + stbir__empty_ring_buffer(stbir_info, out_first_scanline); + + stbir__decode_and_resample_downsample(stbir_info, y); + + // Load in new ones. + if (stbir_info->ring_buffer_begin_index < 0) + stbir__add_empty_ring_buffer_entry(stbir_info, out_first_scanline); + + while (out_last_scanline > stbir_info->ring_buffer_last_scanline) + stbir__add_empty_ring_buffer_entry(stbir_info, stbir_info->ring_buffer_last_scanline + 1); + + // Now the horizontal buffer is ready to write to all ring buffer rows. + stbir__resample_vertical_downsample(stbir_info, y); + } + + stbir__empty_ring_buffer(stbir_info, stbir_info->output_h); +} + +static void stbir__setup(stbir__info *info, int input_w, int input_h, int output_w, int output_h, int channels) +{ + info->input_w = input_w; + info->input_h = input_h; + info->output_w = output_w; + info->output_h = output_h; + info->channels = channels; +} + +static void stbir__calculate_transform(stbir__info *info, float s0, float t0, float s1, float t1, float *transform) +{ + info->s0 = s0; + info->t0 = t0; + info->s1 = s1; + info->t1 = t1; + + if (transform) + { + info->horizontal_scale = transform[0]; + info->vertical_scale = transform[1]; + info->horizontal_shift = transform[2]; + info->vertical_shift = transform[3]; + } + else + { + info->horizontal_scale = ((float)info->output_w / info->input_w) / (s1 - s0); + info->vertical_scale = ((float)info->output_h / info->input_h) / (t1 - t0); + + info->horizontal_shift = s0 * info->output_w / (s1 - s0); + info->vertical_shift = t0 * info->output_h / (t1 - t0); + } +} + +static void stbir__choose_filter(stbir__info *info, stbir_filter h_filter, stbir_filter v_filter) +{ + if (h_filter == 0) + h_filter = stbir__use_upsampling(info->horizontal_scale) ? STBIR_DEFAULT_FILTER_UPSAMPLE : STBIR_DEFAULT_FILTER_DOWNSAMPLE; + if (v_filter == 0) + v_filter = stbir__use_upsampling(info->vertical_scale) ? STBIR_DEFAULT_FILTER_UPSAMPLE : STBIR_DEFAULT_FILTER_DOWNSAMPLE; + info->horizontal_filter = h_filter; + info->vertical_filter = v_filter; +} + +static stbir_uint32 stbir__calculate_memory(stbir__info *info) +{ + int pixel_margin = stbir__get_filter_pixel_margin(info->horizontal_filter, info->horizontal_scale); + int filter_height = stbir__get_filter_pixel_width(info->vertical_filter, info->vertical_scale); + + info->horizontal_num_contributors = stbir__get_contributors(info->horizontal_scale, info->horizontal_filter, info->input_w, info->output_w); + info->vertical_num_contributors = stbir__get_contributors(info->vertical_scale , info->vertical_filter , info->input_h, info->output_h); + + // One extra entry because floating point precision problems sometimes cause an extra to be necessary. + info->ring_buffer_num_entries = filter_height + 1; + + info->horizontal_contributors_size = info->horizontal_num_contributors * sizeof(stbir__contributors); + info->horizontal_coefficients_size = stbir__get_total_horizontal_coefficients(info) * sizeof(float); + info->vertical_contributors_size = info->vertical_num_contributors * sizeof(stbir__contributors); + info->vertical_coefficients_size = stbir__get_total_vertical_coefficients(info) * sizeof(float); + info->decode_buffer_size = (info->input_w + pixel_margin * 2) * info->channels * sizeof(float); + info->horizontal_buffer_size = info->output_w * info->channels * sizeof(float); + info->ring_buffer_size = info->output_w * info->channels * info->ring_buffer_num_entries * sizeof(float); + info->encode_buffer_size = info->output_w * info->channels * sizeof(float); + + STBIR_ASSERT(info->horizontal_filter != 0); + STBIR_ASSERT(info->horizontal_filter < STBIR__ARRAY_SIZE(stbir__filter_info_table)); // this now happens too late + STBIR_ASSERT(info->vertical_filter != 0); + STBIR_ASSERT(info->vertical_filter < STBIR__ARRAY_SIZE(stbir__filter_info_table)); // this now happens too late + + if (stbir__use_height_upsampling(info)) + // The horizontal buffer is for when we're downsampling the height and we + // can't output the result of sampling the decode buffer directly into the + // ring buffers. + info->horizontal_buffer_size = 0; + else + // The encode buffer is to retain precision in the height upsampling method + // and isn't used when height downsampling. + info->encode_buffer_size = 0; + + return info->horizontal_contributors_size + info->horizontal_coefficients_size + + info->vertical_contributors_size + info->vertical_coefficients_size + + info->decode_buffer_size + info->horizontal_buffer_size + + info->ring_buffer_size + info->encode_buffer_size; +} + +static int stbir__resize_allocated(stbir__info *info, + const void* input_data, int input_stride_in_bytes, + void* output_data, int output_stride_in_bytes, + int alpha_channel, stbir_uint32 flags, stbir_datatype type, + stbir_edge edge_horizontal, stbir_edge edge_vertical, stbir_colorspace colorspace, + void* tempmem, size_t tempmem_size_in_bytes) +{ + size_t memory_required = stbir__calculate_memory(info); + + int width_stride_input = input_stride_in_bytes ? input_stride_in_bytes : info->channels * info->input_w * stbir__type_size[type]; + int width_stride_output = output_stride_in_bytes ? output_stride_in_bytes : info->channels * info->output_w * stbir__type_size[type]; + +#ifdef STBIR_DEBUG_OVERWRITE_TEST +#define OVERWRITE_ARRAY_SIZE 8 + unsigned char overwrite_output_before_pre[OVERWRITE_ARRAY_SIZE]; + unsigned char overwrite_tempmem_before_pre[OVERWRITE_ARRAY_SIZE]; + unsigned char overwrite_output_after_pre[OVERWRITE_ARRAY_SIZE]; + unsigned char overwrite_tempmem_after_pre[OVERWRITE_ARRAY_SIZE]; + + size_t begin_forbidden = width_stride_output * (info->output_h - 1) + info->output_w * info->channels * stbir__type_size[type]; + memcpy(overwrite_output_before_pre, &((unsigned char*)output_data)[-OVERWRITE_ARRAY_SIZE], OVERWRITE_ARRAY_SIZE); + memcpy(overwrite_output_after_pre, &((unsigned char*)output_data)[begin_forbidden], OVERWRITE_ARRAY_SIZE); + memcpy(overwrite_tempmem_before_pre, &((unsigned char*)tempmem)[-OVERWRITE_ARRAY_SIZE], OVERWRITE_ARRAY_SIZE); + memcpy(overwrite_tempmem_after_pre, &((unsigned char*)tempmem)[tempmem_size_in_bytes], OVERWRITE_ARRAY_SIZE); +#endif + + STBIR_ASSERT(info->channels >= 0); + STBIR_ASSERT(info->channels <= STBIR_MAX_CHANNELS); + + if (info->channels < 0 || info->channels > STBIR_MAX_CHANNELS) + return 0; + + STBIR_ASSERT(info->horizontal_filter < STBIR__ARRAY_SIZE(stbir__filter_info_table)); + STBIR_ASSERT(info->vertical_filter < STBIR__ARRAY_SIZE(stbir__filter_info_table)); + + if (info->horizontal_filter >= STBIR__ARRAY_SIZE(stbir__filter_info_table)) + return 0; + if (info->vertical_filter >= STBIR__ARRAY_SIZE(stbir__filter_info_table)) + return 0; + + if (alpha_channel < 0) + flags |= STBIR_FLAG_ALPHA_USES_COLORSPACE | STBIR_FLAG_ALPHA_PREMULTIPLIED; + + if (!(flags&STBIR_FLAG_ALPHA_USES_COLORSPACE) || !(flags&STBIR_FLAG_ALPHA_PREMULTIPLIED)) { + STBIR_ASSERT(alpha_channel >= 0 && alpha_channel < info->channels); + } + + if (alpha_channel >= info->channels) + return 0; + + STBIR_ASSERT(tempmem); + + if (!tempmem) + return 0; + + STBIR_ASSERT(tempmem_size_in_bytes >= memory_required); + + if (tempmem_size_in_bytes < memory_required) + return 0; + + memset(tempmem, 0, tempmem_size_in_bytes); + + info->input_data = input_data; + info->input_stride_bytes = width_stride_input; + + info->output_data = output_data; + info->output_stride_bytes = width_stride_output; + + info->alpha_channel = alpha_channel; + info->flags = flags; + info->type = type; + info->edge_horizontal = edge_horizontal; + info->edge_vertical = edge_vertical; + info->colorspace = colorspace; + + info->horizontal_coefficient_width = stbir__get_coefficient_width (info->horizontal_filter, info->horizontal_scale); + info->vertical_coefficient_width = stbir__get_coefficient_width (info->vertical_filter , info->vertical_scale ); + info->horizontal_filter_pixel_width = stbir__get_filter_pixel_width (info->horizontal_filter, info->horizontal_scale); + info->vertical_filter_pixel_width = stbir__get_filter_pixel_width (info->vertical_filter , info->vertical_scale ); + info->horizontal_filter_pixel_margin = stbir__get_filter_pixel_margin(info->horizontal_filter, info->horizontal_scale); + info->vertical_filter_pixel_margin = stbir__get_filter_pixel_margin(info->vertical_filter , info->vertical_scale ); + + info->ring_buffer_length_bytes = info->output_w * info->channels * sizeof(float); + info->decode_buffer_pixels = info->input_w + info->horizontal_filter_pixel_margin * 2; + +#define STBIR__NEXT_MEMPTR(current, newtype) (newtype*)(((unsigned char*)current) + current##_size) + + info->horizontal_contributors = (stbir__contributors *) tempmem; + info->horizontal_coefficients = STBIR__NEXT_MEMPTR(info->horizontal_contributors, float); + info->vertical_contributors = STBIR__NEXT_MEMPTR(info->horizontal_coefficients, stbir__contributors); + info->vertical_coefficients = STBIR__NEXT_MEMPTR(info->vertical_contributors, float); + info->decode_buffer = STBIR__NEXT_MEMPTR(info->vertical_coefficients, float); + + if (stbir__use_height_upsampling(info)) + { + info->horizontal_buffer = NULL; + info->ring_buffer = STBIR__NEXT_MEMPTR(info->decode_buffer, float); + info->encode_buffer = STBIR__NEXT_MEMPTR(info->ring_buffer, float); + + STBIR_ASSERT((size_t)STBIR__NEXT_MEMPTR(info->encode_buffer, unsigned char) == (size_t)tempmem + tempmem_size_in_bytes); + } + else + { + info->horizontal_buffer = STBIR__NEXT_MEMPTR(info->decode_buffer, float); + info->ring_buffer = STBIR__NEXT_MEMPTR(info->horizontal_buffer, float); + info->encode_buffer = NULL; + + STBIR_ASSERT((size_t)STBIR__NEXT_MEMPTR(info->ring_buffer, unsigned char) == (size_t)tempmem + tempmem_size_in_bytes); + } + +#undef STBIR__NEXT_MEMPTR + + // This signals that the ring buffer is empty + info->ring_buffer_begin_index = -1; + + stbir__calculate_filters(info->horizontal_contributors, info->horizontal_coefficients, info->horizontal_filter, info->horizontal_scale, info->horizontal_shift, info->input_w, info->output_w); + stbir__calculate_filters(info->vertical_contributors, info->vertical_coefficients, info->vertical_filter, info->vertical_scale, info->vertical_shift, info->input_h, info->output_h); + + STBIR_PROGRESS_REPORT(0); + + if (stbir__use_height_upsampling(info)) + stbir__buffer_loop_upsample(info); + else + stbir__buffer_loop_downsample(info); + + STBIR_PROGRESS_REPORT(1); + +#ifdef STBIR_DEBUG_OVERWRITE_TEST + STBIR_ASSERT(memcmp(overwrite_output_before_pre, &((unsigned char*)output_data)[-OVERWRITE_ARRAY_SIZE], OVERWRITE_ARRAY_SIZE) == 0); + STBIR_ASSERT(memcmp(overwrite_output_after_pre, &((unsigned char*)output_data)[begin_forbidden], OVERWRITE_ARRAY_SIZE) == 0); + STBIR_ASSERT(memcmp(overwrite_tempmem_before_pre, &((unsigned char*)tempmem)[-OVERWRITE_ARRAY_SIZE], OVERWRITE_ARRAY_SIZE) == 0); + STBIR_ASSERT(memcmp(overwrite_tempmem_after_pre, &((unsigned char*)tempmem)[tempmem_size_in_bytes], OVERWRITE_ARRAY_SIZE) == 0); +#endif + + return 1; +} + + +static int stbir__resize_arbitrary( + void *alloc_context, + const void* input_data, int input_w, int input_h, int input_stride_in_bytes, + void* output_data, int output_w, int output_h, int output_stride_in_bytes, + float s0, float t0, float s1, float t1, float *transform, + int channels, int alpha_channel, stbir_uint32 flags, stbir_datatype type, + stbir_filter h_filter, stbir_filter v_filter, + stbir_edge edge_horizontal, stbir_edge edge_vertical, stbir_colorspace colorspace) +{ + stbir__info info; + int result; + size_t memory_required; + void* extra_memory; + + stbir__setup(&info, input_w, input_h, output_w, output_h, channels); + stbir__calculate_transform(&info, s0,t0,s1,t1,transform); + stbir__choose_filter(&info, h_filter, v_filter); + memory_required = stbir__calculate_memory(&info); + extra_memory = STBIR_MALLOC(memory_required, alloc_context); + + if (!extra_memory) + return 0; + + result = stbir__resize_allocated(&info, input_data, input_stride_in_bytes, + output_data, output_stride_in_bytes, + alpha_channel, flags, type, + edge_horizontal, edge_vertical, + colorspace, extra_memory, memory_required); + + STBIR_FREE(extra_memory, alloc_context); + + return result; +} + +STBIRDEF int stbir_resize_uint8( const unsigned char *input_pixels , int input_w , int input_h , int input_stride_in_bytes, + unsigned char *output_pixels, int output_w, int output_h, int output_stride_in_bytes, + int num_channels) +{ + return stbir__resize_arbitrary(NULL, input_pixels, input_w, input_h, input_stride_in_bytes, + output_pixels, output_w, output_h, output_stride_in_bytes, + 0,0,1,1,NULL,num_channels,-1,0, STBIR_TYPE_UINT8, STBIR_FILTER_DEFAULT, STBIR_FILTER_DEFAULT, + STBIR_EDGE_CLAMP, STBIR_EDGE_CLAMP, STBIR_COLORSPACE_LINEAR); +} + +STBIRDEF int stbir_resize_float( const float *input_pixels , int input_w , int input_h , int input_stride_in_bytes, + float *output_pixels, int output_w, int output_h, int output_stride_in_bytes, + int num_channels) +{ + return stbir__resize_arbitrary(NULL, input_pixels, input_w, input_h, input_stride_in_bytes, + output_pixels, output_w, output_h, output_stride_in_bytes, + 0,0,1,1,NULL,num_channels,-1,0, STBIR_TYPE_FLOAT, STBIR_FILTER_DEFAULT, STBIR_FILTER_DEFAULT, + STBIR_EDGE_CLAMP, STBIR_EDGE_CLAMP, STBIR_COLORSPACE_LINEAR); +} + +STBIRDEF int stbir_resize_uint8_srgb(const unsigned char *input_pixels , int input_w , int input_h , int input_stride_in_bytes, + unsigned char *output_pixels, int output_w, int output_h, int output_stride_in_bytes, + int num_channels, int alpha_channel, int flags) +{ + return stbir__resize_arbitrary(NULL, input_pixels, input_w, input_h, input_stride_in_bytes, + output_pixels, output_w, output_h, output_stride_in_bytes, + 0,0,1,1,NULL,num_channels,alpha_channel,flags, STBIR_TYPE_UINT8, STBIR_FILTER_DEFAULT, STBIR_FILTER_DEFAULT, + STBIR_EDGE_CLAMP, STBIR_EDGE_CLAMP, STBIR_COLORSPACE_SRGB); +} + +STBIRDEF int stbir_resize_uint8_srgb_edgemode(const unsigned char *input_pixels , int input_w , int input_h , int input_stride_in_bytes, + unsigned char *output_pixels, int output_w, int output_h, int output_stride_in_bytes, + int num_channels, int alpha_channel, int flags, + stbir_edge edge_wrap_mode) +{ + return stbir__resize_arbitrary(NULL, input_pixels, input_w, input_h, input_stride_in_bytes, + output_pixels, output_w, output_h, output_stride_in_bytes, + 0,0,1,1,NULL,num_channels,alpha_channel,flags, STBIR_TYPE_UINT8, STBIR_FILTER_DEFAULT, STBIR_FILTER_DEFAULT, + edge_wrap_mode, edge_wrap_mode, STBIR_COLORSPACE_SRGB); +} + +STBIRDEF int stbir_resize_uint8_generic( const unsigned char *input_pixels , int input_w , int input_h , int input_stride_in_bytes, + unsigned char *output_pixels, int output_w, int output_h, int output_stride_in_bytes, + int num_channels, int alpha_channel, int flags, + stbir_edge edge_wrap_mode, stbir_filter filter, stbir_colorspace space, + void *alloc_context) +{ + return stbir__resize_arbitrary(alloc_context, input_pixels, input_w, input_h, input_stride_in_bytes, + output_pixels, output_w, output_h, output_stride_in_bytes, + 0,0,1,1,NULL,num_channels,alpha_channel,flags, STBIR_TYPE_UINT8, filter, filter, + edge_wrap_mode, edge_wrap_mode, space); +} + +STBIRDEF int stbir_resize_uint16_generic(const stbir_uint16 *input_pixels , int input_w , int input_h , int input_stride_in_bytes, + stbir_uint16 *output_pixels , int output_w, int output_h, int output_stride_in_bytes, + int num_channels, int alpha_channel, int flags, + stbir_edge edge_wrap_mode, stbir_filter filter, stbir_colorspace space, + void *alloc_context) +{ + return stbir__resize_arbitrary(alloc_context, input_pixels, input_w, input_h, input_stride_in_bytes, + output_pixels, output_w, output_h, output_stride_in_bytes, + 0,0,1,1,NULL,num_channels,alpha_channel,flags, STBIR_TYPE_UINT16, filter, filter, + edge_wrap_mode, edge_wrap_mode, space); +} + + +STBIRDEF int stbir_resize_float_generic( const float *input_pixels , int input_w , int input_h , int input_stride_in_bytes, + float *output_pixels , int output_w, int output_h, int output_stride_in_bytes, + int num_channels, int alpha_channel, int flags, + stbir_edge edge_wrap_mode, stbir_filter filter, stbir_colorspace space, + void *alloc_context) +{ + return stbir__resize_arbitrary(alloc_context, input_pixels, input_w, input_h, input_stride_in_bytes, + output_pixels, output_w, output_h, output_stride_in_bytes, + 0,0,1,1,NULL,num_channels,alpha_channel,flags, STBIR_TYPE_FLOAT, filter, filter, + edge_wrap_mode, edge_wrap_mode, space); +} + + +STBIRDEF int stbir_resize( const void *input_pixels , int input_w , int input_h , int input_stride_in_bytes, + void *output_pixels, int output_w, int output_h, int output_stride_in_bytes, + stbir_datatype datatype, + int num_channels, int alpha_channel, int flags, + stbir_edge edge_mode_horizontal, stbir_edge edge_mode_vertical, + stbir_filter filter_horizontal, stbir_filter filter_vertical, + stbir_colorspace space, void *alloc_context) +{ + return stbir__resize_arbitrary(alloc_context, input_pixels, input_w, input_h, input_stride_in_bytes, + output_pixels, output_w, output_h, output_stride_in_bytes, + 0,0,1,1,NULL,num_channels,alpha_channel,flags, datatype, filter_horizontal, filter_vertical, + edge_mode_horizontal, edge_mode_vertical, space); +} + + +STBIRDEF int stbir_resize_subpixel(const void *input_pixels , int input_w , int input_h , int input_stride_in_bytes, + void *output_pixels, int output_w, int output_h, int output_stride_in_bytes, + stbir_datatype datatype, + int num_channels, int alpha_channel, int flags, + stbir_edge edge_mode_horizontal, stbir_edge edge_mode_vertical, + stbir_filter filter_horizontal, stbir_filter filter_vertical, + stbir_colorspace space, void *alloc_context, + float x_scale, float y_scale, + float x_offset, float y_offset) +{ + float transform[4]; + transform[0] = x_scale; + transform[1] = y_scale; + transform[2] = x_offset; + transform[3] = y_offset; + return stbir__resize_arbitrary(alloc_context, input_pixels, input_w, input_h, input_stride_in_bytes, + output_pixels, output_w, output_h, output_stride_in_bytes, + 0,0,1,1,transform,num_channels,alpha_channel,flags, datatype, filter_horizontal, filter_vertical, + edge_mode_horizontal, edge_mode_vertical, space); +} + +STBIRDEF int stbir_resize_region( const void *input_pixels , int input_w , int input_h , int input_stride_in_bytes, + void *output_pixels, int output_w, int output_h, int output_stride_in_bytes, + stbir_datatype datatype, + int num_channels, int alpha_channel, int flags, + stbir_edge edge_mode_horizontal, stbir_edge edge_mode_vertical, + stbir_filter filter_horizontal, stbir_filter filter_vertical, + stbir_colorspace space, void *alloc_context, + float s0, float t0, float s1, float t1) +{ + return stbir__resize_arbitrary(alloc_context, input_pixels, input_w, input_h, input_stride_in_bytes, + output_pixels, output_w, output_h, output_stride_in_bytes, + s0,t0,s1,t1,NULL,num_channels,alpha_channel,flags, datatype, filter_horizontal, filter_vertical, + edge_mode_horizontal, edge_mode_vertical, space); +} + +#endif // STB_IMAGE_RESIZE_IMPLEMENTATION + +/* +------------------------------------------------------------------------------ +This software is available under 2 licenses -- choose whichever you prefer. +------------------------------------------------------------------------------ +ALTERNATIVE A - MIT License +Copyright (c) 2017 Sean Barrett +Permission is hereby granted, free of charge, to any person obtaining a copy of +this software and associated documentation files (the "Software"), to deal in +the Software without restriction, including without limitation the rights to +use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies +of the Software, and to permit persons to whom the Software is furnished to do +so, subject to the following conditions: +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. +------------------------------------------------------------------------------ +ALTERNATIVE B - Public Domain (www.unlicense.org) +This is free and unencumbered software released into the public domain. +Anyone is free to copy, modify, publish, use, compile, sell, or distribute this +software, either in source code form or as a compiled binary, for any purpose, +commercial or non-commercial, and by any means. +In jurisdictions that recognize copyright laws, the author or authors of this +software dedicate any and all copyright interest in the software to the public +domain. We make this dedication for the benefit of the public at large and to +the detriment of our heirs and successors. We intend this dedication to be an +overt act of relinquishment in perpetuity of all present and future rights to +this software under copyright law. +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN +ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION +WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. +------------------------------------------------------------------------------ +*/ From 14c79d878e4c44fae1a97ef58490580ffea13fe6 Mon Sep 17 00:00:00 2001 From: Jiseong Oh Date: Tue, 13 Oct 2026 19:39:29 +0900 Subject: [PATCH 13/17] Support enn_llm_stop in the LLM library Let callers stop an in-flight generation. Signed-off-by: Jiseong Oh Co-authored-by: brucekk.kim --- examples/samsung/executor_runner/ennllm.cpp | 4 ++++ examples/samsung/executor_runner/ennllm.h | 3 +++ examples/samsung/executor_runner/llm/runner/runner.cpp | 4 +--- 3 files changed, 8 insertions(+), 3 deletions(-) diff --git a/examples/samsung/executor_runner/ennllm.cpp b/examples/samsung/executor_runner/ennllm.cpp index 426fd2be14e..0601da0e1d4 100644 --- a/examples/samsung/executor_runner/ennllm.cpp +++ b/examples/samsung/executor_runner/ennllm.cpp @@ -75,6 +75,10 @@ void enn_llm_destroy(enn_runner* runner) { delete runner; } +void enn_llm_stop(enn_runner* runner) { + runner->impl->stop(); +} + int32_t enn_llm_generate( enn_runner* runner, const char* prompt, diff --git a/examples/samsung/executor_runner/ennllm.h b/examples/samsung/executor_runner/ennllm.h index d62b3671b05..34227595826 100644 --- a/examples/samsung/executor_runner/ennllm.h +++ b/examples/samsung/executor_runner/ennllm.h @@ -70,6 +70,9 @@ ENN_API enn_runner* enn_llm_create( /** Destroy the runner and release all resources. Safe to call with NULL. */ ENN_API void enn_llm_destroy(enn_runner* runner); +/** Request the runner to stop current generation. */ +ENN_API void enn_llm_stop(enn_runner* runner); + /** * Generate a response for prompt (optionally with an image). The models stay * prepared across calls; KV-cache is reset before each generation so every diff --git a/examples/samsung/executor_runner/llm/runner/runner.cpp b/examples/samsung/executor_runner/llm/runner/runner.cpp index 3014bea7d64..2650db17769 100644 --- a/examples/samsung/executor_runner/llm/runner/runner.cpp +++ b/examples/samsung/executor_runner/llm/runner/runner.cpp @@ -503,10 +503,8 @@ runtime::Error LLMRunner::warmup( } void LLMRunner::stop() { - if (is_loaded()) { + if (text_token_generator_ != nullptr) { text_token_generator_->stop(); - } else { - ET_LOG(Error, "Token generator is not loaded, cannot stop"); } } From 0af447b1e1408081a57706e9b7a8f2783d4e78e1 Mon Sep 17 00:00:00 2001 From: Jiseong Oh Date: Tue, 13 Oct 2026 19:39:30 +0900 Subject: [PATCH 14/17] Add exynos_llm_wrapper and a sample CLI over its C API Add a sample CLI that consumes the public C API, and rename the ennllm library to exynos_llm_wrapper. Signed-off-by: Jiseong Oh Co-authored-by: MJ Kim --- examples/samsung/CMakeLists.txt | 31 +++-- examples/samsung/executor_runner/README.md | 9 +- .../{ennllm.cpp => exynos_llm_wrapper.cpp} | 2 +- .../{ennllm.h => exynos_llm_wrapper.h} | 25 ++-- .../exynos_llm_wrapper_sample_main.cpp | 116 ++++++++++++++++++ 5 files changed, 163 insertions(+), 20 deletions(-) rename examples/samsung/executor_runner/{ennllm.cpp => exynos_llm_wrapper.cpp} (97%) rename examples/samsung/executor_runner/{ennllm.h => exynos_llm_wrapper.h} (84%) create mode 100644 examples/samsung/executor_runner/exynos_llm_wrapper_sample_main.cpp diff --git a/examples/samsung/CMakeLists.txt b/examples/samsung/CMakeLists.txt index 1138327e2e6..9892f2a245b 100644 --- a/examples/samsung/CMakeLists.txt +++ b/examples/samsung/CMakeLists.txt @@ -211,18 +211,20 @@ target_link_libraries( -Wl,--no-whole-archive ) -# Public shared library exposing the C API in ennllm.h. Bundles all +# Public shared library exposing the C API in exynos_llm_wrapper.h. Bundles all # dependencies so only the .so + header need to be shipped to third parties. -add_library(ennllm SHARED ${CMAKE_CURRENT_LIST_DIR}/executor_runner/ennllm.cpp) +add_library( + exynos_llm_wrapper SHARED + ${CMAKE_CURRENT_LIST_DIR}/executor_runner/exynos_llm_wrapper.cpp +) -target_compile_definitions(ennllm PRIVATE EXECUTORCH_VERSION=2) +target_compile_definitions(exynos_llm_wrapper PRIVATE EXECUTORCH_VERSION=2) target_include_directories( - ennllm - PUBLIC ${CMAKE_CURRENT_LIST_DIR}/executor_runner + exynos_llm_wrapper PUBLIC ${CMAKE_CURRENT_LIST_DIR}/executor_runner ) target_link_libraries( - ennllm + exynos_llm_wrapper PUBLIC enn_llm_runner PRIVATE log gflags @@ -243,5 +245,20 @@ target_link_libraries( ) set_target_properties( - ennllm PROPERTIES CXX_VISIBILITY_PRESET hidden VISIBILITY_INLINES_HIDDEN ON + exynos_llm_wrapper PROPERTIES CXX_VISIBILITY_PRESET hidden + VISIBILITY_INLINES_HIDDEN ON +) + +# Sample CLI consuming the public exynos_llm_wrapper C API (same shape as an app +# using libexynos_llm_wrapper.so). Links the shared lib so it coexists with the +# all-in-one enn_executor_llm_runner. +add_executable( + exynos_llm_wrapper_sample + ${CMAKE_CURRENT_LIST_DIR}/executor_runner/exynos_llm_wrapper_sample_main.cpp +) + +target_link_libraries(exynos_llm_wrapper_sample PRIVATE exynos_llm_wrapper) + +set_target_properties( + exynos_llm_wrapper_sample PROPERTIES CXX_VISIBILITY_PRESET hidden ) diff --git a/examples/samsung/executor_runner/README.md b/examples/samsung/executor_runner/README.md index 8c41ae394ac..ce84b88aaad 100644 --- a/examples/samsung/executor_runner/README.md +++ b/examples/samsung/executor_runner/README.md @@ -24,4 +24,11 @@ For example: - Gemma 3 1B: `examples/models/gemma3_1b/exynos_gemma3_1b.json` - Llama 3 1B: `examples/models/llama3_1b/exynos_llama3_1b.json` -Ensure these files are pushed to the target device (e.g., to `/data/local/tmp/`) before running the script. \ No newline at end of file +Ensure these files are pushed to the target device (e.g., to `/data/local/tmp/`) before running the script. + +## `exynos_llm_wrapper` library (app-porting reference) + +`libexynos_llm_wrapper.so` is a reference for porting the Exynos LLM runner into +an Android app: a small public C API (`exynos_llm_wrapper.h`) that hides the +runner internals behind an opaque handle. `exynos_llm_wrapper_sample` is a CLI +that exercises the same API on-device. \ No newline at end of file diff --git a/examples/samsung/executor_runner/ennllm.cpp b/examples/samsung/executor_runner/exynos_llm_wrapper.cpp similarity index 97% rename from examples/samsung/executor_runner/ennllm.cpp rename to examples/samsung/executor_runner/exynos_llm_wrapper.cpp index 0601da0e1d4..3548e52546b 100644 --- a/examples/samsung/executor_runner/ennllm.cpp +++ b/examples/samsung/executor_runner/exynos_llm_wrapper.cpp @@ -6,7 +6,7 @@ * across generate() calls so consecutive inferences do not re-prepare. */ -#include +#include #include diff --git a/examples/samsung/executor_runner/ennllm.h b/examples/samsung/executor_runner/exynos_llm_wrapper.h similarity index 84% rename from examples/samsung/executor_runner/ennllm.h rename to examples/samsung/executor_runner/exynos_llm_wrapper.h index 34227595826..229477f4bfb 100644 --- a/examples/samsung/executor_runner/ennllm.h +++ b/examples/samsung/executor_runner/exynos_llm_wrapper.h @@ -2,9 +2,12 @@ * Copyright (c) 2026 Samsung Electronics Co. LTD * All rights reserved. * - * Public C API for the Exynos LLM runner. + * Public C API wrapping the Exynos LLM runner. * - * Only these symbols are exported from libennllm.so. The underlying + * Provided as a reference for porting the Exynos LLM runner into an Android + * app. + * + * Only these symbols are exported from libexynos_llm_wrapper.so. The underlying * implementation (model loading, ENN backend, tokenizer, KV-cache) is kept * internal to the library. * @@ -13,8 +16,8 @@ * reloading or re-preparing the models. */ -#ifndef ENNLLM_H_ -#define ENNLLM_H_ +#ifndef EXYNOS_LLM_WRAPPER_H_ +#define EXYNOS_LLM_WRAPPER_H_ #include @@ -23,9 +26,9 @@ extern "C" { #endif #if defined(__GNUC__) || defined(__clang__) -#define ENN_API __attribute__((visibility("default"))) +#define EXYNOS_LLM_WRAPPER_API __attribute__((visibility("default"))) #else -#define ENN_API +#define EXYNOS_LLM_WRAPPER_API #endif /** Values returned by the API functions. 0 = OK, non-zero = error. */ @@ -59,7 +62,7 @@ typedef struct enn_runner enn_runner; * @param vision_model_path Vision encoder .pte, or NULL => text-only runner. * @return A handle on success, NULL on failure. */ -ENN_API enn_runner* enn_llm_create( +EXYNOS_LLM_WRAPPER_API enn_runner* enn_llm_create( const char* model_path, const char* data_path, const char* tokenizer_path, @@ -68,10 +71,10 @@ ENN_API enn_runner* enn_llm_create( const char* vision_model_path); /** Destroy the runner and release all resources. Safe to call with NULL. */ -ENN_API void enn_llm_destroy(enn_runner* runner); +EXYNOS_LLM_WRAPPER_API void enn_llm_destroy(enn_runner* runner); /** Request the runner to stop current generation. */ -ENN_API void enn_llm_stop(enn_runner* runner); +EXYNOS_LLM_WRAPPER_API void enn_llm_stop(enn_runner* runner); /** * Generate a response for prompt (optionally with an image). The models stay @@ -93,7 +96,7 @@ ENN_API void enn_llm_stop(enn_runner* runner); * @param userdata Opaque pointer passed through to on_token. * @return ENN_OK on success, an ENN_ERR_* code otherwise. */ -ENN_API int32_t enn_llm_generate( +EXYNOS_LLM_WRAPPER_API int32_t enn_llm_generate( enn_runner* runner, const char* prompt, const char* image_path, @@ -110,4 +113,4 @@ ENN_API int32_t enn_llm_generate( } #endif -#endif // ENNLLM_H_ +#endif // EXYNOS_LLM_WRAPPER_H_ diff --git a/examples/samsung/executor_runner/exynos_llm_wrapper_sample_main.cpp b/examples/samsung/executor_runner/exynos_llm_wrapper_sample_main.cpp new file mode 100644 index 00000000000..31bebd3bbbd --- /dev/null +++ b/examples/samsung/executor_runner/exynos_llm_wrapper_sample_main.cpp @@ -0,0 +1,116 @@ +/* + * Copyright (c) 2026 Samsung Electronics Co. LTD + * All rights reserved. + * + * Sample CLI consuming the public exynos_llm_wrapper C API + * (exynos_llm_wrapper.h). This mirrors how an Android app drives + * libexynos_llm_wrapper.so: create once, generate repeatedly, destroy at + * shutdown. Models stay prepared across generate() calls. + * + * Usage (paths are on-device): + * exynos_llm_wrapper_sample [vision] + * + * Reads prompts from stdin (one per line) and generates on the SAME handle, + * so the model stays prepared across lines. Each line may be prefixed with + * "IMAGE:" to feed an image followed by the rest as the text prompt. + * + * STREAMING=1 (default): the on_token callback prints each token as generated. + * STREAMING=0: the on_token callback appends to a std::string and the full + * answer is printed once when generation ends. Override with -DSTREAMING=0. + */ +#include +#include +#include +#include +#include +#include + +#include "exynos_llm_wrapper.h" + +#ifndef STREAMING +#define STREAMING 1 +#endif + +static void on_token(const char* text, int32_t len, void* userdata) { +#if STREAMING + (void)userdata; + fwrite(text, 1, (size_t)len, stdout); + fflush(stdout); +#else + auto* acc = static_cast(userdata); + acc->append(text, (size_t)len); +#endif +} + +static int generate_one(enn_runner* r, const char* line) { + const char* img = NULL; + const char* prompt = line; + if (strncmp(line, "IMAGE:", 6) == 0) { + img = line + 6; + prompt = "Describe this image."; + } + printf("----- generate (image=%s) -----\n", img ? img : "none"); + +#if STREAMING + int rc = enn_llm_generate( + r, prompt, img, NULL, 196, 0.0f, 0, 0, 0, on_token, NULL); +#else + std::string acc; + int rc = enn_llm_generate( + r, prompt, img, NULL, 196, 0.0f, 0, 0, 0, on_token, &acc); + printf("%s\n", acc.c_str()); + fflush(stdout); +#endif + printf("\n[rc=%d]\n", rc); + return rc; +} + +int main(int argc, char** argv) { + if (argc < 6) { + fprintf(stderr, "usage: %s model ptd tok pos cfg [vision]\n", argv[0]); + fprintf( + stderr, "then read prompts from stdin; 'IMAGE:' for an image.\n"); + return 2; + } + const char* vision = (argc > 6 && argv[6][0]) ? argv[6] : NULL; + + printf("[sample] creating runner...\n"); + enn_runner* r = + enn_llm_create(argv[1], argv[2], argv[3], argv[4], argv[5], vision); + if (!r) { + fprintf(stderr, "enn_llm_create FAILED\n"); + return 1; + } + printf("[sample] created. enter prompts (Ctrl-D to end):\n"); + + // Stopper thread that invokes enn_llm_stop to imitate a user interrupt. + std::thread stopper_thread([r]() { + std::cout << "[Stopper Thread] Waiting 3 seconds before interrupting..." + << std::endl; + std::this_thread::sleep_for(std::chrono::seconds(3)); + std::cout << "[Stopper Thread] Triggering enn_llm_stop()..." << std::endl; + enn_llm_stop(r); + }); + + std::cout << "[CLI] Starting token generation loop..." << std::endl; + int rc_total = 0; + char buf[4096]; + while (fgets(buf, sizeof(buf), stdin)) { + size_t n = strlen(buf); + while (n > 0 && (buf[n - 1] == '\n' || buf[n - 1] == '\r')) { + buf[--n] = '\0'; + } + if (n == 0) { + continue; + } + rc_total |= generate_one(r, buf); + } + + if (stopper_thread.joinable()) { + stopper_thread.join(); + } + + printf("[sample] destroying runner...\n"); + enn_llm_destroy(r); + return rc_total ? 1 : 0; +} From 2953a062c0dc5c0266b60db2f879aa1c9597e09f Mon Sep 17 00:00:00 2001 From: Jiseong Oh Date: Tue, 13 Oct 2026 19:39:46 +0900 Subject: [PATCH 15/17] Set the ENN perf mode from the SoC and extend the LLM runner CLI The ENN perf mode is set from the detected SoC. enn_executor_llm_runner drives the runner through the exynos_llm_wrapper C API and gains a system prompt, --prompt_file, --system_prompt_file, --max_token_limit, configurable prompt echo and a thread-safe stop. The Gemma 4 prefill sliding-window mask is fixed, BOS is prepended to Gemma 4 text prompts, and the Samsung Android build gets a --vulkan option. Signed-off-by: Jiseong Oh Co-authored-by: Hoon Choi --- backends/samsung/build.sh | 8 + .../runtime/enn_api_implementation.cpp | 102 +++++++++++ .../samsung/runtime/enn_api_implementation.h | 8 + backends/samsung/runtime/enn_executor.cpp | 3 + backends/samsung/runtime/enn_type.h | 7 + examples/samsung/CMakeLists.txt | 141 +++++++-------- examples/samsung/build.sh | 20 +++ examples/samsung/executor_runner/README.md | 9 +- .../enn_executor_llm_runner.cpp | 160 +++++++++++------- .../executor_runner/exynos_llm_wrapper.cpp | 16 +- .../executor_runner/exynos_llm_wrapper.h | 17 +- .../exynos_llm_wrapper_sample_main.cpp | 115 ++++++++++--- .../executor_runner/llm/runner/config.h | 3 + .../executor_runner/llm/runner/io_manager.cpp | 23 +-- .../executor_runner/llm/runner/runner.cpp | 23 ++- .../executor_runner/llm/runner/runner.h | 2 - .../llm/runner/token_generator.h | 20 ++- 17 files changed, 486 insertions(+), 191 deletions(-) diff --git a/backends/samsung/build.sh b/backends/samsung/build.sh index a75c0c7bba3..4c268a55ef9 100755 --- a/backends/samsung/build.sh +++ b/backends/samsung/build.sh @@ -7,6 +7,7 @@ X86_64_BUILD_DIR=${PROJECT_DIR}/build_samsung ANDROID_BUILD_DIR=${PROJECT_DIR}/build_samsung_android BUILD_ARCH=all CLEAN_BUILD_DIR=false +BUILD_VULKAN=OFF EXYNOS_AI_LITECORE_ROOT=${EXYNOS_AI_LITECORE_ROOT} ANDROID_NDK_ROOT=${ANDROID_NDK_ROOT} @@ -26,6 +27,9 @@ function usage() { --build, -b [x86_64, android, all] Default is all, x86_64 target to offline compilation, android target to online execution. + --vulkan Build the Vulkan backend into executor_runner (android, default is off). + Requires glslc from the Vulkan SDK (1.4.341.1+) in PATH, e.g. + export PATH=\$VULKAN_SDK/bin:\$PATH --clean, -c Clean the build cache. --help, -h Print the usage information. " @@ -104,6 +108,7 @@ function build_android() { -DEXECUTORCH_BUILD_KERNELS_OPTIMIZED_PORTABLE=ON \ -DEXECUTORCH_BUILD_KERNELS_QUANTIZED=ON \ -DEXECUTORCH_BUILD_XNNPACK=ON \ + -DEXECUTORCH_BUILD_VULKAN=${BUILD_VULKAN} \ -DEXECUTORCH_ENABLE_LOGGING=1 \ -DEXECUTORCH_BUILD_DEVTOOLS=ON \ -DEXECUTORCH_ENABLE_EVENT_TRACER=ON \ @@ -126,6 +131,9 @@ do let i++ ANDROID_NDK_ROOT="${!i}" ;; + "--vulkan") + BUILD_VULKAN=ON + ;; "--clean"|"-c") CLEAN_BUILD_DIR=true ;; diff --git a/backends/samsung/runtime/enn_api_implementation.cpp b/backends/samsung/runtime/enn_api_implementation.cpp index 9ea5317b870..35e866e5741 100644 --- a/backends/samsung/runtime/enn_api_implementation.cpp +++ b/backends/samsung/runtime/enn_api_implementation.cpp @@ -7,9 +7,15 @@ * */ #include +#include + #include #include +#include +#include +#include + #define ENN_LOAD_API_FUNC(handle, name, enn_api_ptr) \ enn_api_ptr->name = \ reinterpret_cast(loadApiFunction(handle, #name, false)); \ @@ -38,6 +44,101 @@ EnnApi* EnnApi::getEnnApiInstance() { return &enn_api; } +namespace { + +constexpr uint32_t kPerfConfigIdDefault = 0; + +struct PerfConfig { + PerfModePreference mode; + uint32_t config_id; +}; + +// Primary source: ro.hardware. Falls back to ro.soc.model if unavailable. +std::string getHwName() { + char prop_value[PROP_VALUE_MAX] = {}; + if (__system_property_get("ro.hardware", prop_value) > 0) { + return std::string(prop_value); + } + return ""; +} + +std::string getSocName() { + char prop_value[PROP_VALUE_MAX] = {}; + if (__system_property_get("ro.soc.model", prop_value) > 0) { + return std::string(prop_value); + } + return ""; +} + +// Extracts trailing digits, e.g. "exynos9955" -> "9955", "s5e9965" -> "9965". +bool trimToTrailingDigits(std::string& value) { + const char* digits = "0123456789"; + auto last = value.find_last_of(digits); + if (last == std::string::npos) { + return false; + } + auto first = last; + while (first > 0 && std::strchr(digits, value[first - 1]) != nullptr) { + first--; + } + value = value.substr(first, last - first + 1); + return true; +} + +const std::unordered_map& getSocPerfConfigMap() { + static const std::unordered_map kSocPerfConfig = { + {"9955", {ENN_PREF_MODE_PERFORMANCE, kPerfConfigIdDefault}}, + {"9965", {ENN_PREF_MODE_CUSTOM, 100111}}, + }; + return kSocPerfConfig; +} + +PerfConfig getPerfConfigForSoc() { + std::string key = getHwName(); + if (key.empty()) { + key = getSocName(); + } + if (!key.empty()) { + trimToTrailingDigits(key); + } + + const auto& map = getSocPerfConfigMap(); + auto it = map.find(key); + if (it != map.end()) { + return it->second; + } + return {ENN_PREF_MODE_PERFORMANCE, kPerfConfigIdDefault}; +} + +// Mirrors LiteRT's litert::samsung::SetGenAiPerfConfigFromSoc: detect the +// SoC from Android system properties and apply its known-good perf mode, +// instead of leaving the ENN runtime on its own default. +void setPerfConfigFromSoc(const EnnApi* enn_api) { + PerfConfig config = getPerfConfigForSoc(); + if (enn_api->EnnSetPreferencePerfMode(static_cast(config.mode)) != + ENN_RET_SUCCESS) { + ENN_LOG_ERROR("Failed to set ENN preference perf mode."); + return; + } + if (config.config_id == kPerfConfigIdDefault) { + return; + } + if (enn_api->EnnSetPreferencePerfConfigId(config.config_id) != + ENN_RET_SUCCESS) { + ENN_LOG_WARN( + "Failed to set ENN perf config id %u, falling back to PERFORMANCE mode.", + config.config_id); + enn_api->EnnSetPreferencePerfMode( + static_cast(ENN_PREF_MODE_PERFORMANCE)); + } +} + +} // namespace + +void EnnApi::applyPerfConfigFromSoc() const { + setPerfConfigFromSoc(this); +} + EnnApi::EnnApi() { auto status = loadApiLib(); if (status != Error::Ok) { @@ -74,6 +175,7 @@ Error EnnApi::loadApiLib() { ENN_LOAD_API_FUNC(libenn_public_api_, EnnInitialize, this); ENN_LOAD_API_FUNC(libenn_public_api_, EnnSetPreferencePerfMode, this); ENN_LOAD_API_FUNC(libenn_public_api_, EnnGetPreferencePerfMode, this); + ENN_LOAD_API_FUNC(libenn_public_api_, EnnSetPreferencePerfConfigId, this); ENN_LOAD_API_FUNC(libenn_public_api_, EnnOpenModel, this); ENN_LOAD_API_FUNC(libenn_public_api_, EnnOpenModelFromMemory, this); ENN_LOAD_API_FUNC(libenn_public_api_, EnnSetFastIpc, this); diff --git a/backends/samsung/runtime/enn_api_implementation.h b/backends/samsung/runtime/enn_api_implementation.h index 675d877565b..2c8e6577012 100644 --- a/backends/samsung/runtime/enn_api_implementation.h +++ b/backends/samsung/runtime/enn_api_implementation.h @@ -30,9 +30,16 @@ class EnnApi { static EnnApi* getEnnApiInstance(); bool isInitialized() const; + // Detects the SoC and applies its known-good ENN perf mode/config id. + // Must be called on the same thread that is about to open a model, and + // as the last preference-mutating call before that open, since + // EnnInitialize() resets the calling thread's preference on every call. + void applyPerfConfigFromSoc() const; + EnnReturn (*EnnInitialize)(void); EnnReturn (*EnnSetPreferencePerfMode)(const uint32_t val); EnnReturn (*EnnGetPreferencePerfMode)(uint32_t* val_ptr); + EnnReturn (*EnnSetPreferencePerfConfigId)(uint32_t val); EnnReturn (*EnnOpenModel)(const char* model_file, EnnModelId* model_id); EnnReturn (*EnnOpenModelFromMemory)( const char* va, @@ -99,6 +106,7 @@ class EnnApi { typedef EnnReturn (*EnnInitialize_fn)(void); typedef EnnReturn (*EnnSetPreferencePerfMode_fn)(const uint32_t val); typedef EnnReturn (*EnnGetPreferencePerfMode_fn)(uint32_t* val_ptr); +typedef EnnReturn (*EnnSetPreferencePerfConfigId_fn)(uint32_t val); typedef EnnReturn ( *EnnOpenModel_fn)(const char* model_file, EnnModelId* model_id); typedef EnnReturn (*EnnOpenModelFromMemory_fn)( diff --git a/backends/samsung/runtime/enn_executor.cpp b/backends/samsung/runtime/enn_executor.cpp index 7572f278510..a27d7a4df30 100644 --- a/backends/samsung/runtime/enn_executor.cpp +++ b/backends/samsung/runtime/enn_executor.cpp @@ -35,6 +35,9 @@ Error EnnExecutor::initialize( const EnnApi* enn_api_inst = EnnApi::getEnnApiInstance(); ET_CHECK_OR_RETURN_ERROR( enn_api_inst->isInitialized(), Internal, "Enn initialize failed."); + // EnnInitialize() resets the calling thread's preference on every call, so + // re-apply it right before opening the model. + enn_api_inst->applyPerfConfigFromSoc(); EnnReturn ret = ENN_RET_SUCCESS; ET_LOG(Info, "Start to open model %p, %zu", binary_buf_addr, buf_size); diff --git a/backends/samsung/runtime/enn_type.h b/backends/samsung/runtime/enn_type.h index 4b91fe2d08a..f4ae534b020 100644 --- a/backends/samsung/runtime/enn_type.h +++ b/backends/samsung/runtime/enn_type.h @@ -53,3 +53,10 @@ typedef struct _NumberOfBuffersInfo { uint32_t n_in_buf; uint32_t n_out_buf; } NumberOfBuffersInfo; + +typedef enum _PerfModePreference { + ENN_PREF_MODE_LOW_POWER = 0, + ENN_PREF_MODE_BALANCED = 4, + ENN_PREF_MODE_CUSTOM = 5, + ENN_PREF_MODE_PERFORMANCE = 7, +} PerfModePreference; diff --git a/examples/samsung/CMakeLists.txt b/examples/samsung/CMakeLists.txt index 9892f2a245b..d1dfbc26fff 100644 --- a/examples/samsung/CMakeLists.txt +++ b/examples/samsung/CMakeLists.txt @@ -145,28 +145,11 @@ set_target_properties( "${EXECUTORCH_ROOT}/build_samsung_android/lib/libcustom_ops.a" ) -set(enn_executor_llm_runner_srcs - ${CMAKE_CURRENT_LIST_DIR}/executor_runner/enn_executor_llm_runner.cpp -) - set(enn_executor_llm_runner_deps executorch_core extension_data_loader extension_module extension_tensor extension_flat_tensor lib_extension_llm_runner ) -add_executable(enn_executor_llm_runner ${enn_executor_llm_runner_srcs}) - -target_link_libraries( - enn_executor_llm_runner PRIVATE ${enn_executor_llm_runner_deps} -) - -target_include_directories( - enn_executor_llm_runner - PUBLIC ${_common_include_directories} - ${EXECUTORCH_ROOT}/backends/samsung/runtime - ${EXECUTORCH_ROOT}/extension/llm/tokenizers/include -) - set(xnnpack_backend_libs xnnpack_backend XNNPACK xnnpack-microkernels-prod extension_threadpool pthreadpool kleidiai ) @@ -191,74 +174,80 @@ list(APPEND _common_include_directories # quantized_ops_lib: Register quantized op kernels into the runtime list(APPEND xnnpack_backend_libs quantized_kernels quantized_ops_lib) -target_link_libraries( - enn_executor_llm_runner - PUBLIC enn_llm_runner - PRIVATE log - gflags - android - lib_portable_ops - ${enn_executor_llm_runner_deps} - ${xnnpack_backend_libs} - sentencepiece - re2 - lib_tokenizer - -Wl,--whole-archive - lib_executorch - lib_custom_op - lib_enn_backend - enn_logging - -Wl,--no-whole-archive +set(exynos_llm_link_libs + log + gflags + android + lib_portable_ops + ${enn_executor_llm_runner_deps} + ${xnnpack_backend_libs} + sentencepiece + re2 + lib_tokenizer + -Wl,--whole-archive + lib_executorch + lib_custom_op + lib_enn_backend + enn_logging + -Wl,--no-whole-archive +) + +# C API in exynos_llm_wrapper.h. Always compiled into enn_executor_llm_runner; +# additionally shipped as libexynos_llm_wrapper.so when +# EXYNOS_BUILD_LLM_WRAPPER_SO is ON. +option(EXYNOS_BUILD_LLM_WRAPPER_SO + "Build libexynos_llm_wrapper.so and exynos_llm_wrapper_sample" OFF ) -# Public shared library exposing the C API in exynos_llm_wrapper.h. Bundles all -# dependencies so only the .so + header need to be shipped to third parties. add_library( - exynos_llm_wrapper SHARED + exynos_llm_wrapper_obj OBJECT ${CMAKE_CURRENT_LIST_DIR}/executor_runner/exynos_llm_wrapper.cpp ) - -target_compile_definitions(exynos_llm_wrapper PRIVATE EXECUTORCH_VERSION=2) -target_include_directories( - exynos_llm_wrapper PUBLIC ${CMAKE_CURRENT_LIST_DIR}/executor_runner -) - -target_link_libraries( - exynos_llm_wrapper - PUBLIC enn_llm_runner - PRIVATE log - gflags - android - lib_portable_ops - ${enn_executor_llm_runner_deps} - ${xnnpack_backend_libs} - sentencepiece - re2 - lib_tokenizer - -Wl,--whole-archive - lib_executorch - lib_custom_op - lib_enn_backend - enn_logging - -Wl,--no-whole-archive - -Wl,--exclude-libs,ALL -) - +target_compile_definitions(exynos_llm_wrapper_obj PRIVATE EXECUTORCH_VERSION=2) +target_link_libraries(exynos_llm_wrapper_obj PRIVATE enn_llm_runner) set_target_properties( - exynos_llm_wrapper PROPERTIES CXX_VISIBILITY_PRESET hidden - VISIBILITY_INLINES_HIDDEN ON + exynos_llm_wrapper_obj PROPERTIES CXX_VISIBILITY_PRESET hidden + VISIBILITY_INLINES_HIDDEN ON ) -# Sample CLI consuming the public exynos_llm_wrapper C API (same shape as an app -# using libexynos_llm_wrapper.so). Links the shared lib so it coexists with the -# all-in-one enn_executor_llm_runner. add_executable( - exynos_llm_wrapper_sample - ${CMAKE_CURRENT_LIST_DIR}/executor_runner/exynos_llm_wrapper_sample_main.cpp + enn_executor_llm_runner + ${CMAKE_CURRENT_LIST_DIR}/executor_runner/enn_executor_llm_runner.cpp ) -target_link_libraries(exynos_llm_wrapper_sample PRIVATE exynos_llm_wrapper) - -set_target_properties( - exynos_llm_wrapper_sample PROPERTIES CXX_VISIBILITY_PRESET hidden +target_include_directories( + enn_executor_llm_runner + PUBLIC ${_common_include_directories} + ${EXECUTORCH_ROOT}/backends/samsung/runtime + ${EXECUTORCH_ROOT}/extension/llm/tokenizers/include ) + +target_link_libraries( + enn_executor_llm_runner PRIVATE exynos_llm_wrapper_obj enn_llm_runner + ${exynos_llm_link_libs} +) + +if(EXYNOS_BUILD_LLM_WRAPPER_SO) + # Bundles all dependencies so only the .so + header need to be shipped. + add_library( + exynos_llm_wrapper SHARED $ + ) + set_target_properties(exynos_llm_wrapper PROPERTIES LINKER_LANGUAGE CXX) + target_include_directories( + exynos_llm_wrapper PUBLIC ${CMAKE_CURRENT_LIST_DIR}/executor_runner + ) + target_link_libraries( + exynos_llm_wrapper PRIVATE enn_llm_runner ${exynos_llm_link_libs} + -Wl,--exclude-libs,ALL + ) + + # Sample CLI consuming the public C API the way an app uses the .so. + add_executable( + exynos_llm_wrapper_sample + ${CMAKE_CURRENT_LIST_DIR}/executor_runner/exynos_llm_wrapper_sample_main.cpp + ) + target_link_libraries(exynos_llm_wrapper_sample PRIVATE exynos_llm_wrapper) + set_target_properties( + exynos_llm_wrapper_sample PROPERTIES CXX_VISIBILITY_PRESET hidden + ) +endif() diff --git a/examples/samsung/build.sh b/examples/samsung/build.sh index 0a0a8b0e1cf..6fb0d38c136 100755 --- a/examples/samsung/build.sh +++ b/examples/samsung/build.sh @@ -19,7 +19,21 @@ echo ANDROID_NDK_ROOT=${ANDROID_NDK_ROOT} echo ANDROID_ABI=${ANDROID_ABI} echo ANDROID_PLATFORM=${ANDROID_PLATFORM} +usage() { + echo "Usage: $0 [--llm-wrapper-so]" + echo " --llm-wrapper-so Also build libexynos_llm_wrapper.so and exynos_llm_wrapper_sample." +} + main() { + local build_llm_wrapper_so=OFF + for arg in "$@"; do + case "$arg" in + "--llm-wrapper-so") build_llm_wrapper_so=ON ;; + "--help"|"-h") usage; exit 0 ;; + *) echo "Unknown option: $arg"; usage; exit 1 ;; + esac + done + cd "$PROJECT_DIR" if [[ -z ${ANDROID_NDK_ROOT} ]]; then @@ -48,9 +62,15 @@ main() { -DANDROID_PLATFORM=$ANDROID_PLATFORM \ -DCMAKE_FIND_ROOT_PATH_MODE_PACKAGE=BOTH \ -DCMAKE_BUILD_TYPE=Release \ + -DEXYNOS_BUILD_LLM_WRAPPER_SO=${build_llm_wrapper_so} \ -B"${build_dir_example}" \ "$PROJECT_DIR/${example_root}" cmake --build ${build_dir_example} --config Release + + if [[ ${build_llm_wrapper_so} == OFF ]]; then + rm -f "${build_dir_example}/libexynos_llm_wrapper.so" \ + "${build_dir_example}/exynos_llm_wrapper_sample" + fi } main "$@" diff --git a/examples/samsung/executor_runner/README.md b/examples/samsung/executor_runner/README.md index ce84b88aaad..ad10fe320cd 100644 --- a/examples/samsung/executor_runner/README.md +++ b/examples/samsung/executor_runner/README.md @@ -31,4 +31,11 @@ Ensure these files are pushed to the target device (e.g., to `/data/local/tmp/`) `libexynos_llm_wrapper.so` is a reference for porting the Exynos LLM runner into an Android app: a small public C API (`exynos_llm_wrapper.h`) that hides the runner internals behind an opaque handle. `exynos_llm_wrapper_sample` is a CLI -that exercises the same API on-device. \ No newline at end of file +that exercises the same API on-device. + +`enn_executor_llm_runner` always drives the runner through this C API, but the +`.so` and the sample are only produced when requested: + +```bash +./examples/samsung/build.sh --llm-wrapper-so +``` \ No newline at end of file diff --git a/examples/samsung/executor_runner/enn_executor_llm_runner.cpp b/examples/samsung/executor_runner/enn_executor_llm_runner.cpp index e527ace4eab..37219efcc41 100644 --- a/examples/samsung/executor_runner/enn_executor_llm_runner.cpp +++ b/examples/samsung/executor_runner/enn_executor_llm_runner.cpp @@ -8,15 +8,20 @@ #include -#include +#include +#include +#include +#include + +#include +#include +#include #if defined(ET_USE_THREADPOOL) #include #include #endif -namespace llm = ::executorch::extension::llm; - DEFINE_string( model_path, "model.pte", @@ -41,6 +46,18 @@ DEFINE_string( DEFINE_string(prompt, "The answer to the ultimate question is", "Prompt."); +DEFINE_string(system_prompt, "", "System prompt. Empty means none."); + +DEFINE_string( + prompt_file, + "", + "Read the prompt from this file instead of --prompt."); + +DEFINE_string( + system_prompt_file, + "", + "Read the system prompt from this file instead of --system_prompt."); + DEFINE_double( temperature, 0.8, @@ -70,49 +87,69 @@ DEFINE_bool(is_warmup, false, "Whether to run a warmup run."); DEFINE_bool(is_decode_only, false, "Whether to run with decode method only."); -int32_t main(int32_t argc, char** argv) { - gflags::ParseCommandLineFlags(&argc, &argv, true); - const char* model_path = FLAGS_model_path.c_str(); - - std::optional data_path = std::nullopt; - if (!FLAGS_data_path.empty()) { - data_path = FLAGS_data_path.c_str(); +DEFINE_int32( + max_token_limit, + 0, + "If > 0, print each token from the callback and call enn_llm_stop() right after the N-th generated token. Used to verify the stop API."); + +struct StopContext { + enn_runner* runner; + int32_t limit; + int32_t num_pieces = 0; + int32_t num_tokens = 0; +}; + +static void on_token_with_limit(const char* text, int32_t len, void* userdata) { + auto* ctx = static_cast(userdata); + fwrite(text, 1, static_cast(len), stdout); + fflush(stdout); + // The first callback is the echoed input prompt, not a generated token. + if (ctx->num_pieces++ == 0) { + return; } - const char* tokenizer_path = FLAGS_tokenizer_path.c_str(); - - std::optional pos_embedding_path = std::nullopt; - if (!FLAGS_pos_embedding_path.empty()) { - pos_embedding_path = FLAGS_pos_embedding_path; + // The runner emits a trailing "\n" after generation ends; ignore it. + if (ctx->num_tokens >= ctx->limit) { + return; } + if (++ctx->num_tokens == ctx->limit) { + enn_llm_stop(ctx->runner); + printf("\n[STOP] enn_llm_stop called at token %d\n", ctx->num_tokens); + fflush(stdout); + } +} - const char* config_path = FLAGS_config_path.c_str(); - std::optional vision_model_path = std::nullopt; - if (!FLAGS_vision_model_path.empty()) { - vision_model_path = FLAGS_vision_model_path; +static bool read_file(const std::string& path, std::string* out) { + std::ifstream file(path, std::ios::binary); + if (!file.is_open()) { + ET_LOG(Error, "Failed to open %s", path.c_str()); + return false; } + std::stringstream ss; + ss << file.rdbuf(); + *out = ss.str(); + return true; +} - const char* prompt = FLAGS_prompt.c_str(); +int32_t main(int32_t argc, char** argv) { + gflags::ParseCommandLineFlags(&argc, &argv, true); + + if ((!FLAGS_prompt_file.empty() && + !read_file(FLAGS_prompt_file, &FLAGS_prompt)) || + (!FLAGS_system_prompt_file.empty() && + !read_file(FLAGS_system_prompt_file, &FLAGS_system_prompt))) { + return ENN_ERR_INVALID_ARG; + } if (FLAGS_prompt.empty()) { ET_LOG(Error, "Input prompt cannot be empty."); - return static_cast(executorch::runtime::Error::InvalidArgument); + return ENN_ERR_INVALID_ARG; } - ET_LOG(Info, "Input prompt message: %s", prompt); - - float temperature = static_cast(FLAGS_temperature); - int32_t seq_len = FLAGS_seq_len; - int32_t cpu_threads = FLAGS_cpu_threads; - - bool is_warmup = FLAGS_is_warmup; - bool is_decode_only = FLAGS_is_decode_only; - - int num_bos = FLAGS_num_bos; - int num_eos = FLAGS_num_eos; + ET_LOG(Info, "Input prompt message: %s", FLAGS_prompt.c_str()); #if defined(ET_USE_THREADPOOL) - uint32_t num_performant_cores = cpu_threads == -1 + uint32_t num_performant_cores = FLAGS_cpu_threads == -1 ? ::executorch::extension::cpuinfo::get_num_performant_cores() - : static_cast(cpu_threads); + : static_cast(FLAGS_cpu_threads); ET_LOG( Info, "Resetting threadpool with num threads = %d", num_performant_cores); if (num_performant_cores > 0) { @@ -121,33 +158,42 @@ int32_t main(int32_t argc, char** argv) { } #endif - example::RunnerPaths runner_paths{ - .model_path = model_path, - .tokenizer_path = tokenizer_path, - .pos_embedding_path = pos_embedding_path, - .data_path = data_path, - .vision_model_path = vision_model_path}; - example::RunnerConfig runner_config{ - .config_path = config_path, - .is_decode_only = is_decode_only, - .verbose = true, - .image_path = FLAGS_image_path, - .prompt_before_image = FLAGS_prompt_before_image}; - std::unique_ptr runner = - example::create_llm_runner(runner_paths, std::move(runner_config)); + enn_runner* runner = enn_llm_create( + FLAGS_model_path.c_str(), + FLAGS_data_path.c_str(), + FLAGS_tokenizer_path.c_str(), + FLAGS_pos_embedding_path.c_str(), + FLAGS_config_path.c_str(), + FLAGS_vision_model_path.c_str(), + FLAGS_is_decode_only); ET_CHECK_MSG(runner != nullptr, "Failed to create LLM runner"); // TODO: support warmup interface - executorch::extension::llm::GenerationConfig config{ - .seq_len = seq_len, - .temperature = temperature, - .num_bos = num_bos, - .num_eos = num_eos}; - auto error = runner->generate(prompt, config, {}, {}); - ET_CHECK_MSG( - error == executorch::runtime::Error::Ok, - "Failed to generate token from LLM runner"); + const bool stop_test = FLAGS_max_token_limit > 0; + StopContext stop_ctx{runner, FLAGS_max_token_limit}; + int32_t error = enn_llm_generate( + runner, + FLAGS_prompt.c_str(), + FLAGS_system_prompt.c_str(), + FLAGS_image_path.c_str(), + FLAGS_prompt_before_image.c_str(), + FLAGS_seq_len, + static_cast(FLAGS_temperature), + FLAGS_num_bos, + FLAGS_num_eos, + /*echo=*/1, + /*verbose=*/!stop_test, + stop_test ? on_token_with_limit : nullptr, + stop_test ? &stop_ctx : nullptr); + enn_llm_destroy(runner); + if (stop_test) { + printf( + "\n[STOP] limit=%d, generated=%d\n", + stop_ctx.limit, + stop_ctx.num_tokens); + } + ET_CHECK_MSG(error == ENN_OK, "Failed to generate token from LLM runner"); return 0; } diff --git a/examples/samsung/executor_runner/exynos_llm_wrapper.cpp b/examples/samsung/executor_runner/exynos_llm_wrapper.cpp index 3548e52546b..2982d4d484e 100644 --- a/examples/samsung/executor_runner/exynos_llm_wrapper.cpp +++ b/examples/samsung/executor_runner/exynos_llm_wrapper.cpp @@ -46,7 +46,8 @@ enn_runner* enn_llm_create( const char* tokenizer_path, const char* pos_embedding_path, const char* config_path, - const char* vision_model_path) { + const char* vision_model_path, + int32_t is_decode_only) { if (model_path == nullptr || tokenizer_path == nullptr || config_path == nullptr) { ET_LOG(Error, "enn_llm_create: model/tokenizer/config path required"); @@ -60,7 +61,8 @@ enn_runner* enn_llm_create( .data_path = opt_or_null(data_path), .vision_model_path = opt_or_null(vision_model_path)}; - example::RunnerConfig cfg{.config_path = config_path}; + example::RunnerConfig cfg{ + .config_path = config_path, .is_decode_only = is_decode_only != 0}; auto runner = example::create_llm_runner(paths, std::move(cfg)); if (runner == nullptr) { @@ -76,18 +78,23 @@ void enn_llm_destroy(enn_runner* runner) { } void enn_llm_stop(enn_runner* runner) { + if (runner == nullptr) { + return; + } runner->impl->stop(); } int32_t enn_llm_generate( enn_runner* runner, const char* prompt, + const char* system_prompt, const char* image_path, const char* prompt_before_image, int32_t seq_len, float temperature, int32_t num_bos, int32_t num_eos, + int32_t echo, int32_t verbose, enn_token_cb on_token, void* userdata) { @@ -105,12 +112,15 @@ int32_t enn_llm_generate( // Start each generation from a clean KV-cache. lrunner->reset(); - // Image inputs are threaded through the runner's config (read by generate). + // System prompt and image inputs are threaded through the runner's config + // (read by generate). + lrunner->get_config().system_prompt = str_or_empty(system_prompt); lrunner->get_config().image_path = str_or_empty(image_path); lrunner->get_config().prompt_before_image = str_or_empty(prompt_before_image); lrunner->get_config().verbose = verbose != 0; llm::GenerationConfig gen_cfg{ + .echo = echo != 0, .seq_len = seq_len >= 0 ? seq_len : -1, .temperature = temperature, .num_bos = num_bos, diff --git a/examples/samsung/executor_runner/exynos_llm_wrapper.h b/examples/samsung/executor_runner/exynos_llm_wrapper.h index 229477f4bfb..d9764de93bc 100644 --- a/examples/samsung/executor_runner/exynos_llm_wrapper.h +++ b/examples/samsung/executor_runner/exynos_llm_wrapper.h @@ -60,6 +60,8 @@ typedef struct enn_runner enn_runner; * @param pos_embedding_path Precomputed rotary embeddings, or NULL. * @param config_path Runner config file (required). * @param vision_model_path Vision encoder .pte, or NULL => text-only runner. + * @param is_decode_only Non-zero runs the prompt through the decode method only + * (no prefill method). * @return A handle on success, NULL on failure. */ EXYNOS_LLM_WRAPPER_API enn_runner* enn_llm_create( @@ -68,12 +70,18 @@ EXYNOS_LLM_WRAPPER_API enn_runner* enn_llm_create( const char* tokenizer_path, const char* pos_embedding_path, const char* config_path, - const char* vision_model_path); + const char* vision_model_path, + int32_t is_decode_only); /** Destroy the runner and release all resources. Safe to call with NULL. */ EXYNOS_LLM_WRAPPER_API void enn_llm_destroy(enn_runner* runner); -/** Request the runner to stop current generation. */ +/** + * Request the in-flight enn_llm_generate on runner to stop. Thread-safe; may + * be called from any thread, including from on_token. Generation returns + * ENN_OK before the next decode step. A no-op when no generation is running + * or runner is NULL. + */ EXYNOS_LLM_WRAPPER_API void enn_llm_stop(enn_runner* runner); /** @@ -83,6 +91,7 @@ EXYNOS_LLM_WRAPPER_API void enn_llm_stop(enn_runner* runner); * * @param runner Handle from enn_llm_create. * @param prompt User prompt text (required, non-empty). + * @param system_prompt System prompt text, or NULL/"" => none. * @param image_path Image file for this turn, or NULL => text-only. * @param prompt_before_image Text placed before the image in the user turn * (Gemma 4 multimodal); ignored when image_path is NULL. @@ -91,6 +100,8 @@ EXYNOS_LLM_WRAPPER_API void enn_llm_stop(enn_runner* runner); * @param temperature Sampling temperature. * @param num_bos BOS tokens to prepend to the prompt. * @param num_eos EOS tokens to append to the prompt. + * @param echo Non-zero also passes the rendered input prompt to on_token + * before the generated tokens. * @param verbose Non-zero also streams tokens to stdout. * @param on_token Per-piece callback, or NULL to discard output. * @param userdata Opaque pointer passed through to on_token. @@ -99,12 +110,14 @@ EXYNOS_LLM_WRAPPER_API void enn_llm_stop(enn_runner* runner); EXYNOS_LLM_WRAPPER_API int32_t enn_llm_generate( enn_runner* runner, const char* prompt, + const char* system_prompt, const char* image_path, const char* prompt_before_image, int32_t seq_len, float temperature, int32_t num_bos, int32_t num_eos, + int32_t echo, int32_t verbose, enn_token_cb on_token, void* userdata); diff --git a/examples/samsung/executor_runner/exynos_llm_wrapper_sample_main.cpp b/examples/samsung/executor_runner/exynos_llm_wrapper_sample_main.cpp index 31bebd3bbbd..ad2c9140966 100644 --- a/examples/samsung/executor_runner/exynos_llm_wrapper_sample_main.cpp +++ b/examples/samsung/executor_runner/exynos_llm_wrapper_sample_main.cpp @@ -9,6 +9,7 @@ * * Usage (paths are on-device): * exynos_llm_wrapper_sample [vision] + * [max_output_tokens] * * Reads prompts from stdin (one per line) and generates on the SAME handle, * so the model stays prepared across lines. Each line may be prefixed with @@ -17,8 +18,14 @@ * STREAMING=1 (default): the on_token callback prints each token as generated. * STREAMING=0: the on_token callback appends to a std::string and the full * answer is printed once when generation ends. Override with -DSTREAMING=0. + * + * Stop test: pass max_output_tokens=N (> 0; use "" for vision to skip it) to + * call enn_llm_stop() right after the N-th output token of each prompt and + * check that no more tokens arrive. Use prompts whose answer is longer than N. + * Without it, a stopper thread calls enn_llm_stop() once after 3 seconds. */ #include +#include #include #include #include @@ -31,18 +38,50 @@ #define STREAMING 1 #endif +struct GenContext { + enn_runner* runner; + int32_t max_output_tokens; + int32_t num_pieces = 0; + int32_t num_tokens = 0; + int32_t tokens_after_stop = 0; + std::string acc; +}; + +static void +count_for_stop_test(GenContext* ctx, const char* text, int32_t len) { + // The first callback is the echoed input prompt, not an output token. + if (ctx->num_pieces++ == 0) { + return; + } + // The runner emits a trailing "\n" after generation ends; ignore it. + if (ctx->num_tokens >= ctx->max_output_tokens) { + if (!(len == 1 && text[0] == '\n')) { + ++ctx->tokens_after_stop; + } + return; + } + if (++ctx->num_tokens == ctx->max_output_tokens) { + enn_llm_stop(ctx->runner); + printf("\n[STOP TEST] enn_llm_stop called at token %d\n", ctx->num_tokens); + fflush(stdout); + } +} + static void on_token(const char* text, int32_t len, void* userdata) { + auto* ctx = static_cast(userdata); #if STREAMING - (void)userdata; fwrite(text, 1, (size_t)len, stdout); fflush(stdout); #else - auto* acc = static_cast(userdata); - acc->append(text, (size_t)len); + ctx->acc.append(text, (size_t)len); #endif + if (ctx->max_output_tokens > 0) { + count_for_stop_test(ctx, text, len); + } } -static int generate_one(enn_runner* r, const char* line) { +static int +generate_one(enn_runner* r, const char* line, int32_t max_output_tokens) { const char* img = NULL; const char* prompt = line; if (strncmp(line, "IMAGE:", 6) == 0) { @@ -51,32 +90,62 @@ static int generate_one(enn_runner* r, const char* line) { } printf("----- generate (image=%s) -----\n", img ? img : "none"); -#if STREAMING + GenContext ctx{r, max_output_tokens}; int rc = enn_llm_generate( - r, prompt, img, NULL, 196, 0.0f, 0, 0, 0, on_token, NULL); -#else - std::string acc; - int rc = enn_llm_generate( - r, prompt, img, NULL, 196, 0.0f, 0, 0, 0, on_token, &acc); - printf("%s\n", acc.c_str()); + r, + prompt, + NULL, + img, + NULL, + 196, + 0.0f, + 0, + 0, + /*echo=*/1, + 0, + on_token, + &ctx); +#if !STREAMING + printf("%s\n", ctx.acc.c_str()); fflush(stdout); #endif printf("\n[rc=%d]\n", rc); + if (rc == ENN_OK && max_output_tokens > 0) { + bool pass = + ctx.num_tokens == max_output_tokens && ctx.tokens_after_stop == 0; + printf( + "[STOP TEST] %s: requested stop at %d, received %d, after stop %d\n", + pass ? "PASS" : "FAIL", + max_output_tokens, + ctx.num_tokens, + ctx.tokens_after_stop); + if (!pass) { + rc = ENN_ERR_GENERATE; + } + } return rc; } int main(int argc, char** argv) { if (argc < 6) { - fprintf(stderr, "usage: %s model ptd tok pos cfg [vision]\n", argv[0]); + fprintf( + stderr, + "usage: %s model ptd tok pos cfg [vision] [max_output_tokens]\n", + argv[0]); fprintf( stderr, "then read prompts from stdin; 'IMAGE:' for an image.\n"); + fprintf( + stderr, + "max_output_tokens=N (> 0): stop test, call enn_llm_stop() right after " + "the N-th output token and check nothing more arrives.\n"); return 2; } const char* vision = (argc > 6 && argv[6][0]) ? argv[6] : NULL; + const int32_t max_output_tokens = argc > 7 ? atoi(argv[7]) : 0; printf("[sample] creating runner...\n"); enn_runner* r = - enn_llm_create(argv[1], argv[2], argv[3], argv[4], argv[5], vision); + enn_llm_create(argv[1], argv[2], argv[3], argv[4], argv[5], vision, 0); if (!r) { fprintf(stderr, "enn_llm_create FAILED\n"); return 1; @@ -84,13 +153,17 @@ int main(int argc, char** argv) { printf("[sample] created. enter prompts (Ctrl-D to end):\n"); // Stopper thread that invokes enn_llm_stop to imitate a user interrupt. - std::thread stopper_thread([r]() { - std::cout << "[Stopper Thread] Waiting 3 seconds before interrupting..." - << std::endl; - std::this_thread::sleep_for(std::chrono::seconds(3)); - std::cout << "[Stopper Thread] Triggering enn_llm_stop()..." << std::endl; - enn_llm_stop(r); - }); + // Replaced by the token-count stop test when max_output_tokens is given. + std::thread stopper_thread; + if (max_output_tokens <= 0) { + stopper_thread = std::thread([r]() { + std::cout << "[Stopper Thread] Waiting 3 seconds before interrupting..." + << std::endl; + std::this_thread::sleep_for(std::chrono::seconds(3)); + std::cout << "[Stopper Thread] Triggering enn_llm_stop()..." << std::endl; + enn_llm_stop(r); + }); + } std::cout << "[CLI] Starting token generation loop..." << std::endl; int rc_total = 0; @@ -103,7 +176,7 @@ int main(int argc, char** argv) { if (n == 0) { continue; } - rc_total |= generate_one(r, buf); + rc_total |= generate_one(r, buf, max_output_tokens); } if (stopper_thread.joinable()) { diff --git a/examples/samsung/executor_runner/llm/runner/config.h b/examples/samsung/executor_runner/llm/runner/config.h index 98be7b42e0b..288a358d443 100644 --- a/examples/samsung/executor_runner/llm/runner/config.h +++ b/examples/samsung/executor_runner/llm/runner/config.h @@ -494,6 +494,9 @@ class RunnerConfig { // Stream generated tokens to stdout (in addition to the token callback). bool verbose = false; + // Empty => no system prompt. + std::string system_prompt; + // Vision inputs (Gemma 4 multimodal). Empty image_path => text-only. std::string image_path; std::string prompt_before_image; diff --git a/examples/samsung/executor_runner/llm/runner/io_manager.cpp b/examples/samsung/executor_runner/llm/runner/io_manager.cpp index fda72beeff5..1dfc426f11d 100644 --- a/examples/samsung/executor_runner/llm/runner/io_manager.cpp +++ b/examples/samsung/executor_runner/llm/runner/io_manager.cpp @@ -513,22 +513,15 @@ void LlmIOManager::update_local_attn_mask(size_t pos, size_t seq_len) { local_attn_mask_.data(), attn_mask_.data(), total_size * sizeof(float)); if (seq_len > 1) { - // Prefill phase with sliding window - if ((pos + seq_len) >= sliding_window) { - // Mask out positions before the sliding window - const size_t mask_start = sliding_window - seq_len; - const size_t mask_end = pos; // 512 - const size_t mask_len = mask_end - mask_start; - - if (mask_len > 0) { - for (size_t i = 0; i < rows; ++i) { - fill_float_region( - local_attn_mask_.data(), - i * cols + mask_start, - mask_len, - maskFalse); - } + // Prefill phase with sliding window: row i sits at position pos + i and + // sees the last sliding_window keys, so mask the oldest cached columns. + for (size_t i = 0; i < rows; ++i) { + const size_t query_pos = pos + i; + if (query_pos + 1 <= sliding_window) { + continue; } + const size_t mask_len = std::min(query_pos + 1 - sliding_window, pos); + fill_float_region(local_attn_mask_.data(), i * cols, mask_len, maskFalse); } } else { // Decode phase (single token) diff --git a/examples/samsung/executor_runner/llm/runner/runner.cpp b/examples/samsung/executor_runner/llm/runner/runner.cpp index 2650db17769..7f8555e0b71 100644 --- a/examples/samsung/executor_runner/llm/runner/runner.cpp +++ b/examples/samsung/executor_runner/llm/runner/runner.cpp @@ -56,13 +56,14 @@ std::string LLMRunner::get_formatted_prompt( break; case DecoderModelVersion::kGemma3: formatted_prompt.append("user\n"); - formatted_prompt.append(prompt); - formatted_prompt.append("\n"); - formatted_prompt.append("model\n"); + // Gemma 3 has no system role; the template prepends it to the user turn. if (!system_prompt.empty()) { formatted_prompt.append(system_prompt); - formatted_prompt.append("\n"); + formatted_prompt.append("\n\n"); } + formatted_prompt.append(prompt); + formatted_prompt.append("\n"); + formatted_prompt.append("model\n"); break; case DecoderModelVersion::kGemma4: if (!system_prompt.empty()) { @@ -241,14 +242,15 @@ runtime::Error LLMRunner::generate_from_pos( // First token time only measures the time it takes to encode the prompt and // return a response token. stats_->inference_start_ms = llm::time_in_ms(); - shouldStop_ = false; + text_token_generator_->clear_stop(); ET_LOG( Info, "################ input prompt ################\n %s\n", prompt.c_str()); - std::string input_prompt = get_formatted_prompt(prompt, "", model_version_); + std::string input_prompt = + get_formatted_prompt(prompt, config_.system_prompt, model_version_); if (model_version_ == DecoderModelVersion::kGemma3) { replace_all_newlines(input_prompt); } @@ -264,6 +266,14 @@ runtime::Error LLMRunner::generate_from_pos( // encode the (string) prompt into tokens sequence std::vector prompt_tokens = encode_res.get(); + // The Gemma 4 tokenizer's post-processor does not add BOS, so num_bos alone + // leaves the prompt without one. + if (model_version_ == DecoderModelVersion::kGemma4 && config.num_bos > 0 && + (prompt_tokens.empty() || + prompt_tokens.front() != tokenizer_->bos_tok())) { + prompt_tokens.insert(prompt_tokens.begin(), tokenizer_->bos_tok()); + } + // Gemma 4 multimodal: run the vision encoder and inject the image token // block. The user turn becomes: // <|turn>user\n xN @@ -503,6 +513,7 @@ runtime::Error LLMRunner::warmup( } void LLMRunner::stop() { + // Not gated on is_loaded(): it checks method "forward", absent here. if (text_token_generator_ != nullptr) { text_token_generator_->stop(); } diff --git a/examples/samsung/executor_runner/llm/runner/runner.h b/examples/samsung/executor_runner/llm/runner/runner.h index 4bf26ee76fb..0feeb0f523a 100644 --- a/examples/samsung/executor_runner/llm/runner/runner.h +++ b/examples/samsung/executor_runner/llm/runner/runner.h @@ -109,8 +109,6 @@ class LLMRunner : public llm::IRunner { void reset(); private: - bool shouldStop_{false}; - RunnerConfig config_; std::unique_ptr<::tokenizers::Tokenizer> tokenizer_; std::unordered_map metadata_; diff --git a/examples/samsung/executor_runner/llm/runner/token_generator.h b/examples/samsung/executor_runner/llm/runner/token_generator.h index a44e781227f..ce81f3413bd 100644 --- a/examples/samsung/executor_runner/llm/runner/token_generator.h +++ b/examples/samsung/executor_runner/llm/runner/token_generator.h @@ -14,6 +14,8 @@ #include #include +#include + namespace example { namespace llm = ::executorch::extension::llm; @@ -82,10 +84,12 @@ class ET_EXPERIMENTAL TokenGenerator { auto tokens_managed = extension::from_blob( token_data.data(), token_shape, aten::ScalarType::Long); - should_stop_ = false; - // Generate our tokens while (pos < start_pos + max_new_tokens) { + if (should_stop_.load()) { + break; + } + // Run the model auto logits_res = text_decoder_runner_->step(tokens_managed, pos); @@ -116,10 +120,6 @@ class ET_EXPERIMENTAL TokenGenerator { token_callback( ET_UNWRAP_TOKENIZER(tokenizer_->decode(prev_token, cur_token))); - if (should_stop_) { - break; - } - // data-dependent terminating condition: we have n_eos_ number of EOS if (eos_ids_->find(cur_token) != eos_ids_->end()) { // ET_LOG(Info, "Reached to the end of generation"); @@ -133,7 +133,11 @@ class ET_EXPERIMENTAL TokenGenerator { * Stop the generation loop. */ inline void stop() { - should_stop_ = true; + should_stop_.store(true); + } + + inline void clear_stop() { + should_stop_.store(false); } /** @@ -167,7 +171,7 @@ class ET_EXPERIMENTAL TokenGenerator { bool use_kv_cache_; // state machine - bool should_stop_ = false; + std::atomic should_stop_{false}; // stats llm::Stats* stats_; From 790684faf42d023f9e845a93f7fdef11fa2c0b82 Mon Sep 17 00:00:00 2001 From: Jiseong Oh Date: Tue, 13 Oct 2026 19:39:55 +0900 Subject: [PATCH 16/17] Add raw_prompt and assistant_prefill_prompt to the LLM runner Let callers feed a preformatted prompt verbatim or prefill the start of the reply. Signed-off-by: Jiseong Oh Co-authored-by: MJ Kim --- .../enn_executor_llm_runner.cpp | 37 ++- .../executor_runner/exynos_llm_wrapper.cpp | 5 + .../executor_runner/exynos_llm_wrapper.h | 14 +- .../exynos_llm_wrapper_sample_main.cpp | 273 ++++++++++++++---- .../executor_runner/llm/runner/config.h | 9 + .../executor_runner/llm/runner/runner.cpp | 23 +- .../executor_runner/llm/runner/runner.h | 1 + 7 files changed, 292 insertions(+), 70 deletions(-) diff --git a/examples/samsung/executor_runner/enn_executor_llm_runner.cpp b/examples/samsung/executor_runner/enn_executor_llm_runner.cpp index 37219efcc41..5d5a21569a0 100644 --- a/examples/samsung/executor_runner/enn_executor_llm_runner.cpp +++ b/examples/samsung/executor_runner/enn_executor_llm_runner.cpp @@ -48,6 +48,12 @@ DEFINE_string(prompt, "The answer to the ultimate question is", "Prompt."); DEFINE_string(system_prompt, "", "System prompt. Empty means none."); +DEFINE_string( + assistant_prefill_prompt, + "", + "Assistant pre-fill appended verbatim to the very end of the rendered " + "prompt (after the assistant turn header). Ignored with --raw_prompt."); + DEFINE_string( prompt_file, "", @@ -58,6 +64,17 @@ DEFINE_string( "", "Read the system prompt from this file instead of --system_prompt."); +// Raw mode: an already-formatted prompt (special tokens, system turn, tool +// declaration) fed verbatim, skipping chat-template wrapping. Supply it inline +// with --raw_prompt or from a file with --raw_prompt_file; either overrides +// --prompt / --system_prompt. +DEFINE_string(raw_prompt, "", "Inline raw (pre-formatted) prompt."); + +DEFINE_string( + raw_prompt_file, + "", + "Read the raw (pre-formatted) prompt from this file instead of --raw_prompt."); + DEFINE_double( temperature, 0.8, @@ -136,15 +153,23 @@ int32_t main(int32_t argc, char** argv) { if ((!FLAGS_prompt_file.empty() && !read_file(FLAGS_prompt_file, &FLAGS_prompt)) || (!FLAGS_system_prompt_file.empty() && - !read_file(FLAGS_system_prompt_file, &FLAGS_system_prompt))) { + !read_file(FLAGS_system_prompt_file, &FLAGS_system_prompt)) || + (!FLAGS_raw_prompt_file.empty() && + !read_file(FLAGS_raw_prompt_file, &FLAGS_raw_prompt))) { return ENN_ERR_INVALID_ARG; } - if (FLAGS_prompt.empty()) { + // A non-empty raw prompt (inline or from file) is fed verbatim and overrides + // the templated --prompt / --system_prompt path. + const bool raw_mode = !FLAGS_raw_prompt.empty(); + const std::string& effective_prompt = + raw_mode ? FLAGS_raw_prompt : FLAGS_prompt; + + if (effective_prompt.empty()) { ET_LOG(Error, "Input prompt cannot be empty."); return ENN_ERR_INVALID_ARG; } - ET_LOG(Info, "Input prompt message: %s", FLAGS_prompt.c_str()); + ET_LOG(Info, "Input prompt message: %s", effective_prompt.c_str()); #if defined(ET_USE_THREADPOOL) uint32_t num_performant_cores = FLAGS_cpu_threads == -1 @@ -174,10 +199,12 @@ int32_t main(int32_t argc, char** argv) { StopContext stop_ctx{runner, FLAGS_max_token_limit}; int32_t error = enn_llm_generate( runner, - FLAGS_prompt.c_str(), - FLAGS_system_prompt.c_str(), + effective_prompt.c_str(), + raw_mode ? "" : FLAGS_system_prompt.c_str(), + raw_mode ? "" : FLAGS_assistant_prefill_prompt.c_str(), FLAGS_image_path.c_str(), FLAGS_prompt_before_image.c_str(), + /*raw_prompt=*/raw_mode ? 1 : 0, FLAGS_seq_len, static_cast(FLAGS_temperature), FLAGS_num_bos, diff --git a/examples/samsung/executor_runner/exynos_llm_wrapper.cpp b/examples/samsung/executor_runner/exynos_llm_wrapper.cpp index 2982d4d484e..2eab0620402 100644 --- a/examples/samsung/executor_runner/exynos_llm_wrapper.cpp +++ b/examples/samsung/executor_runner/exynos_llm_wrapper.cpp @@ -88,8 +88,10 @@ int32_t enn_llm_generate( enn_runner* runner, const char* prompt, const char* system_prompt, + const char* assistant_prefill_prompt, const char* image_path, const char* prompt_before_image, + int32_t raw_prompt, int32_t seq_len, float temperature, int32_t num_bos, @@ -115,8 +117,11 @@ int32_t enn_llm_generate( // System prompt and image inputs are threaded through the runner's config // (read by generate). lrunner->get_config().system_prompt = str_or_empty(system_prompt); + lrunner->get_config().assistant_prefill_prompt = + str_or_empty(assistant_prefill_prompt); lrunner->get_config().image_path = str_or_empty(image_path); lrunner->get_config().prompt_before_image = str_or_empty(prompt_before_image); + lrunner->get_config().raw_prompt = raw_prompt != 0; lrunner->get_config().verbose = verbose != 0; llm::GenerationConfig gen_cfg{ diff --git a/examples/samsung/executor_runner/exynos_llm_wrapper.h b/examples/samsung/executor_runner/exynos_llm_wrapper.h index d9764de93bc..ce91a8391c6 100644 --- a/examples/samsung/executor_runner/exynos_llm_wrapper.h +++ b/examples/samsung/executor_runner/exynos_llm_wrapper.h @@ -90,11 +90,19 @@ EXYNOS_LLM_WRAPPER_API void enn_llm_stop(enn_runner* runner); * prompt starts fresh. * * @param runner Handle from enn_llm_create. - * @param prompt User prompt text (required, non-empty). - * @param system_prompt System prompt text, or NULL/"" => none. + * @param prompt User prompt text (required, non-empty). When raw_prompt is + * non-zero, this is the already-formatted prompt (BOS, turn/tool special + * tokens, system turn) and is fed to the model verbatim. + * @param system_prompt System prompt text, or NULL/"" => none. Ignored when + * raw_prompt is non-zero. + * @param assistant_prefill_prompt Text appended verbatim to the very end of the + * rendered prompt, after the assistant turn header (a pre-filled start of the + * model's reply), or NULL/"" => none. Ignored when raw_prompt is non-zero. * @param image_path Image file for this turn, or NULL => text-only. * @param prompt_before_image Text placed before the image in the user turn * (Gemma 4 multimodal); ignored when image_path is NULL. + * @param raw_prompt Non-zero => prompt is already fully formatted; skip the + * internal chat-template wrapping and system_prompt. * @param seq_len Total token budget (prompt + generated), capped by the model's * max context length; -1 => use the model's max context length. * @param temperature Sampling temperature. @@ -111,8 +119,10 @@ EXYNOS_LLM_WRAPPER_API int32_t enn_llm_generate( enn_runner* runner, const char* prompt, const char* system_prompt, + const char* assistant_prefill_prompt, const char* image_path, const char* prompt_before_image, + int32_t raw_prompt, int32_t seq_len, float temperature, int32_t num_bos, diff --git a/examples/samsung/executor_runner/exynos_llm_wrapper_sample_main.cpp b/examples/samsung/executor_runner/exynos_llm_wrapper_sample_main.cpp index ad2c9140966..b5fc1cea27a 100644 --- a/examples/samsung/executor_runner/exynos_llm_wrapper_sample_main.cpp +++ b/examples/samsung/executor_runner/exynos_llm_wrapper_sample_main.cpp @@ -9,26 +9,42 @@ * * Usage (paths are on-device): * exynos_llm_wrapper_sample [vision] - * [max_output_tokens] + * [--raw_prompt ] [--raw_prompt_file ] + * [--system_prompt ] [--system_prompt_file ] + * [--user_prompt ] [--user_prompt_file ] + * [--assistant_prefill_prompt ] [--assistant_prefill_prompt_file + * ] + * [--image_file ] + * [--max_output_tokens ] [--stop_after_sec ] [--echo] * - * Reads prompts from stdin (one per line) and generates on the SAME handle, - * so the model stays prepared across lines. Each line may be prefixed with - * "IMAGE:" to feed an image followed by the rest as the text prompt. + * Each generation is one en-invocation: + * - `--raw_prompt ` / `--raw_prompt_file ` : feed verbatim + * (raw_prompt=1). + * - `--*_prompt ` : the next argument IS the text itself. + * - `--*_prompt_file `: read the file and use its contents. + * - `--image_file ` : the value IS the image path passed as image_path. + * The system/user/prefill texts are fed together (raw_prompt=0) as one + * generate call, along with any image. Any may be omitted. Repeating runs + * sequentially on the same handle, so the model stays prepared across calls. * * STREAMING=1 (default): the on_token callback prints each token as generated. * STREAMING=0: the on_token callback appends to a std::string and the full * answer is printed once when generation ends. Override with -DSTREAMING=0. * - * Stop test: pass max_output_tokens=N (> 0; use "" for vision to skip it) to - * call enn_llm_stop() right after the N-th output token of each prompt and - * check that no more tokens arrive. Use prompts whose answer is longer than N. - * Without it, a stopper thread calls enn_llm_stop() once after 3 seconds. + * Stop test (exercises enn_llm_stop(); both apply to every generation): + * --max_output_tokens N (> 0): call enn_llm_stop() right after the N-th + * output token of each generation and check that no more tokens arrive. + * Use prompts whose answer is longer than N. Prints PASS/FAIL per run. + * --stop_after_sec N (> 0): spawn a stopper thread that calls enn_llm_stop() + * N seconds into each generation, imitating a user interrupt. */ +#include #include #include #include #include -#include +#include +#include #include #include @@ -38,9 +54,27 @@ #define STREAMING 1 #endif +// Option codes for the long-option table (start past ASCII so getopt_long +// distinguishes them from any single-char option). +enum { + OPT_RAW = 256, + OPT_RAW_FILE, + OPT_SYS, + OPT_SYS_FILE, + OPT_USR, + OPT_USR_FILE, + OPT_PRE, + OPT_PRE_FILE, + OPT_IMG_FILE, + OPT_MAX_TOKENS, + OPT_STOP_SEC, + OPT_ECHO, +}; + struct GenContext { enn_runner* runner; int32_t max_output_tokens; + bool echo; int32_t num_pieces = 0; int32_t num_tokens = 0; int32_t tokens_after_stop = 0; @@ -49,8 +83,9 @@ struct GenContext { static void count_for_stop_test(GenContext* ctx, const char* text, int32_t len) { - // The first callback is the echoed input prompt, not an output token. - if (ctx->num_pieces++ == 0) { + // When echo is on, the first callback is the echoed input prompt, not an + // output token, so skip it. With echo off every piece is an output token. + if (ctx->echo && ctx->num_pieces++ == 0) { return; } // The runner emits a trailing "\n" after generation ends; ignore it. @@ -80,36 +115,76 @@ static void on_token(const char* text, int32_t len, void* userdata) { } } -static int -generate_one(enn_runner* r, const char* line, int32_t max_output_tokens) { - const char* img = NULL; - const char* prompt = line; - if (strncmp(line, "IMAGE:", 6) == 0) { - img = line + 6; - prompt = "Describe this image."; +static bool read_file(const char* path, std::string* out) { + std::ifstream f(path); + if (!f.is_open()) { + fprintf(stderr, "failed to open %s\n", path); + return false; + } + std::stringstream ss; + ss << f.rdbuf(); + *out = ss.str(); + return true; +} + +static int run_one( + enn_runner* r, + const std::string& raw, + const std::string& sys, + const std::string& usr, + const std::string& pre, + const std::string& image, + int32_t max_output_tokens, + int32_t stop_after_sec, + bool echo) { + int32_t raw_prompt = !raw.empty(); + const char* prompt = raw_prompt ? raw.c_str() : usr.c_str(); + // raw prompts already contain the leading token; the templated path + // must add it so both modes feed the model the same prompt. + int32_t num_bos = raw_prompt ? 0 : 1; + + GenContext ctx{r, max_output_tokens, echo}; + + // Timer-based stopper thread imitating a user interrupt. Replaced by the + // token-count stop test when max_output_tokens is given. + std::thread stopper; + if (stop_after_sec > 0) { + stopper = std::thread([r, stop_after_sec]() { + std::this_thread::sleep_for(std::chrono::seconds(stop_after_sec)); + printf( + "\n[Stopper] triggering enn_llm_stop() after %d s\n", stop_after_sec); + fflush(stdout); + enn_llm_stop(r); + }); } - printf("----- generate (image=%s) -----\n", img ? img : "none"); - GenContext ctx{r, max_output_tokens}; int rc = enn_llm_generate( r, prompt, + raw_prompt ? NULL : (sys.empty() ? NULL : sys.c_str()), + raw_prompt ? NULL : (pre.empty() ? NULL : pre.c_str()), + image.empty() ? NULL : image.c_str(), NULL, - img, - NULL, - 196, + raw_prompt, + 3840, 0.0f, + num_bos, 0, - 0, - /*echo=*/1, + echo ? 1 : 0, 0, on_token, &ctx); + + if (stopper.joinable()) { + stopper.join(); + } + #if !STREAMING printf("%s\n", ctx.acc.c_str()); fflush(stdout); #endif printf("\n[rc=%d]\n", rc); + if (rc == ENN_OK && max_output_tokens > 0) { bool pass = ctx.num_tokens == max_output_tokens && ctx.tokens_after_stop == 0; @@ -130,18 +205,35 @@ int main(int argc, char** argv) { if (argc < 6) { fprintf( stderr, - "usage: %s model ptd tok pos cfg [vision] [max_output_tokens]\n", + "usage: %s model ptd tok pos cfg [vision] " + "[--raw_prompt STR | --raw_prompt_file F | " + "--system_prompt STR | --system_prompt_file F | " + "--user_prompt STR | --user_prompt_file F | " + "--assistant_prefill_prompt STR | --assistant_prefill_prompt_file F | " + "--image_file PATH]* " + "[--max_output_tokens N] [--stop_after_sec N] [--echo]\n", argv[0]); - fprintf( - stderr, "then read prompts from stdin; 'IMAGE:' for an image.\n"); - fprintf( - stderr, - "max_output_tokens=N (> 0): stop test, call enn_llm_stop() right after " - "the N-th output token and check nothing more arrives.\n"); return 2; } - const char* vision = (argc > 6 && argv[6][0]) ? argv[6] : NULL; - const int32_t max_output_tokens = argc > 7 ? atoi(argv[7]) : 0; + const char* vision = + (argc > 6 && argv[6][0] && argv[6][0] != '-') ? argv[6] : NULL; + int idx = (argc > 6 && argv[6][0] && argv[6][0] != '-') ? 7 : 6; + + static const struct option long_opts[] = { + {"raw_prompt", required_argument, 0, OPT_RAW}, + {"raw_prompt_file", required_argument, 0, OPT_RAW_FILE}, + {"system_prompt", required_argument, 0, OPT_SYS}, + {"system_prompt_file", required_argument, 0, OPT_SYS_FILE}, + {"user_prompt", required_argument, 0, OPT_USR}, + {"user_prompt_file", required_argument, 0, OPT_USR_FILE}, + {"assistant_prefill_prompt", required_argument, 0, OPT_PRE}, + {"assistant_prefill_prompt_file", required_argument, 0, OPT_PRE_FILE}, + {"image_file", required_argument, 0, OPT_IMG_FILE}, + {"max_output_tokens", required_argument, 0, OPT_MAX_TOKENS}, + {"stop_after_sec", required_argument, 0, OPT_STOP_SEC}, + {"echo", no_argument, 0, OPT_ECHO}, + {0, 0, 0, 0}, + }; printf("[sample] creating runner...\n"); enn_runner* r = @@ -150,37 +242,98 @@ int main(int argc, char** argv) { fprintf(stderr, "enn_llm_create FAILED\n"); return 1; } - printf("[sample] created. enter prompts (Ctrl-D to end):\n"); - - // Stopper thread that invokes enn_llm_stop to imitate a user interrupt. - // Replaced by the token-count stop test when max_output_tokens is given. - std::thread stopper_thread; - if (max_output_tokens <= 0) { - stopper_thread = std::thread([r]() { - std::cout << "[Stopper Thread] Waiting 3 seconds before interrupting..." - << std::endl; - std::this_thread::sleep_for(std::chrono::seconds(3)); - std::cout << "[Stopper Thread] Triggering enn_llm_stop()..." << std::endl; - enn_llm_stop(r); - }); - } - std::cout << "[CLI] Starting token generation loop..." << std::endl; + // Options accumulate into one generation spec. A spec is either a raw + // (standalone) or the system/user/prefill triple combined into ONE generate + // call; each may come as a bare string or a file. Flush the accumulated spec + // only when a new raw begins (a raw group never shares a call with others). + // --max_output_tokens / --stop_after_sec / --echo are global and apply to + // every generation. + std::string raw, sys, usr, pre, image; + bool have = false; int rc_total = 0; - char buf[4096]; - while (fgets(buf, sizeof(buf), stdin)) { - size_t n = strlen(buf); - while (n > 0 && (buf[n - 1] == '\n' || buf[n - 1] == '\r')) { - buf[--n] = '\0'; + int32_t max_output_tokens = 0; + int32_t stop_after_sec = 0; + bool echo = false; + + optind = idx; // resume scanning after the positional model arguments + int opt; + while ((opt = getopt_long(argc, argv, "", long_opts, NULL)) != -1) { + const bool is_file = opt == OPT_RAW_FILE || opt == OPT_SYS_FILE || + opt == OPT_USR_FILE || opt == OPT_PRE_FILE; + std::string content; + if (is_file) { + if (!read_file(optarg, &content)) { + rc_total = 1; + break; + } + } else if (opt != OPT_ECHO) { + content = optarg; } - if (n == 0) { - continue; + + switch (opt) { + case OPT_RAW: + case OPT_RAW_FILE: + if (have) { + rc_total |= run_one( + r, + raw, + sys, + usr, + pre, + image, + max_output_tokens, + stop_after_sec, + echo); + raw.clear(); + sys.clear(); + usr.clear(); + pre.clear(); + image.clear(); + have = false; + } + raw = content; + have = true; + break; + case OPT_SYS: + case OPT_SYS_FILE: + sys = content; + have = true; + break; + case OPT_USR: + case OPT_USR_FILE: + usr = content; + have = true; + break; + case OPT_PRE: + case OPT_PRE_FILE: + pre = content; + have = true; + break; + case OPT_IMG_FILE: + image = content; + have = true; + break; + case OPT_MAX_TOKENS: + max_output_tokens = atoi(content.c_str()); + break; + case OPT_STOP_SEC: + stop_after_sec = atoi(content.c_str()); + break; + case OPT_ECHO: + echo = true; + break; + default: + fprintf(stderr, "unknown argument: %s\n", argv[optind - 1]); + rc_total = 1; + break; } - rc_total |= generate_one(r, buf, max_output_tokens); + if (rc_total) + break; } - - if (stopper_thread.joinable()) { - stopper_thread.join(); + if (have && !rc_total) { + rc_total |= run_one( + r, raw, sys, usr, pre, image, max_output_tokens, stop_after_sec, echo); } printf("[sample] destroying runner...\n"); diff --git a/examples/samsung/executor_runner/llm/runner/config.h b/examples/samsung/executor_runner/llm/runner/config.h index 288a358d443..a1ab48a6f4e 100644 --- a/examples/samsung/executor_runner/llm/runner/config.h +++ b/examples/samsung/executor_runner/llm/runner/config.h @@ -497,6 +497,15 @@ class RunnerConfig { // Empty => no system prompt. std::string system_prompt; + // Text appended verbatim to the very end of the rendered prompt (after the + // assistant turn header), i.e. a pre-filled assistant response. Ignored in + // raw_prompt mode where the caller supplies the full prompt already. + std::string assistant_prefill_prompt; + + // Treat the prompt as already fully formatted: skip chat-template wrapping + // and system_prompt so a caller can feed raw special-token input verbatim. + bool raw_prompt = false; + // Vision inputs (Gemma 4 multimodal). Empty image_path => text-only. std::string image_path; std::string prompt_before_image; diff --git a/examples/samsung/executor_runner/llm/runner/runner.cpp b/examples/samsung/executor_runner/llm/runner/runner.cpp index 7f8555e0b71..807ef910911 100644 --- a/examples/samsung/executor_runner/llm/runner/runner.cpp +++ b/examples/samsung/executor_runner/llm/runner/runner.cpp @@ -39,6 +39,7 @@ using json = nlohmann::json; std::string LLMRunner::get_formatted_prompt( const std::string& prompt, const std::string& system_prompt, + const std::string& assistant_prefill_prompt, DecoderModelVersion version) { std::string formatted_prompt; switch (version) { @@ -53,6 +54,9 @@ std::string LLMRunner::get_formatted_prompt( formatted_prompt.append(prompt); formatted_prompt.append( "<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n"); + // Pre-fill: text placed right after the assistant header as the start of + // the model's reply. + formatted_prompt.append(assistant_prefill_prompt); break; case DecoderModelVersion::kGemma3: formatted_prompt.append("user\n"); @@ -64,6 +68,8 @@ std::string LLMRunner::get_formatted_prompt( formatted_prompt.append(prompt); formatted_prompt.append("\n"); formatted_prompt.append("model\n"); + // Pre-fill: text placed right after the model turn header. + formatted_prompt.append(assistant_prefill_prompt); break; case DecoderModelVersion::kGemma4: if (!system_prompt.empty()) { @@ -75,12 +81,16 @@ std::string LLMRunner::get_formatted_prompt( formatted_prompt.append(prompt); formatted_prompt.append("\n"); formatted_prompt.append("<|turn>model\n"); + // Pre-fill: text placed right after the model turn header. + formatted_prompt.append(assistant_prefill_prompt); break; case DecoderModelVersion::kGeneral: formatted_prompt.append(prompt); + formatted_prompt.append(assistant_prefill_prompt); break; default: formatted_prompt.append(prompt); + formatted_prompt.append(assistant_prefill_prompt); break; } return formatted_prompt; @@ -249,9 +259,16 @@ runtime::Error LLMRunner::generate_from_pos( "################ input prompt ################\n %s\n", prompt.c_str()); - std::string input_prompt = - get_formatted_prompt(prompt, config_.system_prompt, model_version_); - if (model_version_ == DecoderModelVersion::kGemma3) { + // Raw prompts are already fully formatted (special tokens, system turn, tool + // declaration); feed them verbatim and skip every template transformation. + std::string input_prompt = config_.raw_prompt + ? prompt + : get_formatted_prompt( + prompt, + config_.system_prompt, + config_.assistant_prefill_prompt, + model_version_); + if (!config_.raw_prompt && model_version_ == DecoderModelVersion::kGemma3) { replace_all_newlines(input_prompt); } diff --git a/examples/samsung/executor_runner/llm/runner/runner.h b/examples/samsung/executor_runner/llm/runner/runner.h index 0feeb0f523a..e823c13cf86 100644 --- a/examples/samsung/executor_runner/llm/runner/runner.h +++ b/examples/samsung/executor_runner/llm/runner/runner.h @@ -104,6 +104,7 @@ class LLMRunner : public llm::IRunner { static std::string get_formatted_prompt( const std::string& prompt, const std::string& system_prompt, + const std::string& assistant_prefill_prompt, DecoderModelVersion version); void reset(); From a22d8e2d59f8b6f1c4c250096e0c81fa6a063fd8 Mon Sep 17 00:00:00 2001 From: Jiseong Oh Date: Tue, 13 Oct 2026 19:40:08 +0900 Subject: [PATCH 17/17] Refactor ENN backend and decoder for stability and memory safety - Fix memory corruption and implement refcounting in enn_backend.cpp - Resolve buffer overflow issues in IO manager and RoPE embeddings - Fix Tensor lifetime and model loading logic in decoder.cpp - Improve config validation and optimize build system for LLM runner - General cleanup: remove debug logs and fix CLI bugs Signed-off-by: Jiseong Oh --- .../runtime/enn_api_implementation.cpp | 14 +- backends/samsung/runtime/enn_backend.cpp | 154 +- backends/samsung/runtime/enn_executor.cpp | 14 +- backends/samsung/runtime/enn_executor.h | 9 +- examples/samsung/CMakeLists.txt | 85 +- examples/samsung/build.sh | 16 +- examples/samsung/executor_runner/README.md | 30 +- .../configs}/exynos_gemma3_1b.json | 2 +- .../configs}/exynos_gemma4_e2b.json | 2 +- .../configs}/exynos_llama3_1b.json | 2 +- .../enn_executor_llm_runner.cpp | 23 +- .../executor_runner/exynos_llm_wrapper.cpp | 3 + .../executor_runner/exynos_llm_wrapper.h | 15 +- .../exynos_llm_wrapper_sample_main.cpp | 70 +- .../executor_runner/llm/runner/CMakeLists.txt | 44 +- .../executor_runner/llm/runner/config.cpp | 96 +- .../executor_runner/llm/runner/config.h | 141 +- .../executor_runner/llm/runner/decoder.cpp | 394 +- .../executor_runner/llm/runner/decoder.h | 7 + .../executor_runner/llm/runner/io_manager.cpp | 202 +- .../executor_runner/llm/runner/io_manager.h | 68 +- .../executor_runner/llm/runner/prefiller.cpp | 5 +- .../llm/runner/rotary_embedding.cpp | 23 +- .../executor_runner/llm/runner/runner.cpp | 61 +- .../executor_runner/llm/runner/runner.h | 12 +- .../llm/runner/token_generator.h | 6 +- .../executor_runner/run_gemma3_sample.sh | 1 + .../executor_runner/run_llama3_sample.sh | 4 +- .../scripts/gemma4_pos_embedding.py} | 10 +- extension/llm/runner/constants.h | 1 - extension/llm/runner/text_decoder_runner.cpp | 18 +- third-party/stb/stb_image.h | 7988 ----------------- third-party/stb/stb_image_resize.h | 2634 ------ 33 files changed, 645 insertions(+), 11509 deletions(-) rename examples/{models/gemma3_1b => samsung/executor_runner/configs}/exynos_gemma3_1b.json (99%) rename examples/{models/gemma4/config => samsung/executor_runner/configs}/exynos_gemma4_e2b.json (99%) rename examples/{models/llama3_1b => samsung/executor_runner/configs}/exynos_llama3_1b.json (99%) mode change 100644 => 100755 examples/samsung/executor_runner/run_gemma3_sample.sh mode change 100644 => 100755 examples/samsung/executor_runner/run_llama3_sample.sh rename examples/{models/gemma4/text_decoder/generate_pos_embedding_bin.py => samsung/scripts/gemma4_pos_embedding.py} (89%) delete mode 100644 third-party/stb/stb_image.h delete mode 100644 third-party/stb/stb_image_resize.h diff --git a/backends/samsung/runtime/enn_api_implementation.cpp b/backends/samsung/runtime/enn_api_implementation.cpp index 35e866e5741..560d6b6a1d4 100644 --- a/backends/samsung/runtime/enn_api_implementation.cpp +++ b/backends/samsung/runtime/enn_api_implementation.cpp @@ -123,8 +123,9 @@ void setPerfConfigFromSoc(const EnnApi* enn_api) { if (config.config_id == kPerfConfigIdDefault) { return; } - if (enn_api->EnnSetPreferencePerfConfigId(config.config_id) != - ENN_RET_SUCCESS) { + if (enn_api->EnnSetPreferencePerfConfigId == nullptr || + enn_api->EnnSetPreferencePerfConfigId(config.config_id) != + ENN_RET_SUCCESS) { ENN_LOG_WARN( "Failed to set ENN perf config id %u, falling back to PERFORMANCE mode.", config.config_id); @@ -175,7 +176,6 @@ Error EnnApi::loadApiLib() { ENN_LOAD_API_FUNC(libenn_public_api_, EnnInitialize, this); ENN_LOAD_API_FUNC(libenn_public_api_, EnnSetPreferencePerfMode, this); ENN_LOAD_API_FUNC(libenn_public_api_, EnnGetPreferencePerfMode, this); - ENN_LOAD_API_FUNC(libenn_public_api_, EnnSetPreferencePerfConfigId, this); ENN_LOAD_API_FUNC(libenn_public_api_, EnnOpenModel, this); ENN_LOAD_API_FUNC(libenn_public_api_, EnnOpenModelFromMemory, this); ENN_LOAD_API_FUNC(libenn_public_api_, EnnSetFastIpc, this); @@ -198,7 +198,13 @@ Error EnnApi::loadApiLib() { ENN_LOAD_API_FUNC( libenn_public_api_, EnnGetFileDescriptorFromEnnBuffer, this); ENN_LOAD_API_FUNC(libenn_public_api_, EnnOpenModelFromFd, this); - ENN_LOAD_API_FUNC(libenn_public_api_, EnnOpenModelFromFdWithWeight, this); + // Optional: only newer ENN releases export these. + EnnSetPreferencePerfConfigId = + reinterpret_cast(loadApiFunction( + libenn_public_api_, "EnnSetPreferencePerfConfigId", true)); + EnnOpenModelFromFdWithWeight = + reinterpret_cast(loadApiFunction( + libenn_public_api_, "EnnOpenModelFromFdWithWeight", true)); return Error::Ok; } diff --git a/backends/samsung/runtime/enn_backend.cpp b/backends/samsung/runtime/enn_backend.cpp index 98899f74255..0054926da7d 100644 --- a/backends/samsung/runtime/enn_backend.cpp +++ b/backends/samsung/runtime/enn_backend.cpp @@ -16,8 +16,11 @@ #include #include +#include +#include #include -#include +#include +#include #include #include #include @@ -31,8 +34,47 @@ using executorch::runtime::NamedDataMap; const char ennWeightBufferKey[] = "enn_weight_buffer_hash_id"; -// Weight cache LUT for reusing FreeableBuffer across multiple init calls -static std::unordered_map g_weight_cache_lut; +namespace { +// Methods sharing weights (e.g. prefill and decode) reference the same buffer, +// so it is refcounted and released when the last delegate using it is +// destroyed. +struct SharedWeight { + FreeableBuffer buffer; + int32_t refs = 0; +}; + +std::mutex g_weight_cache_mutex; +std::unordered_map g_weight_cache; + +Result acquire_weight( + const NamedDataMap* named_data_map, + const std::string& key) { + std::lock_guard lock(g_weight_cache_mutex); + auto it = g_weight_cache.find(key); + if (it == g_weight_cache.end()) { + ET_CHECK_OR_RETURN_ERROR( + named_data_map != nullptr, + InvalidArgument, + "Weight-sharing delegate needs a NamedDataMap (.ptd) for key %s", + key.c_str()); + Result loaded = named_data_map->get_data(key.c_str()); + ET_CHECK_OK_OR_RETURN_ERROR( + loaded.error(), "Failed to load shared weight %s", key.c_str()); + it = g_weight_cache.emplace(key, SharedWeight{std::move(loaded.get()), 0}) + .first; + } + it->second.refs++; + return &it->second.buffer; +} + +void release_weight(const std::string& key) { + std::lock_guard lock(g_weight_cache_mutex); + auto it = g_weight_cache.find(key); + if (it != g_weight_cache.end() && --it->second.refs == 0) { + g_weight_cache.erase(it); + } +} +} // namespace class EnnBackend final : public PyTorchBackendInterface { public: @@ -54,10 +96,14 @@ class EnnBackend final : public PyTorchBackendInterface { EXYNOS_ATRACE_NAME_LINE("backend_init"); MemoryAllocator* runtime_allocator = context.get_runtime_allocator(); auto executor = runtime_allocator->allocateInstance(); + ET_CHECK_OR_RETURN_ERROR( + executor != nullptr, + MemoryAllocationFailed, + "Failed to allocate EnnExecutor"); + new (executor) enn::EnnExecutor(); const char* binary_buf_addr = reinterpret_cast(processed->data()); size_t buf_size = processed->size(); - const NamedDataMap* named_data_map = context.get_named_data_map(); auto it = std::find_if( compile_specs.begin(), @@ -66,70 +112,35 @@ class EnnBackend final : public PyTorchBackendInterface { return std::strcmp(spec.key, ennWeightBufferKey) == 0; }); - const char* weight_buf_addr; - size_t weight_size; - + Error err = Error::Ok; if (it != compile_specs.end()) { std::string weight_key = decode_spec_value(*it); - - // Check if weight buffer is already cached - FreeableBuffer* wgt = nullptr; - auto cache_it = g_weight_cache_lut.find(weight_key); - - if (cache_it != g_weight_cache_lut.end()) { - weight_buf_addr = - reinterpret_cast((cache_it->second)->data()); - weight_size = (cache_it->second)->size(); - wgt = cache_it->second; - - ET_LOG( - Info, - "EnnBackend, [%s] Weight Buffer from Cache for %s size(%ld), addr(%p)", - context.get_method_name(), - weight_key.c_str(), - weight_size, - weight_buf_addr); - - // For cached buffers, we don't own it, so don't store weight_key - } else { - // Load weight from named_data_map and cache it - Result weight = - named_data_map->get_data(weight_key.c_str()); - - weight_buf_addr = reinterpret_cast(weight->data()); - weight_size = weight->size(); - - ET_LOG( - Info, - "EnnBackend, [%s] Weight Buffer Loaded for %s size(%ld), addr(%p)", - context.get_method_name(), - weight_key.c_str(), - weight_size, - weight_buf_addr); - - wgt = new FreeableBuffer(std::move(weight.get())); - - // Cache the weight buffer - g_weight_cache_lut[weight_key] = wgt; - - // Store the weight_key for cleanup in destroy - executor->set_weight_key(weight_key); - } - - Error err = executor->initialize( - binary_buf_addr, buf_size, weight_buf_addr, weight_size, true); - if (err != Error::Ok) { - ENN_LOG_ERROR("Exynos backend initialize failed."); + auto weight = acquire_weight(context.get_named_data_map(), weight_key); + if (!weight.ok()) { executor->~EnnExecutor(); - return err; + return weight.error(); } + executor->set_weight_key(weight_key); + ET_LOG( + Info, + "EnnBackend, [%s] shared weight %s size(%zu), addr(%p)", + context.get_method_name(), + weight_key.c_str(), + weight.get()->size(), + weight.get()->data()); + err = executor->initialize( + binary_buf_addr, + buf_size, + reinterpret_cast(weight.get()->data()), + weight.get()->size(), + true); } else { - Error err = executor->initialize(binary_buf_addr, buf_size); - if (err != Error::Ok) { - ENN_LOG_ERROR("Exynos backend initialize failed."); - executor->~EnnExecutor(); - return err; - } + err = executor->initialize(binary_buf_addr, buf_size); + } + if (err != Error::Ok) { + ENN_LOG_ERROR("Exynos backend initialize failed."); + destroy(executor); + return err; } return executor; @@ -170,24 +181,13 @@ class EnnBackend final : public PyTorchBackendInterface { EXYNOS_ATRACE_NAME_LINE("backend_destroy"); if (handle != nullptr) { auto executor = static_cast(handle); - - // Find and remove the cached buffer that belongs to this executor - const std::string& weight_key = executor->get_weight_key(); + std::string weight_key = executor->get_weight_key(); + executor->~EnnExecutor(); if (!weight_key.empty()) { - auto it = g_weight_cache_lut.find(weight_key); - if (it != g_weight_cache_lut.end()) { - ET_LOG( - Info, "Deleting cached buffer for key: %s", weight_key.c_str()); - delete it->second; - g_weight_cache_lut.erase(it); - ET_LOG(Info, "Cached buffer removed from LUT"); - } + release_weight(weight_key); } - - executor->~EnnExecutor(); } } - }; // namespace executor namespace { diff --git a/backends/samsung/runtime/enn_executor.cpp b/backends/samsung/runtime/enn_executor.cpp index a27d7a4df30..91318ada03b 100644 --- a/backends/samsung/runtime/enn_executor.cpp +++ b/backends/samsung/runtime/enn_executor.cpp @@ -38,6 +38,12 @@ Error EnnExecutor::initialize( // EnnInitialize() resets the calling thread's preference on every call, so // re-apply it right before opening the model. enn_api_inst->applyPerfConfigFromSoc(); + ET_CHECK_OR_RETURN_ERROR( + buf_size <= UINT32_MAX && weight_size <= UINT32_MAX, + InvalidArgument, + "Model (%zu) or weight (%zu) buffer exceeds the 4GiB ENN limit", + buf_size, + weight_size); EnnReturn ret = ENN_RET_SUCCESS; ET_LOG(Info, "Start to open model %p, %zu", binary_buf_addr, buf_size); @@ -48,6 +54,10 @@ Error EnnExecutor::initialize( if (pte_buffer->va == binary_buf_addr && !enn_api_inst->EnnGetFileDescriptorFromEnnBuffer(pte_buffer, &fd)) { if (use_weight_fd) { + ET_CHECK_OR_RETURN_ERROR( + enn_api_inst->EnnOpenModelFromFdWithWeight != nullptr, + NotSupported, + "This ENN runtime cannot open models with external weights."); EnnBufferPtr ptd_buffer = nullptr; ET_CHECK_OR_RETURN_ERROR( sm_instance->query(&ptd_buffer, weight_buf_addr, weight_size), @@ -63,7 +73,9 @@ Error EnnExecutor::initialize( ret = enn_api_inst->EnnOpenModelFromFdWithWeight( fd, buf_size, pte_buffer->offset, &ptd_buffer, 1, &model_id_); ET_CHECK_OR_RETURN_ERROR( - ret == ENN_RET_SUCCESS, Internal, "Open model with weights failed."); + ret == ENN_RET_SUCCESS, + Internal, + "Open model with weights failed."); ET_LOG(Info, "Opened model from file descriptor with weights"); sm_instance->free(pte_buffer->va); } else { diff --git a/backends/samsung/runtime/enn_executor.h b/backends/samsung/runtime/enn_executor.h index f7429f4f8b4..cbfe5613e6b 100644 --- a/backends/samsung/runtime/enn_executor.h +++ b/backends/samsung/runtime/enn_executor.h @@ -7,13 +7,13 @@ * */ #pragma once + #include #include + #include #include -#include - namespace torch { namespace executor { namespace enn { @@ -44,10 +44,6 @@ class EnnExecutor { return num_of_outputs_; } - executorch::runtime::FreeableBuffer* get_weight() const { - return weight_; - } - void set_weight_key(const std::string& key) { weight_key_ = key; } @@ -63,7 +59,6 @@ class EnnExecutor { EnnBufferPtr* alloc_buffer_ = nullptr; int32_t num_of_inputs_ = 0; int32_t num_of_outputs_ = 0; - executorch::runtime::FreeableBuffer* weight_ = nullptr; std::string weight_key_; }; diff --git a/examples/samsung/CMakeLists.txt b/examples/samsung/CMakeLists.txt index d1dfbc26fff..d4fe9130eb9 100644 --- a/examples/samsung/CMakeLists.txt +++ b/examples/samsung/CMakeLists.txt @@ -9,8 +9,6 @@ set(CMAKE_CXX_STANDARD 17) set(CMAKE_CXX_STANDARD_REQUIRED ON) project(samsung_example) -message(EXECUTORCH_ROOT ${EXECUTORCH_ROOT}) - # Source root directory for executorch. if(NOT EXECUTORCH_ROOT) set(EXECUTORCH_ROOT ${CMAKE_CURRENT_SOURCE_DIR}/../..) @@ -33,49 +31,16 @@ endif() # prebuilt libraries. executorch package should contain portable_ops_lib, # etdump, bundled_program. -find_package(tokenizers REQUIRED) find_package(executorch CONFIG REQUIRED) target_compile_options(executorch INTERFACE -DET_EVENT_TRACER_ENABLED) find_package(gflags REQUIRED) set(_common_compile_options -Wno-deprecated-declarations -fPIC) -# -# The `__srcs` lists are defined by executorch_load_build_variables. -# -executorch_load_build_variables() - -set(EXECUTORCH_SRCS_FILE - "${CMAKE_CURRENT_BINARY_DIR}/../../executorch_srcs.cmake" -) - -get_filename_component( - EXECUTORCH_SOURCE_DIR "${CMAKE_CURRENT_LIST_DIR}/../.." ABSOLUTE -) - -add_compile_options( - -Wall - -Werror - -Wno-deprecated-declarations - -Wno-unused-const-variable - -Wno-unused-variable - -Wno-inconsistent-missing-override - -fPIC -) +add_compile_options(-Wall -Werror -fPIC) message("Build Samsung Android Examples") -link_directories(${EXECUTORCH_ROOT}/build_samsung_android/lib) - -add_library(lib_portable_ops STATIC IMPORTED) - -set_target_properties( - lib_portable_ops - PROPERTIES - IMPORTED_LOCATION - "${EXECUTORCH_ROOT}/build_samsung_android/lib/libportable_ops_lib.a" -) - set(__enn_executor_runner__srcs ${CMAKE_CURRENT_LIST_DIR}/executor_runner/enn_executor_runner.cpp ) @@ -95,19 +60,43 @@ set_target_properties( enn_executor_runner PROPERTIES CXX_VISIBILITY_PRESET hidden ) -# LLM runner library -add_subdirectory(executor_runner/llm/runner) +option(EXYNOS_BUILD_LLM_RUNNER + "Build enn_executor_llm_runner and the Exynos LLM runner library" OFF +) +# C API in exynos_llm_wrapper.h. Always compiled into enn_executor_llm_runner; +# additionally shipped as libexynos_llm_wrapper.so when +# EXYNOS_BUILD_LLM_WRAPPER_SO is ON. +option(EXYNOS_BUILD_LLM_WRAPPER_SO + "Build libexynos_llm_wrapper.so and exynos_llm_wrapper_sample" OFF +) +if(EXYNOS_BUILD_LLM_WRAPPER_SO AND NOT EXYNOS_BUILD_LLM_RUNNER) + message( + FATAL_ERROR "EXYNOS_BUILD_LLM_WRAPPER_SO requires EXYNOS_BUILD_LLM_RUNNER" + ) +endif() +if(NOT EXYNOS_BUILD_LLM_RUNNER) + return() +endif() -target_link_directories(enn_llm_runner PRIVATE) +find_package(tokenizers REQUIRED) + +set(_llm_compile_options + -Wno-deprecated-declarations -Wno-unused-const-variable + -Wno-unused-variable -Wno-inconsistent-missing-override +) -add_library(lib_enn_llm_runner STATIC IMPORTED) +link_directories(${EXECUTORCH_ROOT}/build_samsung_android/lib) + +add_library(lib_portable_ops STATIC IMPORTED) set_target_properties( - lib_enn_llm_runner + lib_portable_ops PROPERTIES IMPORTED_LOCATION - "${EXECUTORCH_ROOT}/build_samsung_android/examples/samsung/executor_runner/llm/runner/libenn_llm_runner.a" + "${EXECUTORCH_ROOT}/build_samsung_android/lib/libportable_ops_lib.a" ) +add_subdirectory(executor_runner/llm/runner) + add_library(lib_tokenizer STATIC IMPORTED) set_target_properties( lib_tokenizer @@ -192,18 +181,12 @@ set(exynos_llm_link_libs -Wl,--no-whole-archive ) -# C API in exynos_llm_wrapper.h. Always compiled into enn_executor_llm_runner; -# additionally shipped as libexynos_llm_wrapper.so when -# EXYNOS_BUILD_LLM_WRAPPER_SO is ON. -option(EXYNOS_BUILD_LLM_WRAPPER_SO - "Build libexynos_llm_wrapper.so and exynos_llm_wrapper_sample" OFF -) - add_library( exynos_llm_wrapper_obj OBJECT ${CMAKE_CURRENT_LIST_DIR}/executor_runner/exynos_llm_wrapper.cpp ) target_compile_definitions(exynos_llm_wrapper_obj PRIVATE EXECUTORCH_VERSION=2) +target_compile_options(exynos_llm_wrapper_obj PRIVATE ${_llm_compile_options}) target_link_libraries(exynos_llm_wrapper_obj PRIVATE enn_llm_runner) set_target_properties( exynos_llm_wrapper_obj PROPERTIES CXX_VISIBILITY_PRESET hidden @@ -222,6 +205,7 @@ target_include_directories( ${EXECUTORCH_ROOT}/extension/llm/tokenizers/include ) +target_compile_options(enn_executor_llm_runner PRIVATE ${_llm_compile_options}) target_link_libraries( enn_executor_llm_runner PRIVATE exynos_llm_wrapper_obj enn_llm_runner ${exynos_llm_link_libs} @@ -246,6 +230,9 @@ if(EXYNOS_BUILD_LLM_WRAPPER_SO) exynos_llm_wrapper_sample ${CMAKE_CURRENT_LIST_DIR}/executor_runner/exynos_llm_wrapper_sample_main.cpp ) + target_compile_options( + exynos_llm_wrapper_sample PRIVATE ${_llm_compile_options} + ) target_link_libraries(exynos_llm_wrapper_sample PRIVATE exynos_llm_wrapper) set_target_properties( exynos_llm_wrapper_sample PROPERTIES CXX_VISIBILITY_PRESET hidden diff --git a/examples/samsung/build.sh b/examples/samsung/build.sh index 6fb0d38c136..c5ed57ee8cc 100755 --- a/examples/samsung/build.sh +++ b/examples/samsung/build.sh @@ -20,15 +20,18 @@ echo ANDROID_ABI=${ANDROID_ABI} echo ANDROID_PLATFORM=${ANDROID_PLATFORM} usage() { - echo "Usage: $0 [--llm-wrapper-so]" - echo " --llm-wrapper-so Also build libexynos_llm_wrapper.so and exynos_llm_wrapper_sample." + echo "Usage: $0 [--llm] [--llm-wrapper-so]" + echo " --llm Also build enn_executor_llm_runner." + echo " --llm-wrapper-so Also build libexynos_llm_wrapper.so and exynos_llm_wrapper_sample (implies --llm)." } main() { + local build_llm_runner=OFF local build_llm_wrapper_so=OFF for arg in "$@"; do case "$arg" in - "--llm-wrapper-so") build_llm_wrapper_so=ON ;; + "--llm") build_llm_runner=ON ;; + "--llm-wrapper-so") build_llm_runner=ON; build_llm_wrapper_so=ON ;; "--help"|"-h") usage; exit 0 ;; *) echo "Unknown option: $arg"; usage; exit 1 ;; esac @@ -36,12 +39,6 @@ main() { cd "$PROJECT_DIR" - if [[ -z ${ANDROID_NDK_ROOT} ]]; then - echo "Please export ANDROID_NDK_ROOT by setting andrid NDK path" - echo "$ export ANDROID_NDK_ROOT = path of android NDK" - exit 1 - fi - if [[ -z ${ANDROID_NDK} ]]; then ANDROID_NDK=${ANDROID_NDK_ROOT} fi @@ -62,6 +59,7 @@ main() { -DANDROID_PLATFORM=$ANDROID_PLATFORM \ -DCMAKE_FIND_ROOT_PATH_MODE_PACKAGE=BOTH \ -DCMAKE_BUILD_TYPE=Release \ + -DEXYNOS_BUILD_LLM_RUNNER=${build_llm_runner} \ -DEXYNOS_BUILD_LLM_WRAPPER_SO=${build_llm_wrapper_so} \ -B"${build_dir_example}" \ "$PROJECT_DIR/${example_root}" diff --git a/examples/samsung/executor_runner/README.md b/examples/samsung/executor_runner/README.md index ad10fe320cd..19d5ac0f2c2 100644 --- a/examples/samsung/executor_runner/README.md +++ b/examples/samsung/executor_runner/README.md @@ -2,6 +2,14 @@ This directory contains the executor runner for LLMs on Exynos platforms. The runner requires a JSON configuration file to correctly set up the model architecture and runtime parameters. +## Building + +The LLM runner is not part of the default example build. After building the backend with `./backends/samsung/build.sh`, build it with: + +```bash +./examples/samsung/build.sh --llm +``` + ## Using JSON Configuration Files The `enn_executor_llm_runner` binary accepts a path to a JSON configuration file via the `--config_path` argument. This file defines critical model dimensions and strategies used during execution. @@ -18,14 +26,26 @@ You can refer to `run_llama3_sample.sh` for a complete execution example. The ke ... ``` +For Gemma 4 multimodal inference, also pass the vision encoder and the image with `--vision_model_path`, `--image_path` and, optionally, `--prompt_before_image` for the text placed before the image. + ### Where to find config files -Configuration files for supported models are located in `examples/models/`. -For example: -- Gemma 3 1B: `examples/models/gemma3_1b/exynos_gemma3_1b.json` -- Llama 3 1B: `examples/models/llama3_1b/exynos_llama3_1b.json` + +Configuration files for supported models are located in `configs/`: +- Gemma 3 1B: `configs/exynos_gemma3_1b.json` +- Gemma 4 E2B: `configs/exynos_gemma4_e2b.json` +- Llama 3 1B: `configs/exynos_llama3_1b.json` Ensure these files are pushed to the target device (e.g., to `/data/local/tmp/`) before running the script. +### Position embeddings + +`--pos_embedding_path` is optional: when the file is missing, the runner precomputes the RoPE tables from the config. For Gemma 4, generate the table that matches the exported model with: + +```bash +python -m executorch.examples.samsung.scripts.gemma4_pos_embedding \ + --config --max-seq-len 4096 --out pos_embedding.bin +``` + ## `exynos_llm_wrapper` library (app-porting reference) `libexynos_llm_wrapper.so` is a reference for porting the Exynos LLM runner into @@ -38,4 +58,4 @@ that exercises the same API on-device. ```bash ./examples/samsung/build.sh --llm-wrapper-so -``` \ No newline at end of file +``` diff --git a/examples/models/gemma3_1b/exynos_gemma3_1b.json b/examples/samsung/executor_runner/configs/exynos_gemma3_1b.json similarity index 99% rename from examples/models/gemma3_1b/exynos_gemma3_1b.json rename to examples/samsung/executor_runner/configs/exynos_gemma3_1b.json index 554d03ec126..3adcedde9c8 100644 --- a/examples/models/gemma3_1b/exynos_gemma3_1b.json +++ b/examples/samsung/executor_runner/configs/exynos_gemma3_1b.json @@ -18,4 +18,4 @@ "attn_mask_rows": 128, "attn_mask_cols": 1024, "sliding_window": 512 -} \ No newline at end of file +} diff --git a/examples/models/gemma4/config/exynos_gemma4_e2b.json b/examples/samsung/executor_runner/configs/exynos_gemma4_e2b.json similarity index 99% rename from examples/models/gemma4/config/exynos_gemma4_e2b.json rename to examples/samsung/executor_runner/configs/exynos_gemma4_e2b.json index 62577ea953e..aaa5d379b91 100644 --- a/examples/models/gemma4/config/exynos_gemma4_e2b.json +++ b/examples/samsung/executor_runner/configs/exynos_gemma4_e2b.json @@ -66,4 +66,4 @@ "image_token_id": 258880, "video_token_id": 258884 } - } \ No newline at end of file + } diff --git a/examples/models/llama3_1b/exynos_llama3_1b.json b/examples/samsung/executor_runner/configs/exynos_llama3_1b.json similarity index 99% rename from examples/models/llama3_1b/exynos_llama3_1b.json rename to examples/samsung/executor_runner/configs/exynos_llama3_1b.json index c5050553a94..b15807af27b 100644 --- a/examples/models/llama3_1b/exynos_llama3_1b.json +++ b/examples/samsung/executor_runner/configs/exynos_llama3_1b.json @@ -17,4 +17,4 @@ "decode_seq_len": 1, "attn_mask_rows": 128, "attn_mask_cols": 1024 -} \ No newline at end of file +} diff --git a/examples/samsung/executor_runner/enn_executor_llm_runner.cpp b/examples/samsung/executor_runner/enn_executor_llm_runner.cpp index 5d5a21569a0..475d531faac 100644 --- a/examples/samsung/executor_runner/enn_executor_llm_runner.cpp +++ b/examples/samsung/executor_runner/enn_executor_llm_runner.cpp @@ -83,7 +83,7 @@ DEFINE_double( DEFINE_int32( seq_len, 128, - "Total number of tokens to generate (prompt + output). Defaults to max_seq_len. If the number of input tokens + seq_len > max_seq_len, the output will be truncated to max_seq_len tokens."); + "Total number of tokens to generate (prompt + output). If the number of input tokens + seq_len > max_seq_len, the output will be truncated to max_seq_len tokens."); DEFINE_int32( cpu_threads, @@ -100,8 +100,6 @@ DEFINE_int32( 0, "Number of EOS tokens to append to the prompt. Defaults to 0. If > 0, the prompt will be appended with EOS tokens. This is useful for models that expect one or more EOS token at the end."); -DEFINE_bool(is_warmup, false, "Whether to run a warmup run."); - DEFINE_bool(is_decode_only, false, "Whether to run with decode method only."); DEFINE_int32( @@ -169,7 +167,10 @@ int32_t main(int32_t argc, char** argv) { ET_LOG(Error, "Input prompt cannot be empty."); return ENN_ERR_INVALID_ARG; } - ET_LOG(Info, "Input prompt message: %s", effective_prompt.c_str()); + if (FLAGS_cpu_threads < -1) { + ET_LOG(Error, "--cpu_threads must be -1 or a non-negative count."); + return ENN_ERR_INVALID_ARG; + } #if defined(ET_USE_THREADPOOL) uint32_t num_performant_cores = FLAGS_cpu_threads == -1 @@ -191,9 +192,10 @@ int32_t main(int32_t argc, char** argv) { FLAGS_config_path.c_str(), FLAGS_vision_model_path.c_str(), FLAGS_is_decode_only); - ET_CHECK_MSG(runner != nullptr, "Failed to create LLM runner"); - - // TODO: support warmup interface + if (runner == nullptr) { + ET_LOG(Error, "Failed to create LLM runner"); + return ENN_ERR_LOAD; + } const bool stop_test = FLAGS_max_token_limit > 0; StopContext stop_ctx{runner, FLAGS_max_token_limit}; @@ -220,7 +222,8 @@ int32_t main(int32_t argc, char** argv) { stop_ctx.limit, stop_ctx.num_tokens); } - ET_CHECK_MSG(error == ENN_OK, "Failed to generate token from LLM runner"); - - return 0; + if (error != ENN_OK) { + ET_LOG(Error, "Failed to generate token from LLM runner"); + } + return error; } diff --git a/examples/samsung/executor_runner/exynos_llm_wrapper.cpp b/examples/samsung/executor_runner/exynos_llm_wrapper.cpp index 2eab0620402..3b9cbd221f5 100644 --- a/examples/samsung/executor_runner/exynos_llm_wrapper.cpp +++ b/examples/samsung/executor_runner/exynos_llm_wrapper.cpp @@ -2,6 +2,9 @@ * Copyright (c) 2026 Samsung Electronics Co. LTD * All rights reserved. * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + * * C API implementation over the Exynos LLM runner. Keeps the model(s) loaded * across generate() calls so consecutive inferences do not re-prepare. */ diff --git a/examples/samsung/executor_runner/exynos_llm_wrapper.h b/examples/samsung/executor_runner/exynos_llm_wrapper.h index ce91a8391c6..94535c74167 100644 --- a/examples/samsung/executor_runner/exynos_llm_wrapper.h +++ b/examples/samsung/executor_runner/exynos_llm_wrapper.h @@ -2,6 +2,9 @@ * Copyright (c) 2026 Samsung Electronics Co. LTD * All rights reserved. * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + * * Public C API wrapping the Exynos LLM runner. * * Provided as a reference for porting the Exynos LLM runner into an Android @@ -73,7 +76,11 @@ EXYNOS_LLM_WRAPPER_API enn_runner* enn_llm_create( const char* vision_model_path, int32_t is_decode_only); -/** Destroy the runner and release all resources. Safe to call with NULL. */ +/** + * Destroy the runner and release all resources. Safe to call with NULL. Must + * not run concurrently with enn_llm_generate or enn_llm_stop on the same + * handle. + */ EXYNOS_LLM_WRAPPER_API void enn_llm_destroy(enn_runner* runner); /** @@ -87,7 +94,8 @@ EXYNOS_LLM_WRAPPER_API void enn_llm_stop(enn_runner* runner); /** * Generate a response for prompt (optionally with an image). The models stay * prepared across calls; KV-cache is reset before each generation so every - * prompt starts fresh. + * prompt starts fresh. Not reentrant: at most one generation may run on a + * handle at a time. * * @param runner Handle from enn_llm_create. * @param prompt User prompt text (required, non-empty). When raw_prompt is @@ -104,7 +112,8 @@ EXYNOS_LLM_WRAPPER_API void enn_llm_stop(enn_runner* runner); * @param raw_prompt Non-zero => prompt is already fully formatted; skip the * internal chat-template wrapping and system_prompt. * @param seq_len Total token budget (prompt + generated), capped by the model's - * max context length; -1 => use the model's max context length. + * max sequence length; -1 => use the model's max sequence length. Must be + * -1 or larger than the prompt length. * @param temperature Sampling temperature. * @param num_bos BOS tokens to prepend to the prompt. * @param num_eos EOS tokens to append to the prompt. diff --git a/examples/samsung/executor_runner/exynos_llm_wrapper_sample_main.cpp b/examples/samsung/executor_runner/exynos_llm_wrapper_sample_main.cpp index b5fc1cea27a..074114c589a 100644 --- a/examples/samsung/executor_runner/exynos_llm_wrapper_sample_main.cpp +++ b/examples/samsung/executor_runner/exynos_llm_wrapper_sample_main.cpp @@ -2,6 +2,9 @@ * Copyright (c) 2026 Samsung Electronics Co. LTD * All rights reserved. * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + * * Sample CLI consuming the public exynos_llm_wrapper C API * (exynos_llm_wrapper.h). This mirrors how an Android app drives * libexynos_llm_wrapper.so: create once, generate repeatedly, destroy at @@ -24,8 +27,11 @@ * - `--*_prompt_file `: read the file and use its contents. * - `--image_file ` : the value IS the image path passed as image_path. * The system/user/prefill texts are fed together (raw_prompt=0) as one - * generate call, along with any image. Any may be omitted. Repeating runs - * sequentially on the same handle, so the model stays prepared across calls. + * generate call, along with any image. The user prompt is required; system + * and prefill may be omitted. A raw prompt always runs on its own, so a + * following system/user/prefill option starts the next generation. Repeating + * runs sequentially on the same handle, so the model stays prepared across + * calls. * * STREAMING=1 (default): the on_token callback prints each token as generated. * STREAMING=0: the on_token callback appends to a std::string and the full @@ -43,7 +49,9 @@ #include #include #include +#include #include +#include #include #include #include @@ -147,10 +155,18 @@ static int run_one( // Timer-based stopper thread imitating a user interrupt. Replaced by the // token-count stop test when max_output_tokens is given. + std::mutex stop_mutex; + std::condition_variable stop_cv; + bool generation_done = false; std::thread stopper; if (stop_after_sec > 0) { - stopper = std::thread([r, stop_after_sec]() { - std::this_thread::sleep_for(std::chrono::seconds(stop_after_sec)); + stopper = std::thread([&]() { + std::unique_lock lock(stop_mutex); + if (stop_cv.wait_for(lock, std::chrono::seconds(stop_after_sec), [&]() { + return generation_done; + })) { + return; + } printf( "\n[Stopper] triggering enn_llm_stop() after %d s\n", stop_after_sec); fflush(stdout); @@ -176,6 +192,11 @@ static int run_one( &ctx); if (stopper.joinable()) { + { + std::lock_guard lock(stop_mutex); + generation_done = true; + } + stop_cv.notify_one(); stopper.join(); } @@ -271,27 +292,32 @@ int main(int argc, char** argv) { content = optarg; } + const bool starts_templated = opt == OPT_SYS || opt == OPT_SYS_FILE || + opt == OPT_USR || opt == OPT_USR_FILE || opt == OPT_PRE || + opt == OPT_PRE_FILE; + const bool starts_raw = opt == OPT_RAW || opt == OPT_RAW_FILE; + if (have && (starts_raw || (starts_templated && !raw.empty()))) { + rc_total |= run_one( + r, + raw, + sys, + usr, + pre, + image, + max_output_tokens, + stop_after_sec, + echo); + raw.clear(); + sys.clear(); + usr.clear(); + pre.clear(); + image.clear(); + have = false; + } + switch (opt) { case OPT_RAW: case OPT_RAW_FILE: - if (have) { - rc_total |= run_one( - r, - raw, - sys, - usr, - pre, - image, - max_output_tokens, - stop_after_sec, - echo); - raw.clear(); - sys.clear(); - usr.clear(); - pre.clear(); - image.clear(); - have = false; - } raw = content; have = true; break; diff --git a/examples/samsung/executor_runner/llm/runner/CMakeLists.txt b/examples/samsung/executor_runner/llm/runner/CMakeLists.txt index 87c02561dfc..842e9eab6f7 100644 --- a/examples/samsung/executor_runner/llm/runner/CMakeLists.txt +++ b/examples/samsung/executor_runner/llm/runner/CMakeLists.txt @@ -4,31 +4,29 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. -message("###### Build ENN LLM runner ######") -message(ROOT ${CMAKE_CURRENT_SOURCE_DIR}/../../../../..) -message(EXECUTORCH_ROOT ${EXECUTORCH_ROOT}) -message(CMAKE_CURRENT_SOURCE_DIR ${CMAKE_CURRENT_SOURCE_DIR}) -message(CMAKE_CURRENT_BINARY_DIR ${CMAKE_CURRENT_BINARY_DIR}) -message(CMAKE_CURRENT_LIST_DIR ${CMAKE_CURRENT_LIST_DIR}) - include(${CMAKE_CURRENT_SOURCE_DIR}/../../../../../tools/cmake/Utils.cmake) include(${CMAKE_CURRENT_SOURCE_DIR}/../../../../../tools/cmake/Codegen.cmake) -list(TRANSFORM _enn_llm_runner__srcs - PREPEND "${CMAKE_CURRENT_SOURCE_DIR}/../../../../.." +# stb_image: same pinned revision as examples/models/gemma3. deprecated/ holds +# stb_image_resize.h. +include(FetchContent) +FetchContent_Declare( + stb + GIT_REPOSITORY https://github.com/nothings/stb.git + GIT_TAG f0569113c93ad095470c54bf34a17b36646bbbb5 ) -list( - PREPEND - _enn_llm_runner__srcs - ${CMAKE_CURRENT_LIST_DIR}/runner.cpp - ${CMAKE_CURRENT_LIST_DIR}/config.cpp - ${CMAKE_CURRENT_LIST_DIR}/decoder.cpp - ${CMAKE_CURRENT_LIST_DIR}/prefiller.cpp - ${CMAKE_CURRENT_LIST_DIR}/io_manager.cpp - ${CMAKE_CURRENT_LIST_DIR}/rotary_embedding.cpp - ${CMAKE_CURRENT_LIST_DIR}/vision_encoder.cpp - ${CMAKE_CURRENT_LIST_DIR}/stb_image_impl.cpp - ${CMAKE_CURRENT_SOURCE_DIR}/../../../../../examples/models/llama/tokenizer/llama_tiktoken.cpp +FetchContent_MakeAvailable(stb) + +set(_enn_llm_runner__srcs + ${CMAKE_CURRENT_LIST_DIR}/runner.cpp + ${CMAKE_CURRENT_LIST_DIR}/config.cpp + ${CMAKE_CURRENT_LIST_DIR}/decoder.cpp + ${CMAKE_CURRENT_LIST_DIR}/prefiller.cpp + ${CMAKE_CURRENT_LIST_DIR}/io_manager.cpp + ${CMAKE_CURRENT_LIST_DIR}/rotary_embedding.cpp + ${CMAKE_CURRENT_LIST_DIR}/vision_encoder.cpp + ${CMAKE_CURRENT_LIST_DIR}/stb_image_impl.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/../../../../../examples/models/llama/tokenizer/llama_tiktoken.cpp ) target_include_directories( @@ -75,6 +73,7 @@ set(enn_llm_runner_deps extension_flat_tensor lib_extension_llm_runner ) +target_compile_options(enn_llm_runner PRIVATE ${_llm_compile_options}) target_link_libraries(enn_llm_runner PUBLIC ${enn_llm_runner_deps}) target_link_libraries(enn_llm_runner PUBLIC tokenizers::tokenizers) target_link_libraries(enn_llm_runner PUBLIC lib_tokenizer) @@ -89,6 +88,7 @@ target_include_directories( ${CMAKE_CURRENT_SOURCE_DIR}/../../../../../extension/llm/tokenizers/third-party/json/include ${CMAKE_CURRENT_SOURCE_DIR}/../../../../../examples/models/llama/tokenizer ${CMAKE_CURRENT_SOURCE_DIR}/../../../../../examples/samsung/executor_runner/llm/runner - ${CMAKE_CURRENT_SOURCE_DIR}/../../../../../third-party/stb + ${stb_SOURCE_DIR} + ${stb_SOURCE_DIR}/deprecated ${CMAKE_CURRENT_SOURCE_DIR}/../../../../../extension/llm/runner ) diff --git a/examples/samsung/executor_runner/llm/runner/config.cpp b/examples/samsung/executor_runner/llm/runner/config.cpp index 7f51d1d1e43..0d9f464fd9f 100644 --- a/examples/samsung/executor_runner/llm/runner/config.cpp +++ b/examples/samsung/executor_runner/llm/runner/config.cpp @@ -7,34 +7,36 @@ */ #include -#include +#include #include -#include -#include +#include +#include namespace example { using json = nlohmann::json; +using executorch::runtime::Error; -void RunnerConfig::load_from_json() { - if (config_path.empty()) { - ET_LOG(Info, "Config path is empty, skipping JSON load"); - return; - } +Error RunnerConfig::load_from_json() { + ET_CHECK_OR_RETURN_ERROR( + !config_path.empty(), + InvalidArgument, + "A runner config path is required"); std::ifstream file(config_path); - if (!file.is_open()) { - ET_LOG(Error, "Could not open config file: %s", config_path.c_str()); - return; - } + ET_CHECK_OR_RETURN_ERROR( + file.is_open(), + InvalidArgument, + "Could not open config file: %s", + config_path.c_str()); json j = json::parse(file, nullptr, false); - if (j.is_discarded()) { - ET_LOG( - Error, "Invalid JSON format in config file: %s", config_path.c_str()); - return; - } + ET_CHECK_OR_RETURN_ERROR( + !j.is_discarded(), + InvalidArgument, + "Invalid JSON format in config file: %s", + config_path.c_str()); if (j.contains("embedding_strategy") && j["embedding_strategy"].is_string()) { std::string strategy = j["embedding_strategy"].get(); @@ -325,5 +327,63 @@ void RunnerConfig::load_from_json() { default: break; } + return validate(); +} + +Error RunnerConfig::validate() const { + ET_CHECK_OR_RETURN_ERROR( + model_args != nullptr, + InvalidArgument, + "%s: model_version must be one of llama3, gemma3, gemma4", + config_path.c_str()); + const BaseModelArgs& args = *model_args; + ET_CHECK_OR_RETURN_ERROR( + args.get_n_layers() > 0 && args.get_n_kv_heads() > 0 && + args.get_hidden_dim() > 0, + InvalidArgument, + "%s: n_layers, n_kv_heads and hidden_dim must be positive", + config_path.c_str()); + ET_CHECK_OR_RETURN_ERROR( + args.get_head_dim() > 0 && args.get_head_dim() % 2 == 0 && + args.get_global_head_dim() > 0 && args.get_global_head_dim() % 2 == 0, + InvalidArgument, + "%s: head dims must be positive and even", + config_path.c_str()); + ET_CHECK_OR_RETURN_ERROR( + args.get_num_kv_shared_layers() >= 0 && + args.get_num_kv_shared_layers() < args.get_n_layers(), + InvalidArgument, + "%s: num_kv_shared_layers (%" PRId64 ") must be in [0, n_layers)", + config_path.c_str(), + args.get_num_kv_shared_layers()); + ET_CHECK_OR_RETURN_ERROR( + args.get_prefill_seq_len() >= 1 && + args.get_prefill_seq_len() <= args.get_max_seq_len(), + InvalidArgument, + "%s: prefill_seq_len (%" PRId64 ") must be in [1, max_seq_len (%" PRId64 + ")]", + config_path.c_str(), + args.get_prefill_seq_len(), + args.get_max_seq_len()); + ET_CHECK_OR_RETURN_ERROR( + args.get_decode_seq_len() == 1, + InvalidArgument, + "%s: only decode_seq_len == 1 is supported", + config_path.c_str()); + ET_CHECK_OR_RETURN_ERROR( + args.get_attn_mask_rows() == args.get_prefill_seq_len() && + args.get_attn_mask_cols() == args.get_max_seq_len(), + InvalidArgument, + "%s: attention mask must be [prefill_seq_len, max_seq_len]", + config_path.c_str()); + const bool needs_sliding_window = + masking_strategy == MaskingStrategy::kCausalLocal || + masking_strategy == MaskingStrategy::kHybridConvAttention; + ET_CHECK_OR_RETURN_ERROR( + !needs_sliding_window || args.get_sliding_window().value_or(0) > 0, + InvalidArgument, + "%s: the selected masking strategy needs a positive sliding_window", + config_path.c_str()); + return Error::Ok; } -} // namespace example \ No newline at end of file +} // namespace example diff --git a/examples/samsung/executor_runner/llm/runner/config.h b/examples/samsung/executor_runner/llm/runner/config.h index a1ab48a6f4e..deec0db3afa 100644 --- a/examples/samsung/executor_runner/llm/runner/config.h +++ b/examples/samsung/executor_runner/llm/runner/config.h @@ -9,6 +9,7 @@ #pragma once #include +#include #include #include @@ -25,6 +26,7 @@ inline constexpr auto kNumOfLayers = "get_n_layers"; inline constexpr auto kNumOfHiddenDim = "get_hidden_dim"; inline constexpr auto kNumOfHeads = "get_n_heads"; inline constexpr auto kNumOfKVHeads = "get_n_kv_heads"; +inline constexpr auto kHeadDim = "get_head_dim"; inline constexpr auto default_method_name = "forward"; inline constexpr auto default_method_name_prefill = "prefill_forward"; @@ -46,138 +48,6 @@ enum class DecoderModelVersion { kGeneral, }; -struct ModelVersionConfig { - size_t position_rank; - size_t attention_mask_rank; - size_t kv_cache_rank; - - size_t attention_mask_tensor_index; - size_t attention_mask_tensor_count; - - size_t position_tensor_index; - size_t position_tensor_count; - - size_t k_cache_tensor_index; - size_t v_cache_tensor_index; - size_t kv_cache_tensor_count; - - size_t heads_dim_index; - size_t head_size_dim_index; - - bool has_local_global_rope; - bool has_partial_rope; - bool has_shared_kv; - - size_t sliding_window; -}; - -inline ModelVersionConfig getModelVersionConfig(DecoderModelVersion version) { - switch (version) { - case DecoderModelVersion::kGemma3: - return { - .position_rank = 3, - .attention_mask_rank = 3, - .kv_cache_rank = 4, - - // Flat PTE: - // 0 input_ids - // 1-2 attention_setup - // 3-4 masks - // 5-6 position - // 7-32 k_cache_0..25 - // 33-58 v_cache_0..25 - .attention_mask_tensor_index = 3, - .attention_mask_tensor_count = 2, - - .position_tensor_index = 5, - .position_tensor_count = 2, - - .k_cache_tensor_index = 7, - .v_cache_tensor_index = 33, - .kv_cache_tensor_count = 26, - - .heads_dim_index = 1, - .head_size_dim_index = 3, - - .has_local_global_rope = true, - .has_partial_rope = false, - .has_shared_kv = false, - .sliding_window = 512, - }; - - case DecoderModelVersion::kGemma4: - return { - .position_rank = 3, - .attention_mask_rank = 4, - .kv_cache_rank = 4, - - // Flat PTE: - // 0 input_ids - // 1 inputs_embeds - // 2 mm_mask - // 3 attn_mask_full - // 4 attn_mask_sliding - // 5 pos_emb_global_cos - // 6 pos_emb_global_sin - // 7 pos_emb_local_cos - // 8 pos_emb_local_sin - // 9-43 k_cache_0..34 - // 44-78 v_cache_0..34 - .attention_mask_tensor_index = 3, - .attention_mask_tensor_count = 2, - - .position_tensor_index = 5, - .position_tensor_count = 4, - - .k_cache_tensor_index = 9, - .v_cache_tensor_index = 44, - .kv_cache_tensor_count = 35, - - .heads_dim_index = 1, - .head_size_dim_index = 3, - - .has_local_global_rope = true, - .has_partial_rope = true, - .has_shared_kv = true, - .sliding_window = 512, - }; - - case DecoderModelVersion::kLlama3: - return { - .position_rank = 3, - .attention_mask_rank = 4, - .kv_cache_rank = 4, - - .attention_mask_tensor_index = 1, - .attention_mask_tensor_count = 1, - - .position_tensor_index = 2, - .position_tensor_count = 2, - - .k_cache_tensor_index = 4, - .v_cache_tensor_index = 20, - .kv_cache_tensor_count = 16, - - .heads_dim_index = 1, - .head_size_dim_index = 3, - - .has_local_global_rope = false, - .has_partial_rope = false, - .has_shared_kv = false, - .sliding_window = 0, - }; - default: - return { - .position_rank = 3, - .kv_cache_rank = 4, - .has_local_global_rope = false, - .has_partial_rope = false, - .has_shared_kv = false, - .sliding_window = 0, - }; - } -} - enum class EmbeddingStrategy { kGlobalOnly, kGlobalAndLocal, @@ -516,10 +386,13 @@ class RunnerConfig { std::string tokenizer_type; - DecoderModelVersion model_version; + DecoderModelVersion model_version = DecoderModelVersion::kGeneral; std::unique_ptr model_args; - void load_from_json(); + ET_NODISCARD executorch::runtime::Error load_from_json(); + + private: + ET_NODISCARD executorch::runtime::Error validate() const; }; } // namespace example \ No newline at end of file diff --git a/examples/samsung/executor_runner/llm/runner/decoder.cpp b/examples/samsung/executor_runner/llm/runner/decoder.cpp index 5d6c17edfc2..5d33f99b4f3 100644 --- a/examples/samsung/executor_runner/llm/runner/decoder.cpp +++ b/examples/samsung/executor_runner/llm/runner/decoder.cpp @@ -11,66 +11,79 @@ #include #include #include -#include -#include +#include #include +#include namespace example { namespace runtime = executorch::runtime; -class IExecutionIOBinder { - virtual ~IExecutionIOBinder() = default; - virtual void beforeExecution() = 0; - virtual void afterExecution() = 0; -}; +namespace { -runtime::Result slice_kth_token( - const runtime::EValue& output_evalue, +runtime::Error validate_inputs( + const runtime::MethodMeta& method_meta, + const std::vector& inputs) { + ET_CHECK_OR_RETURN_ERROR( + inputs.size() == method_meta.num_inputs(), + InvalidProgram, + "Method %s expects %zu inputs but the runner prepared %zu", + method_meta.name(), + method_meta.num_inputs(), + inputs.size()); + for (size_t i = 0; i < inputs.size(); ++i) { + if (!inputs[i].isTensor()) { + continue; + } + auto expected = ET_UNWRAP(method_meta.input_tensor_meta(i)); + const auto prepared_rank = inputs[i].toTensor().dim(); + ET_CHECK_OR_RETURN_ERROR( + static_cast(prepared_rank) == expected.sizes().size(), + InvalidProgram, + "Method %s input %zu: expected rank %zu, prepared rank %zu", + method_meta.name(), + i, + expected.sizes().size(), + static_cast(prepared_rank)); + } + return runtime::Error::Ok; +} + +} // namespace + +runtime::Result Decoder::slice_kth_token( + const runtime::EValue& output, int64_t k) { - ET_CHECK_MSG( - output_evalue.isTensor(), "The provided EValue is not a tensor."); - aten::Tensor logits = output_evalue.toTensor(); + ET_CHECK_OR_RETURN_ERROR( + output.isTensor(), InvalidProgram, "Logits output is not a tensor"); + aten::Tensor logits = output.toTensor(); constexpr int64_t kExpectedNumDim = 3; constexpr int64_t kSequenceDim = 1; - const int64_t max_seq_len = logits.size(kSequenceDim); - const int64_t actual_dim = logits.dim(); - - const auto strides = logits.strides(); - const auto sizes = logits.sizes(); - - ET_CHECK_MSG( - k >= 0 && k < max_seq_len, - "Index k (%ld) out of bounds (max: %ld)", - k, - max_seq_len); - - ET_CHECK_MSG( - actual_dim == kExpectedNumDim, - "num dim shape of logit(%ld) is not expected(%ld)", - actual_dim, + ET_CHECK_OR_RETURN_ERROR( + logits.dim() == kExpectedNumDim, + InvalidProgram, + "Logits rank %zu is not %" PRId64, + static_cast(logits.dim()), kExpectedNumDim); + const int64_t seq_len = logits.size(kSequenceDim); + ET_CHECK_OR_RETURN_ERROR( + k >= 0 && k < seq_len, + InvalidArgument, + "Token index %" PRId64 " out of bounds for %" PRId64 " logits rows", + k, + seq_len); - // Calculate Memory Offset (zero-copy) - int64_t offset = k * strides[kSequenceDim]; - - // Define New Shape for the k-th token (1, 1, VocabSize) - // Using a fixed-size array or vector that matches SizesType + const auto sizes = logits.sizes(); std::vector new_sizes = { static_cast(sizes[0]), - 1, // Target single token + 1, static_cast(sizes[2])}; - - // Create TensorPtr and dereference it to return Result - auto tensor_ptr = extension::from_blob( - logits.mutable_data_ptr() + offset, - new_sizes, - logits.scalar_type()); - - // Dereference the shared_ptr to get the Tensor object - return *tensor_ptr; + uint8_t* row = logits.mutable_data_ptr() + + k * logits.strides()[kSequenceDim] * logits.element_size(); + last_logits_ = extension::from_blob(row, new_sizes, logits.scalar_type()); + return *last_logits_; } runtime::Error Decoder::load() { @@ -126,249 +139,60 @@ runtime::Error Decoder::load() { return ret; } +bool Decoder::is_method_loaded() { + return !prefill_method_.empty() && !decode_method_.empty() && + module_->is_method_loaded(prefill_method_) && + module_->is_method_loaded(decode_method_); +} + runtime::Result Decoder::step( extension::TensorPtr& tokens, int64_t start_pos) { - std::string method_stage; - if (tokens->numel() > 1) { - method_stage = prefill_method_; - } else { - method_stage = decode_method_; - } + const bool is_prefill = tokens->numel() > 1; + const std::string& method_stage = + is_prefill ? prefill_method_ : decode_method_; auto method_meta = ET_UNWRAP(module_->method_meta(method_stage.c_str())); - const size_t total_input_count = method_meta.num_inputs(); - - bool use_kv_cache = total_input_count > 1; - if (use_kv_cache) { - auto* llm_io_manager = static_cast(io_manager_); - const auto& config = llm_io_manager->get_config(); - auto model_version_config = getModelVersionConfig(config.model_version); - - // - // Validate attention-mask tensors. - // - for (size_t i = 0; i < model_version_config.attention_mask_tensor_count; - ++i) { - const size_t tensor_index = - model_version_config.attention_mask_tensor_index + i; - - auto tensor_meta = ET_UNWRAP(method_meta.input_tensor_meta(tensor_index)); - - const auto tensor_shape = tensor_meta.sizes(); - - ET_CHECK_OR_RETURN_ERROR( - tensor_shape.size() == model_version_config.attention_mask_rank, - InvalidProgram, - "Unexpected attention mask rank at input %zu. " - "Expected %zu, got %zu.", - tensor_index, - model_version_config.attention_mask_rank, - tensor_shape.size()); - } - - // - // Validate position / RoPE tensors. - // - // Gemma3: - // position[0], position[1] - // - // Gemma4: - // global_cos, global_sin, local_cos, local_sin - // - for (size_t i = 0; i < model_version_config.position_tensor_count; ++i) { - const size_t tensor_index = - model_version_config.position_tensor_index + i; - - auto tensor_meta = ET_UNWRAP(method_meta.input_tensor_meta(tensor_index)); - - const auto tensor_shape = tensor_meta.sizes(); - - ET_CHECK_OR_RETURN_ERROR( - tensor_shape.size() == model_version_config.position_rank, - InvalidProgram, - "Unexpected position tensor rank at input %zu. " - "Expected %zu, got %zu.", - tensor_index, - model_version_config.position_rank, - tensor_shape.size()); - } - - // - // Validate KV-cache tensors. - // - // We only need one cache tensor to determine/validate the common - // KV-cache layout [B, H, S, D], assuming all cache leaves use - // the same rank/layout. - // - auto k_cache_tensor_meta = ET_UNWRAP(method_meta.input_tensor_meta( - model_version_config.k_cache_tensor_index)); - - const auto k_cache_shape = k_cache_tensor_meta.sizes(); - - ET_CHECK_OR_RETURN_ERROR( - k_cache_shape.size() == model_version_config.kv_cache_rank, - InvalidProgram, - "Unexpected K-cache rank. Expected %zu, got %zu.", - model_version_config.kv_cache_rank, - k_cache_shape.size()); - - auto v_cache_tensor_meta = ET_UNWRAP(method_meta.input_tensor_meta( - model_version_config.v_cache_tensor_index)); - const auto v_cache_shape = v_cache_tensor_meta.sizes(); - - ET_CHECK_OR_RETURN_ERROR( - v_cache_shape.size() == model_version_config.kv_cache_rank, - InvalidProgram, - "Unexpected V-cache rank. Expected %zu, got %zu.", - model_version_config.kv_cache_rank, - v_cache_shape.size()); - - // - // Validate the configured KV-cache axes before indexing them. - // - ET_CHECK_OR_RETURN_ERROR( - model_version_config.heads_dim_index < k_cache_shape.size(), - InvalidProgram, - "KV heads dimension index %zu exceeds cache rank %zu.", - model_version_config.heads_dim_index, - k_cache_shape.size()); - - ET_CHECK_OR_RETURN_ERROR( - model_version_config.head_size_dim_index < k_cache_shape.size(), - InvalidProgram, - "KV head-size dimension index %zu exceeds cache rank %zu.", - model_version_config.head_size_dim_index, - k_cache_shape.size()); - - const auto num_kv_heads = - k_cache_shape[model_version_config.heads_dim_index]; - - const auto kv_head_size = - k_cache_shape[model_version_config.head_size_dim_index]; - - ET_LOG( - Debug, - "KV cache: K base=%zu, V base=%zu, count=%zu, " - "heads=%d, head_size=%d", - model_version_config.k_cache_tensor_index, - model_version_config.v_cache_tensor_index, - model_version_config.kv_cache_tensor_count, - num_kv_heads, - kv_head_size); - - // - // Optional sanity check that the final V-cache input actually - // exists in this PTE method. - // - const size_t last_v_cache_index = - model_version_config.v_cache_tensor_index + - model_version_config.kv_cache_tensor_count - 1; - - ET_CHECK_OR_RETURN_ERROR( - last_v_cache_index < method_meta.num_inputs(), - InvalidProgram, - "KV cache layout exceeds PTE input count. " - "Last V-cache input=%zu, num_inputs=%zu.", - last_v_cache_index, - method_meta.num_inputs()); - - // Assuming model is exported with input_pos, create a tensor with size 1 - extension::TensorPtr start_pos_tensor; - std::vector start_pos_sizes = {1}; - start_pos_tensor = extension::from_blob( - &start_pos, start_pos_sizes, aten::ScalarType::Long); - std::vector inputs; - if (method_stage == prefill_method_) { - auto inputs_res = io_manager_->prepare_prefill(tokens, start_pos_tensor); - ET_CHECK_OK_OR_RETURN_ERROR(inputs_res.error()); - inputs = inputs_res.get(); - - ET_LOG( - Info, - "Prepared %zu inputs, PTE expects %zu inputs", - inputs.size(), - method_meta.num_inputs()); - - const size_t count = std::min(inputs.size(), method_meta.num_inputs()); - - for (size_t i = 0; i < count; ++i) { - if (!inputs[i].isTensor()) { - ET_LOG(Info, "input[%zu]: prepared EValue is not Tensor", i); - continue; - } - - const auto prepared_tensor = inputs[i].toTensor(); - const auto prepared_sizes = prepared_tensor.sizes(); - - auto expected_meta = ET_UNWRAP(method_meta.input_tensor_meta(i)); - const auto expected_sizes = expected_meta.sizes(); - - ET_LOG( - Info, - "input[%zu]: prepared_rank=%zu expected_rank=%zu", - i, - prepared_sizes.size(), - expected_sizes.size()); - - for (size_t d = 0; d < prepared_sizes.size(); ++d) { - ET_LOG(Info, " prepared[%zu]=%d", d, prepared_sizes[d]); - } - - for (size_t d = 0; d < expected_sizes.size(); ++d) { - ET_LOG(Info, " expected[%zu]=%d", d, expected_sizes[d]); - } - } - - } else if (method_stage == decode_method_) { - auto inputs_res = io_manager_->prepare_decode(tokens, start_pos_tensor); - ET_CHECK_OK_OR_RETURN_ERROR(inputs_res.error()); - inputs = inputs_res.get(); - } else { - ET_LOG( - Debug, - "Skipping IO prepare for non-execution method stage: %s", - method_stage.c_str()); - } - - auto outputs_res = module_->execute(method_stage.c_str(), inputs); - ET_CHECK_OK_OR_RETURN_ERROR(outputs_res.error()); - ET_CHECK_MSG( - outputs_res.get()[0].isTensor(), - "Non Tensor Output returned from executing LLM"); - - if (method_stage == prefill_method_) { - ET_CHECK_OK_OR_RETURN_ERROR( - io_manager_->update_prefill(outputs_res.get()), - "Failed to update prefill output by io_manager"); - } else if (method_stage == decode_method_) { - ET_CHECK_OK_OR_RETURN_ERROR( - io_manager_->update_decode(outputs_res.get()), - "Failed to update decode output by io_manager"); - } else { - ET_LOG( - Debug, - "Skipping IO update for non-execution method stage: %s", - method_stage.c_str()); - } - // Return the logits tensor - return slice_kth_token(outputs_res.get()[0], (tokens->numel() - 1)); - } else { // kv cache is not used, when input is 1 - (void)start_pos; + if (method_meta.num_inputs() <= 1) { auto outputs_res = module_->execute(method_stage.c_str(), tokens); - ET_LOG(Info, "After module forward"); ET_CHECK_OK_OR_RETURN_ERROR(outputs_res.error()); - ET_CHECK_MSG( - outputs_res.get().size() == 1, - "More then one output returned from executing LLM."); - ET_CHECK_MSG( - outputs_res.get()[0].isTensor(), - "Non Tensor Output returned from executing LLM"); - - // Return the logits tensor + ET_CHECK_OR_RETURN_ERROR( + outputs_res.get().size() == 1 && outputs_res.get()[0].isTensor(), + InvalidProgram, + "Expected a single logits tensor from %s", + method_stage.c_str()); return outputs_res.get()[0].toTensor(); } + + std::vector start_pos_sizes = {1}; + auto start_pos_tensor = + extension::from_blob(&start_pos, start_pos_sizes, aten::ScalarType::Long); + auto inputs_res = is_prefill + ? io_manager_->prepare_prefill(tokens, start_pos_tensor, method_stage) + : io_manager_->prepare_decode(tokens, start_pos_tensor, method_stage); + ET_CHECK_OK_OR_RETURN_ERROR(inputs_res.error()); + std::vector inputs = std::move(inputs_res.get()); + ET_CHECK_OK_OR_RETURN_ERROR(validate_inputs(method_meta, inputs)); + + auto outputs_res = module_->execute(method_stage.c_str(), inputs); + ET_CHECK_OK_OR_RETURN_ERROR(outputs_res.error()); + ET_CHECK_OR_RETURN_ERROR( + !outputs_res.get().empty(), + InvalidProgram, + "%s returned no outputs", + method_stage.c_str()); + + if (is_prefill) { + ET_CHECK_OK_OR_RETURN_ERROR( + io_manager_->update_prefill(outputs_res.get(), method_stage), + "Failed to update prefill output by io_manager"); + } else { + ET_CHECK_OK_OR_RETURN_ERROR( + io_manager_->update_decode(outputs_res.get(), method_stage), + "Failed to update decode output by io_manager"); + } + return slice_kth_token(outputs_res.get()[0], tokens->numel() - 1); } void Decoder::extract_llm_methods() { @@ -419,7 +243,7 @@ runtime::Error Decoder::create_shared_planned_memory() { ET_LOG( Info, - "Prefill_bufs(%ld), decode_bufs(%ld) num_buffers(%ld)", + "Prefill_bufs(%zu), decode_bufs(%zu) num_buffers(%zu)", prefill_bufs, decode_bufs, num_buffers); @@ -429,16 +253,18 @@ runtime::Error Decoder::create_shared_planned_memory() { shared_method_memory_->spans.reserve(num_buffers); for (size_t i = 0; i < num_buffers; ++i) { - int64_t prefill_size = (i < prefill_bufs) - ? prefill_meta->memory_planned_buffer_size(i).get() - : 0; - int64_t decode_size = (i < decode_bufs) - ? decode_meta->memory_planned_buffer_size(i).get() - : 0; + int64_t prefill_size = 0; + if (i < prefill_bufs) { + prefill_size = ET_UNWRAP(prefill_meta->memory_planned_buffer_size(i)); + } + int64_t decode_size = 0; + if (i < decode_bufs) { + decode_size = ET_UNWRAP(decode_meta->memory_planned_buffer_size(i)); + } size_t max_size = static_cast(std::max(prefill_size, decode_size)); ET_LOG( Info, - "Prefill_size(%ld), decode_size(%ld) max_size(%ld)", + "Prefill_size(%" PRId64 "), decode_size(%" PRId64 ") max_size(%zu)", prefill_size, decode_size, max_size); diff --git a/examples/samsung/executor_runner/llm/runner/decoder.h b/examples/samsung/executor_runner/llm/runner/decoder.h index 6625d2fe98f..30e8f704113 100644 --- a/examples/samsung/executor_runner/llm/runner/decoder.h +++ b/examples/samsung/executor_runner/llm/runner/decoder.h @@ -43,6 +43,8 @@ class ET_EXPERIMENTAL Decoder : public llm::TextDecoderRunner { runtime::Error load() override; + bool is_method_loaded() override; + inline void set_method(std::string& method) { method_names_.insert(method); } @@ -52,7 +54,12 @@ class ET_EXPERIMENTAL Decoder : public llm::TextDecoderRunner { std::string decode_method_; std::unordered_set method_names_; std::unique_ptr shared_method_memory_; + // Keeps the TensorImpl of the logits view returned by step() alive. + extension::TensorPtr last_logits_; + runtime::Result slice_kth_token( + const runtime::EValue& output, + int64_t k); void extract_llm_methods(); runtime::Error load_io_manager(); runtime::Error create_shared_planned_memory(); diff --git a/examples/samsung/executor_runner/llm/runner/io_manager.cpp b/examples/samsung/executor_runner/llm/runner/io_manager.cpp index 1dfc426f11d..b4c2fe05f1b 100644 --- a/examples/samsung/executor_runner/llm/runner/io_manager.cpp +++ b/examples/samsung/executor_runner/llm/runner/io_manager.cpp @@ -7,6 +7,7 @@ */ #include +#include #include #include @@ -46,18 +47,7 @@ LlmIOManager::LlmIOManager( : IOManager(module), config_(config) { rotary_embedding_ = std::make_unique(config); - ET_CHECK_MSG( - config_.model_args->get_max_seq_len() >= - config_.model_args->get_prefill_seq_len(), - "Inconsistent configuration: max_seq_len (%zu) must be greater than or equal to prefill_seq_len (%zu)", - config_.model_args->get_max_seq_len(), - config_.model_args->get_prefill_seq_len()); - - ET_CHECK_MSG( - config_.model_args->get_max_seq_len() >= 1, - "Inconsistent configuration: max_seq_len must be at least 1, but got %zu", - config_.model_args->get_max_seq_len()); - + // RunnerConfig::validate() guarantees 1 <= prefill_seq_len <= max_seq_len. prefill_cache_len_ = config_.model_args->get_max_seq_len() - config_.model_args->get_prefill_seq_len(); decode_cache_len_ = config_.model_args->get_max_seq_len() - 1; @@ -138,7 +128,6 @@ ET_NODISCARD runtime::Error LlmIOManager::reset( (void)prefill_method; (void)decode_method; - dump_step_ = 0; current_pos_ = 0; prefill_done_ = false; @@ -190,8 +179,18 @@ LlmIOManager::prepare_prefill( append_prelude_inputs(prefill_seq_len, inputs); - const int64_t* start_position = start_pos->const_data_ptr(); - current_pos_ = *start_position; + const int64_t start_position = *start_pos->const_data_ptr(); + ET_CHECK_OR_RETURN_ERROR( + start_position >= 0 && + static_cast(start_position) + prompt_token_len <= + prefill_cache_len_, + InvalidArgument, + "Prompt does not fit: chunk at position %" PRId64 + " with %zu tokens exceeds the %zu-token prefill cache", + start_position, + prompt_token_len, + prefill_cache_len_); + current_pos_ = static_cast(start_position); pos_ids_.resize(prefill_seq_len); for (size_t i = 0; i < prefill_seq_len; ++i) { pos_ids_[i] = current_pos_ + static_cast(i); @@ -246,8 +245,6 @@ LlmIOManager::prepare_prefill( prefill_cache_len_, prefill_dummy_kv_, inputs); - dump_kv_input("prefill"); - return inputs; } executorch::runtime::Result> @@ -256,12 +253,19 @@ LlmIOManager::prepare_decode( const executorch::extension::TensorPtr& start_pos, const std::string& decode_method) { (void)decode_method; + const int64_t start_position = *start_pos->const_data_ptr(); + ET_CHECK_OR_RETURN_ERROR( + start_position >= 0 && + static_cast(start_position) <= decode_cache_len_, + InvalidArgument, + "Decode position %" PRId64 " exceeds max_seq_len %" PRId64, + start_position, + config_.model_args->get_max_seq_len()); if (!prefill_done_) { copy_prefill_to_decode_input(); prefill_done_ = true; } - const int64_t* start_position = start_pos->mutable_data_ptr(); - current_pos_ = *start_position; + current_pos_ = static_cast(start_position); update_attn_mask(current_pos_, config_.model_args->get_decode_seq_len()); std::vector inputs; inputs.reserve(get_input_reserve_size()); @@ -314,37 +318,26 @@ LlmIOManager::prepare_decode( decode_cache_len_, decode_dummy_kv_, inputs); - dump_kv_input("decode"); - return inputs; } executorch::runtime::Error LlmIOManager::update_prefill( const std::vector& model_outputs, const std::string& prefill_method) { (void)prefill_method; - update_kv_cache( + return update_kv_cache( model_outputs, prefill_k_input_, prefill_v_input_, prefill_cache_len_, - config_.model_args->get_prefill_seq_len(), - "prefill"); - return executorch::runtime::Error::Ok; + config_.model_args->get_prefill_seq_len()); } executorch::runtime::Error LlmIOManager::update_decode( const std::vector& model_outputs, const std::string& decode_method) { (void)decode_method; - update_kv_cache( - model_outputs, - decode_k_input_, - decode_v_input_, - decode_cache_len_, - 1, - "decode"); - return executorch::runtime::Error::Ok; + return update_kv_cache( + model_outputs, decode_k_input_, decode_v_input_, decode_cache_len_, 1); } -// TOOD: sizeof(float) dtype hard-codes the buffer dtype to be float. void LlmIOManager::allocate_buffers() { const size_t n_kv_heads = config_.model_args->get_n_kv_heads(); const size_t num_cache_layers = config_.model_args->get_num_cache_layers(); @@ -459,32 +452,20 @@ void LlmIOManager::update_attn_mask(size_t pos, size_t seq_len) { const size_t rows = config_.model_args->get_attn_mask_rows(); const size_t cols = config_.model_args->get_attn_mask_cols(); - const size_t start_col = cols - rows; const size_t total_size = rows * cols; - // Initialize entire mask to maskFalse (uses memset for 0.0f) fill_float_region(attn_mask_.data(), 0, total_size, maskFalse); if (seq_len > 1) { - // Prefill phase - const size_t rows = config_.model_args->get_attn_mask_rows(); - const size_t cols = config_.model_args->get_attn_mask_cols(); - const size_t start_col = config_.model_args->get_attn_mask_cols() - - config_.model_args->get_attn_mask_rows(); - - for (int i = 0; i < seq_len; ++i) { - int start_idx = cols - rows; - int end_idx = start_idx + seq_len; - - for (int j = start_idx; j < end_idx && j <= start_idx + i; ++j) { - attn_mask_[i * cols + j] = maskTrue; - } + // Prefill: past positions [0, pos) and the causal block for the chunk, + // which occupies the last `rows` columns. + const size_t start_col = cols - rows; + for (size_t i = 0; i < seq_len; ++i) { + fill_float_region( + attn_mask_.data(), i * cols + start_col, i + 1, maskTrue); } - - for (int i = 0; i < rows; ++i) { - for (int j = 0; j < pos; ++j) { - attn_mask_[i * cols + j] = maskTrue; - } + for (size_t i = 0; i < rows; ++i) { + fill_float_region(attn_mask_.data(), i * cols, pos, maskTrue); } } else { const size_t left_len = std::min(pos, cols); @@ -551,100 +532,57 @@ void LlmIOManager::update_local_attn_mask(size_t pos, size_t seq_len) { } if (cols > 0) { - attn_mask_[cols - 1] = maskTrue; + local_attn_mask_[cols - 1] = maskTrue; } } } -void LlmIOManager::dump_buffer( - const std::string& tag, - size_t layer, - const char* data, - size_t size) { - if (!dump_enabled_) { - return; - } - char filename[512]; - std::snprintf( - filename, - sizeof(filename), - "%s/%s_layer%zu_pos%zu_step%zu.bin", - dump_dir_.c_str(), - tag.c_str(), - layer, - current_pos_, - dump_step_); - FILE* f = std::fopen(filename, "wb"); - if (f) { - std::fwrite(data, 1, size, f); - std::fclose(f); - } -} -void LlmIOManager::dump_kv_input(const std::string& phase) { - if (!dump_enabled_) { - return; - } - const bool is_prefill = (phase == "prefill"); - const auto& k_bufs = is_prefill ? prefill_k_input_ : decode_k_input_; - const auto& v_bufs = is_prefill ? prefill_v_input_ : decode_v_input_; - const size_t num_cache_layers = config_.model_args->get_num_cache_layers(); - for (size_t l = 0; l < num_cache_layers; l++) { - dump_buffer(phase + "_k_input", l, k_bufs[l].data(), k_bufs[l].size()); - dump_buffer(phase + "_v_input", l, v_bufs[l].data(), v_bufs[l].size()); - } -} -void LlmIOManager::update_kv_cache( +executorch::runtime::Error LlmIOManager::update_kv_cache( const std::vector& model_outputs, std::vector>& k_bufs, std::vector>& v_bufs, size_t cache_len, - size_t new_len, - const std::string& phase) { + size_t new_len) { // model_outputs: [logits, k0..k_{C-1}, v0..v_{C-1}] where C is the number of // cached (self-decoder) layers; KV-shared layers emit no new cache. const size_t num_cache_layers = config_.model_args->get_num_cache_layers(); - for (size_t l = 0; l < num_cache_layers; l++) { - copy_tensor_to_cache( - model_outputs[1 + l].toTensor(), - new_len, - k_bufs[l].data(), - cache_len, - current_pos_, - config_.model_args->get_head_dim_for_layer(l)); + ET_CHECK_OR_RETURN_ERROR( + model_outputs.size() == 1 + 2 * num_cache_layers, + InvalidProgram, + "Expected %zu outputs (logits + K/V for %zu layers), got %zu", + 1 + 2 * num_cache_layers, + num_cache_layers, + model_outputs.size()); + // The last decode step may run at pos == cache_len; its KV is never read. + if (current_pos_ >= cache_len) { + return executorch::runtime::Error::Ok; } - for (size_t l = 0; l < num_cache_layers; l++) { - copy_tensor_to_cache( - model_outputs[1 + num_cache_layers + l].toTensor(), + const size_t rows = std::min(new_len, cache_len - current_pos_); + const size_t n_kv_heads = config_.model_args->get_n_kv_heads(); + for (size_t i = 0; i < 2 * num_cache_layers; i++) { + const size_t l = i % num_cache_layers; + const size_t head_dim = config_.model_args->get_head_dim_for_layer(l); + const runtime::EValue& out = model_outputs[1 + i]; + ET_CHECK_OR_RETURN_ERROR( + out.isTensor() && + out.toTensor().scalar_type() == aten::ScalarType::Float && + static_cast(out.toTensor().numel()) == + n_kv_heads * new_len * head_dim, + InvalidProgram, + "KV output %zu must be a float tensor of %zu elements", + 1 + i, + n_kv_heads * new_len * head_dim); + auto& bufs = i < num_cache_layers ? k_bufs : v_bufs; + copy_to_cache( + out.toTensor().const_data_ptr(), new_len, - v_bufs[l].data(), + rows, + bufs[l].data(), cache_len, current_pos_, - config_.model_args->get_head_dim_for_layer(l)); - } - dump_kv_output(phase, model_outputs); - dump_step_++; -} -void LlmIOManager::dump_kv_output( - const std::string& phase, - const std::vector& model_outputs) { - if (!dump_enabled_) { - return; - } - const size_t num_cache_layers = config_.model_args->get_num_cache_layers(); - for (size_t l = 0; l < num_cache_layers; l++) { - const auto& k_out = model_outputs[1 + l].toTensor(); - const auto& v_out = model_outputs[1 + num_cache_layers + l].toTensor(); - dump_buffer( - phase + "_k_output", - l, - reinterpret_cast(k_out.const_data_ptr()), - k_out.numel() * sizeof(float)); - dump_buffer( - phase + "_v_output", - l, - reinterpret_cast(v_out.const_data_ptr()), - v_out.numel() * sizeof(float)); + head_dim); } + return executorch::runtime::Error::Ok; } executorch::runtime::Error LlmIOManager::prepare_rotary_embedding_inputs( diff --git a/examples/samsung/executor_runner/llm/runner/io_manager.h b/examples/samsung/executor_runner/llm/runner/io_manager.h index 2b410746752..493cee16980 100644 --- a/examples/samsung/executor_runner/llm/runner/io_manager.h +++ b/examples/samsung/executor_runner/llm/runner/io_manager.h @@ -127,20 +127,14 @@ class ET_EXPERIMENTAL LlmIOManager : public llm::IOManager { } size_t get_attention_setup_tensor_count() const; - /** - * Calculates the total number of tensor slots required in the model's input - * buffer. - * - * Memory Layout Layout: - * [0] : Token/Prompt IDs - * [1 ... N] : Attention Setup Overheads (Masks, RoPE Sin/Cos - * Tensors) [N+1 ... End] : Alternating Key and Value Cache Tensors for - * all Layers - */ + + // [tokens, (Gemma 4: inputs_embeds, mm_mask), masks, RoPE cos/sin, K x L, + // V x L] size_t get_input_reserve_size() const { - auto model_config = getModelVersionConfig(config_.model_version); - return model_config.v_cache_tensor_index + - model_config.kv_cache_tensor_count; + const size_t prelude = + config_.model_version == DecoderModelVersion::kGemma4 ? 2 : 0; + return 1 + prelude + get_attention_setup_tensor_count() + + 2 * static_cast(config_.model_args->get_n_layers()); } size_t current_pos() const { @@ -148,22 +142,8 @@ class ET_EXPERIMENTAL LlmIOManager : public llm::IOManager { } private: - void dump_buffer( - const std::string& tag, - size_t layer, - const char* data, - size_t size); - void dump_kv_input(const std::string& phase); - void dump_kv_output( - const std::string& phase, - const std::vector& model_outputs); - void allocate_buffers(); - bool dump_enabled_ = false; - mutable size_t dump_step_ = 0; - std::string dump_dir_ = "/data/local/tmp/kv_cache_dump"; - /** * Copy prefill output KV cache (layout [1,H,prefill_cache_len,D]) into * the decode input buffer (layout [1,H,decode_cache_len,D]). @@ -173,35 +153,30 @@ class ET_EXPERIMENTAL LlmIOManager : public llm::IOManager { void copy_prefill_to_decode_input(); /** - * Copy model KV-cache outputs into the given buffers and dump them. + * Copy model KV-cache outputs into the given buffers. * * @param model_outputs [logits, k0..k_{L-1}, v0..v_{L-1}] * @param k_bufs per-layer K cache buffers (input/output) * @param v_bufs per-layer V cache buffers (input/output) * @param cache_len length of each cache buffer * @param new_len number of new tokens to copy (prefill_seq_len or 1) - * @param phase dump label ("prefill" or "decode") */ - void update_kv_cache( + ET_NODISCARD runtime::Error update_kv_cache( const std::vector& model_outputs, std::vector>& k_bufs, std::vector>& v_bufs, size_t cache_len, - size_t new_len, - const std::string& phase); + size_t new_len); /** - * Copy new cache entries from source data (layout [1,H,new_len,D]) into the - * cache buffer (layout [1,H,cache_len,D]) at the given position. - * - * For each head h: - * cache[h * cache_len * D + pos * D .. + (pos + new_len) * D] - * = src[h * new_len * D .. + new_len * D] + * Copy the first `rows` of new cache entries from source data (layout + * [1,H,new_len,D]) into the cache buffer (layout [1,H,cache_len,D]) at pos. */ template void copy_to_cache( const T* src, size_t new_len, + size_t rows, char* cache, size_t cache_len, size_t pos, @@ -210,25 +185,10 @@ class ET_EXPERIMENTAL LlmIOManager : public llm::IOManager { const T* src_head = src + h * new_len * head_dim; char* dst_head = cache + h * cache_len * head_dim * sizeof(T) + pos * head_dim * sizeof(T); - std::memcpy(dst_head, src_head, new_len * head_dim * sizeof(T)); + std::memcpy(dst_head, src_head, rows * head_dim * sizeof(T)); } } - /** - * Copy tensor data to cache buffer, dispatching based on kv_scalar_type. - * Handles both Half (float16) and Float (float32) types. - */ - void copy_tensor_to_cache( - const executorch::aten::Tensor& tensor, - size_t new_len, - char* cache, - size_t cache_len, - size_t pos, - size_t head_dim) { - const float* data = tensor.const_data_ptr(); - copy_to_cache(data, new_len, cache, cache_len, pos, head_dim); - } - /** * Append KV-cache input tensors for all layers. * diff --git a/examples/samsung/executor_runner/llm/runner/prefiller.cpp b/examples/samsung/executor_runner/llm/runner/prefiller.cpp index 57eefb96ef9..4f78405cd59 100644 --- a/examples/samsung/executor_runner/llm/runner/prefiller.cpp +++ b/examples/samsung/executor_runner/llm/runner/prefiller.cpp @@ -37,7 +37,8 @@ Prefiller::Prefiller( runtime::Result Prefiller::prefill( std::vector& prompt_tokens, int64_t& start_pos) { - ET_CHECK_MSG(!prompt_tokens.empty(), "Prompt cannot be null"); + ET_CHECK_OR_RETURN_ERROR( + !prompt_tokens.empty(), InvalidArgument, "Prompt cannot be empty"); if (!text_decoder_runner_->is_method_loaded()) { ET_CHECK_OK_OR_RETURN_ERROR(text_decoder_runner_->load()); } @@ -117,8 +118,6 @@ runtime::Result Prefiller::prefill_chunk( cur_token, start_pos); - // initialize tensor wrappers - // TODO(sicheon.oh) : Set tensor size by metadata::get_prefill_ar_len auto tokens = extension::from_blob(&cur_token, {1, 1}, aten::ScalarType::Long); diff --git a/examples/samsung/executor_runner/llm/runner/rotary_embedding.cpp b/examples/samsung/executor_runner/llm/runner/rotary_embedding.cpp index a085e4552ca..ae8755586fb 100644 --- a/examples/samsung/executor_runner/llm/runner/rotary_embedding.cpp +++ b/examples/samsung/executor_runner/llm/runner/rotary_embedding.cpp @@ -7,10 +7,10 @@ */ #include +#include #include #include #include -#include #include @@ -57,7 +57,6 @@ RotaryEmbedding::RotaryEmbedding(const RunnerConfig& config) } } - // TODO(brucekk.kim): size should be configurable global_bytes_ = config_.model_args->get_max_seq_len() * normalized_global_head_dim_ * sizeof(float); local_bytes_ = config_.model_args->get_max_seq_len() * @@ -132,6 +131,8 @@ runtime::Error RotaryEmbedding::load_from_file(const std::string& path) { file.read(local_cos_.get(), local_bytes_); file.read(local_sin_.get(), local_bytes_); } + ET_CHECK_OR_RETURN_ERROR( + file.good(), InvalidArgument, "Failed to read '%s'.", path.c_str()); loaded_ = true; return runtime::Error::Ok; @@ -148,13 +149,13 @@ runtime::Error RotaryEmbedding::load_by_precompute_rope_tables() { reinterpret_cast(global_sin_ptr) % alignof(float) == 0 && "global_sin_ is misaligned for float access!"); - compute_rope_table( + ET_CHECK_OK_OR_RETURN_ERROR(compute_rope_table( global_cos_ptr, global_sin_ptr, seq_len_, normalized_global_head_dim_, base_freq_global_, - rope_percentage_); + rope_percentage_)); if (config_.embedding_strategy == EmbeddingStrategy::kGlobalAndLocal) { float* local_cos_ptr = reinterpret_cast(local_cos_.get()); @@ -169,13 +170,13 @@ runtime::Error RotaryEmbedding::load_by_precompute_rope_tables() { // Sliding/local layers use full RoPE; the partial rotary factor applies // only to the global (full-attention) layers. - compute_rope_table( + ET_CHECK_OK_OR_RETURN_ERROR(compute_rope_table( local_cos_ptr, local_sin_ptr, seq_len_, normalized_local_head_dim_, base_freq_local_, - /*rope_pct=*/1.0f); + /*rope_pct=*/1.0f)); } loaded_ = true; ET_LOG(Info, "Rotary Embedding tables precomputed successfully."); @@ -293,6 +294,14 @@ runtime::Error RotaryEmbedding::update( return runtime::Error::InvalidArgument; } + ET_CHECK_OR_RETURN_ERROR( + pos_ids[0] >= 0 && + static_cast(pos_ids[0]) + pos_ids.size() <= seq_len_, + InvalidArgument, + "RoPE positions [%" PRId64 ", +%zu) exceed the %zu-entry table", + pos_ids[0], + pos_ids.size(), + seq_len_); size_t pos = static_cast(pos_ids[0]); size_t length = pos_ids.size(); @@ -307,7 +316,7 @@ runtime::Error RotaryEmbedding::update( // Local Copy if (config_.embedding_strategy == EmbeddingStrategy::kGlobalAndLocal) { - size_t l_head_dim = config_.model_args->get_head_dim(); + size_t l_head_dim = normalized_local_head_dim_; size_t l_bytes_to_copy = length * l_head_dim * sizeof(float); size_t l_element_offset = pos * l_head_dim * sizeof(float); std::memcpy( diff --git a/examples/samsung/executor_runner/llm/runner/runner.cpp b/examples/samsung/executor_runner/llm/runner/runner.cpp index 807ef910911..8a0dccf9180 100644 --- a/examples/samsung/executor_runner/llm/runner/runner.cpp +++ b/examples/samsung/executor_runner/llm/runner/runner.cpp @@ -1,5 +1,5 @@ /* - * Copyright (c) Samsung Platforms, Inc. and affiliates. + * Copyright (c) 2025 Samsung Electronics Co. LTD * All rights reserved. * * This source code is licensed under the BSD-style license found in the @@ -9,7 +9,6 @@ // A simple LLM runner that includes preprocessing and post processing logic. // The module takes in a string as input and emits a string as output. -#include #include #include #include @@ -25,6 +24,8 @@ #include #include #include + +#include #include namespace example { @@ -114,7 +115,9 @@ std::unique_ptr create_llm_runner( runner_paths.model_path.c_str(), runner_paths.tokenizer_path.c_str()); - runner_config.load_from_json(); + if (runner_config.load_from_json() != runtime::Error::Ok) { + return nullptr; + } std::unique_ptr<::tokenizers::Tokenizer> tokenizer = load_tokenizer( runner_config.tokenizer_type, @@ -148,7 +151,7 @@ static const std::unordered_map kMetadataMap{ {llm::kVocabSize, "get_vocab_size"}, {llm::kUseKVCache, "use_kv_cache"}, {llm::kUseSDPAWithKVCache, "use_sdpa_with_kv_cache"}, - {llm::kHeadDim, "get_head_dim"}, + {example::kHeadDim, "get_head_dim"}, {example::kPrefillArLen, "get_prefill_ar_len"}, {example::kNumOfHiddenDim, "get_hidden_dim"}, {example::kNumOfHeads, "get_n_heads"}, @@ -221,7 +224,8 @@ runtime::Error LLMRunner::generate_from_pos( std::function stats_callback) { // Prepare the inputs. // Use ones-initialized inputs. - ET_CHECK_MSG(!prompt.empty(), "Prompt cannot be null"); + ET_CHECK_OR_RETURN_ERROR( + !prompt.empty(), InvalidArgument, "Prompt cannot be empty"); if (!is_loaded()) { stats_->model_load_start_ms = llm::time_in_ms(); ET_CHECK_OK_OR_RETURN_ERROR(load()); @@ -254,11 +258,6 @@ runtime::Error LLMRunner::generate_from_pos( stats_->inference_start_ms = llm::time_in_ms(); text_token_generator_->clear_stop(); - ET_LOG( - Info, - "################ input prompt ################\n %s\n", - prompt.c_str()); - // Raw prompts are already fully formatted (special tokens, system turn, tool // declaration); feed them verbatim and skip every template transformation. std::string input_prompt = config_.raw_prompt @@ -300,7 +299,6 @@ runtime::Error LLMRunner::generate_from_pos( if (model_version_ == DecoderModelVersion::kGemma4 && vision_encoder_ != nullptr && !config_.image_path.empty()) { auto* llm_io = static_cast(io_manager_.get()); - ET_CHECK_MSG(llm_io != nullptr, "IOManager is null"); // Vision embeddings are injected only by the multi-token prefill method; // single-token steps go through decode and would silently drop them. @@ -357,9 +355,12 @@ runtime::Error LLMRunner::generate_from_pos( int num_prompt_tokens = prompt_tokens.size(); - // Reduce max_context_len by start_pos - // TODO(brucekk.kim): Check where to fetch from: metatdata or runner_config - int64_t max_context_len = metadata_.at(llm::kMaxContextLen) - start_pos; + // The KV cache and RoPE tables only hold max_seq_len positions, whatever + // context length the model was trained with. + int64_t max_context_len = + std::min( + metadata_.at(llm::kMaxContextLen), metadata_.at(llm::kMaxSeqLen)) - + start_pos; ET_CHECK_OR_RETURN_ERROR( num_prompt_tokens >= 1, InvalidArgument, @@ -424,8 +425,6 @@ runtime::Error LLMRunner::generate_from_pos( ::tokenizers::Result> newline_token = tokenizer_->encode("\n"); - ET_LOG(Info, "[ENN] first token %lu\n", cur_token); - ET_LOG( Info, "################ The start of token generation ################\n"); @@ -456,7 +455,7 @@ runtime::Error LLMRunner::generate_from_pos( }; int64_t num_generated_tokens = ET_UNWRAP(text_token_generator_->generate( prompt_tokens, - num_prompt_tokens, + pos, max_new_tokens - 1, temperature_ == -1.0f ? config.temperature : temperature_, collect_callback)); @@ -541,10 +540,11 @@ void LLMRunner::reset() { if (llm_io == nullptr) { return; } - auto err = llm_io->reset( - example::default_method_name_prefill, - example::default_method_name_decode); - (void)err; + if (llm_io->reset( + example::default_method_name_prefill, + example::default_method_name_decode) != runtime::Error::Ok) { + ET_LOG(Error, "Failed to reset the KV cache"); + } } std::unique_ptr load_tokenizer( @@ -596,7 +596,7 @@ std::unordered_map get_llm_metadata( config.model_args ? config.model_args->get_n_kv_heads() : 0}, {kMetadataMap.at(example::kNumOfLayers), config.model_args ? config.model_args->get_n_layers() : 0}, - {kMetadataMap.at(llm::kHeadDim), + {kMetadataMap.at(example::kHeadDim), config.model_args ? config.model_args->get_head_dim() : 0}, }); @@ -613,6 +613,10 @@ std::unordered_map get_llm_metadata( if (method_names.count(method_name)) { auto get_result = module->get(method_name); + if (!get_result.ok()) { + ET_LOG(Error, "Failed to read metadata %s", method_name.c_str()); + continue; + } value = get_result.get().toScalar().to(); } else { ET_LOG( @@ -684,7 +688,11 @@ std::unique_ptr create_llm_runner( executorch::runtime::Result> methods_res = module->method_names(); if (methods_res.error() != executorch::runtime::Error::Ok) { - ET_CHECK_MSG(false, "Failed to get method names"); + ET_LOG( + Error, + "Failed to get method names from %s", + runner_paths.model_path.c_str()); + return nullptr; } auto methods = methods_res.get(); @@ -746,7 +754,7 @@ std::unique_ptr create_llm_runner( } bool separated_prefill_decode = - (config.model_args->get_prefill_seq_len() > 1); + !is_decode_only && config.model_args->get_prefill_seq_len() > 1; ET_LOG( Info, "separated_prefill_decode(%s)", @@ -785,7 +793,10 @@ std::unique_ptr create_llm_runner( runner->set_vision_encoder(std::move(vision_encoder)); } - runner->load(); + if (runner->load() != executorch::runtime::Error::Ok) { + ET_LOG(Error, "Failed to load the LLM runner."); + return nullptr; + } return runner; } diff --git a/examples/samsung/executor_runner/llm/runner/runner.h b/examples/samsung/executor_runner/llm/runner/runner.h index e823c13cf86..3e006af868d 100644 --- a/examples/samsung/executor_runner/llm/runner/runner.h +++ b/examples/samsung/executor_runner/llm/runner/runner.h @@ -42,17 +42,16 @@ class LLMRunner : public llm::IRunner { DecoderModelVersion model_version, float temperature = -1.0f); - bool is_loaded() const; + bool is_loaded() const override; - runtime::Error load(); + runtime::Error load() override; runtime::Error generate( const std::string& prompt, const llm::GenerationConfig& config, std::function token_callback = {}, - std::function stats_callback = {}); + std::function stats_callback = {}) override; - // TODO (brucekk.kim): consider removing unncessary args runtime::Error generate_from_pos( const std::string& prompt, int64_t start_pos, @@ -62,7 +61,7 @@ class LLMRunner : public llm::IRunner { runtime::Error warmup(const std::string& prompt, int32_t max_new_tokens); - void stop(); + void stop() override; const RunnerConfig& get_config() const { return config_; @@ -100,14 +99,13 @@ class LLMRunner : public llm::IRunner { return vision_encoder_.get(); } - // TODO (brucekk.kim): should it be static? static std::string get_formatted_prompt( const std::string& prompt, const std::string& system_prompt, const std::string& assistant_prefill_prompt, DecoderModelVersion version); - void reset(); + void reset() override; private: RunnerConfig config_; diff --git a/examples/samsung/executor_runner/llm/runner/token_generator.h b/examples/samsung/executor_runner/llm/runner/token_generator.h index ce81f3413bd..04af9e7f457 100644 --- a/examples/samsung/executor_runner/llm/runner/token_generator.h +++ b/examples/samsung/executor_runner/llm/runner/token_generator.h @@ -59,8 +59,10 @@ class ET_EXPERIMENTAL TokenGenerator { int32_t max_new_tokens, float temperature = 0.0f, const std::function& token_callback = {}) { - ET_CHECK_MSG( - !tokens.empty(), "Token generation loop shouldn't take empty tokens"); + ET_CHECK_OR_RETURN_ERROR( + !tokens.empty(), + InvalidArgument, + "Token generation loop shouldn't take empty tokens"); int64_t pos = start_pos; // position in the sequence std::vector token_data; // allocate space for the tokens diff --git a/examples/samsung/executor_runner/run_gemma3_sample.sh b/examples/samsung/executor_runner/run_gemma3_sample.sh old mode 100644 new mode 100755 index f60ef702d12..3a88c64a62b --- a/examples/samsung/executor_runner/run_gemma3_sample.sh +++ b/examples/samsung/executor_runner/run_gemma3_sample.sh @@ -1,3 +1,4 @@ +#!/bin/bash # Copyright (c) 2026 Samsung Electronics Co. LTD # All rights reserved. # diff --git a/examples/samsung/executor_runner/run_llama3_sample.sh b/examples/samsung/executor_runner/run_llama3_sample.sh old mode 100644 new mode 100755 index da4af40d4d5..d2d02a5bd04 --- a/examples/samsung/executor_runner/run_llama3_sample.sh +++ b/examples/samsung/executor_runner/run_llama3_sample.sh @@ -1,3 +1,4 @@ +#!/bin/bash # Copyright (c) 2026 Samsung Electronics Co. LTD # All rights reserved. # @@ -26,5 +27,4 @@ chmod +x enn_executor_llm_runner --prompt="$INPUT_PROMPT" \ --temperature=$TEMPERATURE \ --seq_len=$SEQ_LEN \ - --num_bos=$NUM_BOS \ - + --num_bos=$NUM_BOS diff --git a/examples/models/gemma4/text_decoder/generate_pos_embedding_bin.py b/examples/samsung/scripts/gemma4_pos_embedding.py similarity index 89% rename from examples/models/gemma4/text_decoder/generate_pos_embedding_bin.py rename to examples/samsung/scripts/gemma4_pos_embedding.py index 2138c80a30d..7a8ab838256 100644 --- a/examples/models/gemma4/text_decoder/generate_pos_embedding_bin.py +++ b/examples/samsung/scripts/gemma4_pos_embedding.py @@ -1,4 +1,9 @@ -#!/usr/bin/env python3 +# Copyright (c) 2026 Samsung Electronics Co. LTD +# All rights reserved +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + """Precompute the Gemma4 RoPE position-embedding .bin for the Exynos LLM runner. Reproduces the RoPE convention from gemma4_attention.py: inv_freq denominator is @@ -15,7 +20,8 @@ import argparse import torch -from gemma4_config import Gemma4Config + +from executorch.examples.models.gemma4.text_decoder.gemma4_config import Gemma4Config def rope_table( diff --git a/extension/llm/runner/constants.h b/extension/llm/runner/constants.h index b1887948cbf..60a6383a72e 100644 --- a/extension/llm/runner/constants.h +++ b/extension/llm/runner/constants.h @@ -24,7 +24,6 @@ inline constexpr auto kHeadDims = "get_head_dims"; inline constexpr auto kWindows = "get_windows"; inline constexpr auto kUseKVCache = "use_kv_cache"; inline constexpr auto kUseSDPAWithKVCache = "use_sdpa_with_kv_cache"; -inline constexpr auto kHeadDim = "get_head_dim"; // Multimodal method name conventions inline constexpr auto kVisionEncoderMethod = "vision_encoder"; diff --git a/extension/llm/runner/text_decoder_runner.cpp b/extension/llm/runner/text_decoder_runner.cpp index 25f3fc234ad..bb64d79dd2e 100644 --- a/extension/llm/runner/text_decoder_runner.cpp +++ b/extension/llm/runner/text_decoder_runner.cpp @@ -40,22 +40,32 @@ TextDecoderRunner::TextDecoderRunner( ::executorch::runtime::Result TextDecoderRunner::step( TensorPtr& tokens, int64_t start_pos) { - auto method_meta = ET_UNWRAP(module_->method_meta(method_name_.c_str())); + // ET_LOG(Info, "Input token %" PRIu64, input_token); + auto method_meta_result = module_->method_meta(method_name_); + if (!method_meta_result.ok()) { + return method_meta_result.error(); + } + auto method_meta = std::move(*method_meta_result); // If only 1 input, we are not using kv cache bool use_kv_cache = method_meta.num_inputs() > 1; std::vector cache_positions; if (use_kv_cache) { - auto start_pos_tensor = ET_UNWRAP(populate_start_pos_or_cache_position( + auto start_pos_tensor_result = populate_start_pos_or_cache_position( module_, start_pos, cache_positions, tokens->numel(), - method_name_.c_str())); + method_name_.c_str()); + if (!start_pos_tensor_result.ok()) { + return start_pos_tensor_result.error(); + } + auto start_pos_tensor = std::move(*start_pos_tensor_result); std::vector inputs; - auto inputs_res = io_manager_->prepare_decode(tokens, start_pos_tensor); + auto inputs_res = + io_manager_->prepare_decode(tokens, start_pos_tensor, method_name_); ET_CHECK_OK_OR_RETURN_ERROR(inputs_res.error()); inputs = inputs_res.get(); if (stats_ != nullptr) { diff --git a/third-party/stb/stb_image.h b/third-party/stb/stb_image.h deleted file mode 100644 index 9eedabedc45..00000000000 --- a/third-party/stb/stb_image.h +++ /dev/null @@ -1,7988 +0,0 @@ -/* stb_image - v2.30 - public domain image loader - http://nothings.org/stb - no warranty implied; use at your own risk - - Do this: - #define STB_IMAGE_IMPLEMENTATION - before you include this file in *one* C or C++ file to create the implementation. - - // i.e. it should look like this: - #include ... - #include ... - #include ... - #define STB_IMAGE_IMPLEMENTATION - #include "stb_image.h" - - You can #define STBI_ASSERT(x) before the #include to avoid using assert.h. - And #define STBI_MALLOC, STBI_REALLOC, and STBI_FREE to avoid using malloc,realloc,free - - - QUICK NOTES: - Primarily of interest to game developers and other people who can - avoid problematic images and only need the trivial interface - - JPEG baseline & progressive (12 bpc/arithmetic not supported, same as stock IJG lib) - PNG 1/2/4/8/16-bit-per-channel - - TGA (not sure what subset, if a subset) - BMP non-1bpp, non-RLE - PSD (composited view only, no extra channels, 8/16 bit-per-channel) - - GIF (*comp always reports as 4-channel) - HDR (radiance rgbE format) - PIC (Softimage PIC) - PNM (PPM and PGM binary only) - - Animated GIF still needs a proper API, but here's one way to do it: - http://gist.github.com/urraka/685d9a6340b26b830d49 - - - decode from memory or through FILE (define STBI_NO_STDIO to remove code) - - decode from arbitrary I/O callbacks - - SIMD acceleration on x86/x64 (SSE2) and ARM (NEON) - - Full documentation under "DOCUMENTATION" below. - - -LICENSE - - See end of file for license information. - -RECENT REVISION HISTORY: - - 2.30 (2024-05-31) avoid erroneous gcc warning - 2.29 (2023-05-xx) optimizations - 2.28 (2023-01-29) many error fixes, security errors, just tons of stuff - 2.27 (2021-07-11) document stbi_info better, 16-bit PNM support, bug fixes - 2.26 (2020-07-13) many minor fixes - 2.25 (2020-02-02) fix warnings - 2.24 (2020-02-02) fix warnings; thread-local failure_reason and flip_vertically - 2.23 (2019-08-11) fix clang static analysis warning - 2.22 (2019-03-04) gif fixes, fix warnings - 2.21 (2019-02-25) fix typo in comment - 2.20 (2019-02-07) support utf8 filenames in Windows; fix warnings and platform ifdefs - 2.19 (2018-02-11) fix warning - 2.18 (2018-01-30) fix warnings - 2.17 (2018-01-29) bugfix, 1-bit BMP, 16-bitness query, fix warnings - 2.16 (2017-07-23) all functions have 16-bit variants; optimizations; bugfixes - 2.15 (2017-03-18) fix png-1,2,4; all Imagenet JPGs; no runtime SSE detection on GCC - 2.14 (2017-03-03) remove deprecated STBI_JPEG_OLD; fixes for Imagenet JPGs - 2.13 (2016-12-04) experimental 16-bit API, only for PNG so far; fixes - 2.12 (2016-04-02) fix typo in 2.11 PSD fix that caused crashes - 2.11 (2016-04-02) 16-bit PNGS; enable SSE2 in non-gcc x64 - RGB-format JPEG; remove white matting in PSD; - allocate large structures on the stack; - correct channel count for PNG & BMP - 2.10 (2016-01-22) avoid warning introduced in 2.09 - 2.09 (2016-01-16) 16-bit TGA; comments in PNM files; STBI_REALLOC_SIZED - - See end of file for full revision history. - - - ============================ Contributors ========================= - - Image formats Extensions, features - Sean Barrett (jpeg, png, bmp) Jetro Lauha (stbi_info) - Nicolas Schulz (hdr, psd) Martin "SpartanJ" Golini (stbi_info) - Jonathan Dummer (tga) James "moose2000" Brown (iPhone PNG) - Jean-Marc Lienher (gif) Ben "Disch" Wenger (io callbacks) - Tom Seddon (pic) Omar Cornut (1/2/4-bit PNG) - Thatcher Ulrich (psd) Nicolas Guillemot (vertical flip) - Ken Miller (pgm, ppm) Richard Mitton (16-bit PSD) - github:urraka (animated gif) Junggon Kim (PNM comments) - Christopher Forseth (animated gif) Daniel Gibson (16-bit TGA) - socks-the-fox (16-bit PNG) - Jeremy Sawicki (handle all ImageNet JPGs) - Optimizations & bugfixes Mikhail Morozov (1-bit BMP) - Fabian "ryg" Giesen Anael Seghezzi (is-16-bit query) - Arseny Kapoulkine Simon Breuss (16-bit PNM) - John-Mark Allen - Carmelo J Fdez-Aguera - - Bug & warning fixes - Marc LeBlanc David Woo Guillaume George Martins Mozeiko - Christpher Lloyd Jerry Jansson Joseph Thomson Blazej Dariusz Roszkowski - Phil Jordan Dave Moore Roy Eltham - Hayaki Saito Nathan Reed Won Chun - Luke Graham Johan Duparc Nick Verigakis the Horde3D community - Thomas Ruf Ronny Chevalier github:rlyeh - Janez Zemva John Bartholomew Michal Cichon github:romigrou - Jonathan Blow Ken Hamada Tero Hanninen github:svdijk - Eugene Golushkov Laurent Gomila Cort Stratton github:snagar - Aruelien Pocheville Sergio Gonzalez Thibault Reuille github:Zelex - Cass Everitt Ryamond Barbiero github:grim210 - Paul Du Bois Engin Manap Aldo Culquicondor github:sammyhw - Philipp Wiesemann Dale Weiler Oriol Ferrer Mesia github:phprus - Josh Tobin Neil Bickford Matthew Gregan github:poppolopoppo - Julian Raschke Gregory Mullen Christian Floisand github:darealshinji - Baldur Karlsson Kevin Schmidt JR Smith github:Michaelangel007 - Brad Weinberger Matvey Cherevko github:mosra - Luca Sas Alexander Veselov Zack Middleton [reserved] - Ryan C. Gordon [reserved] [reserved] - DO NOT ADD YOUR NAME HERE - - Jacko Dirks - - To add your name to the credits, pick a random blank space in the middle and fill it. - 80% of merge conflicts on stb PRs are due to people adding their name at the end - of the credits. -*/ - -#ifndef STBI_INCLUDE_STB_IMAGE_H -#define STBI_INCLUDE_STB_IMAGE_H - -// DOCUMENTATION -// -// Limitations: -// - no 12-bit-per-channel JPEG -// - no JPEGs with arithmetic coding -// - GIF always returns *comp=4 -// -// Basic usage (see HDR discussion below for HDR usage): -// int x,y,n; -// unsigned char *data = stbi_load(filename, &x, &y, &n, 0); -// // ... process data if not NULL ... -// // ... x = width, y = height, n = # 8-bit components per pixel ... -// // ... replace '0' with '1'..'4' to force that many components per pixel -// // ... but 'n' will always be the number that it would have been if you said 0 -// stbi_image_free(data); -// -// Standard parameters: -// int *x -- outputs image width in pixels -// int *y -- outputs image height in pixels -// int *channels_in_file -- outputs # of image components in image file -// int desired_channels -- if non-zero, # of image components requested in result -// -// The return value from an image loader is an 'unsigned char *' which points -// to the pixel data, or NULL on an allocation failure or if the image is -// corrupt or invalid. The pixel data consists of *y scanlines of *x pixels, -// with each pixel consisting of N interleaved 8-bit components; the first -// pixel pointed to is top-left-most in the image. There is no padding between -// image scanlines or between pixels, regardless of format. The number of -// components N is 'desired_channels' if desired_channels is non-zero, or -// *channels_in_file otherwise. If desired_channels is non-zero, -// *channels_in_file has the number of components that _would_ have been -// output otherwise. E.g. if you set desired_channels to 4, you will always -// get RGBA output, but you can check *channels_in_file to see if it's trivially -// opaque because e.g. there were only 3 channels in the source image. -// -// An output image with N components has the following components interleaved -// in this order in each pixel: -// -// N=#comp components -// 1 grey -// 2 grey, alpha -// 3 red, green, blue -// 4 red, green, blue, alpha -// -// If image loading fails for any reason, the return value will be NULL, -// and *x, *y, *channels_in_file will be unchanged. The function -// stbi_failure_reason() can be queried for an extremely brief, end-user -// unfriendly explanation of why the load failed. Define STBI_NO_FAILURE_STRINGS -// to avoid compiling these strings at all, and STBI_FAILURE_USERMSG to get slightly -// more user-friendly ones. -// -// Paletted PNG, BMP, GIF, and PIC images are automatically depalettized. -// -// To query the width, height and component count of an image without having to -// decode the full file, you can use the stbi_info family of functions: -// -// int x,y,n,ok; -// ok = stbi_info(filename, &x, &y, &n); -// // returns ok=1 and sets x, y, n if image is a supported format, -// // 0 otherwise. -// -// Note that stb_image pervasively uses ints in its public API for sizes, -// including sizes of memory buffers. This is now part of the API and thus -// hard to change without causing breakage. As a result, the various image -// loaders all have certain limits on image size; these differ somewhat -// by format but generally boil down to either just under 2GB or just under -// 1GB. When the decoded image would be larger than this, stb_image decoding -// will fail. -// -// Additionally, stb_image will reject image files that have any of their -// dimensions set to a larger value than the configurable STBI_MAX_DIMENSIONS, -// which defaults to 2**24 = 16777216 pixels. Due to the above memory limit, -// the only way to have an image with such dimensions load correctly -// is for it to have a rather extreme aspect ratio. Either way, the -// assumption here is that such larger images are likely to be malformed -// or malicious. If you do need to load an image with individual dimensions -// larger than that, and it still fits in the overall size limit, you can -// #define STBI_MAX_DIMENSIONS on your own to be something larger. -// -// =========================================================================== -// -// UNICODE: -// -// If compiling for Windows and you wish to use Unicode filenames, compile -// with -// #define STBI_WINDOWS_UTF8 -// and pass utf8-encoded filenames. Call stbi_convert_wchar_to_utf8 to convert -// Windows wchar_t filenames to utf8. -// -// =========================================================================== -// -// Philosophy -// -// stb libraries are designed with the following priorities: -// -// 1. easy to use -// 2. easy to maintain -// 3. good performance -// -// Sometimes I let "good performance" creep up in priority over "easy to maintain", -// and for best performance I may provide less-easy-to-use APIs that give higher -// performance, in addition to the easy-to-use ones. Nevertheless, it's important -// to keep in mind that from the standpoint of you, a client of this library, -// all you care about is #1 and #3, and stb libraries DO NOT emphasize #3 above all. -// -// Some secondary priorities arise directly from the first two, some of which -// provide more explicit reasons why performance can't be emphasized. -// -// - Portable ("ease of use") -// - Small source code footprint ("easy to maintain") -// - No dependencies ("ease of use") -// -// =========================================================================== -// -// I/O callbacks -// -// I/O callbacks allow you to read from arbitrary sources, like packaged -// files or some other source. Data read from callbacks are processed -// through a small internal buffer (currently 128 bytes) to try to reduce -// overhead. -// -// The three functions you must define are "read" (reads some bytes of data), -// "skip" (skips some bytes of data), "eof" (reports if the stream is at the end). -// -// =========================================================================== -// -// SIMD support -// -// The JPEG decoder will try to automatically use SIMD kernels on x86 when -// supported by the compiler. For ARM Neon support, you must explicitly -// request it. -// -// (The old do-it-yourself SIMD API is no longer supported in the current -// code.) -// -// On x86, SSE2 will automatically be used when available based on a run-time -// test; if not, the generic C versions are used as a fall-back. On ARM targets, -// the typical path is to have separate builds for NEON and non-NEON devices -// (at least this is true for iOS and Android). Therefore, the NEON support is -// toggled by a build flag: define STBI_NEON to get NEON loops. -// -// If for some reason you do not want to use any of SIMD code, or if -// you have issues compiling it, you can disable it entirely by -// defining STBI_NO_SIMD. -// -// =========================================================================== -// -// HDR image support (disable by defining STBI_NO_HDR) -// -// stb_image supports loading HDR images in general, and currently the Radiance -// .HDR file format specifically. You can still load any file through the existing -// interface; if you attempt to load an HDR file, it will be automatically remapped -// to LDR, assuming gamma 2.2 and an arbitrary scale factor defaulting to 1; -// both of these constants can be reconfigured through this interface: -// -// stbi_hdr_to_ldr_gamma(2.2f); -// stbi_hdr_to_ldr_scale(1.0f); -// -// (note, do not use _inverse_ constants; stbi_image will invert them -// appropriately). -// -// Additionally, there is a new, parallel interface for loading files as -// (linear) floats to preserve the full dynamic range: -// -// float *data = stbi_loadf(filename, &x, &y, &n, 0); -// -// If you load LDR images through this interface, those images will -// be promoted to floating point values, run through the inverse of -// constants corresponding to the above: -// -// stbi_ldr_to_hdr_scale(1.0f); -// stbi_ldr_to_hdr_gamma(2.2f); -// -// Finally, given a filename (or an open file or memory block--see header -// file for details) containing image data, you can query for the "most -// appropriate" interface to use (that is, whether the image is HDR or -// not), using: -// -// stbi_is_hdr(char *filename); -// -// =========================================================================== -// -// iPhone PNG support: -// -// We optionally support converting iPhone-formatted PNGs (which store -// premultiplied BGRA) back to RGB, even though they're internally encoded -// differently. To enable this conversion, call -// stbi_convert_iphone_png_to_rgb(1). -// -// Call stbi_set_unpremultiply_on_load(1) as well to force a divide per -// pixel to remove any premultiplied alpha *only* if the image file explicitly -// says there's premultiplied data (currently only happens in iPhone images, -// and only if iPhone convert-to-rgb processing is on). -// -// =========================================================================== -// -// ADDITIONAL CONFIGURATION -// -// - You can suppress implementation of any of the decoders to reduce -// your code footprint by #defining one or more of the following -// symbols before creating the implementation. -// -// STBI_NO_JPEG -// STBI_NO_PNG -// STBI_NO_BMP -// STBI_NO_PSD -// STBI_NO_TGA -// STBI_NO_GIF -// STBI_NO_HDR -// STBI_NO_PIC -// STBI_NO_PNM (.ppm and .pgm) -// -// - You can request *only* certain decoders and suppress all other ones -// (this will be more forward-compatible, as addition of new decoders -// doesn't require you to disable them explicitly): -// -// STBI_ONLY_JPEG -// STBI_ONLY_PNG -// STBI_ONLY_BMP -// STBI_ONLY_PSD -// STBI_ONLY_TGA -// STBI_ONLY_GIF -// STBI_ONLY_HDR -// STBI_ONLY_PIC -// STBI_ONLY_PNM (.ppm and .pgm) -// -// - If you use STBI_NO_PNG (or _ONLY_ without PNG), and you still -// want the zlib decoder to be available, #define STBI_SUPPORT_ZLIB -// -// - If you define STBI_MAX_DIMENSIONS, stb_image will reject images greater -// than that size (in either width or height) without further processing. -// This is to let programs in the wild set an upper bound to prevent -// denial-of-service attacks on untrusted data, as one could generate a -// valid image of gigantic dimensions and force stb_image to allocate a -// huge block of memory and spend disproportionate time decoding it. By -// default this is set to (1 << 24), which is 16777216, but that's still -// very big. - -#ifndef STBI_NO_STDIO -#include -#endif // STBI_NO_STDIO - -#define STBI_VERSION 1 - -enum -{ - STBI_default = 0, // only used for desired_channels - - STBI_grey = 1, - STBI_grey_alpha = 2, - STBI_rgb = 3, - STBI_rgb_alpha = 4 -}; - -#include -typedef unsigned char stbi_uc; -typedef unsigned short stbi_us; - -#ifdef __cplusplus -extern "C" { -#endif - -#ifndef STBIDEF -#ifdef STB_IMAGE_STATIC -#define STBIDEF static -#else -#define STBIDEF extern -#endif -#endif - -////////////////////////////////////////////////////////////////////////////// -// -// PRIMARY API - works on images of any type -// - -// -// load image by filename, open file, or memory buffer -// - -typedef struct -{ - int (*read) (void *user,char *data,int size); // fill 'data' with 'size' bytes. return number of bytes actually read - void (*skip) (void *user,int n); // skip the next 'n' bytes, or 'unget' the last -n bytes if negative - int (*eof) (void *user); // returns nonzero if we are at end of file/data -} stbi_io_callbacks; - -//////////////////////////////////// -// -// 8-bits-per-channel interface -// - -STBIDEF stbi_uc *stbi_load_from_memory (stbi_uc const *buffer, int len , int *x, int *y, int *channels_in_file, int desired_channels); -STBIDEF stbi_uc *stbi_load_from_callbacks(stbi_io_callbacks const *clbk , void *user, int *x, int *y, int *channels_in_file, int desired_channels); - -#ifndef STBI_NO_STDIO -STBIDEF stbi_uc *stbi_load (char const *filename, int *x, int *y, int *channels_in_file, int desired_channels); -STBIDEF stbi_uc *stbi_load_from_file (FILE *f, int *x, int *y, int *channels_in_file, int desired_channels); -// for stbi_load_from_file, file pointer is left pointing immediately after image -#endif - -#ifndef STBI_NO_GIF -STBIDEF stbi_uc *stbi_load_gif_from_memory(stbi_uc const *buffer, int len, int **delays, int *x, int *y, int *z, int *comp, int req_comp); -#endif - -#ifdef STBI_WINDOWS_UTF8 -STBIDEF int stbi_convert_wchar_to_utf8(char *buffer, size_t bufferlen, const wchar_t* input); -#endif - -//////////////////////////////////// -// -// 16-bits-per-channel interface -// - -STBIDEF stbi_us *stbi_load_16_from_memory (stbi_uc const *buffer, int len, int *x, int *y, int *channels_in_file, int desired_channels); -STBIDEF stbi_us *stbi_load_16_from_callbacks(stbi_io_callbacks const *clbk, void *user, int *x, int *y, int *channels_in_file, int desired_channels); - -#ifndef STBI_NO_STDIO -STBIDEF stbi_us *stbi_load_16 (char const *filename, int *x, int *y, int *channels_in_file, int desired_channels); -STBIDEF stbi_us *stbi_load_from_file_16(FILE *f, int *x, int *y, int *channels_in_file, int desired_channels); -#endif - -//////////////////////////////////// -// -// float-per-channel interface -// -#ifndef STBI_NO_LINEAR - STBIDEF float *stbi_loadf_from_memory (stbi_uc const *buffer, int len, int *x, int *y, int *channels_in_file, int desired_channels); - STBIDEF float *stbi_loadf_from_callbacks (stbi_io_callbacks const *clbk, void *user, int *x, int *y, int *channels_in_file, int desired_channels); - - #ifndef STBI_NO_STDIO - STBIDEF float *stbi_loadf (char const *filename, int *x, int *y, int *channels_in_file, int desired_channels); - STBIDEF float *stbi_loadf_from_file (FILE *f, int *x, int *y, int *channels_in_file, int desired_channels); - #endif -#endif - -#ifndef STBI_NO_HDR - STBIDEF void stbi_hdr_to_ldr_gamma(float gamma); - STBIDEF void stbi_hdr_to_ldr_scale(float scale); -#endif // STBI_NO_HDR - -#ifndef STBI_NO_LINEAR - STBIDEF void stbi_ldr_to_hdr_gamma(float gamma); - STBIDEF void stbi_ldr_to_hdr_scale(float scale); -#endif // STBI_NO_LINEAR - -// stbi_is_hdr is always defined, but always returns false if STBI_NO_HDR -STBIDEF int stbi_is_hdr_from_callbacks(stbi_io_callbacks const *clbk, void *user); -STBIDEF int stbi_is_hdr_from_memory(stbi_uc const *buffer, int len); -#ifndef STBI_NO_STDIO -STBIDEF int stbi_is_hdr (char const *filename); -STBIDEF int stbi_is_hdr_from_file(FILE *f); -#endif // STBI_NO_STDIO - - -// get a VERY brief reason for failure -// on most compilers (and ALL modern mainstream compilers) this is threadsafe -STBIDEF const char *stbi_failure_reason (void); - -// free the loaded image -- this is just free() -STBIDEF void stbi_image_free (void *retval_from_stbi_load); - -// get image dimensions & components without fully decoding -STBIDEF int stbi_info_from_memory(stbi_uc const *buffer, int len, int *x, int *y, int *comp); -STBIDEF int stbi_info_from_callbacks(stbi_io_callbacks const *clbk, void *user, int *x, int *y, int *comp); -STBIDEF int stbi_is_16_bit_from_memory(stbi_uc const *buffer, int len); -STBIDEF int stbi_is_16_bit_from_callbacks(stbi_io_callbacks const *clbk, void *user); - -#ifndef STBI_NO_STDIO -STBIDEF int stbi_info (char const *filename, int *x, int *y, int *comp); -STBIDEF int stbi_info_from_file (FILE *f, int *x, int *y, int *comp); -STBIDEF int stbi_is_16_bit (char const *filename); -STBIDEF int stbi_is_16_bit_from_file(FILE *f); -#endif - - - -// for image formats that explicitly notate that they have premultiplied alpha, -// we just return the colors as stored in the file. set this flag to force -// unpremultiplication. results are undefined if the unpremultiply overflow. -STBIDEF void stbi_set_unpremultiply_on_load(int flag_true_if_should_unpremultiply); - -// indicate whether we should process iphone images back to canonical format, -// or just pass them through "as-is" -STBIDEF void stbi_convert_iphone_png_to_rgb(int flag_true_if_should_convert); - -// flip the image vertically, so the first pixel in the output array is the bottom left -STBIDEF void stbi_set_flip_vertically_on_load(int flag_true_if_should_flip); - -// as above, but only applies to images loaded on the thread that calls the function -// this function is only available if your compiler supports thread-local variables; -// calling it will fail to link if your compiler doesn't -STBIDEF void stbi_set_unpremultiply_on_load_thread(int flag_true_if_should_unpremultiply); -STBIDEF void stbi_convert_iphone_png_to_rgb_thread(int flag_true_if_should_convert); -STBIDEF void stbi_set_flip_vertically_on_load_thread(int flag_true_if_should_flip); - -// ZLIB client - used by PNG, available for other purposes - -STBIDEF char *stbi_zlib_decode_malloc_guesssize(const char *buffer, int len, int initial_size, int *outlen); -STBIDEF char *stbi_zlib_decode_malloc_guesssize_headerflag(const char *buffer, int len, int initial_size, int *outlen, int parse_header); -STBIDEF char *stbi_zlib_decode_malloc(const char *buffer, int len, int *outlen); -STBIDEF int stbi_zlib_decode_buffer(char *obuffer, int olen, const char *ibuffer, int ilen); - -STBIDEF char *stbi_zlib_decode_noheader_malloc(const char *buffer, int len, int *outlen); -STBIDEF int stbi_zlib_decode_noheader_buffer(char *obuffer, int olen, const char *ibuffer, int ilen); - - -#ifdef __cplusplus -} -#endif - -// -// -//// end header file ///////////////////////////////////////////////////// -#endif // STBI_INCLUDE_STB_IMAGE_H - -#ifdef STB_IMAGE_IMPLEMENTATION - -#if defined(STBI_ONLY_JPEG) || defined(STBI_ONLY_PNG) || defined(STBI_ONLY_BMP) \ - || defined(STBI_ONLY_TGA) || defined(STBI_ONLY_GIF) || defined(STBI_ONLY_PSD) \ - || defined(STBI_ONLY_HDR) || defined(STBI_ONLY_PIC) || defined(STBI_ONLY_PNM) \ - || defined(STBI_ONLY_ZLIB) - #ifndef STBI_ONLY_JPEG - #define STBI_NO_JPEG - #endif - #ifndef STBI_ONLY_PNG - #define STBI_NO_PNG - #endif - #ifndef STBI_ONLY_BMP - #define STBI_NO_BMP - #endif - #ifndef STBI_ONLY_PSD - #define STBI_NO_PSD - #endif - #ifndef STBI_ONLY_TGA - #define STBI_NO_TGA - #endif - #ifndef STBI_ONLY_GIF - #define STBI_NO_GIF - #endif - #ifndef STBI_ONLY_HDR - #define STBI_NO_HDR - #endif - #ifndef STBI_ONLY_PIC - #define STBI_NO_PIC - #endif - #ifndef STBI_ONLY_PNM - #define STBI_NO_PNM - #endif -#endif - -#if defined(STBI_NO_PNG) && !defined(STBI_SUPPORT_ZLIB) && !defined(STBI_NO_ZLIB) -#define STBI_NO_ZLIB -#endif - - -#include -#include // ptrdiff_t on osx -#include -#include -#include - -#if !defined(STBI_NO_LINEAR) || !defined(STBI_NO_HDR) -#include // ldexp, pow -#endif - -#ifndef STBI_NO_STDIO -#include -#endif - -#ifndef STBI_ASSERT -#include -#define STBI_ASSERT(x) assert(x) -#endif - -#ifdef __cplusplus -#define STBI_EXTERN extern "C" -#else -#define STBI_EXTERN extern -#endif - - -#ifndef _MSC_VER - #ifdef __cplusplus - #define stbi_inline inline - #else - #define stbi_inline - #endif -#else - #define stbi_inline __forceinline -#endif - -#ifndef STBI_NO_THREAD_LOCALS - #if defined(__cplusplus) && __cplusplus >= 201103L - #define STBI_THREAD_LOCAL thread_local - #elif defined(__GNUC__) && __GNUC__ < 5 - #define STBI_THREAD_LOCAL __thread - #elif defined(_MSC_VER) - #define STBI_THREAD_LOCAL __declspec(thread) - #elif defined (__STDC_VERSION__) && __STDC_VERSION__ >= 201112L && !defined(__STDC_NO_THREADS__) - #define STBI_THREAD_LOCAL _Thread_local - #endif - - #ifndef STBI_THREAD_LOCAL - #if defined(__GNUC__) - #define STBI_THREAD_LOCAL __thread - #endif - #endif -#endif - -#if defined(_MSC_VER) || defined(__SYMBIAN32__) -typedef unsigned short stbi__uint16; -typedef signed short stbi__int16; -typedef unsigned int stbi__uint32; -typedef signed int stbi__int32; -#else -#include -typedef uint16_t stbi__uint16; -typedef int16_t stbi__int16; -typedef uint32_t stbi__uint32; -typedef int32_t stbi__int32; -#endif - -// should produce compiler error if size is wrong -typedef unsigned char validate_uint32[sizeof(stbi__uint32)==4 ? 1 : -1]; - -#ifdef _MSC_VER -#define STBI_NOTUSED(v) (void)(v) -#else -#define STBI_NOTUSED(v) (void)sizeof(v) -#endif - -#ifdef _MSC_VER -#define STBI_HAS_LROTL -#endif - -#ifdef STBI_HAS_LROTL - #define stbi_lrot(x,y) _lrotl(x,y) -#else - #define stbi_lrot(x,y) (((x) << (y)) | ((x) >> (-(y) & 31))) -#endif - -#if defined(STBI_MALLOC) && defined(STBI_FREE) && (defined(STBI_REALLOC) || defined(STBI_REALLOC_SIZED)) -// ok -#elif !defined(STBI_MALLOC) && !defined(STBI_FREE) && !defined(STBI_REALLOC) && !defined(STBI_REALLOC_SIZED) -// ok -#else -#error "Must define all or none of STBI_MALLOC, STBI_FREE, and STBI_REALLOC (or STBI_REALLOC_SIZED)." -#endif - -#ifndef STBI_MALLOC -#define STBI_MALLOC(sz) malloc(sz) -#define STBI_REALLOC(p,newsz) realloc(p,newsz) -#define STBI_FREE(p) free(p) -#endif - -#ifndef STBI_REALLOC_SIZED -#define STBI_REALLOC_SIZED(p,oldsz,newsz) STBI_REALLOC(p,newsz) -#endif - -// x86/x64 detection -#if defined(__x86_64__) || defined(_M_X64) -#define STBI__X64_TARGET -#elif defined(__i386) || defined(_M_IX86) -#define STBI__X86_TARGET -#endif - -#if defined(__GNUC__) && defined(STBI__X86_TARGET) && !defined(__SSE2__) && !defined(STBI_NO_SIMD) -// gcc doesn't support sse2 intrinsics unless you compile with -msse2, -// which in turn means it gets to use SSE2 everywhere. This is unfortunate, -// but previous attempts to provide the SSE2 functions with runtime -// detection caused numerous issues. The way architecture extensions are -// exposed in GCC/Clang is, sadly, not really suited for one-file libs. -// New behavior: if compiled with -msse2, we use SSE2 without any -// detection; if not, we don't use it at all. -#define STBI_NO_SIMD -#endif - -#if defined(__MINGW32__) && defined(STBI__X86_TARGET) && !defined(STBI_MINGW_ENABLE_SSE2) && !defined(STBI_NO_SIMD) -// Note that __MINGW32__ doesn't actually mean 32-bit, so we have to avoid STBI__X64_TARGET -// -// 32-bit MinGW wants ESP to be 16-byte aligned, but this is not in the -// Windows ABI and VC++ as well as Windows DLLs don't maintain that invariant. -// As a result, enabling SSE2 on 32-bit MinGW is dangerous when not -// simultaneously enabling "-mstackrealign". -// -// See https://github.com/nothings/stb/issues/81 for more information. -// -// So default to no SSE2 on 32-bit MinGW. If you've read this far and added -// -mstackrealign to your build settings, feel free to #define STBI_MINGW_ENABLE_SSE2. -#define STBI_NO_SIMD -#endif - -#if !defined(STBI_NO_SIMD) && (defined(STBI__X86_TARGET) || defined(STBI__X64_TARGET)) -#define STBI_SSE2 -#include - -#ifdef _MSC_VER - -#if _MSC_VER >= 1400 // not VC6 -#include // __cpuid -static int stbi__cpuid3(void) -{ - int info[4]; - __cpuid(info,1); - return info[3]; -} -#else -static int stbi__cpuid3(void) -{ - int res; - __asm { - mov eax,1 - cpuid - mov res,edx - } - return res; -} -#endif - -#define STBI_SIMD_ALIGN(type, name) __declspec(align(16)) type name - -#if !defined(STBI_NO_JPEG) && defined(STBI_SSE2) -static int stbi__sse2_available(void) -{ - int info3 = stbi__cpuid3(); - return ((info3 >> 26) & 1) != 0; -} -#endif - -#else // assume GCC-style if not VC++ -#define STBI_SIMD_ALIGN(type, name) type name __attribute__((aligned(16))) - -#if !defined(STBI_NO_JPEG) && defined(STBI_SSE2) -static int stbi__sse2_available(void) -{ - // If we're even attempting to compile this on GCC/Clang, that means - // -msse2 is on, which means the compiler is allowed to use SSE2 - // instructions at will, and so are we. - return 1; -} -#endif - -#endif -#endif - -// ARM NEON -#if defined(STBI_NO_SIMD) && defined(STBI_NEON) -#undef STBI_NEON -#endif - -#ifdef STBI_NEON -#include -#ifdef _MSC_VER -#define STBI_SIMD_ALIGN(type, name) __declspec(align(16)) type name -#else -#define STBI_SIMD_ALIGN(type, name) type name __attribute__((aligned(16))) -#endif -#endif - -#ifndef STBI_SIMD_ALIGN -#define STBI_SIMD_ALIGN(type, name) type name -#endif - -#ifndef STBI_MAX_DIMENSIONS -#define STBI_MAX_DIMENSIONS (1 << 24) -#endif - -/////////////////////////////////////////////// -// -// stbi__context struct and start_xxx functions - -// stbi__context structure is our basic context used by all images, so it -// contains all the IO context, plus some basic image information -typedef struct -{ - stbi__uint32 img_x, img_y; - int img_n, img_out_n; - - stbi_io_callbacks io; - void *io_user_data; - - int read_from_callbacks; - int buflen; - stbi_uc buffer_start[128]; - int callback_already_read; - - stbi_uc *img_buffer, *img_buffer_end; - stbi_uc *img_buffer_original, *img_buffer_original_end; -} stbi__context; - - -static void stbi__refill_buffer(stbi__context *s); - -// initialize a memory-decode context -static void stbi__start_mem(stbi__context *s, stbi_uc const *buffer, int len) -{ - s->io.read = NULL; - s->read_from_callbacks = 0; - s->callback_already_read = 0; - s->img_buffer = s->img_buffer_original = (stbi_uc *) buffer; - s->img_buffer_end = s->img_buffer_original_end = (stbi_uc *) buffer+len; -} - -// initialize a callback-based context -static void stbi__start_callbacks(stbi__context *s, stbi_io_callbacks *c, void *user) -{ - s->io = *c; - s->io_user_data = user; - s->buflen = sizeof(s->buffer_start); - s->read_from_callbacks = 1; - s->callback_already_read = 0; - s->img_buffer = s->img_buffer_original = s->buffer_start; - stbi__refill_buffer(s); - s->img_buffer_original_end = s->img_buffer_end; -} - -#ifndef STBI_NO_STDIO - -static int stbi__stdio_read(void *user, char *data, int size) -{ - return (int) fread(data,1,size,(FILE*) user); -} - -static void stbi__stdio_skip(void *user, int n) -{ - int ch; - fseek((FILE*) user, n, SEEK_CUR); - ch = fgetc((FILE*) user); /* have to read a byte to reset feof()'s flag */ - if (ch != EOF) { - ungetc(ch, (FILE *) user); /* push byte back onto stream if valid. */ - } -} - -static int stbi__stdio_eof(void *user) -{ - return feof((FILE*) user) || ferror((FILE *) user); -} - -static stbi_io_callbacks stbi__stdio_callbacks = -{ - stbi__stdio_read, - stbi__stdio_skip, - stbi__stdio_eof, -}; - -static void stbi__start_file(stbi__context *s, FILE *f) -{ - stbi__start_callbacks(s, &stbi__stdio_callbacks, (void *) f); -} - -//static void stop_file(stbi__context *s) { } - -#endif // !STBI_NO_STDIO - -static void stbi__rewind(stbi__context *s) -{ - // conceptually rewind SHOULD rewind to the beginning of the stream, - // but we just rewind to the beginning of the initial buffer, because - // we only use it after doing 'test', which only ever looks at at most 92 bytes - s->img_buffer = s->img_buffer_original; - s->img_buffer_end = s->img_buffer_original_end; -} - -enum -{ - STBI_ORDER_RGB, - STBI_ORDER_BGR -}; - -typedef struct -{ - int bits_per_channel; - int num_channels; - int channel_order; -} stbi__result_info; - -#ifndef STBI_NO_JPEG -static int stbi__jpeg_test(stbi__context *s); -static void *stbi__jpeg_load(stbi__context *s, int *x, int *y, int *comp, int req_comp, stbi__result_info *ri); -static int stbi__jpeg_info(stbi__context *s, int *x, int *y, int *comp); -#endif - -#ifndef STBI_NO_PNG -static int stbi__png_test(stbi__context *s); -static void *stbi__png_load(stbi__context *s, int *x, int *y, int *comp, int req_comp, stbi__result_info *ri); -static int stbi__png_info(stbi__context *s, int *x, int *y, int *comp); -static int stbi__png_is16(stbi__context *s); -#endif - -#ifndef STBI_NO_BMP -static int stbi__bmp_test(stbi__context *s); -static void *stbi__bmp_load(stbi__context *s, int *x, int *y, int *comp, int req_comp, stbi__result_info *ri); -static int stbi__bmp_info(stbi__context *s, int *x, int *y, int *comp); -#endif - -#ifndef STBI_NO_TGA -static int stbi__tga_test(stbi__context *s); -static void *stbi__tga_load(stbi__context *s, int *x, int *y, int *comp, int req_comp, stbi__result_info *ri); -static int stbi__tga_info(stbi__context *s, int *x, int *y, int *comp); -#endif - -#ifndef STBI_NO_PSD -static int stbi__psd_test(stbi__context *s); -static void *stbi__psd_load(stbi__context *s, int *x, int *y, int *comp, int req_comp, stbi__result_info *ri, int bpc); -static int stbi__psd_info(stbi__context *s, int *x, int *y, int *comp); -static int stbi__psd_is16(stbi__context *s); -#endif - -#ifndef STBI_NO_HDR -static int stbi__hdr_test(stbi__context *s); -static float *stbi__hdr_load(stbi__context *s, int *x, int *y, int *comp, int req_comp, stbi__result_info *ri); -static int stbi__hdr_info(stbi__context *s, int *x, int *y, int *comp); -#endif - -#ifndef STBI_NO_PIC -static int stbi__pic_test(stbi__context *s); -static void *stbi__pic_load(stbi__context *s, int *x, int *y, int *comp, int req_comp, stbi__result_info *ri); -static int stbi__pic_info(stbi__context *s, int *x, int *y, int *comp); -#endif - -#ifndef STBI_NO_GIF -static int stbi__gif_test(stbi__context *s); -static void *stbi__gif_load(stbi__context *s, int *x, int *y, int *comp, int req_comp, stbi__result_info *ri); -static void *stbi__load_gif_main(stbi__context *s, int **delays, int *x, int *y, int *z, int *comp, int req_comp); -static int stbi__gif_info(stbi__context *s, int *x, int *y, int *comp); -#endif - -#ifndef STBI_NO_PNM -static int stbi__pnm_test(stbi__context *s); -static void *stbi__pnm_load(stbi__context *s, int *x, int *y, int *comp, int req_comp, stbi__result_info *ri); -static int stbi__pnm_info(stbi__context *s, int *x, int *y, int *comp); -static int stbi__pnm_is16(stbi__context *s); -#endif - -static -#ifdef STBI_THREAD_LOCAL -STBI_THREAD_LOCAL -#endif -const char *stbi__g_failure_reason; - -STBIDEF const char *stbi_failure_reason(void) -{ - return stbi__g_failure_reason; -} - -#ifndef STBI_NO_FAILURE_STRINGS -static int stbi__err(const char *str) -{ - stbi__g_failure_reason = str; - return 0; -} -#endif - -static void *stbi__malloc(size_t size) -{ - return STBI_MALLOC(size); -} - -// stb_image uses ints pervasively, including for offset calculations. -// therefore the largest decoded image size we can support with the -// current code, even on 64-bit targets, is INT_MAX. this is not a -// significant limitation for the intended use case. -// -// we do, however, need to make sure our size calculations don't -// overflow. hence a few helper functions for size calculations that -// multiply integers together, making sure that they're non-negative -// and no overflow occurs. - -// return 1 if the sum is valid, 0 on overflow. -// negative terms are considered invalid. -static int stbi__addsizes_valid(int a, int b) -{ - if (b < 0) return 0; - // now 0 <= b <= INT_MAX, hence also - // 0 <= INT_MAX - b <= INTMAX. - // And "a + b <= INT_MAX" (which might overflow) is the - // same as a <= INT_MAX - b (no overflow) - return a <= INT_MAX - b; -} - -// returns 1 if the product is valid, 0 on overflow. -// negative factors are considered invalid. -static int stbi__mul2sizes_valid(int a, int b) -{ - if (a < 0 || b < 0) return 0; - if (b == 0) return 1; // mul-by-0 is always safe - // portable way to check for no overflows in a*b - return a <= INT_MAX/b; -} - -#if !defined(STBI_NO_JPEG) || !defined(STBI_NO_PNG) || !defined(STBI_NO_TGA) || !defined(STBI_NO_HDR) -// returns 1 if "a*b + add" has no negative terms/factors and doesn't overflow -static int stbi__mad2sizes_valid(int a, int b, int add) -{ - return stbi__mul2sizes_valid(a, b) && stbi__addsizes_valid(a*b, add); -} -#endif - -// returns 1 if "a*b*c + add" has no negative terms/factors and doesn't overflow -static int stbi__mad3sizes_valid(int a, int b, int c, int add) -{ - return stbi__mul2sizes_valid(a, b) && stbi__mul2sizes_valid(a*b, c) && - stbi__addsizes_valid(a*b*c, add); -} - -// returns 1 if "a*b*c*d + add" has no negative terms/factors and doesn't overflow -#if !defined(STBI_NO_LINEAR) || !defined(STBI_NO_HDR) || !defined(STBI_NO_PNM) -static int stbi__mad4sizes_valid(int a, int b, int c, int d, int add) -{ - return stbi__mul2sizes_valid(a, b) && stbi__mul2sizes_valid(a*b, c) && - stbi__mul2sizes_valid(a*b*c, d) && stbi__addsizes_valid(a*b*c*d, add); -} -#endif - -#if !defined(STBI_NO_JPEG) || !defined(STBI_NO_PNG) || !defined(STBI_NO_TGA) || !defined(STBI_NO_HDR) -// mallocs with size overflow checking -static void *stbi__malloc_mad2(int a, int b, int add) -{ - if (!stbi__mad2sizes_valid(a, b, add)) return NULL; - return stbi__malloc(a*b + add); -} -#endif - -static void *stbi__malloc_mad3(int a, int b, int c, int add) -{ - if (!stbi__mad3sizes_valid(a, b, c, add)) return NULL; - return stbi__malloc(a*b*c + add); -} - -#if !defined(STBI_NO_LINEAR) || !defined(STBI_NO_HDR) || !defined(STBI_NO_PNM) -static void *stbi__malloc_mad4(int a, int b, int c, int d, int add) -{ - if (!stbi__mad4sizes_valid(a, b, c, d, add)) return NULL; - return stbi__malloc(a*b*c*d + add); -} -#endif - -// returns 1 if the sum of two signed ints is valid (between -2^31 and 2^31-1 inclusive), 0 on overflow. -static int stbi__addints_valid(int a, int b) -{ - if ((a >= 0) != (b >= 0)) return 1; // a and b have different signs, so no overflow - if (a < 0 && b < 0) return a >= INT_MIN - b; // same as a + b >= INT_MIN; INT_MIN - b cannot overflow since b < 0. - return a <= INT_MAX - b; -} - -// returns 1 if the product of two ints fits in a signed short, 0 on overflow. -static int stbi__mul2shorts_valid(int a, int b) -{ - if (b == 0 || b == -1) return 1; // multiplication by 0 is always 0; check for -1 so SHRT_MIN/b doesn't overflow - if ((a >= 0) == (b >= 0)) return a <= SHRT_MAX/b; // product is positive, so similar to mul2sizes_valid - if (b < 0) return a <= SHRT_MIN / b; // same as a * b >= SHRT_MIN - return a >= SHRT_MIN / b; -} - -// stbi__err - error -// stbi__errpf - error returning pointer to float -// stbi__errpuc - error returning pointer to unsigned char - -#ifdef STBI_NO_FAILURE_STRINGS - #define stbi__err(x,y) 0 -#elif defined(STBI_FAILURE_USERMSG) - #define stbi__err(x,y) stbi__err(y) -#else - #define stbi__err(x,y) stbi__err(x) -#endif - -#define stbi__errpf(x,y) ((float *)(size_t) (stbi__err(x,y)?NULL:NULL)) -#define stbi__errpuc(x,y) ((unsigned char *)(size_t) (stbi__err(x,y)?NULL:NULL)) - -STBIDEF void stbi_image_free(void *retval_from_stbi_load) -{ - STBI_FREE(retval_from_stbi_load); -} - -#ifndef STBI_NO_LINEAR -static float *stbi__ldr_to_hdr(stbi_uc *data, int x, int y, int comp); -#endif - -#ifndef STBI_NO_HDR -static stbi_uc *stbi__hdr_to_ldr(float *data, int x, int y, int comp); -#endif - -static int stbi__vertically_flip_on_load_global = 0; - -STBIDEF void stbi_set_flip_vertically_on_load(int flag_true_if_should_flip) -{ - stbi__vertically_flip_on_load_global = flag_true_if_should_flip; -} - -#ifndef STBI_THREAD_LOCAL -#define stbi__vertically_flip_on_load stbi__vertically_flip_on_load_global -#else -static STBI_THREAD_LOCAL int stbi__vertically_flip_on_load_local, stbi__vertically_flip_on_load_set; - -STBIDEF void stbi_set_flip_vertically_on_load_thread(int flag_true_if_should_flip) -{ - stbi__vertically_flip_on_load_local = flag_true_if_should_flip; - stbi__vertically_flip_on_load_set = 1; -} - -#define stbi__vertically_flip_on_load (stbi__vertically_flip_on_load_set \ - ? stbi__vertically_flip_on_load_local \ - : stbi__vertically_flip_on_load_global) -#endif // STBI_THREAD_LOCAL - -static void *stbi__load_main(stbi__context *s, int *x, int *y, int *comp, int req_comp, stbi__result_info *ri, int bpc) -{ - memset(ri, 0, sizeof(*ri)); // make sure it's initialized if we add new fields - ri->bits_per_channel = 8; // default is 8 so most paths don't have to be changed - ri->channel_order = STBI_ORDER_RGB; // all current input & output are this, but this is here so we can add BGR order - ri->num_channels = 0; - - // test the formats with a very explicit header first (at least a FOURCC - // or distinctive magic number first) - #ifndef STBI_NO_PNG - if (stbi__png_test(s)) return stbi__png_load(s,x,y,comp,req_comp, ri); - #endif - #ifndef STBI_NO_BMP - if (stbi__bmp_test(s)) return stbi__bmp_load(s,x,y,comp,req_comp, ri); - #endif - #ifndef STBI_NO_GIF - if (stbi__gif_test(s)) return stbi__gif_load(s,x,y,comp,req_comp, ri); - #endif - #ifndef STBI_NO_PSD - if (stbi__psd_test(s)) return stbi__psd_load(s,x,y,comp,req_comp, ri, bpc); - #else - STBI_NOTUSED(bpc); - #endif - #ifndef STBI_NO_PIC - if (stbi__pic_test(s)) return stbi__pic_load(s,x,y,comp,req_comp, ri); - #endif - - // then the formats that can end up attempting to load with just 1 or 2 - // bytes matching expectations; these are prone to false positives, so - // try them later - #ifndef STBI_NO_JPEG - if (stbi__jpeg_test(s)) return stbi__jpeg_load(s,x,y,comp,req_comp, ri); - #endif - #ifndef STBI_NO_PNM - if (stbi__pnm_test(s)) return stbi__pnm_load(s,x,y,comp,req_comp, ri); - #endif - - #ifndef STBI_NO_HDR - if (stbi__hdr_test(s)) { - float *hdr = stbi__hdr_load(s, x,y,comp,req_comp, ri); - return stbi__hdr_to_ldr(hdr, *x, *y, req_comp ? req_comp : *comp); - } - #endif - - #ifndef STBI_NO_TGA - // test tga last because it's a crappy test! - if (stbi__tga_test(s)) - return stbi__tga_load(s,x,y,comp,req_comp, ri); - #endif - - return stbi__errpuc("unknown image type", "Image not of any known type, or corrupt"); -} - -static stbi_uc *stbi__convert_16_to_8(stbi__uint16 *orig, int w, int h, int channels) -{ - int i; - int img_len = w * h * channels; - stbi_uc *reduced; - - reduced = (stbi_uc *) stbi__malloc(img_len); - if (reduced == NULL) return stbi__errpuc("outofmem", "Out of memory"); - - for (i = 0; i < img_len; ++i) - reduced[i] = (stbi_uc)((orig[i] >> 8) & 0xFF); // top half of each byte is sufficient approx of 16->8 bit scaling - - STBI_FREE(orig); - return reduced; -} - -static stbi__uint16 *stbi__convert_8_to_16(stbi_uc *orig, int w, int h, int channels) -{ - int i; - int img_len = w * h * channels; - stbi__uint16 *enlarged; - - enlarged = (stbi__uint16 *) stbi__malloc(img_len*2); - if (enlarged == NULL) return (stbi__uint16 *) stbi__errpuc("outofmem", "Out of memory"); - - for (i = 0; i < img_len; ++i) - enlarged[i] = (stbi__uint16)((orig[i] << 8) + orig[i]); // replicate to high and low byte, maps 0->0, 255->0xffff - - STBI_FREE(orig); - return enlarged; -} - -static void stbi__vertical_flip(void *image, int w, int h, int bytes_per_pixel) -{ - int row; - size_t bytes_per_row = (size_t)w * bytes_per_pixel; - stbi_uc temp[2048]; - stbi_uc *bytes = (stbi_uc *)image; - - for (row = 0; row < (h>>1); row++) { - stbi_uc *row0 = bytes + row*bytes_per_row; - stbi_uc *row1 = bytes + (h - row - 1)*bytes_per_row; - // swap row0 with row1 - size_t bytes_left = bytes_per_row; - while (bytes_left) { - size_t bytes_copy = (bytes_left < sizeof(temp)) ? bytes_left : sizeof(temp); - memcpy(temp, row0, bytes_copy); - memcpy(row0, row1, bytes_copy); - memcpy(row1, temp, bytes_copy); - row0 += bytes_copy; - row1 += bytes_copy; - bytes_left -= bytes_copy; - } - } -} - -#ifndef STBI_NO_GIF -static void stbi__vertical_flip_slices(void *image, int w, int h, int z, int bytes_per_pixel) -{ - int slice; - int slice_size = w * h * bytes_per_pixel; - - stbi_uc *bytes = (stbi_uc *)image; - for (slice = 0; slice < z; ++slice) { - stbi__vertical_flip(bytes, w, h, bytes_per_pixel); - bytes += slice_size; - } -} -#endif - -static unsigned char *stbi__load_and_postprocess_8bit(stbi__context *s, int *x, int *y, int *comp, int req_comp) -{ - stbi__result_info ri; - void *result = stbi__load_main(s, x, y, comp, req_comp, &ri, 8); - - if (result == NULL) - return NULL; - - // it is the responsibility of the loaders to make sure we get either 8 or 16 bit. - STBI_ASSERT(ri.bits_per_channel == 8 || ri.bits_per_channel == 16); - - if (ri.bits_per_channel != 8) { - result = stbi__convert_16_to_8((stbi__uint16 *) result, *x, *y, req_comp == 0 ? *comp : req_comp); - ri.bits_per_channel = 8; - } - - // @TODO: move stbi__convert_format to here - - if (stbi__vertically_flip_on_load) { - int channels = req_comp ? req_comp : *comp; - stbi__vertical_flip(result, *x, *y, channels * sizeof(stbi_uc)); - } - - return (unsigned char *) result; -} - -static stbi__uint16 *stbi__load_and_postprocess_16bit(stbi__context *s, int *x, int *y, int *comp, int req_comp) -{ - stbi__result_info ri; - void *result = stbi__load_main(s, x, y, comp, req_comp, &ri, 16); - - if (result == NULL) - return NULL; - - // it is the responsibility of the loaders to make sure we get either 8 or 16 bit. - STBI_ASSERT(ri.bits_per_channel == 8 || ri.bits_per_channel == 16); - - if (ri.bits_per_channel != 16) { - result = stbi__convert_8_to_16((stbi_uc *) result, *x, *y, req_comp == 0 ? *comp : req_comp); - ri.bits_per_channel = 16; - } - - // @TODO: move stbi__convert_format16 to here - // @TODO: special case RGB-to-Y (and RGBA-to-YA) for 8-bit-to-16-bit case to keep more precision - - if (stbi__vertically_flip_on_load) { - int channels = req_comp ? req_comp : *comp; - stbi__vertical_flip(result, *x, *y, channels * sizeof(stbi__uint16)); - } - - return (stbi__uint16 *) result; -} - -#if !defined(STBI_NO_HDR) && !defined(STBI_NO_LINEAR) -static void stbi__float_postprocess(float *result, int *x, int *y, int *comp, int req_comp) -{ - if (stbi__vertically_flip_on_load && result != NULL) { - int channels = req_comp ? req_comp : *comp; - stbi__vertical_flip(result, *x, *y, channels * sizeof(float)); - } -} -#endif - -#ifndef STBI_NO_STDIO - -#if defined(_WIN32) && defined(STBI_WINDOWS_UTF8) -STBI_EXTERN __declspec(dllimport) int __stdcall MultiByteToWideChar(unsigned int cp, unsigned long flags, const char *str, int cbmb, wchar_t *widestr, int cchwide); -STBI_EXTERN __declspec(dllimport) int __stdcall WideCharToMultiByte(unsigned int cp, unsigned long flags, const wchar_t *widestr, int cchwide, char *str, int cbmb, const char *defchar, int *used_default); -#endif - -#if defined(_WIN32) && defined(STBI_WINDOWS_UTF8) -STBIDEF int stbi_convert_wchar_to_utf8(char *buffer, size_t bufferlen, const wchar_t* input) -{ - return WideCharToMultiByte(65001 /* UTF8 */, 0, input, -1, buffer, (int) bufferlen, NULL, NULL); -} -#endif - -static FILE *stbi__fopen(char const *filename, char const *mode) -{ - FILE *f; -#if defined(_WIN32) && defined(STBI_WINDOWS_UTF8) - wchar_t wMode[64]; - wchar_t wFilename[1024]; - if (0 == MultiByteToWideChar(65001 /* UTF8 */, 0, filename, -1, wFilename, sizeof(wFilename)/sizeof(*wFilename))) - return 0; - - if (0 == MultiByteToWideChar(65001 /* UTF8 */, 0, mode, -1, wMode, sizeof(wMode)/sizeof(*wMode))) - return 0; - -#if defined(_MSC_VER) && _MSC_VER >= 1400 - if (0 != _wfopen_s(&f, wFilename, wMode)) - f = 0; -#else - f = _wfopen(wFilename, wMode); -#endif - -#elif defined(_MSC_VER) && _MSC_VER >= 1400 - if (0 != fopen_s(&f, filename, mode)) - f=0; -#else - f = fopen(filename, mode); -#endif - return f; -} - - -STBIDEF stbi_uc *stbi_load(char const *filename, int *x, int *y, int *comp, int req_comp) -{ - FILE *f = stbi__fopen(filename, "rb"); - unsigned char *result; - if (!f) return stbi__errpuc("can't fopen", "Unable to open file"); - result = stbi_load_from_file(f,x,y,comp,req_comp); - fclose(f); - return result; -} - -STBIDEF stbi_uc *stbi_load_from_file(FILE *f, int *x, int *y, int *comp, int req_comp) -{ - unsigned char *result; - stbi__context s; - stbi__start_file(&s,f); - result = stbi__load_and_postprocess_8bit(&s,x,y,comp,req_comp); - if (result) { - // need to 'unget' all the characters in the IO buffer - fseek(f, - (int) (s.img_buffer_end - s.img_buffer), SEEK_CUR); - } - return result; -} - -STBIDEF stbi__uint16 *stbi_load_from_file_16(FILE *f, int *x, int *y, int *comp, int req_comp) -{ - stbi__uint16 *result; - stbi__context s; - stbi__start_file(&s,f); - result = stbi__load_and_postprocess_16bit(&s,x,y,comp,req_comp); - if (result) { - // need to 'unget' all the characters in the IO buffer - fseek(f, - (int) (s.img_buffer_end - s.img_buffer), SEEK_CUR); - } - return result; -} - -STBIDEF stbi_us *stbi_load_16(char const *filename, int *x, int *y, int *comp, int req_comp) -{ - FILE *f = stbi__fopen(filename, "rb"); - stbi__uint16 *result; - if (!f) return (stbi_us *) stbi__errpuc("can't fopen", "Unable to open file"); - result = stbi_load_from_file_16(f,x,y,comp,req_comp); - fclose(f); - return result; -} - - -#endif //!STBI_NO_STDIO - -STBIDEF stbi_us *stbi_load_16_from_memory(stbi_uc const *buffer, int len, int *x, int *y, int *channels_in_file, int desired_channels) -{ - stbi__context s; - stbi__start_mem(&s,buffer,len); - return stbi__load_and_postprocess_16bit(&s,x,y,channels_in_file,desired_channels); -} - -STBIDEF stbi_us *stbi_load_16_from_callbacks(stbi_io_callbacks const *clbk, void *user, int *x, int *y, int *channels_in_file, int desired_channels) -{ - stbi__context s; - stbi__start_callbacks(&s, (stbi_io_callbacks *)clbk, user); - return stbi__load_and_postprocess_16bit(&s,x,y,channels_in_file,desired_channels); -} - -STBIDEF stbi_uc *stbi_load_from_memory(stbi_uc const *buffer, int len, int *x, int *y, int *comp, int req_comp) -{ - stbi__context s; - stbi__start_mem(&s,buffer,len); - return stbi__load_and_postprocess_8bit(&s,x,y,comp,req_comp); -} - -STBIDEF stbi_uc *stbi_load_from_callbacks(stbi_io_callbacks const *clbk, void *user, int *x, int *y, int *comp, int req_comp) -{ - stbi__context s; - stbi__start_callbacks(&s, (stbi_io_callbacks *) clbk, user); - return stbi__load_and_postprocess_8bit(&s,x,y,comp,req_comp); -} - -#ifndef STBI_NO_GIF -STBIDEF stbi_uc *stbi_load_gif_from_memory(stbi_uc const *buffer, int len, int **delays, int *x, int *y, int *z, int *comp, int req_comp) -{ - unsigned char *result; - stbi__context s; - stbi__start_mem(&s,buffer,len); - - result = (unsigned char*) stbi__load_gif_main(&s, delays, x, y, z, comp, req_comp); - if (stbi__vertically_flip_on_load) { - stbi__vertical_flip_slices( result, *x, *y, *z, *comp ); - } - - return result; -} -#endif - -#ifndef STBI_NO_LINEAR -static float *stbi__loadf_main(stbi__context *s, int *x, int *y, int *comp, int req_comp) -{ - unsigned char *data; - #ifndef STBI_NO_HDR - if (stbi__hdr_test(s)) { - stbi__result_info ri; - float *hdr_data = stbi__hdr_load(s,x,y,comp,req_comp, &ri); - if (hdr_data) - stbi__float_postprocess(hdr_data,x,y,comp,req_comp); - return hdr_data; - } - #endif - data = stbi__load_and_postprocess_8bit(s, x, y, comp, req_comp); - if (data) - return stbi__ldr_to_hdr(data, *x, *y, req_comp ? req_comp : *comp); - return stbi__errpf("unknown image type", "Image not of any known type, or corrupt"); -} - -STBIDEF float *stbi_loadf_from_memory(stbi_uc const *buffer, int len, int *x, int *y, int *comp, int req_comp) -{ - stbi__context s; - stbi__start_mem(&s,buffer,len); - return stbi__loadf_main(&s,x,y,comp,req_comp); -} - -STBIDEF float *stbi_loadf_from_callbacks(stbi_io_callbacks const *clbk, void *user, int *x, int *y, int *comp, int req_comp) -{ - stbi__context s; - stbi__start_callbacks(&s, (stbi_io_callbacks *) clbk, user); - return stbi__loadf_main(&s,x,y,comp,req_comp); -} - -#ifndef STBI_NO_STDIO -STBIDEF float *stbi_loadf(char const *filename, int *x, int *y, int *comp, int req_comp) -{ - float *result; - FILE *f = stbi__fopen(filename, "rb"); - if (!f) return stbi__errpf("can't fopen", "Unable to open file"); - result = stbi_loadf_from_file(f,x,y,comp,req_comp); - fclose(f); - return result; -} - -STBIDEF float *stbi_loadf_from_file(FILE *f, int *x, int *y, int *comp, int req_comp) -{ - stbi__context s; - stbi__start_file(&s,f); - return stbi__loadf_main(&s,x,y,comp,req_comp); -} -#endif // !STBI_NO_STDIO - -#endif // !STBI_NO_LINEAR - -// these is-hdr-or-not is defined independent of whether STBI_NO_LINEAR is -// defined, for API simplicity; if STBI_NO_LINEAR is defined, it always -// reports false! - -STBIDEF int stbi_is_hdr_from_memory(stbi_uc const *buffer, int len) -{ - #ifndef STBI_NO_HDR - stbi__context s; - stbi__start_mem(&s,buffer,len); - return stbi__hdr_test(&s); - #else - STBI_NOTUSED(buffer); - STBI_NOTUSED(len); - return 0; - #endif -} - -#ifndef STBI_NO_STDIO -STBIDEF int stbi_is_hdr (char const *filename) -{ - FILE *f = stbi__fopen(filename, "rb"); - int result=0; - if (f) { - result = stbi_is_hdr_from_file(f); - fclose(f); - } - return result; -} - -STBIDEF int stbi_is_hdr_from_file(FILE *f) -{ - #ifndef STBI_NO_HDR - long pos = ftell(f); - int res; - stbi__context s; - stbi__start_file(&s,f); - res = stbi__hdr_test(&s); - fseek(f, pos, SEEK_SET); - return res; - #else - STBI_NOTUSED(f); - return 0; - #endif -} -#endif // !STBI_NO_STDIO - -STBIDEF int stbi_is_hdr_from_callbacks(stbi_io_callbacks const *clbk, void *user) -{ - #ifndef STBI_NO_HDR - stbi__context s; - stbi__start_callbacks(&s, (stbi_io_callbacks *) clbk, user); - return stbi__hdr_test(&s); - #else - STBI_NOTUSED(clbk); - STBI_NOTUSED(user); - return 0; - #endif -} - -#ifndef STBI_NO_LINEAR -static float stbi__l2h_gamma=2.2f, stbi__l2h_scale=1.0f; - -STBIDEF void stbi_ldr_to_hdr_gamma(float gamma) { stbi__l2h_gamma = gamma; } -STBIDEF void stbi_ldr_to_hdr_scale(float scale) { stbi__l2h_scale = scale; } -#endif - -static float stbi__h2l_gamma_i=1.0f/2.2f, stbi__h2l_scale_i=1.0f; - -STBIDEF void stbi_hdr_to_ldr_gamma(float gamma) { stbi__h2l_gamma_i = 1/gamma; } -STBIDEF void stbi_hdr_to_ldr_scale(float scale) { stbi__h2l_scale_i = 1/scale; } - - -////////////////////////////////////////////////////////////////////////////// -// -// Common code used by all image loaders -// - -enum -{ - STBI__SCAN_load=0, - STBI__SCAN_type, - STBI__SCAN_header -}; - -static void stbi__refill_buffer(stbi__context *s) -{ - int n = (s->io.read)(s->io_user_data,(char*)s->buffer_start,s->buflen); - s->callback_already_read += (int) (s->img_buffer - s->img_buffer_original); - if (n == 0) { - // at end of file, treat same as if from memory, but need to handle case - // where s->img_buffer isn't pointing to safe memory, e.g. 0-byte file - s->read_from_callbacks = 0; - s->img_buffer = s->buffer_start; - s->img_buffer_end = s->buffer_start+1; - *s->img_buffer = 0; - } else { - s->img_buffer = s->buffer_start; - s->img_buffer_end = s->buffer_start + n; - } -} - -stbi_inline static stbi_uc stbi__get8(stbi__context *s) -{ - if (s->img_buffer < s->img_buffer_end) - return *s->img_buffer++; - if (s->read_from_callbacks) { - stbi__refill_buffer(s); - return *s->img_buffer++; - } - return 0; -} - -#if defined(STBI_NO_JPEG) && defined(STBI_NO_HDR) && defined(STBI_NO_PIC) && defined(STBI_NO_PNM) -// nothing -#else -stbi_inline static int stbi__at_eof(stbi__context *s) -{ - if (s->io.read) { - if (!(s->io.eof)(s->io_user_data)) return 0; - // if feof() is true, check if buffer = end - // special case: we've only got the special 0 character at the end - if (s->read_from_callbacks == 0) return 1; - } - - return s->img_buffer >= s->img_buffer_end; -} -#endif - -#if defined(STBI_NO_JPEG) && defined(STBI_NO_PNG) && defined(STBI_NO_BMP) && defined(STBI_NO_PSD) && defined(STBI_NO_TGA) && defined(STBI_NO_GIF) && defined(STBI_NO_PIC) -// nothing -#else -static void stbi__skip(stbi__context *s, int n) -{ - if (n == 0) return; // already there! - if (n < 0) { - s->img_buffer = s->img_buffer_end; - return; - } - if (s->io.read) { - int blen = (int) (s->img_buffer_end - s->img_buffer); - if (blen < n) { - s->img_buffer = s->img_buffer_end; - (s->io.skip)(s->io_user_data, n - blen); - return; - } - } - s->img_buffer += n; -} -#endif - -#if defined(STBI_NO_PNG) && defined(STBI_NO_TGA) && defined(STBI_NO_HDR) && defined(STBI_NO_PNM) -// nothing -#else -static int stbi__getn(stbi__context *s, stbi_uc *buffer, int n) -{ - if (s->io.read) { - int blen = (int) (s->img_buffer_end - s->img_buffer); - if (blen < n) { - int res, count; - - memcpy(buffer, s->img_buffer, blen); - - count = (s->io.read)(s->io_user_data, (char*) buffer + blen, n - blen); - res = (count == (n-blen)); - s->img_buffer = s->img_buffer_end; - return res; - } - } - - if (s->img_buffer+n <= s->img_buffer_end) { - memcpy(buffer, s->img_buffer, n); - s->img_buffer += n; - return 1; - } else - return 0; -} -#endif - -#if defined(STBI_NO_JPEG) && defined(STBI_NO_PNG) && defined(STBI_NO_PSD) && defined(STBI_NO_PIC) -// nothing -#else -static int stbi__get16be(stbi__context *s) -{ - int z = stbi__get8(s); - return (z << 8) + stbi__get8(s); -} -#endif - -#if defined(STBI_NO_PNG) && defined(STBI_NO_PSD) && defined(STBI_NO_PIC) -// nothing -#else -static stbi__uint32 stbi__get32be(stbi__context *s) -{ - stbi__uint32 z = stbi__get16be(s); - return (z << 16) + stbi__get16be(s); -} -#endif - -#if defined(STBI_NO_BMP) && defined(STBI_NO_TGA) && defined(STBI_NO_GIF) -// nothing -#else -static int stbi__get16le(stbi__context *s) -{ - int z = stbi__get8(s); - return z + (stbi__get8(s) << 8); -} -#endif - -#ifndef STBI_NO_BMP -static stbi__uint32 stbi__get32le(stbi__context *s) -{ - stbi__uint32 z = stbi__get16le(s); - z += (stbi__uint32)stbi__get16le(s) << 16; - return z; -} -#endif - -#define STBI__BYTECAST(x) ((stbi_uc) ((x) & 255)) // truncate int to byte without warnings - -#if defined(STBI_NO_JPEG) && defined(STBI_NO_PNG) && defined(STBI_NO_BMP) && defined(STBI_NO_PSD) && defined(STBI_NO_TGA) && defined(STBI_NO_GIF) && defined(STBI_NO_PIC) && defined(STBI_NO_PNM) -// nothing -#else -////////////////////////////////////////////////////////////////////////////// -// -// generic converter from built-in img_n to req_comp -// individual types do this automatically as much as possible (e.g. jpeg -// does all cases internally since it needs to colorspace convert anyway, -// and it never has alpha, so very few cases ). png can automatically -// interleave an alpha=255 channel, but falls back to this for other cases -// -// assume data buffer is malloced, so malloc a new one and free that one -// only failure mode is malloc failing - -static stbi_uc stbi__compute_y(int r, int g, int b) -{ - return (stbi_uc) (((r*77) + (g*150) + (29*b)) >> 8); -} -#endif - -#if defined(STBI_NO_PNG) && defined(STBI_NO_BMP) && defined(STBI_NO_PSD) && defined(STBI_NO_TGA) && defined(STBI_NO_GIF) && defined(STBI_NO_PIC) && defined(STBI_NO_PNM) -// nothing -#else -static unsigned char *stbi__convert_format(unsigned char *data, int img_n, int req_comp, unsigned int x, unsigned int y) -{ - int i,j; - unsigned char *good; - - if (req_comp == img_n) return data; - STBI_ASSERT(req_comp >= 1 && req_comp <= 4); - - good = (unsigned char *) stbi__malloc_mad3(req_comp, x, y, 0); - if (good == NULL) { - STBI_FREE(data); - return stbi__errpuc("outofmem", "Out of memory"); - } - - for (j=0; j < (int) y; ++j) { - unsigned char *src = data + j * x * img_n ; - unsigned char *dest = good + j * x * req_comp; - - #define STBI__COMBO(a,b) ((a)*8+(b)) - #define STBI__CASE(a,b) case STBI__COMBO(a,b): for(i=x-1; i >= 0; --i, src += a, dest += b) - // convert source image with img_n components to one with req_comp components; - // avoid switch per pixel, so use switch per scanline and massive macros - switch (STBI__COMBO(img_n, req_comp)) { - STBI__CASE(1,2) { dest[0]=src[0]; dest[1]=255; } break; - STBI__CASE(1,3) { dest[0]=dest[1]=dest[2]=src[0]; } break; - STBI__CASE(1,4) { dest[0]=dest[1]=dest[2]=src[0]; dest[3]=255; } break; - STBI__CASE(2,1) { dest[0]=src[0]; } break; - STBI__CASE(2,3) { dest[0]=dest[1]=dest[2]=src[0]; } break; - STBI__CASE(2,4) { dest[0]=dest[1]=dest[2]=src[0]; dest[3]=src[1]; } break; - STBI__CASE(3,4) { dest[0]=src[0];dest[1]=src[1];dest[2]=src[2];dest[3]=255; } break; - STBI__CASE(3,1) { dest[0]=stbi__compute_y(src[0],src[1],src[2]); } break; - STBI__CASE(3,2) { dest[0]=stbi__compute_y(src[0],src[1],src[2]); dest[1] = 255; } break; - STBI__CASE(4,1) { dest[0]=stbi__compute_y(src[0],src[1],src[2]); } break; - STBI__CASE(4,2) { dest[0]=stbi__compute_y(src[0],src[1],src[2]); dest[1] = src[3]; } break; - STBI__CASE(4,3) { dest[0]=src[0];dest[1]=src[1];dest[2]=src[2]; } break; - default: STBI_ASSERT(0); STBI_FREE(data); STBI_FREE(good); return stbi__errpuc("unsupported", "Unsupported format conversion"); - } - #undef STBI__CASE - } - - STBI_FREE(data); - return good; -} -#endif - -#if defined(STBI_NO_PNG) && defined(STBI_NO_PSD) -// nothing -#else -static stbi__uint16 stbi__compute_y_16(int r, int g, int b) -{ - return (stbi__uint16) (((r*77) + (g*150) + (29*b)) >> 8); -} -#endif - -#if defined(STBI_NO_PNG) && defined(STBI_NO_PSD) -// nothing -#else -static stbi__uint16 *stbi__convert_format16(stbi__uint16 *data, int img_n, int req_comp, unsigned int x, unsigned int y) -{ - int i,j; - stbi__uint16 *good; - - if (req_comp == img_n) return data; - STBI_ASSERT(req_comp >= 1 && req_comp <= 4); - - good = (stbi__uint16 *) stbi__malloc(req_comp * x * y * 2); - if (good == NULL) { - STBI_FREE(data); - return (stbi__uint16 *) stbi__errpuc("outofmem", "Out of memory"); - } - - for (j=0; j < (int) y; ++j) { - stbi__uint16 *src = data + j * x * img_n ; - stbi__uint16 *dest = good + j * x * req_comp; - - #define STBI__COMBO(a,b) ((a)*8+(b)) - #define STBI__CASE(a,b) case STBI__COMBO(a,b): for(i=x-1; i >= 0; --i, src += a, dest += b) - // convert source image with img_n components to one with req_comp components; - // avoid switch per pixel, so use switch per scanline and massive macros - switch (STBI__COMBO(img_n, req_comp)) { - STBI__CASE(1,2) { dest[0]=src[0]; dest[1]=0xffff; } break; - STBI__CASE(1,3) { dest[0]=dest[1]=dest[2]=src[0]; } break; - STBI__CASE(1,4) { dest[0]=dest[1]=dest[2]=src[0]; dest[3]=0xffff; } break; - STBI__CASE(2,1) { dest[0]=src[0]; } break; - STBI__CASE(2,3) { dest[0]=dest[1]=dest[2]=src[0]; } break; - STBI__CASE(2,4) { dest[0]=dest[1]=dest[2]=src[0]; dest[3]=src[1]; } break; - STBI__CASE(3,4) { dest[0]=src[0];dest[1]=src[1];dest[2]=src[2];dest[3]=0xffff; } break; - STBI__CASE(3,1) { dest[0]=stbi__compute_y_16(src[0],src[1],src[2]); } break; - STBI__CASE(3,2) { dest[0]=stbi__compute_y_16(src[0],src[1],src[2]); dest[1] = 0xffff; } break; - STBI__CASE(4,1) { dest[0]=stbi__compute_y_16(src[0],src[1],src[2]); } break; - STBI__CASE(4,2) { dest[0]=stbi__compute_y_16(src[0],src[1],src[2]); dest[1] = src[3]; } break; - STBI__CASE(4,3) { dest[0]=src[0];dest[1]=src[1];dest[2]=src[2]; } break; - default: STBI_ASSERT(0); STBI_FREE(data); STBI_FREE(good); return (stbi__uint16*) stbi__errpuc("unsupported", "Unsupported format conversion"); - } - #undef STBI__CASE - } - - STBI_FREE(data); - return good; -} -#endif - -#ifndef STBI_NO_LINEAR -static float *stbi__ldr_to_hdr(stbi_uc *data, int x, int y, int comp) -{ - int i,k,n; - float *output; - if (!data) return NULL; - output = (float *) stbi__malloc_mad4(x, y, comp, sizeof(float), 0); - if (output == NULL) { STBI_FREE(data); return stbi__errpf("outofmem", "Out of memory"); } - // compute number of non-alpha components - if (comp & 1) n = comp; else n = comp-1; - for (i=0; i < x*y; ++i) { - for (k=0; k < n; ++k) { - output[i*comp + k] = (float) (pow(data[i*comp+k]/255.0f, stbi__l2h_gamma) * stbi__l2h_scale); - } - } - if (n < comp) { - for (i=0; i < x*y; ++i) { - output[i*comp + n] = data[i*comp + n]/255.0f; - } - } - STBI_FREE(data); - return output; -} -#endif - -#ifndef STBI_NO_HDR -#define stbi__float2int(x) ((int) (x)) -static stbi_uc *stbi__hdr_to_ldr(float *data, int x, int y, int comp) -{ - int i,k,n; - stbi_uc *output; - if (!data) return NULL; - output = (stbi_uc *) stbi__malloc_mad3(x, y, comp, 0); - if (output == NULL) { STBI_FREE(data); return stbi__errpuc("outofmem", "Out of memory"); } - // compute number of non-alpha components - if (comp & 1) n = comp; else n = comp-1; - for (i=0; i < x*y; ++i) { - for (k=0; k < n; ++k) { - float z = (float) pow(data[i*comp+k]*stbi__h2l_scale_i, stbi__h2l_gamma_i) * 255 + 0.5f; - if (z < 0) z = 0; - if (z > 255) z = 255; - output[i*comp + k] = (stbi_uc) stbi__float2int(z); - } - if (k < comp) { - float z = data[i*comp+k] * 255 + 0.5f; - if (z < 0) z = 0; - if (z > 255) z = 255; - output[i*comp + k] = (stbi_uc) stbi__float2int(z); - } - } - STBI_FREE(data); - return output; -} -#endif - -////////////////////////////////////////////////////////////////////////////// -// -// "baseline" JPEG/JFIF decoder -// -// simple implementation -// - doesn't support delayed output of y-dimension -// - simple interface (only one output format: 8-bit interleaved RGB) -// - doesn't try to recover corrupt jpegs -// - doesn't allow partial loading, loading multiple at once -// - still fast on x86 (copying globals into locals doesn't help x86) -// - allocates lots of intermediate memory (full size of all components) -// - non-interleaved case requires this anyway -// - allows good upsampling (see next) -// high-quality -// - upsampled channels are bilinearly interpolated, even across blocks -// - quality integer IDCT derived from IJG's 'slow' -// performance -// - fast huffman; reasonable integer IDCT -// - some SIMD kernels for common paths on targets with SSE2/NEON -// - uses a lot of intermediate memory, could cache poorly - -#ifndef STBI_NO_JPEG - -// huffman decoding acceleration -#define FAST_BITS 9 // larger handles more cases; smaller stomps less cache - -typedef struct -{ - stbi_uc fast[1 << FAST_BITS]; - // weirdly, repacking this into AoS is a 10% speed loss, instead of a win - stbi__uint16 code[256]; - stbi_uc values[256]; - stbi_uc size[257]; - unsigned int maxcode[18]; - int delta[17]; // old 'firstsymbol' - old 'firstcode' -} stbi__huffman; - -typedef struct -{ - stbi__context *s; - stbi__huffman huff_dc[4]; - stbi__huffman huff_ac[4]; - stbi__uint16 dequant[4][64]; - stbi__int16 fast_ac[4][1 << FAST_BITS]; - -// sizes for components, interleaved MCUs - int img_h_max, img_v_max; - int img_mcu_x, img_mcu_y; - int img_mcu_w, img_mcu_h; - -// definition of jpeg image component - struct - { - int id; - int h,v; - int tq; - int hd,ha; - int dc_pred; - - int x,y,w2,h2; - stbi_uc *data; - void *raw_data, *raw_coeff; - stbi_uc *linebuf; - short *coeff; // progressive only - int coeff_w, coeff_h; // number of 8x8 coefficient blocks - } img_comp[4]; - - stbi__uint32 code_buffer; // jpeg entropy-coded buffer - int code_bits; // number of valid bits - unsigned char marker; // marker seen while filling entropy buffer - int nomore; // flag if we saw a marker so must stop - - int progressive; - int spec_start; - int spec_end; - int succ_high; - int succ_low; - int eob_run; - int jfif; - int app14_color_transform; // Adobe APP14 tag - int rgb; - - int scan_n, order[4]; - int restart_interval, todo; - -// kernels - void (*idct_block_kernel)(stbi_uc *out, int out_stride, short data[64]); - void (*YCbCr_to_RGB_kernel)(stbi_uc *out, const stbi_uc *y, const stbi_uc *pcb, const stbi_uc *pcr, int count, int step); - stbi_uc *(*resample_row_hv_2_kernel)(stbi_uc *out, stbi_uc *in_near, stbi_uc *in_far, int w, int hs); -} stbi__jpeg; - -static int stbi__build_huffman(stbi__huffman *h, int *count) -{ - int i,j,k=0; - unsigned int code; - // build size list for each symbol (from JPEG spec) - for (i=0; i < 16; ++i) { - for (j=0; j < count[i]; ++j) { - h->size[k++] = (stbi_uc) (i+1); - if(k >= 257) return stbi__err("bad size list","Corrupt JPEG"); - } - } - h->size[k] = 0; - - // compute actual symbols (from jpeg spec) - code = 0; - k = 0; - for(j=1; j <= 16; ++j) { - // compute delta to add to code to compute symbol id - h->delta[j] = k - code; - if (h->size[k] == j) { - while (h->size[k] == j) - h->code[k++] = (stbi__uint16) (code++); - if (code-1 >= (1u << j)) return stbi__err("bad code lengths","Corrupt JPEG"); - } - // compute largest code + 1 for this size, preshifted as needed later - h->maxcode[j] = code << (16-j); - code <<= 1; - } - h->maxcode[j] = 0xffffffff; - - // build non-spec acceleration table; 255 is flag for not-accelerated - memset(h->fast, 255, 1 << FAST_BITS); - for (i=0; i < k; ++i) { - int s = h->size[i]; - if (s <= FAST_BITS) { - int c = h->code[i] << (FAST_BITS-s); - int m = 1 << (FAST_BITS-s); - for (j=0; j < m; ++j) { - h->fast[c+j] = (stbi_uc) i; - } - } - } - return 1; -} - -// build a table that decodes both magnitude and value of small ACs in -// one go. -static void stbi__build_fast_ac(stbi__int16 *fast_ac, stbi__huffman *h) -{ - int i; - for (i=0; i < (1 << FAST_BITS); ++i) { - stbi_uc fast = h->fast[i]; - fast_ac[i] = 0; - if (fast < 255) { - int rs = h->values[fast]; - int run = (rs >> 4) & 15; - int magbits = rs & 15; - int len = h->size[fast]; - - if (magbits && len + magbits <= FAST_BITS) { - // magnitude code followed by receive_extend code - int k = ((i << len) & ((1 << FAST_BITS) - 1)) >> (FAST_BITS - magbits); - int m = 1 << (magbits - 1); - if (k < m) k += (~0U << magbits) + 1; - // if the result is small enough, we can fit it in fast_ac table - if (k >= -128 && k <= 127) - fast_ac[i] = (stbi__int16) ((k * 256) + (run * 16) + (len + magbits)); - } - } - } -} - -static void stbi__grow_buffer_unsafe(stbi__jpeg *j) -{ - do { - unsigned int b = j->nomore ? 0 : stbi__get8(j->s); - if (b == 0xff) { - int c = stbi__get8(j->s); - while (c == 0xff) c = stbi__get8(j->s); // consume fill bytes - if (c != 0) { - j->marker = (unsigned char) c; - j->nomore = 1; - return; - } - } - j->code_buffer |= b << (24 - j->code_bits); - j->code_bits += 8; - } while (j->code_bits <= 24); -} - -// (1 << n) - 1 -static const stbi__uint32 stbi__bmask[17]={0,1,3,7,15,31,63,127,255,511,1023,2047,4095,8191,16383,32767,65535}; - -// decode a jpeg huffman value from the bitstream -stbi_inline static int stbi__jpeg_huff_decode(stbi__jpeg *j, stbi__huffman *h) -{ - unsigned int temp; - int c,k; - - if (j->code_bits < 16) stbi__grow_buffer_unsafe(j); - - // look at the top FAST_BITS and determine what symbol ID it is, - // if the code is <= FAST_BITS - c = (j->code_buffer >> (32 - FAST_BITS)) & ((1 << FAST_BITS)-1); - k = h->fast[c]; - if (k < 255) { - int s = h->size[k]; - if (s > j->code_bits) - return -1; - j->code_buffer <<= s; - j->code_bits -= s; - return h->values[k]; - } - - // naive test is to shift the code_buffer down so k bits are - // valid, then test against maxcode. To speed this up, we've - // preshifted maxcode left so that it has (16-k) 0s at the - // end; in other words, regardless of the number of bits, it - // wants to be compared against something shifted to have 16; - // that way we don't need to shift inside the loop. - temp = j->code_buffer >> 16; - for (k=FAST_BITS+1 ; ; ++k) - if (temp < h->maxcode[k]) - break; - if (k == 17) { - // error! code not found - j->code_bits -= 16; - return -1; - } - - if (k > j->code_bits) - return -1; - - // convert the huffman code to the symbol id - c = ((j->code_buffer >> (32 - k)) & stbi__bmask[k]) + h->delta[k]; - if(c < 0 || c >= 256) // symbol id out of bounds! - return -1; - STBI_ASSERT((((j->code_buffer) >> (32 - h->size[c])) & stbi__bmask[h->size[c]]) == h->code[c]); - - // convert the id to a symbol - j->code_bits -= k; - j->code_buffer <<= k; - return h->values[c]; -} - -// bias[n] = (-1<code_bits < n) stbi__grow_buffer_unsafe(j); - if (j->code_bits < n) return 0; // ran out of bits from stream, return 0s intead of continuing - - sgn = j->code_buffer >> 31; // sign bit always in MSB; 0 if MSB clear (positive), 1 if MSB set (negative) - k = stbi_lrot(j->code_buffer, n); - j->code_buffer = k & ~stbi__bmask[n]; - k &= stbi__bmask[n]; - j->code_bits -= n; - return k + (stbi__jbias[n] & (sgn - 1)); -} - -// get some unsigned bits -stbi_inline static int stbi__jpeg_get_bits(stbi__jpeg *j, int n) -{ - unsigned int k; - if (j->code_bits < n) stbi__grow_buffer_unsafe(j); - if (j->code_bits < n) return 0; // ran out of bits from stream, return 0s intead of continuing - k = stbi_lrot(j->code_buffer, n); - j->code_buffer = k & ~stbi__bmask[n]; - k &= stbi__bmask[n]; - j->code_bits -= n; - return k; -} - -stbi_inline static int stbi__jpeg_get_bit(stbi__jpeg *j) -{ - unsigned int k; - if (j->code_bits < 1) stbi__grow_buffer_unsafe(j); - if (j->code_bits < 1) return 0; // ran out of bits from stream, return 0s intead of continuing - k = j->code_buffer; - j->code_buffer <<= 1; - --j->code_bits; - return k & 0x80000000; -} - -// given a value that's at position X in the zigzag stream, -// where does it appear in the 8x8 matrix coded as row-major? -static const stbi_uc stbi__jpeg_dezigzag[64+15] = -{ - 0, 1, 8, 16, 9, 2, 3, 10, - 17, 24, 32, 25, 18, 11, 4, 5, - 12, 19, 26, 33, 40, 48, 41, 34, - 27, 20, 13, 6, 7, 14, 21, 28, - 35, 42, 49, 56, 57, 50, 43, 36, - 29, 22, 15, 23, 30, 37, 44, 51, - 58, 59, 52, 45, 38, 31, 39, 46, - 53, 60, 61, 54, 47, 55, 62, 63, - // let corrupt input sample past end - 63, 63, 63, 63, 63, 63, 63, 63, - 63, 63, 63, 63, 63, 63, 63 -}; - -// decode one 64-entry block-- -static int stbi__jpeg_decode_block(stbi__jpeg *j, short data[64], stbi__huffman *hdc, stbi__huffman *hac, stbi__int16 *fac, int b, stbi__uint16 *dequant) -{ - int diff,dc,k; - int t; - - if (j->code_bits < 16) stbi__grow_buffer_unsafe(j); - t = stbi__jpeg_huff_decode(j, hdc); - if (t < 0 || t > 15) return stbi__err("bad huffman code","Corrupt JPEG"); - - // 0 all the ac values now so we can do it 32-bits at a time - memset(data,0,64*sizeof(data[0])); - - diff = t ? stbi__extend_receive(j, t) : 0; - if (!stbi__addints_valid(j->img_comp[b].dc_pred, diff)) return stbi__err("bad delta","Corrupt JPEG"); - dc = j->img_comp[b].dc_pred + diff; - j->img_comp[b].dc_pred = dc; - if (!stbi__mul2shorts_valid(dc, dequant[0])) return stbi__err("can't merge dc and ac", "Corrupt JPEG"); - data[0] = (short) (dc * dequant[0]); - - // decode AC components, see JPEG spec - k = 1; - do { - unsigned int zig; - int c,r,s; - if (j->code_bits < 16) stbi__grow_buffer_unsafe(j); - c = (j->code_buffer >> (32 - FAST_BITS)) & ((1 << FAST_BITS)-1); - r = fac[c]; - if (r) { // fast-AC path - k += (r >> 4) & 15; // run - s = r & 15; // combined length - if (s > j->code_bits) return stbi__err("bad huffman code", "Combined length longer than code bits available"); - j->code_buffer <<= s; - j->code_bits -= s; - // decode into unzigzag'd location - zig = stbi__jpeg_dezigzag[k++]; - data[zig] = (short) ((r >> 8) * dequant[zig]); - } else { - int rs = stbi__jpeg_huff_decode(j, hac); - if (rs < 0) return stbi__err("bad huffman code","Corrupt JPEG"); - s = rs & 15; - r = rs >> 4; - if (s == 0) { - if (rs != 0xf0) break; // end block - k += 16; - } else { - k += r; - // decode into unzigzag'd location - zig = stbi__jpeg_dezigzag[k++]; - data[zig] = (short) (stbi__extend_receive(j,s) * dequant[zig]); - } - } - } while (k < 64); - return 1; -} - -static int stbi__jpeg_decode_block_prog_dc(stbi__jpeg *j, short data[64], stbi__huffman *hdc, int b) -{ - int diff,dc; - int t; - if (j->spec_end != 0) return stbi__err("can't merge dc and ac", "Corrupt JPEG"); - - if (j->code_bits < 16) stbi__grow_buffer_unsafe(j); - - if (j->succ_high == 0) { - // first scan for DC coefficient, must be first - memset(data,0,64*sizeof(data[0])); // 0 all the ac values now - t = stbi__jpeg_huff_decode(j, hdc); - if (t < 0 || t > 15) return stbi__err("can't merge dc and ac", "Corrupt JPEG"); - diff = t ? stbi__extend_receive(j, t) : 0; - - if (!stbi__addints_valid(j->img_comp[b].dc_pred, diff)) return stbi__err("bad delta", "Corrupt JPEG"); - dc = j->img_comp[b].dc_pred + diff; - j->img_comp[b].dc_pred = dc; - if (!stbi__mul2shorts_valid(dc, 1 << j->succ_low)) return stbi__err("can't merge dc and ac", "Corrupt JPEG"); - data[0] = (short) (dc * (1 << j->succ_low)); - } else { - // refinement scan for DC coefficient - if (stbi__jpeg_get_bit(j)) - data[0] += (short) (1 << j->succ_low); - } - return 1; -} - -// @OPTIMIZE: store non-zigzagged during the decode passes, -// and only de-zigzag when dequantizing -static int stbi__jpeg_decode_block_prog_ac(stbi__jpeg *j, short data[64], stbi__huffman *hac, stbi__int16 *fac) -{ - int k; - if (j->spec_start == 0) return stbi__err("can't merge dc and ac", "Corrupt JPEG"); - - if (j->succ_high == 0) { - int shift = j->succ_low; - - if (j->eob_run) { - --j->eob_run; - return 1; - } - - k = j->spec_start; - do { - unsigned int zig; - int c,r,s; - if (j->code_bits < 16) stbi__grow_buffer_unsafe(j); - c = (j->code_buffer >> (32 - FAST_BITS)) & ((1 << FAST_BITS)-1); - r = fac[c]; - if (r) { // fast-AC path - k += (r >> 4) & 15; // run - s = r & 15; // combined length - if (s > j->code_bits) return stbi__err("bad huffman code", "Combined length longer than code bits available"); - j->code_buffer <<= s; - j->code_bits -= s; - zig = stbi__jpeg_dezigzag[k++]; - data[zig] = (short) ((r >> 8) * (1 << shift)); - } else { - int rs = stbi__jpeg_huff_decode(j, hac); - if (rs < 0) return stbi__err("bad huffman code","Corrupt JPEG"); - s = rs & 15; - r = rs >> 4; - if (s == 0) { - if (r < 15) { - j->eob_run = (1 << r); - if (r) - j->eob_run += stbi__jpeg_get_bits(j, r); - --j->eob_run; - break; - } - k += 16; - } else { - k += r; - zig = stbi__jpeg_dezigzag[k++]; - data[zig] = (short) (stbi__extend_receive(j,s) * (1 << shift)); - } - } - } while (k <= j->spec_end); - } else { - // refinement scan for these AC coefficients - - short bit = (short) (1 << j->succ_low); - - if (j->eob_run) { - --j->eob_run; - for (k = j->spec_start; k <= j->spec_end; ++k) { - short *p = &data[stbi__jpeg_dezigzag[k]]; - if (*p != 0) - if (stbi__jpeg_get_bit(j)) - if ((*p & bit)==0) { - if (*p > 0) - *p += bit; - else - *p -= bit; - } - } - } else { - k = j->spec_start; - do { - int r,s; - int rs = stbi__jpeg_huff_decode(j, hac); // @OPTIMIZE see if we can use the fast path here, advance-by-r is so slow, eh - if (rs < 0) return stbi__err("bad huffman code","Corrupt JPEG"); - s = rs & 15; - r = rs >> 4; - if (s == 0) { - if (r < 15) { - j->eob_run = (1 << r) - 1; - if (r) - j->eob_run += stbi__jpeg_get_bits(j, r); - r = 64; // force end of block - } else { - // r=15 s=0 should write 16 0s, so we just do - // a run of 15 0s and then write s (which is 0), - // so we don't have to do anything special here - } - } else { - if (s != 1) return stbi__err("bad huffman code", "Corrupt JPEG"); - // sign bit - if (stbi__jpeg_get_bit(j)) - s = bit; - else - s = -bit; - } - - // advance by r - while (k <= j->spec_end) { - short *p = &data[stbi__jpeg_dezigzag[k++]]; - if (*p != 0) { - if (stbi__jpeg_get_bit(j)) - if ((*p & bit)==0) { - if (*p > 0) - *p += bit; - else - *p -= bit; - } - } else { - if (r == 0) { - *p = (short) s; - break; - } - --r; - } - } - } while (k <= j->spec_end); - } - } - return 1; -} - -// take a -128..127 value and stbi__clamp it and convert to 0..255 -stbi_inline static stbi_uc stbi__clamp(int x) -{ - // trick to use a single test to catch both cases - if ((unsigned int) x > 255) { - if (x < 0) return 0; - if (x > 255) return 255; - } - return (stbi_uc) x; -} - -#define stbi__f2f(x) ((int) (((x) * 4096 + 0.5))) -#define stbi__fsh(x) ((x) * 4096) - -// derived from jidctint -- DCT_ISLOW -#define STBI__IDCT_1D(s0,s1,s2,s3,s4,s5,s6,s7) \ - int t0,t1,t2,t3,p1,p2,p3,p4,p5,x0,x1,x2,x3; \ - p2 = s2; \ - p3 = s6; \ - p1 = (p2+p3) * stbi__f2f(0.5411961f); \ - t2 = p1 + p3*stbi__f2f(-1.847759065f); \ - t3 = p1 + p2*stbi__f2f( 0.765366865f); \ - p2 = s0; \ - p3 = s4; \ - t0 = stbi__fsh(p2+p3); \ - t1 = stbi__fsh(p2-p3); \ - x0 = t0+t3; \ - x3 = t0-t3; \ - x1 = t1+t2; \ - x2 = t1-t2; \ - t0 = s7; \ - t1 = s5; \ - t2 = s3; \ - t3 = s1; \ - p3 = t0+t2; \ - p4 = t1+t3; \ - p1 = t0+t3; \ - p2 = t1+t2; \ - p5 = (p3+p4)*stbi__f2f( 1.175875602f); \ - t0 = t0*stbi__f2f( 0.298631336f); \ - t1 = t1*stbi__f2f( 2.053119869f); \ - t2 = t2*stbi__f2f( 3.072711026f); \ - t3 = t3*stbi__f2f( 1.501321110f); \ - p1 = p5 + p1*stbi__f2f(-0.899976223f); \ - p2 = p5 + p2*stbi__f2f(-2.562915447f); \ - p3 = p3*stbi__f2f(-1.961570560f); \ - p4 = p4*stbi__f2f(-0.390180644f); \ - t3 += p1+p4; \ - t2 += p2+p3; \ - t1 += p2+p4; \ - t0 += p1+p3; - -static void stbi__idct_block(stbi_uc *out, int out_stride, short data[64]) -{ - int i,val[64],*v=val; - stbi_uc *o; - short *d = data; - - // columns - for (i=0; i < 8; ++i,++d, ++v) { - // if all zeroes, shortcut -- this avoids dequantizing 0s and IDCTing - if (d[ 8]==0 && d[16]==0 && d[24]==0 && d[32]==0 - && d[40]==0 && d[48]==0 && d[56]==0) { - // no shortcut 0 seconds - // (1|2|3|4|5|6|7)==0 0 seconds - // all separate -0.047 seconds - // 1 && 2|3 && 4|5 && 6|7: -0.047 seconds - int dcterm = d[0]*4; - v[0] = v[8] = v[16] = v[24] = v[32] = v[40] = v[48] = v[56] = dcterm; - } else { - STBI__IDCT_1D(d[ 0],d[ 8],d[16],d[24],d[32],d[40],d[48],d[56]) - // constants scaled things up by 1<<12; let's bring them back - // down, but keep 2 extra bits of precision - x0 += 512; x1 += 512; x2 += 512; x3 += 512; - v[ 0] = (x0+t3) >> 10; - v[56] = (x0-t3) >> 10; - v[ 8] = (x1+t2) >> 10; - v[48] = (x1-t2) >> 10; - v[16] = (x2+t1) >> 10; - v[40] = (x2-t1) >> 10; - v[24] = (x3+t0) >> 10; - v[32] = (x3-t0) >> 10; - } - } - - for (i=0, v=val, o=out; i < 8; ++i,v+=8,o+=out_stride) { - // no fast case since the first 1D IDCT spread components out - STBI__IDCT_1D(v[0],v[1],v[2],v[3],v[4],v[5],v[6],v[7]) - // constants scaled things up by 1<<12, plus we had 1<<2 from first - // loop, plus horizontal and vertical each scale by sqrt(8) so together - // we've got an extra 1<<3, so 1<<17 total we need to remove. - // so we want to round that, which means adding 0.5 * 1<<17, - // aka 65536. Also, we'll end up with -128 to 127 that we want - // to encode as 0..255 by adding 128, so we'll add that before the shift - x0 += 65536 + (128<<17); - x1 += 65536 + (128<<17); - x2 += 65536 + (128<<17); - x3 += 65536 + (128<<17); - // tried computing the shifts into temps, or'ing the temps to see - // if any were out of range, but that was slower - o[0] = stbi__clamp((x0+t3) >> 17); - o[7] = stbi__clamp((x0-t3) >> 17); - o[1] = stbi__clamp((x1+t2) >> 17); - o[6] = stbi__clamp((x1-t2) >> 17); - o[2] = stbi__clamp((x2+t1) >> 17); - o[5] = stbi__clamp((x2-t1) >> 17); - o[3] = stbi__clamp((x3+t0) >> 17); - o[4] = stbi__clamp((x3-t0) >> 17); - } -} - -#ifdef STBI_SSE2 -// sse2 integer IDCT. not the fastest possible implementation but it -// produces bit-identical results to the generic C version so it's -// fully "transparent". -static void stbi__idct_simd(stbi_uc *out, int out_stride, short data[64]) -{ - // This is constructed to match our regular (generic) integer IDCT exactly. - __m128i row0, row1, row2, row3, row4, row5, row6, row7; - __m128i tmp; - - // dot product constant: even elems=x, odd elems=y - #define dct_const(x,y) _mm_setr_epi16((x),(y),(x),(y),(x),(y),(x),(y)) - - // out(0) = c0[even]*x + c0[odd]*y (c0, x, y 16-bit, out 32-bit) - // out(1) = c1[even]*x + c1[odd]*y - #define dct_rot(out0,out1, x,y,c0,c1) \ - __m128i c0##lo = _mm_unpacklo_epi16((x),(y)); \ - __m128i c0##hi = _mm_unpackhi_epi16((x),(y)); \ - __m128i out0##_l = _mm_madd_epi16(c0##lo, c0); \ - __m128i out0##_h = _mm_madd_epi16(c0##hi, c0); \ - __m128i out1##_l = _mm_madd_epi16(c0##lo, c1); \ - __m128i out1##_h = _mm_madd_epi16(c0##hi, c1) - - // out = in << 12 (in 16-bit, out 32-bit) - #define dct_widen(out, in) \ - __m128i out##_l = _mm_srai_epi32(_mm_unpacklo_epi16(_mm_setzero_si128(), (in)), 4); \ - __m128i out##_h = _mm_srai_epi32(_mm_unpackhi_epi16(_mm_setzero_si128(), (in)), 4) - - // wide add - #define dct_wadd(out, a, b) \ - __m128i out##_l = _mm_add_epi32(a##_l, b##_l); \ - __m128i out##_h = _mm_add_epi32(a##_h, b##_h) - - // wide sub - #define dct_wsub(out, a, b) \ - __m128i out##_l = _mm_sub_epi32(a##_l, b##_l); \ - __m128i out##_h = _mm_sub_epi32(a##_h, b##_h) - - // butterfly a/b, add bias, then shift by "s" and pack - #define dct_bfly32o(out0, out1, a,b,bias,s) \ - { \ - __m128i abiased_l = _mm_add_epi32(a##_l, bias); \ - __m128i abiased_h = _mm_add_epi32(a##_h, bias); \ - dct_wadd(sum, abiased, b); \ - dct_wsub(dif, abiased, b); \ - out0 = _mm_packs_epi32(_mm_srai_epi32(sum_l, s), _mm_srai_epi32(sum_h, s)); \ - out1 = _mm_packs_epi32(_mm_srai_epi32(dif_l, s), _mm_srai_epi32(dif_h, s)); \ - } - - // 8-bit interleave step (for transposes) - #define dct_interleave8(a, b) \ - tmp = a; \ - a = _mm_unpacklo_epi8(a, b); \ - b = _mm_unpackhi_epi8(tmp, b) - - // 16-bit interleave step (for transposes) - #define dct_interleave16(a, b) \ - tmp = a; \ - a = _mm_unpacklo_epi16(a, b); \ - b = _mm_unpackhi_epi16(tmp, b) - - #define dct_pass(bias,shift) \ - { \ - /* even part */ \ - dct_rot(t2e,t3e, row2,row6, rot0_0,rot0_1); \ - __m128i sum04 = _mm_add_epi16(row0, row4); \ - __m128i dif04 = _mm_sub_epi16(row0, row4); \ - dct_widen(t0e, sum04); \ - dct_widen(t1e, dif04); \ - dct_wadd(x0, t0e, t3e); \ - dct_wsub(x3, t0e, t3e); \ - dct_wadd(x1, t1e, t2e); \ - dct_wsub(x2, t1e, t2e); \ - /* odd part */ \ - dct_rot(y0o,y2o, row7,row3, rot2_0,rot2_1); \ - dct_rot(y1o,y3o, row5,row1, rot3_0,rot3_1); \ - __m128i sum17 = _mm_add_epi16(row1, row7); \ - __m128i sum35 = _mm_add_epi16(row3, row5); \ - dct_rot(y4o,y5o, sum17,sum35, rot1_0,rot1_1); \ - dct_wadd(x4, y0o, y4o); \ - dct_wadd(x5, y1o, y5o); \ - dct_wadd(x6, y2o, y5o); \ - dct_wadd(x7, y3o, y4o); \ - dct_bfly32o(row0,row7, x0,x7,bias,shift); \ - dct_bfly32o(row1,row6, x1,x6,bias,shift); \ - dct_bfly32o(row2,row5, x2,x5,bias,shift); \ - dct_bfly32o(row3,row4, x3,x4,bias,shift); \ - } - - __m128i rot0_0 = dct_const(stbi__f2f(0.5411961f), stbi__f2f(0.5411961f) + stbi__f2f(-1.847759065f)); - __m128i rot0_1 = dct_const(stbi__f2f(0.5411961f) + stbi__f2f( 0.765366865f), stbi__f2f(0.5411961f)); - __m128i rot1_0 = dct_const(stbi__f2f(1.175875602f) + stbi__f2f(-0.899976223f), stbi__f2f(1.175875602f)); - __m128i rot1_1 = dct_const(stbi__f2f(1.175875602f), stbi__f2f(1.175875602f) + stbi__f2f(-2.562915447f)); - __m128i rot2_0 = dct_const(stbi__f2f(-1.961570560f) + stbi__f2f( 0.298631336f), stbi__f2f(-1.961570560f)); - __m128i rot2_1 = dct_const(stbi__f2f(-1.961570560f), stbi__f2f(-1.961570560f) + stbi__f2f( 3.072711026f)); - __m128i rot3_0 = dct_const(stbi__f2f(-0.390180644f) + stbi__f2f( 2.053119869f), stbi__f2f(-0.390180644f)); - __m128i rot3_1 = dct_const(stbi__f2f(-0.390180644f), stbi__f2f(-0.390180644f) + stbi__f2f( 1.501321110f)); - - // rounding biases in column/row passes, see stbi__idct_block for explanation. - __m128i bias_0 = _mm_set1_epi32(512); - __m128i bias_1 = _mm_set1_epi32(65536 + (128<<17)); - - // load - row0 = _mm_load_si128((const __m128i *) (data + 0*8)); - row1 = _mm_load_si128((const __m128i *) (data + 1*8)); - row2 = _mm_load_si128((const __m128i *) (data + 2*8)); - row3 = _mm_load_si128((const __m128i *) (data + 3*8)); - row4 = _mm_load_si128((const __m128i *) (data + 4*8)); - row5 = _mm_load_si128((const __m128i *) (data + 5*8)); - row6 = _mm_load_si128((const __m128i *) (data + 6*8)); - row7 = _mm_load_si128((const __m128i *) (data + 7*8)); - - // column pass - dct_pass(bias_0, 10); - - { - // 16bit 8x8 transpose pass 1 - dct_interleave16(row0, row4); - dct_interleave16(row1, row5); - dct_interleave16(row2, row6); - dct_interleave16(row3, row7); - - // transpose pass 2 - dct_interleave16(row0, row2); - dct_interleave16(row1, row3); - dct_interleave16(row4, row6); - dct_interleave16(row5, row7); - - // transpose pass 3 - dct_interleave16(row0, row1); - dct_interleave16(row2, row3); - dct_interleave16(row4, row5); - dct_interleave16(row6, row7); - } - - // row pass - dct_pass(bias_1, 17); - - { - // pack - __m128i p0 = _mm_packus_epi16(row0, row1); // a0a1a2a3...a7b0b1b2b3...b7 - __m128i p1 = _mm_packus_epi16(row2, row3); - __m128i p2 = _mm_packus_epi16(row4, row5); - __m128i p3 = _mm_packus_epi16(row6, row7); - - // 8bit 8x8 transpose pass 1 - dct_interleave8(p0, p2); // a0e0a1e1... - dct_interleave8(p1, p3); // c0g0c1g1... - - // transpose pass 2 - dct_interleave8(p0, p1); // a0c0e0g0... - dct_interleave8(p2, p3); // b0d0f0h0... - - // transpose pass 3 - dct_interleave8(p0, p2); // a0b0c0d0... - dct_interleave8(p1, p3); // a4b4c4d4... - - // store - _mm_storel_epi64((__m128i *) out, p0); out += out_stride; - _mm_storel_epi64((__m128i *) out, _mm_shuffle_epi32(p0, 0x4e)); out += out_stride; - _mm_storel_epi64((__m128i *) out, p2); out += out_stride; - _mm_storel_epi64((__m128i *) out, _mm_shuffle_epi32(p2, 0x4e)); out += out_stride; - _mm_storel_epi64((__m128i *) out, p1); out += out_stride; - _mm_storel_epi64((__m128i *) out, _mm_shuffle_epi32(p1, 0x4e)); out += out_stride; - _mm_storel_epi64((__m128i *) out, p3); out += out_stride; - _mm_storel_epi64((__m128i *) out, _mm_shuffle_epi32(p3, 0x4e)); - } - -#undef dct_const -#undef dct_rot -#undef dct_widen -#undef dct_wadd -#undef dct_wsub -#undef dct_bfly32o -#undef dct_interleave8 -#undef dct_interleave16 -#undef dct_pass -} - -#endif // STBI_SSE2 - -#ifdef STBI_NEON - -// NEON integer IDCT. should produce bit-identical -// results to the generic C version. -static void stbi__idct_simd(stbi_uc *out, int out_stride, short data[64]) -{ - int16x8_t row0, row1, row2, row3, row4, row5, row6, row7; - - int16x4_t rot0_0 = vdup_n_s16(stbi__f2f(0.5411961f)); - int16x4_t rot0_1 = vdup_n_s16(stbi__f2f(-1.847759065f)); - int16x4_t rot0_2 = vdup_n_s16(stbi__f2f( 0.765366865f)); - int16x4_t rot1_0 = vdup_n_s16(stbi__f2f( 1.175875602f)); - int16x4_t rot1_1 = vdup_n_s16(stbi__f2f(-0.899976223f)); - int16x4_t rot1_2 = vdup_n_s16(stbi__f2f(-2.562915447f)); - int16x4_t rot2_0 = vdup_n_s16(stbi__f2f(-1.961570560f)); - int16x4_t rot2_1 = vdup_n_s16(stbi__f2f(-0.390180644f)); - int16x4_t rot3_0 = vdup_n_s16(stbi__f2f( 0.298631336f)); - int16x4_t rot3_1 = vdup_n_s16(stbi__f2f( 2.053119869f)); - int16x4_t rot3_2 = vdup_n_s16(stbi__f2f( 3.072711026f)); - int16x4_t rot3_3 = vdup_n_s16(stbi__f2f( 1.501321110f)); - -#define dct_long_mul(out, inq, coeff) \ - int32x4_t out##_l = vmull_s16(vget_low_s16(inq), coeff); \ - int32x4_t out##_h = vmull_s16(vget_high_s16(inq), coeff) - -#define dct_long_mac(out, acc, inq, coeff) \ - int32x4_t out##_l = vmlal_s16(acc##_l, vget_low_s16(inq), coeff); \ - int32x4_t out##_h = vmlal_s16(acc##_h, vget_high_s16(inq), coeff) - -#define dct_widen(out, inq) \ - int32x4_t out##_l = vshll_n_s16(vget_low_s16(inq), 12); \ - int32x4_t out##_h = vshll_n_s16(vget_high_s16(inq), 12) - -// wide add -#define dct_wadd(out, a, b) \ - int32x4_t out##_l = vaddq_s32(a##_l, b##_l); \ - int32x4_t out##_h = vaddq_s32(a##_h, b##_h) - -// wide sub -#define dct_wsub(out, a, b) \ - int32x4_t out##_l = vsubq_s32(a##_l, b##_l); \ - int32x4_t out##_h = vsubq_s32(a##_h, b##_h) - -// butterfly a/b, then shift using "shiftop" by "s" and pack -#define dct_bfly32o(out0,out1, a,b,shiftop,s) \ - { \ - dct_wadd(sum, a, b); \ - dct_wsub(dif, a, b); \ - out0 = vcombine_s16(shiftop(sum_l, s), shiftop(sum_h, s)); \ - out1 = vcombine_s16(shiftop(dif_l, s), shiftop(dif_h, s)); \ - } - -#define dct_pass(shiftop, shift) \ - { \ - /* even part */ \ - int16x8_t sum26 = vaddq_s16(row2, row6); \ - dct_long_mul(p1e, sum26, rot0_0); \ - dct_long_mac(t2e, p1e, row6, rot0_1); \ - dct_long_mac(t3e, p1e, row2, rot0_2); \ - int16x8_t sum04 = vaddq_s16(row0, row4); \ - int16x8_t dif04 = vsubq_s16(row0, row4); \ - dct_widen(t0e, sum04); \ - dct_widen(t1e, dif04); \ - dct_wadd(x0, t0e, t3e); \ - dct_wsub(x3, t0e, t3e); \ - dct_wadd(x1, t1e, t2e); \ - dct_wsub(x2, t1e, t2e); \ - /* odd part */ \ - int16x8_t sum15 = vaddq_s16(row1, row5); \ - int16x8_t sum17 = vaddq_s16(row1, row7); \ - int16x8_t sum35 = vaddq_s16(row3, row5); \ - int16x8_t sum37 = vaddq_s16(row3, row7); \ - int16x8_t sumodd = vaddq_s16(sum17, sum35); \ - dct_long_mul(p5o, sumodd, rot1_0); \ - dct_long_mac(p1o, p5o, sum17, rot1_1); \ - dct_long_mac(p2o, p5o, sum35, rot1_2); \ - dct_long_mul(p3o, sum37, rot2_0); \ - dct_long_mul(p4o, sum15, rot2_1); \ - dct_wadd(sump13o, p1o, p3o); \ - dct_wadd(sump24o, p2o, p4o); \ - dct_wadd(sump23o, p2o, p3o); \ - dct_wadd(sump14o, p1o, p4o); \ - dct_long_mac(x4, sump13o, row7, rot3_0); \ - dct_long_mac(x5, sump24o, row5, rot3_1); \ - dct_long_mac(x6, sump23o, row3, rot3_2); \ - dct_long_mac(x7, sump14o, row1, rot3_3); \ - dct_bfly32o(row0,row7, x0,x7,shiftop,shift); \ - dct_bfly32o(row1,row6, x1,x6,shiftop,shift); \ - dct_bfly32o(row2,row5, x2,x5,shiftop,shift); \ - dct_bfly32o(row3,row4, x3,x4,shiftop,shift); \ - } - - // load - row0 = vld1q_s16(data + 0*8); - row1 = vld1q_s16(data + 1*8); - row2 = vld1q_s16(data + 2*8); - row3 = vld1q_s16(data + 3*8); - row4 = vld1q_s16(data + 4*8); - row5 = vld1q_s16(data + 5*8); - row6 = vld1q_s16(data + 6*8); - row7 = vld1q_s16(data + 7*8); - - // add DC bias - row0 = vaddq_s16(row0, vsetq_lane_s16(1024, vdupq_n_s16(0), 0)); - - // column pass - dct_pass(vrshrn_n_s32, 10); - - // 16bit 8x8 transpose - { -// these three map to a single VTRN.16, VTRN.32, and VSWP, respectively. -// whether compilers actually get this is another story, sadly. -#define dct_trn16(x, y) { int16x8x2_t t = vtrnq_s16(x, y); x = t.val[0]; y = t.val[1]; } -#define dct_trn32(x, y) { int32x4x2_t t = vtrnq_s32(vreinterpretq_s32_s16(x), vreinterpretq_s32_s16(y)); x = vreinterpretq_s16_s32(t.val[0]); y = vreinterpretq_s16_s32(t.val[1]); } -#define dct_trn64(x, y) { int16x8_t x0 = x; int16x8_t y0 = y; x = vcombine_s16(vget_low_s16(x0), vget_low_s16(y0)); y = vcombine_s16(vget_high_s16(x0), vget_high_s16(y0)); } - - // pass 1 - dct_trn16(row0, row1); // a0b0a2b2a4b4a6b6 - dct_trn16(row2, row3); - dct_trn16(row4, row5); - dct_trn16(row6, row7); - - // pass 2 - dct_trn32(row0, row2); // a0b0c0d0a4b4c4d4 - dct_trn32(row1, row3); - dct_trn32(row4, row6); - dct_trn32(row5, row7); - - // pass 3 - dct_trn64(row0, row4); // a0b0c0d0e0f0g0h0 - dct_trn64(row1, row5); - dct_trn64(row2, row6); - dct_trn64(row3, row7); - -#undef dct_trn16 -#undef dct_trn32 -#undef dct_trn64 - } - - // row pass - // vrshrn_n_s32 only supports shifts up to 16, we need - // 17. so do a non-rounding shift of 16 first then follow - // up with a rounding shift by 1. - dct_pass(vshrn_n_s32, 16); - - { - // pack and round - uint8x8_t p0 = vqrshrun_n_s16(row0, 1); - uint8x8_t p1 = vqrshrun_n_s16(row1, 1); - uint8x8_t p2 = vqrshrun_n_s16(row2, 1); - uint8x8_t p3 = vqrshrun_n_s16(row3, 1); - uint8x8_t p4 = vqrshrun_n_s16(row4, 1); - uint8x8_t p5 = vqrshrun_n_s16(row5, 1); - uint8x8_t p6 = vqrshrun_n_s16(row6, 1); - uint8x8_t p7 = vqrshrun_n_s16(row7, 1); - - // again, these can translate into one instruction, but often don't. -#define dct_trn8_8(x, y) { uint8x8x2_t t = vtrn_u8(x, y); x = t.val[0]; y = t.val[1]; } -#define dct_trn8_16(x, y) { uint16x4x2_t t = vtrn_u16(vreinterpret_u16_u8(x), vreinterpret_u16_u8(y)); x = vreinterpret_u8_u16(t.val[0]); y = vreinterpret_u8_u16(t.val[1]); } -#define dct_trn8_32(x, y) { uint32x2x2_t t = vtrn_u32(vreinterpret_u32_u8(x), vreinterpret_u32_u8(y)); x = vreinterpret_u8_u32(t.val[0]); y = vreinterpret_u8_u32(t.val[1]); } - - // sadly can't use interleaved stores here since we only write - // 8 bytes to each scan line! - - // 8x8 8-bit transpose pass 1 - dct_trn8_8(p0, p1); - dct_trn8_8(p2, p3); - dct_trn8_8(p4, p5); - dct_trn8_8(p6, p7); - - // pass 2 - dct_trn8_16(p0, p2); - dct_trn8_16(p1, p3); - dct_trn8_16(p4, p6); - dct_trn8_16(p5, p7); - - // pass 3 - dct_trn8_32(p0, p4); - dct_trn8_32(p1, p5); - dct_trn8_32(p2, p6); - dct_trn8_32(p3, p7); - - // store - vst1_u8(out, p0); out += out_stride; - vst1_u8(out, p1); out += out_stride; - vst1_u8(out, p2); out += out_stride; - vst1_u8(out, p3); out += out_stride; - vst1_u8(out, p4); out += out_stride; - vst1_u8(out, p5); out += out_stride; - vst1_u8(out, p6); out += out_stride; - vst1_u8(out, p7); - -#undef dct_trn8_8 -#undef dct_trn8_16 -#undef dct_trn8_32 - } - -#undef dct_long_mul -#undef dct_long_mac -#undef dct_widen -#undef dct_wadd -#undef dct_wsub -#undef dct_bfly32o -#undef dct_pass -} - -#endif // STBI_NEON - -#define STBI__MARKER_none 0xff -// if there's a pending marker from the entropy stream, return that -// otherwise, fetch from the stream and get a marker. if there's no -// marker, return 0xff, which is never a valid marker value -static stbi_uc stbi__get_marker(stbi__jpeg *j) -{ - stbi_uc x; - if (j->marker != STBI__MARKER_none) { x = j->marker; j->marker = STBI__MARKER_none; return x; } - x = stbi__get8(j->s); - if (x != 0xff) return STBI__MARKER_none; - while (x == 0xff) - x = stbi__get8(j->s); // consume repeated 0xff fill bytes - return x; -} - -// in each scan, we'll have scan_n components, and the order -// of the components is specified by order[] -#define STBI__RESTART(x) ((x) >= 0xd0 && (x) <= 0xd7) - -// after a restart interval, stbi__jpeg_reset the entropy decoder and -// the dc prediction -static void stbi__jpeg_reset(stbi__jpeg *j) -{ - j->code_bits = 0; - j->code_buffer = 0; - j->nomore = 0; - j->img_comp[0].dc_pred = j->img_comp[1].dc_pred = j->img_comp[2].dc_pred = j->img_comp[3].dc_pred = 0; - j->marker = STBI__MARKER_none; - j->todo = j->restart_interval ? j->restart_interval : 0x7fffffff; - j->eob_run = 0; - // no more than 1<<31 MCUs if no restart_interal? that's plenty safe, - // since we don't even allow 1<<30 pixels -} - -static int stbi__parse_entropy_coded_data(stbi__jpeg *z) -{ - stbi__jpeg_reset(z); - if (!z->progressive) { - if (z->scan_n == 1) { - int i,j; - STBI_SIMD_ALIGN(short, data[64]); - int n = z->order[0]; - // non-interleaved data, we just need to process one block at a time, - // in trivial scanline order - // number of blocks to do just depends on how many actual "pixels" this - // component has, independent of interleaved MCU blocking and such - int w = (z->img_comp[n].x+7) >> 3; - int h = (z->img_comp[n].y+7) >> 3; - for (j=0; j < h; ++j) { - for (i=0; i < w; ++i) { - int ha = z->img_comp[n].ha; - if (!stbi__jpeg_decode_block(z, data, z->huff_dc+z->img_comp[n].hd, z->huff_ac+ha, z->fast_ac[ha], n, z->dequant[z->img_comp[n].tq])) return 0; - z->idct_block_kernel(z->img_comp[n].data+z->img_comp[n].w2*j*8+i*8, z->img_comp[n].w2, data); - // every data block is an MCU, so countdown the restart interval - if (--z->todo <= 0) { - if (z->code_bits < 24) stbi__grow_buffer_unsafe(z); - // if it's NOT a restart, then just bail, so we get corrupt data - // rather than no data - if (!STBI__RESTART(z->marker)) return 1; - stbi__jpeg_reset(z); - } - } - } - return 1; - } else { // interleaved - int i,j,k,x,y; - STBI_SIMD_ALIGN(short, data[64]); - for (j=0; j < z->img_mcu_y; ++j) { - for (i=0; i < z->img_mcu_x; ++i) { - // scan an interleaved mcu... process scan_n components in order - for (k=0; k < z->scan_n; ++k) { - int n = z->order[k]; - // scan out an mcu's worth of this component; that's just determined - // by the basic H and V specified for the component - for (y=0; y < z->img_comp[n].v; ++y) { - for (x=0; x < z->img_comp[n].h; ++x) { - int x2 = (i*z->img_comp[n].h + x)*8; - int y2 = (j*z->img_comp[n].v + y)*8; - int ha = z->img_comp[n].ha; - if (!stbi__jpeg_decode_block(z, data, z->huff_dc+z->img_comp[n].hd, z->huff_ac+ha, z->fast_ac[ha], n, z->dequant[z->img_comp[n].tq])) return 0; - z->idct_block_kernel(z->img_comp[n].data+z->img_comp[n].w2*y2+x2, z->img_comp[n].w2, data); - } - } - } - // after all interleaved components, that's an interleaved MCU, - // so now count down the restart interval - if (--z->todo <= 0) { - if (z->code_bits < 24) stbi__grow_buffer_unsafe(z); - if (!STBI__RESTART(z->marker)) return 1; - stbi__jpeg_reset(z); - } - } - } - return 1; - } - } else { - if (z->scan_n == 1) { - int i,j; - int n = z->order[0]; - // non-interleaved data, we just need to process one block at a time, - // in trivial scanline order - // number of blocks to do just depends on how many actual "pixels" this - // component has, independent of interleaved MCU blocking and such - int w = (z->img_comp[n].x+7) >> 3; - int h = (z->img_comp[n].y+7) >> 3; - for (j=0; j < h; ++j) { - for (i=0; i < w; ++i) { - short *data = z->img_comp[n].coeff + 64 * (i + j * z->img_comp[n].coeff_w); - if (z->spec_start == 0) { - if (!stbi__jpeg_decode_block_prog_dc(z, data, &z->huff_dc[z->img_comp[n].hd], n)) - return 0; - } else { - int ha = z->img_comp[n].ha; - if (!stbi__jpeg_decode_block_prog_ac(z, data, &z->huff_ac[ha], z->fast_ac[ha])) - return 0; - } - // every data block is an MCU, so countdown the restart interval - if (--z->todo <= 0) { - if (z->code_bits < 24) stbi__grow_buffer_unsafe(z); - if (!STBI__RESTART(z->marker)) return 1; - stbi__jpeg_reset(z); - } - } - } - return 1; - } else { // interleaved - int i,j,k,x,y; - for (j=0; j < z->img_mcu_y; ++j) { - for (i=0; i < z->img_mcu_x; ++i) { - // scan an interleaved mcu... process scan_n components in order - for (k=0; k < z->scan_n; ++k) { - int n = z->order[k]; - // scan out an mcu's worth of this component; that's just determined - // by the basic H and V specified for the component - for (y=0; y < z->img_comp[n].v; ++y) { - for (x=0; x < z->img_comp[n].h; ++x) { - int x2 = (i*z->img_comp[n].h + x); - int y2 = (j*z->img_comp[n].v + y); - short *data = z->img_comp[n].coeff + 64 * (x2 + y2 * z->img_comp[n].coeff_w); - if (!stbi__jpeg_decode_block_prog_dc(z, data, &z->huff_dc[z->img_comp[n].hd], n)) - return 0; - } - } - } - // after all interleaved components, that's an interleaved MCU, - // so now count down the restart interval - if (--z->todo <= 0) { - if (z->code_bits < 24) stbi__grow_buffer_unsafe(z); - if (!STBI__RESTART(z->marker)) return 1; - stbi__jpeg_reset(z); - } - } - } - return 1; - } - } -} - -static void stbi__jpeg_dequantize(short *data, stbi__uint16 *dequant) -{ - int i; - for (i=0; i < 64; ++i) - data[i] *= dequant[i]; -} - -static void stbi__jpeg_finish(stbi__jpeg *z) -{ - if (z->progressive) { - // dequantize and idct the data - int i,j,n; - for (n=0; n < z->s->img_n; ++n) { - int w = (z->img_comp[n].x+7) >> 3; - int h = (z->img_comp[n].y+7) >> 3; - for (j=0; j < h; ++j) { - for (i=0; i < w; ++i) { - short *data = z->img_comp[n].coeff + 64 * (i + j * z->img_comp[n].coeff_w); - stbi__jpeg_dequantize(data, z->dequant[z->img_comp[n].tq]); - z->idct_block_kernel(z->img_comp[n].data+z->img_comp[n].w2*j*8+i*8, z->img_comp[n].w2, data); - } - } - } - } -} - -static int stbi__process_marker(stbi__jpeg *z, int m) -{ - int L; - switch (m) { - case STBI__MARKER_none: // no marker found - return stbi__err("expected marker","Corrupt JPEG"); - - case 0xDD: // DRI - specify restart interval - if (stbi__get16be(z->s) != 4) return stbi__err("bad DRI len","Corrupt JPEG"); - z->restart_interval = stbi__get16be(z->s); - return 1; - - case 0xDB: // DQT - define quantization table - L = stbi__get16be(z->s)-2; - while (L > 0) { - int q = stbi__get8(z->s); - int p = q >> 4, sixteen = (p != 0); - int t = q & 15,i; - if (p != 0 && p != 1) return stbi__err("bad DQT type","Corrupt JPEG"); - if (t > 3) return stbi__err("bad DQT table","Corrupt JPEG"); - - for (i=0; i < 64; ++i) - z->dequant[t][stbi__jpeg_dezigzag[i]] = (stbi__uint16)(sixteen ? stbi__get16be(z->s) : stbi__get8(z->s)); - L -= (sixteen ? 129 : 65); - } - return L==0; - - case 0xC4: // DHT - define huffman table - L = stbi__get16be(z->s)-2; - while (L > 0) { - stbi_uc *v; - int sizes[16],i,n=0; - int q = stbi__get8(z->s); - int tc = q >> 4; - int th = q & 15; - if (tc > 1 || th > 3) return stbi__err("bad DHT header","Corrupt JPEG"); - for (i=0; i < 16; ++i) { - sizes[i] = stbi__get8(z->s); - n += sizes[i]; - } - if(n > 256) return stbi__err("bad DHT header","Corrupt JPEG"); // Loop over i < n would write past end of values! - L -= 17; - if (tc == 0) { - if (!stbi__build_huffman(z->huff_dc+th, sizes)) return 0; - v = z->huff_dc[th].values; - } else { - if (!stbi__build_huffman(z->huff_ac+th, sizes)) return 0; - v = z->huff_ac[th].values; - } - for (i=0; i < n; ++i) - v[i] = stbi__get8(z->s); - if (tc != 0) - stbi__build_fast_ac(z->fast_ac[th], z->huff_ac + th); - L -= n; - } - return L==0; - } - - // check for comment block or APP blocks - if ((m >= 0xE0 && m <= 0xEF) || m == 0xFE) { - L = stbi__get16be(z->s); - if (L < 2) { - if (m == 0xFE) - return stbi__err("bad COM len","Corrupt JPEG"); - else - return stbi__err("bad APP len","Corrupt JPEG"); - } - L -= 2; - - if (m == 0xE0 && L >= 5) { // JFIF APP0 segment - static const unsigned char tag[5] = {'J','F','I','F','\0'}; - int ok = 1; - int i; - for (i=0; i < 5; ++i) - if (stbi__get8(z->s) != tag[i]) - ok = 0; - L -= 5; - if (ok) - z->jfif = 1; - } else if (m == 0xEE && L >= 12) { // Adobe APP14 segment - static const unsigned char tag[6] = {'A','d','o','b','e','\0'}; - int ok = 1; - int i; - for (i=0; i < 6; ++i) - if (stbi__get8(z->s) != tag[i]) - ok = 0; - L -= 6; - if (ok) { - stbi__get8(z->s); // version - stbi__get16be(z->s); // flags0 - stbi__get16be(z->s); // flags1 - z->app14_color_transform = stbi__get8(z->s); // color transform - L -= 6; - } - } - - stbi__skip(z->s, L); - return 1; - } - - return stbi__err("unknown marker","Corrupt JPEG"); -} - -// after we see SOS -static int stbi__process_scan_header(stbi__jpeg *z) -{ - int i; - int Ls = stbi__get16be(z->s); - z->scan_n = stbi__get8(z->s); - if (z->scan_n < 1 || z->scan_n > 4 || z->scan_n > (int) z->s->img_n) return stbi__err("bad SOS component count","Corrupt JPEG"); - if (Ls != 6+2*z->scan_n) return stbi__err("bad SOS len","Corrupt JPEG"); - for (i=0; i < z->scan_n; ++i) { - int id = stbi__get8(z->s), which; - int q = stbi__get8(z->s); - for (which = 0; which < z->s->img_n; ++which) - if (z->img_comp[which].id == id) - break; - if (which == z->s->img_n) return 0; // no match - z->img_comp[which].hd = q >> 4; if (z->img_comp[which].hd > 3) return stbi__err("bad DC huff","Corrupt JPEG"); - z->img_comp[which].ha = q & 15; if (z->img_comp[which].ha > 3) return stbi__err("bad AC huff","Corrupt JPEG"); - z->order[i] = which; - } - - { - int aa; - z->spec_start = stbi__get8(z->s); - z->spec_end = stbi__get8(z->s); // should be 63, but might be 0 - aa = stbi__get8(z->s); - z->succ_high = (aa >> 4); - z->succ_low = (aa & 15); - if (z->progressive) { - if (z->spec_start > 63 || z->spec_end > 63 || z->spec_start > z->spec_end || z->succ_high > 13 || z->succ_low > 13) - return stbi__err("bad SOS", "Corrupt JPEG"); - } else { - if (z->spec_start != 0) return stbi__err("bad SOS","Corrupt JPEG"); - if (z->succ_high != 0 || z->succ_low != 0) return stbi__err("bad SOS","Corrupt JPEG"); - z->spec_end = 63; - } - } - - return 1; -} - -static int stbi__free_jpeg_components(stbi__jpeg *z, int ncomp, int why) -{ - int i; - for (i=0; i < ncomp; ++i) { - if (z->img_comp[i].raw_data) { - STBI_FREE(z->img_comp[i].raw_data); - z->img_comp[i].raw_data = NULL; - z->img_comp[i].data = NULL; - } - if (z->img_comp[i].raw_coeff) { - STBI_FREE(z->img_comp[i].raw_coeff); - z->img_comp[i].raw_coeff = 0; - z->img_comp[i].coeff = 0; - } - if (z->img_comp[i].linebuf) { - STBI_FREE(z->img_comp[i].linebuf); - z->img_comp[i].linebuf = NULL; - } - } - return why; -} - -static int stbi__process_frame_header(stbi__jpeg *z, int scan) -{ - stbi__context *s = z->s; - int Lf,p,i,q, h_max=1,v_max=1,c; - Lf = stbi__get16be(s); if (Lf < 11) return stbi__err("bad SOF len","Corrupt JPEG"); // JPEG - p = stbi__get8(s); if (p != 8) return stbi__err("only 8-bit","JPEG format not supported: 8-bit only"); // JPEG baseline - s->img_y = stbi__get16be(s); if (s->img_y == 0) return stbi__err("no header height", "JPEG format not supported: delayed height"); // Legal, but we don't handle it--but neither does IJG - s->img_x = stbi__get16be(s); if (s->img_x == 0) return stbi__err("0 width","Corrupt JPEG"); // JPEG requires - if (s->img_y > STBI_MAX_DIMENSIONS) return stbi__err("too large","Very large image (corrupt?)"); - if (s->img_x > STBI_MAX_DIMENSIONS) return stbi__err("too large","Very large image (corrupt?)"); - c = stbi__get8(s); - if (c != 3 && c != 1 && c != 4) return stbi__err("bad component count","Corrupt JPEG"); - s->img_n = c; - for (i=0; i < c; ++i) { - z->img_comp[i].data = NULL; - z->img_comp[i].linebuf = NULL; - } - - if (Lf != 8+3*s->img_n) return stbi__err("bad SOF len","Corrupt JPEG"); - - z->rgb = 0; - for (i=0; i < s->img_n; ++i) { - static const unsigned char rgb[3] = { 'R', 'G', 'B' }; - z->img_comp[i].id = stbi__get8(s); - if (s->img_n == 3 && z->img_comp[i].id == rgb[i]) - ++z->rgb; - q = stbi__get8(s); - z->img_comp[i].h = (q >> 4); if (!z->img_comp[i].h || z->img_comp[i].h > 4) return stbi__err("bad H","Corrupt JPEG"); - z->img_comp[i].v = q & 15; if (!z->img_comp[i].v || z->img_comp[i].v > 4) return stbi__err("bad V","Corrupt JPEG"); - z->img_comp[i].tq = stbi__get8(s); if (z->img_comp[i].tq > 3) return stbi__err("bad TQ","Corrupt JPEG"); - } - - if (scan != STBI__SCAN_load) return 1; - - if (!stbi__mad3sizes_valid(s->img_x, s->img_y, s->img_n, 0)) return stbi__err("too large", "Image too large to decode"); - - for (i=0; i < s->img_n; ++i) { - if (z->img_comp[i].h > h_max) h_max = z->img_comp[i].h; - if (z->img_comp[i].v > v_max) v_max = z->img_comp[i].v; - } - - // check that plane subsampling factors are integer ratios; our resamplers can't deal with fractional ratios - // and I've never seen a non-corrupted JPEG file actually use them - for (i=0; i < s->img_n; ++i) { - if (h_max % z->img_comp[i].h != 0) return stbi__err("bad H","Corrupt JPEG"); - if (v_max % z->img_comp[i].v != 0) return stbi__err("bad V","Corrupt JPEG"); - } - - // compute interleaved mcu info - z->img_h_max = h_max; - z->img_v_max = v_max; - z->img_mcu_w = h_max * 8; - z->img_mcu_h = v_max * 8; - // these sizes can't be more than 17 bits - z->img_mcu_x = (s->img_x + z->img_mcu_w-1) / z->img_mcu_w; - z->img_mcu_y = (s->img_y + z->img_mcu_h-1) / z->img_mcu_h; - - for (i=0; i < s->img_n; ++i) { - // number of effective pixels (e.g. for non-interleaved MCU) - z->img_comp[i].x = (s->img_x * z->img_comp[i].h + h_max-1) / h_max; - z->img_comp[i].y = (s->img_y * z->img_comp[i].v + v_max-1) / v_max; - // to simplify generation, we'll allocate enough memory to decode - // the bogus oversized data from using interleaved MCUs and their - // big blocks (e.g. a 16x16 iMCU on an image of width 33); we won't - // discard the extra data until colorspace conversion - // - // img_mcu_x, img_mcu_y: <=17 bits; comp[i].h and .v are <=4 (checked earlier) - // so these muls can't overflow with 32-bit ints (which we require) - z->img_comp[i].w2 = z->img_mcu_x * z->img_comp[i].h * 8; - z->img_comp[i].h2 = z->img_mcu_y * z->img_comp[i].v * 8; - z->img_comp[i].coeff = 0; - z->img_comp[i].raw_coeff = 0; - z->img_comp[i].linebuf = NULL; - z->img_comp[i].raw_data = stbi__malloc_mad2(z->img_comp[i].w2, z->img_comp[i].h2, 15); - if (z->img_comp[i].raw_data == NULL) - return stbi__free_jpeg_components(z, i+1, stbi__err("outofmem", "Out of memory")); - // align blocks for idct using mmx/sse - z->img_comp[i].data = (stbi_uc*) (((size_t) z->img_comp[i].raw_data + 15) & ~15); - if (z->progressive) { - // w2, h2 are multiples of 8 (see above) - z->img_comp[i].coeff_w = z->img_comp[i].w2 / 8; - z->img_comp[i].coeff_h = z->img_comp[i].h2 / 8; - z->img_comp[i].raw_coeff = stbi__malloc_mad3(z->img_comp[i].w2, z->img_comp[i].h2, sizeof(short), 15); - if (z->img_comp[i].raw_coeff == NULL) - return stbi__free_jpeg_components(z, i+1, stbi__err("outofmem", "Out of memory")); - z->img_comp[i].coeff = (short*) (((size_t) z->img_comp[i].raw_coeff + 15) & ~15); - } - } - - return 1; -} - -// use comparisons since in some cases we handle more than one case (e.g. SOF) -#define stbi__DNL(x) ((x) == 0xdc) -#define stbi__SOI(x) ((x) == 0xd8) -#define stbi__EOI(x) ((x) == 0xd9) -#define stbi__SOF(x) ((x) == 0xc0 || (x) == 0xc1 || (x) == 0xc2) -#define stbi__SOS(x) ((x) == 0xda) - -#define stbi__SOF_progressive(x) ((x) == 0xc2) - -static int stbi__decode_jpeg_header(stbi__jpeg *z, int scan) -{ - int m; - z->jfif = 0; - z->app14_color_transform = -1; // valid values are 0,1,2 - z->marker = STBI__MARKER_none; // initialize cached marker to empty - m = stbi__get_marker(z); - if (!stbi__SOI(m)) return stbi__err("no SOI","Corrupt JPEG"); - if (scan == STBI__SCAN_type) return 1; - m = stbi__get_marker(z); - while (!stbi__SOF(m)) { - if (!stbi__process_marker(z,m)) return 0; - m = stbi__get_marker(z); - while (m == STBI__MARKER_none) { - // some files have extra padding after their blocks, so ok, we'll scan - if (stbi__at_eof(z->s)) return stbi__err("no SOF", "Corrupt JPEG"); - m = stbi__get_marker(z); - } - } - z->progressive = stbi__SOF_progressive(m); - if (!stbi__process_frame_header(z, scan)) return 0; - return 1; -} - -static stbi_uc stbi__skip_jpeg_junk_at_end(stbi__jpeg *j) -{ - // some JPEGs have junk at end, skip over it but if we find what looks - // like a valid marker, resume there - while (!stbi__at_eof(j->s)) { - stbi_uc x = stbi__get8(j->s); - while (x == 0xff) { // might be a marker - if (stbi__at_eof(j->s)) return STBI__MARKER_none; - x = stbi__get8(j->s); - if (x != 0x00 && x != 0xff) { - // not a stuffed zero or lead-in to another marker, looks - // like an actual marker, return it - return x; - } - // stuffed zero has x=0 now which ends the loop, meaning we go - // back to regular scan loop. - // repeated 0xff keeps trying to read the next byte of the marker. - } - } - return STBI__MARKER_none; -} - -// decode image to YCbCr format -static int stbi__decode_jpeg_image(stbi__jpeg *j) -{ - int m; - for (m = 0; m < 4; m++) { - j->img_comp[m].raw_data = NULL; - j->img_comp[m].raw_coeff = NULL; - } - j->restart_interval = 0; - if (!stbi__decode_jpeg_header(j, STBI__SCAN_load)) return 0; - m = stbi__get_marker(j); - while (!stbi__EOI(m)) { - if (stbi__SOS(m)) { - if (!stbi__process_scan_header(j)) return 0; - if (!stbi__parse_entropy_coded_data(j)) return 0; - if (j->marker == STBI__MARKER_none ) { - j->marker = stbi__skip_jpeg_junk_at_end(j); - // if we reach eof without hitting a marker, stbi__get_marker() below will fail and we'll eventually return 0 - } - m = stbi__get_marker(j); - if (STBI__RESTART(m)) - m = stbi__get_marker(j); - } else if (stbi__DNL(m)) { - int Ld = stbi__get16be(j->s); - stbi__uint32 NL = stbi__get16be(j->s); - if (Ld != 4) return stbi__err("bad DNL len", "Corrupt JPEG"); - if (NL != j->s->img_y) return stbi__err("bad DNL height", "Corrupt JPEG"); - m = stbi__get_marker(j); - } else { - if (!stbi__process_marker(j, m)) return 1; - m = stbi__get_marker(j); - } - } - if (j->progressive) - stbi__jpeg_finish(j); - return 1; -} - -// static jfif-centered resampling (across block boundaries) - -typedef stbi_uc *(*resample_row_func)(stbi_uc *out, stbi_uc *in0, stbi_uc *in1, - int w, int hs); - -#define stbi__div4(x) ((stbi_uc) ((x) >> 2)) - -static stbi_uc *resample_row_1(stbi_uc *out, stbi_uc *in_near, stbi_uc *in_far, int w, int hs) -{ - STBI_NOTUSED(out); - STBI_NOTUSED(in_far); - STBI_NOTUSED(w); - STBI_NOTUSED(hs); - return in_near; -} - -static stbi_uc* stbi__resample_row_v_2(stbi_uc *out, stbi_uc *in_near, stbi_uc *in_far, int w, int hs) -{ - // need to generate two samples vertically for every one in input - int i; - STBI_NOTUSED(hs); - for (i=0; i < w; ++i) - out[i] = stbi__div4(3*in_near[i] + in_far[i] + 2); - return out; -} - -static stbi_uc* stbi__resample_row_h_2(stbi_uc *out, stbi_uc *in_near, stbi_uc *in_far, int w, int hs) -{ - // need to generate two samples horizontally for every one in input - int i; - stbi_uc *input = in_near; - - if (w == 1) { - // if only one sample, can't do any interpolation - out[0] = out[1] = input[0]; - return out; - } - - out[0] = input[0]; - out[1] = stbi__div4(input[0]*3 + input[1] + 2); - for (i=1; i < w-1; ++i) { - int n = 3*input[i]+2; - out[i*2+0] = stbi__div4(n+input[i-1]); - out[i*2+1] = stbi__div4(n+input[i+1]); - } - out[i*2+0] = stbi__div4(input[w-2]*3 + input[w-1] + 2); - out[i*2+1] = input[w-1]; - - STBI_NOTUSED(in_far); - STBI_NOTUSED(hs); - - return out; -} - -#define stbi__div16(x) ((stbi_uc) ((x) >> 4)) - -static stbi_uc *stbi__resample_row_hv_2(stbi_uc *out, stbi_uc *in_near, stbi_uc *in_far, int w, int hs) -{ - // need to generate 2x2 samples for every one in input - int i,t0,t1; - if (w == 1) { - out[0] = out[1] = stbi__div4(3*in_near[0] + in_far[0] + 2); - return out; - } - - t1 = 3*in_near[0] + in_far[0]; - out[0] = stbi__div4(t1+2); - for (i=1; i < w; ++i) { - t0 = t1; - t1 = 3*in_near[i]+in_far[i]; - out[i*2-1] = stbi__div16(3*t0 + t1 + 8); - out[i*2 ] = stbi__div16(3*t1 + t0 + 8); - } - out[w*2-1] = stbi__div4(t1+2); - - STBI_NOTUSED(hs); - - return out; -} - -#if defined(STBI_SSE2) || defined(STBI_NEON) -static stbi_uc *stbi__resample_row_hv_2_simd(stbi_uc *out, stbi_uc *in_near, stbi_uc *in_far, int w, int hs) -{ - // need to generate 2x2 samples for every one in input - int i=0,t0,t1; - - if (w == 1) { - out[0] = out[1] = stbi__div4(3*in_near[0] + in_far[0] + 2); - return out; - } - - t1 = 3*in_near[0] + in_far[0]; - // process groups of 8 pixels for as long as we can. - // note we can't handle the last pixel in a row in this loop - // because we need to handle the filter boundary conditions. - for (; i < ((w-1) & ~7); i += 8) { -#if defined(STBI_SSE2) - // load and perform the vertical filtering pass - // this uses 3*x + y = 4*x + (y - x) - __m128i zero = _mm_setzero_si128(); - __m128i farb = _mm_loadl_epi64((__m128i *) (in_far + i)); - __m128i nearb = _mm_loadl_epi64((__m128i *) (in_near + i)); - __m128i farw = _mm_unpacklo_epi8(farb, zero); - __m128i nearw = _mm_unpacklo_epi8(nearb, zero); - __m128i diff = _mm_sub_epi16(farw, nearw); - __m128i nears = _mm_slli_epi16(nearw, 2); - __m128i curr = _mm_add_epi16(nears, diff); // current row - - // horizontal filter works the same based on shifted vers of current - // row. "prev" is current row shifted right by 1 pixel; we need to - // insert the previous pixel value (from t1). - // "next" is current row shifted left by 1 pixel, with first pixel - // of next block of 8 pixels added in. - __m128i prv0 = _mm_slli_si128(curr, 2); - __m128i nxt0 = _mm_srli_si128(curr, 2); - __m128i prev = _mm_insert_epi16(prv0, t1, 0); - __m128i next = _mm_insert_epi16(nxt0, 3*in_near[i+8] + in_far[i+8], 7); - - // horizontal filter, polyphase implementation since it's convenient: - // even pixels = 3*cur + prev = cur*4 + (prev - cur) - // odd pixels = 3*cur + next = cur*4 + (next - cur) - // note the shared term. - __m128i bias = _mm_set1_epi16(8); - __m128i curs = _mm_slli_epi16(curr, 2); - __m128i prvd = _mm_sub_epi16(prev, curr); - __m128i nxtd = _mm_sub_epi16(next, curr); - __m128i curb = _mm_add_epi16(curs, bias); - __m128i even = _mm_add_epi16(prvd, curb); - __m128i odd = _mm_add_epi16(nxtd, curb); - - // interleave even and odd pixels, then undo scaling. - __m128i int0 = _mm_unpacklo_epi16(even, odd); - __m128i int1 = _mm_unpackhi_epi16(even, odd); - __m128i de0 = _mm_srli_epi16(int0, 4); - __m128i de1 = _mm_srli_epi16(int1, 4); - - // pack and write output - __m128i outv = _mm_packus_epi16(de0, de1); - _mm_storeu_si128((__m128i *) (out + i*2), outv); -#elif defined(STBI_NEON) - // load and perform the vertical filtering pass - // this uses 3*x + y = 4*x + (y - x) - uint8x8_t farb = vld1_u8(in_far + i); - uint8x8_t nearb = vld1_u8(in_near + i); - int16x8_t diff = vreinterpretq_s16_u16(vsubl_u8(farb, nearb)); - int16x8_t nears = vreinterpretq_s16_u16(vshll_n_u8(nearb, 2)); - int16x8_t curr = vaddq_s16(nears, diff); // current row - - // horizontal filter works the same based on shifted vers of current - // row. "prev" is current row shifted right by 1 pixel; we need to - // insert the previous pixel value (from t1). - // "next" is current row shifted left by 1 pixel, with first pixel - // of next block of 8 pixels added in. - int16x8_t prv0 = vextq_s16(curr, curr, 7); - int16x8_t nxt0 = vextq_s16(curr, curr, 1); - int16x8_t prev = vsetq_lane_s16(t1, prv0, 0); - int16x8_t next = vsetq_lane_s16(3*in_near[i+8] + in_far[i+8], nxt0, 7); - - // horizontal filter, polyphase implementation since it's convenient: - // even pixels = 3*cur + prev = cur*4 + (prev - cur) - // odd pixels = 3*cur + next = cur*4 + (next - cur) - // note the shared term. - int16x8_t curs = vshlq_n_s16(curr, 2); - int16x8_t prvd = vsubq_s16(prev, curr); - int16x8_t nxtd = vsubq_s16(next, curr); - int16x8_t even = vaddq_s16(curs, prvd); - int16x8_t odd = vaddq_s16(curs, nxtd); - - // undo scaling and round, then store with even/odd phases interleaved - uint8x8x2_t o; - o.val[0] = vqrshrun_n_s16(even, 4); - o.val[1] = vqrshrun_n_s16(odd, 4); - vst2_u8(out + i*2, o); -#endif - - // "previous" value for next iter - t1 = 3*in_near[i+7] + in_far[i+7]; - } - - t0 = t1; - t1 = 3*in_near[i] + in_far[i]; - out[i*2] = stbi__div16(3*t1 + t0 + 8); - - for (++i; i < w; ++i) { - t0 = t1; - t1 = 3*in_near[i]+in_far[i]; - out[i*2-1] = stbi__div16(3*t0 + t1 + 8); - out[i*2 ] = stbi__div16(3*t1 + t0 + 8); - } - out[w*2-1] = stbi__div4(t1+2); - - STBI_NOTUSED(hs); - - return out; -} -#endif - -static stbi_uc *stbi__resample_row_generic(stbi_uc *out, stbi_uc *in_near, stbi_uc *in_far, int w, int hs) -{ - // resample with nearest-neighbor - int i,j; - STBI_NOTUSED(in_far); - for (i=0; i < w; ++i) - for (j=0; j < hs; ++j) - out[i*hs+j] = in_near[i]; - return out; -} - -// this is a reduced-precision calculation of YCbCr-to-RGB introduced -// to make sure the code produces the same results in both SIMD and scalar -#define stbi__float2fixed(x) (((int) ((x) * 4096.0f + 0.5f)) << 8) -static void stbi__YCbCr_to_RGB_row(stbi_uc *out, const stbi_uc *y, const stbi_uc *pcb, const stbi_uc *pcr, int count, int step) -{ - int i; - for (i=0; i < count; ++i) { - int y_fixed = (y[i] << 20) + (1<<19); // rounding - int r,g,b; - int cr = pcr[i] - 128; - int cb = pcb[i] - 128; - r = y_fixed + cr* stbi__float2fixed(1.40200f); - g = y_fixed + (cr*-stbi__float2fixed(0.71414f)) + ((cb*-stbi__float2fixed(0.34414f)) & 0xffff0000); - b = y_fixed + cb* stbi__float2fixed(1.77200f); - r >>= 20; - g >>= 20; - b >>= 20; - if ((unsigned) r > 255) { if (r < 0) r = 0; else r = 255; } - if ((unsigned) g > 255) { if (g < 0) g = 0; else g = 255; } - if ((unsigned) b > 255) { if (b < 0) b = 0; else b = 255; } - out[0] = (stbi_uc)r; - out[1] = (stbi_uc)g; - out[2] = (stbi_uc)b; - out[3] = 255; - out += step; - } -} - -#if defined(STBI_SSE2) || defined(STBI_NEON) -static void stbi__YCbCr_to_RGB_simd(stbi_uc *out, stbi_uc const *y, stbi_uc const *pcb, stbi_uc const *pcr, int count, int step) -{ - int i = 0; - -#ifdef STBI_SSE2 - // step == 3 is pretty ugly on the final interleave, and i'm not convinced - // it's useful in practice (you wouldn't use it for textures, for example). - // so just accelerate step == 4 case. - if (step == 4) { - // this is a fairly straightforward implementation and not super-optimized. - __m128i signflip = _mm_set1_epi8(-0x80); - __m128i cr_const0 = _mm_set1_epi16( (short) ( 1.40200f*4096.0f+0.5f)); - __m128i cr_const1 = _mm_set1_epi16( - (short) ( 0.71414f*4096.0f+0.5f)); - __m128i cb_const0 = _mm_set1_epi16( - (short) ( 0.34414f*4096.0f+0.5f)); - __m128i cb_const1 = _mm_set1_epi16( (short) ( 1.77200f*4096.0f+0.5f)); - __m128i y_bias = _mm_set1_epi8((char) (unsigned char) 128); - __m128i xw = _mm_set1_epi16(255); // alpha channel - - for (; i+7 < count; i += 8) { - // load - __m128i y_bytes = _mm_loadl_epi64((__m128i *) (y+i)); - __m128i cr_bytes = _mm_loadl_epi64((__m128i *) (pcr+i)); - __m128i cb_bytes = _mm_loadl_epi64((__m128i *) (pcb+i)); - __m128i cr_biased = _mm_xor_si128(cr_bytes, signflip); // -128 - __m128i cb_biased = _mm_xor_si128(cb_bytes, signflip); // -128 - - // unpack to short (and left-shift cr, cb by 8) - __m128i yw = _mm_unpacklo_epi8(y_bias, y_bytes); - __m128i crw = _mm_unpacklo_epi8(_mm_setzero_si128(), cr_biased); - __m128i cbw = _mm_unpacklo_epi8(_mm_setzero_si128(), cb_biased); - - // color transform - __m128i yws = _mm_srli_epi16(yw, 4); - __m128i cr0 = _mm_mulhi_epi16(cr_const0, crw); - __m128i cb0 = _mm_mulhi_epi16(cb_const0, cbw); - __m128i cb1 = _mm_mulhi_epi16(cbw, cb_const1); - __m128i cr1 = _mm_mulhi_epi16(crw, cr_const1); - __m128i rws = _mm_add_epi16(cr0, yws); - __m128i gwt = _mm_add_epi16(cb0, yws); - __m128i bws = _mm_add_epi16(yws, cb1); - __m128i gws = _mm_add_epi16(gwt, cr1); - - // descale - __m128i rw = _mm_srai_epi16(rws, 4); - __m128i bw = _mm_srai_epi16(bws, 4); - __m128i gw = _mm_srai_epi16(gws, 4); - - // back to byte, set up for transpose - __m128i brb = _mm_packus_epi16(rw, bw); - __m128i gxb = _mm_packus_epi16(gw, xw); - - // transpose to interleave channels - __m128i t0 = _mm_unpacklo_epi8(brb, gxb); - __m128i t1 = _mm_unpackhi_epi8(brb, gxb); - __m128i o0 = _mm_unpacklo_epi16(t0, t1); - __m128i o1 = _mm_unpackhi_epi16(t0, t1); - - // store - _mm_storeu_si128((__m128i *) (out + 0), o0); - _mm_storeu_si128((__m128i *) (out + 16), o1); - out += 32; - } - } -#endif - -#ifdef STBI_NEON - // in this version, step=3 support would be easy to add. but is there demand? - if (step == 4) { - // this is a fairly straightforward implementation and not super-optimized. - uint8x8_t signflip = vdup_n_u8(0x80); - int16x8_t cr_const0 = vdupq_n_s16( (short) ( 1.40200f*4096.0f+0.5f)); - int16x8_t cr_const1 = vdupq_n_s16( - (short) ( 0.71414f*4096.0f+0.5f)); - int16x8_t cb_const0 = vdupq_n_s16( - (short) ( 0.34414f*4096.0f+0.5f)); - int16x8_t cb_const1 = vdupq_n_s16( (short) ( 1.77200f*4096.0f+0.5f)); - - for (; i+7 < count; i += 8) { - // load - uint8x8_t y_bytes = vld1_u8(y + i); - uint8x8_t cr_bytes = vld1_u8(pcr + i); - uint8x8_t cb_bytes = vld1_u8(pcb + i); - int8x8_t cr_biased = vreinterpret_s8_u8(vsub_u8(cr_bytes, signflip)); - int8x8_t cb_biased = vreinterpret_s8_u8(vsub_u8(cb_bytes, signflip)); - - // expand to s16 - int16x8_t yws = vreinterpretq_s16_u16(vshll_n_u8(y_bytes, 4)); - int16x8_t crw = vshll_n_s8(cr_biased, 7); - int16x8_t cbw = vshll_n_s8(cb_biased, 7); - - // color transform - int16x8_t cr0 = vqdmulhq_s16(crw, cr_const0); - int16x8_t cb0 = vqdmulhq_s16(cbw, cb_const0); - int16x8_t cr1 = vqdmulhq_s16(crw, cr_const1); - int16x8_t cb1 = vqdmulhq_s16(cbw, cb_const1); - int16x8_t rws = vaddq_s16(yws, cr0); - int16x8_t gws = vaddq_s16(vaddq_s16(yws, cb0), cr1); - int16x8_t bws = vaddq_s16(yws, cb1); - - // undo scaling, round, convert to byte - uint8x8x4_t o; - o.val[0] = vqrshrun_n_s16(rws, 4); - o.val[1] = vqrshrun_n_s16(gws, 4); - o.val[2] = vqrshrun_n_s16(bws, 4); - o.val[3] = vdup_n_u8(255); - - // store, interleaving r/g/b/a - vst4_u8(out, o); - out += 8*4; - } - } -#endif - - for (; i < count; ++i) { - int y_fixed = (y[i] << 20) + (1<<19); // rounding - int r,g,b; - int cr = pcr[i] - 128; - int cb = pcb[i] - 128; - r = y_fixed + cr* stbi__float2fixed(1.40200f); - g = y_fixed + cr*-stbi__float2fixed(0.71414f) + ((cb*-stbi__float2fixed(0.34414f)) & 0xffff0000); - b = y_fixed + cb* stbi__float2fixed(1.77200f); - r >>= 20; - g >>= 20; - b >>= 20; - if ((unsigned) r > 255) { if (r < 0) r = 0; else r = 255; } - if ((unsigned) g > 255) { if (g < 0) g = 0; else g = 255; } - if ((unsigned) b > 255) { if (b < 0) b = 0; else b = 255; } - out[0] = (stbi_uc)r; - out[1] = (stbi_uc)g; - out[2] = (stbi_uc)b; - out[3] = 255; - out += step; - } -} -#endif - -// set up the kernels -static void stbi__setup_jpeg(stbi__jpeg *j) -{ - j->idct_block_kernel = stbi__idct_block; - j->YCbCr_to_RGB_kernel = stbi__YCbCr_to_RGB_row; - j->resample_row_hv_2_kernel = stbi__resample_row_hv_2; - -#ifdef STBI_SSE2 - if (stbi__sse2_available()) { - j->idct_block_kernel = stbi__idct_simd; - j->YCbCr_to_RGB_kernel = stbi__YCbCr_to_RGB_simd; - j->resample_row_hv_2_kernel = stbi__resample_row_hv_2_simd; - } -#endif - -#ifdef STBI_NEON - j->idct_block_kernel = stbi__idct_simd; - j->YCbCr_to_RGB_kernel = stbi__YCbCr_to_RGB_simd; - j->resample_row_hv_2_kernel = stbi__resample_row_hv_2_simd; -#endif -} - -// clean up the temporary component buffers -static void stbi__cleanup_jpeg(stbi__jpeg *j) -{ - stbi__free_jpeg_components(j, j->s->img_n, 0); -} - -typedef struct -{ - resample_row_func resample; - stbi_uc *line0,*line1; - int hs,vs; // expansion factor in each axis - int w_lores; // horizontal pixels pre-expansion - int ystep; // how far through vertical expansion we are - int ypos; // which pre-expansion row we're on -} stbi__resample; - -// fast 0..255 * 0..255 => 0..255 rounded multiplication -static stbi_uc stbi__blinn_8x8(stbi_uc x, stbi_uc y) -{ - unsigned int t = x*y + 128; - return (stbi_uc) ((t + (t >>8)) >> 8); -} - -static stbi_uc *load_jpeg_image(stbi__jpeg *z, int *out_x, int *out_y, int *comp, int req_comp) -{ - int n, decode_n, is_rgb; - z->s->img_n = 0; // make stbi__cleanup_jpeg safe - - // validate req_comp - if (req_comp < 0 || req_comp > 4) return stbi__errpuc("bad req_comp", "Internal error"); - - // load a jpeg image from whichever source, but leave in YCbCr format - if (!stbi__decode_jpeg_image(z)) { stbi__cleanup_jpeg(z); return NULL; } - - // determine actual number of components to generate - n = req_comp ? req_comp : z->s->img_n >= 3 ? 3 : 1; - - is_rgb = z->s->img_n == 3 && (z->rgb == 3 || (z->app14_color_transform == 0 && !z->jfif)); - - if (z->s->img_n == 3 && n < 3 && !is_rgb) - decode_n = 1; - else - decode_n = z->s->img_n; - - // nothing to do if no components requested; check this now to avoid - // accessing uninitialized coutput[0] later - if (decode_n <= 0) { stbi__cleanup_jpeg(z); return NULL; } - - // resample and color-convert - { - int k; - unsigned int i,j; - stbi_uc *output; - stbi_uc *coutput[4] = { NULL, NULL, NULL, NULL }; - - stbi__resample res_comp[4]; - - for (k=0; k < decode_n; ++k) { - stbi__resample *r = &res_comp[k]; - - // allocate line buffer big enough for upsampling off the edges - // with upsample factor of 4 - z->img_comp[k].linebuf = (stbi_uc *) stbi__malloc(z->s->img_x + 3); - if (!z->img_comp[k].linebuf) { stbi__cleanup_jpeg(z); return stbi__errpuc("outofmem", "Out of memory"); } - - r->hs = z->img_h_max / z->img_comp[k].h; - r->vs = z->img_v_max / z->img_comp[k].v; - r->ystep = r->vs >> 1; - r->w_lores = (z->s->img_x + r->hs-1) / r->hs; - r->ypos = 0; - r->line0 = r->line1 = z->img_comp[k].data; - - if (r->hs == 1 && r->vs == 1) r->resample = resample_row_1; - else if (r->hs == 1 && r->vs == 2) r->resample = stbi__resample_row_v_2; - else if (r->hs == 2 && r->vs == 1) r->resample = stbi__resample_row_h_2; - else if (r->hs == 2 && r->vs == 2) r->resample = z->resample_row_hv_2_kernel; - else r->resample = stbi__resample_row_generic; - } - - // can't error after this so, this is safe - output = (stbi_uc *) stbi__malloc_mad3(n, z->s->img_x, z->s->img_y, 1); - if (!output) { stbi__cleanup_jpeg(z); return stbi__errpuc("outofmem", "Out of memory"); } - - // now go ahead and resample - for (j=0; j < z->s->img_y; ++j) { - stbi_uc *out = output + n * z->s->img_x * j; - for (k=0; k < decode_n; ++k) { - stbi__resample *r = &res_comp[k]; - int y_bot = r->ystep >= (r->vs >> 1); - coutput[k] = r->resample(z->img_comp[k].linebuf, - y_bot ? r->line1 : r->line0, - y_bot ? r->line0 : r->line1, - r->w_lores, r->hs); - if (++r->ystep >= r->vs) { - r->ystep = 0; - r->line0 = r->line1; - if (++r->ypos < z->img_comp[k].y) - r->line1 += z->img_comp[k].w2; - } - } - if (n >= 3) { - stbi_uc *y = coutput[0]; - if (z->s->img_n == 3) { - if (is_rgb) { - for (i=0; i < z->s->img_x; ++i) { - out[0] = y[i]; - out[1] = coutput[1][i]; - out[2] = coutput[2][i]; - out[3] = 255; - out += n; - } - } else { - z->YCbCr_to_RGB_kernel(out, y, coutput[1], coutput[2], z->s->img_x, n); - } - } else if (z->s->img_n == 4) { - if (z->app14_color_transform == 0) { // CMYK - for (i=0; i < z->s->img_x; ++i) { - stbi_uc m = coutput[3][i]; - out[0] = stbi__blinn_8x8(coutput[0][i], m); - out[1] = stbi__blinn_8x8(coutput[1][i], m); - out[2] = stbi__blinn_8x8(coutput[2][i], m); - out[3] = 255; - out += n; - } - } else if (z->app14_color_transform == 2) { // YCCK - z->YCbCr_to_RGB_kernel(out, y, coutput[1], coutput[2], z->s->img_x, n); - for (i=0; i < z->s->img_x; ++i) { - stbi_uc m = coutput[3][i]; - out[0] = stbi__blinn_8x8(255 - out[0], m); - out[1] = stbi__blinn_8x8(255 - out[1], m); - out[2] = stbi__blinn_8x8(255 - out[2], m); - out += n; - } - } else { // YCbCr + alpha? Ignore the fourth channel for now - z->YCbCr_to_RGB_kernel(out, y, coutput[1], coutput[2], z->s->img_x, n); - } - } else - for (i=0; i < z->s->img_x; ++i) { - out[0] = out[1] = out[2] = y[i]; - out[3] = 255; // not used if n==3 - out += n; - } - } else { - if (is_rgb) { - if (n == 1) - for (i=0; i < z->s->img_x; ++i) - *out++ = stbi__compute_y(coutput[0][i], coutput[1][i], coutput[2][i]); - else { - for (i=0; i < z->s->img_x; ++i, out += 2) { - out[0] = stbi__compute_y(coutput[0][i], coutput[1][i], coutput[2][i]); - out[1] = 255; - } - } - } else if (z->s->img_n == 4 && z->app14_color_transform == 0) { - for (i=0; i < z->s->img_x; ++i) { - stbi_uc m = coutput[3][i]; - stbi_uc r = stbi__blinn_8x8(coutput[0][i], m); - stbi_uc g = stbi__blinn_8x8(coutput[1][i], m); - stbi_uc b = stbi__blinn_8x8(coutput[2][i], m); - out[0] = stbi__compute_y(r, g, b); - out[1] = 255; - out += n; - } - } else if (z->s->img_n == 4 && z->app14_color_transform == 2) { - for (i=0; i < z->s->img_x; ++i) { - out[0] = stbi__blinn_8x8(255 - coutput[0][i], coutput[3][i]); - out[1] = 255; - out += n; - } - } else { - stbi_uc *y = coutput[0]; - if (n == 1) - for (i=0; i < z->s->img_x; ++i) out[i] = y[i]; - else - for (i=0; i < z->s->img_x; ++i) { *out++ = y[i]; *out++ = 255; } - } - } - } - stbi__cleanup_jpeg(z); - *out_x = z->s->img_x; - *out_y = z->s->img_y; - if (comp) *comp = z->s->img_n >= 3 ? 3 : 1; // report original components, not output - return output; - } -} - -static void *stbi__jpeg_load(stbi__context *s, int *x, int *y, int *comp, int req_comp, stbi__result_info *ri) -{ - unsigned char* result; - stbi__jpeg* j = (stbi__jpeg*) stbi__malloc(sizeof(stbi__jpeg)); - if (!j) return stbi__errpuc("outofmem", "Out of memory"); - memset(j, 0, sizeof(stbi__jpeg)); - STBI_NOTUSED(ri); - j->s = s; - stbi__setup_jpeg(j); - result = load_jpeg_image(j, x,y,comp,req_comp); - STBI_FREE(j); - return result; -} - -static int stbi__jpeg_test(stbi__context *s) -{ - int r; - stbi__jpeg* j = (stbi__jpeg*)stbi__malloc(sizeof(stbi__jpeg)); - if (!j) return stbi__err("outofmem", "Out of memory"); - memset(j, 0, sizeof(stbi__jpeg)); - j->s = s; - stbi__setup_jpeg(j); - r = stbi__decode_jpeg_header(j, STBI__SCAN_type); - stbi__rewind(s); - STBI_FREE(j); - return r; -} - -static int stbi__jpeg_info_raw(stbi__jpeg *j, int *x, int *y, int *comp) -{ - if (!stbi__decode_jpeg_header(j, STBI__SCAN_header)) { - stbi__rewind( j->s ); - return 0; - } - if (x) *x = j->s->img_x; - if (y) *y = j->s->img_y; - if (comp) *comp = j->s->img_n >= 3 ? 3 : 1; - return 1; -} - -static int stbi__jpeg_info(stbi__context *s, int *x, int *y, int *comp) -{ - int result; - stbi__jpeg* j = (stbi__jpeg*) (stbi__malloc(sizeof(stbi__jpeg))); - if (!j) return stbi__err("outofmem", "Out of memory"); - memset(j, 0, sizeof(stbi__jpeg)); - j->s = s; - result = stbi__jpeg_info_raw(j, x, y, comp); - STBI_FREE(j); - return result; -} -#endif - -// public domain zlib decode v0.2 Sean Barrett 2006-11-18 -// simple implementation -// - all input must be provided in an upfront buffer -// - all output is written to a single output buffer (can malloc/realloc) -// performance -// - fast huffman - -#ifndef STBI_NO_ZLIB - -// fast-way is faster to check than jpeg huffman, but slow way is slower -#define STBI__ZFAST_BITS 9 // accelerate all cases in default tables -#define STBI__ZFAST_MASK ((1 << STBI__ZFAST_BITS) - 1) -#define STBI__ZNSYMS 288 // number of symbols in literal/length alphabet - -// zlib-style huffman encoding -// (jpegs packs from left, zlib from right, so can't share code) -typedef struct -{ - stbi__uint16 fast[1 << STBI__ZFAST_BITS]; - stbi__uint16 firstcode[16]; - int maxcode[17]; - stbi__uint16 firstsymbol[16]; - stbi_uc size[STBI__ZNSYMS]; - stbi__uint16 value[STBI__ZNSYMS]; -} stbi__zhuffman; - -stbi_inline static int stbi__bitreverse16(int n) -{ - n = ((n & 0xAAAA) >> 1) | ((n & 0x5555) << 1); - n = ((n & 0xCCCC) >> 2) | ((n & 0x3333) << 2); - n = ((n & 0xF0F0) >> 4) | ((n & 0x0F0F) << 4); - n = ((n & 0xFF00) >> 8) | ((n & 0x00FF) << 8); - return n; -} - -stbi_inline static int stbi__bit_reverse(int v, int bits) -{ - STBI_ASSERT(bits <= 16); - // to bit reverse n bits, reverse 16 and shift - // e.g. 11 bits, bit reverse and shift away 5 - return stbi__bitreverse16(v) >> (16-bits); -} - -static int stbi__zbuild_huffman(stbi__zhuffman *z, const stbi_uc *sizelist, int num) -{ - int i,k=0; - int code, next_code[16], sizes[17]; - - // DEFLATE spec for generating codes - memset(sizes, 0, sizeof(sizes)); - memset(z->fast, 0, sizeof(z->fast)); - for (i=0; i < num; ++i) - ++sizes[sizelist[i]]; - sizes[0] = 0; - for (i=1; i < 16; ++i) - if (sizes[i] > (1 << i)) - return stbi__err("bad sizes", "Corrupt PNG"); - code = 0; - for (i=1; i < 16; ++i) { - next_code[i] = code; - z->firstcode[i] = (stbi__uint16) code; - z->firstsymbol[i] = (stbi__uint16) k; - code = (code + sizes[i]); - if (sizes[i]) - if (code-1 >= (1 << i)) return stbi__err("bad codelengths","Corrupt PNG"); - z->maxcode[i] = code << (16-i); // preshift for inner loop - code <<= 1; - k += sizes[i]; - } - z->maxcode[16] = 0x10000; // sentinel - for (i=0; i < num; ++i) { - int s = sizelist[i]; - if (s) { - int c = next_code[s] - z->firstcode[s] + z->firstsymbol[s]; - stbi__uint16 fastv = (stbi__uint16) ((s << 9) | i); - z->size [c] = (stbi_uc ) s; - z->value[c] = (stbi__uint16) i; - if (s <= STBI__ZFAST_BITS) { - int j = stbi__bit_reverse(next_code[s],s); - while (j < (1 << STBI__ZFAST_BITS)) { - z->fast[j] = fastv; - j += (1 << s); - } - } - ++next_code[s]; - } - } - return 1; -} - -// zlib-from-memory implementation for PNG reading -// because PNG allows splitting the zlib stream arbitrarily, -// and it's annoying structurally to have PNG call ZLIB call PNG, -// we require PNG read all the IDATs and combine them into a single -// memory buffer - -typedef struct -{ - stbi_uc *zbuffer, *zbuffer_end; - int num_bits; - int hit_zeof_once; - stbi__uint32 code_buffer; - - char *zout; - char *zout_start; - char *zout_end; - int z_expandable; - - stbi__zhuffman z_length, z_distance; -} stbi__zbuf; - -stbi_inline static int stbi__zeof(stbi__zbuf *z) -{ - return (z->zbuffer >= z->zbuffer_end); -} - -stbi_inline static stbi_uc stbi__zget8(stbi__zbuf *z) -{ - return stbi__zeof(z) ? 0 : *z->zbuffer++; -} - -static void stbi__fill_bits(stbi__zbuf *z) -{ - do { - if (z->code_buffer >= (1U << z->num_bits)) { - z->zbuffer = z->zbuffer_end; /* treat this as EOF so we fail. */ - return; - } - z->code_buffer |= (unsigned int) stbi__zget8(z) << z->num_bits; - z->num_bits += 8; - } while (z->num_bits <= 24); -} - -stbi_inline static unsigned int stbi__zreceive(stbi__zbuf *z, int n) -{ - unsigned int k; - if (z->num_bits < n) stbi__fill_bits(z); - k = z->code_buffer & ((1 << n) - 1); - z->code_buffer >>= n; - z->num_bits -= n; - return k; -} - -static int stbi__zhuffman_decode_slowpath(stbi__zbuf *a, stbi__zhuffman *z) -{ - int b,s,k; - // not resolved by fast table, so compute it the slow way - // use jpeg approach, which requires MSbits at top - k = stbi__bit_reverse(a->code_buffer, 16); - for (s=STBI__ZFAST_BITS+1; ; ++s) - if (k < z->maxcode[s]) - break; - if (s >= 16) return -1; // invalid code! - // code size is s, so: - b = (k >> (16-s)) - z->firstcode[s] + z->firstsymbol[s]; - if (b >= STBI__ZNSYMS) return -1; // some data was corrupt somewhere! - if (z->size[b] != s) return -1; // was originally an assert, but report failure instead. - a->code_buffer >>= s; - a->num_bits -= s; - return z->value[b]; -} - -stbi_inline static int stbi__zhuffman_decode(stbi__zbuf *a, stbi__zhuffman *z) -{ - int b,s; - if (a->num_bits < 16) { - if (stbi__zeof(a)) { - if (!a->hit_zeof_once) { - // This is the first time we hit eof, insert 16 extra padding btis - // to allow us to keep going; if we actually consume any of them - // though, that is invalid data. This is caught later. - a->hit_zeof_once = 1; - a->num_bits += 16; // add 16 implicit zero bits - } else { - // We already inserted our extra 16 padding bits and are again - // out, this stream is actually prematurely terminated. - return -1; - } - } else { - stbi__fill_bits(a); - } - } - b = z->fast[a->code_buffer & STBI__ZFAST_MASK]; - if (b) { - s = b >> 9; - a->code_buffer >>= s; - a->num_bits -= s; - return b & 511; - } - return stbi__zhuffman_decode_slowpath(a, z); -} - -static int stbi__zexpand(stbi__zbuf *z, char *zout, int n) // need to make room for n bytes -{ - char *q; - unsigned int cur, limit, old_limit; - z->zout = zout; - if (!z->z_expandable) return stbi__err("output buffer limit","Corrupt PNG"); - cur = (unsigned int) (z->zout - z->zout_start); - limit = old_limit = (unsigned) (z->zout_end - z->zout_start); - if (UINT_MAX - cur < (unsigned) n) return stbi__err("outofmem", "Out of memory"); - while (cur + n > limit) { - if(limit > UINT_MAX / 2) return stbi__err("outofmem", "Out of memory"); - limit *= 2; - } - q = (char *) STBI_REALLOC_SIZED(z->zout_start, old_limit, limit); - STBI_NOTUSED(old_limit); - if (q == NULL) return stbi__err("outofmem", "Out of memory"); - z->zout_start = q; - z->zout = q + cur; - z->zout_end = q + limit; - return 1; -} - -static const int stbi__zlength_base[31] = { - 3,4,5,6,7,8,9,10,11,13, - 15,17,19,23,27,31,35,43,51,59, - 67,83,99,115,131,163,195,227,258,0,0 }; - -static const int stbi__zlength_extra[31]= -{ 0,0,0,0,0,0,0,0,1,1,1,1,2,2,2,2,3,3,3,3,4,4,4,4,5,5,5,5,0,0,0 }; - -static const int stbi__zdist_base[32] = { 1,2,3,4,5,7,9,13,17,25,33,49,65,97,129,193, -257,385,513,769,1025,1537,2049,3073,4097,6145,8193,12289,16385,24577,0,0}; - -static const int stbi__zdist_extra[32] = -{ 0,0,0,0,1,1,2,2,3,3,4,4,5,5,6,6,7,7,8,8,9,9,10,10,11,11,12,12,13,13}; - -static int stbi__parse_huffman_block(stbi__zbuf *a) -{ - char *zout = a->zout; - for(;;) { - int z = stbi__zhuffman_decode(a, &a->z_length); - if (z < 256) { - if (z < 0) return stbi__err("bad huffman code","Corrupt PNG"); // error in huffman codes - if (zout >= a->zout_end) { - if (!stbi__zexpand(a, zout, 1)) return 0; - zout = a->zout; - } - *zout++ = (char) z; - } else { - stbi_uc *p; - int len,dist; - if (z == 256) { - a->zout = zout; - if (a->hit_zeof_once && a->num_bits < 16) { - // The first time we hit zeof, we inserted 16 extra zero bits into our bit - // buffer so the decoder can just do its speculative decoding. But if we - // actually consumed any of those bits (which is the case when num_bits < 16), - // the stream actually read past the end so it is malformed. - return stbi__err("unexpected end","Corrupt PNG"); - } - return 1; - } - if (z >= 286) return stbi__err("bad huffman code","Corrupt PNG"); // per DEFLATE, length codes 286 and 287 must not appear in compressed data - z -= 257; - len = stbi__zlength_base[z]; - if (stbi__zlength_extra[z]) len += stbi__zreceive(a, stbi__zlength_extra[z]); - z = stbi__zhuffman_decode(a, &a->z_distance); - if (z < 0 || z >= 30) return stbi__err("bad huffman code","Corrupt PNG"); // per DEFLATE, distance codes 30 and 31 must not appear in compressed data - dist = stbi__zdist_base[z]; - if (stbi__zdist_extra[z]) dist += stbi__zreceive(a, stbi__zdist_extra[z]); - if (zout - a->zout_start < dist) return stbi__err("bad dist","Corrupt PNG"); - if (len > a->zout_end - zout) { - if (!stbi__zexpand(a, zout, len)) return 0; - zout = a->zout; - } - p = (stbi_uc *) (zout - dist); - if (dist == 1) { // run of one byte; common in images. - stbi_uc v = *p; - if (len) { do *zout++ = v; while (--len); } - } else { - if (len) { do *zout++ = *p++; while (--len); } - } - } - } -} - -static int stbi__compute_huffman_codes(stbi__zbuf *a) -{ - static const stbi_uc length_dezigzag[19] = { 16,17,18,0,8,7,9,6,10,5,11,4,12,3,13,2,14,1,15 }; - stbi__zhuffman z_codelength; - stbi_uc lencodes[286+32+137];//padding for maximum single op - stbi_uc codelength_sizes[19]; - int i,n; - - int hlit = stbi__zreceive(a,5) + 257; - int hdist = stbi__zreceive(a,5) + 1; - int hclen = stbi__zreceive(a,4) + 4; - int ntot = hlit + hdist; - - memset(codelength_sizes, 0, sizeof(codelength_sizes)); - for (i=0; i < hclen; ++i) { - int s = stbi__zreceive(a,3); - codelength_sizes[length_dezigzag[i]] = (stbi_uc) s; - } - if (!stbi__zbuild_huffman(&z_codelength, codelength_sizes, 19)) return 0; - - n = 0; - while (n < ntot) { - int c = stbi__zhuffman_decode(a, &z_codelength); - if (c < 0 || c >= 19) return stbi__err("bad codelengths", "Corrupt PNG"); - if (c < 16) - lencodes[n++] = (stbi_uc) c; - else { - stbi_uc fill = 0; - if (c == 16) { - c = stbi__zreceive(a,2)+3; - if (n == 0) return stbi__err("bad codelengths", "Corrupt PNG"); - fill = lencodes[n-1]; - } else if (c == 17) { - c = stbi__zreceive(a,3)+3; - } else if (c == 18) { - c = stbi__zreceive(a,7)+11; - } else { - return stbi__err("bad codelengths", "Corrupt PNG"); - } - if (ntot - n < c) return stbi__err("bad codelengths", "Corrupt PNG"); - memset(lencodes+n, fill, c); - n += c; - } - } - if (n != ntot) return stbi__err("bad codelengths","Corrupt PNG"); - if (!stbi__zbuild_huffman(&a->z_length, lencodes, hlit)) return 0; - if (!stbi__zbuild_huffman(&a->z_distance, lencodes+hlit, hdist)) return 0; - return 1; -} - -static int stbi__parse_uncompressed_block(stbi__zbuf *a) -{ - stbi_uc header[4]; - int len,nlen,k; - if (a->num_bits & 7) - stbi__zreceive(a, a->num_bits & 7); // discard - // drain the bit-packed data into header - k = 0; - while (a->num_bits > 0) { - header[k++] = (stbi_uc) (a->code_buffer & 255); // suppress MSVC run-time check - a->code_buffer >>= 8; - a->num_bits -= 8; - } - if (a->num_bits < 0) return stbi__err("zlib corrupt","Corrupt PNG"); - // now fill header the normal way - while (k < 4) - header[k++] = stbi__zget8(a); - len = header[1] * 256 + header[0]; - nlen = header[3] * 256 + header[2]; - if (nlen != (len ^ 0xffff)) return stbi__err("zlib corrupt","Corrupt PNG"); - if (a->zbuffer + len > a->zbuffer_end) return stbi__err("read past buffer","Corrupt PNG"); - if (a->zout + len > a->zout_end) - if (!stbi__zexpand(a, a->zout, len)) return 0; - memcpy(a->zout, a->zbuffer, len); - a->zbuffer += len; - a->zout += len; - return 1; -} - -static int stbi__parse_zlib_header(stbi__zbuf *a) -{ - int cmf = stbi__zget8(a); - int cm = cmf & 15; - /* int cinfo = cmf >> 4; */ - int flg = stbi__zget8(a); - if (stbi__zeof(a)) return stbi__err("bad zlib header","Corrupt PNG"); // zlib spec - if ((cmf*256+flg) % 31 != 0) return stbi__err("bad zlib header","Corrupt PNG"); // zlib spec - if (flg & 32) return stbi__err("no preset dict","Corrupt PNG"); // preset dictionary not allowed in png - if (cm != 8) return stbi__err("bad compression","Corrupt PNG"); // DEFLATE required for png - // window = 1 << (8 + cinfo)... but who cares, we fully buffer output - return 1; -} - -static const stbi_uc stbi__zdefault_length[STBI__ZNSYMS] = -{ - 8,8,8,8,8,8,8,8,8,8,8,8,8,8,8,8, 8,8,8,8,8,8,8,8,8,8,8,8,8,8,8,8, - 8,8,8,8,8,8,8,8,8,8,8,8,8,8,8,8, 8,8,8,8,8,8,8,8,8,8,8,8,8,8,8,8, - 8,8,8,8,8,8,8,8,8,8,8,8,8,8,8,8, 8,8,8,8,8,8,8,8,8,8,8,8,8,8,8,8, - 8,8,8,8,8,8,8,8,8,8,8,8,8,8,8,8, 8,8,8,8,8,8,8,8,8,8,8,8,8,8,8,8, - 8,8,8,8,8,8,8,8,8,8,8,8,8,8,8,8, 9,9,9,9,9,9,9,9,9,9,9,9,9,9,9,9, - 9,9,9,9,9,9,9,9,9,9,9,9,9,9,9,9, 9,9,9,9,9,9,9,9,9,9,9,9,9,9,9,9, - 9,9,9,9,9,9,9,9,9,9,9,9,9,9,9,9, 9,9,9,9,9,9,9,9,9,9,9,9,9,9,9,9, - 9,9,9,9,9,9,9,9,9,9,9,9,9,9,9,9, 9,9,9,9,9,9,9,9,9,9,9,9,9,9,9,9, - 7,7,7,7,7,7,7,7,7,7,7,7,7,7,7,7, 7,7,7,7,7,7,7,7,8,8,8,8,8,8,8,8 -}; -static const stbi_uc stbi__zdefault_distance[32] = -{ - 5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5,5 -}; -/* -Init algorithm: -{ - int i; // use <= to match clearly with spec - for (i=0; i <= 143; ++i) stbi__zdefault_length[i] = 8; - for ( ; i <= 255; ++i) stbi__zdefault_length[i] = 9; - for ( ; i <= 279; ++i) stbi__zdefault_length[i] = 7; - for ( ; i <= 287; ++i) stbi__zdefault_length[i] = 8; - - for (i=0; i <= 31; ++i) stbi__zdefault_distance[i] = 5; -} -*/ - -static int stbi__parse_zlib(stbi__zbuf *a, int parse_header) -{ - int final, type; - if (parse_header) - if (!stbi__parse_zlib_header(a)) return 0; - a->num_bits = 0; - a->code_buffer = 0; - a->hit_zeof_once = 0; - do { - final = stbi__zreceive(a,1); - type = stbi__zreceive(a,2); - if (type == 0) { - if (!stbi__parse_uncompressed_block(a)) return 0; - } else if (type == 3) { - return 0; - } else { - if (type == 1) { - // use fixed code lengths - if (!stbi__zbuild_huffman(&a->z_length , stbi__zdefault_length , STBI__ZNSYMS)) return 0; - if (!stbi__zbuild_huffman(&a->z_distance, stbi__zdefault_distance, 32)) return 0; - } else { - if (!stbi__compute_huffman_codes(a)) return 0; - } - if (!stbi__parse_huffman_block(a)) return 0; - } - } while (!final); - return 1; -} - -static int stbi__do_zlib(stbi__zbuf *a, char *obuf, int olen, int exp, int parse_header) -{ - a->zout_start = obuf; - a->zout = obuf; - a->zout_end = obuf + olen; - a->z_expandable = exp; - - return stbi__parse_zlib(a, parse_header); -} - -STBIDEF char *stbi_zlib_decode_malloc_guesssize(const char *buffer, int len, int initial_size, int *outlen) -{ - stbi__zbuf a; - char *p = (char *) stbi__malloc(initial_size); - if (p == NULL) return NULL; - a.zbuffer = (stbi_uc *) buffer; - a.zbuffer_end = (stbi_uc *) buffer + len; - if (stbi__do_zlib(&a, p, initial_size, 1, 1)) { - if (outlen) *outlen = (int) (a.zout - a.zout_start); - return a.zout_start; - } else { - STBI_FREE(a.zout_start); - return NULL; - } -} - -STBIDEF char *stbi_zlib_decode_malloc(char const *buffer, int len, int *outlen) -{ - return stbi_zlib_decode_malloc_guesssize(buffer, len, 16384, outlen); -} - -STBIDEF char *stbi_zlib_decode_malloc_guesssize_headerflag(const char *buffer, int len, int initial_size, int *outlen, int parse_header) -{ - stbi__zbuf a; - char *p = (char *) stbi__malloc(initial_size); - if (p == NULL) return NULL; - a.zbuffer = (stbi_uc *) buffer; - a.zbuffer_end = (stbi_uc *) buffer + len; - if (stbi__do_zlib(&a, p, initial_size, 1, parse_header)) { - if (outlen) *outlen = (int) (a.zout - a.zout_start); - return a.zout_start; - } else { - STBI_FREE(a.zout_start); - return NULL; - } -} - -STBIDEF int stbi_zlib_decode_buffer(char *obuffer, int olen, char const *ibuffer, int ilen) -{ - stbi__zbuf a; - a.zbuffer = (stbi_uc *) ibuffer; - a.zbuffer_end = (stbi_uc *) ibuffer + ilen; - if (stbi__do_zlib(&a, obuffer, olen, 0, 1)) - return (int) (a.zout - a.zout_start); - else - return -1; -} - -STBIDEF char *stbi_zlib_decode_noheader_malloc(char const *buffer, int len, int *outlen) -{ - stbi__zbuf a; - char *p = (char *) stbi__malloc(16384); - if (p == NULL) return NULL; - a.zbuffer = (stbi_uc *) buffer; - a.zbuffer_end = (stbi_uc *) buffer+len; - if (stbi__do_zlib(&a, p, 16384, 1, 0)) { - if (outlen) *outlen = (int) (a.zout - a.zout_start); - return a.zout_start; - } else { - STBI_FREE(a.zout_start); - return NULL; - } -} - -STBIDEF int stbi_zlib_decode_noheader_buffer(char *obuffer, int olen, const char *ibuffer, int ilen) -{ - stbi__zbuf a; - a.zbuffer = (stbi_uc *) ibuffer; - a.zbuffer_end = (stbi_uc *) ibuffer + ilen; - if (stbi__do_zlib(&a, obuffer, olen, 0, 0)) - return (int) (a.zout - a.zout_start); - else - return -1; -} -#endif - -// public domain "baseline" PNG decoder v0.10 Sean Barrett 2006-11-18 -// simple implementation -// - only 8-bit samples -// - no CRC checking -// - allocates lots of intermediate memory -// - avoids problem of streaming data between subsystems -// - avoids explicit window management -// performance -// - uses stb_zlib, a PD zlib implementation with fast huffman decoding - -#ifndef STBI_NO_PNG -typedef struct -{ - stbi__uint32 length; - stbi__uint32 type; -} stbi__pngchunk; - -static stbi__pngchunk stbi__get_chunk_header(stbi__context *s) -{ - stbi__pngchunk c; - c.length = stbi__get32be(s); - c.type = stbi__get32be(s); - return c; -} - -static int stbi__check_png_header(stbi__context *s) -{ - static const stbi_uc png_sig[8] = { 137,80,78,71,13,10,26,10 }; - int i; - for (i=0; i < 8; ++i) - if (stbi__get8(s) != png_sig[i]) return stbi__err("bad png sig","Not a PNG"); - return 1; -} - -typedef struct -{ - stbi__context *s; - stbi_uc *idata, *expanded, *out; - int depth; -} stbi__png; - - -enum { - STBI__F_none=0, - STBI__F_sub=1, - STBI__F_up=2, - STBI__F_avg=3, - STBI__F_paeth=4, - // synthetic filter used for first scanline to avoid needing a dummy row of 0s - STBI__F_avg_first -}; - -static stbi_uc first_row_filter[5] = -{ - STBI__F_none, - STBI__F_sub, - STBI__F_none, - STBI__F_avg_first, - STBI__F_sub // Paeth with b=c=0 turns out to be equivalent to sub -}; - -static int stbi__paeth(int a, int b, int c) -{ - // This formulation looks very different from the reference in the PNG spec, but is - // actually equivalent and has favorable data dependencies and admits straightforward - // generation of branch-free code, which helps performance significantly. - int thresh = c*3 - (a + b); - int lo = a < b ? a : b; - int hi = a < b ? b : a; - int t0 = (hi <= thresh) ? lo : c; - int t1 = (thresh <= lo) ? hi : t0; - return t1; -} - -static const stbi_uc stbi__depth_scale_table[9] = { 0, 0xff, 0x55, 0, 0x11, 0,0,0, 0x01 }; - -// adds an extra all-255 alpha channel -// dest == src is legal -// img_n must be 1 or 3 -static void stbi__create_png_alpha_expand8(stbi_uc *dest, stbi_uc *src, stbi__uint32 x, int img_n) -{ - int i; - // must process data backwards since we allow dest==src - if (img_n == 1) { - for (i=x-1; i >= 0; --i) { - dest[i*2+1] = 255; - dest[i*2+0] = src[i]; - } - } else { - STBI_ASSERT(img_n == 3); - for (i=x-1; i >= 0; --i) { - dest[i*4+3] = 255; - dest[i*4+2] = src[i*3+2]; - dest[i*4+1] = src[i*3+1]; - dest[i*4+0] = src[i*3+0]; - } - } -} - -// create the png data from post-deflated data -static int stbi__create_png_image_raw(stbi__png *a, stbi_uc *raw, stbi__uint32 raw_len, int out_n, stbi__uint32 x, stbi__uint32 y, int depth, int color) -{ - int bytes = (depth == 16 ? 2 : 1); - stbi__context *s = a->s; - stbi__uint32 i,j,stride = x*out_n*bytes; - stbi__uint32 img_len, img_width_bytes; - stbi_uc *filter_buf; - int all_ok = 1; - int k; - int img_n = s->img_n; // copy it into a local for later - - int output_bytes = out_n*bytes; - int filter_bytes = img_n*bytes; - int width = x; - - STBI_ASSERT(out_n == s->img_n || out_n == s->img_n+1); - a->out = (stbi_uc *) stbi__malloc_mad3(x, y, output_bytes, 0); // extra bytes to write off the end into - if (!a->out) return stbi__err("outofmem", "Out of memory"); - - // note: error exits here don't need to clean up a->out individually, - // stbi__do_png always does on error. - if (!stbi__mad3sizes_valid(img_n, x, depth, 7)) return stbi__err("too large", "Corrupt PNG"); - img_width_bytes = (((img_n * x * depth) + 7) >> 3); - if (!stbi__mad2sizes_valid(img_width_bytes, y, img_width_bytes)) return stbi__err("too large", "Corrupt PNG"); - img_len = (img_width_bytes + 1) * y; - - // we used to check for exact match between raw_len and img_len on non-interlaced PNGs, - // but issue #276 reported a PNG in the wild that had extra data at the end (all zeros), - // so just check for raw_len < img_len always. - if (raw_len < img_len) return stbi__err("not enough pixels","Corrupt PNG"); - - // Allocate two scan lines worth of filter workspace buffer. - filter_buf = (stbi_uc *) stbi__malloc_mad2(img_width_bytes, 2, 0); - if (!filter_buf) return stbi__err("outofmem", "Out of memory"); - - // Filtering for low-bit-depth images - if (depth < 8) { - filter_bytes = 1; - width = img_width_bytes; - } - - for (j=0; j < y; ++j) { - // cur/prior filter buffers alternate - stbi_uc *cur = filter_buf + (j & 1)*img_width_bytes; - stbi_uc *prior = filter_buf + (~j & 1)*img_width_bytes; - stbi_uc *dest = a->out + stride*j; - int nk = width * filter_bytes; - int filter = *raw++; - - // check filter type - if (filter > 4) { - all_ok = stbi__err("invalid filter","Corrupt PNG"); - break; - } - - // if first row, use special filter that doesn't sample previous row - if (j == 0) filter = first_row_filter[filter]; - - // perform actual filtering - switch (filter) { - case STBI__F_none: - memcpy(cur, raw, nk); - break; - case STBI__F_sub: - memcpy(cur, raw, filter_bytes); - for (k = filter_bytes; k < nk; ++k) - cur[k] = STBI__BYTECAST(raw[k] + cur[k-filter_bytes]); - break; - case STBI__F_up: - for (k = 0; k < nk; ++k) - cur[k] = STBI__BYTECAST(raw[k] + prior[k]); - break; - case STBI__F_avg: - for (k = 0; k < filter_bytes; ++k) - cur[k] = STBI__BYTECAST(raw[k] + (prior[k]>>1)); - for (k = filter_bytes; k < nk; ++k) - cur[k] = STBI__BYTECAST(raw[k] + ((prior[k] + cur[k-filter_bytes])>>1)); - break; - case STBI__F_paeth: - for (k = 0; k < filter_bytes; ++k) - cur[k] = STBI__BYTECAST(raw[k] + prior[k]); // prior[k] == stbi__paeth(0,prior[k],0) - for (k = filter_bytes; k < nk; ++k) - cur[k] = STBI__BYTECAST(raw[k] + stbi__paeth(cur[k-filter_bytes], prior[k], prior[k-filter_bytes])); - break; - case STBI__F_avg_first: - memcpy(cur, raw, filter_bytes); - for (k = filter_bytes; k < nk; ++k) - cur[k] = STBI__BYTECAST(raw[k] + (cur[k-filter_bytes] >> 1)); - break; - } - - raw += nk; - - // expand decoded bits in cur to dest, also adding an extra alpha channel if desired - if (depth < 8) { - stbi_uc scale = (color == 0) ? stbi__depth_scale_table[depth] : 1; // scale grayscale values to 0..255 range - stbi_uc *in = cur; - stbi_uc *out = dest; - stbi_uc inb = 0; - stbi__uint32 nsmp = x*img_n; - - // expand bits to bytes first - if (depth == 4) { - for (i=0; i < nsmp; ++i) { - if ((i & 1) == 0) inb = *in++; - *out++ = scale * (inb >> 4); - inb <<= 4; - } - } else if (depth == 2) { - for (i=0; i < nsmp; ++i) { - if ((i & 3) == 0) inb = *in++; - *out++ = scale * (inb >> 6); - inb <<= 2; - } - } else { - STBI_ASSERT(depth == 1); - for (i=0; i < nsmp; ++i) { - if ((i & 7) == 0) inb = *in++; - *out++ = scale * (inb >> 7); - inb <<= 1; - } - } - - // insert alpha=255 values if desired - if (img_n != out_n) - stbi__create_png_alpha_expand8(dest, dest, x, img_n); - } else if (depth == 8) { - if (img_n == out_n) - memcpy(dest, cur, x*img_n); - else - stbi__create_png_alpha_expand8(dest, cur, x, img_n); - } else if (depth == 16) { - // convert the image data from big-endian to platform-native - stbi__uint16 *dest16 = (stbi__uint16*)dest; - stbi__uint32 nsmp = x*img_n; - - if (img_n == out_n) { - for (i = 0; i < nsmp; ++i, ++dest16, cur += 2) - *dest16 = (cur[0] << 8) | cur[1]; - } else { - STBI_ASSERT(img_n+1 == out_n); - if (img_n == 1) { - for (i = 0; i < x; ++i, dest16 += 2, cur += 2) { - dest16[0] = (cur[0] << 8) | cur[1]; - dest16[1] = 0xffff; - } - } else { - STBI_ASSERT(img_n == 3); - for (i = 0; i < x; ++i, dest16 += 4, cur += 6) { - dest16[0] = (cur[0] << 8) | cur[1]; - dest16[1] = (cur[2] << 8) | cur[3]; - dest16[2] = (cur[4] << 8) | cur[5]; - dest16[3] = 0xffff; - } - } - } - } - } - - STBI_FREE(filter_buf); - if (!all_ok) return 0; - - return 1; -} - -static int stbi__create_png_image(stbi__png *a, stbi_uc *image_data, stbi__uint32 image_data_len, int out_n, int depth, int color, int interlaced) -{ - int bytes = (depth == 16 ? 2 : 1); - int out_bytes = out_n * bytes; - stbi_uc *final; - int p; - if (!interlaced) - return stbi__create_png_image_raw(a, image_data, image_data_len, out_n, a->s->img_x, a->s->img_y, depth, color); - - // de-interlacing - final = (stbi_uc *) stbi__malloc_mad3(a->s->img_x, a->s->img_y, out_bytes, 0); - if (!final) return stbi__err("outofmem", "Out of memory"); - for (p=0; p < 7; ++p) { - int xorig[] = { 0,4,0,2,0,1,0 }; - int yorig[] = { 0,0,4,0,2,0,1 }; - int xspc[] = { 8,8,4,4,2,2,1 }; - int yspc[] = { 8,8,8,4,4,2,2 }; - int i,j,x,y; - // pass1_x[4] = 0, pass1_x[5] = 1, pass1_x[12] = 1 - x = (a->s->img_x - xorig[p] + xspc[p]-1) / xspc[p]; - y = (a->s->img_y - yorig[p] + yspc[p]-1) / yspc[p]; - if (x && y) { - stbi__uint32 img_len = ((((a->s->img_n * x * depth) + 7) >> 3) + 1) * y; - if (!stbi__create_png_image_raw(a, image_data, image_data_len, out_n, x, y, depth, color)) { - STBI_FREE(final); - return 0; - } - for (j=0; j < y; ++j) { - for (i=0; i < x; ++i) { - int out_y = j*yspc[p]+yorig[p]; - int out_x = i*xspc[p]+xorig[p]; - memcpy(final + out_y*a->s->img_x*out_bytes + out_x*out_bytes, - a->out + (j*x+i)*out_bytes, out_bytes); - } - } - STBI_FREE(a->out); - image_data += img_len; - image_data_len -= img_len; - } - } - a->out = final; - - return 1; -} - -static int stbi__compute_transparency(stbi__png *z, stbi_uc tc[3], int out_n) -{ - stbi__context *s = z->s; - stbi__uint32 i, pixel_count = s->img_x * s->img_y; - stbi_uc *p = z->out; - - // compute color-based transparency, assuming we've - // already got 255 as the alpha value in the output - STBI_ASSERT(out_n == 2 || out_n == 4); - - if (out_n == 2) { - for (i=0; i < pixel_count; ++i) { - p[1] = (p[0] == tc[0] ? 0 : 255); - p += 2; - } - } else { - for (i=0; i < pixel_count; ++i) { - if (p[0] == tc[0] && p[1] == tc[1] && p[2] == tc[2]) - p[3] = 0; - p += 4; - } - } - return 1; -} - -static int stbi__compute_transparency16(stbi__png *z, stbi__uint16 tc[3], int out_n) -{ - stbi__context *s = z->s; - stbi__uint32 i, pixel_count = s->img_x * s->img_y; - stbi__uint16 *p = (stbi__uint16*) z->out; - - // compute color-based transparency, assuming we've - // already got 65535 as the alpha value in the output - STBI_ASSERT(out_n == 2 || out_n == 4); - - if (out_n == 2) { - for (i = 0; i < pixel_count; ++i) { - p[1] = (p[0] == tc[0] ? 0 : 65535); - p += 2; - } - } else { - for (i = 0; i < pixel_count; ++i) { - if (p[0] == tc[0] && p[1] == tc[1] && p[2] == tc[2]) - p[3] = 0; - p += 4; - } - } - return 1; -} - -static int stbi__expand_png_palette(stbi__png *a, stbi_uc *palette, int len, int pal_img_n) -{ - stbi__uint32 i, pixel_count = a->s->img_x * a->s->img_y; - stbi_uc *p, *temp_out, *orig = a->out; - - p = (stbi_uc *) stbi__malloc_mad2(pixel_count, pal_img_n, 0); - if (p == NULL) return stbi__err("outofmem", "Out of memory"); - - // between here and free(out) below, exitting would leak - temp_out = p; - - if (pal_img_n == 3) { - for (i=0; i < pixel_count; ++i) { - int n = orig[i]*4; - p[0] = palette[n ]; - p[1] = palette[n+1]; - p[2] = palette[n+2]; - p += 3; - } - } else { - for (i=0; i < pixel_count; ++i) { - int n = orig[i]*4; - p[0] = palette[n ]; - p[1] = palette[n+1]; - p[2] = palette[n+2]; - p[3] = palette[n+3]; - p += 4; - } - } - STBI_FREE(a->out); - a->out = temp_out; - - STBI_NOTUSED(len); - - return 1; -} - -static int stbi__unpremultiply_on_load_global = 0; -static int stbi__de_iphone_flag_global = 0; - -STBIDEF void stbi_set_unpremultiply_on_load(int flag_true_if_should_unpremultiply) -{ - stbi__unpremultiply_on_load_global = flag_true_if_should_unpremultiply; -} - -STBIDEF void stbi_convert_iphone_png_to_rgb(int flag_true_if_should_convert) -{ - stbi__de_iphone_flag_global = flag_true_if_should_convert; -} - -#ifndef STBI_THREAD_LOCAL -#define stbi__unpremultiply_on_load stbi__unpremultiply_on_load_global -#define stbi__de_iphone_flag stbi__de_iphone_flag_global -#else -static STBI_THREAD_LOCAL int stbi__unpremultiply_on_load_local, stbi__unpremultiply_on_load_set; -static STBI_THREAD_LOCAL int stbi__de_iphone_flag_local, stbi__de_iphone_flag_set; - -STBIDEF void stbi_set_unpremultiply_on_load_thread(int flag_true_if_should_unpremultiply) -{ - stbi__unpremultiply_on_load_local = flag_true_if_should_unpremultiply; - stbi__unpremultiply_on_load_set = 1; -} - -STBIDEF void stbi_convert_iphone_png_to_rgb_thread(int flag_true_if_should_convert) -{ - stbi__de_iphone_flag_local = flag_true_if_should_convert; - stbi__de_iphone_flag_set = 1; -} - -#define stbi__unpremultiply_on_load (stbi__unpremultiply_on_load_set \ - ? stbi__unpremultiply_on_load_local \ - : stbi__unpremultiply_on_load_global) -#define stbi__de_iphone_flag (stbi__de_iphone_flag_set \ - ? stbi__de_iphone_flag_local \ - : stbi__de_iphone_flag_global) -#endif // STBI_THREAD_LOCAL - -static void stbi__de_iphone(stbi__png *z) -{ - stbi__context *s = z->s; - stbi__uint32 i, pixel_count = s->img_x * s->img_y; - stbi_uc *p = z->out; - - if (s->img_out_n == 3) { // convert bgr to rgb - for (i=0; i < pixel_count; ++i) { - stbi_uc t = p[0]; - p[0] = p[2]; - p[2] = t; - p += 3; - } - } else { - STBI_ASSERT(s->img_out_n == 4); - if (stbi__unpremultiply_on_load) { - // convert bgr to rgb and unpremultiply - for (i=0; i < pixel_count; ++i) { - stbi_uc a = p[3]; - stbi_uc t = p[0]; - if (a) { - stbi_uc half = a / 2; - p[0] = (p[2] * 255 + half) / a; - p[1] = (p[1] * 255 + half) / a; - p[2] = ( t * 255 + half) / a; - } else { - p[0] = p[2]; - p[2] = t; - } - p += 4; - } - } else { - // convert bgr to rgb - for (i=0; i < pixel_count; ++i) { - stbi_uc t = p[0]; - p[0] = p[2]; - p[2] = t; - p += 4; - } - } - } -} - -#define STBI__PNG_TYPE(a,b,c,d) (((unsigned) (a) << 24) + ((unsigned) (b) << 16) + ((unsigned) (c) << 8) + (unsigned) (d)) - -static int stbi__parse_png_file(stbi__png *z, int scan, int req_comp) -{ - stbi_uc palette[1024], pal_img_n=0; - stbi_uc has_trans=0, tc[3]={0}; - stbi__uint16 tc16[3]; - stbi__uint32 ioff=0, idata_limit=0, i, pal_len=0; - int first=1,k,interlace=0, color=0, is_iphone=0; - stbi__context *s = z->s; - - z->expanded = NULL; - z->idata = NULL; - z->out = NULL; - - if (!stbi__check_png_header(s)) return 0; - - if (scan == STBI__SCAN_type) return 1; - - for (;;) { - stbi__pngchunk c = stbi__get_chunk_header(s); - switch (c.type) { - case STBI__PNG_TYPE('C','g','B','I'): - is_iphone = 1; - stbi__skip(s, c.length); - break; - case STBI__PNG_TYPE('I','H','D','R'): { - int comp,filter; - if (!first) return stbi__err("multiple IHDR","Corrupt PNG"); - first = 0; - if (c.length != 13) return stbi__err("bad IHDR len","Corrupt PNG"); - s->img_x = stbi__get32be(s); - s->img_y = stbi__get32be(s); - if (s->img_y > STBI_MAX_DIMENSIONS) return stbi__err("too large","Very large image (corrupt?)"); - if (s->img_x > STBI_MAX_DIMENSIONS) return stbi__err("too large","Very large image (corrupt?)"); - z->depth = stbi__get8(s); if (z->depth != 1 && z->depth != 2 && z->depth != 4 && z->depth != 8 && z->depth != 16) return stbi__err("1/2/4/8/16-bit only","PNG not supported: 1/2/4/8/16-bit only"); - color = stbi__get8(s); if (color > 6) return stbi__err("bad ctype","Corrupt PNG"); - if (color == 3 && z->depth == 16) return stbi__err("bad ctype","Corrupt PNG"); - if (color == 3) pal_img_n = 3; else if (color & 1) return stbi__err("bad ctype","Corrupt PNG"); - comp = stbi__get8(s); if (comp) return stbi__err("bad comp method","Corrupt PNG"); - filter= stbi__get8(s); if (filter) return stbi__err("bad filter method","Corrupt PNG"); - interlace = stbi__get8(s); if (interlace>1) return stbi__err("bad interlace method","Corrupt PNG"); - if (!s->img_x || !s->img_y) return stbi__err("0-pixel image","Corrupt PNG"); - if (!pal_img_n) { - s->img_n = (color & 2 ? 3 : 1) + (color & 4 ? 1 : 0); - if ((1 << 30) / s->img_x / s->img_n < s->img_y) return stbi__err("too large", "Image too large to decode"); - } else { - // if paletted, then pal_n is our final components, and - // img_n is # components to decompress/filter. - s->img_n = 1; - if ((1 << 30) / s->img_x / 4 < s->img_y) return stbi__err("too large","Corrupt PNG"); - } - // even with SCAN_header, have to scan to see if we have a tRNS - break; - } - - case STBI__PNG_TYPE('P','L','T','E'): { - if (first) return stbi__err("first not IHDR", "Corrupt PNG"); - if (c.length > 256*3) return stbi__err("invalid PLTE","Corrupt PNG"); - pal_len = c.length / 3; - if (pal_len * 3 != c.length) return stbi__err("invalid PLTE","Corrupt PNG"); - for (i=0; i < pal_len; ++i) { - palette[i*4+0] = stbi__get8(s); - palette[i*4+1] = stbi__get8(s); - palette[i*4+2] = stbi__get8(s); - palette[i*4+3] = 255; - } - break; - } - - case STBI__PNG_TYPE('t','R','N','S'): { - if (first) return stbi__err("first not IHDR", "Corrupt PNG"); - if (z->idata) return stbi__err("tRNS after IDAT","Corrupt PNG"); - if (pal_img_n) { - if (scan == STBI__SCAN_header) { s->img_n = 4; return 1; } - if (pal_len == 0) return stbi__err("tRNS before PLTE","Corrupt PNG"); - if (c.length > pal_len) return stbi__err("bad tRNS len","Corrupt PNG"); - pal_img_n = 4; - for (i=0; i < c.length; ++i) - palette[i*4+3] = stbi__get8(s); - } else { - if (!(s->img_n & 1)) return stbi__err("tRNS with alpha","Corrupt PNG"); - if (c.length != (stbi__uint32) s->img_n*2) return stbi__err("bad tRNS len","Corrupt PNG"); - has_trans = 1; - // non-paletted with tRNS = constant alpha. if header-scanning, we can stop now. - if (scan == STBI__SCAN_header) { ++s->img_n; return 1; } - if (z->depth == 16) { - for (k = 0; k < s->img_n && k < 3; ++k) // extra loop test to suppress false GCC warning - tc16[k] = (stbi__uint16)stbi__get16be(s); // copy the values as-is - } else { - for (k = 0; k < s->img_n && k < 3; ++k) - tc[k] = (stbi_uc)(stbi__get16be(s) & 255) * stbi__depth_scale_table[z->depth]; // non 8-bit images will be larger - } - } - break; - } - - case STBI__PNG_TYPE('I','D','A','T'): { - if (first) return stbi__err("first not IHDR", "Corrupt PNG"); - if (pal_img_n && !pal_len) return stbi__err("no PLTE","Corrupt PNG"); - if (scan == STBI__SCAN_header) { - // header scan definitely stops at first IDAT - if (pal_img_n) - s->img_n = pal_img_n; - return 1; - } - if (c.length > (1u << 30)) return stbi__err("IDAT size limit", "IDAT section larger than 2^30 bytes"); - if ((int)(ioff + c.length) < (int)ioff) return 0; - if (ioff + c.length > idata_limit) { - stbi__uint32 idata_limit_old = idata_limit; - stbi_uc *p; - if (idata_limit == 0) idata_limit = c.length > 4096 ? c.length : 4096; - while (ioff + c.length > idata_limit) - idata_limit *= 2; - STBI_NOTUSED(idata_limit_old); - p = (stbi_uc *) STBI_REALLOC_SIZED(z->idata, idata_limit_old, idata_limit); if (p == NULL) return stbi__err("outofmem", "Out of memory"); - z->idata = p; - } - if (!stbi__getn(s, z->idata+ioff,c.length)) return stbi__err("outofdata","Corrupt PNG"); - ioff += c.length; - break; - } - - case STBI__PNG_TYPE('I','E','N','D'): { - stbi__uint32 raw_len, bpl; - if (first) return stbi__err("first not IHDR", "Corrupt PNG"); - if (scan != STBI__SCAN_load) return 1; - if (z->idata == NULL) return stbi__err("no IDAT","Corrupt PNG"); - // initial guess for decoded data size to avoid unnecessary reallocs - bpl = (s->img_x * z->depth + 7) / 8; // bytes per line, per component - raw_len = bpl * s->img_y * s->img_n /* pixels */ + s->img_y /* filter mode per row */; - z->expanded = (stbi_uc *) stbi_zlib_decode_malloc_guesssize_headerflag((char *) z->idata, ioff, raw_len, (int *) &raw_len, !is_iphone); - if (z->expanded == NULL) return 0; // zlib should set error - STBI_FREE(z->idata); z->idata = NULL; - if ((req_comp == s->img_n+1 && req_comp != 3 && !pal_img_n) || has_trans) - s->img_out_n = s->img_n+1; - else - s->img_out_n = s->img_n; - if (!stbi__create_png_image(z, z->expanded, raw_len, s->img_out_n, z->depth, color, interlace)) return 0; - if (has_trans) { - if (z->depth == 16) { - if (!stbi__compute_transparency16(z, tc16, s->img_out_n)) return 0; - } else { - if (!stbi__compute_transparency(z, tc, s->img_out_n)) return 0; - } - } - if (is_iphone && stbi__de_iphone_flag && s->img_out_n > 2) - stbi__de_iphone(z); - if (pal_img_n) { - // pal_img_n == 3 or 4 - s->img_n = pal_img_n; // record the actual colors we had - s->img_out_n = pal_img_n; - if (req_comp >= 3) s->img_out_n = req_comp; - if (!stbi__expand_png_palette(z, palette, pal_len, s->img_out_n)) - return 0; - } else if (has_trans) { - // non-paletted image with tRNS -> source image has (constant) alpha - ++s->img_n; - } - STBI_FREE(z->expanded); z->expanded = NULL; - // end of PNG chunk, read and skip CRC - stbi__get32be(s); - return 1; - } - - default: - // if critical, fail - if (first) return stbi__err("first not IHDR", "Corrupt PNG"); - if ((c.type & (1 << 29)) == 0) { - #ifndef STBI_NO_FAILURE_STRINGS - // not threadsafe - static char invalid_chunk[] = "XXXX PNG chunk not known"; - invalid_chunk[0] = STBI__BYTECAST(c.type >> 24); - invalid_chunk[1] = STBI__BYTECAST(c.type >> 16); - invalid_chunk[2] = STBI__BYTECAST(c.type >> 8); - invalid_chunk[3] = STBI__BYTECAST(c.type >> 0); - #endif - return stbi__err(invalid_chunk, "PNG not supported: unknown PNG chunk type"); - } - stbi__skip(s, c.length); - break; - } - // end of PNG chunk, read and skip CRC - stbi__get32be(s); - } -} - -static void *stbi__do_png(stbi__png *p, int *x, int *y, int *n, int req_comp, stbi__result_info *ri) -{ - void *result=NULL; - if (req_comp < 0 || req_comp > 4) return stbi__errpuc("bad req_comp", "Internal error"); - if (stbi__parse_png_file(p, STBI__SCAN_load, req_comp)) { - if (p->depth <= 8) - ri->bits_per_channel = 8; - else if (p->depth == 16) - ri->bits_per_channel = 16; - else - return stbi__errpuc("bad bits_per_channel", "PNG not supported: unsupported color depth"); - result = p->out; - p->out = NULL; - if (req_comp && req_comp != p->s->img_out_n) { - if (ri->bits_per_channel == 8) - result = stbi__convert_format((unsigned char *) result, p->s->img_out_n, req_comp, p->s->img_x, p->s->img_y); - else - result = stbi__convert_format16((stbi__uint16 *) result, p->s->img_out_n, req_comp, p->s->img_x, p->s->img_y); - p->s->img_out_n = req_comp; - if (result == NULL) return result; - } - *x = p->s->img_x; - *y = p->s->img_y; - if (n) *n = p->s->img_n; - } - STBI_FREE(p->out); p->out = NULL; - STBI_FREE(p->expanded); p->expanded = NULL; - STBI_FREE(p->idata); p->idata = NULL; - - return result; -} - -static void *stbi__png_load(stbi__context *s, int *x, int *y, int *comp, int req_comp, stbi__result_info *ri) -{ - stbi__png p; - p.s = s; - return stbi__do_png(&p, x,y,comp,req_comp, ri); -} - -static int stbi__png_test(stbi__context *s) -{ - int r; - r = stbi__check_png_header(s); - stbi__rewind(s); - return r; -} - -static int stbi__png_info_raw(stbi__png *p, int *x, int *y, int *comp) -{ - if (!stbi__parse_png_file(p, STBI__SCAN_header, 0)) { - stbi__rewind( p->s ); - return 0; - } - if (x) *x = p->s->img_x; - if (y) *y = p->s->img_y; - if (comp) *comp = p->s->img_n; - return 1; -} - -static int stbi__png_info(stbi__context *s, int *x, int *y, int *comp) -{ - stbi__png p; - p.s = s; - return stbi__png_info_raw(&p, x, y, comp); -} - -static int stbi__png_is16(stbi__context *s) -{ - stbi__png p; - p.s = s; - if (!stbi__png_info_raw(&p, NULL, NULL, NULL)) - return 0; - if (p.depth != 16) { - stbi__rewind(p.s); - return 0; - } - return 1; -} -#endif - -// Microsoft/Windows BMP image - -#ifndef STBI_NO_BMP -static int stbi__bmp_test_raw(stbi__context *s) -{ - int r; - int sz; - if (stbi__get8(s) != 'B') return 0; - if (stbi__get8(s) != 'M') return 0; - stbi__get32le(s); // discard filesize - stbi__get16le(s); // discard reserved - stbi__get16le(s); // discard reserved - stbi__get32le(s); // discard data offset - sz = stbi__get32le(s); - r = (sz == 12 || sz == 40 || sz == 56 || sz == 108 || sz == 124); - return r; -} - -static int stbi__bmp_test(stbi__context *s) -{ - int r = stbi__bmp_test_raw(s); - stbi__rewind(s); - return r; -} - - -// returns 0..31 for the highest set bit -static int stbi__high_bit(unsigned int z) -{ - int n=0; - if (z == 0) return -1; - if (z >= 0x10000) { n += 16; z >>= 16; } - if (z >= 0x00100) { n += 8; z >>= 8; } - if (z >= 0x00010) { n += 4; z >>= 4; } - if (z >= 0x00004) { n += 2; z >>= 2; } - if (z >= 0x00002) { n += 1;/* >>= 1;*/ } - return n; -} - -static int stbi__bitcount(unsigned int a) -{ - a = (a & 0x55555555) + ((a >> 1) & 0x55555555); // max 2 - a = (a & 0x33333333) + ((a >> 2) & 0x33333333); // max 4 - a = (a + (a >> 4)) & 0x0f0f0f0f; // max 8 per 4, now 8 bits - a = (a + (a >> 8)); // max 16 per 8 bits - a = (a + (a >> 16)); // max 32 per 8 bits - return a & 0xff; -} - -// extract an arbitrarily-aligned N-bit value (N=bits) -// from v, and then make it 8-bits long and fractionally -// extend it to full full range. -static int stbi__shiftsigned(unsigned int v, int shift, int bits) -{ - static unsigned int mul_table[9] = { - 0, - 0xff/*0b11111111*/, 0x55/*0b01010101*/, 0x49/*0b01001001*/, 0x11/*0b00010001*/, - 0x21/*0b00100001*/, 0x41/*0b01000001*/, 0x81/*0b10000001*/, 0x01/*0b00000001*/, - }; - static unsigned int shift_table[9] = { - 0, 0,0,1,0,2,4,6,0, - }; - if (shift < 0) - v <<= -shift; - else - v >>= shift; - STBI_ASSERT(v < 256); - v >>= (8-bits); - STBI_ASSERT(bits >= 0 && bits <= 8); - return (int) ((unsigned) v * mul_table[bits]) >> shift_table[bits]; -} - -typedef struct -{ - int bpp, offset, hsz; - unsigned int mr,mg,mb,ma, all_a; - int extra_read; -} stbi__bmp_data; - -static int stbi__bmp_set_mask_defaults(stbi__bmp_data *info, int compress) -{ - // BI_BITFIELDS specifies masks explicitly, don't override - if (compress == 3) - return 1; - - if (compress == 0) { - if (info->bpp == 16) { - info->mr = 31u << 10; - info->mg = 31u << 5; - info->mb = 31u << 0; - } else if (info->bpp == 32) { - info->mr = 0xffu << 16; - info->mg = 0xffu << 8; - info->mb = 0xffu << 0; - info->ma = 0xffu << 24; - info->all_a = 0; // if all_a is 0 at end, then we loaded alpha channel but it was all 0 - } else { - // otherwise, use defaults, which is all-0 - info->mr = info->mg = info->mb = info->ma = 0; - } - return 1; - } - return 0; // error -} - -static void *stbi__bmp_parse_header(stbi__context *s, stbi__bmp_data *info) -{ - int hsz; - if (stbi__get8(s) != 'B' || stbi__get8(s) != 'M') return stbi__errpuc("not BMP", "Corrupt BMP"); - stbi__get32le(s); // discard filesize - stbi__get16le(s); // discard reserved - stbi__get16le(s); // discard reserved - info->offset = stbi__get32le(s); - info->hsz = hsz = stbi__get32le(s); - info->mr = info->mg = info->mb = info->ma = 0; - info->extra_read = 14; - - if (info->offset < 0) return stbi__errpuc("bad BMP", "bad BMP"); - - if (hsz != 12 && hsz != 40 && hsz != 56 && hsz != 108 && hsz != 124) return stbi__errpuc("unknown BMP", "BMP type not supported: unknown"); - if (hsz == 12) { - s->img_x = stbi__get16le(s); - s->img_y = stbi__get16le(s); - } else { - s->img_x = stbi__get32le(s); - s->img_y = stbi__get32le(s); - } - if (stbi__get16le(s) != 1) return stbi__errpuc("bad BMP", "bad BMP"); - info->bpp = stbi__get16le(s); - if (hsz != 12) { - int compress = stbi__get32le(s); - if (compress == 1 || compress == 2) return stbi__errpuc("BMP RLE", "BMP type not supported: RLE"); - if (compress >= 4) return stbi__errpuc("BMP JPEG/PNG", "BMP type not supported: unsupported compression"); // this includes PNG/JPEG modes - if (compress == 3 && info->bpp != 16 && info->bpp != 32) return stbi__errpuc("bad BMP", "bad BMP"); // bitfields requires 16 or 32 bits/pixel - stbi__get32le(s); // discard sizeof - stbi__get32le(s); // discard hres - stbi__get32le(s); // discard vres - stbi__get32le(s); // discard colorsused - stbi__get32le(s); // discard max important - if (hsz == 40 || hsz == 56) { - if (hsz == 56) { - stbi__get32le(s); - stbi__get32le(s); - stbi__get32le(s); - stbi__get32le(s); - } - if (info->bpp == 16 || info->bpp == 32) { - if (compress == 0) { - stbi__bmp_set_mask_defaults(info, compress); - } else if (compress == 3) { - info->mr = stbi__get32le(s); - info->mg = stbi__get32le(s); - info->mb = stbi__get32le(s); - info->extra_read += 12; - // not documented, but generated by photoshop and handled by mspaint - if (info->mr == info->mg && info->mg == info->mb) { - // ?!?!? - return stbi__errpuc("bad BMP", "bad BMP"); - } - } else - return stbi__errpuc("bad BMP", "bad BMP"); - } - } else { - // V4/V5 header - int i; - if (hsz != 108 && hsz != 124) - return stbi__errpuc("bad BMP", "bad BMP"); - info->mr = stbi__get32le(s); - info->mg = stbi__get32le(s); - info->mb = stbi__get32le(s); - info->ma = stbi__get32le(s); - if (compress != 3) // override mr/mg/mb unless in BI_BITFIELDS mode, as per docs - stbi__bmp_set_mask_defaults(info, compress); - stbi__get32le(s); // discard color space - for (i=0; i < 12; ++i) - stbi__get32le(s); // discard color space parameters - if (hsz == 124) { - stbi__get32le(s); // discard rendering intent - stbi__get32le(s); // discard offset of profile data - stbi__get32le(s); // discard size of profile data - stbi__get32le(s); // discard reserved - } - } - } - return (void *) 1; -} - - -static void *stbi__bmp_load(stbi__context *s, int *x, int *y, int *comp, int req_comp, stbi__result_info *ri) -{ - stbi_uc *out; - unsigned int mr=0,mg=0,mb=0,ma=0, all_a; - stbi_uc pal[256][4]; - int psize=0,i,j,width; - int flip_vertically, pad, target; - stbi__bmp_data info; - STBI_NOTUSED(ri); - - info.all_a = 255; - if (stbi__bmp_parse_header(s, &info) == NULL) - return NULL; // error code already set - - flip_vertically = ((int) s->img_y) > 0; - s->img_y = abs((int) s->img_y); - - if (s->img_y > STBI_MAX_DIMENSIONS) return stbi__errpuc("too large","Very large image (corrupt?)"); - if (s->img_x > STBI_MAX_DIMENSIONS) return stbi__errpuc("too large","Very large image (corrupt?)"); - - mr = info.mr; - mg = info.mg; - mb = info.mb; - ma = info.ma; - all_a = info.all_a; - - if (info.hsz == 12) { - if (info.bpp < 24) - psize = (info.offset - info.extra_read - 24) / 3; - } else { - if (info.bpp < 16) - psize = (info.offset - info.extra_read - info.hsz) >> 2; - } - if (psize == 0) { - // accept some number of extra bytes after the header, but if the offset points either to before - // the header ends or implies a large amount of extra data, reject the file as malformed - int bytes_read_so_far = s->callback_already_read + (int)(s->img_buffer - s->img_buffer_original); - int header_limit = 1024; // max we actually read is below 256 bytes currently. - int extra_data_limit = 256*4; // what ordinarily goes here is a palette; 256 entries*4 bytes is its max size. - if (bytes_read_so_far <= 0 || bytes_read_so_far > header_limit) { - return stbi__errpuc("bad header", "Corrupt BMP"); - } - // we established that bytes_read_so_far is positive and sensible. - // the first half of this test rejects offsets that are either too small positives, or - // negative, and guarantees that info.offset >= bytes_read_so_far > 0. this in turn - // ensures the number computed in the second half of the test can't overflow. - if (info.offset < bytes_read_so_far || info.offset - bytes_read_so_far > extra_data_limit) { - return stbi__errpuc("bad offset", "Corrupt BMP"); - } else { - stbi__skip(s, info.offset - bytes_read_so_far); - } - } - - if (info.bpp == 24 && ma == 0xff000000) - s->img_n = 3; - else - s->img_n = ma ? 4 : 3; - if (req_comp && req_comp >= 3) // we can directly decode 3 or 4 - target = req_comp; - else - target = s->img_n; // if they want monochrome, we'll post-convert - - // sanity-check size - if (!stbi__mad3sizes_valid(target, s->img_x, s->img_y, 0)) - return stbi__errpuc("too large", "Corrupt BMP"); - - out = (stbi_uc *) stbi__malloc_mad3(target, s->img_x, s->img_y, 0); - if (!out) return stbi__errpuc("outofmem", "Out of memory"); - if (info.bpp < 16) { - int z=0; - if (psize == 0 || psize > 256) { STBI_FREE(out); return stbi__errpuc("invalid", "Corrupt BMP"); } - for (i=0; i < psize; ++i) { - pal[i][2] = stbi__get8(s); - pal[i][1] = stbi__get8(s); - pal[i][0] = stbi__get8(s); - if (info.hsz != 12) stbi__get8(s); - pal[i][3] = 255; - } - stbi__skip(s, info.offset - info.extra_read - info.hsz - psize * (info.hsz == 12 ? 3 : 4)); - if (info.bpp == 1) width = (s->img_x + 7) >> 3; - else if (info.bpp == 4) width = (s->img_x + 1) >> 1; - else if (info.bpp == 8) width = s->img_x; - else { STBI_FREE(out); return stbi__errpuc("bad bpp", "Corrupt BMP"); } - pad = (-width)&3; - if (info.bpp == 1) { - for (j=0; j < (int) s->img_y; ++j) { - int bit_offset = 7, v = stbi__get8(s); - for (i=0; i < (int) s->img_x; ++i) { - int color = (v>>bit_offset)&0x1; - out[z++] = pal[color][0]; - out[z++] = pal[color][1]; - out[z++] = pal[color][2]; - if (target == 4) out[z++] = 255; - if (i+1 == (int) s->img_x) break; - if((--bit_offset) < 0) { - bit_offset = 7; - v = stbi__get8(s); - } - } - stbi__skip(s, pad); - } - } else { - for (j=0; j < (int) s->img_y; ++j) { - for (i=0; i < (int) s->img_x; i += 2) { - int v=stbi__get8(s),v2=0; - if (info.bpp == 4) { - v2 = v & 15; - v >>= 4; - } - out[z++] = pal[v][0]; - out[z++] = pal[v][1]; - out[z++] = pal[v][2]; - if (target == 4) out[z++] = 255; - if (i+1 == (int) s->img_x) break; - v = (info.bpp == 8) ? stbi__get8(s) : v2; - out[z++] = pal[v][0]; - out[z++] = pal[v][1]; - out[z++] = pal[v][2]; - if (target == 4) out[z++] = 255; - } - stbi__skip(s, pad); - } - } - } else { - int rshift=0,gshift=0,bshift=0,ashift=0,rcount=0,gcount=0,bcount=0,acount=0; - int z = 0; - int easy=0; - stbi__skip(s, info.offset - info.extra_read - info.hsz); - if (info.bpp == 24) width = 3 * s->img_x; - else if (info.bpp == 16) width = 2*s->img_x; - else /* bpp = 32 and pad = 0 */ width=0; - pad = (-width) & 3; - if (info.bpp == 24) { - easy = 1; - } else if (info.bpp == 32) { - if (mb == 0xff && mg == 0xff00 && mr == 0x00ff0000 && ma == 0xff000000) - easy = 2; - } - if (!easy) { - if (!mr || !mg || !mb) { STBI_FREE(out); return stbi__errpuc("bad masks", "Corrupt BMP"); } - // right shift amt to put high bit in position #7 - rshift = stbi__high_bit(mr)-7; rcount = stbi__bitcount(mr); - gshift = stbi__high_bit(mg)-7; gcount = stbi__bitcount(mg); - bshift = stbi__high_bit(mb)-7; bcount = stbi__bitcount(mb); - ashift = stbi__high_bit(ma)-7; acount = stbi__bitcount(ma); - if (rcount > 8 || gcount > 8 || bcount > 8 || acount > 8) { STBI_FREE(out); return stbi__errpuc("bad masks", "Corrupt BMP"); } - } - for (j=0; j < (int) s->img_y; ++j) { - if (easy) { - for (i=0; i < (int) s->img_x; ++i) { - unsigned char a; - out[z+2] = stbi__get8(s); - out[z+1] = stbi__get8(s); - out[z+0] = stbi__get8(s); - z += 3; - a = (easy == 2 ? stbi__get8(s) : 255); - all_a |= a; - if (target == 4) out[z++] = a; - } - } else { - int bpp = info.bpp; - for (i=0; i < (int) s->img_x; ++i) { - stbi__uint32 v = (bpp == 16 ? (stbi__uint32) stbi__get16le(s) : stbi__get32le(s)); - unsigned int a; - out[z++] = STBI__BYTECAST(stbi__shiftsigned(v & mr, rshift, rcount)); - out[z++] = STBI__BYTECAST(stbi__shiftsigned(v & mg, gshift, gcount)); - out[z++] = STBI__BYTECAST(stbi__shiftsigned(v & mb, bshift, bcount)); - a = (ma ? stbi__shiftsigned(v & ma, ashift, acount) : 255); - all_a |= a; - if (target == 4) out[z++] = STBI__BYTECAST(a); - } - } - stbi__skip(s, pad); - } - } - - // if alpha channel is all 0s, replace with all 255s - if (target == 4 && all_a == 0) - for (i=4*s->img_x*s->img_y-1; i >= 0; i -= 4) - out[i] = 255; - - if (flip_vertically) { - stbi_uc t; - for (j=0; j < (int) s->img_y>>1; ++j) { - stbi_uc *p1 = out + j *s->img_x*target; - stbi_uc *p2 = out + (s->img_y-1-j)*s->img_x*target; - for (i=0; i < (int) s->img_x*target; ++i) { - t = p1[i]; p1[i] = p2[i]; p2[i] = t; - } - } - } - - if (req_comp && req_comp != target) { - out = stbi__convert_format(out, target, req_comp, s->img_x, s->img_y); - if (out == NULL) return out; // stbi__convert_format frees input on failure - } - - *x = s->img_x; - *y = s->img_y; - if (comp) *comp = s->img_n; - return out; -} -#endif - -// Targa Truevision - TGA -// by Jonathan Dummer -#ifndef STBI_NO_TGA -// returns STBI_rgb or whatever, 0 on error -static int stbi__tga_get_comp(int bits_per_pixel, int is_grey, int* is_rgb16) -{ - // only RGB or RGBA (incl. 16bit) or grey allowed - if (is_rgb16) *is_rgb16 = 0; - switch(bits_per_pixel) { - case 8: return STBI_grey; - case 16: if(is_grey) return STBI_grey_alpha; - // fallthrough - case 15: if(is_rgb16) *is_rgb16 = 1; - return STBI_rgb; - case 24: // fallthrough - case 32: return bits_per_pixel/8; - default: return 0; - } -} - -static int stbi__tga_info(stbi__context *s, int *x, int *y, int *comp) -{ - int tga_w, tga_h, tga_comp, tga_image_type, tga_bits_per_pixel, tga_colormap_bpp; - int sz, tga_colormap_type; - stbi__get8(s); // discard Offset - tga_colormap_type = stbi__get8(s); // colormap type - if( tga_colormap_type > 1 ) { - stbi__rewind(s); - return 0; // only RGB or indexed allowed - } - tga_image_type = stbi__get8(s); // image type - if ( tga_colormap_type == 1 ) { // colormapped (paletted) image - if (tga_image_type != 1 && tga_image_type != 9) { - stbi__rewind(s); - return 0; - } - stbi__skip(s,4); // skip index of first colormap entry and number of entries - sz = stbi__get8(s); // check bits per palette color entry - if ( (sz != 8) && (sz != 15) && (sz != 16) && (sz != 24) && (sz != 32) ) { - stbi__rewind(s); - return 0; - } - stbi__skip(s,4); // skip image x and y origin - tga_colormap_bpp = sz; - } else { // "normal" image w/o colormap - only RGB or grey allowed, +/- RLE - if ( (tga_image_type != 2) && (tga_image_type != 3) && (tga_image_type != 10) && (tga_image_type != 11) ) { - stbi__rewind(s); - return 0; // only RGB or grey allowed, +/- RLE - } - stbi__skip(s,9); // skip colormap specification and image x/y origin - tga_colormap_bpp = 0; - } - tga_w = stbi__get16le(s); - if( tga_w < 1 ) { - stbi__rewind(s); - return 0; // test width - } - tga_h = stbi__get16le(s); - if( tga_h < 1 ) { - stbi__rewind(s); - return 0; // test height - } - tga_bits_per_pixel = stbi__get8(s); // bits per pixel - stbi__get8(s); // ignore alpha bits - if (tga_colormap_bpp != 0) { - if((tga_bits_per_pixel != 8) && (tga_bits_per_pixel != 16)) { - // when using a colormap, tga_bits_per_pixel is the size of the indexes - // I don't think anything but 8 or 16bit indexes makes sense - stbi__rewind(s); - return 0; - } - tga_comp = stbi__tga_get_comp(tga_colormap_bpp, 0, NULL); - } else { - tga_comp = stbi__tga_get_comp(tga_bits_per_pixel, (tga_image_type == 3) || (tga_image_type == 11), NULL); - } - if(!tga_comp) { - stbi__rewind(s); - return 0; - } - if (x) *x = tga_w; - if (y) *y = tga_h; - if (comp) *comp = tga_comp; - return 1; // seems to have passed everything -} - -static int stbi__tga_test(stbi__context *s) -{ - int res = 0; - int sz, tga_color_type; - stbi__get8(s); // discard Offset - tga_color_type = stbi__get8(s); // color type - if ( tga_color_type > 1 ) goto errorEnd; // only RGB or indexed allowed - sz = stbi__get8(s); // image type - if ( tga_color_type == 1 ) { // colormapped (paletted) image - if (sz != 1 && sz != 9) goto errorEnd; // colortype 1 demands image type 1 or 9 - stbi__skip(s,4); // skip index of first colormap entry and number of entries - sz = stbi__get8(s); // check bits per palette color entry - if ( (sz != 8) && (sz != 15) && (sz != 16) && (sz != 24) && (sz != 32) ) goto errorEnd; - stbi__skip(s,4); // skip image x and y origin - } else { // "normal" image w/o colormap - if ( (sz != 2) && (sz != 3) && (sz != 10) && (sz != 11) ) goto errorEnd; // only RGB or grey allowed, +/- RLE - stbi__skip(s,9); // skip colormap specification and image x/y origin - } - if ( stbi__get16le(s) < 1 ) goto errorEnd; // test width - if ( stbi__get16le(s) < 1 ) goto errorEnd; // test height - sz = stbi__get8(s); // bits per pixel - if ( (tga_color_type == 1) && (sz != 8) && (sz != 16) ) goto errorEnd; // for colormapped images, bpp is size of an index - if ( (sz != 8) && (sz != 15) && (sz != 16) && (sz != 24) && (sz != 32) ) goto errorEnd; - - res = 1; // if we got this far, everything's good and we can return 1 instead of 0 - -errorEnd: - stbi__rewind(s); - return res; -} - -// read 16bit value and convert to 24bit RGB -static void stbi__tga_read_rgb16(stbi__context *s, stbi_uc* out) -{ - stbi__uint16 px = (stbi__uint16)stbi__get16le(s); - stbi__uint16 fiveBitMask = 31; - // we have 3 channels with 5bits each - int r = (px >> 10) & fiveBitMask; - int g = (px >> 5) & fiveBitMask; - int b = px & fiveBitMask; - // Note that this saves the data in RGB(A) order, so it doesn't need to be swapped later - out[0] = (stbi_uc)((r * 255)/31); - out[1] = (stbi_uc)((g * 255)/31); - out[2] = (stbi_uc)((b * 255)/31); - - // some people claim that the most significant bit might be used for alpha - // (possibly if an alpha-bit is set in the "image descriptor byte") - // but that only made 16bit test images completely translucent.. - // so let's treat all 15 and 16bit TGAs as RGB with no alpha. -} - -static void *stbi__tga_load(stbi__context *s, int *x, int *y, int *comp, int req_comp, stbi__result_info *ri) -{ - // read in the TGA header stuff - int tga_offset = stbi__get8(s); - int tga_indexed = stbi__get8(s); - int tga_image_type = stbi__get8(s); - int tga_is_RLE = 0; - int tga_palette_start = stbi__get16le(s); - int tga_palette_len = stbi__get16le(s); - int tga_palette_bits = stbi__get8(s); - int tga_x_origin = stbi__get16le(s); - int tga_y_origin = stbi__get16le(s); - int tga_width = stbi__get16le(s); - int tga_height = stbi__get16le(s); - int tga_bits_per_pixel = stbi__get8(s); - int tga_comp, tga_rgb16=0; - int tga_inverted = stbi__get8(s); - // int tga_alpha_bits = tga_inverted & 15; // the 4 lowest bits - unused (useless?) - // image data - unsigned char *tga_data; - unsigned char *tga_palette = NULL; - int i, j; - unsigned char raw_data[4] = {0}; - int RLE_count = 0; - int RLE_repeating = 0; - int read_next_pixel = 1; - STBI_NOTUSED(ri); - STBI_NOTUSED(tga_x_origin); // @TODO - STBI_NOTUSED(tga_y_origin); // @TODO - - if (tga_height > STBI_MAX_DIMENSIONS) return stbi__errpuc("too large","Very large image (corrupt?)"); - if (tga_width > STBI_MAX_DIMENSIONS) return stbi__errpuc("too large","Very large image (corrupt?)"); - - // do a tiny bit of precessing - if ( tga_image_type >= 8 ) - { - tga_image_type -= 8; - tga_is_RLE = 1; - } - tga_inverted = 1 - ((tga_inverted >> 5) & 1); - - // If I'm paletted, then I'll use the number of bits from the palette - if ( tga_indexed ) tga_comp = stbi__tga_get_comp(tga_palette_bits, 0, &tga_rgb16); - else tga_comp = stbi__tga_get_comp(tga_bits_per_pixel, (tga_image_type == 3), &tga_rgb16); - - if(!tga_comp) // shouldn't really happen, stbi__tga_test() should have ensured basic consistency - return stbi__errpuc("bad format", "Can't find out TGA pixelformat"); - - // tga info - *x = tga_width; - *y = tga_height; - if (comp) *comp = tga_comp; - - if (!stbi__mad3sizes_valid(tga_width, tga_height, tga_comp, 0)) - return stbi__errpuc("too large", "Corrupt TGA"); - - tga_data = (unsigned char*)stbi__malloc_mad3(tga_width, tga_height, tga_comp, 0); - if (!tga_data) return stbi__errpuc("outofmem", "Out of memory"); - - // skip to the data's starting position (offset usually = 0) - stbi__skip(s, tga_offset ); - - if ( !tga_indexed && !tga_is_RLE && !tga_rgb16 ) { - for (i=0; i < tga_height; ++i) { - int row = tga_inverted ? tga_height -i - 1 : i; - stbi_uc *tga_row = tga_data + row*tga_width*tga_comp; - stbi__getn(s, tga_row, tga_width * tga_comp); - } - } else { - // do I need to load a palette? - if ( tga_indexed) - { - if (tga_palette_len == 0) { /* you have to have at least one entry! */ - STBI_FREE(tga_data); - return stbi__errpuc("bad palette", "Corrupt TGA"); - } - - // any data to skip? (offset usually = 0) - stbi__skip(s, tga_palette_start ); - // load the palette - tga_palette = (unsigned char*)stbi__malloc_mad2(tga_palette_len, tga_comp, 0); - if (!tga_palette) { - STBI_FREE(tga_data); - return stbi__errpuc("outofmem", "Out of memory"); - } - if (tga_rgb16) { - stbi_uc *pal_entry = tga_palette; - STBI_ASSERT(tga_comp == STBI_rgb); - for (i=0; i < tga_palette_len; ++i) { - stbi__tga_read_rgb16(s, pal_entry); - pal_entry += tga_comp; - } - } else if (!stbi__getn(s, tga_palette, tga_palette_len * tga_comp)) { - STBI_FREE(tga_data); - STBI_FREE(tga_palette); - return stbi__errpuc("bad palette", "Corrupt TGA"); - } - } - // load the data - for (i=0; i < tga_width * tga_height; ++i) - { - // if I'm in RLE mode, do I need to get a RLE stbi__pngchunk? - if ( tga_is_RLE ) - { - if ( RLE_count == 0 ) - { - // yep, get the next byte as a RLE command - int RLE_cmd = stbi__get8(s); - RLE_count = 1 + (RLE_cmd & 127); - RLE_repeating = RLE_cmd >> 7; - read_next_pixel = 1; - } else if ( !RLE_repeating ) - { - read_next_pixel = 1; - } - } else - { - read_next_pixel = 1; - } - // OK, if I need to read a pixel, do it now - if ( read_next_pixel ) - { - // load however much data we did have - if ( tga_indexed ) - { - // read in index, then perform the lookup - int pal_idx = (tga_bits_per_pixel == 8) ? stbi__get8(s) : stbi__get16le(s); - if ( pal_idx >= tga_palette_len ) { - // invalid index - pal_idx = 0; - } - pal_idx *= tga_comp; - for (j = 0; j < tga_comp; ++j) { - raw_data[j] = tga_palette[pal_idx+j]; - } - } else if(tga_rgb16) { - STBI_ASSERT(tga_comp == STBI_rgb); - stbi__tga_read_rgb16(s, raw_data); - } else { - // read in the data raw - for (j = 0; j < tga_comp; ++j) { - raw_data[j] = stbi__get8(s); - } - } - // clear the reading flag for the next pixel - read_next_pixel = 0; - } // end of reading a pixel - - // copy data - for (j = 0; j < tga_comp; ++j) - tga_data[i*tga_comp+j] = raw_data[j]; - - // in case we're in RLE mode, keep counting down - --RLE_count; - } - // do I need to invert the image? - if ( tga_inverted ) - { - for (j = 0; j*2 < tga_height; ++j) - { - int index1 = j * tga_width * tga_comp; - int index2 = (tga_height - 1 - j) * tga_width * tga_comp; - for (i = tga_width * tga_comp; i > 0; --i) - { - unsigned char temp = tga_data[index1]; - tga_data[index1] = tga_data[index2]; - tga_data[index2] = temp; - ++index1; - ++index2; - } - } - } - // clear my palette, if I had one - if ( tga_palette != NULL ) - { - STBI_FREE( tga_palette ); - } - } - - // swap RGB - if the source data was RGB16, it already is in the right order - if (tga_comp >= 3 && !tga_rgb16) - { - unsigned char* tga_pixel = tga_data; - for (i=0; i < tga_width * tga_height; ++i) - { - unsigned char temp = tga_pixel[0]; - tga_pixel[0] = tga_pixel[2]; - tga_pixel[2] = temp; - tga_pixel += tga_comp; - } - } - - // convert to target component count - if (req_comp && req_comp != tga_comp) - tga_data = stbi__convert_format(tga_data, tga_comp, req_comp, tga_width, tga_height); - - // the things I do to get rid of an error message, and yet keep - // Microsoft's C compilers happy... [8^( - tga_palette_start = tga_palette_len = tga_palette_bits = - tga_x_origin = tga_y_origin = 0; - STBI_NOTUSED(tga_palette_start); - // OK, done - return tga_data; -} -#endif - -// ************************************************************************************************* -// Photoshop PSD loader -- PD by Thatcher Ulrich, integration by Nicolas Schulz, tweaked by STB - -#ifndef STBI_NO_PSD -static int stbi__psd_test(stbi__context *s) -{ - int r = (stbi__get32be(s) == 0x38425053); - stbi__rewind(s); - return r; -} - -static int stbi__psd_decode_rle(stbi__context *s, stbi_uc *p, int pixelCount) -{ - int count, nleft, len; - - count = 0; - while ((nleft = pixelCount - count) > 0) { - len = stbi__get8(s); - if (len == 128) { - // No-op. - } else if (len < 128) { - // Copy next len+1 bytes literally. - len++; - if (len > nleft) return 0; // corrupt data - count += len; - while (len) { - *p = stbi__get8(s); - p += 4; - len--; - } - } else if (len > 128) { - stbi_uc val; - // Next -len+1 bytes in the dest are replicated from next source byte. - // (Interpret len as a negative 8-bit int.) - len = 257 - len; - if (len > nleft) return 0; // corrupt data - val = stbi__get8(s); - count += len; - while (len) { - *p = val; - p += 4; - len--; - } - } - } - - return 1; -} - -static void *stbi__psd_load(stbi__context *s, int *x, int *y, int *comp, int req_comp, stbi__result_info *ri, int bpc) -{ - int pixelCount; - int channelCount, compression; - int channel, i; - int bitdepth; - int w,h; - stbi_uc *out; - STBI_NOTUSED(ri); - - // Check identifier - if (stbi__get32be(s) != 0x38425053) // "8BPS" - return stbi__errpuc("not PSD", "Corrupt PSD image"); - - // Check file type version. - if (stbi__get16be(s) != 1) - return stbi__errpuc("wrong version", "Unsupported version of PSD image"); - - // Skip 6 reserved bytes. - stbi__skip(s, 6 ); - - // Read the number of channels (R, G, B, A, etc). - channelCount = stbi__get16be(s); - if (channelCount < 0 || channelCount > 16) - return stbi__errpuc("wrong channel count", "Unsupported number of channels in PSD image"); - - // Read the rows and columns of the image. - h = stbi__get32be(s); - w = stbi__get32be(s); - - if (h > STBI_MAX_DIMENSIONS) return stbi__errpuc("too large","Very large image (corrupt?)"); - if (w > STBI_MAX_DIMENSIONS) return stbi__errpuc("too large","Very large image (corrupt?)"); - - // Make sure the depth is 8 bits. - bitdepth = stbi__get16be(s); - if (bitdepth != 8 && bitdepth != 16) - return stbi__errpuc("unsupported bit depth", "PSD bit depth is not 8 or 16 bit"); - - // Make sure the color mode is RGB. - // Valid options are: - // 0: Bitmap - // 1: Grayscale - // 2: Indexed color - // 3: RGB color - // 4: CMYK color - // 7: Multichannel - // 8: Duotone - // 9: Lab color - if (stbi__get16be(s) != 3) - return stbi__errpuc("wrong color format", "PSD is not in RGB color format"); - - // Skip the Mode Data. (It's the palette for indexed color; other info for other modes.) - stbi__skip(s,stbi__get32be(s) ); - - // Skip the image resources. (resolution, pen tool paths, etc) - stbi__skip(s, stbi__get32be(s) ); - - // Skip the reserved data. - stbi__skip(s, stbi__get32be(s) ); - - // Find out if the data is compressed. - // Known values: - // 0: no compression - // 1: RLE compressed - compression = stbi__get16be(s); - if (compression > 1) - return stbi__errpuc("bad compression", "PSD has an unknown compression format"); - - // Check size - if (!stbi__mad3sizes_valid(4, w, h, 0)) - return stbi__errpuc("too large", "Corrupt PSD"); - - // Create the destination image. - - if (!compression && bitdepth == 16 && bpc == 16) { - out = (stbi_uc *) stbi__malloc_mad3(8, w, h, 0); - ri->bits_per_channel = 16; - } else - out = (stbi_uc *) stbi__malloc(4 * w*h); - - if (!out) return stbi__errpuc("outofmem", "Out of memory"); - pixelCount = w*h; - - // Initialize the data to zero. - //memset( out, 0, pixelCount * 4 ); - - // Finally, the image data. - if (compression) { - // RLE as used by .PSD and .TIFF - // Loop until you get the number of unpacked bytes you are expecting: - // Read the next source byte into n. - // If n is between 0 and 127 inclusive, copy the next n+1 bytes literally. - // Else if n is between -127 and -1 inclusive, copy the next byte -n+1 times. - // Else if n is 128, noop. - // Endloop - - // The RLE-compressed data is preceded by a 2-byte data count for each row in the data, - // which we're going to just skip. - stbi__skip(s, h * channelCount * 2 ); - - // Read the RLE data by channel. - for (channel = 0; channel < 4; channel++) { - stbi_uc *p; - - p = out+channel; - if (channel >= channelCount) { - // Fill this channel with default data. - for (i = 0; i < pixelCount; i++, p += 4) - *p = (channel == 3 ? 255 : 0); - } else { - // Read the RLE data. - if (!stbi__psd_decode_rle(s, p, pixelCount)) { - STBI_FREE(out); - return stbi__errpuc("corrupt", "bad RLE data"); - } - } - } - - } else { - // We're at the raw image data. It's each channel in order (Red, Green, Blue, Alpha, ...) - // where each channel consists of an 8-bit (or 16-bit) value for each pixel in the image. - - // Read the data by channel. - for (channel = 0; channel < 4; channel++) { - if (channel >= channelCount) { - // Fill this channel with default data. - if (bitdepth == 16 && bpc == 16) { - stbi__uint16 *q = ((stbi__uint16 *) out) + channel; - stbi__uint16 val = channel == 3 ? 65535 : 0; - for (i = 0; i < pixelCount; i++, q += 4) - *q = val; - } else { - stbi_uc *p = out+channel; - stbi_uc val = channel == 3 ? 255 : 0; - for (i = 0; i < pixelCount; i++, p += 4) - *p = val; - } - } else { - if (ri->bits_per_channel == 16) { // output bpc - stbi__uint16 *q = ((stbi__uint16 *) out) + channel; - for (i = 0; i < pixelCount; i++, q += 4) - *q = (stbi__uint16) stbi__get16be(s); - } else { - stbi_uc *p = out+channel; - if (bitdepth == 16) { // input bpc - for (i = 0; i < pixelCount; i++, p += 4) - *p = (stbi_uc) (stbi__get16be(s) >> 8); - } else { - for (i = 0; i < pixelCount; i++, p += 4) - *p = stbi__get8(s); - } - } - } - } - } - - // remove weird white matte from PSD - if (channelCount >= 4) { - if (ri->bits_per_channel == 16) { - for (i=0; i < w*h; ++i) { - stbi__uint16 *pixel = (stbi__uint16 *) out + 4*i; - if (pixel[3] != 0 && pixel[3] != 65535) { - float a = pixel[3] / 65535.0f; - float ra = 1.0f / a; - float inv_a = 65535.0f * (1 - ra); - pixel[0] = (stbi__uint16) (pixel[0]*ra + inv_a); - pixel[1] = (stbi__uint16) (pixel[1]*ra + inv_a); - pixel[2] = (stbi__uint16) (pixel[2]*ra + inv_a); - } - } - } else { - for (i=0; i < w*h; ++i) { - unsigned char *pixel = out + 4*i; - if (pixel[3] != 0 && pixel[3] != 255) { - float a = pixel[3] / 255.0f; - float ra = 1.0f / a; - float inv_a = 255.0f * (1 - ra); - pixel[0] = (unsigned char) (pixel[0]*ra + inv_a); - pixel[1] = (unsigned char) (pixel[1]*ra + inv_a); - pixel[2] = (unsigned char) (pixel[2]*ra + inv_a); - } - } - } - } - - // convert to desired output format - if (req_comp && req_comp != 4) { - if (ri->bits_per_channel == 16) - out = (stbi_uc *) stbi__convert_format16((stbi__uint16 *) out, 4, req_comp, w, h); - else - out = stbi__convert_format(out, 4, req_comp, w, h); - if (out == NULL) return out; // stbi__convert_format frees input on failure - } - - if (comp) *comp = 4; - *y = h; - *x = w; - - return out; -} -#endif - -// ************************************************************************************************* -// Softimage PIC loader -// by Tom Seddon -// -// See http://softimage.wiki.softimage.com/index.php/INFO:_PIC_file_format -// See http://ozviz.wasp.uwa.edu.au/~pbourke/dataformats/softimagepic/ - -#ifndef STBI_NO_PIC -static int stbi__pic_is4(stbi__context *s,const char *str) -{ - int i; - for (i=0; i<4; ++i) - if (stbi__get8(s) != (stbi_uc)str[i]) - return 0; - - return 1; -} - -static int stbi__pic_test_core(stbi__context *s) -{ - int i; - - if (!stbi__pic_is4(s,"\x53\x80\xF6\x34")) - return 0; - - for(i=0;i<84;++i) - stbi__get8(s); - - if (!stbi__pic_is4(s,"PICT")) - return 0; - - return 1; -} - -typedef struct -{ - stbi_uc size,type,channel; -} stbi__pic_packet; - -static stbi_uc *stbi__readval(stbi__context *s, int channel, stbi_uc *dest) -{ - int mask=0x80, i; - - for (i=0; i<4; ++i, mask>>=1) { - if (channel & mask) { - if (stbi__at_eof(s)) return stbi__errpuc("bad file","PIC file too short"); - dest[i]=stbi__get8(s); - } - } - - return dest; -} - -static void stbi__copyval(int channel,stbi_uc *dest,const stbi_uc *src) -{ - int mask=0x80,i; - - for (i=0;i<4; ++i, mask>>=1) - if (channel&mask) - dest[i]=src[i]; -} - -static stbi_uc *stbi__pic_load_core(stbi__context *s,int width,int height,int *comp, stbi_uc *result) -{ - int act_comp=0,num_packets=0,y,chained; - stbi__pic_packet packets[10]; - - // this will (should...) cater for even some bizarre stuff like having data - // for the same channel in multiple packets. - do { - stbi__pic_packet *packet; - - if (num_packets==sizeof(packets)/sizeof(packets[0])) - return stbi__errpuc("bad format","too many packets"); - - packet = &packets[num_packets++]; - - chained = stbi__get8(s); - packet->size = stbi__get8(s); - packet->type = stbi__get8(s); - packet->channel = stbi__get8(s); - - act_comp |= packet->channel; - - if (stbi__at_eof(s)) return stbi__errpuc("bad file","file too short (reading packets)"); - if (packet->size != 8) return stbi__errpuc("bad format","packet isn't 8bpp"); - } while (chained); - - *comp = (act_comp & 0x10 ? 4 : 3); // has alpha channel? - - for(y=0; ytype) { - default: - return stbi__errpuc("bad format","packet has bad compression type"); - - case 0: {//uncompressed - int x; - - for(x=0;xchannel,dest)) - return 0; - break; - } - - case 1://Pure RLE - { - int left=width, i; - - while (left>0) { - stbi_uc count,value[4]; - - count=stbi__get8(s); - if (stbi__at_eof(s)) return stbi__errpuc("bad file","file too short (pure read count)"); - - if (count > left) - count = (stbi_uc) left; - - if (!stbi__readval(s,packet->channel,value)) return 0; - - for(i=0; ichannel,dest,value); - left -= count; - } - } - break; - - case 2: {//Mixed RLE - int left=width; - while (left>0) { - int count = stbi__get8(s), i; - if (stbi__at_eof(s)) return stbi__errpuc("bad file","file too short (mixed read count)"); - - if (count >= 128) { // Repeated - stbi_uc value[4]; - - if (count==128) - count = stbi__get16be(s); - else - count -= 127; - if (count > left) - return stbi__errpuc("bad file","scanline overrun"); - - if (!stbi__readval(s,packet->channel,value)) - return 0; - - for(i=0;ichannel,dest,value); - } else { // Raw - ++count; - if (count>left) return stbi__errpuc("bad file","scanline overrun"); - - for(i=0;ichannel,dest)) - return 0; - } - left-=count; - } - break; - } - } - } - } - - return result; -} - -static void *stbi__pic_load(stbi__context *s,int *px,int *py,int *comp,int req_comp, stbi__result_info *ri) -{ - stbi_uc *result; - int i, x,y, internal_comp; - STBI_NOTUSED(ri); - - if (!comp) comp = &internal_comp; - - for (i=0; i<92; ++i) - stbi__get8(s); - - x = stbi__get16be(s); - y = stbi__get16be(s); - - if (y > STBI_MAX_DIMENSIONS) return stbi__errpuc("too large","Very large image (corrupt?)"); - if (x > STBI_MAX_DIMENSIONS) return stbi__errpuc("too large","Very large image (corrupt?)"); - - if (stbi__at_eof(s)) return stbi__errpuc("bad file","file too short (pic header)"); - if (!stbi__mad3sizes_valid(x, y, 4, 0)) return stbi__errpuc("too large", "PIC image too large to decode"); - - stbi__get32be(s); //skip `ratio' - stbi__get16be(s); //skip `fields' - stbi__get16be(s); //skip `pad' - - // intermediate buffer is RGBA - result = (stbi_uc *) stbi__malloc_mad3(x, y, 4, 0); - if (!result) return stbi__errpuc("outofmem", "Out of memory"); - memset(result, 0xff, x*y*4); - - if (!stbi__pic_load_core(s,x,y,comp, result)) { - STBI_FREE(result); - result=0; - } - *px = x; - *py = y; - if (req_comp == 0) req_comp = *comp; - result=stbi__convert_format(result,4,req_comp,x,y); - - return result; -} - -static int stbi__pic_test(stbi__context *s) -{ - int r = stbi__pic_test_core(s); - stbi__rewind(s); - return r; -} -#endif - -// ************************************************************************************************* -// GIF loader -- public domain by Jean-Marc Lienher -- simplified/shrunk by stb - -#ifndef STBI_NO_GIF -typedef struct -{ - stbi__int16 prefix; - stbi_uc first; - stbi_uc suffix; -} stbi__gif_lzw; - -typedef struct -{ - int w,h; - stbi_uc *out; // output buffer (always 4 components) - stbi_uc *background; // The current "background" as far as a gif is concerned - stbi_uc *history; - int flags, bgindex, ratio, transparent, eflags; - stbi_uc pal[256][4]; - stbi_uc lpal[256][4]; - stbi__gif_lzw codes[8192]; - stbi_uc *color_table; - int parse, step; - int lflags; - int start_x, start_y; - int max_x, max_y; - int cur_x, cur_y; - int line_size; - int delay; -} stbi__gif; - -static int stbi__gif_test_raw(stbi__context *s) -{ - int sz; - if (stbi__get8(s) != 'G' || stbi__get8(s) != 'I' || stbi__get8(s) != 'F' || stbi__get8(s) != '8') return 0; - sz = stbi__get8(s); - if (sz != '9' && sz != '7') return 0; - if (stbi__get8(s) != 'a') return 0; - return 1; -} - -static int stbi__gif_test(stbi__context *s) -{ - int r = stbi__gif_test_raw(s); - stbi__rewind(s); - return r; -} - -static void stbi__gif_parse_colortable(stbi__context *s, stbi_uc pal[256][4], int num_entries, int transp) -{ - int i; - for (i=0; i < num_entries; ++i) { - pal[i][2] = stbi__get8(s); - pal[i][1] = stbi__get8(s); - pal[i][0] = stbi__get8(s); - pal[i][3] = transp == i ? 0 : 255; - } -} - -static int stbi__gif_header(stbi__context *s, stbi__gif *g, int *comp, int is_info) -{ - stbi_uc version; - if (stbi__get8(s) != 'G' || stbi__get8(s) != 'I' || stbi__get8(s) != 'F' || stbi__get8(s) != '8') - return stbi__err("not GIF", "Corrupt GIF"); - - version = stbi__get8(s); - if (version != '7' && version != '9') return stbi__err("not GIF", "Corrupt GIF"); - if (stbi__get8(s) != 'a') return stbi__err("not GIF", "Corrupt GIF"); - - stbi__g_failure_reason = ""; - g->w = stbi__get16le(s); - g->h = stbi__get16le(s); - g->flags = stbi__get8(s); - g->bgindex = stbi__get8(s); - g->ratio = stbi__get8(s); - g->transparent = -1; - - if (g->w > STBI_MAX_DIMENSIONS) return stbi__err("too large","Very large image (corrupt?)"); - if (g->h > STBI_MAX_DIMENSIONS) return stbi__err("too large","Very large image (corrupt?)"); - - if (comp != 0) *comp = 4; // can't actually tell whether it's 3 or 4 until we parse the comments - - if (is_info) return 1; - - if (g->flags & 0x80) - stbi__gif_parse_colortable(s,g->pal, 2 << (g->flags & 7), -1); - - return 1; -} - -static int stbi__gif_info_raw(stbi__context *s, int *x, int *y, int *comp) -{ - stbi__gif* g = (stbi__gif*) stbi__malloc(sizeof(stbi__gif)); - if (!g) return stbi__err("outofmem", "Out of memory"); - if (!stbi__gif_header(s, g, comp, 1)) { - STBI_FREE(g); - stbi__rewind( s ); - return 0; - } - if (x) *x = g->w; - if (y) *y = g->h; - STBI_FREE(g); - return 1; -} - -static void stbi__out_gif_code(stbi__gif *g, stbi__uint16 code) -{ - stbi_uc *p, *c; - int idx; - - // recurse to decode the prefixes, since the linked-list is backwards, - // and working backwards through an interleaved image would be nasty - if (g->codes[code].prefix >= 0) - stbi__out_gif_code(g, g->codes[code].prefix); - - if (g->cur_y >= g->max_y) return; - - idx = g->cur_x + g->cur_y; - p = &g->out[idx]; - g->history[idx / 4] = 1; - - c = &g->color_table[g->codes[code].suffix * 4]; - if (c[3] > 128) { // don't render transparent pixels; - p[0] = c[2]; - p[1] = c[1]; - p[2] = c[0]; - p[3] = c[3]; - } - g->cur_x += 4; - - if (g->cur_x >= g->max_x) { - g->cur_x = g->start_x; - g->cur_y += g->step; - - while (g->cur_y >= g->max_y && g->parse > 0) { - g->step = (1 << g->parse) * g->line_size; - g->cur_y = g->start_y + (g->step >> 1); - --g->parse; - } - } -} - -static stbi_uc *stbi__process_gif_raster(stbi__context *s, stbi__gif *g) -{ - stbi_uc lzw_cs; - stbi__int32 len, init_code; - stbi__uint32 first; - stbi__int32 codesize, codemask, avail, oldcode, bits, valid_bits, clear; - stbi__gif_lzw *p; - - lzw_cs = stbi__get8(s); - if (lzw_cs > 12) return NULL; - clear = 1 << lzw_cs; - first = 1; - codesize = lzw_cs + 1; - codemask = (1 << codesize) - 1; - bits = 0; - valid_bits = 0; - for (init_code = 0; init_code < clear; init_code++) { - g->codes[init_code].prefix = -1; - g->codes[init_code].first = (stbi_uc) init_code; - g->codes[init_code].suffix = (stbi_uc) init_code; - } - - // support no starting clear code - avail = clear+2; - oldcode = -1; - - len = 0; - for(;;) { - if (valid_bits < codesize) { - if (len == 0) { - len = stbi__get8(s); // start new block - if (len == 0) - return g->out; - } - --len; - bits |= (stbi__int32) stbi__get8(s) << valid_bits; - valid_bits += 8; - } else { - stbi__int32 code = bits & codemask; - bits >>= codesize; - valid_bits -= codesize; - // @OPTIMIZE: is there some way we can accelerate the non-clear path? - if (code == clear) { // clear code - codesize = lzw_cs + 1; - codemask = (1 << codesize) - 1; - avail = clear + 2; - oldcode = -1; - first = 0; - } else if (code == clear + 1) { // end of stream code - stbi__skip(s, len); - while ((len = stbi__get8(s)) > 0) - stbi__skip(s,len); - return g->out; - } else if (code <= avail) { - if (first) { - return stbi__errpuc("no clear code", "Corrupt GIF"); - } - - if (oldcode >= 0) { - p = &g->codes[avail++]; - if (avail > 8192) { - return stbi__errpuc("too many codes", "Corrupt GIF"); - } - - p->prefix = (stbi__int16) oldcode; - p->first = g->codes[oldcode].first; - p->suffix = (code == avail) ? p->first : g->codes[code].first; - } else if (code == avail) - return stbi__errpuc("illegal code in raster", "Corrupt GIF"); - - stbi__out_gif_code(g, (stbi__uint16) code); - - if ((avail & codemask) == 0 && avail <= 0x0FFF) { - codesize++; - codemask = (1 << codesize) - 1; - } - - oldcode = code; - } else { - return stbi__errpuc("illegal code in raster", "Corrupt GIF"); - } - } - } -} - -// this function is designed to support animated gifs, although stb_image doesn't support it -// two back is the image from two frames ago, used for a very specific disposal format -static stbi_uc *stbi__gif_load_next(stbi__context *s, stbi__gif *g, int *comp, int req_comp, stbi_uc *two_back) -{ - int dispose; - int first_frame; - int pi; - int pcount; - STBI_NOTUSED(req_comp); - - // on first frame, any non-written pixels get the background colour (non-transparent) - first_frame = 0; - if (g->out == 0) { - if (!stbi__gif_header(s, g, comp,0)) return 0; // stbi__g_failure_reason set by stbi__gif_header - if (!stbi__mad3sizes_valid(4, g->w, g->h, 0)) - return stbi__errpuc("too large", "GIF image is too large"); - pcount = g->w * g->h; - g->out = (stbi_uc *) stbi__malloc(4 * pcount); - g->background = (stbi_uc *) stbi__malloc(4 * pcount); - g->history = (stbi_uc *) stbi__malloc(pcount); - if (!g->out || !g->background || !g->history) - return stbi__errpuc("outofmem", "Out of memory"); - - // image is treated as "transparent" at the start - ie, nothing overwrites the current background; - // background colour is only used for pixels that are not rendered first frame, after that "background" - // color refers to the color that was there the previous frame. - memset(g->out, 0x00, 4 * pcount); - memset(g->background, 0x00, 4 * pcount); // state of the background (starts transparent) - memset(g->history, 0x00, pcount); // pixels that were affected previous frame - first_frame = 1; - } else { - // second frame - how do we dispose of the previous one? - dispose = (g->eflags & 0x1C) >> 2; - pcount = g->w * g->h; - - if ((dispose == 3) && (two_back == 0)) { - dispose = 2; // if I don't have an image to revert back to, default to the old background - } - - if (dispose == 3) { // use previous graphic - for (pi = 0; pi < pcount; ++pi) { - if (g->history[pi]) { - memcpy( &g->out[pi * 4], &two_back[pi * 4], 4 ); - } - } - } else if (dispose == 2) { - // restore what was changed last frame to background before that frame; - for (pi = 0; pi < pcount; ++pi) { - if (g->history[pi]) { - memcpy( &g->out[pi * 4], &g->background[pi * 4], 4 ); - } - } - } else { - // This is a non-disposal case eithe way, so just - // leave the pixels as is, and they will become the new background - // 1: do not dispose - // 0: not specified. - } - - // background is what out is after the undoing of the previou frame; - memcpy( g->background, g->out, 4 * g->w * g->h ); - } - - // clear my history; - memset( g->history, 0x00, g->w * g->h ); // pixels that were affected previous frame - - for (;;) { - int tag = stbi__get8(s); - switch (tag) { - case 0x2C: /* Image Descriptor */ - { - stbi__int32 x, y, w, h; - stbi_uc *o; - - x = stbi__get16le(s); - y = stbi__get16le(s); - w = stbi__get16le(s); - h = stbi__get16le(s); - if (((x + w) > (g->w)) || ((y + h) > (g->h))) - return stbi__errpuc("bad Image Descriptor", "Corrupt GIF"); - - g->line_size = g->w * 4; - g->start_x = x * 4; - g->start_y = y * g->line_size; - g->max_x = g->start_x + w * 4; - g->max_y = g->start_y + h * g->line_size; - g->cur_x = g->start_x; - g->cur_y = g->start_y; - - // if the width of the specified rectangle is 0, that means - // we may not see *any* pixels or the image is malformed; - // to make sure this is caught, move the current y down to - // max_y (which is what out_gif_code checks). - if (w == 0) - g->cur_y = g->max_y; - - g->lflags = stbi__get8(s); - - if (g->lflags & 0x40) { - g->step = 8 * g->line_size; // first interlaced spacing - g->parse = 3; - } else { - g->step = g->line_size; - g->parse = 0; - } - - if (g->lflags & 0x80) { - stbi__gif_parse_colortable(s,g->lpal, 2 << (g->lflags & 7), g->eflags & 0x01 ? g->transparent : -1); - g->color_table = (stbi_uc *) g->lpal; - } else if (g->flags & 0x80) { - g->color_table = (stbi_uc *) g->pal; - } else - return stbi__errpuc("missing color table", "Corrupt GIF"); - - o = stbi__process_gif_raster(s, g); - if (!o) return NULL; - - // if this was the first frame, - pcount = g->w * g->h; - if (first_frame && (g->bgindex > 0)) { - // if first frame, any pixel not drawn to gets the background color - for (pi = 0; pi < pcount; ++pi) { - if (g->history[pi] == 0) { - g->pal[g->bgindex][3] = 255; // just in case it was made transparent, undo that; It will be reset next frame if need be; - memcpy( &g->out[pi * 4], &g->pal[g->bgindex], 4 ); - } - } - } - - return o; - } - - case 0x21: // Comment Extension. - { - int len; - int ext = stbi__get8(s); - if (ext == 0xF9) { // Graphic Control Extension. - len = stbi__get8(s); - if (len == 4) { - g->eflags = stbi__get8(s); - g->delay = 10 * stbi__get16le(s); // delay - 1/100th of a second, saving as 1/1000ths. - - // unset old transparent - if (g->transparent >= 0) { - g->pal[g->transparent][3] = 255; - } - if (g->eflags & 0x01) { - g->transparent = stbi__get8(s); - if (g->transparent >= 0) { - g->pal[g->transparent][3] = 0; - } - } else { - // don't need transparent - stbi__skip(s, 1); - g->transparent = -1; - } - } else { - stbi__skip(s, len); - break; - } - } - while ((len = stbi__get8(s)) != 0) { - stbi__skip(s, len); - } - break; - } - - case 0x3B: // gif stream termination code - return (stbi_uc *) s; // using '1' causes warning on some compilers - - default: - return stbi__errpuc("unknown code", "Corrupt GIF"); - } - } -} - -static void *stbi__load_gif_main_outofmem(stbi__gif *g, stbi_uc *out, int **delays) -{ - STBI_FREE(g->out); - STBI_FREE(g->history); - STBI_FREE(g->background); - - if (out) STBI_FREE(out); - if (delays && *delays) STBI_FREE(*delays); - return stbi__errpuc("outofmem", "Out of memory"); -} - -static void *stbi__load_gif_main(stbi__context *s, int **delays, int *x, int *y, int *z, int *comp, int req_comp) -{ - if (stbi__gif_test(s)) { - int layers = 0; - stbi_uc *u = 0; - stbi_uc *out = 0; - stbi_uc *two_back = 0; - stbi__gif g; - int stride; - int out_size = 0; - int delays_size = 0; - - STBI_NOTUSED(out_size); - STBI_NOTUSED(delays_size); - - memset(&g, 0, sizeof(g)); - if (delays) { - *delays = 0; - } - - do { - u = stbi__gif_load_next(s, &g, comp, req_comp, two_back); - if (u == (stbi_uc *) s) u = 0; // end of animated gif marker - - if (u) { - *x = g.w; - *y = g.h; - ++layers; - stride = g.w * g.h * 4; - - if (out) { - void *tmp = (stbi_uc*) STBI_REALLOC_SIZED( out, out_size, layers * stride ); - if (!tmp) - return stbi__load_gif_main_outofmem(&g, out, delays); - else { - out = (stbi_uc*) tmp; - out_size = layers * stride; - } - - if (delays) { - int *new_delays = (int*) STBI_REALLOC_SIZED( *delays, delays_size, sizeof(int) * layers ); - if (!new_delays) - return stbi__load_gif_main_outofmem(&g, out, delays); - *delays = new_delays; - delays_size = layers * sizeof(int); - } - } else { - out = (stbi_uc*)stbi__malloc( layers * stride ); - if (!out) - return stbi__load_gif_main_outofmem(&g, out, delays); - out_size = layers * stride; - if (delays) { - *delays = (int*) stbi__malloc( layers * sizeof(int) ); - if (!*delays) - return stbi__load_gif_main_outofmem(&g, out, delays); - delays_size = layers * sizeof(int); - } - } - memcpy( out + ((layers - 1) * stride), u, stride ); - if (layers >= 2) { - two_back = out - 2 * stride; - } - - if (delays) { - (*delays)[layers - 1U] = g.delay; - } - } - } while (u != 0); - - // free temp buffer; - STBI_FREE(g.out); - STBI_FREE(g.history); - STBI_FREE(g.background); - - // do the final conversion after loading everything; - if (req_comp && req_comp != 4) - out = stbi__convert_format(out, 4, req_comp, layers * g.w, g.h); - - *z = layers; - return out; - } else { - return stbi__errpuc("not GIF", "Image was not as a gif type."); - } -} - -static void *stbi__gif_load(stbi__context *s, int *x, int *y, int *comp, int req_comp, stbi__result_info *ri) -{ - stbi_uc *u = 0; - stbi__gif g; - memset(&g, 0, sizeof(g)); - STBI_NOTUSED(ri); - - u = stbi__gif_load_next(s, &g, comp, req_comp, 0); - if (u == (stbi_uc *) s) u = 0; // end of animated gif marker - if (u) { - *x = g.w; - *y = g.h; - - // moved conversion to after successful load so that the same - // can be done for multiple frames. - if (req_comp && req_comp != 4) - u = stbi__convert_format(u, 4, req_comp, g.w, g.h); - } else if (g.out) { - // if there was an error and we allocated an image buffer, free it! - STBI_FREE(g.out); - } - - // free buffers needed for multiple frame loading; - STBI_FREE(g.history); - STBI_FREE(g.background); - - return u; -} - -static int stbi__gif_info(stbi__context *s, int *x, int *y, int *comp) -{ - return stbi__gif_info_raw(s,x,y,comp); -} -#endif - -// ************************************************************************************************* -// Radiance RGBE HDR loader -// originally by Nicolas Schulz -#ifndef STBI_NO_HDR -static int stbi__hdr_test_core(stbi__context *s, const char *signature) -{ - int i; - for (i=0; signature[i]; ++i) - if (stbi__get8(s) != signature[i]) - return 0; - stbi__rewind(s); - return 1; -} - -static int stbi__hdr_test(stbi__context* s) -{ - int r = stbi__hdr_test_core(s, "#?RADIANCE\n"); - stbi__rewind(s); - if(!r) { - r = stbi__hdr_test_core(s, "#?RGBE\n"); - stbi__rewind(s); - } - return r; -} - -#define STBI__HDR_BUFLEN 1024 -static char *stbi__hdr_gettoken(stbi__context *z, char *buffer) -{ - int len=0; - char c = '\0'; - - c = (char) stbi__get8(z); - - while (!stbi__at_eof(z) && c != '\n') { - buffer[len++] = c; - if (len == STBI__HDR_BUFLEN-1) { - // flush to end of line - while (!stbi__at_eof(z) && stbi__get8(z) != '\n') - ; - break; - } - c = (char) stbi__get8(z); - } - - buffer[len] = 0; - return buffer; -} - -static void stbi__hdr_convert(float *output, stbi_uc *input, int req_comp) -{ - if ( input[3] != 0 ) { - float f1; - // Exponent - f1 = (float) ldexp(1.0f, input[3] - (int)(128 + 8)); - if (req_comp <= 2) - output[0] = (input[0] + input[1] + input[2]) * f1 / 3; - else { - output[0] = input[0] * f1; - output[1] = input[1] * f1; - output[2] = input[2] * f1; - } - if (req_comp == 2) output[1] = 1; - if (req_comp == 4) output[3] = 1; - } else { - switch (req_comp) { - case 4: output[3] = 1; /* fallthrough */ - case 3: output[0] = output[1] = output[2] = 0; - break; - case 2: output[1] = 1; /* fallthrough */ - case 1: output[0] = 0; - break; - } - } -} - -static float *stbi__hdr_load(stbi__context *s, int *x, int *y, int *comp, int req_comp, stbi__result_info *ri) -{ - char buffer[STBI__HDR_BUFLEN]; - char *token; - int valid = 0; - int width, height; - stbi_uc *scanline; - float *hdr_data; - int len; - unsigned char count, value; - int i, j, k, c1,c2, z; - const char *headerToken; - STBI_NOTUSED(ri); - - // Check identifier - headerToken = stbi__hdr_gettoken(s,buffer); - if (strcmp(headerToken, "#?RADIANCE") != 0 && strcmp(headerToken, "#?RGBE") != 0) - return stbi__errpf("not HDR", "Corrupt HDR image"); - - // Parse header - for(;;) { - token = stbi__hdr_gettoken(s,buffer); - if (token[0] == 0) break; - if (strcmp(token, "FORMAT=32-bit_rle_rgbe") == 0) valid = 1; - } - - if (!valid) return stbi__errpf("unsupported format", "Unsupported HDR format"); - - // Parse width and height - // can't use sscanf() if we're not using stdio! - token = stbi__hdr_gettoken(s,buffer); - if (strncmp(token, "-Y ", 3)) return stbi__errpf("unsupported data layout", "Unsupported HDR format"); - token += 3; - height = (int) strtol(token, &token, 10); - while (*token == ' ') ++token; - if (strncmp(token, "+X ", 3)) return stbi__errpf("unsupported data layout", "Unsupported HDR format"); - token += 3; - width = (int) strtol(token, NULL, 10); - - if (height > STBI_MAX_DIMENSIONS) return stbi__errpf("too large","Very large image (corrupt?)"); - if (width > STBI_MAX_DIMENSIONS) return stbi__errpf("too large","Very large image (corrupt?)"); - - *x = width; - *y = height; - - if (comp) *comp = 3; - if (req_comp == 0) req_comp = 3; - - if (!stbi__mad4sizes_valid(width, height, req_comp, sizeof(float), 0)) - return stbi__errpf("too large", "HDR image is too large"); - - // Read data - hdr_data = (float *) stbi__malloc_mad4(width, height, req_comp, sizeof(float), 0); - if (!hdr_data) - return stbi__errpf("outofmem", "Out of memory"); - - // Load image data - // image data is stored as some number of sca - if ( width < 8 || width >= 32768) { - // Read flat data - for (j=0; j < height; ++j) { - for (i=0; i < width; ++i) { - stbi_uc rgbe[4]; - main_decode_loop: - stbi__getn(s, rgbe, 4); - stbi__hdr_convert(hdr_data + j * width * req_comp + i * req_comp, rgbe, req_comp); - } - } - } else { - // Read RLE-encoded data - scanline = NULL; - - for (j = 0; j < height; ++j) { - c1 = stbi__get8(s); - c2 = stbi__get8(s); - len = stbi__get8(s); - if (c1 != 2 || c2 != 2 || (len & 0x80)) { - // not run-length encoded, so we have to actually use THIS data as a decoded - // pixel (note this can't be a valid pixel--one of RGB must be >= 128) - stbi_uc rgbe[4]; - rgbe[0] = (stbi_uc) c1; - rgbe[1] = (stbi_uc) c2; - rgbe[2] = (stbi_uc) len; - rgbe[3] = (stbi_uc) stbi__get8(s); - stbi__hdr_convert(hdr_data, rgbe, req_comp); - i = 1; - j = 0; - STBI_FREE(scanline); - goto main_decode_loop; // yes, this makes no sense - } - len <<= 8; - len |= stbi__get8(s); - if (len != width) { STBI_FREE(hdr_data); STBI_FREE(scanline); return stbi__errpf("invalid decoded scanline length", "corrupt HDR"); } - if (scanline == NULL) { - scanline = (stbi_uc *) stbi__malloc_mad2(width, 4, 0); - if (!scanline) { - STBI_FREE(hdr_data); - return stbi__errpf("outofmem", "Out of memory"); - } - } - - for (k = 0; k < 4; ++k) { - int nleft; - i = 0; - while ((nleft = width - i) > 0) { - count = stbi__get8(s); - if (count > 128) { - // Run - value = stbi__get8(s); - count -= 128; - if ((count == 0) || (count > nleft)) { STBI_FREE(hdr_data); STBI_FREE(scanline); return stbi__errpf("corrupt", "bad RLE data in HDR"); } - for (z = 0; z < count; ++z) - scanline[i++ * 4 + k] = value; - } else { - // Dump - if ((count == 0) || (count > nleft)) { STBI_FREE(hdr_data); STBI_FREE(scanline); return stbi__errpf("corrupt", "bad RLE data in HDR"); } - for (z = 0; z < count; ++z) - scanline[i++ * 4 + k] = stbi__get8(s); - } - } - } - for (i=0; i < width; ++i) - stbi__hdr_convert(hdr_data+(j*width + i)*req_comp, scanline + i*4, req_comp); - } - if (scanline) - STBI_FREE(scanline); - } - - return hdr_data; -} - -static int stbi__hdr_info(stbi__context *s, int *x, int *y, int *comp) -{ - char buffer[STBI__HDR_BUFLEN]; - char *token; - int valid = 0; - int dummy; - - if (!x) x = &dummy; - if (!y) y = &dummy; - if (!comp) comp = &dummy; - - if (stbi__hdr_test(s) == 0) { - stbi__rewind( s ); - return 0; - } - - for(;;) { - token = stbi__hdr_gettoken(s,buffer); - if (token[0] == 0) break; - if (strcmp(token, "FORMAT=32-bit_rle_rgbe") == 0) valid = 1; - } - - if (!valid) { - stbi__rewind( s ); - return 0; - } - token = stbi__hdr_gettoken(s,buffer); - if (strncmp(token, "-Y ", 3)) { - stbi__rewind( s ); - return 0; - } - token += 3; - *y = (int) strtol(token, &token, 10); - while (*token == ' ') ++token; - if (strncmp(token, "+X ", 3)) { - stbi__rewind( s ); - return 0; - } - token += 3; - *x = (int) strtol(token, NULL, 10); - *comp = 3; - return 1; -} -#endif // STBI_NO_HDR - -#ifndef STBI_NO_BMP -static int stbi__bmp_info(stbi__context *s, int *x, int *y, int *comp) -{ - void *p; - stbi__bmp_data info; - - info.all_a = 255; - p = stbi__bmp_parse_header(s, &info); - if (p == NULL) { - stbi__rewind( s ); - return 0; - } - if (x) *x = s->img_x; - if (y) *y = s->img_y; - if (comp) { - if (info.bpp == 24 && info.ma == 0xff000000) - *comp = 3; - else - *comp = info.ma ? 4 : 3; - } - return 1; -} -#endif - -#ifndef STBI_NO_PSD -static int stbi__psd_info(stbi__context *s, int *x, int *y, int *comp) -{ - int channelCount, dummy, depth; - if (!x) x = &dummy; - if (!y) y = &dummy; - if (!comp) comp = &dummy; - if (stbi__get32be(s) != 0x38425053) { - stbi__rewind( s ); - return 0; - } - if (stbi__get16be(s) != 1) { - stbi__rewind( s ); - return 0; - } - stbi__skip(s, 6); - channelCount = stbi__get16be(s); - if (channelCount < 0 || channelCount > 16) { - stbi__rewind( s ); - return 0; - } - *y = stbi__get32be(s); - *x = stbi__get32be(s); - depth = stbi__get16be(s); - if (depth != 8 && depth != 16) { - stbi__rewind( s ); - return 0; - } - if (stbi__get16be(s) != 3) { - stbi__rewind( s ); - return 0; - } - *comp = 4; - return 1; -} - -static int stbi__psd_is16(stbi__context *s) -{ - int channelCount, depth; - if (stbi__get32be(s) != 0x38425053) { - stbi__rewind( s ); - return 0; - } - if (stbi__get16be(s) != 1) { - stbi__rewind( s ); - return 0; - } - stbi__skip(s, 6); - channelCount = stbi__get16be(s); - if (channelCount < 0 || channelCount > 16) { - stbi__rewind( s ); - return 0; - } - STBI_NOTUSED(stbi__get32be(s)); - STBI_NOTUSED(stbi__get32be(s)); - depth = stbi__get16be(s); - if (depth != 16) { - stbi__rewind( s ); - return 0; - } - return 1; -} -#endif - -#ifndef STBI_NO_PIC -static int stbi__pic_info(stbi__context *s, int *x, int *y, int *comp) -{ - int act_comp=0,num_packets=0,chained,dummy; - stbi__pic_packet packets[10]; - - if (!x) x = &dummy; - if (!y) y = &dummy; - if (!comp) comp = &dummy; - - if (!stbi__pic_is4(s,"\x53\x80\xF6\x34")) { - stbi__rewind(s); - return 0; - } - - stbi__skip(s, 88); - - *x = stbi__get16be(s); - *y = stbi__get16be(s); - if (stbi__at_eof(s)) { - stbi__rewind( s); - return 0; - } - if ( (*x) != 0 && (1 << 28) / (*x) < (*y)) { - stbi__rewind( s ); - return 0; - } - - stbi__skip(s, 8); - - do { - stbi__pic_packet *packet; - - if (num_packets==sizeof(packets)/sizeof(packets[0])) - return 0; - - packet = &packets[num_packets++]; - chained = stbi__get8(s); - packet->size = stbi__get8(s); - packet->type = stbi__get8(s); - packet->channel = stbi__get8(s); - act_comp |= packet->channel; - - if (stbi__at_eof(s)) { - stbi__rewind( s ); - return 0; - } - if (packet->size != 8) { - stbi__rewind( s ); - return 0; - } - } while (chained); - - *comp = (act_comp & 0x10 ? 4 : 3); - - return 1; -} -#endif - -// ************************************************************************************************* -// Portable Gray Map and Portable Pixel Map loader -// by Ken Miller -// -// PGM: http://netpbm.sourceforge.net/doc/pgm.html -// PPM: http://netpbm.sourceforge.net/doc/ppm.html -// -// Known limitations: -// Does not support comments in the header section -// Does not support ASCII image data (formats P2 and P3) - -#ifndef STBI_NO_PNM - -static int stbi__pnm_test(stbi__context *s) -{ - char p, t; - p = (char) stbi__get8(s); - t = (char) stbi__get8(s); - if (p != 'P' || (t != '5' && t != '6')) { - stbi__rewind( s ); - return 0; - } - return 1; -} - -static void *stbi__pnm_load(stbi__context *s, int *x, int *y, int *comp, int req_comp, stbi__result_info *ri) -{ - stbi_uc *out; - STBI_NOTUSED(ri); - - ri->bits_per_channel = stbi__pnm_info(s, (int *)&s->img_x, (int *)&s->img_y, (int *)&s->img_n); - if (ri->bits_per_channel == 0) - return 0; - - if (s->img_y > STBI_MAX_DIMENSIONS) return stbi__errpuc("too large","Very large image (corrupt?)"); - if (s->img_x > STBI_MAX_DIMENSIONS) return stbi__errpuc("too large","Very large image (corrupt?)"); - - *x = s->img_x; - *y = s->img_y; - if (comp) *comp = s->img_n; - - if (!stbi__mad4sizes_valid(s->img_n, s->img_x, s->img_y, ri->bits_per_channel / 8, 0)) - return stbi__errpuc("too large", "PNM too large"); - - out = (stbi_uc *) stbi__malloc_mad4(s->img_n, s->img_x, s->img_y, ri->bits_per_channel / 8, 0); - if (!out) return stbi__errpuc("outofmem", "Out of memory"); - if (!stbi__getn(s, out, s->img_n * s->img_x * s->img_y * (ri->bits_per_channel / 8))) { - STBI_FREE(out); - return stbi__errpuc("bad PNM", "PNM file truncated"); - } - - if (req_comp && req_comp != s->img_n) { - if (ri->bits_per_channel == 16) { - out = (stbi_uc *) stbi__convert_format16((stbi__uint16 *) out, s->img_n, req_comp, s->img_x, s->img_y); - } else { - out = stbi__convert_format(out, s->img_n, req_comp, s->img_x, s->img_y); - } - if (out == NULL) return out; // stbi__convert_format frees input on failure - } - return out; -} - -static int stbi__pnm_isspace(char c) -{ - return c == ' ' || c == '\t' || c == '\n' || c == '\v' || c == '\f' || c == '\r'; -} - -static void stbi__pnm_skip_whitespace(stbi__context *s, char *c) -{ - for (;;) { - while (!stbi__at_eof(s) && stbi__pnm_isspace(*c)) - *c = (char) stbi__get8(s); - - if (stbi__at_eof(s) || *c != '#') - break; - - while (!stbi__at_eof(s) && *c != '\n' && *c != '\r' ) - *c = (char) stbi__get8(s); - } -} - -static int stbi__pnm_isdigit(char c) -{ - return c >= '0' && c <= '9'; -} - -static int stbi__pnm_getinteger(stbi__context *s, char *c) -{ - int value = 0; - - while (!stbi__at_eof(s) && stbi__pnm_isdigit(*c)) { - value = value*10 + (*c - '0'); - *c = (char) stbi__get8(s); - if((value > 214748364) || (value == 214748364 && *c > '7')) - return stbi__err("integer parse overflow", "Parsing an integer in the PPM header overflowed a 32-bit int"); - } - - return value; -} - -static int stbi__pnm_info(stbi__context *s, int *x, int *y, int *comp) -{ - int maxv, dummy; - char c, p, t; - - if (!x) x = &dummy; - if (!y) y = &dummy; - if (!comp) comp = &dummy; - - stbi__rewind(s); - - // Get identifier - p = (char) stbi__get8(s); - t = (char) stbi__get8(s); - if (p != 'P' || (t != '5' && t != '6')) { - stbi__rewind(s); - return 0; - } - - *comp = (t == '6') ? 3 : 1; // '5' is 1-component .pgm; '6' is 3-component .ppm - - c = (char) stbi__get8(s); - stbi__pnm_skip_whitespace(s, &c); - - *x = stbi__pnm_getinteger(s, &c); // read width - if(*x == 0) - return stbi__err("invalid width", "PPM image header had zero or overflowing width"); - stbi__pnm_skip_whitespace(s, &c); - - *y = stbi__pnm_getinteger(s, &c); // read height - if (*y == 0) - return stbi__err("invalid width", "PPM image header had zero or overflowing width"); - stbi__pnm_skip_whitespace(s, &c); - - maxv = stbi__pnm_getinteger(s, &c); // read max value - if (maxv > 65535) - return stbi__err("max value > 65535", "PPM image supports only 8-bit and 16-bit images"); - else if (maxv > 255) - return 16; - else - return 8; -} - -static int stbi__pnm_is16(stbi__context *s) -{ - if (stbi__pnm_info(s, NULL, NULL, NULL) == 16) - return 1; - return 0; -} -#endif - -static int stbi__info_main(stbi__context *s, int *x, int *y, int *comp) -{ - #ifndef STBI_NO_JPEG - if (stbi__jpeg_info(s, x, y, comp)) return 1; - #endif - - #ifndef STBI_NO_PNG - if (stbi__png_info(s, x, y, comp)) return 1; - #endif - - #ifndef STBI_NO_GIF - if (stbi__gif_info(s, x, y, comp)) return 1; - #endif - - #ifndef STBI_NO_BMP - if (stbi__bmp_info(s, x, y, comp)) return 1; - #endif - - #ifndef STBI_NO_PSD - if (stbi__psd_info(s, x, y, comp)) return 1; - #endif - - #ifndef STBI_NO_PIC - if (stbi__pic_info(s, x, y, comp)) return 1; - #endif - - #ifndef STBI_NO_PNM - if (stbi__pnm_info(s, x, y, comp)) return 1; - #endif - - #ifndef STBI_NO_HDR - if (stbi__hdr_info(s, x, y, comp)) return 1; - #endif - - // test tga last because it's a crappy test! - #ifndef STBI_NO_TGA - if (stbi__tga_info(s, x, y, comp)) - return 1; - #endif - return stbi__err("unknown image type", "Image not of any known type, or corrupt"); -} - -static int stbi__is_16_main(stbi__context *s) -{ - #ifndef STBI_NO_PNG - if (stbi__png_is16(s)) return 1; - #endif - - #ifndef STBI_NO_PSD - if (stbi__psd_is16(s)) return 1; - #endif - - #ifndef STBI_NO_PNM - if (stbi__pnm_is16(s)) return 1; - #endif - return 0; -} - -#ifndef STBI_NO_STDIO -STBIDEF int stbi_info(char const *filename, int *x, int *y, int *comp) -{ - FILE *f = stbi__fopen(filename, "rb"); - int result; - if (!f) return stbi__err("can't fopen", "Unable to open file"); - result = stbi_info_from_file(f, x, y, comp); - fclose(f); - return result; -} - -STBIDEF int stbi_info_from_file(FILE *f, int *x, int *y, int *comp) -{ - int r; - stbi__context s; - long pos = ftell(f); - stbi__start_file(&s, f); - r = stbi__info_main(&s,x,y,comp); - fseek(f,pos,SEEK_SET); - return r; -} - -STBIDEF int stbi_is_16_bit(char const *filename) -{ - FILE *f = stbi__fopen(filename, "rb"); - int result; - if (!f) return stbi__err("can't fopen", "Unable to open file"); - result = stbi_is_16_bit_from_file(f); - fclose(f); - return result; -} - -STBIDEF int stbi_is_16_bit_from_file(FILE *f) -{ - int r; - stbi__context s; - long pos = ftell(f); - stbi__start_file(&s, f); - r = stbi__is_16_main(&s); - fseek(f,pos,SEEK_SET); - return r; -} -#endif // !STBI_NO_STDIO - -STBIDEF int stbi_info_from_memory(stbi_uc const *buffer, int len, int *x, int *y, int *comp) -{ - stbi__context s; - stbi__start_mem(&s,buffer,len); - return stbi__info_main(&s,x,y,comp); -} - -STBIDEF int stbi_info_from_callbacks(stbi_io_callbacks const *c, void *user, int *x, int *y, int *comp) -{ - stbi__context s; - stbi__start_callbacks(&s, (stbi_io_callbacks *) c, user); - return stbi__info_main(&s,x,y,comp); -} - -STBIDEF int stbi_is_16_bit_from_memory(stbi_uc const *buffer, int len) -{ - stbi__context s; - stbi__start_mem(&s,buffer,len); - return stbi__is_16_main(&s); -} - -STBIDEF int stbi_is_16_bit_from_callbacks(stbi_io_callbacks const *c, void *user) -{ - stbi__context s; - stbi__start_callbacks(&s, (stbi_io_callbacks *) c, user); - return stbi__is_16_main(&s); -} - -#endif // STB_IMAGE_IMPLEMENTATION - -/* - revision history: - 2.20 (2019-02-07) support utf8 filenames in Windows; fix warnings and platform ifdefs - 2.19 (2018-02-11) fix warning - 2.18 (2018-01-30) fix warnings - 2.17 (2018-01-29) change sbti__shiftsigned to avoid clang -O2 bug - 1-bit BMP - *_is_16_bit api - avoid warnings - 2.16 (2017-07-23) all functions have 16-bit variants; - STBI_NO_STDIO works again; - compilation fixes; - fix rounding in unpremultiply; - optimize vertical flip; - disable raw_len validation; - documentation fixes - 2.15 (2017-03-18) fix png-1,2,4 bug; now all Imagenet JPGs decode; - warning fixes; disable run-time SSE detection on gcc; - uniform handling of optional "return" values; - thread-safe initialization of zlib tables - 2.14 (2017-03-03) remove deprecated STBI_JPEG_OLD; fixes for Imagenet JPGs - 2.13 (2016-11-29) add 16-bit API, only supported for PNG right now - 2.12 (2016-04-02) fix typo in 2.11 PSD fix that caused crashes - 2.11 (2016-04-02) allocate large structures on the stack - remove white matting for transparent PSD - fix reported channel count for PNG & BMP - re-enable SSE2 in non-gcc 64-bit - support RGB-formatted JPEG - read 16-bit PNGs (only as 8-bit) - 2.10 (2016-01-22) avoid warning introduced in 2.09 by STBI_REALLOC_SIZED - 2.09 (2016-01-16) allow comments in PNM files - 16-bit-per-pixel TGA (not bit-per-component) - info() for TGA could break due to .hdr handling - info() for BMP to shares code instead of sloppy parse - can use STBI_REALLOC_SIZED if allocator doesn't support realloc - code cleanup - 2.08 (2015-09-13) fix to 2.07 cleanup, reading RGB PSD as RGBA - 2.07 (2015-09-13) fix compiler warnings - partial animated GIF support - limited 16-bpc PSD support - #ifdef unused functions - bug with < 92 byte PIC,PNM,HDR,TGA - 2.06 (2015-04-19) fix bug where PSD returns wrong '*comp' value - 2.05 (2015-04-19) fix bug in progressive JPEG handling, fix warning - 2.04 (2015-04-15) try to re-enable SIMD on MinGW 64-bit - 2.03 (2015-04-12) extra corruption checking (mmozeiko) - stbi_set_flip_vertically_on_load (nguillemot) - fix NEON support; fix mingw support - 2.02 (2015-01-19) fix incorrect assert, fix warning - 2.01 (2015-01-17) fix various warnings; suppress SIMD on gcc 32-bit without -msse2 - 2.00b (2014-12-25) fix STBI_MALLOC in progressive JPEG - 2.00 (2014-12-25) optimize JPG, including x86 SSE2 & NEON SIMD (ryg) - progressive JPEG (stb) - PGM/PPM support (Ken Miller) - STBI_MALLOC,STBI_REALLOC,STBI_FREE - GIF bugfix -- seemingly never worked - STBI_NO_*, STBI_ONLY_* - 1.48 (2014-12-14) fix incorrectly-named assert() - 1.47 (2014-12-14) 1/2/4-bit PNG support, both direct and paletted (Omar Cornut & stb) - optimize PNG (ryg) - fix bug in interlaced PNG with user-specified channel count (stb) - 1.46 (2014-08-26) - fix broken tRNS chunk (colorkey-style transparency) in non-paletted PNG - 1.45 (2014-08-16) - fix MSVC-ARM internal compiler error by wrapping malloc - 1.44 (2014-08-07) - various warning fixes from Ronny Chevalier - 1.43 (2014-07-15) - fix MSVC-only compiler problem in code changed in 1.42 - 1.42 (2014-07-09) - don't define _CRT_SECURE_NO_WARNINGS (affects user code) - fixes to stbi__cleanup_jpeg path - added STBI_ASSERT to avoid requiring assert.h - 1.41 (2014-06-25) - fix search&replace from 1.36 that messed up comments/error messages - 1.40 (2014-06-22) - fix gcc struct-initialization warning - 1.39 (2014-06-15) - fix to TGA optimization when req_comp != number of components in TGA; - fix to GIF loading because BMP wasn't rewinding (whoops, no GIFs in my test suite) - add support for BMP version 5 (more ignored fields) - 1.38 (2014-06-06) - suppress MSVC warnings on integer casts truncating values - fix accidental rename of 'skip' field of I/O - 1.37 (2014-06-04) - remove duplicate typedef - 1.36 (2014-06-03) - convert to header file single-file library - if de-iphone isn't set, load iphone images color-swapped instead of returning NULL - 1.35 (2014-05-27) - various warnings - fix broken STBI_SIMD path - fix bug where stbi_load_from_file no longer left file pointer in correct place - fix broken non-easy path for 32-bit BMP (possibly never used) - TGA optimization by Arseny Kapoulkine - 1.34 (unknown) - use STBI_NOTUSED in stbi__resample_row_generic(), fix one more leak in tga failure case - 1.33 (2011-07-14) - make stbi_is_hdr work in STBI_NO_HDR (as specified), minor compiler-friendly improvements - 1.32 (2011-07-13) - support for "info" function for all supported filetypes (SpartanJ) - 1.31 (2011-06-20) - a few more leak fixes, bug in PNG handling (SpartanJ) - 1.30 (2011-06-11) - added ability to load files via callbacks to accomidate custom input streams (Ben Wenger) - removed deprecated format-specific test/load functions - removed support for installable file formats (stbi_loader) -- would have been broken for IO callbacks anyway - error cases in bmp and tga give messages and don't leak (Raymond Barbiero, grisha) - fix inefficiency in decoding 32-bit BMP (David Woo) - 1.29 (2010-08-16) - various warning fixes from Aurelien Pocheville - 1.28 (2010-08-01) - fix bug in GIF palette transparency (SpartanJ) - 1.27 (2010-08-01) - cast-to-stbi_uc to fix warnings - 1.26 (2010-07-24) - fix bug in file buffering for PNG reported by SpartanJ - 1.25 (2010-07-17) - refix trans_data warning (Won Chun) - 1.24 (2010-07-12) - perf improvements reading from files on platforms with lock-heavy fgetc() - minor perf improvements for jpeg - deprecated type-specific functions so we'll get feedback if they're needed - attempt to fix trans_data warning (Won Chun) - 1.23 fixed bug in iPhone support - 1.22 (2010-07-10) - removed image *writing* support - stbi_info support from Jetro Lauha - GIF support from Jean-Marc Lienher - iPhone PNG-extensions from James Brown - warning-fixes from Nicolas Schulz and Janez Zemva (i.stbi__err. Janez (U+017D)emva) - 1.21 fix use of 'stbi_uc' in header (reported by jon blow) - 1.20 added support for Softimage PIC, by Tom Seddon - 1.19 bug in interlaced PNG corruption check (found by ryg) - 1.18 (2008-08-02) - fix a threading bug (local mutable static) - 1.17 support interlaced PNG - 1.16 major bugfix - stbi__convert_format converted one too many pixels - 1.15 initialize some fields for thread safety - 1.14 fix threadsafe conversion bug - header-file-only version (#define STBI_HEADER_FILE_ONLY before including) - 1.13 threadsafe - 1.12 const qualifiers in the API - 1.11 Support installable IDCT, colorspace conversion routines - 1.10 Fixes for 64-bit (don't use "unsigned long") - optimized upsampling by Fabian "ryg" Giesen - 1.09 Fix format-conversion for PSD code (bad global variables!) - 1.08 Thatcher Ulrich's PSD code integrated by Nicolas Schulz - 1.07 attempt to fix C++ warning/errors again - 1.06 attempt to fix C++ warning/errors again - 1.05 fix TGA loading to return correct *comp and use good luminance calc - 1.04 default float alpha is 1, not 255; use 'void *' for stbi_image_free - 1.03 bugfixes to STBI_NO_STDIO, STBI_NO_HDR - 1.02 support for (subset of) HDR files, float interface for preferred access to them - 1.01 fix bug: possible bug in handling right-side up bmps... not sure - fix bug: the stbi__bmp_load() and stbi__tga_load() functions didn't work at all - 1.00 interface to zlib that skips zlib header - 0.99 correct handling of alpha in palette - 0.98 TGA loader by lonesock; dynamically add loaders (untested) - 0.97 jpeg errors on too large a file; also catch another malloc failure - 0.96 fix detection of invalid v value - particleman@mollyrocket forum - 0.95 during header scan, seek to markers in case of padding - 0.94 STBI_NO_STDIO to disable stdio usage; rename all #defines the same - 0.93 handle jpegtran output; verbose errors - 0.92 read 4,8,16,24,32-bit BMP files of several formats - 0.91 output 24-bit Windows 3.0 BMP files - 0.90 fix a few more warnings; bump version number to approach 1.0 - 0.61 bugfixes due to Marc LeBlanc, Christopher Lloyd - 0.60 fix compiling as c++ - 0.59 fix warnings: merge Dave Moore's -Wall fixes - 0.58 fix bug: zlib uncompressed mode len/nlen was wrong endian - 0.57 fix bug: jpg last huffman symbol before marker was >9 bits but less than 16 available - 0.56 fix bug: zlib uncompressed mode len vs. nlen - 0.55 fix bug: restart_interval not initialized to 0 - 0.54 allow NULL for 'int *comp' - 0.53 fix bug in png 3->4; speedup png decoding - 0.52 png handles req_comp=3,4 directly; minor cleanup; jpeg comments - 0.51 obey req_comp requests, 1-component jpegs return as 1-component, - on 'test' only check type, not whether we support this variant - 0.50 (2006-11-19) - first released version -*/ - - -/* ------------------------------------------------------------------------------- -This software is available under 2 licenses -- choose whichever you prefer. ------------------------------------------------------------------------------- -ALTERNATIVE A - MIT License -Copyright (c) 2017 Sean Barrett -Permission is hereby granted, free of charge, to any person obtaining a copy of -this software and associated documentation files (the "Software"), to deal in -the Software without restriction, including without limitation the rights to -use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies -of the Software, and to permit persons to whom the Software is furnished to do -so, subject to the following conditions: -The above copyright notice and this permission notice shall be included in all -copies or substantial portions of the Software. -THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -SOFTWARE. ------------------------------------------------------------------------------- -ALTERNATIVE B - Public Domain (www.unlicense.org) -This is free and unencumbered software released into the public domain. -Anyone is free to copy, modify, publish, use, compile, sell, or distribute this -software, either in source code form or as a compiled binary, for any purpose, -commercial or non-commercial, and by any means. -In jurisdictions that recognize copyright laws, the author or authors of this -software dedicate any and all copyright interest in the software to the public -domain. We make this dedication for the benefit of the public at large and to -the detriment of our heirs and successors. We intend this dedication to be an -overt act of relinquishment in perpetuity of all present and future rights to -this software under copyright law. -THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -AUTHORS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN -ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION -WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. ------------------------------------------------------------------------------- -*/ diff --git a/third-party/stb/stb_image_resize.h b/third-party/stb/stb_image_resize.h deleted file mode 100644 index ef9e6fe8704..00000000000 --- a/third-party/stb/stb_image_resize.h +++ /dev/null @@ -1,2634 +0,0 @@ -/* stb_image_resize - v0.97 - public domain image resizing - by Jorge L Rodriguez (@VinoBS) - 2014 - http://github.com/nothings/stb - - Written with emphasis on usability, portability, and efficiency. (No - SIMD or threads, so it be easily outperformed by libs that use those.) - Only scaling and translation is supported, no rotations or shears. - Easy API downsamples w/Mitchell filter, upsamples w/cubic interpolation. - - COMPILING & LINKING - In one C/C++ file that #includes this file, do this: - #define STB_IMAGE_RESIZE_IMPLEMENTATION - before the #include. That will create the implementation in that file. - - QUICKSTART - stbir_resize_uint8( input_pixels , in_w , in_h , 0, - output_pixels, out_w, out_h, 0, num_channels) - stbir_resize_float(...) - stbir_resize_uint8_srgb( input_pixels , in_w , in_h , 0, - output_pixels, out_w, out_h, 0, - num_channels , alpha_chan , 0) - stbir_resize_uint8_srgb_edgemode( - input_pixels , in_w , in_h , 0, - output_pixels, out_w, out_h, 0, - num_channels , alpha_chan , 0, STBIR_EDGE_CLAMP) - // WRAP/REFLECT/ZERO - - FULL API - See the "header file" section of the source for API documentation. - - ADDITIONAL DOCUMENTATION - - SRGB & FLOATING POINT REPRESENTATION - The sRGB functions presume IEEE floating point. If you do not have - IEEE floating point, define STBIR_NON_IEEE_FLOAT. This will use - a slower implementation. - - MEMORY ALLOCATION - The resize functions here perform a single memory allocation using - malloc. To control the memory allocation, before the #include that - triggers the implementation, do: - - #define STBIR_MALLOC(size,context) ... - #define STBIR_FREE(ptr,context) ... - - Each resize function makes exactly one call to malloc/free, so to use - temp memory, store the temp memory in the context and return that. - - ASSERT - Define STBIR_ASSERT(boolval) to override assert() and not use assert.h - - OPTIMIZATION - Define STBIR_SATURATE_INT to compute clamp values in-range using - integer operations instead of float operations. This may be faster - on some platforms. - - DEFAULT FILTERS - For functions which don't provide explicit control over what filters - to use, you can change the compile-time defaults with - - #define STBIR_DEFAULT_FILTER_UPSAMPLE STBIR_FILTER_something - #define STBIR_DEFAULT_FILTER_DOWNSAMPLE STBIR_FILTER_something - - See stbir_filter in the header-file section for the list of filters. - - NEW FILTERS - A number of 1D filter kernels are used. For a list of - supported filters see the stbir_filter enum. To add a new filter, - write a filter function and add it to stbir__filter_info_table. - - PROGRESS - For interactive use with slow resize operations, you can install - a progress-report callback: - - #define STBIR_PROGRESS_REPORT(val) some_func(val) - - The parameter val is a float which goes from 0 to 1 as progress is made. - - For example: - - static void my_progress_report(float progress); - #define STBIR_PROGRESS_REPORT(val) my_progress_report(val) - - #define STB_IMAGE_RESIZE_IMPLEMENTATION - #include "stb_image_resize.h" - - static void my_progress_report(float progress) - { - printf("Progress: %f%%\n", progress*100); - } - - MAX CHANNELS - If your image has more than 64 channels, define STBIR_MAX_CHANNELS - to the max you'll have. - - ALPHA CHANNEL - Most of the resizing functions provide the ability to control how - the alpha channel of an image is processed. The important things - to know about this: - - 1. The best mathematically-behaved version of alpha to use is - called "premultiplied alpha", in which the other color channels - have had the alpha value multiplied in. If you use premultiplied - alpha, linear filtering (such as image resampling done by this - library, or performed in texture units on GPUs) does the "right - thing". While premultiplied alpha is standard in the movie CGI - industry, it is still uncommon in the videogame/real-time world. - - If you linearly filter non-premultiplied alpha, strange effects - occur. (For example, the 50/50 average of 99% transparent bright green - and 1% transparent black produces 50% transparent dark green when - non-premultiplied, whereas premultiplied it produces 50% - transparent near-black. The former introduces green energy - that doesn't exist in the source image.) - - 2. Artists should not edit premultiplied-alpha images; artists - want non-premultiplied alpha images. Thus, art tools generally output - non-premultiplied alpha images. - - 3. You will get best results in most cases by converting images - to premultiplied alpha before processing them mathematically. - - 4. If you pass the flag STBIR_FLAG_ALPHA_PREMULTIPLIED, the - resizer does not do anything special for the alpha channel; - it is resampled identically to other channels. This produces - the correct results for premultiplied-alpha images, but produces - less-than-ideal results for non-premultiplied-alpha images. - - 5. If you do not pass the flag STBIR_FLAG_ALPHA_PREMULTIPLIED, - then the resizer weights the contribution of input pixels - based on their alpha values, or, equivalently, it multiplies - the alpha value into the color channels, resamples, then divides - by the resultant alpha value. Input pixels which have alpha=0 do - not contribute at all to output pixels unless _all_ of the input - pixels affecting that output pixel have alpha=0, in which case - the result for that pixel is the same as it would be without - STBIR_FLAG_ALPHA_PREMULTIPLIED. However, this is only true for - input images in integer formats. For input images in float format, - input pixels with alpha=0 have no effect, and output pixels - which have alpha=0 will be 0 in all channels. (For float images, - you can manually achieve the same result by adding a tiny epsilon - value to the alpha channel of every image, and then subtracting - or clamping it at the end.) - - 6. You can suppress the behavior described in #5 and make - all-0-alpha pixels have 0 in all channels by #defining - STBIR_NO_ALPHA_EPSILON. - - 7. You can separately control whether the alpha channel is - interpreted as linear or affected by the colorspace. By default - it is linear; you almost never want to apply the colorspace. - (For example, graphics hardware does not apply sRGB conversion - to the alpha channel.) - - CONTRIBUTORS - Jorge L Rodriguez: Implementation - Sean Barrett: API design, optimizations - Aras Pranckevicius: bugfix - Nathan Reed: warning fixes - - REVISIONS - 0.97 (2020-02-02) fixed warning - 0.96 (2019-03-04) fixed warnings - 0.95 (2017-07-23) fixed warnings - 0.94 (2017-03-18) fixed warnings - 0.93 (2017-03-03) fixed bug with certain combinations of heights - 0.92 (2017-01-02) fix integer overflow on large (>2GB) images - 0.91 (2016-04-02) fix warnings; fix handling of subpixel regions - 0.90 (2014-09-17) first released version - - LICENSE - See end of file for license information. - - TODO - Don't decode all of the image data when only processing a partial tile - Don't use full-width decode buffers when only processing a partial tile - When processing wide images, break processing into tiles so data fits in L1 cache - Installable filters? - Resize that respects alpha test coverage - (Reference code: FloatImage::alphaTestCoverage and FloatImage::scaleAlphaToCoverage: - https://code.google.com/p/nvidia-texture-tools/source/browse/trunk/src/nvimage/FloatImage.cpp ) -*/ - -#ifndef STBIR_INCLUDE_STB_IMAGE_RESIZE_H -#define STBIR_INCLUDE_STB_IMAGE_RESIZE_H - -#ifdef _MSC_VER -typedef unsigned char stbir_uint8; -typedef unsigned short stbir_uint16; -typedef unsigned int stbir_uint32; -#else -#include -typedef uint8_t stbir_uint8; -typedef uint16_t stbir_uint16; -typedef uint32_t stbir_uint32; -#endif - -#ifndef STBIRDEF -#ifdef STB_IMAGE_RESIZE_STATIC -#define STBIRDEF static -#else -#ifdef __cplusplus -#define STBIRDEF extern "C" -#else -#define STBIRDEF extern -#endif -#endif -#endif - -////////////////////////////////////////////////////////////////////////////// -// -// Easy-to-use API: -// -// * "input pixels" points to an array of image data with 'num_channels' channels (e.g. RGB=3, RGBA=4) -// * input_w is input image width (x-axis), input_h is input image height (y-axis) -// * stride is the offset between successive rows of image data in memory, in bytes. you can -// specify 0 to mean packed continuously in memory -// * alpha channel is treated identically to other channels. -// * colorspace is linear or sRGB as specified by function name -// * returned result is 1 for success or 0 in case of an error. -// #define STBIR_ASSERT() to trigger an assert on parameter validation errors. -// * Memory required grows approximately linearly with input and output size, but with -// discontinuities at input_w == output_w and input_h == output_h. -// * These functions use a "default" resampling filter defined at compile time. To change the filter, -// you can change the compile-time defaults by #defining STBIR_DEFAULT_FILTER_UPSAMPLE -// and STBIR_DEFAULT_FILTER_DOWNSAMPLE, or you can use the medium-complexity API. - -STBIRDEF int stbir_resize_uint8( const unsigned char *input_pixels , int input_w , int input_h , int input_stride_in_bytes, - unsigned char *output_pixels, int output_w, int output_h, int output_stride_in_bytes, - int num_channels); - -STBIRDEF int stbir_resize_float( const float *input_pixels , int input_w , int input_h , int input_stride_in_bytes, - float *output_pixels, int output_w, int output_h, int output_stride_in_bytes, - int num_channels); - - -// The following functions interpret image data as gamma-corrected sRGB. -// Specify STBIR_ALPHA_CHANNEL_NONE if you have no alpha channel, -// or otherwise provide the index of the alpha channel. Flags value -// of 0 will probably do the right thing if you're not sure what -// the flags mean. - -#define STBIR_ALPHA_CHANNEL_NONE -1 - -// Set this flag if your texture has premultiplied alpha. Otherwise, stbir will -// use alpha-weighted resampling (effectively premultiplying, resampling, -// then unpremultiplying). -#define STBIR_FLAG_ALPHA_PREMULTIPLIED (1 << 0) -// The specified alpha channel should be handled as gamma-corrected value even -// when doing sRGB operations. -#define STBIR_FLAG_ALPHA_USES_COLORSPACE (1 << 1) - -STBIRDEF int stbir_resize_uint8_srgb(const unsigned char *input_pixels , int input_w , int input_h , int input_stride_in_bytes, - unsigned char *output_pixels, int output_w, int output_h, int output_stride_in_bytes, - int num_channels, int alpha_channel, int flags); - - -typedef enum -{ - STBIR_EDGE_CLAMP = 1, - STBIR_EDGE_REFLECT = 2, - STBIR_EDGE_WRAP = 3, - STBIR_EDGE_ZERO = 4, -} stbir_edge; - -// This function adds the ability to specify how requests to sample off the edge of the image are handled. -STBIRDEF int stbir_resize_uint8_srgb_edgemode(const unsigned char *input_pixels , int input_w , int input_h , int input_stride_in_bytes, - unsigned char *output_pixels, int output_w, int output_h, int output_stride_in_bytes, - int num_channels, int alpha_channel, int flags, - stbir_edge edge_wrap_mode); - -////////////////////////////////////////////////////////////////////////////// -// -// Medium-complexity API -// -// This extends the easy-to-use API as follows: -// -// * Alpha-channel can be processed separately -// * If alpha_channel is not STBIR_ALPHA_CHANNEL_NONE -// * Alpha channel will not be gamma corrected (unless flags&STBIR_FLAG_GAMMA_CORRECT) -// * Filters will be weighted by alpha channel (unless flags&STBIR_FLAG_ALPHA_PREMULTIPLIED) -// * Filter can be selected explicitly -// * uint16 image type -// * sRGB colorspace available for all types -// * context parameter for passing to STBIR_MALLOC - -typedef enum -{ - STBIR_FILTER_DEFAULT = 0, // use same filter type that easy-to-use API chooses - STBIR_FILTER_BOX = 1, // A trapezoid w/1-pixel wide ramps, same result as box for integer scale ratios - STBIR_FILTER_TRIANGLE = 2, // On upsampling, produces same results as bilinear texture filtering - STBIR_FILTER_CUBICBSPLINE = 3, // The cubic b-spline (aka Mitchell-Netrevalli with B=1,C=0), gaussian-esque - STBIR_FILTER_CATMULLROM = 4, // An interpolating cubic spline - STBIR_FILTER_MITCHELL = 5, // Mitchell-Netrevalli filter with B=1/3, C=1/3 -} stbir_filter; - -typedef enum -{ - STBIR_COLORSPACE_LINEAR, - STBIR_COLORSPACE_SRGB, - - STBIR_MAX_COLORSPACES, -} stbir_colorspace; - -// The following functions are all identical except for the type of the image data - -STBIRDEF int stbir_resize_uint8_generic( const unsigned char *input_pixels , int input_w , int input_h , int input_stride_in_bytes, - unsigned char *output_pixels, int output_w, int output_h, int output_stride_in_bytes, - int num_channels, int alpha_channel, int flags, - stbir_edge edge_wrap_mode, stbir_filter filter, stbir_colorspace space, - void *alloc_context); - -STBIRDEF int stbir_resize_uint16_generic(const stbir_uint16 *input_pixels , int input_w , int input_h , int input_stride_in_bytes, - stbir_uint16 *output_pixels , int output_w, int output_h, int output_stride_in_bytes, - int num_channels, int alpha_channel, int flags, - stbir_edge edge_wrap_mode, stbir_filter filter, stbir_colorspace space, - void *alloc_context); - -STBIRDEF int stbir_resize_float_generic( const float *input_pixels , int input_w , int input_h , int input_stride_in_bytes, - float *output_pixels , int output_w, int output_h, int output_stride_in_bytes, - int num_channels, int alpha_channel, int flags, - stbir_edge edge_wrap_mode, stbir_filter filter, stbir_colorspace space, - void *alloc_context); - - - -////////////////////////////////////////////////////////////////////////////// -// -// Full-complexity API -// -// This extends the medium API as follows: -// -// * uint32 image type -// * not typesafe -// * separate filter types for each axis -// * separate edge modes for each axis -// * can specify scale explicitly for subpixel correctness -// * can specify image source tile using texture coordinates - -typedef enum -{ - STBIR_TYPE_UINT8 , - STBIR_TYPE_UINT16, - STBIR_TYPE_UINT32, - STBIR_TYPE_FLOAT , - - STBIR_MAX_TYPES -} stbir_datatype; - -STBIRDEF int stbir_resize( const void *input_pixels , int input_w , int input_h , int input_stride_in_bytes, - void *output_pixels, int output_w, int output_h, int output_stride_in_bytes, - stbir_datatype datatype, - int num_channels, int alpha_channel, int flags, - stbir_edge edge_mode_horizontal, stbir_edge edge_mode_vertical, - stbir_filter filter_horizontal, stbir_filter filter_vertical, - stbir_colorspace space, void *alloc_context); - -STBIRDEF int stbir_resize_subpixel(const void *input_pixels , int input_w , int input_h , int input_stride_in_bytes, - void *output_pixels, int output_w, int output_h, int output_stride_in_bytes, - stbir_datatype datatype, - int num_channels, int alpha_channel, int flags, - stbir_edge edge_mode_horizontal, stbir_edge edge_mode_vertical, - stbir_filter filter_horizontal, stbir_filter filter_vertical, - stbir_colorspace space, void *alloc_context, - float x_scale, float y_scale, - float x_offset, float y_offset); - -STBIRDEF int stbir_resize_region( const void *input_pixels , int input_w , int input_h , int input_stride_in_bytes, - void *output_pixels, int output_w, int output_h, int output_stride_in_bytes, - stbir_datatype datatype, - int num_channels, int alpha_channel, int flags, - stbir_edge edge_mode_horizontal, stbir_edge edge_mode_vertical, - stbir_filter filter_horizontal, stbir_filter filter_vertical, - stbir_colorspace space, void *alloc_context, - float s0, float t0, float s1, float t1); -// (s0, t0) & (s1, t1) are the top-left and bottom right corner (uv addressing style: [0, 1]x[0, 1]) of a region of the input image to use. - -// -// -//// end header file ///////////////////////////////////////////////////// -#endif // STBIR_INCLUDE_STB_IMAGE_RESIZE_H - - - - - -#ifdef STB_IMAGE_RESIZE_IMPLEMENTATION - -#ifndef STBIR_ASSERT -#include -#define STBIR_ASSERT(x) assert(x) -#endif - -// For memset -#include - -#include - -#ifndef STBIR_MALLOC -#include -// use comma operator to evaluate c, to avoid "unused parameter" warnings -#define STBIR_MALLOC(size,c) ((void)(c), malloc(size)) -#define STBIR_FREE(ptr,c) ((void)(c), free(ptr)) -#endif - -#ifndef _MSC_VER -#ifdef __cplusplus -#define stbir__inline inline -#else -#define stbir__inline -#endif -#else -#define stbir__inline __forceinline -#endif - - -// should produce compiler error if size is wrong -typedef unsigned char stbir__validate_uint32[sizeof(stbir_uint32) == 4 ? 1 : -1]; - -#ifdef _MSC_VER -#define STBIR__NOTUSED(v) (void)(v) -#else -#define STBIR__NOTUSED(v) (void)sizeof(v) -#endif - -#define STBIR__ARRAY_SIZE(a) (sizeof((a))/sizeof((a)[0])) - -#ifndef STBIR_DEFAULT_FILTER_UPSAMPLE -#define STBIR_DEFAULT_FILTER_UPSAMPLE STBIR_FILTER_CATMULLROM -#endif - -#ifndef STBIR_DEFAULT_FILTER_DOWNSAMPLE -#define STBIR_DEFAULT_FILTER_DOWNSAMPLE STBIR_FILTER_MITCHELL -#endif - -#ifndef STBIR_PROGRESS_REPORT -#define STBIR_PROGRESS_REPORT(float_0_to_1) -#endif - -#ifndef STBIR_MAX_CHANNELS -#define STBIR_MAX_CHANNELS 64 -#endif - -#if STBIR_MAX_CHANNELS > 65536 -#error "Too many channels; STBIR_MAX_CHANNELS must be no more than 65536." -// because we store the indices in 16-bit variables -#endif - -// This value is added to alpha just before premultiplication to avoid -// zeroing out color values. It is equivalent to 2^-80. If you don't want -// that behavior (it may interfere if you have floating point images with -// very small alpha values) then you can define STBIR_NO_ALPHA_EPSILON to -// disable it. -#ifndef STBIR_ALPHA_EPSILON -#define STBIR_ALPHA_EPSILON ((float)1 / (1 << 20) / (1 << 20) / (1 << 20) / (1 << 20)) -#endif - - - -#ifdef _MSC_VER -#define STBIR__UNUSED_PARAM(v) (void)(v) -#else -#define STBIR__UNUSED_PARAM(v) (void)sizeof(v) -#endif - -// must match stbir_datatype -static unsigned char stbir__type_size[] = { - 1, // STBIR_TYPE_UINT8 - 2, // STBIR_TYPE_UINT16 - 4, // STBIR_TYPE_UINT32 - 4, // STBIR_TYPE_FLOAT -}; - -// Kernel function centered at 0 -typedef float (stbir__kernel_fn)(float x, float scale); -typedef float (stbir__support_fn)(float scale); - -typedef struct -{ - stbir__kernel_fn* kernel; - stbir__support_fn* support; -} stbir__filter_info; - -// When upsampling, the contributors are which source pixels contribute. -// When downsampling, the contributors are which destination pixels are contributed to. -typedef struct -{ - int n0; // First contributing pixel - int n1; // Last contributing pixel -} stbir__contributors; - -typedef struct -{ - const void* input_data; - int input_w; - int input_h; - int input_stride_bytes; - - void* output_data; - int output_w; - int output_h; - int output_stride_bytes; - - float s0, t0, s1, t1; - - float horizontal_shift; // Units: output pixels - float vertical_shift; // Units: output pixels - float horizontal_scale; - float vertical_scale; - - int channels; - int alpha_channel; - stbir_uint32 flags; - stbir_datatype type; - stbir_filter horizontal_filter; - stbir_filter vertical_filter; - stbir_edge edge_horizontal; - stbir_edge edge_vertical; - stbir_colorspace colorspace; - - stbir__contributors* horizontal_contributors; - float* horizontal_coefficients; - - stbir__contributors* vertical_contributors; - float* vertical_coefficients; - - int decode_buffer_pixels; - float* decode_buffer; - - float* horizontal_buffer; - - // cache these because ceil/floor are inexplicably showing up in profile - int horizontal_coefficient_width; - int vertical_coefficient_width; - int horizontal_filter_pixel_width; - int vertical_filter_pixel_width; - int horizontal_filter_pixel_margin; - int vertical_filter_pixel_margin; - int horizontal_num_contributors; - int vertical_num_contributors; - - int ring_buffer_length_bytes; // The length of an individual entry in the ring buffer. The total number of ring buffers is stbir__get_filter_pixel_width(filter) - int ring_buffer_num_entries; // Total number of entries in the ring buffer. - int ring_buffer_first_scanline; - int ring_buffer_last_scanline; - int ring_buffer_begin_index; // first_scanline is at this index in the ring buffer - float* ring_buffer; - - float* encode_buffer; // A temporary buffer to store floats so we don't lose precision while we do multiply-adds. - - int horizontal_contributors_size; - int horizontal_coefficients_size; - int vertical_contributors_size; - int vertical_coefficients_size; - int decode_buffer_size; - int horizontal_buffer_size; - int ring_buffer_size; - int encode_buffer_size; -} stbir__info; - - -static const float stbir__max_uint8_as_float = 255.0f; -static const float stbir__max_uint16_as_float = 65535.0f; -static const double stbir__max_uint32_as_float = 4294967295.0; - - -static stbir__inline int stbir__min(int a, int b) -{ - return a < b ? a : b; -} - -static stbir__inline float stbir__saturate(float x) -{ - if (x < 0) - return 0; - - if (x > 1) - return 1; - - return x; -} - -#ifdef STBIR_SATURATE_INT -static stbir__inline stbir_uint8 stbir__saturate8(int x) -{ - if ((unsigned int) x <= 255) - return x; - - if (x < 0) - return 0; - - return 255; -} - -static stbir__inline stbir_uint16 stbir__saturate16(int x) -{ - if ((unsigned int) x <= 65535) - return x; - - if (x < 0) - return 0; - - return 65535; -} -#endif - -static float stbir__srgb_uchar_to_linear_float[256] = { - 0.000000f, 0.000304f, 0.000607f, 0.000911f, 0.001214f, 0.001518f, 0.001821f, 0.002125f, 0.002428f, 0.002732f, 0.003035f, - 0.003347f, 0.003677f, 0.004025f, 0.004391f, 0.004777f, 0.005182f, 0.005605f, 0.006049f, 0.006512f, 0.006995f, 0.007499f, - 0.008023f, 0.008568f, 0.009134f, 0.009721f, 0.010330f, 0.010960f, 0.011612f, 0.012286f, 0.012983f, 0.013702f, 0.014444f, - 0.015209f, 0.015996f, 0.016807f, 0.017642f, 0.018500f, 0.019382f, 0.020289f, 0.021219f, 0.022174f, 0.023153f, 0.024158f, - 0.025187f, 0.026241f, 0.027321f, 0.028426f, 0.029557f, 0.030713f, 0.031896f, 0.033105f, 0.034340f, 0.035601f, 0.036889f, - 0.038204f, 0.039546f, 0.040915f, 0.042311f, 0.043735f, 0.045186f, 0.046665f, 0.048172f, 0.049707f, 0.051269f, 0.052861f, - 0.054480f, 0.056128f, 0.057805f, 0.059511f, 0.061246f, 0.063010f, 0.064803f, 0.066626f, 0.068478f, 0.070360f, 0.072272f, - 0.074214f, 0.076185f, 0.078187f, 0.080220f, 0.082283f, 0.084376f, 0.086500f, 0.088656f, 0.090842f, 0.093059f, 0.095307f, - 0.097587f, 0.099899f, 0.102242f, 0.104616f, 0.107023f, 0.109462f, 0.111932f, 0.114435f, 0.116971f, 0.119538f, 0.122139f, - 0.124772f, 0.127438f, 0.130136f, 0.132868f, 0.135633f, 0.138432f, 0.141263f, 0.144128f, 0.147027f, 0.149960f, 0.152926f, - 0.155926f, 0.158961f, 0.162029f, 0.165132f, 0.168269f, 0.171441f, 0.174647f, 0.177888f, 0.181164f, 0.184475f, 0.187821f, - 0.191202f, 0.194618f, 0.198069f, 0.201556f, 0.205079f, 0.208637f, 0.212231f, 0.215861f, 0.219526f, 0.223228f, 0.226966f, - 0.230740f, 0.234551f, 0.238398f, 0.242281f, 0.246201f, 0.250158f, 0.254152f, 0.258183f, 0.262251f, 0.266356f, 0.270498f, - 0.274677f, 0.278894f, 0.283149f, 0.287441f, 0.291771f, 0.296138f, 0.300544f, 0.304987f, 0.309469f, 0.313989f, 0.318547f, - 0.323143f, 0.327778f, 0.332452f, 0.337164f, 0.341914f, 0.346704f, 0.351533f, 0.356400f, 0.361307f, 0.366253f, 0.371238f, - 0.376262f, 0.381326f, 0.386430f, 0.391573f, 0.396755f, 0.401978f, 0.407240f, 0.412543f, 0.417885f, 0.423268f, 0.428691f, - 0.434154f, 0.439657f, 0.445201f, 0.450786f, 0.456411f, 0.462077f, 0.467784f, 0.473532f, 0.479320f, 0.485150f, 0.491021f, - 0.496933f, 0.502887f, 0.508881f, 0.514918f, 0.520996f, 0.527115f, 0.533276f, 0.539480f, 0.545725f, 0.552011f, 0.558340f, - 0.564712f, 0.571125f, 0.577581f, 0.584078f, 0.590619f, 0.597202f, 0.603827f, 0.610496f, 0.617207f, 0.623960f, 0.630757f, - 0.637597f, 0.644480f, 0.651406f, 0.658375f, 0.665387f, 0.672443f, 0.679543f, 0.686685f, 0.693872f, 0.701102f, 0.708376f, - 0.715694f, 0.723055f, 0.730461f, 0.737911f, 0.745404f, 0.752942f, 0.760525f, 0.768151f, 0.775822f, 0.783538f, 0.791298f, - 0.799103f, 0.806952f, 0.814847f, 0.822786f, 0.830770f, 0.838799f, 0.846873f, 0.854993f, 0.863157f, 0.871367f, 0.879622f, - 0.887923f, 0.896269f, 0.904661f, 0.913099f, 0.921582f, 0.930111f, 0.938686f, 0.947307f, 0.955974f, 0.964686f, 0.973445f, - 0.982251f, 0.991102f, 1.0f -}; - -static float stbir__srgb_to_linear(float f) -{ - if (f <= 0.04045f) - return f / 12.92f; - else - return (float)pow((f + 0.055f) / 1.055f, 2.4f); -} - -static float stbir__linear_to_srgb(float f) -{ - if (f <= 0.0031308f) - return f * 12.92f; - else - return 1.055f * (float)pow(f, 1 / 2.4f) - 0.055f; -} - -#ifndef STBIR_NON_IEEE_FLOAT -// From https://gist.github.com/rygorous/2203834 - -typedef union -{ - stbir_uint32 u; - float f; -} stbir__FP32; - -static const stbir_uint32 fp32_to_srgb8_tab4[104] = { - 0x0073000d, 0x007a000d, 0x0080000d, 0x0087000d, 0x008d000d, 0x0094000d, 0x009a000d, 0x00a1000d, - 0x00a7001a, 0x00b4001a, 0x00c1001a, 0x00ce001a, 0x00da001a, 0x00e7001a, 0x00f4001a, 0x0101001a, - 0x010e0033, 0x01280033, 0x01410033, 0x015b0033, 0x01750033, 0x018f0033, 0x01a80033, 0x01c20033, - 0x01dc0067, 0x020f0067, 0x02430067, 0x02760067, 0x02aa0067, 0x02dd0067, 0x03110067, 0x03440067, - 0x037800ce, 0x03df00ce, 0x044600ce, 0x04ad00ce, 0x051400ce, 0x057b00c5, 0x05dd00bc, 0x063b00b5, - 0x06970158, 0x07420142, 0x07e30130, 0x087b0120, 0x090b0112, 0x09940106, 0x0a1700fc, 0x0a9500f2, - 0x0b0f01cb, 0x0bf401ae, 0x0ccb0195, 0x0d950180, 0x0e56016e, 0x0f0d015e, 0x0fbc0150, 0x10630143, - 0x11070264, 0x1238023e, 0x1357021d, 0x14660201, 0x156601e9, 0x165a01d3, 0x174401c0, 0x182401af, - 0x18fe0331, 0x1a9602fe, 0x1c1502d2, 0x1d7e02ad, 0x1ed4028d, 0x201a0270, 0x21520256, 0x227d0240, - 0x239f0443, 0x25c003fe, 0x27bf03c4, 0x29a10392, 0x2b6a0367, 0x2d1d0341, 0x2ebe031f, 0x304d0300, - 0x31d105b0, 0x34a80555, 0x37520507, 0x39d504c5, 0x3c37048b, 0x3e7c0458, 0x40a8042a, 0x42bd0401, - 0x44c20798, 0x488e071e, 0x4c1c06b6, 0x4f76065d, 0x52a50610, 0x55ac05cc, 0x5892058f, 0x5b590559, - 0x5e0c0a23, 0x631c0980, 0x67db08f6, 0x6c55087f, 0x70940818, 0x74a007bd, 0x787d076c, 0x7c330723, -}; - -static stbir_uint8 stbir__linear_to_srgb_uchar(float in) -{ - static const stbir__FP32 almostone = { 0x3f7fffff }; // 1-eps - static const stbir__FP32 minval = { (127-13) << 23 }; - stbir_uint32 tab,bias,scale,t; - stbir__FP32 f; - - // Clamp to [2^(-13), 1-eps]; these two values map to 0 and 1, respectively. - // The tests are carefully written so that NaNs map to 0, same as in the reference - // implementation. - if (!(in > minval.f)) // written this way to catch NaNs - in = minval.f; - if (in > almostone.f) - in = almostone.f; - - // Do the table lookup and unpack bias, scale - f.f = in; - tab = fp32_to_srgb8_tab4[(f.u - minval.u) >> 20]; - bias = (tab >> 16) << 9; - scale = tab & 0xffff; - - // Grab next-highest mantissa bits and perform linear interpolation - t = (f.u >> 12) & 0xff; - return (unsigned char) ((bias + scale*t) >> 16); -} - -#else -// sRGB transition values, scaled by 1<<28 -static int stbir__srgb_offset_to_linear_scaled[256] = -{ - 0, 40738, 122216, 203693, 285170, 366648, 448125, 529603, - 611080, 692557, 774035, 855852, 942009, 1033024, 1128971, 1229926, - 1335959, 1447142, 1563542, 1685229, 1812268, 1944725, 2082664, 2226148, - 2375238, 2529996, 2690481, 2856753, 3028870, 3206888, 3390865, 3580856, - 3776916, 3979100, 4187460, 4402049, 4622919, 4850123, 5083710, 5323731, - 5570236, 5823273, 6082892, 6349140, 6622065, 6901714, 7188133, 7481369, - 7781466, 8088471, 8402427, 8723380, 9051372, 9386448, 9728650, 10078021, - 10434603, 10798439, 11169569, 11548036, 11933879, 12327139, 12727857, 13136073, - 13551826, 13975156, 14406100, 14844697, 15290987, 15745007, 16206795, 16676389, - 17153826, 17639142, 18132374, 18633560, 19142734, 19659934, 20185196, 20718552, - 21260042, 21809696, 22367554, 22933648, 23508010, 24090680, 24681686, 25281066, - 25888850, 26505076, 27129772, 27762974, 28404716, 29055026, 29713942, 30381490, - 31057708, 31742624, 32436272, 33138682, 33849884, 34569912, 35298800, 36036568, - 36783260, 37538896, 38303512, 39077136, 39859796, 40651528, 41452360, 42262316, - 43081432, 43909732, 44747252, 45594016, 46450052, 47315392, 48190064, 49074096, - 49967516, 50870356, 51782636, 52704392, 53635648, 54576432, 55526772, 56486700, - 57456236, 58435408, 59424248, 60422780, 61431036, 62449032, 63476804, 64514376, - 65561776, 66619028, 67686160, 68763192, 69850160, 70947088, 72053992, 73170912, - 74297864, 75434880, 76581976, 77739184, 78906536, 80084040, 81271736, 82469648, - 83677792, 84896192, 86124888, 87363888, 88613232, 89872928, 91143016, 92423512, - 93714432, 95015816, 96327688, 97650056, 98982952, 100326408, 101680440, 103045072, - 104420320, 105806224, 107202800, 108610064, 110028048, 111456776, 112896264, 114346544, - 115807632, 117279552, 118762328, 120255976, 121760536, 123276016, 124802440, 126339832, - 127888216, 129447616, 131018048, 132599544, 134192112, 135795792, 137410592, 139036528, - 140673648, 142321952, 143981456, 145652208, 147334208, 149027488, 150732064, 152447968, - 154175200, 155913792, 157663776, 159425168, 161197984, 162982240, 164777968, 166585184, - 168403904, 170234160, 172075968, 173929344, 175794320, 177670896, 179559120, 181458992, - 183370528, 185293776, 187228736, 189175424, 191133888, 193104112, 195086128, 197079968, - 199085648, 201103184, 203132592, 205173888, 207227120, 209292272, 211369392, 213458480, - 215559568, 217672656, 219797792, 221934976, 224084240, 226245600, 228419056, 230604656, - 232802400, 235012320, 237234432, 239468736, 241715280, 243974080, 246245120, 248528464, - 250824112, 253132064, 255452368, 257785040, 260130080, 262487520, 264857376, 267239664, -}; - -static stbir_uint8 stbir__linear_to_srgb_uchar(float f) -{ - int x = (int) (f * (1 << 28)); // has headroom so you don't need to clamp - int v = 0; - int i; - - // Refine the guess with a short binary search. - i = v + 128; if (x >= stbir__srgb_offset_to_linear_scaled[i]) v = i; - i = v + 64; if (x >= stbir__srgb_offset_to_linear_scaled[i]) v = i; - i = v + 32; if (x >= stbir__srgb_offset_to_linear_scaled[i]) v = i; - i = v + 16; if (x >= stbir__srgb_offset_to_linear_scaled[i]) v = i; - i = v + 8; if (x >= stbir__srgb_offset_to_linear_scaled[i]) v = i; - i = v + 4; if (x >= stbir__srgb_offset_to_linear_scaled[i]) v = i; - i = v + 2; if (x >= stbir__srgb_offset_to_linear_scaled[i]) v = i; - i = v + 1; if (x >= stbir__srgb_offset_to_linear_scaled[i]) v = i; - - return (stbir_uint8) v; -} -#endif - -static float stbir__filter_trapezoid(float x, float scale) -{ - float halfscale = scale / 2; - float t = 0.5f + halfscale; - STBIR_ASSERT(scale <= 1); - - x = (float)fabs(x); - - if (x >= t) - return 0; - else - { - float r = 0.5f - halfscale; - if (x <= r) - return 1; - else - return (t - x) / scale; - } -} - -static float stbir__support_trapezoid(float scale) -{ - STBIR_ASSERT(scale <= 1); - return 0.5f + scale / 2; -} - -static float stbir__filter_triangle(float x, float s) -{ - STBIR__UNUSED_PARAM(s); - - x = (float)fabs(x); - - if (x <= 1.0f) - return 1 - x; - else - return 0; -} - -static float stbir__filter_cubic(float x, float s) -{ - STBIR__UNUSED_PARAM(s); - - x = (float)fabs(x); - - if (x < 1.0f) - return (4 + x*x*(3*x - 6))/6; - else if (x < 2.0f) - return (8 + x*(-12 + x*(6 - x)))/6; - - return (0.0f); -} - -static float stbir__filter_catmullrom(float x, float s) -{ - STBIR__UNUSED_PARAM(s); - - x = (float)fabs(x); - - if (x < 1.0f) - return 1 - x*x*(2.5f - 1.5f*x); - else if (x < 2.0f) - return 2 - x*(4 + x*(0.5f*x - 2.5f)); - - return (0.0f); -} - -static float stbir__filter_mitchell(float x, float s) -{ - STBIR__UNUSED_PARAM(s); - - x = (float)fabs(x); - - if (x < 1.0f) - return (16 + x*x*(21 * x - 36))/18; - else if (x < 2.0f) - return (32 + x*(-60 + x*(36 - 7*x)))/18; - - return (0.0f); -} - -static float stbir__support_zero(float s) -{ - STBIR__UNUSED_PARAM(s); - return 0; -} - -static float stbir__support_one(float s) -{ - STBIR__UNUSED_PARAM(s); - return 1; -} - -static float stbir__support_two(float s) -{ - STBIR__UNUSED_PARAM(s); - return 2; -} - -static stbir__filter_info stbir__filter_info_table[] = { - { NULL, stbir__support_zero }, - { stbir__filter_trapezoid, stbir__support_trapezoid }, - { stbir__filter_triangle, stbir__support_one }, - { stbir__filter_cubic, stbir__support_two }, - { stbir__filter_catmullrom, stbir__support_two }, - { stbir__filter_mitchell, stbir__support_two }, -}; - -stbir__inline static int stbir__use_upsampling(float ratio) -{ - return ratio > 1; -} - -stbir__inline static int stbir__use_width_upsampling(stbir__info* stbir_info) -{ - return stbir__use_upsampling(stbir_info->horizontal_scale); -} - -stbir__inline static int stbir__use_height_upsampling(stbir__info* stbir_info) -{ - return stbir__use_upsampling(stbir_info->vertical_scale); -} - -// This is the maximum number of input samples that can affect an output sample -// with the given filter -static int stbir__get_filter_pixel_width(stbir_filter filter, float scale) -{ - STBIR_ASSERT(filter != 0); - STBIR_ASSERT(filter < STBIR__ARRAY_SIZE(stbir__filter_info_table)); - - if (stbir__use_upsampling(scale)) - return (int)ceil(stbir__filter_info_table[filter].support(1/scale) * 2); - else - return (int)ceil(stbir__filter_info_table[filter].support(scale) * 2 / scale); -} - -// This is how much to expand buffers to account for filters seeking outside -// the image boundaries. -static int stbir__get_filter_pixel_margin(stbir_filter filter, float scale) -{ - return stbir__get_filter_pixel_width(filter, scale) / 2; -} - -static int stbir__get_coefficient_width(stbir_filter filter, float scale) -{ - if (stbir__use_upsampling(scale)) - return (int)ceil(stbir__filter_info_table[filter].support(1 / scale) * 2); - else - return (int)ceil(stbir__filter_info_table[filter].support(scale) * 2); -} - -static int stbir__get_contributors(float scale, stbir_filter filter, int input_size, int output_size) -{ - if (stbir__use_upsampling(scale)) - return output_size; - else - return (input_size + stbir__get_filter_pixel_margin(filter, scale) * 2); -} - -static int stbir__get_total_horizontal_coefficients(stbir__info* info) -{ - return info->horizontal_num_contributors - * stbir__get_coefficient_width (info->horizontal_filter, info->horizontal_scale); -} - -static int stbir__get_total_vertical_coefficients(stbir__info* info) -{ - return info->vertical_num_contributors - * stbir__get_coefficient_width (info->vertical_filter, info->vertical_scale); -} - -static stbir__contributors* stbir__get_contributor(stbir__contributors* contributors, int n) -{ - return &contributors[n]; -} - -// For perf reasons this code is duplicated in stbir__resample_horizontal_upsample/downsample, -// if you change it here change it there too. -static float* stbir__get_coefficient(float* coefficients, stbir_filter filter, float scale, int n, int c) -{ - int width = stbir__get_coefficient_width(filter, scale); - return &coefficients[width*n + c]; -} - -static int stbir__edge_wrap_slow(stbir_edge edge, int n, int max) -{ - switch (edge) - { - case STBIR_EDGE_ZERO: - return 0; // we'll decode the wrong pixel here, and then overwrite with 0s later - - case STBIR_EDGE_CLAMP: - if (n < 0) - return 0; - - if (n >= max) - return max - 1; - - return n; // NOTREACHED - - case STBIR_EDGE_REFLECT: - { - if (n < 0) - { - if (n < max) - return -n; - else - return max - 1; - } - - if (n >= max) - { - int max2 = max * 2; - if (n >= max2) - return 0; - else - return max2 - n - 1; - } - - return n; // NOTREACHED - } - - case STBIR_EDGE_WRAP: - if (n >= 0) - return (n % max); - else - { - int m = (-n) % max; - - if (m != 0) - m = max - m; - - return (m); - } - // NOTREACHED - - default: - STBIR_ASSERT(!"Unimplemented edge type"); - return 0; - } -} - -stbir__inline static int stbir__edge_wrap(stbir_edge edge, int n, int max) -{ - // avoid per-pixel switch - if (n >= 0 && n < max) - return n; - return stbir__edge_wrap_slow(edge, n, max); -} - -// What input pixels contribute to this output pixel? -static void stbir__calculate_sample_range_upsample(int n, float out_filter_radius, float scale_ratio, float out_shift, int* in_first_pixel, int* in_last_pixel, float* in_center_of_out) -{ - float out_pixel_center = (float)n + 0.5f; - float out_pixel_influence_lowerbound = out_pixel_center - out_filter_radius; - float out_pixel_influence_upperbound = out_pixel_center + out_filter_radius; - - float in_pixel_influence_lowerbound = (out_pixel_influence_lowerbound + out_shift) / scale_ratio; - float in_pixel_influence_upperbound = (out_pixel_influence_upperbound + out_shift) / scale_ratio; - - *in_center_of_out = (out_pixel_center + out_shift) / scale_ratio; - *in_first_pixel = (int)(floor(in_pixel_influence_lowerbound + 0.5)); - *in_last_pixel = (int)(floor(in_pixel_influence_upperbound - 0.5)); -} - -// What output pixels does this input pixel contribute to? -static void stbir__calculate_sample_range_downsample(int n, float in_pixels_radius, float scale_ratio, float out_shift, int* out_first_pixel, int* out_last_pixel, float* out_center_of_in) -{ - float in_pixel_center = (float)n + 0.5f; - float in_pixel_influence_lowerbound = in_pixel_center - in_pixels_radius; - float in_pixel_influence_upperbound = in_pixel_center + in_pixels_radius; - - float out_pixel_influence_lowerbound = in_pixel_influence_lowerbound * scale_ratio - out_shift; - float out_pixel_influence_upperbound = in_pixel_influence_upperbound * scale_ratio - out_shift; - - *out_center_of_in = in_pixel_center * scale_ratio - out_shift; - *out_first_pixel = (int)(floor(out_pixel_influence_lowerbound + 0.5)); - *out_last_pixel = (int)(floor(out_pixel_influence_upperbound - 0.5)); -} - -static void stbir__calculate_coefficients_upsample(stbir_filter filter, float scale, int in_first_pixel, int in_last_pixel, float in_center_of_out, stbir__contributors* contributor, float* coefficient_group) -{ - int i; - float total_filter = 0; - float filter_scale; - - STBIR_ASSERT(in_last_pixel - in_first_pixel <= (int)ceil(stbir__filter_info_table[filter].support(1/scale) * 2)); // Taken directly from stbir__get_coefficient_width() which we can't call because we don't know if we're horizontal or vertical. - - contributor->n0 = in_first_pixel; - contributor->n1 = in_last_pixel; - - STBIR_ASSERT(contributor->n1 >= contributor->n0); - - for (i = 0; i <= in_last_pixel - in_first_pixel; i++) - { - float in_pixel_center = (float)(i + in_first_pixel) + 0.5f; - coefficient_group[i] = stbir__filter_info_table[filter].kernel(in_center_of_out - in_pixel_center, 1 / scale); - - // If the coefficient is zero, skip it. (Don't do the <0 check here, we want the influence of those outside pixels.) - if (i == 0 && !coefficient_group[i]) - { - contributor->n0 = ++in_first_pixel; - i--; - continue; - } - - total_filter += coefficient_group[i]; - } - - // NOTE(fg): Not actually true in general, nor is there any reason to expect it should be. - // It would be true in exact math but is at best approximately true in floating-point math, - // and it would not make sense to try and put actual bounds on this here because it depends - // on the image aspect ratio which can get pretty extreme. - //STBIR_ASSERT(stbir__filter_info_table[filter].kernel((float)(in_last_pixel + 1) + 0.5f - in_center_of_out, 1/scale) == 0); - - STBIR_ASSERT(total_filter > 0.9); - STBIR_ASSERT(total_filter < 1.1f); // Make sure it's not way off. - - // Make sure the sum of all coefficients is 1. - filter_scale = 1 / total_filter; - - for (i = 0; i <= in_last_pixel - in_first_pixel; i++) - coefficient_group[i] *= filter_scale; - - for (i = in_last_pixel - in_first_pixel; i >= 0; i--) - { - if (coefficient_group[i]) - break; - - // This line has no weight. We can skip it. - contributor->n1 = contributor->n0 + i - 1; - } -} - -static void stbir__calculate_coefficients_downsample(stbir_filter filter, float scale_ratio, int out_first_pixel, int out_last_pixel, float out_center_of_in, stbir__contributors* contributor, float* coefficient_group) -{ - int i; - - STBIR_ASSERT(out_last_pixel - out_first_pixel <= (int)ceil(stbir__filter_info_table[filter].support(scale_ratio) * 2)); // Taken directly from stbir__get_coefficient_width() which we can't call because we don't know if we're horizontal or vertical. - - contributor->n0 = out_first_pixel; - contributor->n1 = out_last_pixel; - - STBIR_ASSERT(contributor->n1 >= contributor->n0); - - for (i = 0; i <= out_last_pixel - out_first_pixel; i++) - { - float out_pixel_center = (float)(i + out_first_pixel) + 0.5f; - float x = out_pixel_center - out_center_of_in; - coefficient_group[i] = stbir__filter_info_table[filter].kernel(x, scale_ratio) * scale_ratio; - } - - // NOTE(fg): Not actually true in general, nor is there any reason to expect it should be. - // It would be true in exact math but is at best approximately true in floating-point math, - // and it would not make sense to try and put actual bounds on this here because it depends - // on the image aspect ratio which can get pretty extreme. - //STBIR_ASSERT(stbir__filter_info_table[filter].kernel((float)(out_last_pixel + 1) + 0.5f - out_center_of_in, scale_ratio) == 0); - - for (i = out_last_pixel - out_first_pixel; i >= 0; i--) - { - if (coefficient_group[i]) - break; - - // This line has no weight. We can skip it. - contributor->n1 = contributor->n0 + i - 1; - } -} - -static void stbir__normalize_downsample_coefficients(stbir__contributors* contributors, float* coefficients, stbir_filter filter, float scale_ratio, int input_size, int output_size) -{ - int num_contributors = stbir__get_contributors(scale_ratio, filter, input_size, output_size); - int num_coefficients = stbir__get_coefficient_width(filter, scale_ratio); - int i, j; - int skip; - - for (i = 0; i < output_size; i++) - { - float scale; - float total = 0; - - for (j = 0; j < num_contributors; j++) - { - if (i >= contributors[j].n0 && i <= contributors[j].n1) - { - float coefficient = *stbir__get_coefficient(coefficients, filter, scale_ratio, j, i - contributors[j].n0); - total += coefficient; - } - else if (i < contributors[j].n0) - break; - } - - STBIR_ASSERT(total > 0.9f); - STBIR_ASSERT(total < 1.1f); - - scale = 1 / total; - - for (j = 0; j < num_contributors; j++) - { - if (i >= contributors[j].n0 && i <= contributors[j].n1) - *stbir__get_coefficient(coefficients, filter, scale_ratio, j, i - contributors[j].n0) *= scale; - else if (i < contributors[j].n0) - break; - } - } - - // Optimize: Skip zero coefficients and contributions outside of image bounds. - // Do this after normalizing because normalization depends on the n0/n1 values. - for (j = 0; j < num_contributors; j++) - { - int range, max, width; - - skip = 0; - while (*stbir__get_coefficient(coefficients, filter, scale_ratio, j, skip) == 0) - skip++; - - contributors[j].n0 += skip; - - while (contributors[j].n0 < 0) - { - contributors[j].n0++; - skip++; - } - - range = contributors[j].n1 - contributors[j].n0 + 1; - max = stbir__min(num_coefficients, range); - - width = stbir__get_coefficient_width(filter, scale_ratio); - for (i = 0; i < max; i++) - { - if (i + skip >= width) - break; - - *stbir__get_coefficient(coefficients, filter, scale_ratio, j, i) = *stbir__get_coefficient(coefficients, filter, scale_ratio, j, i + skip); - } - - continue; - } - - // Using min to avoid writing into invalid pixels. - for (i = 0; i < num_contributors; i++) - contributors[i].n1 = stbir__min(contributors[i].n1, output_size - 1); -} - -// Each scan line uses the same kernel values so we should calculate the kernel -// values once and then we can use them for every scan line. -static void stbir__calculate_filters(stbir__contributors* contributors, float* coefficients, stbir_filter filter, float scale_ratio, float shift, int input_size, int output_size) -{ - int n; - int total_contributors = stbir__get_contributors(scale_ratio, filter, input_size, output_size); - - if (stbir__use_upsampling(scale_ratio)) - { - float out_pixels_radius = stbir__filter_info_table[filter].support(1 / scale_ratio) * scale_ratio; - - // Looping through out pixels - for (n = 0; n < total_contributors; n++) - { - float in_center_of_out; // Center of the current out pixel in the in pixel space - int in_first_pixel, in_last_pixel; - - stbir__calculate_sample_range_upsample(n, out_pixels_radius, scale_ratio, shift, &in_first_pixel, &in_last_pixel, &in_center_of_out); - - stbir__calculate_coefficients_upsample(filter, scale_ratio, in_first_pixel, in_last_pixel, in_center_of_out, stbir__get_contributor(contributors, n), stbir__get_coefficient(coefficients, filter, scale_ratio, n, 0)); - } - } - else - { - float in_pixels_radius = stbir__filter_info_table[filter].support(scale_ratio) / scale_ratio; - - // Looping through in pixels - for (n = 0; n < total_contributors; n++) - { - float out_center_of_in; // Center of the current out pixel in the in pixel space - int out_first_pixel, out_last_pixel; - int n_adjusted = n - stbir__get_filter_pixel_margin(filter, scale_ratio); - - stbir__calculate_sample_range_downsample(n_adjusted, in_pixels_radius, scale_ratio, shift, &out_first_pixel, &out_last_pixel, &out_center_of_in); - - stbir__calculate_coefficients_downsample(filter, scale_ratio, out_first_pixel, out_last_pixel, out_center_of_in, stbir__get_contributor(contributors, n), stbir__get_coefficient(coefficients, filter, scale_ratio, n, 0)); - } - - stbir__normalize_downsample_coefficients(contributors, coefficients, filter, scale_ratio, input_size, output_size); - } -} - -static float* stbir__get_decode_buffer(stbir__info* stbir_info) -{ - // The 0 index of the decode buffer starts after the margin. This makes - // it okay to use negative indexes on the decode buffer. - return &stbir_info->decode_buffer[stbir_info->horizontal_filter_pixel_margin * stbir_info->channels]; -} - -#define STBIR__DECODE(type, colorspace) ((int)(type) * (STBIR_MAX_COLORSPACES) + (int)(colorspace)) - -static void stbir__decode_scanline(stbir__info* stbir_info, int n) -{ - int c; - int channels = stbir_info->channels; - int alpha_channel = stbir_info->alpha_channel; - int type = stbir_info->type; - int colorspace = stbir_info->colorspace; - int input_w = stbir_info->input_w; - size_t input_stride_bytes = stbir_info->input_stride_bytes; - float* decode_buffer = stbir__get_decode_buffer(stbir_info); - stbir_edge edge_horizontal = stbir_info->edge_horizontal; - stbir_edge edge_vertical = stbir_info->edge_vertical; - size_t in_buffer_row_offset = stbir__edge_wrap(edge_vertical, n, stbir_info->input_h) * input_stride_bytes; - const void* input_data = (char *) stbir_info->input_data + in_buffer_row_offset; - int max_x = input_w + stbir_info->horizontal_filter_pixel_margin; - int decode = STBIR__DECODE(type, colorspace); - - int x = -stbir_info->horizontal_filter_pixel_margin; - - // special handling for STBIR_EDGE_ZERO because it needs to return an item that doesn't appear in the input, - // and we want to avoid paying overhead on every pixel if not STBIR_EDGE_ZERO - if (edge_vertical == STBIR_EDGE_ZERO && (n < 0 || n >= stbir_info->input_h)) - { - for (; x < max_x; x++) - for (c = 0; c < channels; c++) - decode_buffer[x*channels + c] = 0; - return; - } - - switch (decode) - { - case STBIR__DECODE(STBIR_TYPE_UINT8, STBIR_COLORSPACE_LINEAR): - for (; x < max_x; x++) - { - int decode_pixel_index = x * channels; - int input_pixel_index = stbir__edge_wrap(edge_horizontal, x, input_w) * channels; - for (c = 0; c < channels; c++) - decode_buffer[decode_pixel_index + c] = ((float)((const unsigned char*)input_data)[input_pixel_index + c]) / stbir__max_uint8_as_float; - } - break; - - case STBIR__DECODE(STBIR_TYPE_UINT8, STBIR_COLORSPACE_SRGB): - for (; x < max_x; x++) - { - int decode_pixel_index = x * channels; - int input_pixel_index = stbir__edge_wrap(edge_horizontal, x, input_w) * channels; - for (c = 0; c < channels; c++) - decode_buffer[decode_pixel_index + c] = stbir__srgb_uchar_to_linear_float[((const unsigned char*)input_data)[input_pixel_index + c]]; - - if (!(stbir_info->flags&STBIR_FLAG_ALPHA_USES_COLORSPACE)) - decode_buffer[decode_pixel_index + alpha_channel] = ((float)((const unsigned char*)input_data)[input_pixel_index + alpha_channel]) / stbir__max_uint8_as_float; - } - break; - - case STBIR__DECODE(STBIR_TYPE_UINT16, STBIR_COLORSPACE_LINEAR): - for (; x < max_x; x++) - { - int decode_pixel_index = x * channels; - int input_pixel_index = stbir__edge_wrap(edge_horizontal, x, input_w) * channels; - for (c = 0; c < channels; c++) - decode_buffer[decode_pixel_index + c] = ((float)((const unsigned short*)input_data)[input_pixel_index + c]) / stbir__max_uint16_as_float; - } - break; - - case STBIR__DECODE(STBIR_TYPE_UINT16, STBIR_COLORSPACE_SRGB): - for (; x < max_x; x++) - { - int decode_pixel_index = x * channels; - int input_pixel_index = stbir__edge_wrap(edge_horizontal, x, input_w) * channels; - for (c = 0; c < channels; c++) - decode_buffer[decode_pixel_index + c] = stbir__srgb_to_linear(((float)((const unsigned short*)input_data)[input_pixel_index + c]) / stbir__max_uint16_as_float); - - if (!(stbir_info->flags&STBIR_FLAG_ALPHA_USES_COLORSPACE)) - decode_buffer[decode_pixel_index + alpha_channel] = ((float)((const unsigned short*)input_data)[input_pixel_index + alpha_channel]) / stbir__max_uint16_as_float; - } - break; - - case STBIR__DECODE(STBIR_TYPE_UINT32, STBIR_COLORSPACE_LINEAR): - for (; x < max_x; x++) - { - int decode_pixel_index = x * channels; - int input_pixel_index = stbir__edge_wrap(edge_horizontal, x, input_w) * channels; - for (c = 0; c < channels; c++) - decode_buffer[decode_pixel_index + c] = (float)(((double)((const unsigned int*)input_data)[input_pixel_index + c]) / stbir__max_uint32_as_float); - } - break; - - case STBIR__DECODE(STBIR_TYPE_UINT32, STBIR_COLORSPACE_SRGB): - for (; x < max_x; x++) - { - int decode_pixel_index = x * channels; - int input_pixel_index = stbir__edge_wrap(edge_horizontal, x, input_w) * channels; - for (c = 0; c < channels; c++) - decode_buffer[decode_pixel_index + c] = stbir__srgb_to_linear((float)(((double)((const unsigned int*)input_data)[input_pixel_index + c]) / stbir__max_uint32_as_float)); - - if (!(stbir_info->flags&STBIR_FLAG_ALPHA_USES_COLORSPACE)) - decode_buffer[decode_pixel_index + alpha_channel] = (float)(((double)((const unsigned int*)input_data)[input_pixel_index + alpha_channel]) / stbir__max_uint32_as_float); - } - break; - - case STBIR__DECODE(STBIR_TYPE_FLOAT, STBIR_COLORSPACE_LINEAR): - for (; x < max_x; x++) - { - int decode_pixel_index = x * channels; - int input_pixel_index = stbir__edge_wrap(edge_horizontal, x, input_w) * channels; - for (c = 0; c < channels; c++) - decode_buffer[decode_pixel_index + c] = ((const float*)input_data)[input_pixel_index + c]; - } - break; - - case STBIR__DECODE(STBIR_TYPE_FLOAT, STBIR_COLORSPACE_SRGB): - for (; x < max_x; x++) - { - int decode_pixel_index = x * channels; - int input_pixel_index = stbir__edge_wrap(edge_horizontal, x, input_w) * channels; - for (c = 0; c < channels; c++) - decode_buffer[decode_pixel_index + c] = stbir__srgb_to_linear(((const float*)input_data)[input_pixel_index + c]); - - if (!(stbir_info->flags&STBIR_FLAG_ALPHA_USES_COLORSPACE)) - decode_buffer[decode_pixel_index + alpha_channel] = ((const float*)input_data)[input_pixel_index + alpha_channel]; - } - - break; - - default: - STBIR_ASSERT(!"Unknown type/colorspace/channels combination."); - break; - } - - if (!(stbir_info->flags & STBIR_FLAG_ALPHA_PREMULTIPLIED)) - { - for (x = -stbir_info->horizontal_filter_pixel_margin; x < max_x; x++) - { - int decode_pixel_index = x * channels; - - // If the alpha value is 0 it will clobber the color values. Make sure it's not. - float alpha = decode_buffer[decode_pixel_index + alpha_channel]; -#ifndef STBIR_NO_ALPHA_EPSILON - if (stbir_info->type != STBIR_TYPE_FLOAT) { - alpha += STBIR_ALPHA_EPSILON; - decode_buffer[decode_pixel_index + alpha_channel] = alpha; - } -#endif - for (c = 0; c < channels; c++) - { - if (c == alpha_channel) - continue; - - decode_buffer[decode_pixel_index + c] *= alpha; - } - } - } - - if (edge_horizontal == STBIR_EDGE_ZERO) - { - for (x = -stbir_info->horizontal_filter_pixel_margin; x < 0; x++) - { - for (c = 0; c < channels; c++) - decode_buffer[x*channels + c] = 0; - } - for (x = input_w; x < max_x; x++) - { - for (c = 0; c < channels; c++) - decode_buffer[x*channels + c] = 0; - } - } -} - -static float* stbir__get_ring_buffer_entry(float* ring_buffer, int index, int ring_buffer_length) -{ - return &ring_buffer[index * ring_buffer_length]; -} - -static float* stbir__add_empty_ring_buffer_entry(stbir__info* stbir_info, int n) -{ - int ring_buffer_index; - float* ring_buffer; - - stbir_info->ring_buffer_last_scanline = n; - - if (stbir_info->ring_buffer_begin_index < 0) - { - ring_buffer_index = stbir_info->ring_buffer_begin_index = 0; - stbir_info->ring_buffer_first_scanline = n; - } - else - { - ring_buffer_index = (stbir_info->ring_buffer_begin_index + (stbir_info->ring_buffer_last_scanline - stbir_info->ring_buffer_first_scanline)) % stbir_info->ring_buffer_num_entries; - STBIR_ASSERT(ring_buffer_index != stbir_info->ring_buffer_begin_index); - } - - ring_buffer = stbir__get_ring_buffer_entry(stbir_info->ring_buffer, ring_buffer_index, stbir_info->ring_buffer_length_bytes / sizeof(float)); - memset(ring_buffer, 0, stbir_info->ring_buffer_length_bytes); - - return ring_buffer; -} - - -static void stbir__resample_horizontal_upsample(stbir__info* stbir_info, float* output_buffer) -{ - int x, k; - int output_w = stbir_info->output_w; - int channels = stbir_info->channels; - float* decode_buffer = stbir__get_decode_buffer(stbir_info); - stbir__contributors* horizontal_contributors = stbir_info->horizontal_contributors; - float* horizontal_coefficients = stbir_info->horizontal_coefficients; - int coefficient_width = stbir_info->horizontal_coefficient_width; - - for (x = 0; x < output_w; x++) - { - int n0 = horizontal_contributors[x].n0; - int n1 = horizontal_contributors[x].n1; - - int out_pixel_index = x * channels; - int coefficient_group = coefficient_width * x; - int coefficient_counter = 0; - - STBIR_ASSERT(n1 >= n0); - STBIR_ASSERT(n0 >= -stbir_info->horizontal_filter_pixel_margin); - STBIR_ASSERT(n1 >= -stbir_info->horizontal_filter_pixel_margin); - STBIR_ASSERT(n0 < stbir_info->input_w + stbir_info->horizontal_filter_pixel_margin); - STBIR_ASSERT(n1 < stbir_info->input_w + stbir_info->horizontal_filter_pixel_margin); - - switch (channels) { - case 1: - for (k = n0; k <= n1; k++) - { - int in_pixel_index = k * 1; - float coefficient = horizontal_coefficients[coefficient_group + coefficient_counter++]; - STBIR_ASSERT(coefficient != 0); - output_buffer[out_pixel_index + 0] += decode_buffer[in_pixel_index + 0] * coefficient; - } - break; - case 2: - for (k = n0; k <= n1; k++) - { - int in_pixel_index = k * 2; - float coefficient = horizontal_coefficients[coefficient_group + coefficient_counter++]; - STBIR_ASSERT(coefficient != 0); - output_buffer[out_pixel_index + 0] += decode_buffer[in_pixel_index + 0] * coefficient; - output_buffer[out_pixel_index + 1] += decode_buffer[in_pixel_index + 1] * coefficient; - } - break; - case 3: - for (k = n0; k <= n1; k++) - { - int in_pixel_index = k * 3; - float coefficient = horizontal_coefficients[coefficient_group + coefficient_counter++]; - STBIR_ASSERT(coefficient != 0); - output_buffer[out_pixel_index + 0] += decode_buffer[in_pixel_index + 0] * coefficient; - output_buffer[out_pixel_index + 1] += decode_buffer[in_pixel_index + 1] * coefficient; - output_buffer[out_pixel_index + 2] += decode_buffer[in_pixel_index + 2] * coefficient; - } - break; - case 4: - for (k = n0; k <= n1; k++) - { - int in_pixel_index = k * 4; - float coefficient = horizontal_coefficients[coefficient_group + coefficient_counter++]; - STBIR_ASSERT(coefficient != 0); - output_buffer[out_pixel_index + 0] += decode_buffer[in_pixel_index + 0] * coefficient; - output_buffer[out_pixel_index + 1] += decode_buffer[in_pixel_index + 1] * coefficient; - output_buffer[out_pixel_index + 2] += decode_buffer[in_pixel_index + 2] * coefficient; - output_buffer[out_pixel_index + 3] += decode_buffer[in_pixel_index + 3] * coefficient; - } - break; - default: - for (k = n0; k <= n1; k++) - { - int in_pixel_index = k * channels; - float coefficient = horizontal_coefficients[coefficient_group + coefficient_counter++]; - int c; - STBIR_ASSERT(coefficient != 0); - for (c = 0; c < channels; c++) - output_buffer[out_pixel_index + c] += decode_buffer[in_pixel_index + c] * coefficient; - } - break; - } - } -} - -static void stbir__resample_horizontal_downsample(stbir__info* stbir_info, float* output_buffer) -{ - int x, k; - int input_w = stbir_info->input_w; - int channels = stbir_info->channels; - float* decode_buffer = stbir__get_decode_buffer(stbir_info); - stbir__contributors* horizontal_contributors = stbir_info->horizontal_contributors; - float* horizontal_coefficients = stbir_info->horizontal_coefficients; - int coefficient_width = stbir_info->horizontal_coefficient_width; - int filter_pixel_margin = stbir_info->horizontal_filter_pixel_margin; - int max_x = input_w + filter_pixel_margin * 2; - - STBIR_ASSERT(!stbir__use_width_upsampling(stbir_info)); - - switch (channels) { - case 1: - for (x = 0; x < max_x; x++) - { - int n0 = horizontal_contributors[x].n0; - int n1 = horizontal_contributors[x].n1; - - int in_x = x - filter_pixel_margin; - int in_pixel_index = in_x * 1; - int max_n = n1; - int coefficient_group = coefficient_width * x; - - for (k = n0; k <= max_n; k++) - { - int out_pixel_index = k * 1; - float coefficient = horizontal_coefficients[coefficient_group + k - n0]; - output_buffer[out_pixel_index + 0] += decode_buffer[in_pixel_index + 0] * coefficient; - } - } - break; - - case 2: - for (x = 0; x < max_x; x++) - { - int n0 = horizontal_contributors[x].n0; - int n1 = horizontal_contributors[x].n1; - - int in_x = x - filter_pixel_margin; - int in_pixel_index = in_x * 2; - int max_n = n1; - int coefficient_group = coefficient_width * x; - - for (k = n0; k <= max_n; k++) - { - int out_pixel_index = k * 2; - float coefficient = horizontal_coefficients[coefficient_group + k - n0]; - output_buffer[out_pixel_index + 0] += decode_buffer[in_pixel_index + 0] * coefficient; - output_buffer[out_pixel_index + 1] += decode_buffer[in_pixel_index + 1] * coefficient; - } - } - break; - - case 3: - for (x = 0; x < max_x; x++) - { - int n0 = horizontal_contributors[x].n0; - int n1 = horizontal_contributors[x].n1; - - int in_x = x - filter_pixel_margin; - int in_pixel_index = in_x * 3; - int max_n = n1; - int coefficient_group = coefficient_width * x; - - for (k = n0; k <= max_n; k++) - { - int out_pixel_index = k * 3; - float coefficient = horizontal_coefficients[coefficient_group + k - n0]; - output_buffer[out_pixel_index + 0] += decode_buffer[in_pixel_index + 0] * coefficient; - output_buffer[out_pixel_index + 1] += decode_buffer[in_pixel_index + 1] * coefficient; - output_buffer[out_pixel_index + 2] += decode_buffer[in_pixel_index + 2] * coefficient; - } - } - break; - - case 4: - for (x = 0; x < max_x; x++) - { - int n0 = horizontal_contributors[x].n0; - int n1 = horizontal_contributors[x].n1; - - int in_x = x - filter_pixel_margin; - int in_pixel_index = in_x * 4; - int max_n = n1; - int coefficient_group = coefficient_width * x; - - for (k = n0; k <= max_n; k++) - { - int out_pixel_index = k * 4; - float coefficient = horizontal_coefficients[coefficient_group + k - n0]; - output_buffer[out_pixel_index + 0] += decode_buffer[in_pixel_index + 0] * coefficient; - output_buffer[out_pixel_index + 1] += decode_buffer[in_pixel_index + 1] * coefficient; - output_buffer[out_pixel_index + 2] += decode_buffer[in_pixel_index + 2] * coefficient; - output_buffer[out_pixel_index + 3] += decode_buffer[in_pixel_index + 3] * coefficient; - } - } - break; - - default: - for (x = 0; x < max_x; x++) - { - int n0 = horizontal_contributors[x].n0; - int n1 = horizontal_contributors[x].n1; - - int in_x = x - filter_pixel_margin; - int in_pixel_index = in_x * channels; - int max_n = n1; - int coefficient_group = coefficient_width * x; - - for (k = n0; k <= max_n; k++) - { - int c; - int out_pixel_index = k * channels; - float coefficient = horizontal_coefficients[coefficient_group + k - n0]; - for (c = 0; c < channels; c++) - output_buffer[out_pixel_index + c] += decode_buffer[in_pixel_index + c] * coefficient; - } - } - break; - } -} - -static void stbir__decode_and_resample_upsample(stbir__info* stbir_info, int n) -{ - // Decode the nth scanline from the source image into the decode buffer. - stbir__decode_scanline(stbir_info, n); - - // Now resample it into the ring buffer. - if (stbir__use_width_upsampling(stbir_info)) - stbir__resample_horizontal_upsample(stbir_info, stbir__add_empty_ring_buffer_entry(stbir_info, n)); - else - stbir__resample_horizontal_downsample(stbir_info, stbir__add_empty_ring_buffer_entry(stbir_info, n)); - - // Now it's sitting in the ring buffer ready to be used as source for the vertical sampling. -} - -static void stbir__decode_and_resample_downsample(stbir__info* stbir_info, int n) -{ - // Decode the nth scanline from the source image into the decode buffer. - stbir__decode_scanline(stbir_info, n); - - memset(stbir_info->horizontal_buffer, 0, stbir_info->output_w * stbir_info->channels * sizeof(float)); - - // Now resample it into the horizontal buffer. - if (stbir__use_width_upsampling(stbir_info)) - stbir__resample_horizontal_upsample(stbir_info, stbir_info->horizontal_buffer); - else - stbir__resample_horizontal_downsample(stbir_info, stbir_info->horizontal_buffer); - - // Now it's sitting in the horizontal buffer ready to be distributed into the ring buffers. -} - -// Get the specified scan line from the ring buffer. -static float* stbir__get_ring_buffer_scanline(int get_scanline, float* ring_buffer, int begin_index, int first_scanline, int ring_buffer_num_entries, int ring_buffer_length) -{ - int ring_buffer_index = (begin_index + (get_scanline - first_scanline)) % ring_buffer_num_entries; - return stbir__get_ring_buffer_entry(ring_buffer, ring_buffer_index, ring_buffer_length); -} - - -static void stbir__encode_scanline(stbir__info* stbir_info, int num_pixels, void *output_buffer, float *encode_buffer, int channels, int alpha_channel, int decode) -{ - int x; - int n; - int num_nonalpha; - stbir_uint16 nonalpha[STBIR_MAX_CHANNELS]; - - if (!(stbir_info->flags&STBIR_FLAG_ALPHA_PREMULTIPLIED)) - { - for (x=0; x < num_pixels; ++x) - { - int pixel_index = x*channels; - - float alpha = encode_buffer[pixel_index + alpha_channel]; - float reciprocal_alpha = alpha ? 1.0f / alpha : 0; - - // unrolling this produced a 1% slowdown upscaling a large RGBA linear-space image on my machine - stb - for (n = 0; n < channels; n++) - if (n != alpha_channel) - encode_buffer[pixel_index + n] *= reciprocal_alpha; - - // We added in a small epsilon to prevent the color channel from being deleted with zero alpha. - // Because we only add it for integer types, it will automatically be discarded on integer - // conversion, so we don't need to subtract it back out (which would be problematic for - // numeric precision reasons). - } - } - - // build a table of all channels that need colorspace correction, so - // we don't perform colorspace correction on channels that don't need it. - for (x = 0, num_nonalpha = 0; x < channels; ++x) - { - if (x != alpha_channel || (stbir_info->flags & STBIR_FLAG_ALPHA_USES_COLORSPACE)) - { - nonalpha[num_nonalpha++] = (stbir_uint16)x; - } - } - - #define STBIR__ROUND_INT(f) ((int) ((f)+0.5)) - #define STBIR__ROUND_UINT(f) ((stbir_uint32) ((f)+0.5)) - - #ifdef STBIR__SATURATE_INT - #define STBIR__ENCODE_LINEAR8(f) stbir__saturate8 (STBIR__ROUND_INT((f) * stbir__max_uint8_as_float )) - #define STBIR__ENCODE_LINEAR16(f) stbir__saturate16(STBIR__ROUND_INT((f) * stbir__max_uint16_as_float)) - #else - #define STBIR__ENCODE_LINEAR8(f) (unsigned char ) STBIR__ROUND_INT(stbir__saturate(f) * stbir__max_uint8_as_float ) - #define STBIR__ENCODE_LINEAR16(f) (unsigned short) STBIR__ROUND_INT(stbir__saturate(f) * stbir__max_uint16_as_float) - #endif - - switch (decode) - { - case STBIR__DECODE(STBIR_TYPE_UINT8, STBIR_COLORSPACE_LINEAR): - for (x=0; x < num_pixels; ++x) - { - int pixel_index = x*channels; - - for (n = 0; n < channels; n++) - { - int index = pixel_index + n; - ((unsigned char*)output_buffer)[index] = STBIR__ENCODE_LINEAR8(encode_buffer[index]); - } - } - break; - - case STBIR__DECODE(STBIR_TYPE_UINT8, STBIR_COLORSPACE_SRGB): - for (x=0; x < num_pixels; ++x) - { - int pixel_index = x*channels; - - for (n = 0; n < num_nonalpha; n++) - { - int index = pixel_index + nonalpha[n]; - ((unsigned char*)output_buffer)[index] = stbir__linear_to_srgb_uchar(encode_buffer[index]); - } - - if (!(stbir_info->flags & STBIR_FLAG_ALPHA_USES_COLORSPACE)) - ((unsigned char *)output_buffer)[pixel_index + alpha_channel] = STBIR__ENCODE_LINEAR8(encode_buffer[pixel_index+alpha_channel]); - } - break; - - case STBIR__DECODE(STBIR_TYPE_UINT16, STBIR_COLORSPACE_LINEAR): - for (x=0; x < num_pixels; ++x) - { - int pixel_index = x*channels; - - for (n = 0; n < channels; n++) - { - int index = pixel_index + n; - ((unsigned short*)output_buffer)[index] = STBIR__ENCODE_LINEAR16(encode_buffer[index]); - } - } - break; - - case STBIR__DECODE(STBIR_TYPE_UINT16, STBIR_COLORSPACE_SRGB): - for (x=0; x < num_pixels; ++x) - { - int pixel_index = x*channels; - - for (n = 0; n < num_nonalpha; n++) - { - int index = pixel_index + nonalpha[n]; - ((unsigned short*)output_buffer)[index] = (unsigned short)STBIR__ROUND_INT(stbir__linear_to_srgb(stbir__saturate(encode_buffer[index])) * stbir__max_uint16_as_float); - } - - if (!(stbir_info->flags&STBIR_FLAG_ALPHA_USES_COLORSPACE)) - ((unsigned short*)output_buffer)[pixel_index + alpha_channel] = STBIR__ENCODE_LINEAR16(encode_buffer[pixel_index + alpha_channel]); - } - - break; - - case STBIR__DECODE(STBIR_TYPE_UINT32, STBIR_COLORSPACE_LINEAR): - for (x=0; x < num_pixels; ++x) - { - int pixel_index = x*channels; - - for (n = 0; n < channels; n++) - { - int index = pixel_index + n; - ((unsigned int*)output_buffer)[index] = (unsigned int)STBIR__ROUND_UINT(((double)stbir__saturate(encode_buffer[index])) * stbir__max_uint32_as_float); - } - } - break; - - case STBIR__DECODE(STBIR_TYPE_UINT32, STBIR_COLORSPACE_SRGB): - for (x=0; x < num_pixels; ++x) - { - int pixel_index = x*channels; - - for (n = 0; n < num_nonalpha; n++) - { - int index = pixel_index + nonalpha[n]; - ((unsigned int*)output_buffer)[index] = (unsigned int)STBIR__ROUND_UINT(((double)stbir__linear_to_srgb(stbir__saturate(encode_buffer[index]))) * stbir__max_uint32_as_float); - } - - if (!(stbir_info->flags&STBIR_FLAG_ALPHA_USES_COLORSPACE)) - ((unsigned int*)output_buffer)[pixel_index + alpha_channel] = (unsigned int)STBIR__ROUND_INT(((double)stbir__saturate(encode_buffer[pixel_index + alpha_channel])) * stbir__max_uint32_as_float); - } - break; - - case STBIR__DECODE(STBIR_TYPE_FLOAT, STBIR_COLORSPACE_LINEAR): - for (x=0; x < num_pixels; ++x) - { - int pixel_index = x*channels; - - for (n = 0; n < channels; n++) - { - int index = pixel_index + n; - ((float*)output_buffer)[index] = encode_buffer[index]; - } - } - break; - - case STBIR__DECODE(STBIR_TYPE_FLOAT, STBIR_COLORSPACE_SRGB): - for (x=0; x < num_pixels; ++x) - { - int pixel_index = x*channels; - - for (n = 0; n < num_nonalpha; n++) - { - int index = pixel_index + nonalpha[n]; - ((float*)output_buffer)[index] = stbir__linear_to_srgb(encode_buffer[index]); - } - - if (!(stbir_info->flags&STBIR_FLAG_ALPHA_USES_COLORSPACE)) - ((float*)output_buffer)[pixel_index + alpha_channel] = encode_buffer[pixel_index + alpha_channel]; - } - break; - - default: - STBIR_ASSERT(!"Unknown type/colorspace/channels combination."); - break; - } -} - -static void stbir__resample_vertical_upsample(stbir__info* stbir_info, int n) -{ - int x, k; - int output_w = stbir_info->output_w; - stbir__contributors* vertical_contributors = stbir_info->vertical_contributors; - float* vertical_coefficients = stbir_info->vertical_coefficients; - int channels = stbir_info->channels; - int alpha_channel = stbir_info->alpha_channel; - int type = stbir_info->type; - int colorspace = stbir_info->colorspace; - int ring_buffer_entries = stbir_info->ring_buffer_num_entries; - void* output_data = stbir_info->output_data; - float* encode_buffer = stbir_info->encode_buffer; - int decode = STBIR__DECODE(type, colorspace); - int coefficient_width = stbir_info->vertical_coefficient_width; - int coefficient_counter; - int contributor = n; - - float* ring_buffer = stbir_info->ring_buffer; - int ring_buffer_begin_index = stbir_info->ring_buffer_begin_index; - int ring_buffer_first_scanline = stbir_info->ring_buffer_first_scanline; - int ring_buffer_length = stbir_info->ring_buffer_length_bytes/sizeof(float); - - int n0,n1, output_row_start; - int coefficient_group = coefficient_width * contributor; - - n0 = vertical_contributors[contributor].n0; - n1 = vertical_contributors[contributor].n1; - - output_row_start = n * stbir_info->output_stride_bytes; - - STBIR_ASSERT(stbir__use_height_upsampling(stbir_info)); - - memset(encode_buffer, 0, output_w * sizeof(float) * channels); - - // I tried reblocking this for better cache usage of encode_buffer - // (using x_outer, k, x_inner), but it lost speed. -- stb - - coefficient_counter = 0; - switch (channels) { - case 1: - for (k = n0; k <= n1; k++) - { - int coefficient_index = coefficient_counter++; - float* ring_buffer_entry = stbir__get_ring_buffer_scanline(k, ring_buffer, ring_buffer_begin_index, ring_buffer_first_scanline, ring_buffer_entries, ring_buffer_length); - float coefficient = vertical_coefficients[coefficient_group + coefficient_index]; - for (x = 0; x < output_w; ++x) - { - int in_pixel_index = x * 1; - encode_buffer[in_pixel_index + 0] += ring_buffer_entry[in_pixel_index + 0] * coefficient; - } - } - break; - case 2: - for (k = n0; k <= n1; k++) - { - int coefficient_index = coefficient_counter++; - float* ring_buffer_entry = stbir__get_ring_buffer_scanline(k, ring_buffer, ring_buffer_begin_index, ring_buffer_first_scanline, ring_buffer_entries, ring_buffer_length); - float coefficient = vertical_coefficients[coefficient_group + coefficient_index]; - for (x = 0; x < output_w; ++x) - { - int in_pixel_index = x * 2; - encode_buffer[in_pixel_index + 0] += ring_buffer_entry[in_pixel_index + 0] * coefficient; - encode_buffer[in_pixel_index + 1] += ring_buffer_entry[in_pixel_index + 1] * coefficient; - } - } - break; - case 3: - for (k = n0; k <= n1; k++) - { - int coefficient_index = coefficient_counter++; - float* ring_buffer_entry = stbir__get_ring_buffer_scanline(k, ring_buffer, ring_buffer_begin_index, ring_buffer_first_scanline, ring_buffer_entries, ring_buffer_length); - float coefficient = vertical_coefficients[coefficient_group + coefficient_index]; - for (x = 0; x < output_w; ++x) - { - int in_pixel_index = x * 3; - encode_buffer[in_pixel_index + 0] += ring_buffer_entry[in_pixel_index + 0] * coefficient; - encode_buffer[in_pixel_index + 1] += ring_buffer_entry[in_pixel_index + 1] * coefficient; - encode_buffer[in_pixel_index + 2] += ring_buffer_entry[in_pixel_index + 2] * coefficient; - } - } - break; - case 4: - for (k = n0; k <= n1; k++) - { - int coefficient_index = coefficient_counter++; - float* ring_buffer_entry = stbir__get_ring_buffer_scanline(k, ring_buffer, ring_buffer_begin_index, ring_buffer_first_scanline, ring_buffer_entries, ring_buffer_length); - float coefficient = vertical_coefficients[coefficient_group + coefficient_index]; - for (x = 0; x < output_w; ++x) - { - int in_pixel_index = x * 4; - encode_buffer[in_pixel_index + 0] += ring_buffer_entry[in_pixel_index + 0] * coefficient; - encode_buffer[in_pixel_index + 1] += ring_buffer_entry[in_pixel_index + 1] * coefficient; - encode_buffer[in_pixel_index + 2] += ring_buffer_entry[in_pixel_index + 2] * coefficient; - encode_buffer[in_pixel_index + 3] += ring_buffer_entry[in_pixel_index + 3] * coefficient; - } - } - break; - default: - for (k = n0; k <= n1; k++) - { - int coefficient_index = coefficient_counter++; - float* ring_buffer_entry = stbir__get_ring_buffer_scanline(k, ring_buffer, ring_buffer_begin_index, ring_buffer_first_scanline, ring_buffer_entries, ring_buffer_length); - float coefficient = vertical_coefficients[coefficient_group + coefficient_index]; - for (x = 0; x < output_w; ++x) - { - int in_pixel_index = x * channels; - int c; - for (c = 0; c < channels; c++) - encode_buffer[in_pixel_index + c] += ring_buffer_entry[in_pixel_index + c] * coefficient; - } - } - break; - } - stbir__encode_scanline(stbir_info, output_w, (char *) output_data + output_row_start, encode_buffer, channels, alpha_channel, decode); -} - -static void stbir__resample_vertical_downsample(stbir__info* stbir_info, int n) -{ - int x, k; - int output_w = stbir_info->output_w; - stbir__contributors* vertical_contributors = stbir_info->vertical_contributors; - float* vertical_coefficients = stbir_info->vertical_coefficients; - int channels = stbir_info->channels; - int ring_buffer_entries = stbir_info->ring_buffer_num_entries; - float* horizontal_buffer = stbir_info->horizontal_buffer; - int coefficient_width = stbir_info->vertical_coefficient_width; - int contributor = n + stbir_info->vertical_filter_pixel_margin; - - float* ring_buffer = stbir_info->ring_buffer; - int ring_buffer_begin_index = stbir_info->ring_buffer_begin_index; - int ring_buffer_first_scanline = stbir_info->ring_buffer_first_scanline; - int ring_buffer_length = stbir_info->ring_buffer_length_bytes/sizeof(float); - int n0,n1; - - n0 = vertical_contributors[contributor].n0; - n1 = vertical_contributors[contributor].n1; - - STBIR_ASSERT(!stbir__use_height_upsampling(stbir_info)); - - for (k = n0; k <= n1; k++) - { - int coefficient_index = k - n0; - int coefficient_group = coefficient_width * contributor; - float coefficient = vertical_coefficients[coefficient_group + coefficient_index]; - - float* ring_buffer_entry = stbir__get_ring_buffer_scanline(k, ring_buffer, ring_buffer_begin_index, ring_buffer_first_scanline, ring_buffer_entries, ring_buffer_length); - - switch (channels) { - case 1: - for (x = 0; x < output_w; x++) - { - int in_pixel_index = x * 1; - ring_buffer_entry[in_pixel_index + 0] += horizontal_buffer[in_pixel_index + 0] * coefficient; - } - break; - case 2: - for (x = 0; x < output_w; x++) - { - int in_pixel_index = x * 2; - ring_buffer_entry[in_pixel_index + 0] += horizontal_buffer[in_pixel_index + 0] * coefficient; - ring_buffer_entry[in_pixel_index + 1] += horizontal_buffer[in_pixel_index + 1] * coefficient; - } - break; - case 3: - for (x = 0; x < output_w; x++) - { - int in_pixel_index = x * 3; - ring_buffer_entry[in_pixel_index + 0] += horizontal_buffer[in_pixel_index + 0] * coefficient; - ring_buffer_entry[in_pixel_index + 1] += horizontal_buffer[in_pixel_index + 1] * coefficient; - ring_buffer_entry[in_pixel_index + 2] += horizontal_buffer[in_pixel_index + 2] * coefficient; - } - break; - case 4: - for (x = 0; x < output_w; x++) - { - int in_pixel_index = x * 4; - ring_buffer_entry[in_pixel_index + 0] += horizontal_buffer[in_pixel_index + 0] * coefficient; - ring_buffer_entry[in_pixel_index + 1] += horizontal_buffer[in_pixel_index + 1] * coefficient; - ring_buffer_entry[in_pixel_index + 2] += horizontal_buffer[in_pixel_index + 2] * coefficient; - ring_buffer_entry[in_pixel_index + 3] += horizontal_buffer[in_pixel_index + 3] * coefficient; - } - break; - default: - for (x = 0; x < output_w; x++) - { - int in_pixel_index = x * channels; - - int c; - for (c = 0; c < channels; c++) - ring_buffer_entry[in_pixel_index + c] += horizontal_buffer[in_pixel_index + c] * coefficient; - } - break; - } - } -} - -static void stbir__buffer_loop_upsample(stbir__info* stbir_info) -{ - int y; - float scale_ratio = stbir_info->vertical_scale; - float out_scanlines_radius = stbir__filter_info_table[stbir_info->vertical_filter].support(1/scale_ratio) * scale_ratio; - - STBIR_ASSERT(stbir__use_height_upsampling(stbir_info)); - - for (y = 0; y < stbir_info->output_h; y++) - { - float in_center_of_out = 0; // Center of the current out scanline in the in scanline space - int in_first_scanline = 0, in_last_scanline = 0; - - stbir__calculate_sample_range_upsample(y, out_scanlines_radius, scale_ratio, stbir_info->vertical_shift, &in_first_scanline, &in_last_scanline, &in_center_of_out); - - STBIR_ASSERT(in_last_scanline - in_first_scanline + 1 <= stbir_info->ring_buffer_num_entries); - - if (stbir_info->ring_buffer_begin_index >= 0) - { - // Get rid of whatever we don't need anymore. - while (in_first_scanline > stbir_info->ring_buffer_first_scanline) - { - if (stbir_info->ring_buffer_first_scanline == stbir_info->ring_buffer_last_scanline) - { - // We just popped the last scanline off the ring buffer. - // Reset it to the empty state. - stbir_info->ring_buffer_begin_index = -1; - stbir_info->ring_buffer_first_scanline = 0; - stbir_info->ring_buffer_last_scanline = 0; - break; - } - else - { - stbir_info->ring_buffer_first_scanline++; - stbir_info->ring_buffer_begin_index = (stbir_info->ring_buffer_begin_index + 1) % stbir_info->ring_buffer_num_entries; - } - } - } - - // Load in new ones. - if (stbir_info->ring_buffer_begin_index < 0) - stbir__decode_and_resample_upsample(stbir_info, in_first_scanline); - - while (in_last_scanline > stbir_info->ring_buffer_last_scanline) - stbir__decode_and_resample_upsample(stbir_info, stbir_info->ring_buffer_last_scanline + 1); - - // Now all buffers should be ready to write a row of vertical sampling. - stbir__resample_vertical_upsample(stbir_info, y); - - STBIR_PROGRESS_REPORT((float)y / stbir_info->output_h); - } -} - -static void stbir__empty_ring_buffer(stbir__info* stbir_info, int first_necessary_scanline) -{ - int output_stride_bytes = stbir_info->output_stride_bytes; - int channels = stbir_info->channels; - int alpha_channel = stbir_info->alpha_channel; - int type = stbir_info->type; - int colorspace = stbir_info->colorspace; - int output_w = stbir_info->output_w; - void* output_data = stbir_info->output_data; - int decode = STBIR__DECODE(type, colorspace); - - float* ring_buffer = stbir_info->ring_buffer; - int ring_buffer_length = stbir_info->ring_buffer_length_bytes/sizeof(float); - - if (stbir_info->ring_buffer_begin_index >= 0) - { - // Get rid of whatever we don't need anymore. - while (first_necessary_scanline > stbir_info->ring_buffer_first_scanline) - { - if (stbir_info->ring_buffer_first_scanline >= 0 && stbir_info->ring_buffer_first_scanline < stbir_info->output_h) - { - int output_row_start = stbir_info->ring_buffer_first_scanline * output_stride_bytes; - float* ring_buffer_entry = stbir__get_ring_buffer_entry(ring_buffer, stbir_info->ring_buffer_begin_index, ring_buffer_length); - stbir__encode_scanline(stbir_info, output_w, (char *) output_data + output_row_start, ring_buffer_entry, channels, alpha_channel, decode); - STBIR_PROGRESS_REPORT((float)stbir_info->ring_buffer_first_scanline / stbir_info->output_h); - } - - if (stbir_info->ring_buffer_first_scanline == stbir_info->ring_buffer_last_scanline) - { - // We just popped the last scanline off the ring buffer. - // Reset it to the empty state. - stbir_info->ring_buffer_begin_index = -1; - stbir_info->ring_buffer_first_scanline = 0; - stbir_info->ring_buffer_last_scanline = 0; - break; - } - else - { - stbir_info->ring_buffer_first_scanline++; - stbir_info->ring_buffer_begin_index = (stbir_info->ring_buffer_begin_index + 1) % stbir_info->ring_buffer_num_entries; - } - } - } -} - -static void stbir__buffer_loop_downsample(stbir__info* stbir_info) -{ - int y; - float scale_ratio = stbir_info->vertical_scale; - int output_h = stbir_info->output_h; - float in_pixels_radius = stbir__filter_info_table[stbir_info->vertical_filter].support(scale_ratio) / scale_ratio; - int pixel_margin = stbir_info->vertical_filter_pixel_margin; - int max_y = stbir_info->input_h + pixel_margin; - - STBIR_ASSERT(!stbir__use_height_upsampling(stbir_info)); - - for (y = -pixel_margin; y < max_y; y++) - { - float out_center_of_in; // Center of the current out scanline in the in scanline space - int out_first_scanline, out_last_scanline; - - stbir__calculate_sample_range_downsample(y, in_pixels_radius, scale_ratio, stbir_info->vertical_shift, &out_first_scanline, &out_last_scanline, &out_center_of_in); - - STBIR_ASSERT(out_last_scanline - out_first_scanline + 1 <= stbir_info->ring_buffer_num_entries); - - if (out_last_scanline < 0 || out_first_scanline >= output_h) - continue; - - stbir__empty_ring_buffer(stbir_info, out_first_scanline); - - stbir__decode_and_resample_downsample(stbir_info, y); - - // Load in new ones. - if (stbir_info->ring_buffer_begin_index < 0) - stbir__add_empty_ring_buffer_entry(stbir_info, out_first_scanline); - - while (out_last_scanline > stbir_info->ring_buffer_last_scanline) - stbir__add_empty_ring_buffer_entry(stbir_info, stbir_info->ring_buffer_last_scanline + 1); - - // Now the horizontal buffer is ready to write to all ring buffer rows. - stbir__resample_vertical_downsample(stbir_info, y); - } - - stbir__empty_ring_buffer(stbir_info, stbir_info->output_h); -} - -static void stbir__setup(stbir__info *info, int input_w, int input_h, int output_w, int output_h, int channels) -{ - info->input_w = input_w; - info->input_h = input_h; - info->output_w = output_w; - info->output_h = output_h; - info->channels = channels; -} - -static void stbir__calculate_transform(stbir__info *info, float s0, float t0, float s1, float t1, float *transform) -{ - info->s0 = s0; - info->t0 = t0; - info->s1 = s1; - info->t1 = t1; - - if (transform) - { - info->horizontal_scale = transform[0]; - info->vertical_scale = transform[1]; - info->horizontal_shift = transform[2]; - info->vertical_shift = transform[3]; - } - else - { - info->horizontal_scale = ((float)info->output_w / info->input_w) / (s1 - s0); - info->vertical_scale = ((float)info->output_h / info->input_h) / (t1 - t0); - - info->horizontal_shift = s0 * info->output_w / (s1 - s0); - info->vertical_shift = t0 * info->output_h / (t1 - t0); - } -} - -static void stbir__choose_filter(stbir__info *info, stbir_filter h_filter, stbir_filter v_filter) -{ - if (h_filter == 0) - h_filter = stbir__use_upsampling(info->horizontal_scale) ? STBIR_DEFAULT_FILTER_UPSAMPLE : STBIR_DEFAULT_FILTER_DOWNSAMPLE; - if (v_filter == 0) - v_filter = stbir__use_upsampling(info->vertical_scale) ? STBIR_DEFAULT_FILTER_UPSAMPLE : STBIR_DEFAULT_FILTER_DOWNSAMPLE; - info->horizontal_filter = h_filter; - info->vertical_filter = v_filter; -} - -static stbir_uint32 stbir__calculate_memory(stbir__info *info) -{ - int pixel_margin = stbir__get_filter_pixel_margin(info->horizontal_filter, info->horizontal_scale); - int filter_height = stbir__get_filter_pixel_width(info->vertical_filter, info->vertical_scale); - - info->horizontal_num_contributors = stbir__get_contributors(info->horizontal_scale, info->horizontal_filter, info->input_w, info->output_w); - info->vertical_num_contributors = stbir__get_contributors(info->vertical_scale , info->vertical_filter , info->input_h, info->output_h); - - // One extra entry because floating point precision problems sometimes cause an extra to be necessary. - info->ring_buffer_num_entries = filter_height + 1; - - info->horizontal_contributors_size = info->horizontal_num_contributors * sizeof(stbir__contributors); - info->horizontal_coefficients_size = stbir__get_total_horizontal_coefficients(info) * sizeof(float); - info->vertical_contributors_size = info->vertical_num_contributors * sizeof(stbir__contributors); - info->vertical_coefficients_size = stbir__get_total_vertical_coefficients(info) * sizeof(float); - info->decode_buffer_size = (info->input_w + pixel_margin * 2) * info->channels * sizeof(float); - info->horizontal_buffer_size = info->output_w * info->channels * sizeof(float); - info->ring_buffer_size = info->output_w * info->channels * info->ring_buffer_num_entries * sizeof(float); - info->encode_buffer_size = info->output_w * info->channels * sizeof(float); - - STBIR_ASSERT(info->horizontal_filter != 0); - STBIR_ASSERT(info->horizontal_filter < STBIR__ARRAY_SIZE(stbir__filter_info_table)); // this now happens too late - STBIR_ASSERT(info->vertical_filter != 0); - STBIR_ASSERT(info->vertical_filter < STBIR__ARRAY_SIZE(stbir__filter_info_table)); // this now happens too late - - if (stbir__use_height_upsampling(info)) - // The horizontal buffer is for when we're downsampling the height and we - // can't output the result of sampling the decode buffer directly into the - // ring buffers. - info->horizontal_buffer_size = 0; - else - // The encode buffer is to retain precision in the height upsampling method - // and isn't used when height downsampling. - info->encode_buffer_size = 0; - - return info->horizontal_contributors_size + info->horizontal_coefficients_size - + info->vertical_contributors_size + info->vertical_coefficients_size - + info->decode_buffer_size + info->horizontal_buffer_size - + info->ring_buffer_size + info->encode_buffer_size; -} - -static int stbir__resize_allocated(stbir__info *info, - const void* input_data, int input_stride_in_bytes, - void* output_data, int output_stride_in_bytes, - int alpha_channel, stbir_uint32 flags, stbir_datatype type, - stbir_edge edge_horizontal, stbir_edge edge_vertical, stbir_colorspace colorspace, - void* tempmem, size_t tempmem_size_in_bytes) -{ - size_t memory_required = stbir__calculate_memory(info); - - int width_stride_input = input_stride_in_bytes ? input_stride_in_bytes : info->channels * info->input_w * stbir__type_size[type]; - int width_stride_output = output_stride_in_bytes ? output_stride_in_bytes : info->channels * info->output_w * stbir__type_size[type]; - -#ifdef STBIR_DEBUG_OVERWRITE_TEST -#define OVERWRITE_ARRAY_SIZE 8 - unsigned char overwrite_output_before_pre[OVERWRITE_ARRAY_SIZE]; - unsigned char overwrite_tempmem_before_pre[OVERWRITE_ARRAY_SIZE]; - unsigned char overwrite_output_after_pre[OVERWRITE_ARRAY_SIZE]; - unsigned char overwrite_tempmem_after_pre[OVERWRITE_ARRAY_SIZE]; - - size_t begin_forbidden = width_stride_output * (info->output_h - 1) + info->output_w * info->channels * stbir__type_size[type]; - memcpy(overwrite_output_before_pre, &((unsigned char*)output_data)[-OVERWRITE_ARRAY_SIZE], OVERWRITE_ARRAY_SIZE); - memcpy(overwrite_output_after_pre, &((unsigned char*)output_data)[begin_forbidden], OVERWRITE_ARRAY_SIZE); - memcpy(overwrite_tempmem_before_pre, &((unsigned char*)tempmem)[-OVERWRITE_ARRAY_SIZE], OVERWRITE_ARRAY_SIZE); - memcpy(overwrite_tempmem_after_pre, &((unsigned char*)tempmem)[tempmem_size_in_bytes], OVERWRITE_ARRAY_SIZE); -#endif - - STBIR_ASSERT(info->channels >= 0); - STBIR_ASSERT(info->channels <= STBIR_MAX_CHANNELS); - - if (info->channels < 0 || info->channels > STBIR_MAX_CHANNELS) - return 0; - - STBIR_ASSERT(info->horizontal_filter < STBIR__ARRAY_SIZE(stbir__filter_info_table)); - STBIR_ASSERT(info->vertical_filter < STBIR__ARRAY_SIZE(stbir__filter_info_table)); - - if (info->horizontal_filter >= STBIR__ARRAY_SIZE(stbir__filter_info_table)) - return 0; - if (info->vertical_filter >= STBIR__ARRAY_SIZE(stbir__filter_info_table)) - return 0; - - if (alpha_channel < 0) - flags |= STBIR_FLAG_ALPHA_USES_COLORSPACE | STBIR_FLAG_ALPHA_PREMULTIPLIED; - - if (!(flags&STBIR_FLAG_ALPHA_USES_COLORSPACE) || !(flags&STBIR_FLAG_ALPHA_PREMULTIPLIED)) { - STBIR_ASSERT(alpha_channel >= 0 && alpha_channel < info->channels); - } - - if (alpha_channel >= info->channels) - return 0; - - STBIR_ASSERT(tempmem); - - if (!tempmem) - return 0; - - STBIR_ASSERT(tempmem_size_in_bytes >= memory_required); - - if (tempmem_size_in_bytes < memory_required) - return 0; - - memset(tempmem, 0, tempmem_size_in_bytes); - - info->input_data = input_data; - info->input_stride_bytes = width_stride_input; - - info->output_data = output_data; - info->output_stride_bytes = width_stride_output; - - info->alpha_channel = alpha_channel; - info->flags = flags; - info->type = type; - info->edge_horizontal = edge_horizontal; - info->edge_vertical = edge_vertical; - info->colorspace = colorspace; - - info->horizontal_coefficient_width = stbir__get_coefficient_width (info->horizontal_filter, info->horizontal_scale); - info->vertical_coefficient_width = stbir__get_coefficient_width (info->vertical_filter , info->vertical_scale ); - info->horizontal_filter_pixel_width = stbir__get_filter_pixel_width (info->horizontal_filter, info->horizontal_scale); - info->vertical_filter_pixel_width = stbir__get_filter_pixel_width (info->vertical_filter , info->vertical_scale ); - info->horizontal_filter_pixel_margin = stbir__get_filter_pixel_margin(info->horizontal_filter, info->horizontal_scale); - info->vertical_filter_pixel_margin = stbir__get_filter_pixel_margin(info->vertical_filter , info->vertical_scale ); - - info->ring_buffer_length_bytes = info->output_w * info->channels * sizeof(float); - info->decode_buffer_pixels = info->input_w + info->horizontal_filter_pixel_margin * 2; - -#define STBIR__NEXT_MEMPTR(current, newtype) (newtype*)(((unsigned char*)current) + current##_size) - - info->horizontal_contributors = (stbir__contributors *) tempmem; - info->horizontal_coefficients = STBIR__NEXT_MEMPTR(info->horizontal_contributors, float); - info->vertical_contributors = STBIR__NEXT_MEMPTR(info->horizontal_coefficients, stbir__contributors); - info->vertical_coefficients = STBIR__NEXT_MEMPTR(info->vertical_contributors, float); - info->decode_buffer = STBIR__NEXT_MEMPTR(info->vertical_coefficients, float); - - if (stbir__use_height_upsampling(info)) - { - info->horizontal_buffer = NULL; - info->ring_buffer = STBIR__NEXT_MEMPTR(info->decode_buffer, float); - info->encode_buffer = STBIR__NEXT_MEMPTR(info->ring_buffer, float); - - STBIR_ASSERT((size_t)STBIR__NEXT_MEMPTR(info->encode_buffer, unsigned char) == (size_t)tempmem + tempmem_size_in_bytes); - } - else - { - info->horizontal_buffer = STBIR__NEXT_MEMPTR(info->decode_buffer, float); - info->ring_buffer = STBIR__NEXT_MEMPTR(info->horizontal_buffer, float); - info->encode_buffer = NULL; - - STBIR_ASSERT((size_t)STBIR__NEXT_MEMPTR(info->ring_buffer, unsigned char) == (size_t)tempmem + tempmem_size_in_bytes); - } - -#undef STBIR__NEXT_MEMPTR - - // This signals that the ring buffer is empty - info->ring_buffer_begin_index = -1; - - stbir__calculate_filters(info->horizontal_contributors, info->horizontal_coefficients, info->horizontal_filter, info->horizontal_scale, info->horizontal_shift, info->input_w, info->output_w); - stbir__calculate_filters(info->vertical_contributors, info->vertical_coefficients, info->vertical_filter, info->vertical_scale, info->vertical_shift, info->input_h, info->output_h); - - STBIR_PROGRESS_REPORT(0); - - if (stbir__use_height_upsampling(info)) - stbir__buffer_loop_upsample(info); - else - stbir__buffer_loop_downsample(info); - - STBIR_PROGRESS_REPORT(1); - -#ifdef STBIR_DEBUG_OVERWRITE_TEST - STBIR_ASSERT(memcmp(overwrite_output_before_pre, &((unsigned char*)output_data)[-OVERWRITE_ARRAY_SIZE], OVERWRITE_ARRAY_SIZE) == 0); - STBIR_ASSERT(memcmp(overwrite_output_after_pre, &((unsigned char*)output_data)[begin_forbidden], OVERWRITE_ARRAY_SIZE) == 0); - STBIR_ASSERT(memcmp(overwrite_tempmem_before_pre, &((unsigned char*)tempmem)[-OVERWRITE_ARRAY_SIZE], OVERWRITE_ARRAY_SIZE) == 0); - STBIR_ASSERT(memcmp(overwrite_tempmem_after_pre, &((unsigned char*)tempmem)[tempmem_size_in_bytes], OVERWRITE_ARRAY_SIZE) == 0); -#endif - - return 1; -} - - -static int stbir__resize_arbitrary( - void *alloc_context, - const void* input_data, int input_w, int input_h, int input_stride_in_bytes, - void* output_data, int output_w, int output_h, int output_stride_in_bytes, - float s0, float t0, float s1, float t1, float *transform, - int channels, int alpha_channel, stbir_uint32 flags, stbir_datatype type, - stbir_filter h_filter, stbir_filter v_filter, - stbir_edge edge_horizontal, stbir_edge edge_vertical, stbir_colorspace colorspace) -{ - stbir__info info; - int result; - size_t memory_required; - void* extra_memory; - - stbir__setup(&info, input_w, input_h, output_w, output_h, channels); - stbir__calculate_transform(&info, s0,t0,s1,t1,transform); - stbir__choose_filter(&info, h_filter, v_filter); - memory_required = stbir__calculate_memory(&info); - extra_memory = STBIR_MALLOC(memory_required, alloc_context); - - if (!extra_memory) - return 0; - - result = stbir__resize_allocated(&info, input_data, input_stride_in_bytes, - output_data, output_stride_in_bytes, - alpha_channel, flags, type, - edge_horizontal, edge_vertical, - colorspace, extra_memory, memory_required); - - STBIR_FREE(extra_memory, alloc_context); - - return result; -} - -STBIRDEF int stbir_resize_uint8( const unsigned char *input_pixels , int input_w , int input_h , int input_stride_in_bytes, - unsigned char *output_pixels, int output_w, int output_h, int output_stride_in_bytes, - int num_channels) -{ - return stbir__resize_arbitrary(NULL, input_pixels, input_w, input_h, input_stride_in_bytes, - output_pixels, output_w, output_h, output_stride_in_bytes, - 0,0,1,1,NULL,num_channels,-1,0, STBIR_TYPE_UINT8, STBIR_FILTER_DEFAULT, STBIR_FILTER_DEFAULT, - STBIR_EDGE_CLAMP, STBIR_EDGE_CLAMP, STBIR_COLORSPACE_LINEAR); -} - -STBIRDEF int stbir_resize_float( const float *input_pixels , int input_w , int input_h , int input_stride_in_bytes, - float *output_pixels, int output_w, int output_h, int output_stride_in_bytes, - int num_channels) -{ - return stbir__resize_arbitrary(NULL, input_pixels, input_w, input_h, input_stride_in_bytes, - output_pixels, output_w, output_h, output_stride_in_bytes, - 0,0,1,1,NULL,num_channels,-1,0, STBIR_TYPE_FLOAT, STBIR_FILTER_DEFAULT, STBIR_FILTER_DEFAULT, - STBIR_EDGE_CLAMP, STBIR_EDGE_CLAMP, STBIR_COLORSPACE_LINEAR); -} - -STBIRDEF int stbir_resize_uint8_srgb(const unsigned char *input_pixels , int input_w , int input_h , int input_stride_in_bytes, - unsigned char *output_pixels, int output_w, int output_h, int output_stride_in_bytes, - int num_channels, int alpha_channel, int flags) -{ - return stbir__resize_arbitrary(NULL, input_pixels, input_w, input_h, input_stride_in_bytes, - output_pixels, output_w, output_h, output_stride_in_bytes, - 0,0,1,1,NULL,num_channels,alpha_channel,flags, STBIR_TYPE_UINT8, STBIR_FILTER_DEFAULT, STBIR_FILTER_DEFAULT, - STBIR_EDGE_CLAMP, STBIR_EDGE_CLAMP, STBIR_COLORSPACE_SRGB); -} - -STBIRDEF int stbir_resize_uint8_srgb_edgemode(const unsigned char *input_pixels , int input_w , int input_h , int input_stride_in_bytes, - unsigned char *output_pixels, int output_w, int output_h, int output_stride_in_bytes, - int num_channels, int alpha_channel, int flags, - stbir_edge edge_wrap_mode) -{ - return stbir__resize_arbitrary(NULL, input_pixels, input_w, input_h, input_stride_in_bytes, - output_pixels, output_w, output_h, output_stride_in_bytes, - 0,0,1,1,NULL,num_channels,alpha_channel,flags, STBIR_TYPE_UINT8, STBIR_FILTER_DEFAULT, STBIR_FILTER_DEFAULT, - edge_wrap_mode, edge_wrap_mode, STBIR_COLORSPACE_SRGB); -} - -STBIRDEF int stbir_resize_uint8_generic( const unsigned char *input_pixels , int input_w , int input_h , int input_stride_in_bytes, - unsigned char *output_pixels, int output_w, int output_h, int output_stride_in_bytes, - int num_channels, int alpha_channel, int flags, - stbir_edge edge_wrap_mode, stbir_filter filter, stbir_colorspace space, - void *alloc_context) -{ - return stbir__resize_arbitrary(alloc_context, input_pixels, input_w, input_h, input_stride_in_bytes, - output_pixels, output_w, output_h, output_stride_in_bytes, - 0,0,1,1,NULL,num_channels,alpha_channel,flags, STBIR_TYPE_UINT8, filter, filter, - edge_wrap_mode, edge_wrap_mode, space); -} - -STBIRDEF int stbir_resize_uint16_generic(const stbir_uint16 *input_pixels , int input_w , int input_h , int input_stride_in_bytes, - stbir_uint16 *output_pixels , int output_w, int output_h, int output_stride_in_bytes, - int num_channels, int alpha_channel, int flags, - stbir_edge edge_wrap_mode, stbir_filter filter, stbir_colorspace space, - void *alloc_context) -{ - return stbir__resize_arbitrary(alloc_context, input_pixels, input_w, input_h, input_stride_in_bytes, - output_pixels, output_w, output_h, output_stride_in_bytes, - 0,0,1,1,NULL,num_channels,alpha_channel,flags, STBIR_TYPE_UINT16, filter, filter, - edge_wrap_mode, edge_wrap_mode, space); -} - - -STBIRDEF int stbir_resize_float_generic( const float *input_pixels , int input_w , int input_h , int input_stride_in_bytes, - float *output_pixels , int output_w, int output_h, int output_stride_in_bytes, - int num_channels, int alpha_channel, int flags, - stbir_edge edge_wrap_mode, stbir_filter filter, stbir_colorspace space, - void *alloc_context) -{ - return stbir__resize_arbitrary(alloc_context, input_pixels, input_w, input_h, input_stride_in_bytes, - output_pixels, output_w, output_h, output_stride_in_bytes, - 0,0,1,1,NULL,num_channels,alpha_channel,flags, STBIR_TYPE_FLOAT, filter, filter, - edge_wrap_mode, edge_wrap_mode, space); -} - - -STBIRDEF int stbir_resize( const void *input_pixels , int input_w , int input_h , int input_stride_in_bytes, - void *output_pixels, int output_w, int output_h, int output_stride_in_bytes, - stbir_datatype datatype, - int num_channels, int alpha_channel, int flags, - stbir_edge edge_mode_horizontal, stbir_edge edge_mode_vertical, - stbir_filter filter_horizontal, stbir_filter filter_vertical, - stbir_colorspace space, void *alloc_context) -{ - return stbir__resize_arbitrary(alloc_context, input_pixels, input_w, input_h, input_stride_in_bytes, - output_pixels, output_w, output_h, output_stride_in_bytes, - 0,0,1,1,NULL,num_channels,alpha_channel,flags, datatype, filter_horizontal, filter_vertical, - edge_mode_horizontal, edge_mode_vertical, space); -} - - -STBIRDEF int stbir_resize_subpixel(const void *input_pixels , int input_w , int input_h , int input_stride_in_bytes, - void *output_pixels, int output_w, int output_h, int output_stride_in_bytes, - stbir_datatype datatype, - int num_channels, int alpha_channel, int flags, - stbir_edge edge_mode_horizontal, stbir_edge edge_mode_vertical, - stbir_filter filter_horizontal, stbir_filter filter_vertical, - stbir_colorspace space, void *alloc_context, - float x_scale, float y_scale, - float x_offset, float y_offset) -{ - float transform[4]; - transform[0] = x_scale; - transform[1] = y_scale; - transform[2] = x_offset; - transform[3] = y_offset; - return stbir__resize_arbitrary(alloc_context, input_pixels, input_w, input_h, input_stride_in_bytes, - output_pixels, output_w, output_h, output_stride_in_bytes, - 0,0,1,1,transform,num_channels,alpha_channel,flags, datatype, filter_horizontal, filter_vertical, - edge_mode_horizontal, edge_mode_vertical, space); -} - -STBIRDEF int stbir_resize_region( const void *input_pixels , int input_w , int input_h , int input_stride_in_bytes, - void *output_pixels, int output_w, int output_h, int output_stride_in_bytes, - stbir_datatype datatype, - int num_channels, int alpha_channel, int flags, - stbir_edge edge_mode_horizontal, stbir_edge edge_mode_vertical, - stbir_filter filter_horizontal, stbir_filter filter_vertical, - stbir_colorspace space, void *alloc_context, - float s0, float t0, float s1, float t1) -{ - return stbir__resize_arbitrary(alloc_context, input_pixels, input_w, input_h, input_stride_in_bytes, - output_pixels, output_w, output_h, output_stride_in_bytes, - s0,t0,s1,t1,NULL,num_channels,alpha_channel,flags, datatype, filter_horizontal, filter_vertical, - edge_mode_horizontal, edge_mode_vertical, space); -} - -#endif // STB_IMAGE_RESIZE_IMPLEMENTATION - -/* ------------------------------------------------------------------------------- -This software is available under 2 licenses -- choose whichever you prefer. ------------------------------------------------------------------------------- -ALTERNATIVE A - MIT License -Copyright (c) 2017 Sean Barrett -Permission is hereby granted, free of charge, to any person obtaining a copy of -this software and associated documentation files (the "Software"), to deal in -the Software without restriction, including without limitation the rights to -use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies -of the Software, and to permit persons to whom the Software is furnished to do -so, subject to the following conditions: -The above copyright notice and this permission notice shall be included in all -copies or substantial portions of the Software. -THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -SOFTWARE. ------------------------------------------------------------------------------- -ALTERNATIVE B - Public Domain (www.unlicense.org) -This is free and unencumbered software released into the public domain. -Anyone is free to copy, modify, publish, use, compile, sell, or distribute this -software, either in source code form or as a compiled binary, for any purpose, -commercial or non-commercial, and by any means. -In jurisdictions that recognize copyright laws, the author or authors of this -software dedicate any and all copyright interest in the software to the public -domain. We make this dedication for the benefit of the public at large and to -the detriment of our heirs and successors. We intend this dedication to be an -overt act of relinquishment in perpetuity of all present and future rights to -this software under copyright law. -THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -AUTHORS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN -ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION -WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. ------------------------------------------------------------------------------- -*/