diff --git a/.ci/scripts/test-native.sh b/.ci/scripts/test-native.sh index 98139a453e2..fadbfe75840 100644 --- a/.ci/scripts/test-native.sh +++ b/.ci/scripts/test-native.sh @@ -36,8 +36,9 @@ case "${component}" in -DEXECUTORCH_BUILD_NATIVE_VISION_RUNNER=ON -DGGML_NATIVE=OFF ) - targets+=(run_hf_llm_ggml native_ggml_engine_test native_llm_runner_utils_test run_classification_ggml_cpu) - test_regex='^native_(ggml_engine|llm_runner_utils)_test$' + # TODO: run the other sample tests once they pass. + targets+=(run_hf_llm_ggml native_ggml_engine_test native_llm_runner_utils_test run_classification_ggml_cpu native_output_check_test native_ggml_sample_test) + test_regex='^native_(ggml_engine|llm_runner_utils|output_check)_test$|^native_ggml_sample_test_act_quant$' ;; vulkan) cmake_args+=( diff --git a/.github/workflows/test-native.yml b/.github/workflows/test-native.yml index c7a7aae2f0c..71d051f1679 100644 --- a/.github/workflows/test-native.yml +++ b/.github/workflows/test-native.yml @@ -17,6 +17,8 @@ on: - backends/cpu/CMakeLists.txt - backends/cpu/runtime/** - backends/native/** + - backends/cadence/utils/FACTO + - backends/test/suite/operators/test_linear.py - backends/vulkan/CMakeLists.txt - backends/vulkan/cmake/** - backends/vulkan/runtime/api/** @@ -113,9 +115,10 @@ jobs: if [[ "${NATIVE_COMPONENT}" == "ggml" ]]; then ./install_requirements.sh --use-pt-pinned-commit - # The LLM runner test imports executorch from src/ without installing the - # wheel, so install its export-time dependencies directly. - pip install flatbuffers ruamel.yaml tabulate \ + # The LLM runner test and the op sample generator import executorch from + # src/ without installing the wheel, so install their export-time + # dependencies directly. + pip install flatbuffers gguf ruamel.yaml tabulate \ "$(grep -E '^transformers\b' requirements-examples.txt)" fi diff --git a/backends/native/README.md b/backends/native/README.md index 8a2d76ce20a..33e5d73cf24 100644 --- a/backends/native/README.md +++ b/backends/native/README.md @@ -145,6 +145,14 @@ an engine through the interfaces in `runtime/engine/Engine.h`: An engine rejects a method whose semantics it cannot fully honor. +Engines need not match eager bit for bit: accumulation precision and order, +fused multiply-add, and the accuracy of functions such as `exp` are up to each +engine. As a guideline, an engine should round to nearest when it narrows a +value for storage, with ties to even between floating-point types (f32 to f16 or +bf16), as eager does. Rounding toward zero, which some GPU APIs allow, biases +every stored value toward zero, and the bias compounds across layers instead of +averaging out. + | Engine | Source | Notes | | --- | --- | --- | | Vulkan | `runtime/vulkan/` | Builds on the ET-VK `ComputeGraph` and shaders. | @@ -209,8 +217,9 @@ enabled, `native_runner_utils_test` covers their shared runner utilities. Run th with `ctest --test-dir cmake-out -R native_` after `cmake --build cmake-out`. The engines' sample tests judge them on the bundled op samples. Building them runs -the sample generator, which needs `torch` and `torchao`. Those with failing cases -are disabled in CTest; run their binaries directly with +the sample generator, which needs `torch`, `torchao`, `gguf` and the FACTO +submodule at `backends/cadence/utils/FACTO`. Those with failing cases are +disabled in CTest; run their binaries directly with `--packages=/backends/native/test/op_samples`. The `Test Native Engines` CI workflow has separate Linux jobs for the shared diff --git a/backends/native/runtime/vulkan/TODO.md b/backends/native/runtime/vulkan/TODO.md index 2b6b1de54df..f4c8d3bd7ad 100644 --- a/backends/native/runtime/vulkan/TODO.md +++ b/backends/native/runtime/vulkan/TODO.md @@ -14,3 +14,23 @@ HuggingFace Llama 3.2 1B 8da4w, 2K export: peak GPU memory is 2794 MiB, against 2265 MiB with buffer SDPA disabled and 2263 MiB for the ET-VK `.pte`. The +528 MiB comes from enabling buffer SDPA. Give these temporaries lifetimes in the memory plan, measure again, and attribute whatever remains. + +## Lower int4 weight-only linear for any N + +`FuseQuantizedLinear` leaves an int4 weight-only linear unlowered when N is not +a multiple of 4, to match `VK_CHECK_COND(N % 4 == 0)` in ET-VK's +`prepack_q4_w_4x8_nc_buffer`, and the engine then rejects the quantized weight. +Round N up to a multiple of 4 there (the pack shader already fills n >= N with +zeros) and mask the partial last tile on store. Until then Vulkan fails the op +sample tests `linear_int4wo_*.linear_feature_sizes_2` and +`linear_int4wo_*.linear_unbatched_0`, where N is 1 and 43. + +## Rank-4 linear input + +Vulkan returns wrong outputs for the op sample test +`linear_float.linear_leading_batch_1`, an `nn.Linear(67, 43)` on a +(2, 4, 8, 67) input: in f32 the fitted rel_bias is -0.48, and f16 returns +values up to 6.1e4. Both are consistent with only one of the two outer batches +being written. The rank-3 inputs (4, 8, 67) and (2, 7, 256) pass. Find where +the outer dimension is lost: the engine's linear lowering, or ET-VK's linear on +a rank-4 tensor. diff --git a/backends/native/test/CMakeLists.txt b/backends/native/test/CMakeLists.txt index 719edc54578..0874426cc68 100644 --- a/backends/native/test/CMakeLists.txt +++ b/backends/native/test/CMakeLists.txt @@ -48,18 +48,20 @@ endforeach() # environment pins the eager kernels' ISA and threads, so the golden is the same # on any x86 host with AVX2. Other hosts can differ. set(_native_op_samples ${CMAKE_CURRENT_BINARY_DIR}/op_samples) -file(GLOB _native_op_samples_srcs CONFIGURE_DEPENDS +file(GLOB_RECURSE _native_op_samples_srcs CONFIGURE_DEPENDS ${CMAKE_CURRENT_SOURCE_DIR}/samples/*.py + ${EXECUTORCH_ROOT}/backends/cadence/utils/FACTO/facto/*.py ) add_custom_command( OUTPUT ${_native_op_samples}.stamp COMMAND ${CMAKE_COMMAND} -E env MKL_CBWR=COMPATIBLE ONEDNN_MAX_CPU_ISA=AVX2 ATEN_CPU_CAPABILITY=avx2 OMP_NUM_THREADS=1 MKL_NUM_THREADS=1 - "PYTHONPATH=${EXECUTORCH_ROOT}/src" ${PYTHON_EXECUTABLE} -m + "PYTHONPATH=${EXECUTORCH_ROOT}/src:${EXECUTORCH_ROOT}/backends/cadence/utils/FACTO" + FLATC_EXECUTABLE=$ ${PYTHON_EXECUTABLE} -m executorch.backends.native.test.samples.generate --out ${_native_op_samples} COMMAND ${CMAKE_COMMAND} -E touch ${_native_op_samples}.stamp - DEPENDS ${_native_op_samples_srcs} + DEPENDS flatc ${_native_op_samples_srcs} ${EXECUTORCH_ROOT}/backends/test/suite/operators/test_linear.py VERBATIM ) @@ -155,6 +157,12 @@ if(EXECUTORCH_BUILD_NATIVE_GGML) ) # TODO: Enable once every case passes. set_tests_properties(native_ggml_sample_test PROPERTIES DISABLED TRUE) + # native_ggml_sample_test_act_quant lets ggml quantize the activations of + # weight-only quantized linears, judged against golden-int8_act. + add_test(NAME native_ggml_sample_test_act_quant + COMMAND native_ggml_sample_test --packages=${_native_op_samples} + --device=cpu --allow_activation_quant_for_weight_only + ) target_compile_options( native_ggml_sample_test PRIVATE $<$:/EHsc> diff --git a/backends/native/test/runtime/PtnRunner.h b/backends/native/test/runtime/PtnRunner.h index de23aeb4da4..7bb03987f7a 100644 --- a/backends/native/test/runtime/PtnRunner.h +++ b/backends/native/test/runtime/PtnRunner.h @@ -12,8 +12,15 @@ namespace ptn::testing { +// Numerics a driver lets an engine trade away. Drivers then judge the engine as +// having made the trade, so an engine that never makes it rejects the option. +struct HostOptions { + // Quantize the activations of linears whose weights alone are quantized. + bool allow_activation_quant_for_weight_only = false; +}; + // Defined once per binary by an engine's host library; drivers such as // ptn_runner_main.cpp run every engine through the host it returns. -std::unique_ptr make_engine_host(); +std::unique_ptr make_engine_host(const HostOptions& options); } // namespace ptn::testing diff --git a/backends/native/test/runtime/ggml/BUCK b/backends/native/test/runtime/ggml/BUCK index 71d998be09e..94105c0a11f 100644 --- a/backends/native/test/runtime/ggml/BUCK +++ b/backends/native/test/runtime/ggml/BUCK @@ -45,6 +45,11 @@ _VARIANTS = [ # Buck does not put the CUDA driver stub's library on the loader path. # @oss-disable[end= ]: "LD_PRELOAD": "$(location //tools/lazify:lazy_stub_binder[shared])", }, + # ggml-CUDA computes f16 BLAS GEMMs in f16 and f32 ones in TF32 by default. + "sample_env": { + "GGML_CUDA_CUBLAS_COMPUTE_TYPE": "f32", + "NVIDIA_TF32_OVERRIDE": "0", + }, }, { "suffix": "_hip", @@ -91,23 +96,29 @@ _VARIANTS = [ for v in _VARIANTS ] +# ggml_sample_test_act_quant lets ggml quantize the activations of weight-only +# quantized linears, judged against golden-int8_act. [ fbcode_target( _kind = runtime.cxx_test, - name = "ggml_sample_test" + v["suffix"], + name = "ggml_sample_test" + name + v["suffix"], args = [ "--packages=$(location //executorch/backends/native/test/samples:op_samples)", "--device=" + v["device"], - ], - env = v.get("env", {}), + ] + flags, + env = v.get("env", {}) | v.get("sample_env", {}), remote_execution = v.get("remote_execution"), - # TODO: run in CI once they pass. - labels = ci.labels(ci.map(ci.skip_test())), + # TODO: run the others in CI once they pass. + labels = [] if v["suffix"] + name == "_act_quant" else ci.labels(ci.map(ci.skip_test())), test_main = "//executorch/backends/native/test/samples:sample_test_main", deps = [":ggml_engine_host" + v["suffix"]], **v.get("compat", {}) ) for v in _VARIANTS + for name, flags in [ + ("", []), + ("_act_quant", ["--allow_activation_quant_for_weight_only"]), + ] ] [ diff --git a/backends/native/test/runtime/ggml/ggml_ptn_runner.cpp b/backends/native/test/runtime/ggml/ggml_ptn_runner.cpp index ad600382f04..9bff34d2f66 100644 --- a/backends/native/test/runtime/ggml/ggml_ptn_runner.cpp +++ b/backends/native/test/runtime/ggml/ggml_ptn_runner.cpp @@ -26,15 +26,11 @@ DEFINE_bool( print_graph_ops, false, "Print each built graph's node count per ggml op to stderr."); -DEFINE_bool( - allow_activation_quant_for_weight_only, - false, - "Run weight-only quantized linears with ggml's quantized activations."); namespace ptn::testing { // cppcheck-suppress unusedFunction -std::unique_ptr make_engine_host() { +std::unique_ptr make_engine_host(const HostOptions& host_options) { GgmlEngineOptions options; if (FLAGS_device == "gpu") { options.device = GgmlDevice::Gpu; @@ -49,7 +45,7 @@ std::unique_ptr make_engine_host() { options.max_cached_graphs = static_cast(FLAGS_max_cached_graphs); options.print_graph_ops = FLAGS_print_graph_ops; options.allow_activation_quant_for_weight_only = - FLAGS_allow_activation_quant_for_weight_only; + host_options.allow_activation_quant_for_weight_only; return GgmlEngineHost::create(options); } diff --git a/backends/native/test/runtime/ptn_runner_main.cpp b/backends/native/test/runtime/ptn_runner_main.cpp index 654ef0ad569..973aeec5025 100644 --- a/backends/native/test/runtime/ptn_runner_main.cpp +++ b/backends/native/test/runtime/ptn_runner_main.cpp @@ -65,6 +65,10 @@ DEFINE_int32( iterations, 1, "Timed executions; above one, prints median and min latency."); +DEFINE_bool( + allow_activation_quant_for_weight_only, + false, + "Let the engine quantize the activations of weight-only quantized linears."); namespace { @@ -201,7 +205,9 @@ void run() { const ptn::ByteSpan program_bytes = package->program_bytes(); const auto program = std::make_shared( ptn::Program::load(program_bytes.data(), program_bytes.size())); - const auto host = ptn::testing::make_engine_host(); + const auto host = ptn::testing::make_engine_host( + {.allow_activation_quant_for_weight_only = + FLAGS_allow_activation_quant_for_weight_only}); const auto context = host->create_context(program, package); const auto executable = context->compile(FLAGS_method); diff --git a/backends/native/test/runtime/vulkan/vulkan_ptn_runner.cpp b/backends/native/test/runtime/vulkan/vulkan_ptn_runner.cpp index 3f8b9e76ee0..51163c0d478 100644 --- a/backends/native/test/runtime/vulkan/vulkan_ptn_runner.cpp +++ b/backends/native/test/runtime/vulkan/vulkan_ptn_runner.cpp @@ -6,13 +6,20 @@ // The Vulkan engine host for PTN test drivers. +#include + #include #include namespace ptn::testing { // cppcheck-suppress unusedFunction -std::unique_ptr make_engine_host() { +std::unique_ptr make_engine_host(const HostOptions& options) { + if (options.allow_activation_quant_for_weight_only) { + throw std::invalid_argument( + "vulkan: never quantizes activations, so it rejects " + "allow_activation_quant_for_weight_only"); + } return VulkanEngineHost::create(); } diff --git a/backends/native/test/samples/BUCK b/backends/native/test/samples/BUCK index 122272f1d3e..6ff24d67abc 100644 --- a/backends/native/test/samples/BUCK +++ b/backends/native/test/samples/BUCK @@ -17,7 +17,11 @@ fbcode_target( "//caffe2:torch", "//executorch/backends/test/suite/operators:test_operators", "//executorch/exir/native:native", + "//executorch/extension/llm/export:gguf", "//pytorch/ao:torchao", + "//pytorch/facto:facto", + "fbsource//third-party/pypi/gguf:gguf", + "fbsource//third-party/pypi/numpy:numpy", ], ) @@ -37,6 +41,7 @@ fbcode_target( ":record", "//caffe2:torch", "//executorch/exir/native:native", + "//pytorch/ao:torchao", ], ) diff --git a/backends/native/test/samples/cases.py b/backends/native/test/samples/cases.py index b08baba963a..82b97316b24 100644 --- a/backends/native/test/samples/cases.py +++ b/backends/native/test/samples/cases.py @@ -7,11 +7,16 @@ """Op cases: a model and the inputs of each sample it is recorded with.""" import inspect +from collections.abc import Callable from dataclasses import dataclass from types import ModuleType +from typing import TYPE_CHECKING import torch +if TYPE_CHECKING: + from facto.inputgen.specs.model import Spec + Inputs = tuple[torch.Tensor, ...] # Caps a suite case's input and output elements over its samples, keeping at @@ -85,3 +90,20 @@ def suite_cases(module: ModuleType, samples: int) -> list[Case]: if not cases: raise ValueError(f"no test in {module.__name__} lowers a model") return cases + + +def facto_cases( + spec: "Spec", build: Callable[..., Case], keep: Callable[[int, tuple], bool] +) -> list[Case]: + """One-sample cases `build(name, *args)` from FACTO's argument tuples for + `spec`, named by position, keeping those where `keep(position, args)`.""" + # FACTO is on the path for the sample generator only, not for pytest. + from facto.inputgen.argtuple.gen import ArgumentTupleGenerator + from facto.inputgen.utils.random_manager import seeded_random_manager + + seeded_random_manager.seed(0) + return [ + build(f"facto_{i}", *args) + for i, (args, _, _) in enumerate(ArgumentTupleGenerator(spec).gen()) + if keep(i, args) + ] diff --git a/backends/native/test/samples/linear.py b/backends/native/test/samples/linear.py index 24c39fe2ccc..9beea6408fa 100644 --- a/backends/native/test/samples/linear.py +++ b/backends/native/test/samples/linear.py @@ -4,12 +4,90 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. -"""Linear cases: the backend test suite's, and shapes and inputs that stress -accumulation and quantization.""" +"""Linear cases: the backend test suite's, FACTO's, and shapes and inputs that +stress accumulation and quantization.""" import torch -from executorch.backends.native.test.samples.cases import Case, draws, suite_cases +from executorch.backends.native.test.samples.cases import ( + Case, + draws, + facto_cases, + suite_cases, +) +from executorch.backends.native.test.samples.schemes import GGUF_SUPER_SCALES from executorch.backends.test.suite.operators import test_linear +from executorch.extension.llm.export.gguf import QK_K +from facto.inputgen.argument.type import ArgType +from facto.inputgen.specs.model import ConstraintProducer as cp, InPosArg, OutArg, Spec + + +def _unit(*constraints) -> list: + """Float32 values in [-1, 1] under `constraints`.""" + return [ + cp.Dtype.Eq(lambda deps: torch.float32), + *constraints, + cp.Value.Ge(lambda deps, dtype, struct: -1), + cp.Value.Le(lambda deps, dtype, struct: 1), + ] + + +# in_features is a multiple of every torchao scheme's group size and, in some +# tuples, of GGUF's block; out_features is large enough to fit an error's +# abs_bias and rel_bias from one sample. +_FACTO_SPEC = Spec( + op="linear.default", + inspec=[ + InPosArg( + ArgType.Tensor, + name="input", + deps=[1], + constraints=_unit( + cp.Rank.Ge(lambda deps: 1), + cp.Rank.Le(lambda deps: 3), + cp.Size.In( + lambda deps, r, d: [deps[0].size(1)] if d == r - 1 else range(1, 9) + ), + ), + ), + InPosArg( + ArgType.Tensor, + name="weight", + constraints=_unit( + cp.Rank.Eq(lambda deps: 2), + cp.Size.In(lambda deps, r, d: [128, 256, 384] if d else [16, 48, 96]), + ), + ), + InPosArg( + ArgType.TensorOpt, + name="bias", + deps=[1], + constraints=_unit( + cp.Rank.Eq(lambda deps: 1), + cp.Size.Eq(lambda deps, r, d: deps[0].size(0)), + ), + ), + ], + outspec=[OutArg(ArgType.Tensor)], +) + + +def _facto_case(name: str, x, weight, bias) -> Case: + model = torch.nn.Linear(weight.size(1), weight.size(0), bias=bias is not None) + model.weight.data = weight + if bias is not None: + model.bias.data = bias + return Case(name, model, [(x,)]) + + +class _Parallel(torch.nn.Module): + """nn.Linears on one input, each giving its own output.""" + + def __init__(self, *linears: torch.nn.Linear) -> None: + super().__init__() + self.linears = torch.nn.ModuleList(linears) + + def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, ...]: + return tuple(m(x) for m in self.linears) def cases() -> list[Case]: @@ -22,7 +100,8 @@ def cases() -> list[Case]: zero_row = draws(2, (8, 256)) for (x,) in zero_row: x[2] = 0 - return suite + [ + return [ + *suite, Case("decode", torch.nn.Linear(2048, 512), draws(8, (1, 2048))), Case("prefill", torch.nn.Linear(1024, 256), draws(2, (64, 1024))), Case("large_k", torch.nn.Linear(8192, 32), draws(4, (4, 8192))), @@ -33,4 +112,19 @@ def cases() -> list[Case]: ), Case("outlier_channel", torch.nn.Linear(256, 128), outlier), Case("zero_row", torch.nn.Linear(256, 64), zero_row), + # One output per family of GGUF super-block scales. + Case( + "scale_families", + _Parallel(*(torch.nn.Linear(256, 32) for _ in GGUF_SUPER_SCALES["q4_k"])), + draws(2, (4, 256)), + ), + # FACTO varies one argument at a time, so only every fifth tuple is kept, + # plus those whose in_features fit GGUF blocks. Its values lie on a 1/8 + # grid, which makes float32 math exact: these cases cover shapes more + # than rounding. + *facto_cases( + _FACTO_SPEC, + _facto_case, + lambda i, args: i % 5 == 0 or args[1].size(1) % QK_K == 0, + ), ] diff --git a/backends/native/test/samples/record.py b/backends/native/test/samples/record.py index 6846890f081..90a57daa399 100644 --- a/backends/native/test/samples/record.py +++ b/backends/native/test/samples/record.py @@ -8,10 +8,12 @@ A package holds one method per case and dtype under one quantization scheme, and two sample sets over the same inputs: `golden`, the eager outputs, and -`precise`, the same math in float64 with weights dequantized and activations -unquantized, floating outputs stored as float32. Nothing in it depends on an -engine or a device: consumers judge an engine by its distance from `precise` -relative to eager's. +`precise`, the same math in float64 with quantized weights decoded in float32, +whatever the method's dtype, and activations unquantized, floating outputs +stored as float32. Weight-only schemes add `golden-int8_act`, the eager outputs +with the same weights and int8 per-token dynamic activations. Nothing in it +depends on an engine or a device: consumers judge an engine by its distance +from `precise` relative to eager's. """ import copy @@ -21,7 +23,7 @@ import torch from executorch.backends.native.test.samples.cases import Case, Inputs, outputs -from executorch.backends.native.test.samples.schemes import Scheme +from executorch.backends.native.test.samples.schemes import int8_activations, Scheme from executorch.exir.native import Call, Sample, SampleSet, to_native, write_samples from torch.func import functional_call @@ -50,7 +52,7 @@ def record( dtypes: Iterable[str] = tuple(DTYPES), ) -> None: """Writes no package, removing any at `path`, when the scheme fits no case.""" - programs, golden, precise = {}, [], [] + programs, golden, precise, golden_act = {}, [], [], [] for case in cases: if not scheme.fits(case.model): logger.info(f"{path.name}: {case.name} does not fit the scheme") @@ -76,6 +78,11 @@ def record( t.float() if t.is_floating_point() else t for t in outputs(ref) ] precise.append(Sample([Call(name, list(x), ref)])) + if scheme.weight_only: + with int8_activations(method): + golden_act.append( + Sample([Call(name, list(x), outputs(method(*x)))]) + ) if not programs: logger.info(f"{path.name}: no case fits the scheme, not written") path.unlink(missing_ok=True) @@ -87,6 +94,14 @@ def record( str(path), "precise", SampleSet( - precise, "float64, weights dequantized, activations unquantized", source + precise, + "float64, quantized weights decoded in float32, activations unquantized", + source, ), ) + if scheme.weight_only: + write_samples( + str(path), + "golden-int8_act", + SampleSet(golden_act, "eager, int8 per-token dynamic activations", source), + ) diff --git a/backends/native/test/samples/sample_test_main.cpp b/backends/native/test/samples/sample_test_main.cpp index 584676a49f6..e44d987343c 100644 --- a/backends/native/test/samples/sample_test_main.cpp +++ b/backends/native/test/samples/sample_test_main.cpp @@ -11,7 +11,8 @@ // // Usage: // --packages=a.ptn,dir [--rmse_ratio --max_ratio --rtol --sigmas] -// [engine host flags] [gtest flags] +// [--allow_activation_quant_for_weight_only] [engine host flags] +// [gtest flags] #include #include @@ -56,6 +57,11 @@ DEFINE_double( sigmas, ptn::testing::Policy{}.sigmas, "Standard errors the fitted biases may deviate by."); +DEFINE_bool( + allow_activation_quant_for_weight_only, + false, + "Let the engine quantize the activations of weight-only quantized linears, " + "judging it against golden-int8_act samples where a package has them."); namespace ptn::testing { namespace { @@ -63,6 +69,7 @@ namespace { struct Bundle { std::shared_ptr package; std::shared_ptr program; + std::string golden_set; SampleSet golden; SampleSet precise; }; @@ -70,7 +77,9 @@ struct Bundle { // Created on first use, so listing tests touches no device. EngineHost& host() { static const std::unique_ptr instance = [] { - std::unique_ptr made = make_engine_host(); + std::unique_ptr made = make_engine_host( + {.allow_activation_quant_for_weight_only = + FLAGS_allow_activation_quant_for_weight_only}); std::printf( "engine %s on %s\n", made->name().c_str(), made->device_name().c_str()); return made; @@ -111,7 +120,8 @@ class SampleTest : public ::testing::Test { .rtol = FLAGS_rtol, .sigmas = FLAGS_sigmas}; for (const auto& [key, check] : checks) { - const std::string name = key.first + " out " + std::to_string(key.second); + const std::string name = key.first + " out " + + std::to_string(key.second) + " vs " + bundle_->golden_set; std::printf("%s: %s\n", name.c_str(), check.summary().c_str()); for (const std::string& failure : check.failures(policy)) { ADD_FAILURE() << name << ": " << failure; @@ -129,7 +139,8 @@ class SampleTest : public ::testing::Test { const SampleCall& call = golden.calls[c]; ASSERT_EQ(call.method, precise.calls[c].method); ASSERT_TRUE(call.inputs == precise.calls[c].inputs) - << call.method << " has golden and precise inputs that differ"; + << call.method << " has " << bundle_->golden_set + << " and precise inputs that differ"; ASSERT_TRUE(call.outputs && precise.calls[c].outputs) << call.method << " has no recorded outputs"; const std::vector& want = *call.outputs; @@ -208,7 +219,12 @@ void register_package(const std::filesystem::path& path) { bundle->program = std::make_shared( Program::load(program.data(), program.size())); const ZipReader zip = ZipReader::open(path.string()); - bundle->golden = read_sample_set(zip, "golden"); + // Eager outputs with int8 per-token dynamic activations, which packages with + // weight-only quantized linears also carry. + const bool act_quant_golden = FLAGS_allow_activation_quant_for_weight_only && + zip.member_size("samples/golden-int8_act/index.json").has_value(); + bundle->golden_set = act_quant_golden ? "golden-int8_act" : "golden"; + bundle->golden = read_sample_set(zip, bundle->golden_set); bundle->precise = read_sample_set(zip, "precise"); std::set methods; diff --git a/backends/native/test/samples/schemes.py b/backends/native/test/samples/schemes.py index 8966e7facad..c0c2f3d1fe8 100644 --- a/backends/native/test/samples/schemes.py +++ b/backends/native/test/samples/schemes.py @@ -7,22 +7,42 @@ """Quantization schemes. Every op is recorded under every scheme, keeping the cases the scheme fits.""" +import contextlib +import tempfile +from collections.abc import Iterator from dataclasses import dataclass +import numpy as np import torch +from executorch.extension.llm.export.gguf import ExportableGGUFTensor, iter_gguf, QK_K +from gguf import GGML_QUANT_SIZES, GGMLQuantizationType, GGUFWriter from torchao.core.config import AOBaseConfig from torchao.quantization import ( + dequantize_affine, Int8DynamicActivationIntxWeightConfig, + IntxUnpackedToInt8Tensor, IntxWeightOnlyConfig, quantize_, ) from torchao.quantization.granularity import PerGroup + +# Private, but it is the input quantize-dequantize that +# Int8DynamicActivationIntxWeightConfig linears run. +from torchao.quantization.quantize_.workflows.intx.intx_unpacked_to_int8_tensor import ( + _apply_int8_act_asym_per_token_quant_dequant, +) from torchao.utils import TorchAOBaseTensor class Scheme: """No quantization.""" + @property + def weight_only(self) -> bool: + """Whether the scheme quantizes only weights; `record` then adds + `golden-int8_act`.""" + return False + def fits(self, model: torch.nn.Module) -> bool: return True @@ -30,7 +50,9 @@ def quantize(self, model: torch.nn.Module) -> None: pass def decode(self, value: torch.Tensor) -> torch.Tensor: - """A state tensor of a quantized model as a plain tensor.""" + """A state tensor of a quantized model as a plain tensor, with quantized + weights decoded in float32 rather than the method's dtype, which can + round them.""" return value @@ -52,6 +74,10 @@ class Torchao(Scheme): config: AOBaseConfig + @property + def weight_only(self) -> bool: + return isinstance(self.config, IntxWeightOnlyConfig) + def fits(self, model: torch.nn.Module) -> bool: granularity = ( self.config.granularity @@ -66,7 +92,122 @@ def quantize(self, model: torch.nn.Module) -> None: quantize_(model, self.config) def decode(self, value: torch.Tensor) -> torch.Tensor: - return value.dequantize() if isinstance(value, TorchAOBaseTensor) else value + if isinstance(value, IntxUnpackedToInt8Tensor): + return dequantize_affine( + value.qdata, value.block_size, value.scale, value.zero_point, torch.int8 + ) + if isinstance(value, TorchAOBaseTensor): + raise TypeError(f"no float32 decode for {type(value).__name__}") + return value + + +# Families of fp16 super-block scales by byte offset in each block: powers of +# two, where q4_k's dmin of 8 d centers random weights near zero, then negative, +# rounded and subnormal ones. A model's i-th nn.Linear takes the i-th family, so +# one nn.Linear per family isolates each in its own output. +# TODO: Add a zero-d family once OutputCheck's bias standard error stops +# assuming independent outputs. +GGUF_SUPER_SCALES = { + "q4_k": [ + {0: 2**-12, 2: 2**-9}, + {0: -(2**-12), 2: -(2**-9)}, + {0: 3.1e-4, 2: 2.3e-3}, + {0: 3 * 2**-21, 2: 3 * 2**-18}, + ], + "q6_k": [ + {208: 2**-12}, + {208: -(2**-12)}, + {208: 3.1e-4}, + {208: 3 * 2**-21}, + ], +} + + +def _random_blocks( + ggml_type: str, + rows: int, + cols: int, + scales: dict[int, float], + rng: np.random.Generator, +) -> np.ndarray: + """Valid `ggml_type` blocks for a [rows, cols] weight from random bytes and + the given super-block scales.""" + _, block_bytes = GGML_QUANT_SIZES[GGMLQuantizationType[ggml_type.upper()]] + raw = rng.integers(0, 256, (rows, cols // QK_K, block_bytes), dtype=np.uint8) + for offset, scale in scales.items(): + raw[..., offset : offset + 2] = np.array([scale], np.float16).view(np.uint8) + return raw.reshape(rows, -1) + + +def _load_gguf(model: torch.nn.Module, ggml_type: str) -> None: + """Writes random `ggml_type` weights for every nn.Linear to a .gguf file and + loads them back as a GGUF checkpoint would be.""" + linears = [m for m in model.modules() if isinstance(m, torch.nn.Linear)] + rng = np.random.default_rng(0) + with tempfile.TemporaryDirectory() as tmp: + path = f"{tmp}/weights.gguf" + writer = GGUFWriter(path, arch="linear") + families = GGUF_SUPER_SCALES[ggml_type] + if len(linears) > len(families): + raise ValueError( + f"{len(linears)} nn.Linears, {len(families)} {ggml_type} scale families" + ) + for i, (m, scales) in enumerate(zip(linears, families)): + writer.add_tensor( + f"linear.{i}.weight", + _random_blocks(ggml_type, *m.weight.shape, scales, rng), + raw_dtype=GGMLQuantizationType[ggml_type.upper()], + ) + writer.write_header_to_file() + writer.write_kv_data_to_file() + writer.write_tensors_to_file() + writer.close() + for m, (_, weight) in zip(linears, iter_gguf(path), strict=True): + m.weight = torch.nn.Parameter( + weight.to(m.weight.dtype), requires_grad=False + ) + + +@dataclass(frozen=True) +class Gguf(Scheme): + """Weight-only GGUF k-quant blocks (`q4_k` or `q6_k`) for every nn.Linear, + loaded from a .gguf file.""" + + ggml_type: str + + @property + def weight_only(self) -> bool: + return True + + def fits(self, model: torch.nn.Module) -> bool: + linears = sum(isinstance(m, torch.nn.Linear) for m in model.modules()) + return _fits(model, QK_K) and linears <= len(GGUF_SUPER_SCALES[self.ggml_type]) + + def quantize(self, model: torch.nn.Module) -> None: + _load_gguf(model, self.ggml_type) + + def decode(self, value: torch.Tensor) -> torch.Tensor: + if isinstance(value, ExportableGGUFTensor): + return value.dequantize(torch.float32) + return value + + +@contextlib.contextmanager +def int8_activations(model: torch.nn.Module) -> Iterator[None]: + """Quantizes every nn.Linear input as torchao's int8 dynamic activation + schemes do: int8, asymmetric, per token.""" + hooks = [ + m.register_forward_pre_hook( + lambda _, args: (_apply_int8_act_asym_per_token_quant_dequant(args[0]),) + ) + for m in model.modules() + if isinstance(m, torch.nn.Linear) + ] + try: + yield + finally: + for hook in hooks: + hook.remove() def schemes() -> dict[str, Scheme]: @@ -88,4 +229,6 @@ def schemes() -> dict[str, Scheme]: weight_dtype=torch.int4, weight_granularity=PerGroup(128) ) ), + "gguf_q4_k": Gguf("q4_k"), + "gguf_q6_k": Gguf("q6_k"), } diff --git a/backends/native/test/samples/test_record.py b/backends/native/test/samples/test_record.py index 391ecd9b52a..fd5d28d3b2f 100644 --- a/backends/native/test/samples/test_record.py +++ b/backends/native/test/samples/test_record.py @@ -4,6 +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. +import copy import tempfile import unittest from pathlib import Path @@ -13,8 +14,16 @@ import torch.nn.functional as F from executorch.backends.native.test.samples.cases import Case, draws, suite_cases from executorch.backends.native.test.samples.record import record -from executorch.backends.native.test.samples.schemes import Scheme, schemes +from executorch.backends.native.test.samples.schemes import ( + Gguf, + GGUF_SUPER_SCALES, + Scheme, + schemes, + Torchao, +) from executorch.exir.native import read_samples +from torchao.quantization import IntxWeightOnlyConfig, quantize_ +from torchao.quantization.granularity import PerGroup class RecordTest(unittest.TestCase): @@ -77,6 +86,25 @@ def test_record_golden_is_eager_and_precise_is_float64(self) -> None: atol=0, ) + def test_record_adds_the_8da4w_golden_to_weight_only_packages(self) -> None: + case = Case("linear", torch.nn.Linear(64, 8), draws(2, (3, 64))) + sets = {} + with tempfile.TemporaryDirectory() as tmp: + for name in ("float", "int4wo_g32", "8da4w_g32"): + path = Path(tmp) / f"{name}.ptn" + record([case], schemes()[name], path, dtypes=["f32"]) + sets[name] = read_samples(str(path)) + self.assertNotIn("golden-int8_act", sets["float"]) + self.assertNotIn("golden-int8_act", sets["8da4w_g32"]) + for got, want in zip( + sets["int4wo_g32"]["golden-int8_act"].samples, + sets["8da4w_g32"]["golden"].samples, + strict=True, + ): + torch.testing.assert_close( + got.calls[0].outputs, want.calls[0].outputs, rtol=0, atol=0 + ) + def test_record_keeps_integer_outputs_precise(self) -> None: class Argmax(torch.nn.Module): def forward(self, x: torch.Tensor) -> torch.Tensor: @@ -88,3 +116,53 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: record([case], Scheme(), path, dtypes=["f32"]) precise = read_samples(str(path))["precise"].samples self.assertEqual(precise[0].calls[0].outputs[0].dtype, torch.int64) + + def test_precise_decodes_gguf_weights_whatever_the_method_dtype(self) -> None: + # Halves are exact in every recorded dtype and the methods share their + # GGUF blocks, so their precise outputs must match. + x = torch.randint(-8, 9, (3, 256)) / 2 + model = torch.nn.Linear(256, 8, bias=False) + with tempfile.TemporaryDirectory() as tmp: + path = Path(tmp) / "linear.ptn" + record([Case("linear", model, [(x,)])], Gguf("q6_k"), path) + samples = read_samples(str(path))["precise"].samples + outputs = {s.calls[0].method: s.calls[0].outputs[0] for s in samples} + for dtype in ("f16", "bf16"): + torch.testing.assert_close( + outputs[f"linear_{dtype}"], outputs["linear_f32"], rtol=0, atol=0 + ) + + def test_gguf_fits_no_more_linears_than_scale_families(self) -> None: + def linears(count: int) -> torch.nn.Module: + return torch.nn.Sequential( + *(torch.nn.Linear(256, 256) for _ in range(count)) + ) + + families = len(GGUF_SUPER_SCALES["q4_k"]) + self.assertTrue(Gguf("q4_k").fits(linears(families))) + self.assertFalse(Gguf("q4_k").fits(linears(families + 1))) + + def test_precise_decodes_intx_weights_in_float32(self) -> None: + # A code times its scale can need more bits than bf16 holds, so decoding + # in the method's dtype would round the weights. + x = torch.randint(-8, 9, (3, 64)) / 2 + model = torch.nn.Linear(64, 8, bias=False) + config = IntxWeightOnlyConfig(weight_dtype=torch.int4, granularity=PerGroup(32)) + with tempfile.TemporaryDirectory() as tmp: + path = Path(tmp) / "linear.ptn" + record([Case("linear", model, [(x,)])], Torchao(config), path) + samples = read_samples(str(path))["precise"].samples + outputs = {s.calls[0].method: s.calls[0].outputs[0] for s in samples} + bf16 = copy.deepcopy(model).to(torch.bfloat16) + quantize_(bf16, config) + w = bf16.weight + codes = ( + w.qdata.double().unflatten(1, (-1, 32)) - w.zero_point.double()[..., None] + ) + weight = (codes * w.scale.double()[..., None]).flatten(1) + torch.testing.assert_close( + outputs["linear_bf16"], + torch.nn.functional.linear(x.double(), weight).float(), + rtol=0, + atol=0, + ) diff --git a/backends/vulkan/runtime/gen_vulkan_spv.py b/backends/vulkan/runtime/gen_vulkan_spv.py index 5701248266e..76999cb0995 100644 --- a/backends/vulkan/runtime/gen_vulkan_spv.py +++ b/backends/vulkan/runtime/gen_vulkan_spv.py @@ -355,6 +355,38 @@ def layout_declare_buffer( return out_str +# Vulkan leaves the rounding of a store to an rgba16f image to the device, and +# some devices round toward zero. Rounding to the nearest f16, ties to even, in +# the shader makes the store exact. Integer texels are stored as is, but stores +# to rgba32f images would be rounded too, so a shader may not write both. +ROUND_F16_IMAGE_STORES = """ +#ifdef WRITES_RGBA32F_IMAGE +#error "A shader cannot write both rgba16f and rgba32f images" +#endif +#ifndef WRITES_RGBA16F_IMAGE +#define WRITES_RGBA16F_IMAGE +vec4 round_to_f16(const vec4 v) { + const uvec4 u = floatBitsToUint(v); + // Keep 10 mantissa bits: add just under half an f16 ulp, plus one when the + // kept lowest bit is odd so ties go to even, then clear the 13 dropped bits. + vec4 r = uintBitsToFloat((u + 0xFFFu + ((u >> 13) & 1u)) & ~0x1FFFu); + // Below the smallest normal f16, round to a multiple of 2^-24. + r = mix(r, roundEven(v * 16777216.0) * 5.9604644775390625e-8, lessThan(abs(v), vec4(6.103515625e-5))); + r = mix(r, uintBitsToFloat((u & 0x80000000u) | 0x7F800000u), greaterThan(abs(r), vec4(65504.0))); + return mix(r, v, isnan(v)); +} +ivec4 round_to_f16(const ivec4 v) { return v; } +uvec4 round_to_f16(const uvec4 v) { return v; } +#define imageStore(image, pos, texel) imageStore(image, pos, round_to_f16(texel)) +#endif""" + +WRITES_RGBA32F_IMAGE = """ +#ifdef WRITES_RGBA16F_IMAGE +#error "A shader cannot write both rgba16f and rgba32f images" +#endif +#define WRITES_RGBA32F_IMAGE""" + + def layout_declare_image( slot: Union[int, List[int]], access_type: str, @@ -367,6 +399,11 @@ def layout_declare_image( image_type = TYPE_MAPPINGS["IMAGE_T"][image_ndim][dtype] ret_str = f"layout(set = 0, binding = {get_slot_val(slot)}, {image_format}) uniform {precision} restrict {get_access_qualifier(access_type)} {image_type} {var_name};" + if access_type.lower() != "r": + if image_format == "rgba16f": + ret_str += ROUND_F16_IMAGE_STORES + elif image_format == "rgba32f": + ret_str += WRITES_RGBA32F_IMAGE if isinstance(slot, list): slot[0] = slot[0] + 1 diff --git a/backends/vulkan/runtime/graph/ComputeGraph.cpp b/backends/vulkan/runtime/graph/ComputeGraph.cpp index a6b3c4f6347..e3997b26988 100644 --- a/backends/vulkan/runtime/graph/ComputeGraph.cpp +++ b/backends/vulkan/runtime/graph/ComputeGraph.cpp @@ -426,10 +426,10 @@ bool ComputeGraph::is_standard_channels_packed_texture_tensor( bool ComputeGraph::is_2d_matrix(const ValueRef idx) const { std::vector sizes = sizes_of(idx); const size_t ndim = sizes.size(); - if (sizes.size() < 2) { + if (ndim == 0) { return false; } - if (sizes.size() == 2) { + if (ndim <= 2) { return true; } diff --git a/backends/vulkan/runtime/graph/ComputeGraph.h b/backends/vulkan/runtime/graph/ComputeGraph.h index 6b78d8b50fb..eb0a0e168d2 100644 --- a/backends/vulkan/runtime/graph/ComputeGraph.h +++ b/backends/vulkan/runtime/graph/ComputeGraph.h @@ -440,9 +440,9 @@ class ComputeGraph final { bool is_standard_channels_packed_texture_tensor(const ValueRef idx) const; /* - * Checks that the value at `idx` is either a 2D tensor, or if the tensor has - * more than 2 dims, the outermost dims have size of 1, i.e. can be squeezed - * to be a 2D tensor. + * Checks that the value at `idx` is a 1D tensor (a single row), a 2D tensor, + * or a tensor whose outermost dims have size of 1, i.e. can be squeezed to + * be a 2D tensor. */ bool is_2d_matrix(const ValueRef idx) const; diff --git a/backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/pack_weights/conv2d_dw_prepack_weights.glsl b/backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/pack_weights/conv2d_dw_prepack_weights.glsl index 4707bcf9f9c..4a16ee05d42 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/pack_weights/conv2d_dw_prepack_weights.glsl +++ b/backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/pack_weights/conv2d_dw_prepack_weights.glsl @@ -21,7 +21,7 @@ ${define_required_extensions("buffer", BUF_DTYPE)} layout(std430) buffer; -layout(set = 0, binding = 0, ${IMAGE_FORMAT[DTYPE]}) uniform PRECISION restrict writeonly ${IMAGE_T[2][DTYPE]} image_out; +${layout_declare_tensor(0, "w", "image_out", DTYPE, "texture2d")} layout(set = 0, binding = 1) buffer PRECISION restrict readonly Buffer { BUF_T buffer_in[]; }; diff --git a/backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/pack_weights/conv2d_prepack_weights.glsl b/backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/pack_weights/conv2d_prepack_weights.glsl index b5bec8fca0b..95c8cc3c4ee 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/pack_weights/conv2d_prepack_weights.glsl +++ b/backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/pack_weights/conv2d_prepack_weights.glsl @@ -29,7 +29,7 @@ $if OUTPUT_BUFFER: f16vec4 buffer_out[]; }; $else: - layout(set = 0, binding = 0, ${IMAGE_FORMAT[DTYPE]}) uniform PRECISION restrict writeonly ${IMAGE_T[2][DTYPE]} image_out; + ${layout_declare_tensor(0, "w", "image_out", DTYPE, "texture2d")} layout(set = 0, binding = 1) buffer PRECISION restrict readonly Buffer { BUF_T buffer_in[]; }; diff --git a/backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/transposed/pack_weights/conv_transpose2d_prepack_weights.glsl b/backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/transposed/pack_weights/conv_transpose2d_prepack_weights.glsl index 65b998a9c5b..4a47892375e 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/transposed/pack_weights/conv_transpose2d_prepack_weights.glsl +++ b/backends/vulkan/runtime/graph/ops/glsl/convolution/conv2d/transposed/pack_weights/conv_transpose2d_prepack_weights.glsl @@ -21,7 +21,7 @@ ${define_required_extensions("buffer", BUF_DTYPE)} layout(std430) buffer; -layout(set = 0, binding = 0, ${IMAGE_FORMAT[DTYPE]}) uniform PRECISION restrict writeonly ${IMAGE_T[2][DTYPE]} image_out; +${layout_declare_tensor(0, "w", "image_out", DTYPE, "texture2d")} layout(set = 0, binding = 1) buffer PRECISION restrict readonly Buffer { BUF_T buffer_in[]; }; diff --git a/backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/dq8ca_q4gsw/linear_dq8ca_q4gsw_tiled.glsl b/backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/dq8ca_q4gsw/linear_dq8ca_q4gsw_tiled.glsl index 5acab8e9a12..d29167e5c9a 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/dq8ca_q4gsw/linear_dq8ca_q4gsw_tiled.glsl +++ b/backends/vulkan/runtime/graph/ops/glsl/gemm/linear/qw/dq8ca_q4gsw/linear_dq8ca_q4gsw_tiled.glsl @@ -130,7 +130,8 @@ void main() { } load_weight_scales_tile_for_group(weight_scales_tile, n4, group_i, N4); - load_weight_sums_tile_for_group(weight_sums_tile, n4, group_i, N4); + // Weight sums rows are padded to a multiple of 8. + load_weight_sums_tile_for_group(weight_sums_tile, n4, group_i, mul_2(N8)); load_int8_input_sums_tile_for_group(int8_input_sums_tile, m4, group_i, M4); accumulate_out_tile_with_int_accum_from_int4_weights( diff --git a/backends/vulkan/runtime/graph/ops/glsl/repeat_channel.glsl b/backends/vulkan/runtime/graph/ops/glsl/repeat_channel.glsl index 42c7f86aea8..96e9f7aac4c 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/repeat_channel.glsl +++ b/backends/vulkan/runtime/graph/ops/glsl/repeat_channel.glsl @@ -16,7 +16,7 @@ layout(std430) buffer; #include "indexing_utils.h" -layout(set = 0, binding = 0, ${IMAGE_FORMAT[DTYPE]}) uniform PRECISION restrict writeonly ${IMAGE_T[NDIM][DTYPE]} image_out; +${layout_declare_image(0, "w", "image_out", DTYPE, image_ndim=NDIM)} layout(set = 0, binding = 1) uniform PRECISION sampler3D image_in; layout(set = 0, binding = 2) uniform PRECISION restrict RepeatArgs { diff --git a/backends/vulkan/runtime/graph/ops/impl/QuantizeDequantize.cpp b/backends/vulkan/runtime/graph/ops/impl/QuantizeDequantize.cpp index 7cf03d98a94..5834f2551e1 100644 --- a/backends/vulkan/runtime/graph/ops/impl/QuantizeDequantize.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/QuantizeDequantize.cpp @@ -32,11 +32,10 @@ bool is_gemv(ComputeGraph* graph, const ValueRef& fp_input) { std::tuple get_quantized_input_num_blocks( ComputeGraph& graph, const ValueRef input) { - std::vector input_sizes = graph.sizes_of(input); - const int64_t ndim = graph.dim_of(input); + const std::vector input_sizes = graph.sizes_of(input); - const int64_t M = input_sizes.at(ndim - 2); - const int64_t K = input_sizes.at(ndim - 1); + const int64_t M = utils::val_at(-2, input_sizes); + const int64_t K = utils::val_at(-1, input_sizes); const int64_t num_blocks_M = utils::div_up(M, int64_t(4)); const int64_t num_blocks_K = utils::div_up(K, int64_t(4)); diff --git a/backends/vulkan/test/glsl/binary_op_nobroadcast__test.glsl b/backends/vulkan/test/glsl/binary_op_nobroadcast__test.glsl index 7f72ac58972..5fcd1513c4c 100644 --- a/backends/vulkan/test/glsl/binary_op_nobroadcast__test.glsl +++ b/backends/vulkan/test/glsl/binary_op_nobroadcast__test.glsl @@ -14,9 +14,7 @@ layout(std430) buffer; -// clang-format off -layout(set = 0, binding = 0, ${IMAGE_FORMAT[DTYPE]}) uniform PRECISION restrict writeonly image3D image_out; -// clang-format on +${layout_declare_image(0, "w", "image_out", DTYPE)} layout(set = 0, binding = 1) uniform PRECISION sampler3D image_in; layout(set = 0, binding = 2) uniform PRECISION sampler3D image_other; diff --git a/backends/vulkan/test/glsl/fill_texture__test.glsl b/backends/vulkan/test/glsl/fill_texture__test.glsl index 76c630de55e..637c12fbb51 100644 --- a/backends/vulkan/test/glsl/fill_texture__test.glsl +++ b/backends/vulkan/test/glsl/fill_texture__test.glsl @@ -12,7 +12,7 @@ layout(std430) buffer; -layout(set = 0, binding = 0, ${IMAGE_FORMAT[DTYPE]}) uniform PRECISION restrict writeonly ${IMAGE_T[NDIM][DTYPE]} uOutput; +${layout_declare_image(0, "w", "uOutput", DTYPE, image_ndim=NDIM)} layout(set = 0, binding = 1) uniform PRECISION restrict Block { ivec3 size; int fill; diff --git a/backends/vulkan/test/spv_codegen/test_gen_vulkan_spv.py b/backends/vulkan/test/spv_codegen/test_gen_vulkan_spv.py index aef464dbf38..e569d6fe27d 100644 --- a/backends/vulkan/test/spv_codegen/test_gen_vulkan_spv.py +++ b/backends/vulkan/test/spv_codegen/test_gen_vulkan_spv.py @@ -13,6 +13,22 @@ from runtime import gen_vulkan_spv +class TestF16ImageStores(unittest.TestCase): + def test_writable_rgba16f_images_round_stores(self) -> None: + for access in ("w", "rw"): + with self.subTest(access=access): + declaration = gen_vulkan_spv.layout_declare_image( + 0, access, "t_out", "half" + ) + self.assertIn("round_to_f16(texel)", declaration) + + def test_other_images_store_unchanged(self) -> None: + for access, dtype in (("r", "half"), ("w", "float"), ("w", "int32")): + with self.subTest(access=access, dtype=dtype): + declaration = gen_vulkan_spv.layout_declare_image(0, access, "t", dtype) + self.assertNotIn("round_to_f16", declaration) + + class TestNestedShaderSources(unittest.TestCase): def _write_fake_compiler( self, directory: Path, assertions: str = "" diff --git a/backends/vulkan/test/vulkan_compute_api_test.cpp b/backends/vulkan/test/vulkan_compute_api_test.cpp index 0374505d0bb..7f4ef389040 100644 --- a/backends/vulkan/test/vulkan_compute_api_test.cpp +++ b/backends/vulkan/test/vulkan_compute_api_test.cpp @@ -1202,6 +1202,71 @@ TEST_F(VulkanComputeAPITest, spec_var_shader_test) { } } +TEST_F(VulkanComputeAPITest, half_texture_store_rounds_to_nearest_even) { + // Ties, their neighbors, f16 subnormals and values at and past the largest + // f16, where rounding toward zero differs from rounding to nearest even. + const std::vector values = { + 1.0f + 0x1p-11f, + 1.0f + 0x3p-11f, + -(1.0f + 0x3p-11f), + std::nextafter(1.0f + 0x1p-11f, 0.0f), + std::nextafter(1.0f + 0x1p-11f, 2.0f), + 0x1p-25f, + 0x3p-25f, + std::nextafter(0x1p-25f, 0.0f), + 0x1.ffcp-15f, + -0x1.0p-24f, + 65504.0f, + 65519.0f, + 65520.0f, + -65520.0f, + 1e6f, + -1e6f, + }; + + struct Params final { + utils::ivec3 size; + int32_t fill; + utils::vec4 values; + }; + + for (size_t i = 0; i < values.size(); i += 4) { + vTensor a( + context(), + {4, 1, 1}, + vkapi::kHalf, + utils::kTexture3D, + utils::kChannelsPacked); + Params block{ + {1, 1, 1}, + 0, + {values[i], values[i + 1], values[i + 2], values[i + 3]}, + }; + ParamsBuffer params(context(), block); + + vkapi::PipelineBarrier pipeline_barrier{}; + context()->submit_compute_job( + VK_KERNEL(fill_texture__test_half), + pipeline_barrier, + GlobalWorkGrid({1u, 1u, 1u}, kExplicitWorkGrid), + LocalWorkGroup(1u, 1u, 1u), + {}, + VK_NULL_HANDLE, + 0, + a.image( + pipeline_barrier, + vkapi::PipelineStage::COMPUTE, + vkapi::MemoryAccessType::WRITE), + params.buffer()); + + const std::vector out = extract_vtensor(a); + for (size_t j = 0; j < 4; ++j) { + EXPECT_EQ(out[j], float(executorch::aten::Half(values[i + j]))) + << "input " << values[i + j]; + } + } +} + TEST_F(VulkanComputeAPITest, update_params_between_submit) { context()->set_cmd(/*reusable = */ true); std::vector sizes = {4, 4, 2}; @@ -2465,6 +2530,84 @@ TEST(VulkanComputeGraphTest, choose_qparams_handles_dynamic_row_counts) { } } +TEST(VulkanComputeGraphTest, choose_qparams_affine_handles_1d_input) { + constexpr int64_t kK = 128; + + ComputeGraph graph(GraphConfig{}); + const IOValueRef input = + graph.add_input_tensor({kK}, vkapi::kFloat, utils::kBuffer); + const ValueRef scales = graph.add_tensor( + {1}, vkapi::kFloat, utils::kTexture3D, utils::kWidthPacked); + const ValueRef zero_points = graph.add_tensor( + {1}, vkapi::kChar, utils::kTexture3D, utils::kWidthPacked); + + VK_GET_OP_FN("torchao.choose_qparams_affine.default") + (graph, + {input.value, + graph.add_none(), // mapping_type + graph.add_scalar_list({kK}), + graph.add_none(), // target_dtype + graph.add_scalar(-128), + graph.add_scalar(127), + graph.add_none(), // eps + graph.add_none(), // scale_dtype + graph.add_none(), // zero_point_dtype + graph.add_value_list({scales, zero_points})}); + + const ValueRef scales_staging = graph.set_output_tensor(scales); + const ValueRef zero_points_staging = graph.set_output_tensor(zero_points); + + graph.prepare(); + graph.prepack(); + + // The row's maximum is its last element, so the whole row must be reduced. + std::vector input_data(kK); + for (int64_t k = 0; k < kK; ++k) { + input_data[k] = 0.5f * float(k); + } + graph.maybe_cast_and_copy_into_staging( + input.staging, input_data.data(), input_data.size(), vkapi::kFloat); + + graph.execute(); + + float scale = 0; + int8_t zero_point = 0; + graph.maybe_cast_and_copy_from_staging( + scales_staging, &scale, 1, vkapi::kFloat); + graph.maybe_cast_and_copy_from_staging( + zero_points_staging, &zero_point, 1, vkapi::kChar); + EXPECT_NEAR(scale, input_data.back() / 255.0f, 1e-6f); + EXPECT_EQ(zero_point, -128); +} + +TEST(VulkanComputeGraphTest, quantize_and_pack_accepts_1d_input) { + constexpr int64_t kK = 128; + + ComputeGraph graph(GraphConfig{}); + const IOValueRef input = + graph.add_input_tensor({kK}, vkapi::kFloat, utils::kBuffer); + const ValueRef scales = graph.add_tensor( + {1}, vkapi::kFloat, utils::kTexture3D, utils::kWidthPacked); + const ValueRef zero_points = graph.add_tensor( + {1}, vkapi::kChar, utils::kTexture3D, utils::kWidthPacked); + const ValueRef packed_input = graph.add_tensor( + {1, kK}, vkapi::kInt8x4, utils::kBuffer, utils::kPackedInt8_4H4W); + const ValueRef input_sums = + graph.add_tensor({4}, vkapi::kInt, utils::kBuffer, utils::kWidthPacked); + const QuantizationConfig input_quant_config( + 8, kPerChannel, {1, kK}, false, true); + + EXPECT_NO_THROW(add_quantize_and_pack_4h4w_with_group_sums_node( + graph, + input_quant_config, + input.value, + input_sums, + scales, + zero_points, + packed_input, + graph.add_scalar(kK))); +} + void test_quantize_and_pack_handles_dynamic_row_counts( const int64_t group_size_value, const utils::uvec3& expected_local_wg_size) {