diff --git a/src/ggml-hsa/AGENTS.md b/src/ggml-hsa/AGENTS.md index 5909a32340..387e953c09 100644 --- a/src/ggml-hsa/AGENTS.md +++ b/src/ggml-hsa/AGENTS.md @@ -52,6 +52,7 @@ src/ggml-hsa/ │ ├── soft_max.py # Top-level softmax op dispatch │ ├── clamp.py # Top-level clamp op dispatch │ ├── mul_mat.py # Top-level matrix multiply dispatch +│ ├── pool_2d.py # Top-level 2D pooling op dispatch │ ├── argmax.py # Top-level argmax op dispatch │ ├── count_equal.py # Top-level count_equal op dispatch │ ├── cross_entropy_loss.py # Top-level cross entropy loss op dispatch @@ -69,6 +70,7 @@ src/ggml-hsa/ │ ├── scale.py/cc # Scale IRON design + AIE core function │ ├── softmax.py/cc # Softmax IRON design + AIE core function (unary/masked/ternary variants) │ ├── clamp.py/cc # Clamp IRON design + AIE core function +│ ├── pool_2d.py/cc # 2D pooling IRON design + AIE core function (MAX/AVG, one channel-plane per tile) │ ├── argmax.py/cc # Argmax IRON design + AIE core function │ ├── count_equal.py/cc # Count equal IRON design + AIE core function │ ├── cross_entropy_loss.py/cc # Cross entropy loss IRON design + AIE core function @@ -537,6 +539,7 @@ These operations have complete AIE kernel implementations: | Binary | `ADD`, `SUB`, `MUL`, `DIV` (with broadcast support) | | Unary (GGML_UNARY_OP) | `ABS`, `SGN`, `NEG`, `STEP`, `RELU`, `HARDSWISH`, `HARDSIGMOID`, `FLOOR`, `CEIL`, `ROUND`, `TRUNC` | | Unary (GGML_OP) | `SQR`, `LOG` | +| Pooling | `POOL_2D` (`MAX` and `AVG`, with padding) | | Other | `SCALE`, `SOFT_MAX`, `CLAMP`, `ARGMAX`, `COUNT_EQUAL`, `CROSS_ENTROPY_LOSS`, `MUL_MAT` | | Host-only | `DUP`, `CPY`, `CONT` (run on CPU, not AIE) | diff --git a/src/ggml-hsa/README.md b/src/ggml-hsa/README.md index bea1cc0e4a..265cf0800c 100644 --- a/src/ggml-hsa/README.md +++ b/src/ggml-hsa/README.md @@ -19,6 +19,7 @@ The GGML HSA (`ggml-hsa`) backend enables GGML tensor operations to run on AMD X | Binary | `ADD`, `SUB`, `MUL`, `DIV` (with multi-dimensional broadcast) | | Unary | `SQR`, `LOG`, `ABS`, `SGN`, `NEG`, `STEP`, `FLOOR`, `CEIL`, `ROUND`, `TRUNC`, `RELU`, `HARDSWISH`, `HARDSIGMOID` | | Matrix | `MUL_MAT` | +| Pooling | `POOL_2D` (`MAX` and `AVG`, with padding) | | Reduction | `ARGMAX`, `COUNT_EQUAL` | | Loss | `CROSS_ENTROPY_LOSS` | | Other | `SCALE`, `SOFT_MAX`, `CLAMP` | diff --git a/src/ggml-hsa/kernels/CMakeLists.txt b/src/ggml-hsa/kernels/CMakeLists.txt index 445b74fc12..bc3b49af47 100644 --- a/src/ggml-hsa/kernels/CMakeLists.txt +++ b/src/ggml-hsa/kernels/CMakeLists.txt @@ -19,6 +19,7 @@ if (GGML_HSA_JIT_COMPILE) ${CMAKE_CURRENT_SOURCE_DIR}/cross_entropy_loss.py ${CMAKE_CURRENT_SOURCE_DIR}/kernel.py ${CMAKE_CURRENT_SOURCE_DIR}/mul_mat.py + ${CMAKE_CURRENT_SOURCE_DIR}/pool_2d.py ${CMAKE_CURRENT_SOURCE_DIR}/scale.py ${CMAKE_CURRENT_SOURCE_DIR}/soft_max.py ${CMAKE_CURRENT_SOURCE_DIR}/tensor_desc.py diff --git a/src/ggml-hsa/kernels/__init__.py b/src/ggml-hsa/kernels/__init__.py index 52292a3472..2c52a997e9 100644 --- a/src/ggml-hsa/kernels/__init__.py +++ b/src/ggml-hsa/kernels/__init__.py @@ -1,11 +1,6 @@ # (c) Copyright 2025-2026 Advanced Micro Devices, Inc. or its affiliates -"""GGML HSA Kernels package. - -This package provides IRON-based kernel implementations for GGML operations -targeting AMD AIE (AI Engine) devices. It exposes the main compilation function -and tensor descriptor utilities needed for JIT kernel compilation. -""" +"""GGML HSA kernels for AMD AIE devices: compilation entry point and tensor descriptors.""" from .build import ggml_compile_op from .kernel import Kernel diff --git a/src/ggml-hsa/kernels/argmax.py b/src/ggml-hsa/kernels/argmax.py index c9453d394e..dee917015f 100644 --- a/src/ggml-hsa/kernels/argmax.py +++ b/src/ggml-hsa/kernels/argmax.py @@ -13,22 +13,18 @@ def ggml_op_argmax( arch: str, input_tensors: list, output_tensor, op_params: bytearray ) -> KernelSpec: - """GGML_OP_ARGMAX dispatch function. + """Return the KernelSpec for GGML_OP_ARGMAX. - Finds the index of the maximum value along the first dimension (ne0) of each row. - For a tensor with shape [ne0, ne1, ne2, ne3], computes argmax over ne0 for each - of the ne1 * ne2 * ne3 rows, producing an I32 output tensor with shape [ne1, ne2, ne3]. + Argmax over ne0 for input [ne0, ne1, ne2, ne3], producing I32 output [ne1, ne2, ne3]. Parameters: arch: Target architecture. - input_tensors: List containing exactly one input tensor. - output_tensor: Output tensor of type I32. Shape is - the input shape with the first dimension removed. - op_params: Operation parameters (unused for ARGMAX, but required - by the dispatch interface). + input_tensors: List of one input tensor. + output_tensor: I32 output tensor. + op_params: Unused. Returns: - KernelSpec: Kernel specification for the ARGMAX operation. + KernelSpec for the ARGMAX operation. """ from functools import partial diff --git a/src/ggml-hsa/kernels/binary_ops.py b/src/ggml-hsa/kernels/binary_ops.py index ebd301608f..6859172d9e 100644 --- a/src/ggml-hsa/kernels/binary_ops.py +++ b/src/ggml-hsa/kernels/binary_ops.py @@ -13,7 +13,7 @@ def _validate_binary_inputs(input_tensors: list) -> None: - """Validate the inputs of a binary operation. + """Validate that a binary operation has exactly two input tensors. Parameters: input_tensors: List of input tensors to validate. @@ -33,16 +33,16 @@ def _make_iron_binary_kernel_spec( output_tensor, op_name: str, ) -> KernelSpec: - """Create a KernelSpec for a binary operation targeting the IRON backend. + """Create an IRON-backend KernelSpec for a binary operation. Parameters: arch: Target architecture. input_tensors: List of two input tensors. output_tensor: Output tensor. - op_name: Name of the operation. + op_name: Name of the GGML operation. Returns: - KernelSpec configured for IRON backend. + KernelSpec configured for the IRON backend. """ from functools import partial @@ -70,19 +70,20 @@ def _make_triton_add_kernel_spec( input_tensors: list, output_tensor, ) -> KernelSpec: - """Create a KernelSpec for ADD operation targeting the TRITON backend. + """Create a TRITON-backend KernelSpec for ADD. Parameters: - arch (str): Target architecture. - input_tensors (list): Two input tensors. - output_tensor (TensorDesc): Output tensor. + arch: Target architecture. + input_tensors: List of two input tensors. + output_tensor: Output tensor. Returns: - KernelSpec configured for TRITON backend. + KernelSpec configured for the TRITON backend. Raises: ValueError: If the tensors require broadcasting or are non-contiguous (raised lazily when the returned compile function is invoked). + """ n_elements = output_tensor.numel() @@ -149,17 +150,17 @@ def _compile( def ggml_op_add( arch: str, input_tensors: list, output_tensor, op_params: bytearray ) -> list[KernelSpec]: - """GGML_OP_ADD implementation. + """Return KernelSpecs for GGML_OP_ADD (IRON primary, Triton fallback). Parameters: arch: Target architecture. input_tensors: List of two input tensors. output_tensor: Output tensor. - op_params: Operation parameters (unused for ADD, but required + op_params: Operation parameters (unused for elementwise ops but required by the dispatch interface). Returns: - KernelSpec for the ADD operation. + List of KernelSpecs for the ADD operation. """ _validate_binary_inputs(input_tensors) @@ -175,13 +176,13 @@ def ggml_op_add( def ggml_op_sub( arch: str, input_tensors: list, output_tensor, op_params: bytearray ) -> KernelSpec: - """GGML_OP_SUB implementation. + """Return the KernelSpec for GGML_OP_SUB. Parameters: arch: Target architecture. input_tensors: List of two input tensors. output_tensor: Output tensor. - op_params: Operation parameters (unused for SUB, but required + op_params: Operation parameters (unused for elementwise ops but required by the dispatch interface). Returns: @@ -198,13 +199,13 @@ def ggml_op_sub( def ggml_op_mul( arch: str, input_tensors: list, output_tensor, op_params: bytearray ) -> KernelSpec: - """GGML_OP_MUL implementation. + """Return the KernelSpec for GGML_OP_MUL. Parameters: arch: Target architecture. input_tensors: List of two input tensors. output_tensor: Output tensor. - op_params: Operation parameters (unused for MUL, but required + op_params: Operation parameters (unused for elementwise ops but required by the dispatch interface). Returns: @@ -221,13 +222,13 @@ def ggml_op_mul( def ggml_op_div( arch: str, input_tensors: list, output_tensor, op_params: bytearray ) -> KernelSpec: - """GGML_OP_DIV implementation. + """Return the KernelSpec for GGML_OP_DIV. Parameters: arch: Target architecture. input_tensors: List of two input tensors. output_tensor: Output tensor. - op_params: Operation parameters (unused for DIV, but required + op_params: Operation parameters (unused for elementwise ops but required by the dispatch interface). Returns: diff --git a/src/ggml-hsa/kernels/build.py b/src/ggml-hsa/kernels/build.py index 79820e2753..a5f4ca30a0 100644 --- a/src/ggml-hsa/kernels/build.py +++ b/src/ggml-hsa/kernels/build.py @@ -1,14 +1,10 @@ # Copyright (c) 2025-2026 Advanced Micro Devices, Inc. All Rights Reserved. -"""GGML HSA backend kernel build system. +"""GGML HSA kernel build system for AMD XDNA / XDNA2 devices. -This module provides the infrastructure for compiling kernels to executable code -for AMD XDNA / XDNA2 devices. It handles mapping GGML operations to their corresponding -kernel implementations, dynamic module loading, and orchestrating the compilation -pipeline. - -The build system supports multiple compilation backends with per-operation dispatch -based on compilation parameters. +Maps GGML operations to kernel implementations, dynamically loads dispatch +modules, and orchestrates compilation across multiple backends with +per-operation dispatch. Usage: As a module: @@ -40,21 +36,17 @@ def _get_compiler(backend: Backend) -> Callable: - """Get the compiler function for the given backend. + """Return the compiler function for the given backend. Parameters: - backend: The compiler backend to use. - - Returns: - The compiler function for the specified backend. + backend: The backend whose compiler function to return. Raises: NotImplementedError: If the backend is not implemented. Note: - Uses backend.name for lookup to handle the case where Backend enums - from dynamically imported modules have different identity than those - in this module. + Looks up by ``backend.name`` because Backend enums from dynamically + imported modules have different identity than those defined here. """ # Lookup by name to handle different enum class identities from dynamic imports @@ -65,8 +57,7 @@ def _get_compiler(backend: Backend) -> Callable: raise NotImplementedError(msg) -# Mapping of GGML operations to kernel source files. -# Each entry maps an operation name to a Kernel that identifies the dispatch module. +# Maps each GGML operation name to a Kernel identifying its dispatch module. _OP_KERNEL_MAP: dict[str, Kernel] = { # unary operation to kernel source mapping "ABS": Kernel("ggml_unary_op_abs", "unary_ops.py"), @@ -100,6 +91,7 @@ def _get_compiler(backend: Backend) -> Callable: "SIN": Kernel("ggml_op_sin", "unary_ops.py"), "COS": Kernel("ggml_op_cos", "unary_ops.py"), "MUL_MAT": Kernel("ggml_op_mul_mat", "mul_mat.py"), + "POOL_2D": Kernel("ggml_op_pool_2d", "pool_2d.py"), "SCALE": Kernel("ggml_op_scale", "scale.py"), "SOFT_MAX": Kernel("ggml_op_soft_max", "soft_max.py"), "CLAMP": Kernel("ggml_op_clamp", "clamp.py"), @@ -110,13 +102,10 @@ def _get_compiler(backend: Backend) -> Callable: def _get_kernel(op_name: str) -> Kernel: - """Get the kernel for the given operation. + """Return the Kernel for the given operation. Parameters: - op_name: Operation name. - - Returns: - The Kernel object associated with the operation. + op_name: Operation name to look up. Raises: NotImplementedError: If the Kernel is not found. @@ -130,17 +119,11 @@ def _get_kernel(op_name: str) -> Kernel: def _import_from_path(module_name: str, path: str | Path): - """Import a module by name from the specified file path. - - This function handles the complexity of importing Python modules dynamically, - including setting up the package structure for relative imports. + """Dynamically import a module from a file path, wiring up the package structure for relative imports. Parameters: - module_name: Name of the module to import. - path: Path to the Python file containing the module. - - Returns: - The imported module object. + module_name: Name to assign the imported module. + path: File path of the module to import. Raises: ImportError: If the module cannot be found or loaded. @@ -194,9 +177,6 @@ def _setup_logger(name: str, verbose: bool) -> logging.Logger: name: Logger name, typically __name__ of the calling module. verbose: If True, enables DEBUG-level output to stderr. - Returns: - Configured Logger instance. - """ logger = logging.getLogger(name) for handler in logger.handlers.copy(): @@ -223,19 +203,18 @@ def ggml_compile_op( ) -> None: """Compile a GGML operation kernel to PDI and instruction files. - This is the main entry point for kernel compilation. It: - 1. Looks up the kernel dispatch module for the operation - 2. Calls the dispatch function to get a KernelSpec (backend + function) - 3. Invokes the appropriate backend compiler + Main entry point for kernel compilation: looks up the dispatch module, + calls it to obtain a KernelSpec (backend + function), then invokes the + matching backend compiler. Parameters: op_name: Operation name (e.g., "ADD", "MUL_MAT"). arch: Target architecture (e.g., "aie2", "aie2p"). - input_tensors: List of input tensor descriptions. + input_tensors: Input tensor descriptions. output_tensor: Output tensor description. - op_params: Operation-specific parameters as a bytearray. + op_params: Operation-specific parameters. exported_name: Name to export the compiled kernel as. - output_directory: Directory to save the compiled PDI and instruction files. + output_directory: Destination for the PDI and instruction files. verbose: If True, enables verbose logging output. Raises: @@ -319,13 +298,10 @@ def ggml_compile_op( def _to_tuple_of_ints(string: str) -> tuple[int, int, int, int]: - """Convert a string of the form "(x,y,z,w)" to a tuple of integers. + """Convert a string of the form "(x,y,z,w)" to a 4-tuple of integers. Parameters: - string: String representation of a 4-element tuple. - - Returns: - A tuple of 4 integers. + string: String of the form "(x,y,z,w)" to convert. Raises: ValueError: If the string does not represent exactly 4 integers. @@ -341,13 +317,10 @@ def _to_tuple_of_ints(string: str) -> tuple[int, int, int, int]: def _to_tensordesc(string: str) -> TensorDesc: - """Create a TensorDesc from a string representation. + """Create a TensorDesc from a string of the form "(shape)/dtype", e.g. "(1024,1,1,1)/f32". Parameters: - string: String of the form "(shape)/dtype", e.g., "(1024,1,1,1)/f32". - - Returns: - A TensorDesc instance with the specified shape and dtype. + string: String of the form "(shape)/dtype" to convert. """ shape_str, dtype = string.split("/") diff --git a/src/ggml-hsa/kernels/build_iron.py b/src/ggml-hsa/kernels/build_iron.py index 0fb024baa5..641f4f1fb7 100644 --- a/src/ggml-hsa/kernels/build_iron.py +++ b/src/ggml-hsa/kernels/build_iron.py @@ -14,15 +14,11 @@ def _compile_aie_core_kernels( functions: Iterable, work_dir: Path, ) -> None: - """Compile AIE core functions to object files. - - This function compiles the C++ source files for external functions - (core compute kernels) into object files that will be linked into - the final PDI. + """Compile external-function C++ sources into object files linked into the final PDI. Parameters: arch: Target architecture (e.g., "aie2", "aie2p"). - functions: Iterable of ExternalFunction objects to compile. + functions: ExternalFunction objects to compile. work_dir: Working directory for intermediate files. """ @@ -46,12 +42,11 @@ def compile_iron_kernel( logger: logging.Logger, verbose: bool, ) -> None: - """Compile an IRON kernel. + """Run the IRON compilation pipeline for a kernel. - This function executes the IRON compilation pipeline: - 1. Executes the kernel's Python function to generate an MLIR module - 2. Compiles any external C++ core functions to object files - 3. Compiles the MLIR module to produce PDI and instructions binaries + Runs the kernel's Python function to generate an MLIR module, compiles any + external C++ core functions to object files, then compiles the module into + PDI and instruction binaries. Parameters: kernel_spec: The KernelSpec containing the IRON kernel function. diff --git a/src/ggml-hsa/kernels/build_triton.py b/src/ggml-hsa/kernels/build_triton.py index e69bb8eb5b..89c2ac9fea 100644 --- a/src/ggml-hsa/kernels/build_triton.py +++ b/src/ggml-hsa/kernels/build_triton.py @@ -17,16 +17,15 @@ class TempEnvSet(ContextDecorator): """Context manager to temporarily set an environment variable. - This ensures that Triton uses a specific cache directory for compiled artifacts, - which helps with organization and cleanup. - - Parameters: - env_var: Name of the environment variable to set. - value: Value to set for the environment variable. - Usage: with TempEnvSet("TRITON_CACHE_DIR", str(Path("/path/to/cache"))): # Triton compilation code here + + Attributes: + env_var: Name of the environment variable to set. + value: Value to set; if None, the variable is left untouched. + old_value: Original value of the variable, restored on exit. + """ env_var: str @@ -34,11 +33,11 @@ class TempEnvSet(ContextDecorator): old_value: str | None = None def __init__(self, env_var: str, value: str | None) -> None: - """Initialize the context manager with the desired environment variable and value. + """Initialize the context manager. Parameters: env_var: Name of the environment variable to set. - value: Value to set for the environment variable. If None, the variable will not be set. + value: Value to set; if None, the variable is left untouched. """ self.env_var = env_var self.value = value @@ -52,7 +51,14 @@ def __enter__(self) -> None: os.environ[self.env_var] = str(self.value) def __exit__(self, exc_type, exc_val, exc_tb) -> None: - """Restore the original environment variable after exiting the context.""" + """Restore the original environment variable after exiting the context. + + Parameters: + exc_type: Exception type raised in the context, if any. + exc_val: Exception instance raised in the context, if any. + exc_tb: Traceback of the exception raised, if any. + + """ if self.value is None: return if self.old_value is not None: @@ -62,19 +68,17 @@ def __exit__(self, exc_type, exc_val, exc_tb) -> None: def _get_triton_target(kernel_spec: KernelSpec) -> str: - """Returns the Triton target string for a given KernelSpec architecture. + """Return the Triton target string for a KernelSpec's architecture. Maps NPU architecture names to their Triton equivalents and passes GPU - architectures through unchanged. Raises for unsupported architectures. + architectures through unchanged (e.g. "npu1", "npu2", "gfx942"). Parameters: - kernel_spec: The KernelSpec containing the architecture information. - - Returns: - Triton target string (e.g. "npu1", "npu2", "gfx942"). + kernel_spec: The KernelSpec whose architecture to map. Raises: ValueError: If the architecture is not a known NPU or GPU target. + """ if kernel_spec.arch in NPU_ARCH_MAP: return NPU_ARCH_MAP[kernel_spec.arch] @@ -93,10 +97,9 @@ def compile_triton_kernel( ) -> None: """Compile a Triton kernel for the target architecture in kernel_spec. - For NPU targets this runs the Triton-XDNA pipeline and extracts a PDI - and instructions binary from the resulting xclbin. - For GPU targets this runs the standard HIP pipeline and copies the - hsaco object from the Triton cache. + NPU targets run the Triton-XDNA pipeline and extract a PDI and instructions + binary from the resulting xclbin; GPU targets run the HIP pipeline and copy + the hsaco object from the Triton cache. Parameters: kernel_spec: The KernelSpec containing the Triton kernel function. diff --git a/src/ggml-hsa/kernels/clamp.py b/src/ggml-hsa/kernels/clamp.py index 569406de3b..9d57dc1350 100644 --- a/src/ggml-hsa/kernels/clamp.py +++ b/src/ggml-hsa/kernels/clamp.py @@ -13,16 +13,15 @@ def ggml_op_clamp( arch: str, input_tensors: list, output_tensor, op_params: bytearray ) -> KernelSpec: - """GGML_OP_CLAMP implementation. + """Return the KernelSpec for GGML_OP_CLAMP. - Clamps each element of the input tensor to the range [min_val, max_val]. - output[i] = max(min_val, min(input[i], max_val)) + Clamps each element to [min_val, max_val]. Parameters: arch: Target architecture. input_tensors: List of one input tensor. output_tensor: Output tensor. - op_params: Operation parameters containing min and max values. + op_params: Min and max values. Returns: KernelSpec for the CLAMP operation. diff --git a/src/ggml-hsa/kernels/count_equal.py b/src/ggml-hsa/kernels/count_equal.py index 90b92c81c8..025ea95fc7 100644 --- a/src/ggml-hsa/kernels/count_equal.py +++ b/src/ggml-hsa/kernels/count_equal.py @@ -13,23 +13,18 @@ def ggml_op_count_equal( arch: str, input_tensors: list, output_tensor, op_params: bytearray ) -> KernelSpec: - """GGML_OP_COUNT_EQUAL dispatch function. + """Return the KernelSpec for GGML_OP_COUNT_EQUAL. - Counts the number of elements that are equal between two input tensors. - Both input tensors must have the same shape and be of type I32. - The output is a single I64 scalar containing the count. + Counts elementwise-equal entries between two same-shaped I32 tensors. Parameters: arch: Target architecture. - input_tensors: List containing exactly two input tensors. Both tensors must be I32 type and contiguous in memory. - output_tensor: Output tensor of type I64 with - shape [1, 1, 1, 1] containing the count of equal elements. - op_params: Operation parameters as a 64-byte buffer (unused - for COUNT_EQUAL, but required by the dispatch interface). + input_tensors: List of two contiguous I32 tensors. + output_tensor: I64 scalar [1, 1, 1, 1] holding the count. + op_params: Unused. Returns: - KernelSpec: Kernel specification with backend=IRON and the count_equal_op - function for generating the MLIR module. + KernelSpec for the COUNT_EQUAL operation. """ from functools import partial diff --git a/src/ggml-hsa/kernels/cross_entropy_loss.py b/src/ggml-hsa/kernels/cross_entropy_loss.py index 37b0df710c..ca0a73ee6c 100644 --- a/src/ggml-hsa/kernels/cross_entropy_loss.py +++ b/src/ggml-hsa/kernels/cross_entropy_loss.py @@ -13,15 +13,13 @@ def ggml_op_cross_entropy_loss( arch: str, input_tensors: list, output_tensor, op_params: bytearray ) -> KernelSpec: - """GGML_OP_CROSS_ENTROPY_LOSS implementation. + """Return the KernelSpec for GGML_OP_CROSS_ENTROPY_LOSS. Parameters: arch: Target architecture. - input_tensors: List of 2 input tensors: - - input_tensors[0]: Logits tensor (predictions before softmax) - - input_tensors[1]: Labels tensor (ground truth, often one-hot encoded) - output_tensor: Output scalar tensor containing the loss value. - op_params: Operation parameters (currently unused). + input_tensors: [logits, labels]. + output_tensor: Scalar loss value. + op_params: Unused. Returns: KernelSpec for the CROSS_ENTROPY_LOSS operation. diff --git a/src/ggml-hsa/kernels/iron_kernels/CMakeLists.txt b/src/ggml-hsa/kernels/iron_kernels/CMakeLists.txt index 2c4e514bb5..e8844e14d0 100644 --- a/src/ggml-hsa/kernels/iron_kernels/CMakeLists.txt +++ b/src/ggml-hsa/kernels/iron_kernels/CMakeLists.txt @@ -17,6 +17,8 @@ set(IRON_FILES ${CMAKE_CURRENT_SOURCE_DIR}/cross_entropy_loss.py ${CMAKE_CURRENT_SOURCE_DIR}/gemm.py ${CMAKE_CURRENT_SOURCE_DIR}/ggml-aie.hpp + ${CMAKE_CURRENT_SOURCE_DIR}/pool_2d.cc + ${CMAKE_CURRENT_SOURCE_DIR}/pool_2d.py ${CMAKE_CURRENT_SOURCE_DIR}/scale.cc ${CMAKE_CURRENT_SOURCE_DIR}/scale.py ${CMAKE_CURRENT_SOURCE_DIR}/softmax.cc diff --git a/src/ggml-hsa/kernels/iron_kernels/__init__.py b/src/ggml-hsa/kernels/iron_kernels/__init__.py index f3c675bcfb..f8fa5b7895 100644 --- a/src/ggml-hsa/kernels/iron_kernels/__init__.py +++ b/src/ggml-hsa/kernels/iron_kernels/__init__.py @@ -5,9 +5,7 @@ # # (c) Copyright 2026 Advanced Micro Devices, Inc. or its affiliates -"""IRON kernel implementations. +"""IRON (Intermediate Representation for Optimized NPU) kernel designs. -This package contains the low-level IRON (Intermediate Representation for -Optimized NPU) kernel implementations for various GGML operations. Each module -provides kernel designs that generate MLIR code for AMD XDNA / XDNA2 NPUs. +Each module generates MLIR for a GGML operation targeting AMD XDNA / XDNA2 NPUs. """ diff --git a/src/ggml-hsa/kernels/iron_kernels/argmax.py b/src/ggml-hsa/kernels/iron_kernels/argmax.py index b04b6d44e8..fc2a79657f 100644 --- a/src/ggml-hsa/kernels/iron_kernels/argmax.py +++ b/src/ggml-hsa/kernels/iron_kernels/argmax.py @@ -5,9 +5,9 @@ # # (c) Copyright 2026 Advanced Micro Devices, Inc. or its affiliates -"""IRON kernel implementation for the argmax operation. +"""IRON design for argmax: index of the max value along dim 0, per row. -Finds the index of the maximum value along the first dimension (columns) for each row. +Each worker iteration processes one row and emits a single I32 index. """ from pathlib import Path @@ -28,29 +28,19 @@ def argmax_op(arch: str, input_tensors: list, output_tensor): - """IRON design for argmax. - - Computes the index of the maximum value along the first dimension for each row. - Uses row-by-row processing where each kernel invocation processes one row and - outputs a single I32 index. + """Build the argmax IRON program. Parameters: arch: Target architecture. - input_tensors: List containing exactly one input tensor. - The tensor must be F32 with shape [ne0, ne1, ne2, ne3] where ne0 is the - row length (dimension over which argmax is computed) and the product - ne1 * ne2 * ne3 is the number of rows. - output_tensor: Output tensor of type I32 with shape [ne1, ne2, ne3] - containing one index per row indicating the position of the maximum value. + input_tensors: One F32 tensor [ne0, ne1, ne2, ne3]; ne0 is the row length + and ne1 * ne2 * ne3 the number of rows. + output_tensor: I32 tensor holding one index per row. Returns: - MLIR module representing the IRON program for argmax. + The resolved IRON program (MLIR module). Raises: - ValueError: If input_tensors does not contain exactly one tensor. - ValueError: If input or output tensors are not contiguous in memory. - ValueError: If output tensor size does not match the number of input rows. - ValueError: If output tensor dtype is not int32. + ValueError: On invalid tensor count, contiguity, output size, or dtype. """ if len(input_tensors) != 1: @@ -128,22 +118,16 @@ def _create_external_function( output_tensor, row_length: int, ) -> ExternalFunction: - """Create an ExternalFunction specification for argmax. - - The external function wraps the C++ kernel that performs the actual argmax - computation on the AIE tile. The kernel receives one row of input data and - outputs a single I32 index. + """Create the ExternalFunction wrapping argmax.cc. Parameters: - op_name: Operation name used for function naming and compile flags. + op_name: Operation name (drives function name and compile flags). input_tensor: Input tensor. output_tensor: Output tensor. - row_length: Number of elements per row (ne0 dimension). + row_length: Number of elements per row (ne0). Returns: - ExternalFunction: Configured external function specification that references - the argmax.cc source file with appropriate compile flags for dtype and - vector size configuration. + The configured ExternalFunction. """ current_dir = Path(__file__).resolve().parent diff --git a/src/ggml-hsa/kernels/iron_kernels/binary_ops.py b/src/ggml-hsa/kernels/iron_kernels/binary_ops.py index 17da682dfa..8c71277435 100644 --- a/src/ggml-hsa/kernels/iron_kernels/binary_ops.py +++ b/src/ggml-hsa/kernels/iron_kernels/binary_ops.py @@ -31,25 +31,14 @@ def _ggml_can_repeat(t0_shape: tuple, t1_shape: tuple) -> bool: - """Python reimplementation of ggml_can_repeat. - - Checks if tensor t0 can be repeated to fill tensor t1. - This is the GGML broadcast semantic: t1->ne[i] % t0->ne[i] == 0 for all dims. - - From ggml.c: - bool ggml_can_repeat(const struct ggml_tensor * t0, const struct ggml_tensor * t1) { - return (t1->ne[0]%t0->ne[0] == 0) && - (t1->ne[1]%t0->ne[1] == 0) && - (t1->ne[2]%t0->ne[2] == 0) && - (t1->ne[3]%t0->ne[3] == 0); - } + """Whether t0 can be repeated to fill t1 (GGML broadcast: t1[i] % t0[i] == 0). Parameters: - t0_shape: Shape of the smaller tensor to be repeated. - t1_shape: Shape of the larger tensor to fill. + t0_shape: Shape of the tensor to be repeated. + t1_shape: Target shape to fill. Returns: - True if t0 can be repeated to fill t1. + True if t0 can be broadcast to t1. """ return all(t1_shape[i] % t0_shape[i] == 0 for i in range(4)) @@ -57,11 +46,11 @@ def _ggml_can_repeat(t0_shape: tuple, t1_shape: tuple) -> bool: @dataclass(frozen=True) class CoreFunctionSpec: - """Specification for a core function to be used in binary operations. + """Core function plus total element count for an element-wise binary op. Attributes: - external_function: The external function to be called for the binary operation. - num_elements: The total number of elements in the input/output tensors. + external_function: External function implementing the operation. + num_elements: Total number of elements to process. """ @@ -70,7 +59,7 @@ class CoreFunctionSpec: @property def tile_size(self) -> int: - """Returns the tile size used by the external function.""" + """Tile size used by the external function.""" return self.external_function.tile_size(0) @@ -80,12 +69,12 @@ def _binary_op( function_spec: CoreFunctionSpec, output_tensor, ): - """Implement output_tensor = op(*input_tensors). + """Element-wise output_tensor = op(*input_tensors). Parameters: arch: Target architecture. - input_tensors: Input tensors. - function_spec: Binary operator specification. + input_tensors: Input tensors [src0, src1]. + function_spec: Core function specification. output_tensor: Output tensor. """ @@ -152,17 +141,14 @@ def _create_external_function( input_tensors: list, output_tensor, ) -> CoreFunctionSpec: - """Create a specification for binary ops. + """Create the CoreFunctionSpec for an element-wise binary op. Parameters: arch: Target architecture. op_name: Name of the operation. - input_tensors: List of input tensors. + input_tensors: Two input tensors [src0, src1]. output_tensor: Output tensor. - Returns: - CoreFunctionSpec: Specification for the core function to be used in binary ops. - """ num_elements = arch_aligned_num_elements(arch=arch, tensor=output_tensor) tile_size = max_tile_size(arch, output_tensor.dtype, num_elements) @@ -190,14 +176,14 @@ def _create_external_function( @dataclass(frozen=True) class BroadcastFunctionSpec: - """Specification for a broadcast binary operation. + """Core function and shapes for a broadcast binary op (src1 repeated). Attributes: - external_function: The external function for broadcast op. - num_elements_out: Total number of elements in output (and src0). - num_elements_src1: Total number of elements in src1 (smaller). - src1_ne: Shape of src1 as 4-element tuple (ne0, ne1, ne2, ne3). - dst_ne: Shape of dst as 4-element tuple (ne0, ne1, ne2, ne3). + external_function: External function implementing the operation. + num_elements_out: Total number of output elements. + num_elements_src1: Total number of src1 elements. + src1_ne: src1 shape as (ne0, ne1, ne2, ne3). + dst_ne: Destination shape as (ne0, ne1, ne2, ne3). """ @@ -209,7 +195,7 @@ class BroadcastFunctionSpec: @property def tile_size(self) -> int: - """Returns the tile size used by the external function.""" + """Tile size used by the external function.""" return self.external_function.tile_size(0) @@ -219,20 +205,17 @@ def _create_broadcast_external_function( input_tensors: list, output_tensor, ) -> BroadcastFunctionSpec: - """Create a specification for broadcast binary ops. + """Create the BroadcastFunctionSpec for a broadcast binary op. - In broadcast mode, src1 is smaller than src0/dst and gets repeated. - The kernel receives the full src1 buffer and uses modulo indexing. + src1 is smaller than src0/dst; the kernel gets the full src1 buffer and + uses modulo indexing to repeat it. Parameters: arch: Target architecture. op_name: Name of the operation. - input_tensors: List of input tensors [src0, src1]. + input_tensors: Two input tensors [src0, src1]. output_tensor: Output tensor. - Returns: - BroadcastFunctionSpec: Specification for broadcast binary ops. - """ num_elements_out = arch_aligned_num_elements(arch=arch, tensor=output_tensor) num_elements_src1 = arch_aligned_num_elements(arch=arch, tensor=input_tensors[1]) @@ -285,12 +268,12 @@ def _binary_op_broadcast( function_spec: BroadcastFunctionSpec, output_tensor, ): - """Binary op with broadcasting - src1 loaded fully once, src0 streamed in tiles. + """Broadcast binary op: src1 loaded fully once, src0 streamed in tiles. Parameters: arch: Target architecture. input_tensors: Input tensors [src0, src1]. - function_spec: Broadcast operation specification. + function_spec: Broadcast function specification. output_tensor: Output tensor. """ @@ -372,15 +355,12 @@ def binary_op( input_tensors: list, output_tensor, ): - """IRON generic design for binary operations. - - Supports both element-wise operations (same shape) and broadcasting - (src1 smaller, gets repeated to match src0/dst). + """IRON design for binary ops (element-wise, or broadcasting src1 to src0/dst). Parameters: op_name: Name of the operation. arch: Target architecture. - input_tensors: List of two input tensors [src0, src1]. + input_tensors: Two input tensors [src0, src1]. output_tensor: Output tensor. """ diff --git a/src/ggml-hsa/kernels/iron_kernels/clamp.py b/src/ggml-hsa/kernels/iron_kernels/clamp.py index a5b8a83ae0..c8437a32e4 100644 --- a/src/ggml-hsa/kernels/iron_kernels/clamp.py +++ b/src/ggml-hsa/kernels/iron_kernels/clamp.py @@ -34,19 +34,16 @@ def _create_external_function( input_tensor, output_tensor, ) -> tuple[ExternalFunction, int, int]: - """Create an ExternalFunction specification for the clamp operation. + """Create the clamp ExternalFunction. Parameters: arch: Target architecture. - op_name: Operation name used for function naming and compile flags. + op_name: Name of the operation. input_tensor: Input tensor. output_tensor: Output tensor. Returns: - Tuple[ExternalFunction, int, int]: A tuple containing: - - func: The configured ExternalFunction specification. - - num_elements: Architecture-aligned number of elements. - - tile_size: Size of each processing tile. + (func, num_elements, tile_size) where num_elements is arch-aligned. """ num_elements = arch_aligned_num_elements(arch=arch, tensor=input_tensor) @@ -73,16 +70,13 @@ def _create_external_function( def clamp(arch: str, input_tensors: list, output_tensor, op_params: bytearray): - """IRON design for clamp. - - Clamps each element of the input tensor to the range [min_val, max_val]. - output[i] = max(min_val, min(input[i], max_val)) + """IRON design for clamp: output = max(min_val, min(input, max_val)). Parameters: arch: Target architecture. input_tensors: List of one input tensor. output_tensor: Output tensor. - op_params: Operation parameters containing min and max values. + op_params: min_val and max_val as 2 x float32. """ if len(input_tensors) != 1: diff --git a/src/ggml-hsa/kernels/iron_kernels/count_equal.py b/src/ggml-hsa/kernels/iron_kernels/count_equal.py index a1bb9ef573..5176a0e3e5 100644 --- a/src/ggml-hsa/kernels/iron_kernels/count_equal.py +++ b/src/ggml-hsa/kernels/iron_kernels/count_equal.py @@ -5,11 +5,10 @@ # # (c) Copyright 2026 Advanced Micro Devices, Inc. or its affiliates -"""IRON kernel implementation for the count_equal operation. +"""IRON design for count_equal: number of equal elements between two I32 tensors. -Counts the number of elements that are equal between two I32 input tensors. -The output is a single I64 value, but since IRON doesn't support I64 in ObjectFifos, -we use two I32 values (low and high parts) for the transfer. +The count is a single I64 scalar. IRON ObjectFifos lack I64, so it is transferred +as two I32 lanes (low and high 32 bits) that bitwise form the I64. """ from pathlib import Path @@ -31,32 +30,18 @@ def count_equal_op(arch: str, input_tensors: list, output_tensor): - """IRON design for count_equal. - - Counts elements that are equal between two I32 input tensors and outputs - a single I64 scalar with the count. Processes data in tiles. - - Since IRON doesn't support I64 types in ObjectFifos, we transfer the count - as two I32 values (low and high 32 bits). The C++ kernel writes the 64-bit - count as these two I32 lanes to the ObjectFifo output buffer, which together - bitwise represent a single I64 value. + """Build the count_equal IRON program. Parameters: arch: Target architecture. - input_tensors: List containing exactly two input tensors. - Both tensors must be I32 with the same shape. - output_tensor: Output tensor of type I64 with shape [1,1,1,1] - containing the count of equal elements. + input_tensors: Two I32 tensors of identical shape. + output_tensor: I64 scalar tensor, shape [1, 1, 1, 1]. Returns: - MLIR module representing the IRON program for count_equal. + The resolved IRON program (MLIR module). Raises: - ValueError: If input_tensors does not contain exactly two tensors. - ValueError: If input tensors have different shapes. - ValueError: If input or output tensors are not contiguous in memory. - ValueError: If input tensor dtype is not int32. - ValueError: If output tensor dtype is not int64. + ValueError: On invalid tensor count, shape mismatch, contiguity, or dtype. """ if len(input_tensors) != 2: @@ -185,19 +170,15 @@ def _create_external_function( input_tensor, tile_size: int, ) -> ExternalFunction: - """Create an ExternalFunction specification for count_equal. - - The external function wraps the C++ kernel that performs the actual count_equal - computation on the AIE tile. + """Create the ExternalFunction wrapping count_equal.cc. Parameters: - op_name: Operation name used for function naming and compile flags. + op_name: Operation name (drives function name and compile flags). input_tensor: Input tensor. - tile_size: Size of each tile in elements. + tile_size: Number of elements per tile. Returns: - ExternalFunction: Configured external function specification that references - the count_equal.cc source file with appropriate compile flags. + The configured ExternalFunction. """ current_dir = Path(__file__).resolve().parent diff --git a/src/ggml-hsa/kernels/iron_kernels/cross_entropy_loss.py b/src/ggml-hsa/kernels/iron_kernels/cross_entropy_loss.py index 2a0697bef8..de78c43a6b 100644 --- a/src/ggml-hsa/kernels/iron_kernels/cross_entropy_loss.py +++ b/src/ggml-hsa/kernels/iron_kernels/cross_entropy_loss.py @@ -5,7 +5,12 @@ # # (c) Copyright 2026 Advanced Micro Devices, Inc. or its affiliates -"""IRON kernel implementation for the cross entropy loss operation.""" +"""IRON design for GGML_OP_CROSS_ENTROPY_LOSS. + +Computes -sum(labels * log_softmax(logits)) / num_rows with numerically stable +log-softmax (max subtraction). One row is processed per worker iteration and +per-row losses are accumulated on-tile into a single scalar. +""" from pathlib import Path @@ -26,18 +31,19 @@ def get_cross_entropy_loss_dimensions(tensor) -> tuple[int, int]: - """Extract cross entropy loss dimensions from tensor shape. + """Return (row_length, num_rows) for a GGML-ordered tensor. - GGML convention: cross entropy loss is computed over dimension 0 (ne00). - GGML shape ordering: (ne00, ne01, ne02, ne03) where ne00 is innermost. + Loss is computed over dim 0 (ne00), so row_length = ne00 and + num_rows = ne01 * ne02 * ne03. Parameters: - tensor: Input tensor with shape in GGML order. + tensor: GGML-ordered tensor of rank 1 to 4. Returns: - Tuple of (row_length, num_rows) where: - - row_length = ne00 (dimension over which loss is computed per row) - - num_rows = ne01 * ne02 * ne03 (number of independent rows) + The (row_length, num_rows) pair. + + Raises: + ValueError: If the tensor rank is unsupported. """ shape = tensor.shape @@ -63,17 +69,12 @@ def get_cross_entropy_loss_dimensions(tensor) -> tuple[int, int]: def cross_entropy_loss(arch: str, input_tensors: list, output_tensor): - """IRON design for GGML_OP_CROSS_ENTROPY_LOSS implementation. - - Cross entropy loss computes: -sum(labels * log(softmax(logits))) / num_rows - where the softmax is computed with numerical stability. + """Build the cross entropy loss IRON program. Parameters: arch: Target architecture. - input_tensors: List of 2 input tensors: - - input_tensors[0]: Logits tensor (predictions before softmax) - - input_tensors[1]: Labels tensor (ground truth, often one-hot encoded) - output_tensor: Output scalar tensor containing the loss value. + input_tensors: [logits, labels] of identical shape. + output_tensor: Output scalar tensor holding the loss. """ if len(input_tensors) != 2: @@ -140,32 +141,22 @@ def create_reduction_program( tile_size: int, num_rows: int, ): - """Create an IRON program for cross entropy loss with on-tile reduction. - - The C++ kernel computes per-row loss: loss_row = -sum(labels * log_softmax). - The worker accumulates all per-row losses on-tile and outputs a single - scalar: total_loss / num_rows, matching the CPU reference behavior. + """Build the IRON program that reduces per-row losses on-tile. - Algorithm: - 1. Acquire the output FIFO element once (single scalar buffer). - 2. For each row: save accumulated value, call kernel (which overwrites - the buffer with this row's loss), then add accumulated + row loss - and store back. - 3. After all rows: divide by num_rows and release. - 4. DMA drains exactly 1 float to the host. + The output FIFO element is acquired once. For each row the accumulated value + is saved, the kernel overwrites the buffer with the row's loss, and the sum + is stored back. After all rows the total is divided by num_rows and released, + so the DMA drains exactly one float. Parameters: - arch: Target architecture (e.g., "aie2", "aie2p"). - function: The external function for per-row loss. + arch: Target architecture. + function: The per-row loss external function. logits_tensor: Logits tensor. labels_tensor: Labels tensor. output_tensor: Output tensor. tile_size: Number of elements per tile (row length). num_rows: Number of rows to process. - Returns: - MLIR module representing the cross entropy loss program. - """ num_tiles = num_rows @@ -262,11 +253,7 @@ def _create_external_function( output_tensor, tile_size: int, ): - """Create an external function specification for cross entropy loss. - - The external function wraps the C++ kernel that computes per-row loss: - loss = -sum(labels * log_softmax(logits)) using numerically stable - log-softmax with max subtraction. + """Create the ExternalFunction wrapping cross_entropy_loss.cc. Parameters: logits_tensor: Logits tensor. @@ -275,8 +262,7 @@ def _create_external_function( tile_size: Number of elements per tile (equals row length). Returns: - ExternalFunction: Configured external function specification that - references cross_entropy_loss.cc with appropriate compile flags. + The configured ExternalFunction. """ arg_types = [ diff --git a/src/ggml-hsa/kernels/iron_kernels/gemm.py b/src/ggml-hsa/kernels/iron_kernels/gemm.py index c5053eea60..64a913cdf8 100644 --- a/src/ggml-hsa/kernels/iron_kernels/gemm.py +++ b/src/ggml-hsa/kernels/iron_kernels/gemm.py @@ -5,7 +5,13 @@ # # (c) Copyright 2025-2026 AMD Inc. -"""IRON kernel implementation for matrix multiplication (GEMM).""" +"""IRON design for matrix multiplication (GEMM). + +A is tiled into (m, k) blocks broadcast across columns and distributed across +rows; B into (k, n) blocks broadcast across rows and distributed across columns. +Each core accumulates C tiles over the K dimension. r, s, t are the microkernel +MAC dimensions (see microkernel_mac_dim_map). +""" import argparse from pathlib import Path @@ -37,11 +43,7 @@ def main(): - """Command-line entry point for generating matrix multiplication MLIR. - - Parses command-line arguments and generates MLIR code for a matrix - multiplication design with the specified dimensions and configuration. - """ + """CLI entry point: parse arguments and print the generated GEMM MLIR.""" argparser = argparse.ArgumentParser( prog="AIE Matrix Multiplication MLIR Design (Whole Array)", description="Emits MLIR code for a matrix multiplication design of the given input size", @@ -107,7 +109,16 @@ def main(): def ceildiv(a, b): - """Return the ceiling of integer division a/b.""" + """Return the ceiling of integer division a/b. + + Parameters: + a: Dividend. + b: Divisor. + + Returns: + The smallest integer >= a / b. + + """ return (a + b - 1) // b @@ -132,35 +143,34 @@ def my_matmul( object_file, generate_taps=False, ): - """Generate MLIR for tiled matrix multiplication across an AIE array. + """Generate MLIR for tiled GEMM across an AIE array. - This function creates the complete AIE design including tile declarations, - object FIFOs for data movement, compute core logic, and runtime DMA sequences. + Builds tile declarations, object FIFOs, compute cores, and runtime DMA + sequences for C = A @ B. Parameters: dev: Device type ("npu" or "npu2"). - M: Number of rows in matrix A and C. + M: Rows of A and C. K: Inner dimension (columns of A, rows of B). - N: Number of columns in matrix B and C. - m: Tile size in M dimension per core. - k: Tile size in K dimension (shared across all cores). - n: Tile size in N dimension per core. + N: Columns of B and C. + m: Per-core tile size in the M dimension. + k: Per-core tile size in the K dimension. + n: Per-core tile size in the N dimension. n_aie_cols: Number of AIE columns to use (1, 2, 4, or 8). - dtype_in_str: Input data type ("bf16", "i8", or "i16"). - dtype_out_str: Output data type ("bf16", "i8", "i16", "f32", or "i32"). - b_col_maj: If True, matrix B is in column-major layout. - c_col_maj: If True, matrix C is in column-major layout. - use_scalar: If True, use scalar kernels (for debugging small sizes). - emulate_bf16_mmul_with_bfp16: If True, use bfp16 emulation for bf16. - trace_size: Size of trace buffer (0 to disable tracing). - zero_fn: Name of the zero initialization function. - matmul_fn: Name of the matrix multiply accumulate function. - object_file: Name of the compiled object file containing kernels. - generate_taps: If True, return TensorAccessPattern objects for visualization. + dtype_in_str: Input dtype ("bf16", "i8", or "i16"). + dtype_out_str: Output dtype ("bf16", "i8", "i16", "f32", or "i32"). + b_col_maj: B is in column-major layout. + c_col_maj: C is in column-major layout. + use_scalar: Use scalar kernels (for debugging small sizes). + emulate_bf16_mmul_with_bfp16: Use bfp16 emulation for bf16. + trace_size: Trace buffer size (0 disables tracing). + zero_fn: Name of the zero-init function. + matmul_fn: Name of the multiply-accumulate function. + object_file: Compiled object file containing the kernels. + generate_taps: Return TensorAccessSequence objects for visualization. Returns: - If generate_taps is True, returns a tuple of TensorAccessSequence objects - for A, B, and C matrices. Otherwise returns None. + A tuple of (A, B, C) TensorAccessSequences if generate_taps, else None. """ n_aie_rows = 4 @@ -211,7 +221,6 @@ def my_matmul( """B must be tileable into (k, n * n_aie_cols)-sized blocks""" ) - # r, s, t are the dimensions required by the microkernel MAC instructions. if not use_scalar: assert m % r == 0 assert k % s == 0 @@ -673,21 +682,18 @@ def create_mat_mul_external_functions( input_tensors: list, output_tensor, ): - """Create parameters and names of external functions for matrix multiplication. + """Create the zero-init and matmul ExternalFunctions for GEMM. - Args: - arch: Target architecture. - input_tensors: List of two input tensors. - output_tensor: Output tensor. + Parameters: + arch: Target architecture ("aie2" or "aie2p"). + input_tensors: Two input tensors [A, B]. + output_tensor: Output tensor C. Returns: - A tuple containing: - - m: The block size in the M dimension. - - n: The block size in the N dimension. - - k: The block size in the K dimension. - - use_scalar: Boolean indicating if scalar multiplication is used. - - mm_fn: The name of the matrix multiplication function. - - zero_fn: The name of the zeroing function. + (m, n, k, use_scalar, num_cols, zero_fn, matmul_fn). + + Raises: + ValueError: If the architecture is unsupported. """ use_scalar = False @@ -752,15 +758,15 @@ def create_mat_mul_external_functions( def gemm(arch: str, input_tensors: list, output_tensor): - """IRON design for matrix multiplication. + """Build the GEMM IRON program (C = A @ B). - Args: - arch: Target architecture (e.g., "aie2", "aie2p"). - input_tensors: List of two input tensors (A and B). - output_tensor: Output tensor (C). + Parameters: + arch: Target architecture ("aie2" or "aie2p"). + input_tensors: Two input tensors [A, B]. + output_tensor: Output tensor C. Returns: - The MLIR module representing the matrix multiplication operation. + The MLIR module. """ if len(input_tensors) != 2: diff --git a/src/ggml-hsa/kernels/iron_kernels/pool_2d.cc b/src/ggml-hsa/kernels/iron_kernels/pool_2d.cc new file mode 100644 index 0000000000..31d9c8c564 --- /dev/null +++ b/src/ggml-hsa/kernels/iron_kernels/pool_2d.cc @@ -0,0 +1,101 @@ +// Copyright (c) 2026 Advanced Micro Devices, Inc. All Rights Reserved. + +/** + * @file pool_2d.cc + * @brief 2D pooling operation for AIE kernels. + */ + +#include +#include + +#include + +#include "ggml-aie.hpp" + +// Pooling op selector (matches enum ggml_op_pool in include/ggml.h). +constexpr int32_t GGML_OP_POOL_MAX = 0; +constexpr int32_t GGML_OP_POOL_AVG = 1; + +extern "C" { + +/** + * @brief Reduces each k1 x k0 window of one input channel-plane to an output element. + * + * Mirrors ggml_compute_forward_pool_2d for a single channel-plane. Padding is + * handled by skipping out-of-bounds taps; AVG divides by the full k0 * k1 window + * area (matching the GGML CPU reference), not the count of valid taps. Taps are + * accumulated in float. + * + * @param[in] in Input channel-plane of iw * ih elements (row-major, width fastest). + * @param[out] out Output channel-plane of ow * oh elements. + * @param[in] iw Input width. + * @param[in] ih Input height. + * @param[in] ow Output width. + * @param[in] oh Output height. + * @param[in] k0 Kernel width. + * @param[in] k1 Kernel height. + * @param[in] s0 Stride along width. + * @param[in] s1 Stride along height. + * @param[in] p0 Padding along width. + * @param[in] p1 Padding along height. + * @param[in] op Pooling op: GGML_OP_POOL_MAX or GGML_OP_POOL_AVG. + */ +void ggml_op_pool_2d(const INPUT_DTYPE * __restrict in, + OUTPUT_DTYPE * __restrict out, + int32_t iw, + int32_t ih, + int32_t ow, + int32_t oh, + int32_t k0, + int32_t k1, + int32_t s0, + int32_t s1, + int32_t p0, + int32_t p1, + int32_t op) { + static_assert(is_floating_point_v, "INPUT_DTYPE must be a floating-point type"); + static_assert(std::is_same::value, "OUTPUT_DTYPE must be float"); + + event0(); + + const int32_t offset0 = -p0; + const int32_t offset1 = -p1; + + for (int32_t oy = 0; oy < oh; ++oy) { + for (int32_t ox = 0; ox < ow; ++ox) { + auto res = (op == GGML_OP_POOL_MAX) ? std::numeric_limits::lowest() : 0.0f; + + const int32_t ix = offset0 + ox * s0; + const int32_t iy = offset1 + oy * s1; + + for (int32_t ky = 0; ky < k1; ++ky) { + const int32_t y = iy + ky; + if (y < 0 || y >= ih) { + continue; + } + const auto * srow = in + static_cast(y) * iw; + for (int32_t kx = 0; kx < k0; ++kx) { + const int32_t x = ix + kx; + if (x < 0 || x >= iw) { + continue; + } + const auto v = static_cast(srow[x]); + if (op == GGML_OP_POOL_MAX) { + res = (v > res) ? v : res; + } else { + res += v; + } + } + } + + if (op == GGML_OP_POOL_AVG) { + res *= 1.0f / static_cast(k0 * k1); + } + out[oy * ow + ox] = res; + } + } + + event1(); +} + +} // extern "C" diff --git a/src/ggml-hsa/kernels/iron_kernels/pool_2d.py b/src/ggml-hsa/kernels/iron_kernels/pool_2d.py new file mode 100644 index 0000000000..140fcea688 --- /dev/null +++ b/src/ggml-hsa/kernels/iron_kernels/pool_2d.py @@ -0,0 +1,199 @@ +# +# This file is licensed under the Apache License v2.0 with LLVM Exceptions. +# See https://llvm.org/LICENSE.txt for license information. +# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +# +# (c) Copyright 2026 Advanced Micro Devices, Inc. or its affiliates + +"""IRON design for GGML_OP_POOL_2D (max/avg pooling). + +Channel-planes are independent, so one input plane [IW, IH] -> output plane +[OW, OH] is processed per worker iteration across all C * N planes. Input is +float32 or bfloat16; output is float32. + +Out-of-bounds taps are skipped. For MAX pooling this is equivalent to +-FLT_MAX padding; for AVG pooling the divisor is the full k0*k1 window area +(not the count of valid taps), matching the GGML CPU reference. +""" + +import struct +from pathlib import Path + +import numpy as np +from aie.iron import ( + ExternalFunction, + ObjectFifo, + Program, + Runtime, + Worker, + dtype_to_str, +) +from aie.iron.controlflow import range_ +from ml_dtypes import bfloat16 + +from .utils import arch_to_device + +# GGML pooling op selector (matches enum ggml_op_pool in include/ggml.h). +_GGML_OP_POOL_MAX = 0 +_GGML_OP_POOL_AVG = 1 + + +def pool_2d(arch: str, input_tensors: list, output_tensor, op_params: bytearray): + """Build the pooling IRON program. + + Parameters: + arch: Target architecture. + input_tensors: List of one input tensor, shape [IW, IH, C, N]. + output_tensor: Output tensor, shape [OW, OH, C, N]. + op_params: {op, k0, k1, s0, s1, p0, p1} as 7 x int32. + + Returns: + The resolved IRON program (MLIR module). + + Raises: + ValueError: On invalid tensor count, dtype, contiguity, op_params, or + pooling op. + + """ + if len(input_tensors) != 1: + msg = "Operation requires exactly one input tensor." + raise ValueError(msg) + + input_tensor = input_tensors[0] + + if ( + input_tensor.dtype not in (np.float32, bfloat16) + or output_tensor.dtype != np.float32 + ): + msg = ( + f"POOL_2D only supports float32/bfloat16 input and float32 output; " + f"got input dtype={input_tensor.dtype}, output dtype={output_tensor.dtype}." + ) + raise ValueError(msg) + + if not input_tensor.contiguous or not output_tensor.contiguous: + msg = "Input and output tensors must be contiguous in memory." + raise ValueError(msg) + + # op_params: {op, k0, k1, s0, s1, p0, p1} as 7 x int32. + _POOL_2D_PARAMS_SIZE = 7 * 4 # 7 int32 fields + if len(op_params) < _POOL_2D_PARAMS_SIZE: + msg = ( + f"op_params too short: expected at least {_POOL_2D_PARAMS_SIZE} bytes, " + f"got {len(op_params)}." + ) + raise ValueError(msg) + op, k0, k1, s0, s1, p0, p1 = struct.unpack_from("7i", op_params, 0) + + if op not in (_GGML_OP_POOL_MAX, _GGML_OP_POOL_AVG): + msg = f"Unsupported pooling op: {op}." + raise ValueError(msg) + + if k0 <= 0 or k1 <= 0: + msg = f"Kernel dimensions must be positive; got k0={k0}, k1={k1}." + raise ValueError(msg) + + if s0 <= 0 or s1 <= 0: + msg = f"Strides must be positive; got s0={s0}, s1={s1}." + raise ValueError(msg) + + if p0 < 0 or p1 < 0: + msg = f"Padding must be non-negative; got p0={p0}, p1={p1}." + raise ValueError(msg) + + iw, ih, in_c, in_n = input_tensor.shape + ow, oh, out_c, out_n = output_tensor.shape + + if (in_c, in_n) != (out_c, out_n): + msg = ( + f"Channel/batch mismatch: input {(in_c, in_n)} vs output {(out_c, out_n)}." + ) + raise ValueError(msg) + + in_plane = iw * ih + out_plane = ow * oh + num_planes = in_c * in_n + + function = _create_external_function( + op_name="GGML_OP_POOL_2D", + input_tensor=input_tensor, + output_tensor=output_tensor, + in_plane=in_plane, + out_plane=out_plane, + ) + + # AIE-array data movement with object fifos: one channel-plane per tile. + input_tile_ty = np.ndarray[(in_plane,), np.dtype[input_tensor.dtype]] + output_tile_ty = np.ndarray[(out_plane,), np.dtype[output_tensor.dtype]] + of_in = ObjectFifo(input_tile_ty, name="in") + of_out = ObjectFifo(output_tile_ty, name="out") + + def ext_core_fn(of_in, of_out, function): + for _ in range_(num_planes): + elem_in = of_in.acquire(1) + elem_out = of_out.acquire(1) + function(elem_in, elem_out, iw, ih, ow, oh, k0, k1, s0, s1, p0, p1, op) + of_in.release(1) + of_out.release(1) + + worker = Worker(ext_core_fn, fn_args=[of_in.cons(), of_out.prod(), function]) + + # Runtime operations to move data to/from the AIE-array. + rt = Runtime() + input_tensor_ty = np.ndarray[(in_plane * num_planes,), np.dtype[input_tensor.dtype]] + output_tensor_ty = np.ndarray[ + (out_plane * num_planes,), np.dtype[output_tensor.dtype] + ] + with rt.sequence(input_tensor_ty, output_tensor_ty) as (a_in, b_out): + rt.start(worker) + rt.fill(of_in.prod(), a_in) + rt.drain(of_out.cons(), b_out, wait=True) + + return Program(arch_to_device(arch), rt).resolve_program() + + +def _create_external_function( + op_name: str, + input_tensor, + output_tensor, + in_plane: int, + out_plane: int, +) -> ExternalFunction: + """Create the ExternalFunction for the pooling core function. + + Parameters: + op_name: Operation name (drives function name and compile flags). + input_tensor: Input tensor. + output_tensor: Output tensor. + in_plane: Input channel-plane size (IW * IH). + out_plane: Output channel-plane size (OW * OH). + + Returns: + The configured ExternalFunction. + + """ + current_dir = Path(__file__).resolve().parent + return ExternalFunction( + name=op_name.lower(), + object_file_name=f"{op_name.lower()}_core_function.o", + source_file=str(current_dir / "pool_2d.cc"), + arg_types=[ + np.ndarray[(in_plane,), np.dtype[input_tensor.dtype]], + np.ndarray[(out_plane,), np.dtype[output_tensor.dtype]], + np.int32, # iw + np.int32, # ih + np.int32, # ow + np.int32, # oh + np.int32, # k0 + np.int32, # k1 + np.int32, # s0 + np.int32, # s1 + np.int32, # p0 + np.int32, # p1 + np.int32, # op + ], + compile_flags=[ + f"-DINPUT_DTYPE={dtype_to_str(input_tensor.dtype)}", + f"-DOUTPUT_DTYPE={dtype_to_str(output_tensor.dtype)}", + ], + ) diff --git a/src/ggml-hsa/kernels/iron_kernels/scale.py b/src/ggml-hsa/kernels/iron_kernels/scale.py index e8d20a89b7..fab74cad4f 100644 --- a/src/ggml-hsa/kernels/iron_kernels/scale.py +++ b/src/ggml-hsa/kernels/iron_kernels/scale.py @@ -25,13 +25,13 @@ def scale(arch: str, input_tensors: list, output_tensor, op_params: bytearray): - """IRON design for scale. + """IRON design for scale: output = input * s + b. Parameters: arch: Target architecture. input_tensors: List of one input tensor. output_tensor: Output tensor. - op_params: Operation parameters. + op_params: s and b as 2 x float32. """ if len(input_tensors) != 1: @@ -99,19 +99,16 @@ def _create_external_function( input_tensor, output_tensor, ) -> tuple[ExternalFunction, int, int]: - """Create an ExternalFunction specification for the scale operation. + """Create the scale ExternalFunction. Parameters: - arch: Target architecture (e.g., "aie2", "aie2p"). - op_name: Operation name used for function naming and compile flags. + arch: Target architecture. + op_name: Name of the operation. input_tensor: Input tensor. output_tensor: Output tensor. Returns: - Tuple[ExternalFunction, int, int]: A tuple containing: - - func: The configured ExternalFunction specification. - - num_elements: Architecture-aligned number of elements. - - tile_size: Size of each processing tile. + (func, num_elements, tile_size) where num_elements is arch-aligned. """ num_elements = arch_aligned_num_elements(arch=arch, tensor=input_tensor) diff --git a/src/ggml-hsa/kernels/iron_kernels/softmax.py b/src/ggml-hsa/kernels/iron_kernels/softmax.py index 3f86d226b3..134427774a 100644 --- a/src/ggml-hsa/kernels/iron_kernels/softmax.py +++ b/src/ggml-hsa/kernels/iron_kernels/softmax.py @@ -5,7 +5,12 @@ # # (c) Copyright 2026 Advanced Micro Devices, Inc. or its affiliates -"""IRON kernel implementation for the softmax operation.""" +"""IRON design for softmax over dim 0, one row per tile. + +Three variants by input count: plain softmax, masked softmax (mask added before +exp, ALiBi slopes when max_bias > 0), and masked softmax with per-head sinks. +The sink array is loaded once and indexed by tile via rows_per_head. +""" import struct from pathlib import Path @@ -28,18 +33,19 @@ def get_softmax_dimensions(tensor) -> tuple[int, int]: - """Extract softmax dimensions from tensor shape. + """Return (row_length, num_rows) for a GGML-ordered tensor. - GGML convention: softmax is over dimension 0 (ne00). - GGML shape ordering: (ne00, ne01, ne02, ne03) where ne00 is innermost. + Softmax is over dim 0 (ne00), so row_length = ne00 and + num_rows = ne01 * ne02 * ne03. Parameters: - tensor: Input tensor with shape in GGML order. + tensor: GGML-ordered tensor of rank 1 to 4. Returns: - Tuple of (row_length, num_rows) where: - - row_length = ne00 (dimension over which softmax is computed) - - num_rows = ne01 * ne02 * ne03 (number of independent rows) + The (row_length, num_rows) pair. + + Raises: + ValueError: If the tensor rank is unsupported. """ shape = tensor.shape @@ -61,16 +67,13 @@ def get_softmax_dimensions(tensor) -> tuple[int, int]: def softmax(arch: str, input_tensors: list, output_tensor, op_params: bytearray): - """IRON design for softmax. + """Build the softmax IRON program, dispatching by input count. Parameters: arch: Target architecture. - input_tensors: List of input tensors: - - input_tensors[0]: Input tensor (required) - - input_tensors[1]: Mask tensor (optional; may be None) - - input_tensors[2]: Sink tensor (optional; may be None) + input_tensors: [input, optional mask, optional sink]. output_tensor: Output tensor. - op_params: Operation parameters (scale, max_bias). + op_params: scale and max_bias as 2 x float32. """ input_tensor_count = len(input_tensors) @@ -146,18 +149,18 @@ def softmax(arch: str, input_tensors: list, output_tensor, op_params: bytearray) def create_unary_program(arch, op_name, input_tensor, output_tensor, scale, max_bias): - """Create an IRON program for basic softmax without mask or sink tensors. + """Plain softmax without mask or sink (max_bias unused). Parameters: arch: Target architecture. - op_name: Operation name for the external function. + op_name: GGML operation name. input_tensor: Input tensor. output_tensor: Output tensor. - scale: Scaling factor applied before exponentiation. - max_bias: Maximum bias (unused in unary variant). + scale: Scale applied before exp. + max_bias: Unused in this variant. Returns: - MLIR module representing the softmax program. + The resolved IRON program. """ function, num_elements, tile_size = _create_external_function( @@ -202,23 +205,19 @@ def ext_core_fn(of_in, of_out, function): def create_binary_program( arch, op_name, input_tensor, mask_tensor, output_tensor, scale, max_bias ): - """Create an IRON program for softmax with a mask tensor. - - This variant supports attention masking where the mask is added to the input - before computing softmax. It also supports ALiBi positional encoding when - max_bias > 0. + """Masked softmax; max_bias > 0 enables ALiBi positional encoding. Parameters: arch: Target architecture. - op_name: Operation name for the external function. + op_name: GGML operation name. input_tensor: Input tensor. - mask_tensor: Mask tensor (added to input before softmax). + mask_tensor: Additive mask applied before exp. output_tensor: Output tensor. - scale: Scaling factor applied before exponentiation. - max_bias: Maximum bias for ALiBi positional encoding. + scale: Scale applied before exp. + max_bias: ALiBi max bias; > 0 enables positional slopes. Returns: - MLIR module representing the masked softmax program. + The resolved IRON program. """ func_result = _create_external_function( @@ -309,10 +308,21 @@ def create_ternary_program( scale, max_bias, ): - """Softmax with mask tensor and sink tensor. + """Masked softmax with a per-head sink array (loaded once, indexed by tile). + + Parameters: + arch: Target architecture. + op_name: GGML operation name. + input_tensor: Input tensor. + mask_tensor: Additive mask applied before exp. + sink_tensor: Per-head sink values, one per head. + output_tensor: Output tensor. + scale: Scale applied before exp. + max_bias: ALiBi max bias; > 0 enables positional slopes. + + Returns: + The resolved IRON program. - Sink tensor contains one value per head. The kernel receives the full - sink array and indexes into it based on tile_idx and rows_per_head. """ func_result = _create_external_function( op_name=op_name, @@ -415,6 +425,13 @@ def _create_external_function( ) -> tuple: """Create an external function specification for softmax variants. + Parameters: + op_name: GGML operation name. + input_tensor: Input tensor. + mask_tensor: Optional additive mask, or None. + sink_tensor: Optional per-head sink array, or None. + output_tensor: Output tensor. + Returns: If no mask or sink tensor: (func, num_elements_in, tile_size_in) diff --git a/src/ggml-hsa/kernels/iron_kernels/unary_ops.py b/src/ggml-hsa/kernels/iron_kernels/unary_ops.py index c7d30396da..9e02a7700d 100644 --- a/src/ggml-hsa/kernels/iron_kernels/unary_ops.py +++ b/src/ggml-hsa/kernels/iron_kernels/unary_ops.py @@ -26,11 +26,11 @@ @dataclass(frozen=True) class CoreFunctionSpec: - """Specification for a core function to be used in unary operations. + """Core function plus total element count for a unary op. Attributes: - external_function: The external function to be called for the unary operation. - num_elements: The total number of elements in the input/output tensors. + external_function: External function implementing the operation. + num_elements: Total number of elements to process. """ @@ -39,7 +39,7 @@ class CoreFunctionSpec: @property def tile_size(self) -> int: - """Returns the tile size used by the external function.""" + """Tile size used by the external function.""" return self.external_function.tile_size(0) @@ -49,12 +49,12 @@ def _unary_op( function_spec: CoreFunctionSpec, output_tensor, ): - """Implement output_tensor = op(input_tensors[0]). + """Element-wise output_tensor = op(input_tensors[0]). Parameters: arch: Target architecture. - input_tensors: Input tensors. - function_spec: Unary operator specification. + input_tensors: List of one input tensor. + function_spec: Core function specification. output_tensor: Output tensor. """ @@ -112,17 +112,14 @@ def _create_external_function( input_tensor, output_tensor, ) -> CoreFunctionSpec: - """Create a specification for unary ops. + """Create the CoreFunctionSpec for a unary op. Parameters: arch: Target architecture. - op_name: Name of the operation. + op_name: Name of the unary operation. input_tensor: Input tensor. output_tensor: Output tensor. - Returns: - CoreFunctionSpec: Specification for the core function to be used in unary ops. - """ num_elements = arch_aligned_num_elements(arch=arch, tensor=input_tensor) tile_size = max_tile_size(arch, input_tensor.dtype, num_elements) @@ -152,7 +149,7 @@ def unary_op( input_tensors: list, output_tensor, ): - """IRON design for unary operations. + """IRON design for unary element-wise operations. Parameters: op_name: Name of the unary operation. diff --git a/src/ggml-hsa/kernels/iron_kernels/utils.py b/src/ggml-hsa/kernels/iron_kernels/utils.py index 47f1820369..8e460caa8b 100644 --- a/src/ggml-hsa/kernels/iron_kernels/utils.py +++ b/src/ggml-hsa/kernels/iron_kernels/utils.py @@ -9,16 +9,16 @@ def align_to_arch( arch: str, size: int, dtype: np.dtype, alignment_bytes: int = 4 ) -> int: - """Align a size to architecture requirements. + """Align an element count so its byte size is a multiple of alignment_bytes. Parameters: arch: Target architecture. - size: Size to align (number of elements). - dtype: Data type of elements. - alignment_bytes: Alignment in bytes. + size: Element count to align. + dtype: Element data type. + alignment_bytes: Byte boundary to align to. Returns: - int: Aligned size. + The aligned element count. """ if arch in ["aie2", "aie2p"]: @@ -36,32 +36,29 @@ def align_to_arch( def arch_aligned_num_elements(arch: str, tensor) -> int: - """Align number of elements to architecture requirements. - - Return the number of elements in the tensor aligned to what the architecture - expects for the data type of the tensor. + """Tensor element count aligned to the architecture for its dtype. Parameters: arch: Target architecture. - tensor: Tensor. + tensor: Tensor whose element count is aligned. Returns: - int: Number of elements aligned to architecture requirements. + The arch-aligned element count. """ return align_to_arch(arch, tensor.numel(), tensor.dtype) def max_tile_size(arch: str, dtype: np.dtype, num_elements: int) -> int: - """Return the maximum tile size based on device, data type and number of elements. + """Largest power-of-two tile within a 512-bit vector dividing num_elements. Parameters: arch: Target architecture. - dtype: Data type of the tensor elements. - num_elements: Total number of elements in the tensor. + dtype: Element data type. + num_elements: Total number of elements to tile. Returns: - int: Maximum tile size. + The chosen tile size. """ vector_register_width = 0 @@ -84,16 +81,13 @@ def max_tile_size(arch: str, dtype: np.dtype, num_elements: int) -> int: def arch_to_device(device): - """Convert an architecture string to an IRON device object. + """Map "aie2" -> NPU1, "aie2p" -> NPU2; pass through existing device objects. Parameters: - device: Architecture string ("aie2" or "aie2p") or an existing device object. + device: Architecture string or an existing device object. Returns: - NPU1 for "aie2", NPU2 for "aie2p", or the input if already a device object. - - Raises: - ValueError: If the architecture string is not supported. + The corresponding device object. """ if isinstance(device, str): diff --git a/src/ggml-hsa/kernels/kernel.py b/src/ggml-hsa/kernels/kernel.py index 0f24915076..18f550b50f 100644 --- a/src/ggml-hsa/kernels/kernel.py +++ b/src/ggml-hsa/kernels/kernel.py @@ -1,13 +1,12 @@ # Copyright (c) 2026 Advanced Micro Devices, Inc. All Rights Reserved. -"""Kernel specification types for the GGML HSA backend. +"""Kernel dispatch and backend-selection types for the GGML HSA backend. -This module defines the core data structures used for kernel dispatch and -compilation backend selection. The two-layer architecture separates: +Two layers separate concerns: -1. Static mapping (Kernel): Maps GGML operation names to dispatch modules -2. Runtime dispatch (KernelSpec): Returned by dispatch functions to specify - which backend and function to use for compilation +1. Kernel (static): maps a GGML operation name to its dispatch module. +2. KernelSpec (runtime): returned by a dispatch function to specify which + backend and function to compile. Example: # In op_to_kernel_map (static) @@ -28,9 +27,8 @@ class Backend(Enum): """Supported kernel compilation backends. - Each backend has its own compilation pipeline: - - IRON: Uses MLIR-AIE/IRON framework for optimized AIE kernels - - TRITON: Uses Triton-XDNA for compiler-driven kernel generation via MLIR-AIR/AIE + - IRON: MLIR-AIE/IRON framework for optimized AIE kernels. + - TRITON: Triton-XDNA for compiler-driven generation via MLIR-AIR/AIE. """ IRON = auto() @@ -39,15 +37,11 @@ class Backend(Enum): @dataclass(frozen=True) class Kernel: - """Static mapping entry from GGML operation to dispatch module. - - This dataclass represents an entry in op_to_kernel_map. It identifies - which Python module contains the dispatch function for a given operation. + """Static op_to_kernel_map entry identifying a dispatch function and its module. Attributes: name: Name of the dispatch function to call (e.g., "ggml_op_add"). - source_file: Path to the Python module containing the - dispatch function. + source_file: Python module containing the dispatch function. """ @@ -57,15 +51,11 @@ class Kernel: @dataclass(frozen=True) class KernelSpec: - """Specification returned by kernel dispatch functions. - - When a kernel dispatch function (e.g., ggml_op_add) is called, it examines - the input parameters and returns a KernelSpec that tells the build system: - 1. Which backend to use for compilation - 2. Which function to call to generate the IR + """Specification returned by a kernel dispatch function. - This enables per-invocation backend selection based on tensor shapes, - dtypes, and other runtime parameters. + Tells the build system which backend to use and which function generates the + IR, enabling per-invocation backend selection based on tensor shapes, dtypes, + and other runtime parameters. Attributes: backend: The compilation backend to use. diff --git a/src/ggml-hsa/kernels/mul_mat.py b/src/ggml-hsa/kernels/mul_mat.py index fafdcb1f2d..4c9938809f 100644 --- a/src/ggml-hsa/kernels/mul_mat.py +++ b/src/ggml-hsa/kernels/mul_mat.py @@ -13,13 +13,13 @@ def ggml_op_mul_mat( arch: str, input_tensors: list, output_tensor, op_params: bytearray ) -> KernelSpec: - """GGML_OP_MUL_MAT implementation. + """Return the KernelSpec for GGML_OP_MUL_MAT. Parameters: - arch: Target architecture (e.g., "aie2", "aie2p"). - input_tensors: List of two input tensors (A and B). - output_tensor: Output tensor (C). - op_params: Operation-specific parameters. + arch: Target architecture. + input_tensors: Input tensors A and B. + output_tensor: Output tensor C. + op_params: Operation parameters. Returns: KernelSpec for the MUL_MAT operation. diff --git a/src/ggml-hsa/kernels/pool_2d.py b/src/ggml-hsa/kernels/pool_2d.py new file mode 100644 index 0000000000..e98e6b7468 --- /dev/null +++ b/src/ggml-hsa/kernels/pool_2d.py @@ -0,0 +1,46 @@ +# +# This file is licensed under the Apache License v2.0 with LLVM Exceptions. +# See https://llvm.org/LICENSE.txt for license information. +# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +# +# (c) Copyright 2026 Advanced Micro Devices, Inc. or its affiliates + +"""Top-level entry point for the GGML 2D pooling operation (GGML_OP_POOL_2D).""" + +from .kernel import Backend, KernelSpec + + +def ggml_op_pool_2d( + arch: str, input_tensors: list, output_tensor, op_params: bytearray +) -> KernelSpec: + """Return the KernelSpec for GGML_OP_POOL_2D. + + Parameters: + arch: Target architecture. + input_tensors: List of one input tensor. + output_tensor: Output tensor. + op_params: {op, k0, k1, s0, s1, p0, p1} as 7 x int32. + + Returns: + KernelSpec for the POOL_2D operation. + + """ + from functools import partial + + from .iron_kernels.pool_2d import pool_2d + + return KernelSpec( + backend=Backend.IRON, + op_name="GGML_OP_POOL_2D", + arch=arch, + input_tensors=input_tensors, + output_tensor=output_tensor, + op_params=op_params, + function=partial( + pool_2d, + arch=arch, + input_tensors=input_tensors, + output_tensor=output_tensor, + op_params=op_params, + ), + ) diff --git a/src/ggml-hsa/kernels/scale.py b/src/ggml-hsa/kernels/scale.py index cc34ea2228..f8d16afbb9 100644 --- a/src/ggml-hsa/kernels/scale.py +++ b/src/ggml-hsa/kernels/scale.py @@ -13,13 +13,13 @@ def ggml_op_scale( arch: str, input_tensors: list, output_tensor, op_params: bytearray ) -> KernelSpec: - """GGML_OP_SCALE implementation. + """Return the KernelSpec for GGML_OP_SCALE. Parameters: arch: Target architecture. input_tensors: List of one input tensor. output_tensor: Output tensor. - op_params: Operation parameters containing the scale factor. + op_params: Scale factor. Returns: KernelSpec for the SCALE operation. diff --git a/src/ggml-hsa/kernels/soft_max.py b/src/ggml-hsa/kernels/soft_max.py index 3b45ceee41..2a98126f78 100644 --- a/src/ggml-hsa/kernels/soft_max.py +++ b/src/ggml-hsa/kernels/soft_max.py @@ -13,16 +13,13 @@ def ggml_op_soft_max( arch: str, input_tensors: list, output_tensor, op_params: bytearray ) -> KernelSpec: - """GGML_OP_SOFT_MAX implementation. + """Return the KernelSpec for GGML_OP_SOFT_MAX. Parameters: arch: Target architecture. - input_tensors: List of 1-3 input tensors: - - input_tensors[0]: Input tensor (required) - - input_tensors[1]: Mask tensor (optional) - - input_tensors[2]: Sink tensor (optional) + input_tensors: [input, mask (optional), sink (optional)]. output_tensor: Output tensor. - op_params: Operation parameters (scale, max_bias). + op_params: scale, max_bias. Returns: KernelSpec for the SOFT_MAX operation. diff --git a/src/ggml-hsa/kernels/tensor_desc.py b/src/ggml-hsa/kernels/tensor_desc.py index 6597a43dc1..84aa27335f 100644 --- a/src/ggml-hsa/kernels/tensor_desc.py +++ b/src/ggml-hsa/kernels/tensor_desc.py @@ -2,12 +2,9 @@ """Tensor descriptor for GGML HSA kernel operations. -This module provides the TensorDesc dataclass used to describe tensors passed -to kernels. It captures the essential properties needed for kernel -compilation: data type, shape, stride, and contiguity information. - -The tensor dimensions follow GGML conventions where dimensions are ordered -from innermost to outermost (reverse of PyTorch). +TensorDesc captures the dtype, shape, stride, and contiguity a kernel needs to +compile. Dimensions follow GGML convention: innermost to outermost (reverse of +PyTorch). """ from dataclasses import dataclass @@ -28,16 +25,13 @@ def str_to_dtype(dtype_str: str): - """Converts a GGML dtype representation to its corresponding np.dtype object. - - Args: - dtype_str: The string representation of the data type. + """Convert a GGML dtype string to its corresponding numpy dtype. - Returns: - The corresponding np.dtype object. + Parameters: + dtype_str: GGML dtype string to convert (e.g. "f32", "bf16"). Raises: - ValueError: If the provided dtype_str is not recognized. + ValueError: If dtype_str is not recognized. """ try: @@ -83,21 +77,11 @@ def __post_init__(self) -> None: @property def size(self): - """Return the number of elements in the tensor. - - Returns: - int: The total number of elements in the tensor. - - """ + """Number of elements in the tensor.""" return int(np.prod(self.shape)) def numel(self): - """Return the number of elements in the tensor. - - Returns: - int: The total number of elements in the tensor. - - """ + """Number of elements in the tensor.""" return self.size @@ -107,18 +91,13 @@ def ggml_tensor_to_tensordesc( nb: tuple[int, int, int, int], contiguous: bool, ) -> TensorDesc: - """Create a TensorDesc from the ggml_tensor parameters. + """Create a TensorDesc from ggml_tensor parameters. Parameters: dtype: Tensor data type. - ne: Number of elements in each dimension. Dimensions - are from innermost to outermost (reverse of PyTorch). - nb: Tensor stride in bytes for each dimension. - Dimensions are from innermost to outermost (reverse of PyTorch). - contiguous: Indicates if the tensor is contiguous in memory. - - Returns: - TensorDesc: A new TensorDesc instance. + ne: Number of elements per dimension (innermost to outermost). + nb: Stride in bytes per dimension (innermost to outermost). + contiguous: Whether the tensor is contiguous in memory. """ return TensorDesc(dtype=dtype, shape=ne, stride=nb, contiguous=contiguous) diff --git a/src/ggml-hsa/kernels/triton_kernels/__init__.py b/src/ggml-hsa/kernels/triton_kernels/__init__.py index 73afbd21c2..9789f895f8 100644 --- a/src/ggml-hsa/kernels/triton_kernels/__init__.py +++ b/src/ggml-hsa/kernels/triton_kernels/__init__.py @@ -5,8 +5,4 @@ # # (c) Copyright 2026 Advanced Micro Devices, Inc. or its affiliates -"""Triton kernel implementations. - -This package contains Triton kernel implementations for various GGML operations. Each module -provides kernel designs that generate executables for AMD XDNA / XDNA2 NPUs. -""" +"""Triton kernel designs for GGML operations, targeting AMD XDNA / XDNA2 NPUs.""" diff --git a/src/ggml-hsa/kernels/triton_kernels/utils.py b/src/ggml-hsa/kernels/triton_kernels/utils.py index 4d4961155a..637845320a 100644 --- a/src/ggml-hsa/kernels/triton_kernels/utils.py +++ b/src/ggml-hsa/kernels/triton_kernels/utils.py @@ -24,28 +24,34 @@ def is_npu_arch(arch: str) -> bool: - """Return True if arch is a raw NPU architecture string (e.g. "aie2", "aie2p").""" + """Return True if arch is a raw NPU architecture string (e.g. "aie2", "aie2p"). + + Parameters: + arch: Architecture string to check. + + """ return arch in NPU_ARCH_MAP def is_gpu_arch(arch: str) -> bool: - """Return True if arch is an AMD GPU architecture string (e.g. "gfx942").""" + """Return True if arch is an AMD GPU architecture string (e.g. "gfx942"). + + Parameters: + arch: Architecture string to check. + + """ return arch.startswith("gfx") def triton_device(arch: str) -> str: - """Return the Triton device string for the given architecture. - - NPU architectures use "cpu"; GPU architectures use "cuda". + """Return the Triton device string: "cpu" for NPU targets, "cuda" for GPU targets. Parameters: - arch: Raw architecture string (e.g. "aie2", "aie2p", "gfx942"). - - Returns: - "cpu" for NPU targets, "cuda" for GPU targets. + arch: Architecture string to map to a Triton device. Raises: ValueError: If arch is not a recognised NPU or GPU architecture. + """ if is_npu_arch(arch): return "cpu" @@ -59,10 +65,7 @@ def numpy_dtype_to_torch(dtype: np.dtype) -> torch.dtype: """Convert a numpy dtype to the corresponding torch dtype. Parameters: - dtype: A numpy dtype instance. - - Returns: - The corresponding torch dtype. + dtype: The numpy dtype to convert. Raises: ValueError: If the numpy dtype has no torch equivalent. diff --git a/src/ggml-hsa/kernels/triton_kernels/vecadd.py b/src/ggml-hsa/kernels/triton_kernels/vecadd.py index 9e43aed4bb..4dfab2289f 100644 --- a/src/ggml-hsa/kernels/triton_kernels/vecadd.py +++ b/src/ggml-hsa/kernels/triton_kernels/vecadd.py @@ -15,14 +15,15 @@ def vecadd( n_elements: tl.constexpr, BLOCK_SIZE_N: tl.constexpr, ): - """Triton kernel for vector addition: C = A + B. + """Compute C = A + B over n_elements, BLOCK_SIZE_N elements per block. Parameters: - A: Pointer to input vector A - B: Pointer to input vector B - C: Pointer to output vector C - n_elements: Total number of elements in the vectors - BLOCK_SIZE_N: Number of elements processed by each block + A: Pointer to the first input vector. + B: Pointer to the second input vector. + C: Pointer to the output vector. + n_elements: Total number of elements in each vector. + BLOCK_SIZE_N: Number of elements processed per block. + """ pid = tl.program_id(0) # block row id block_start = pid * BLOCK_SIZE_N diff --git a/src/ggml-hsa/kernels/unary_ops.py b/src/ggml-hsa/kernels/unary_ops.py index 54a7ac6a43..19afaf35e4 100644 --- a/src/ggml-hsa/kernels/unary_ops.py +++ b/src/ggml-hsa/kernels/unary_ops.py @@ -16,16 +16,16 @@ def _make_iron_unary_kernel_spec( output_tensor, op_name: str, ) -> KernelSpec: - """Create a KernelSpec for a unary operation targeting the IRON backend. + """Create an IRON-backend KernelSpec for a unary operation. Parameters: arch: Target architecture. input_tensors: List of one input tensor. output_tensor: Output tensor. - op_name: Name of the unary operation. + op_name: Name of the GGML operation. Returns: - KernelSpec configured for IRON backend. + KernelSpec configured for the IRON backend. """ from functools import partial @@ -51,13 +51,14 @@ def _make_iron_unary_kernel_spec( def ggml_op_sqr( arch: str, input_tensors: list, output_tensor, op_params: bytearray ) -> KernelSpec: - """GGML_OP_SQR implementation. + """Return the KernelSpec for GGML_OP_SQR. Parameters: arch: Target architecture. input_tensors: List of one input tensor. output_tensor: Output tensor. - op_params: Operation parameters. + op_params: Operation parameters (unused for elementwise ops but required + by the dispatch interface). Returns: KernelSpec for the SQR operation. @@ -71,13 +72,14 @@ def ggml_op_sqr( def ggml_op_sqrt( arch: str, input_tensors: list, output_tensor, op_params: bytearray ) -> KernelSpec: - """GGML_OP_SQRT implementation. + """Return the KernelSpec for GGML_OP_SQRT. Parameters: arch: Target architecture. input_tensors: List of one input tensor. output_tensor: Output tensor. - op_params: Operation parameters. + op_params: Operation parameters (unused for elementwise ops but required + by the dispatch interface). Returns: KernelSpec for the SQRT operation. @@ -89,13 +91,14 @@ def ggml_op_sqrt( def ggml_op_log( arch: str, input_tensors: list, output_tensor, op_params: bytearray ) -> KernelSpec: - """GGML_OP_LOG implementation. + """Return the KernelSpec for GGML_OP_LOG. Parameters: arch: Target architecture. input_tensors: List of one input tensor. output_tensor: Output tensor. - op_params: Operation parameters. + op_params: Operation parameters (unused for elementwise ops but required + by the dispatch interface). Returns: KernelSpec for the LOG operation. @@ -109,13 +112,14 @@ def ggml_op_log( def ggml_op_sin( arch: str, input_tensors: list, output_tensor, op_params: bytearray ) -> KernelSpec: - """GGML_OP_SIN implementation. + """Return the KernelSpec for GGML_OP_SIN. Parameters: arch: Target architecture. input_tensors: List of one input tensor. output_tensor: Output tensor. - op_params: Operation parameters. + op_params: Operation parameters (unused for elementwise ops but required + by the dispatch interface). Returns: KernelSpec for the SIN operation. @@ -127,13 +131,14 @@ def ggml_op_sin( def ggml_op_cos( arch: str, input_tensors: list, output_tensor, op_params: bytearray ) -> KernelSpec: - """GGML_OP_COS implementation. + """Return the KernelSpec for GGML_OP_COS. Parameters: arch: Target architecture. input_tensors: List of one input tensor. output_tensor: Output tensor. - op_params: Operation parameters. + op_params: Operation parameters (unused for elementwise ops but required + by the dispatch interface). Returns: KernelSpec for the COS operation. @@ -145,13 +150,14 @@ def ggml_op_cos( def ggml_unary_op_abs( arch: str, input_tensors: list, output_tensor, op_params: bytearray ) -> KernelSpec: - """GGML_UNARY_OP_ABS implementation. + """Return the KernelSpec for GGML_UNARY_OP_ABS. Parameters: arch: Target architecture. input_tensors: List of one input tensor. output_tensor: Output tensor. - op_params: Operation parameters. + op_params: Operation parameters (unused for elementwise ops but required + by the dispatch interface). Returns: KernelSpec for the ABS operation. @@ -165,13 +171,14 @@ def ggml_unary_op_abs( def ggml_unary_op_sgn( arch: str, input_tensors: list, output_tensor, op_params: bytearray ) -> KernelSpec: - """GGML_UNARY_OP_SGN implementation. + """Return the KernelSpec for GGML_UNARY_OP_SGN. Parameters: arch: Target architecture. input_tensors: List of one input tensor. output_tensor: Output tensor. - op_params: Operation parameters. + op_params: Operation parameters (unused for elementwise ops but required + by the dispatch interface). Returns: KernelSpec for the SGN operation. @@ -185,13 +192,14 @@ def ggml_unary_op_sgn( def ggml_unary_op_neg( arch: str, input_tensors: list, output_tensor, op_params: bytearray ) -> KernelSpec: - """GGML_UNARY_OP_NEG implementation. + """Return the KernelSpec for GGML_UNARY_OP_NEG. Parameters: arch: Target architecture. input_tensors: List of one input tensor. output_tensor: Output tensor. - op_params: Operation parameters. + op_params: Operation parameters (unused for elementwise ops but required + by the dispatch interface). Returns: KernelSpec for the NEG operation. @@ -205,13 +213,14 @@ def ggml_unary_op_neg( def ggml_unary_op_step( arch: str, input_tensors: list, output_tensor, op_params: bytearray ) -> KernelSpec: - """GGML_UNARY_OP_STEP implementation. + """Return the KernelSpec for GGML_UNARY_OP_STEP. Parameters: arch: Target architecture. input_tensors: List of one input tensor. output_tensor: Output tensor. - op_params: Operation parameters. + op_params: Operation parameters (unused for elementwise ops but required + by the dispatch interface). Returns: KernelSpec for the STEP operation. @@ -225,13 +234,14 @@ def ggml_unary_op_step( def ggml_unary_op_tanh( arch: str, input_tensors: list, output_tensor, op_params: bytearray ) -> KernelSpec: - """GGML_UNARY_OP_TANH implementation. + """Return the KernelSpec for GGML_UNARY_OP_TANH. Parameters: arch: Target architecture. input_tensors: List of one input tensor. output_tensor: Output tensor. - op_params: Operation parameters. + op_params: Operation parameters (unused for elementwise ops but required + by the dispatch interface). Returns: KernelSpec for the TANH operation. @@ -243,13 +253,14 @@ def ggml_unary_op_tanh( def ggml_unary_op_elu( arch: str, input_tensors: list, output_tensor, op_params: bytearray ) -> KernelSpec: - """GGML_UNARY_OP_ELU implementation. + """Return the KernelSpec for GGML_UNARY_OP_ELU. Parameters: arch: Target architecture. input_tensors: List of one input tensor. output_tensor: Output tensor. - op_params: Operation parameters. + op_params: Operation parameters (unused for elementwise ops but required + by the dispatch interface). Returns: KernelSpec for the ELU operation. @@ -261,13 +272,14 @@ def ggml_unary_op_elu( def ggml_unary_op_relu( arch: str, input_tensors: list, output_tensor, op_params: bytearray ) -> KernelSpec: - """GGML_UNARY_OP_RELU implementation. + """Return the KernelSpec for GGML_UNARY_OP_RELU. Parameters: arch: Target architecture. input_tensors: List of one input tensor. output_tensor: Output tensor. - op_params: Operation parameters. + op_params: Operation parameters (unused for elementwise ops but required + by the dispatch interface). Returns: KernelSpec for the RELU operation. @@ -281,13 +293,14 @@ def ggml_unary_op_relu( def ggml_unary_op_sigmoid( arch: str, input_tensors: list, output_tensor, op_params: bytearray ) -> KernelSpec: - """GGML_UNARY_OP_SIGMOID implementation. + """Return the KernelSpec for GGML_UNARY_OP_SIGMOID. Parameters: arch: Target architecture. input_tensors: List of one input tensor. output_tensor: Output tensor. - op_params: Operation parameters. + op_params: Operation parameters (unused for elementwise ops but required + by the dispatch interface). Returns: KernelSpec for the SIGMOID operation. @@ -299,13 +312,14 @@ def ggml_unary_op_sigmoid( def ggml_unary_op_gelu( arch: str, input_tensors: list, output_tensor, op_params: bytearray ) -> KernelSpec: - """GGML_UNARY_OP_GELU implementation. + """Return the KernelSpec for GGML_UNARY_OP_GELU. Parameters: arch: Target architecture. input_tensors: List of one input tensor. output_tensor: Output tensor. - op_params: Operation parameters. + op_params: Operation parameters (unused for elementwise ops but required + by the dispatch interface). Returns: KernelSpec for the GELU operation. @@ -317,13 +331,14 @@ def ggml_unary_op_gelu( def ggml_unary_op_gelu_quick( arch: str, input_tensors: list, output_tensor, op_params: bytearray ) -> KernelSpec: - """GGML_UNARY_OP_GELU_QUICK implementation. + """Return the KernelSpec for GGML_UNARY_OP_GELU_QUICK. Parameters: arch: Target architecture. input_tensors: List of one input tensor. output_tensor: Output tensor. - op_params: Operation parameters. + op_params: Operation parameters (unused for elementwise ops but required + by the dispatch interface). Returns: KernelSpec for the GELU_QUICK operation. @@ -335,13 +350,14 @@ def ggml_unary_op_gelu_quick( def ggml_unary_op_silu( arch: str, input_tensors: list, output_tensor, op_params: bytearray ) -> KernelSpec: - """GGML_UNARY_OP_SILU implementation. + """Return the KernelSpec for GGML_UNARY_OP_SILU. Parameters: arch: Target architecture. input_tensors: List of one input tensor. output_tensor: Output tensor. - op_params: Operation parameters. + op_params: Operation parameters (unused for elementwise ops but required + by the dispatch interface). Returns: KernelSpec for the SILU operation. @@ -353,13 +369,14 @@ def ggml_unary_op_silu( def ggml_unary_op_hardswish( arch: str, input_tensors: list, output_tensor, op_params: bytearray ) -> KernelSpec: - """GGML_UNARY_OP_HARDSWISH implementation. + """Return the KernelSpec for GGML_UNARY_OP_HARDSWISH. Parameters: arch: Target architecture. input_tensors: List of one input tensor. output_tensor: Output tensor. - op_params: Operation parameters. + op_params: Operation parameters (unused for elementwise ops but required + by the dispatch interface). Returns: KernelSpec for the HARDSWISH operation. @@ -373,13 +390,14 @@ def ggml_unary_op_hardswish( def ggml_unary_op_hardsigmoid( arch: str, input_tensors: list, output_tensor, op_params: bytearray ) -> KernelSpec: - """GGML_UNARY_OP_HARDSIGMOID implementation. + """Return the KernelSpec for GGML_UNARY_OP_HARDSIGMOID. Parameters: arch: Target architecture. input_tensors: List of one input tensor. output_tensor: Output tensor. - op_params: Operation parameters. + op_params: Operation parameters (unused for elementwise ops but required + by the dispatch interface). Returns: KernelSpec for the HARDSIGMOID operation. @@ -393,13 +411,14 @@ def ggml_unary_op_hardsigmoid( def ggml_unary_op_exp( arch: str, input_tensors: list, output_tensor, op_params: bytearray ) -> KernelSpec: - """GGML_UNARY_OP_EXP implementation. + """Return the KernelSpec for GGML_UNARY_OP_EXP. Parameters: arch: Target architecture. input_tensors: List of one input tensor. - output_tensor : Output tensor. - op_params: Operation parameters. + output_tensor: Output tensor. + op_params: Operation parameters (unused for elementwise ops but required + by the dispatch interface). Returns: KernelSpec for the EXP operation. @@ -411,13 +430,14 @@ def ggml_unary_op_exp( def ggml_unary_op_gelu_erf( arch: str, input_tensors: list, output_tensor, op_params: bytearray ) -> KernelSpec: - """GGML_UNARY_OP_GELU_ERF implementation. + """Return the KernelSpec for GGML_UNARY_OP_GELU_ERF. Parameters: arch: Target architecture. input_tensors: List of one input tensor. output_tensor: Output tensor. - op_params: Operation parameters. + op_params: Operation parameters (unused for elementwise ops but required + by the dispatch interface). Returns: KernelSpec for the GELU_ERF operation. @@ -429,13 +449,14 @@ def ggml_unary_op_gelu_erf( def ggml_unary_op_xielu( arch: str, input_tensors: list, output_tensor, op_params: bytearray ) -> KernelSpec: - """GGML_UNARY_OP_XIELU implementation. + """Return the KernelSpec for GGML_UNARY_OP_XIELU. Parameters: arch: Target architecture. input_tensors: List of one input tensor. output_tensor: Output tensor. - op_params: Operation parameters. + op_params: Operation parameters (unused for elementwise ops but required + by the dispatch interface). Returns: KernelSpec for the XIELU operation. @@ -447,13 +468,14 @@ def ggml_unary_op_xielu( def ggml_unary_op_floor( arch: str, input_tensors: list, output_tensor, op_params: bytearray ) -> KernelSpec: - """GGML_UNARY_OP_FLOOR implementation. + """Return the KernelSpec for GGML_UNARY_OP_FLOOR. Parameters: arch: Target architecture. input_tensors: List of one input tensor. output_tensor: Output tensor. - op_params: Operation parameters. + op_params: Operation parameters (unused for elementwise ops but required + by the dispatch interface). Returns: KernelSpec for the FLOOR operation. @@ -467,13 +489,14 @@ def ggml_unary_op_floor( def ggml_unary_op_ceil( arch: str, input_tensors: list, output_tensor, op_params: bytearray ) -> KernelSpec: - """GGML_UNARY_OP_CEIL implementation. + """Return the KernelSpec for GGML_UNARY_OP_CEIL. Parameters: arch: Target architecture. input_tensors: List of one input tensor. output_tensor: Output tensor. - op_params: Operation parameters. + op_params: Operation parameters (unused for elementwise ops but required + by the dispatch interface). Returns: KernelSpec for the CEIL operation. @@ -487,13 +510,14 @@ def ggml_unary_op_ceil( def ggml_unary_op_round( arch: str, input_tensors: list, output_tensor, op_params: bytearray ) -> KernelSpec: - """GGML_UNARY_OP_ROUND implementation. + """Return the KernelSpec for GGML_UNARY_OP_ROUND. Parameters: arch: Target architecture. input_tensors: List of one input tensor. output_tensor: Output tensor. - op_params: Operation parameters. + op_params: Operation parameters (unused for elementwise ops but required + by the dispatch interface). Returns: KernelSpec for the ROUND operation. @@ -507,13 +531,14 @@ def ggml_unary_op_round( def ggml_unary_op_trunc( arch: str, input_tensors: list, output_tensor, op_params: bytearray ) -> KernelSpec: - """GGML_UNARY_OP_TRUNC implementation. + """Return the KernelSpec for GGML_UNARY_OP_TRUNC. - Args: + Parameters: arch: Target architecture. input_tensors: List of one input tensor. output_tensor: Output tensor. - op_params: Operation parameters. + op_params: Operation parameters (unused for elementwise ops but required + by the dispatch interface). Returns: KernelSpec for the TRUNC operation. diff --git a/tests/ggml-hsa/test-backend-ops-mnist.cpp b/tests/ggml-hsa/test-backend-ops-mnist.cpp index a4af3a4750..2536d55e9a 100644 --- a/tests/ggml-hsa/test-backend-ops-mnist.cpp +++ b/tests/ggml-hsa/test-backend-ops-mnist.cpp @@ -1,30 +1,27 @@ // This file defines tests for various GGML ops and backends. -// For the forward pass it asserts that the results of multiple backends computing the same GGML ops are consistent. -// For the backward pass it asserts that the gradients from backpropagation are consistent -// with the gradients obtained via the method of finite differences ("grad" mode, this is optional). -// It is also possible to check the performance ("perf" mode). +// For the forward pass it asserts that the results of multiple backends computing the same GGML ops +// are consistent. For the backward pass it asserts that the gradients from backpropagation are +// consistent with the gradients obtained via the method of finite differences ("grad" mode, this is +// optional). It is also possible to check the performance ("perf" mode). // -// this file has three sections: Section 1 does general setup, section 2 defines the GGML ops to be tested, -// and section 3 defines which tests to run. -// Quick start for adding a new GGML op: Go to section 2 and create a struct that inherits from test_case, -// then go to section 3 and add an instantiation of your struct. - +// this file has three sections: Section 1 does general setup, section 2 defines the GGML ops to be +// tested, and section 3 defines which tests to run. Quick start for adding a new GGML op: Go to +// section 2 and create a struct that inherits from test_case, then go to section 3 and add an +// instantiation of your struct. // ############################## // ## Section 1: General Setup ## // ############################## - -#include #include #include #include +#include #include #include #include #include -#include #include #include #include @@ -37,14 +34,13 @@ #include #include #include -#include -#include #include +#include #ifdef __EMSCRIPTEN__ -# define N_THREADS 1 +#define N_THREADS 1 #else -# define N_THREADS std::thread::hardware_concurrency() +#define N_THREADS std::thread::hardware_concurrency() #endif static void init_tensor_uniform(ggml_tensor * tensor, float min = -1.0f, float max = 1.0f) { @@ -58,8 +54,10 @@ static void init_tensor_uniform(ggml_tensor * tensor, float min = -1.0f, float m std::random_device rd; std::vector vec; vec.reserve(n_threads); - //for (size_t i = 0; i < n_threads; i++) { vec.emplace_back(1234 + i); } // fixed seed - for (size_t i = 0; i < n_threads; i++) { vec.emplace_back(rd()); } + // for (size_t i = 0; i < n_threads; i++) { vec.emplace_back(1234 + i); } // fixed seed + for (size_t i = 0; i < n_threads; i++) { + vec.emplace_back(rd()); + } return vec; }(); @@ -77,8 +75,8 @@ static void init_tensor_uniform(ggml_tensor * tensor, float min = -1.0f, float m std::vector> tasks; tasks.reserve(n_threads); for (size_t i = 0; i < n_threads; i++) { - size_t start = i*nels/n_threads; - size_t end = (i+1)*nels/n_threads; + size_t start = i * nels / n_threads; + size_t end = (i + 1) * nels / n_threads; tasks.push_back(std::async(std::launch::async, init_thread, i, start, end)); } for (auto & t : tasks) { @@ -89,16 +87,17 @@ static void init_tensor_uniform(ggml_tensor * tensor, float min = -1.0f, float m if (tensor->type == GGML_TYPE_F32 || tensor->type == GGML_TYPE_I32) { ggml_backend_tensor_set(tensor, data.data(), 0, nels * sizeof(float)); - } else if (ggml_is_quantized(tensor->type) || tensor->type == GGML_TYPE_F16 || tensor->type == GGML_TYPE_BF16) { + } else if (ggml_is_quantized(tensor->type) || tensor->type == GGML_TYPE_F16 || + tensor->type == GGML_TYPE_BF16) { GGML_ASSERT(nels % ggml_blck_size(tensor->type) == 0); - // dummy importance matrix + // dummy importance matrix std::vector imatrix(tensor->ne[0], 1.0f); const float * im = imatrix.data(); if (!ggml_quantize_requires_imatrix(tensor->type)) { - // when the imatrix is optional, we want to test both quantization with and without imatrix - // use one of the random numbers to decide - if (data[0] > 0.5f*(min + max)) { + // when the imatrix is optional, we want to test both quantization with and without + // imatrix use one of the random numbers to decide + if (data[0] > 0.5f * (min + max)) { im = nullptr; } } @@ -110,13 +109,14 @@ static void init_tensor_uniform(ggml_tensor * tensor, float min = -1.0f, float m size_t n_blocks = nels / blck_size; auto quantize_thread = [&](size_t start, size_t end) { - ggml_quantize_chunk(tensor->type, data.data(), dataq.data(), - start * blck_size, end - start, blck_size, im); + ggml_quantize_chunk(tensor->type, data.data(), dataq.data(), start * blck_size, + end - start, blck_size, im); }; const size_t min_blocks_per_thread = 1; - const size_t n_quant_threads = std::min(std::max(N_THREADS/2, 1), - std::max(1, n_blocks / min_blocks_per_thread)); + const size_t n_quant_threads = + std::min(std::max(N_THREADS / 2, 1), + std::max(1, n_blocks / min_blocks_per_thread)); if (n_quant_threads == 1) { // single-threaded quantization: do all blocks in the current thread @@ -125,8 +125,8 @@ static void init_tensor_uniform(ggml_tensor * tensor, float min = -1.0f, float m std::vector> tasks; tasks.reserve(n_quant_threads); for (size_t i = 0; i < n_quant_threads; i++) { - size_t start = i*n_blocks/n_quant_threads; - size_t end = (i+1)*n_blocks/n_quant_threads; + size_t start = i * n_blocks / n_quant_threads; + size_t end = (i + 1) * n_blocks / n_quant_threads; tasks.push_back(std::async(std::launch::async, quantize_thread, start, end)); } for (auto & t : tasks) { @@ -135,97 +135,21 @@ static void init_tensor_uniform(ggml_tensor * tensor, float min = -1.0f, float m } } ggml_backend_tensor_set(tensor, dataq.data(), 0, dataq.size()); - } else if (tensor->type == GGML_TYPE_I8 || tensor->type == GGML_TYPE_I16 || tensor->type == GGML_TYPE_I32) { + } else if (tensor->type == GGML_TYPE_I8 || tensor->type == GGML_TYPE_I16 || + tensor->type == GGML_TYPE_I32) { // This is going to create some weird integers though. ggml_backend_tensor_set(tensor, data.data(), 0, ggml_nbytes(tensor)); } else if (tensor->type == GGML_TYPE_I64) { - // Integers with a size of 8 bytes can be set by mirroring the float data, the specific values are again not really meaningful. - const size_t nbytes_half = ggml_nbytes(tensor)/2; - ggml_backend_tensor_set(tensor, data.data(), 0*nbytes_half, nbytes_half); - ggml_backend_tensor_set(tensor, data.data(), 1*nbytes_half, nbytes_half); + // Integers with a size of 8 bytes can be set by mirroring the float data, the specific + // values are again not really meaningful. + const size_t nbytes_half = ggml_nbytes(tensor) / 2; + ggml_backend_tensor_set(tensor, data.data(), 0 * nbytes_half, nbytes_half); + ggml_backend_tensor_set(tensor, data.data(), 1 * nbytes_half, nbytes_half); } else { GGML_ABORT("fatal error"); } } -// generate an F16 mask where certain blocks are randomly masked with -INF value -static void init_tensor_kq_mask(ggml_tensor * tensor, float min = -1.0f, float max = 1.0f) { - GGML_ASSERT(tensor->type == GGML_TYPE_F16); - - GGML_TENSOR_LOCALS( int32_t, ne, tensor, ne); - - std::vector data_f32(ne0*ne1*ne2*ne3); - std::vector data_f16(ne0*ne1*ne2*ne3); - - std::random_device rd; - std::mt19937 gen(rd()); - std::uniform_real_distribution dis(min, max); - - for (size_t i = 0; i < data_f32.size(); i++) { - data_f32[i] = dis(gen); - } - - // block size - const int blck0 = 128; - const int blck1 = 64; - - // number of INF/zero blocks - const int n_inf_zero_blocks = 0.2*(ne0*ne1*ne2*ne3)/(blck0*blck1); - - for (int b = 0; b < n_inf_zero_blocks; b++) { - const int p3 = (rd() % ne3); - const int p2 = (rd() % ne2); - const int p1 = (rd() % ne1); - const int p0 = (rd() % ne0); - - bool inf = rd() & 1; - - for (int i1 = 0; i1 < blck1 && p1 + i1 < ne1; i1++) { - const int idx = p3*ne2*ne1*ne0 + p2*ne1*ne0 + (p1 + i1)*ne0 + p0; - - for (int i0 = 0; i0 < blck0 && p0 + i0 < ne0; i0++) { - data_f32[idx + i0] = inf ? -INFINITY : 0.0f; - } - } - } - - ggml_fp32_to_fp16_row(data_f32.data(), data_f16.data(), ne0*ne1*ne2*ne3); - - ggml_backend_tensor_set(tensor, data_f16.data(), 0, data_f16.size()*sizeof(ggml_fp16_t)); -} - -// generate a lower triangular matrix -static void init_tensor_tril(ggml_tensor * tensor, float min = -1.0f, float max = 1.0f) { - GGML_ASSERT(tensor->type == GGML_TYPE_F32); - GGML_ASSERT(tensor->ne[0] == tensor->ne[1]); - - GGML_TENSOR_LOCALS(int32_t, ne, tensor, ne); - GGML_TENSOR_LOCALS(size_t, nb, tensor, nb); - - std::vector data_f32(ne0*ne1*ne2*ne3); - - std::random_device rd; - std::mt19937 gen(rd()); - std::uniform_real_distribution dis(min, max); - - for (int64_t i3 = 0; i3 < ne3; i3++) { - for (int64_t i2 = 0; i2 < ne2; i2++) { - for (int64_t i1 = 0; i1 < ne1; i1++) { - for (int64_t i0 = 0; i0 < ne0; i0++) { - int64_t idx = (i0 * nb0 + i1 * nb1 + i2 * nb2 + i3 * nb3) / sizeof(float); - if (i0 <= i1) { - data_f32[idx] = dis(gen); - } else { - data_f32[idx] = 0.0f; - } - } - } - } - } - - ggml_backend_tensor_set(tensor, data_f32.data(), 0, ggml_nbytes(tensor)); -} - static std::vector tensor_to_float(const ggml_tensor * t) { std::vector tv; tv.reserve(ggml_nelements(t)); @@ -243,21 +167,21 @@ static std::vector tensor_to_float(const ggml_tensor * t) { for (int64_t i2 = 0; i2 < t->ne[2]; i2++) { for (int64_t i1 = 0; i1 < t->ne[1]; i1++) { for (int64_t i0 = 0; i0 < t->ne[0]; i0 += bs) { - size_t i = i3*t->nb[3] + i2*t->nb[2] + i1*t->nb[1] + i0/bs*t->nb[0]; + size_t i = i3 * t->nb[3] + i2 * t->nb[2] + i1 * t->nb[1] + i0 / bs * t->nb[0]; if (t->type == GGML_TYPE_F16) { - tv.push_back(ggml_fp16_to_fp32(*(ggml_fp16_t*)&buf[i])); + tv.push_back(ggml_fp16_to_fp32(*(ggml_fp16_t *)&buf[i])); } else if (t->type == GGML_TYPE_BF16) { - tv.push_back(ggml_bf16_to_fp32(*(ggml_bf16_t*)&buf[i])); + tv.push_back(ggml_bf16_to_fp32(*(ggml_bf16_t *)&buf[i])); } else if (t->type == GGML_TYPE_F32) { - tv.push_back(*(float *) &buf[i]); + tv.push_back(*(float *)&buf[i]); } else if (t->type == GGML_TYPE_I64) { - tv.push_back((float)*(int64_t *) &buf[i]); + tv.push_back((float)*(int64_t *)&buf[i]); } else if (t->type == GGML_TYPE_I32) { - tv.push_back((float)*(int32_t *) &buf[i]); + tv.push_back((float)*(int32_t *)&buf[i]); } else if (t->type == GGML_TYPE_I16) { - tv.push_back((float)*(int16_t *) &buf[i]); + tv.push_back((float)*(int16_t *)&buf[i]); } else if (t->type == GGML_TYPE_I8) { - tv.push_back((float)*(int8_t *) &buf[i]); + tv.push_back((float)*(int8_t *)&buf[i]); } else if (quantized) { tt->to_float(&buf[i], vq.data(), bs); tv.insert(tv.end(), vq.begin(), vq.end()); @@ -288,42 +212,18 @@ static double nmse(const float * a, const float * b, size_t n) { return mse_a_b / mse_a_0; } -// difference between 2 sets (Jaccard distance, 0 - no difference, 1 - no overlap) -template -static double jdst(const T * a, const T * b, size_t n) { - std::unordered_map set_a; - std::unordered_map set_b; - - for (size_t i = 0; i < n; ++i) { - set_a[a[i]]++; - set_b[b[i]]++; - } - - size_t diff = 0; - - for (const auto & p : set_a) { - const int64_t na = p.second; - const int64_t nb = set_b.find(p.first) != set_b.end() ? set_b.at(p.first) : 0; - - diff += std::abs(na - nb); - } - - for (const auto & p : set_b) { - if (set_a.find(p.first) == set_a.end()) { - diff += p.second; - } - } - - return (double) diff / (2*n); -} - // maximum absolute asymmetry between a and b // asymmetry: (a - b) / (a + b) // This is more stable than relative error if one of the values fluctuates towards zero. // n: number of values to compare. -// expected_vals: optional vector of expected values for a. If expected_vals is not empty, filter out all comparisons where -// a does not match any of the expected values. Needed for noncontinuous gradients where the numerical calculation can fail. -static double mean_abs_asymm(const float * a, const float * b, const size_t n, const std::vector & expected_vals) { +// expected_vals: optional vector of expected values for a. If expected_vals is not empty, filter +// out all comparisons where +// a does not match any of the expected values. Needed for noncontinuous gradients where the +// numerical calculation can fail. +static double mean_abs_asymm(const float * a, + const float * b, + const size_t n, + const std::vector & expected_vals) { double sum = 0.0f; size_t nvalid = 0; @@ -347,21 +247,19 @@ static double mean_abs_asymm(const float * a, const float * b, const size_t n, c nvalid++; } - return sum/nvalid; + return sum / nvalid; } // utils for printing the variables of the test cases -static std::string var_to_str(const std::string & x) { - return x; -} +static std::string var_to_str(const std::string & x) { return x; } -template +template static std::string var_to_str(const T & x) { return std::to_string(x); } -template +template static std::string var_to_str(const T (&x)[N]) { std::string s = "["; for (size_t i = 0; i < N; i++) { @@ -374,7 +272,7 @@ static std::string var_to_str(const T (&x)[N]) { return s; } -template +template static std::string var_to_str(const std::array & x) { std::string s = "["; for (size_t i = 0; i < N; i++) { @@ -387,37 +285,19 @@ static std::string var_to_str(const std::array & x) { return s; } -static std::string var_to_str(ggml_type type) { - return ggml_type_name(type); -} +static std::string var_to_str(ggml_type type) { return ggml_type_name(type); } -static std::string var_to_str(ggml_prec prec) { - return prec == GGML_PREC_F32 ? "f32" : "def"; -} +static std::string var_to_str(ggml_prec prec) { return prec == GGML_PREC_F32 ? "f32" : "def"; } static std::string var_to_str(ggml_op_pool pool) { switch (pool) { - case GGML_OP_POOL_AVG: return "avg"; - case GGML_OP_POOL_MAX: return "max"; - default: return std::to_string(pool); - } -} - -static std::string var_to_str(ggml_scale_mode mode) { - std::string str; - switch (mode & 0xFF) { - case GGML_SCALE_MODE_NEAREST: str = "nearest"; break; - case GGML_SCALE_MODE_BILINEAR: str = "bilinear"; break; - case GGML_SCALE_MODE_BICUBIC: str = "bicubic"; break; - default: str = std::to_string(mode); break; - } - if (mode & GGML_SCALE_FLAG_ALIGN_CORNERS) { - str += "|align_corners"; - } - if (mode & GGML_SCALE_FLAG_ANTIALIAS) { - str += "|antialias"; + case GGML_OP_POOL_AVG: + return "avg"; + case GGML_OP_POOL_MAX: + return "max"; + default: + return std::to_string(pool); } - return str; } #define VAR_TO_STR(x) (#x "=" + var_to_str(x)) @@ -430,37 +310,43 @@ static std::string var_to_str(ggml_scale_mode mode) { #define VARS_TO_STR6(a, b, c, d, e, f) VAR_TO_STR(a) + "," + VARS_TO_STR5(b, c, d, e, f) #define VARS_TO_STR7(a, b, c, d, e, f, g) VAR_TO_STR(a) + "," + VARS_TO_STR6(b, c, d, e, f, g) #define VARS_TO_STR8(a, b, c, d, e, f, g, h) VAR_TO_STR(a) + "," + VARS_TO_STR7(b, c, d, e, f, g, h) -#define VARS_TO_STR9(a, b, c, d, e, f, g, h, i) VAR_TO_STR(a) + "," + VARS_TO_STR8(b, c, d, e, f, g, h, i) -#define VARS_TO_STR10(a, b, c, d, e, f, g, h, i, j) VAR_TO_STR(a) + "," + VARS_TO_STR9(b, c, d, e, f, g, h, i, j) -#define VARS_TO_STR11(a, b, c, d, e, f, g, h, i, j, k) VAR_TO_STR(a) + "," + VARS_TO_STR10(b, c, d, e, f, g, h, i, j, k) -#define VARS_TO_STR12(a, b, c, d, e, f, g, h, i, j, k, l) VAR_TO_STR(a) + "," + VARS_TO_STR11(b, c, d, e, f, g, h, i, j, k, l) -#define VARS_TO_STR13(a, b, c, d, e, f, g, h, i, j, k, l, m) VAR_TO_STR(a) + "," + VARS_TO_STR12(b, c, d, e, f, g, h, i, j, k, l, m) -#define VARS_TO_STR14(a, b, c, d, e, f, g, h, i, j, k, l, m, n) VAR_TO_STR(a) + "," + VARS_TO_STR13(b, c, d, e, f, g, h, i, j, k, l, m, n) -#define VARS_TO_STR15(a, b, c, d, e, f, g, h, i, j, k, l, m, n, o) VAR_TO_STR(a) + "," + VARS_TO_STR14(b, c, d, e, f, g, h, i, j, k, l, m, n, o) -#define VARS_TO_STR16(a, b, c, d, e, f, g, h, i, j, k, l, m, n, o, p) VAR_TO_STR(a) + "," + VARS_TO_STR15(b, c, d, e, f, g, h, i, j, k, l, m, n, o, p) +#define VARS_TO_STR9(a, b, c, d, e, f, g, h, i) \ + VAR_TO_STR(a) + "," + VARS_TO_STR8(b, c, d, e, f, g, h, i) +#define VARS_TO_STR10(a, b, c, d, e, f, g, h, i, j) \ + VAR_TO_STR(a) + "," + VARS_TO_STR9(b, c, d, e, f, g, h, i, j) +#define VARS_TO_STR11(a, b, c, d, e, f, g, h, i, j, k) \ + VAR_TO_STR(a) + "," + VARS_TO_STR10(b, c, d, e, f, g, h, i, j, k) +#define VARS_TO_STR12(a, b, c, d, e, f, g, h, i, j, k, l) \ + VAR_TO_STR(a) + "," + VARS_TO_STR11(b, c, d, e, f, g, h, i, j, k, l) +#define VARS_TO_STR13(a, b, c, d, e, f, g, h, i, j, k, l, m) \ + VAR_TO_STR(a) + "," + VARS_TO_STR12(b, c, d, e, f, g, h, i, j, k, l, m) +#define VARS_TO_STR14(a, b, c, d, e, f, g, h, i, j, k, l, m, n) \ + VAR_TO_STR(a) + "," + VARS_TO_STR13(b, c, d, e, f, g, h, i, j, k, l, m, n) +#define VARS_TO_STR15(a, b, c, d, e, f, g, h, i, j, k, l, m, n, o) \ + VAR_TO_STR(a) + "," + VARS_TO_STR14(b, c, d, e, f, g, h, i, j, k, l, m, n, o) +#define VARS_TO_STR16(a, b, c, d, e, f, g, h, i, j, k, l, m, n, o, p) \ + VAR_TO_STR(a) + "," + VARS_TO_STR15(b, c, d, e, f, g, h, i, j, k, l, m, n, o, p) #ifdef GGML_USE_SYCL -static bool inline _isinf(float f) { - return (*(uint32_t *)&f & 0x7fffffff) == 0x7f800000; -} +static bool inline _isinf(float f) { return (*(uint32_t *)&f & 0x7fffffff) == 0x7f800000; } #else static bool inline _isinf(float f) { return std::isinf(f); } #endif // accept FLT_MAX as infinity -static bool isinf_or_max(float f) { - return _isinf(f) || f == FLT_MAX || f == -FLT_MAX; -} +static bool isinf_or_max(float f) { return _isinf(f) || f == FLT_MAX || f == -FLT_MAX; } static bool ggml_is_view_op(enum ggml_op op) { - return op == GGML_OP_VIEW || op == GGML_OP_RESHAPE || op == GGML_OP_PERMUTE || op == GGML_OP_TRANSPOSE; + return op == GGML_OP_VIEW || op == GGML_OP_RESHAPE || op == GGML_OP_PERMUTE || + op == GGML_OP_TRANSPOSE; } static bool backend_has_feature(ggml_backend_t backend, const char * feature_name) { ggml_backend_dev_t dev = ggml_backend_get_device(backend); ggml_backend_reg_t reg = ggml_backend_dev_backend_reg(dev); - auto get_features = (ggml_backend_get_features_t) ggml_backend_reg_get_proc_address(reg, "ggml_backend_get_features"); + auto get_features = (ggml_backend_get_features_t)ggml_backend_reg_get_proc_address( + reg, "ggml_backend_get_features"); if (!get_features) { return false; } @@ -522,30 +408,30 @@ struct test_result { std::string op_name; std::string op_params; std::string test_mode; - bool supported; - bool passed; + bool supported; + bool passed; std::string error_message; - double time_us; - double flops; - double bandwidth_gb_s; - size_t memory_kb; - int n_runs; + double time_us; + double flops; + double bandwidth_gb_s; + size_t memory_kb; + int n_runs; std::string device_description; std::string backend_reg_name; test_result() { // Initialize with default values - time_us = 0.0; - flops = 0.0; + time_us = 0.0; + flops = 0.0; bandwidth_gb_s = 0.0; - memory_kb = 0; - n_runs = 0; - supported = false; - passed = false; + memory_kb = 0; + n_runs = 0; + supported = false; + passed = false; // Set test time time_t t = time(NULL); - char buf[32]; + char buf[32]; std::strftime(buf, sizeof(buf), "%FT%TZ", gmtime(&t)); test_time = buf; @@ -553,10 +439,20 @@ struct test_result { build_commit = ggml_commit(); } - test_result(const std::string & backend_name, const std::string & op_name, const std::string & op_params, - const std::string & test_mode, bool supported, bool passed, const std::string & error_message = "", - double time_us = 0.0, double flops = 0.0, double bandwidth_gb_s = 0.0, size_t memory_kb = 0, - int n_runs = 0, const std::string & device_description = "", const std::string & backend_reg_name = "") : + test_result(const std::string & backend_name, + const std::string & op_name, + const std::string & op_params, + const std::string & test_mode, + bool supported, + bool passed, + const std::string & error_message = "", + double time_us = 0.0, + double flops = 0.0, + double bandwidth_gb_s = 0.0, + size_t memory_kb = 0, + int n_runs = 0, + const std::string & device_description = "", + const std::string & backend_reg_name = "") : backend_name(backend_name), op_name(op_name), op_params(op_params), @@ -573,7 +469,7 @@ struct test_result { backend_reg_name(backend_reg_name) { // Set test time time_t t = time(NULL); - char buf[32]; + char buf[32]; std::strftime(buf, sizeof(buf), "%FT%TZ", gmtime(&t)); test_time = buf; @@ -583,10 +479,12 @@ struct test_result { static const std::vector & get_fields() { static const std::vector fields = { - "test_time", "build_commit", "backend_name", "op_name", "op_params", "test_mode", "supported", - "passed", "error_message", "time_us", "flops", "bandwidth_gb_s", "memory_kb", "n_runs", - "device_description", "backend_reg_name" - }; + "test_time", "build_commit", "backend_name", + "op_name", "op_params", "test_mode", + "supported", "passed", "error_message", + "time_us", "flops", "bandwidth_gb_s", + "memory_kb", "n_runs", "device_description", + "backend_reg_name"}; return fields; } @@ -606,22 +504,22 @@ struct test_result { } std::vector get_values() const { - return { test_time, - build_commit, - backend_name, - op_name, - op_params, - test_mode, - std::to_string(supported), - std::to_string(passed), - error_message, - std::to_string(time_us), - std::to_string(flops), - std::to_string(bandwidth_gb_s), - std::to_string(memory_kb), - std::to_string(n_runs), - device_description, - backend_reg_name }; + return {test_time, + build_commit, + backend_name, + op_name, + op_params, + test_mode, + std::to_string(supported), + std::to_string(passed), + error_message, + std::to_string(time_us), + std::to_string(flops), + std::to_string(bandwidth_gb_s), + std::to_string(memory_kb), + std::to_string(n_runs), + device_description, + backend_reg_name}; } }; @@ -629,36 +527,39 @@ struct test_result { enum class test_status_t { NOT_SUPPORTED, OK, FAIL, SKIPPED }; struct test_operation_info { - std::string op_name; - std::string op_params; - std::string backend_name; + std::string op_name; + std::string op_params; + std::string backend_name; test_status_t status = test_status_t::OK; - std::string failure_reason; + std::string failure_reason; // Additional information fields that were previously in separate structs std::string error_component; std::string error_details; // Gradient info - int64_t gradient_index = -1; + int64_t gradient_index = -1; std::string gradient_param_name; - float gradient_value = 0.0f; + float gradient_value = 0.0f; // MAA error info - double maa_error = 0.0; + double maa_error = 0.0; double maa_threshold = 0.0; // Flags for different types of information - bool has_error = false; - bool has_gradient_info = false; - bool has_maa_error = false; - bool is_compare_failure = false; + bool has_error = false; + bool has_gradient_info = false; + bool has_maa_error = false; + bool is_compare_failure = false; bool is_large_tensor_skip = false; test_operation_info() = default; - test_operation_info(const std::string & op_name, const std::string & op_params, const std::string & backend_name, - test_status_t status = test_status_t::OK, const std::string & failure_reason = "") : + test_operation_info(const std::string & op_name, + const std::string & op_params, + const std::string & backend_name, + test_status_t status = test_status_t::OK, + const std::string & failure_reason = "") : op_name(op_name), op_params(op_params), backend_name(backend_name), @@ -667,9 +568,9 @@ struct test_operation_info { // Set error information void set_error(const std::string & component, const std::string & details) { - has_error = true; + has_error = true; error_component = component; - error_details = details; + error_details = details; if (status == test_status_t::OK) { status = test_status_t::FAIL; } @@ -677,10 +578,10 @@ struct test_operation_info { // Set gradient information void set_gradient_info(int64_t index, const std::string & param_name, float value) { - has_gradient_info = true; - gradient_index = index; + has_gradient_info = true; + gradient_index = index; gradient_param_name = param_name; - gradient_value = value; + gradient_value = value; if (status == test_status_t::OK) { status = test_status_t::FAIL; } @@ -689,7 +590,7 @@ struct test_operation_info { // Set MAA error information void set_maa_error(double error, double threshold) { has_maa_error = true; - maa_error = error; + maa_error = error; maa_threshold = threshold; if (status == test_status_t::OK) { status = test_status_t::FAIL; @@ -711,7 +612,7 @@ struct test_operation_info { struct test_summary_info { size_t tests_passed; size_t tests_total; - bool is_backend_summary = false; // true for backend summary, false for test summary + bool is_backend_summary = false; // true for backend summary, false for test summary test_summary_info() = default; @@ -730,21 +631,27 @@ struct testing_start_info { }; struct backend_init_info { - size_t device_index; - size_t total_devices; + size_t device_index; + size_t total_devices; std::string device_name; - bool skipped = false; + bool skipped = false; std::string skip_reason; std::string description; - size_t memory_total_mb = 0; - size_t memory_free_mb = 0; - bool has_memory_info = false; + size_t memory_total_mb = 0; + size_t memory_free_mb = 0; + bool has_memory_info = false; backend_init_info() = default; - backend_init_info(size_t device_index, size_t total_devices, const std::string & device_name, bool skipped = false, - const std::string & skip_reason = "", const std::string & description = "", - size_t memory_total_mb = 0, size_t memory_free_mb = 0, bool has_memory_info = false) : + backend_init_info(size_t device_index, + size_t total_devices, + const std::string & device_name, + bool skipped = false, + const std::string & skip_reason = "", + const std::string & description = "", + size_t memory_total_mb = 0, + size_t memory_free_mb = 0, + bool has_memory_info = false) : device_index(device_index), total_devices(total_devices), device_name(device_name), @@ -757,27 +664,24 @@ struct backend_init_info { }; struct backend_status_info { - std::string backend_name; + std::string backend_name; test_status_t status; backend_status_info() = default; backend_status_info(const std::string & backend_name, test_status_t status) : - backend_name(backend_name), - status(status) {} + backend_name(backend_name), status(status) {} }; struct overall_summary_info { size_t backends_passed; size_t backends_total; - bool all_passed; + bool all_passed; overall_summary_info() = default; overall_summary_info(size_t backends_passed, size_t backends_total, bool all_passed) : - backends_passed(backends_passed), - backends_total(backends_total), - all_passed(all_passed) {} + backends_passed(backends_passed), backends_total(backends_total), all_passed(all_passed) {} }; struct printer { @@ -791,19 +695,21 @@ struct printer { virtual void print_footer() {} - virtual void print_operation(const test_operation_info & info) { (void) info; } + virtual void print_operation(const test_operation_info & info) { (void)info; } - virtual void print_summary(const test_summary_info & info) { (void) info; } + virtual void print_summary(const test_summary_info & info) { (void)info; } - virtual void print_testing_start(const testing_start_info & info) { (void) info; } + virtual void print_testing_start(const testing_start_info & info) { (void)info; } - virtual void print_backend_init(const backend_init_info & info) { (void) info; } + virtual void print_backend_init(const backend_init_info & info) { (void)info; } - virtual void print_backend_status(const backend_status_info & info) { (void) info; } + virtual void print_backend_status(const backend_status_info & info) { (void)info; } - virtual void print_overall_summary(const overall_summary_info & info) { (void) info; } + virtual void print_overall_summary(const overall_summary_info & info) { (void)info; } - virtual void print_failed_tests(const std::vector & failed_tests) { (void) failed_tests; } + virtual void print_failed_tests(const std::vector & failed_tests) { + (void)failed_tests; + } }; struct console_printer : public printer { @@ -844,19 +750,21 @@ struct console_printer : public printer { } else if (info.error_component == "backend") { fprintf(stderr, " Failed to initialize %s backend\n", info.backend_name.c_str()); } else { - fprintf(stderr, "Error in %s: %s\n", info.error_component.c_str(), info.error_details.c_str()); + fprintf(stderr, "Error in %s: %s\n", info.error_component.c_str(), + info.error_details.c_str()); } } // Handle gradient info if (info.has_gradient_info) { - printf("[%s] nonfinite gradient at index %" PRId64 " (%s=%f) ", info.op_name.c_str(), info.gradient_index, - info.gradient_param_name.c_str(), info.gradient_value); + printf("[%s] nonfinite gradient at index %" PRId64 " (%s=%f) ", info.op_name.c_str(), + info.gradient_index, info.gradient_param_name.c_str(), info.gradient_value); } // Handle MAA error if (info.has_maa_error) { - printf("[%s] MAA = %.9f > %.9f ", info.op_name.c_str(), info.maa_error, info.maa_threshold); + printf("[%s] MAA = %.9f > %.9f ", info.op_name.c_str(), info.maa_error, + info.maa_threshold); } // Handle compare failure @@ -894,7 +802,8 @@ struct console_printer : public printer { } void print_backend_init(const backend_init_info & info) override { - printf("Backend %zu/%zu: %s\n", info.device_index + 1, info.total_devices, info.device_name.c_str()); + printf("Backend %zu/%zu: %s\n", info.device_index + 1, info.total_devices, + info.device_name.c_str()); if (info.skipped) { printf(" %s\n", info.skip_reason.c_str()); @@ -906,7 +815,8 @@ struct console_printer : public printer { } if (info.has_memory_info) { - printf(" Device memory: %zu MB (%zu MB free)\n", info.memory_total_mb, info.memory_free_mb); + printf(" Device memory: %zu MB (%zu MB free)\n", info.memory_total_mb, + info.memory_free_mb); } printf("\n"); @@ -961,7 +871,7 @@ struct console_printer : public printer { // align while also leaving some margin for variations in parameters int align = 8; - int last = (len + align - 1) / align * align; + int last = (len + align - 1) / align * align; if (last - len < 5) { last += align; } @@ -987,7 +897,8 @@ struct console_printer : public printer { printf("%s/run - \033[1;34m%sS\033[0m", format_flops(op_flops_per_run).c_str(), format_flops(result.flops).c_str()); } else { - printf("%8zu kB/run - \033[1;34m%7.2f GB/s\033[0m", result.memory_kb, result.bandwidth_gb_s); + printf("%8zu kB/run - \033[1;34m%7.2f GB/s\033[0m", result.memory_kb, + result.bandwidth_gb_s); } printf("\n"); } @@ -1047,10 +958,11 @@ struct sql_printer : public printer { struct csv_printer : public printer { void print_header() override { - std::vector fields = test_result::get_fields(); + std::vector fields = test_result::get_fields(); std::vector fields_csv = get_fields_csv(); for (size_t i = 0; i < fields.size(); i++) { - if (std::find(std::begin(fields_csv), std::end(fields_csv), fields[i]) == std::end(fields_csv)) { + if (std::find(std::begin(fields_csv), std::end(fields_csv), fields[i]) == + std::end(fields_csv)) { continue; } printf("\"%s\"%s", fields[i].c_str(), i < fields.size() - 1 ? "," : ""); @@ -1060,13 +972,14 @@ struct csv_printer : public printer { void print_test_result(const test_result & result) override { - std::vector values = result.get_values(); - std::vector fields = test_result::get_fields(); + std::vector values = result.get_values(); + std::vector fields = test_result::get_fields(); std::vector fields_csv = get_fields_csv(); for (size_t i = 0; i < values.size(); i++) { - if (std::find(std::begin(fields_csv), std::end(fields_csv), fields[i]) == std::end(fields_csv)) { + if (std::find(std::begin(fields_csv), std::end(fields_csv), fields[i]) == + std::end(fields_csv)) { continue; } @@ -1084,16 +997,10 @@ struct csv_printer : public printer { static std::vector get_fields_csv() { return { - "op_name", - "op_params", - "supported", - "error_message", - "test_mode", - "backend_reg_name", - "backend_name", + "op_name", "op_params", "supported", "error_message", + "test_mode", "backend_reg_name", "backend_name", }; } - }; static std::unique_ptr create_printer(output_formats format) { @@ -1111,65 +1018,46 @@ static std::unique_ptr create_printer(output_formats format) { struct test_case { virtual ~test_case() {} - virtual std::string op_desc(ggml_tensor * t) { - return ggml_op_desc(t); - } + virtual std::string op_desc(ggml_tensor * t) { return ggml_op_desc(t); } - virtual std::string vars() { - return ""; - } + virtual std::string vars() { return ""; } virtual ggml_tensor * build_graph(ggml_context * ctx) = 0; - virtual double max_nmse_err() { - return 1e-7; - } + virtual double max_nmse_err() { return 1e-7; } virtual double max_nmse_err(ggml_backend_t backend) { GGML_UNUSED(backend); return max_nmse_err(); } - virtual double max_maa_err() { - return 1e-4; - } + virtual double max_maa_err() { return 1e-4; } - virtual double max_err() { - return max_nmse_err(); - } + virtual double max_err() { return max_nmse_err(); } - virtual double max_err(ggml_backend_t backend) { - return max_nmse_err(backend); - } + virtual double max_err(ggml_backend_t backend) { return max_nmse_err(backend); } - virtual double err(const float * a, const float * b, size_t n) { - return nmse(a, b, n); - } + virtual double err(const float * a, const float * b, size_t n) { return nmse(a, b, n); } - virtual float grad_eps() { - return 1e-1f; - } + virtual float grad_eps() { return 1e-1f; } // If false, estimate gradient with 2 points, neglects 3rd order derivative and higher. // If true, estimate gradient with 4 points, neglects 5th order derivative and higher. - virtual bool grad_precise() { - return false; - } + virtual bool grad_precise() { return false; } - // Skip gradient checks if total number of gradients to be checked is larger than this (to speed up the tests). - virtual int64_t grad_nmax() { - return 10000; - } + // Skip gradient checks if total number of gradients to be checked is larger than this (to speed + // up the tests). + virtual int64_t grad_nmax() { return 10000; } // No effect if empty. - // If not empty, skip all gradient checks where the numerical result does not match any of the values. - // Needed for dealing with noncontinuous gradients (e.g. ReLU) where estimation using finite differences is unreliable. - virtual std::vector grad_expect() { - return {}; - } + // If not empty, skip all gradient checks where the numerical result does not match any of the + // values. Needed for dealing with noncontinuous gradients (e.g. ReLU) where estimation using + // finite differences is unreliable. + virtual std::vector grad_expect() { return {}; } virtual void initialize_tensors(ggml_context * ctx) { - for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != nullptr; t = ggml_get_next_tensor(ctx, t)) { + for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != nullptr; + t = ggml_get_next_tensor(ctx, t)) { init_tensor_uniform(t); } } @@ -1213,9 +1101,11 @@ struct test_case { sentinels.push_back(sentinel); } - // hijack ggml_new_tensor to add sentinels after each tensor to check for overflows in the backend + // hijack ggml_new_tensor to add sentinels after each tensor to check for overflows in the + // backend - ggml_tensor * ggml_new_tensor(ggml_context * ctx, ggml_type type, int n_dims, const int64_t * ne) { + ggml_tensor * + ggml_new_tensor(ggml_context * ctx, ggml_type type, int n_dims, const int64_t * ne) { ggml_tensor * t = ::ggml_new_tensor(ctx, type, n_dims, ne); add_sentinel(ctx); return t; @@ -1233,19 +1123,22 @@ struct test_case { return t; } - ggml_tensor * ggml_new_tensor_3d(ggml_context * ctx, ggml_type type, int64_t ne0, int64_t ne1, int64_t ne2) { + ggml_tensor * + ggml_new_tensor_3d(ggml_context * ctx, ggml_type type, int64_t ne0, int64_t ne1, int64_t ne2) { ggml_tensor * t = ::ggml_new_tensor_3d(ctx, type, ne0, ne1, ne2); add_sentinel(ctx); return t; } - ggml_tensor * ggml_new_tensor_4d(ggml_context * ctx, ggml_type type, int64_t ne0, int64_t ne1, int64_t ne2, int64_t ne3) { + ggml_tensor * ggml_new_tensor_4d( + ggml_context * ctx, ggml_type type, int64_t ne0, int64_t ne1, int64_t ne2, int64_t ne3) { ggml_tensor * t = ::ggml_new_tensor_4d(ctx, type, ne0, ne1, ne2, ne3); add_sentinel(ctx); return t; } - // Checks an op against the test filter, which is a comma separated list of OP names or specific variations + // Checks an op against the test filter, which is a comma separated list of OP names or specific + // variations bool matches_filter(ggml_tensor * op, const char * op_names_filter) { if (op_names_filter) { const auto op_name = op_desc(op); @@ -1277,12 +1170,12 @@ struct test_case { test_status_t eval(ggml_backend_t backend1, ggml_backend_t backend2, - const char * op_names_filter, - printer * output_printer) { + const char * op_names_filter, + printer * output_printer) { mode = MODE_TEST; ggml_init_params params = { - /* .mem_size = */ ggml_tensor_overhead()*128 + ggml_graph_overhead(), + /* .mem_size = */ ggml_tensor_overhead() * 128 + ggml_graph_overhead(), /* .mem_base = */ NULL, /* .no_alloc = */ true, }; @@ -1295,10 +1188,10 @@ struct test_case { add_sentinel(ctx); ggml_tensor * out = build_graph(ctx); - current_op_name = op_desc(out); + current_op_name = op_desc(out); if (!matches_filter(out, op_names_filter)) { - //printf(" %s: skipping\n", op_desc(out).c_str()); + // printf(" %s: skipping\n", op_desc(out).c_str()); ggml_free(ctx); return test_status_t::SKIPPED; } @@ -1306,7 +1199,8 @@ struct test_case { // check if the backends support the ops bool supported = true; for (ggml_backend_t backend : {backend1, backend2}) { - for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != NULL; t = ggml_get_next_tensor(ctx, t)) { + for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != NULL; + t = ggml_get_next_tensor(ctx, t)) { if (!ggml_backend_supports_op(backend, t)) { supported = false; break; @@ -1316,8 +1210,8 @@ struct test_case { if (!supported) { // Create test result for unsupported operation - test_result result(ggml_backend_name(backend1), current_op_name, vars(), "test", - false, false, "not supported"); + test_result result(ggml_backend_name(backend1), current_op_name, vars(), "test", false, + false, "not supported"); if (output_printer) { output_printer->print_test_result(result); @@ -1352,21 +1246,22 @@ struct test_case { // compare struct callback_userdata { - bool ok; + bool ok; test_case * tc; ggml_backend_t backend1; ggml_backend_t backend2; }; - callback_userdata ud { + callback_userdata ud{ true, this, backend1, backend2, }; - auto callback = [](int index, ggml_tensor * t1, ggml_tensor * t2, void * user_data) -> bool { - callback_userdata * ud = (callback_userdata *) user_data; + auto callback = [](int index, ggml_tensor * t1, ggml_tensor * t2, + void * user_data) -> bool { + callback_userdata * ud = (callback_userdata *)user_data; const char * bn1 = ggml_backend_name(ud->backend1); const char * bn2 = ggml_backend_name(ud->backend2); @@ -1390,7 +1285,8 @@ struct test_case { for (size_t i = 0; i < f1.size(); i++) { // check for nans if (std::isnan(f1[i]) || std::isnan(f2[i])) { - printf("[%s] NaN at index %zu (%s=%f %s=%f) ", ggml_op_desc(t1), i, bn1, f1[i], bn2, f2[i]); + printf("[%s] NaN at index %zu (%s=%f %s=%f) ", ggml_op_desc(t1), i, bn1, f1[i], + bn2, f2[i]); ud->ok = false; return true; } @@ -1398,12 +1294,14 @@ struct test_case { if (isinf_or_max(f1[i]) || isinf_or_max(f2[i])) { if (isinf_or_max(f1[i]) && isinf_or_max(f2[i])) { if (std::signbit(f1[i]) != std::signbit(f2[i])) { - printf("[%s] inf sign mismatch: %s=%f %s=%f ", ggml_op_desc(t1), bn1, f1[i], bn2, f2[i]); + printf("[%s] inf sign mismatch: %s=%f %s=%f ", ggml_op_desc(t1), bn1, + f1[i], bn2, f2[i]); ud->ok = false; return true; } } else { - printf("[%s] inf mismatch: %s=%f %s=%f ", ggml_op_desc(t1), bn1, f1[i], bn2, f2[i]); + printf("[%s] inf mismatch: %s=%f %s=%f ", ggml_op_desc(t1), bn1, f1[i], bn2, + f2[i]); ud->ok = false; return true; } @@ -1412,12 +1310,13 @@ struct test_case { double err = ud->tc->err(f1.data(), f2.data(), f1.size()); if (err > ud->tc->max_err(ud->backend1)) { - printf("[%s] ERR = %.9f > %.9f ", ggml_op_desc(t1), err, ud->tc->max_err(ud->backend1)); - //for (int i = 0; i < (int) f1.size(); i++) { - // printf("%5d %9.6f %9.6f, diff = %9.6f\n", i, f1[i], f2[i], f1[i] - f2[i]); - //} - //printf("\n"); - //exit(1); + printf("[%s] ERR = %.9f > %.9f ", ggml_op_desc(t1), err, + ud->tc->max_err(ud->backend1)); + // for (int i = 0; i < (int) f1.size(); i++) { + // printf("%5d %9.6f %9.6f, diff = %9.6f\n", i, f1[i], f2[i], f1[i] - f2[i]); + // } + // printf("\n"); + // exit(1); ud->ok = false; } return true; @@ -1429,19 +1328,20 @@ struct test_case { if (fused_nodes_to_verify.size() == 0 && run_whole_graph()) { fused_nodes_to_verify.push_back(out); } - const bool cmp_ok = ggml_backend_compare_graph_backend(backend1, backend2, gf, callback, &ud, - run_whole_graph() ? fused_nodes_to_verify.data() : nullptr, - fused_nodes_to_verify.size()); + const bool cmp_ok = ggml_backend_compare_graph_backend( + backend1, backend2, gf, callback, &ud, + run_whole_graph() ? fused_nodes_to_verify.data() : nullptr, + fused_nodes_to_verify.size()); ggml_backend_buffer_free(buf); ggml_free(ctx); // Create test result - bool test_passed = ud.ok && cmp_ok; - std::string error_msg = test_passed ? "" : (!cmp_ok ? "compare failed" : "test failed"); - test_result result(ggml_backend_name(backend1), current_op_name, vars(), "test", supported, test_passed, - error_msg); + bool test_passed = ud.ok && cmp_ok; + std::string error_msg = test_passed ? "" : (!cmp_ok ? "compare failed" : "test failed"); + test_result result(ggml_backend_name(backend1), current_op_name, vars(), "test", supported, + test_passed, error_msg); if (output_printer) { output_printer->print_test_result(result); @@ -1456,24 +1356,25 @@ struct test_case { static const size_t graph_nodes = 8192; ggml_init_params params = { - /* .mem_size = */ ggml_tensor_overhead()*128 + ggml_graph_overhead_custom(graph_nodes, false), + /* .mem_size = */ ggml_tensor_overhead() * 128 + + ggml_graph_overhead_custom(graph_nodes, false), /* .mem_base = */ NULL, /* .no_alloc = */ true, }; ggml_context_ptr ctx(ggml_init(params)); // smart ptr GGML_ASSERT(ctx); - ggml_tensor * out = build_graph(ctx.get()); - current_op_name = op_desc(out); + ggml_tensor * out = build_graph(ctx.get()); + current_op_name = op_desc(out); if (!matches_filter(out, op_names_filter)) { - //printf(" %s: skipping\n", op_desc(out).c_str()); + // printf(" %s: skipping\n", op_desc(out).c_str()); return true; } if (!ggml_backend_supports_op(backend, out)) { // Create test result for unsupported performance test - test_result result(ggml_backend_name(backend), current_op_name, vars(), "perf", false, false, - "not supported"); + test_result result(ggml_backend_name(backend), current_op_name, vars(), "perf", false, + false, "not supported"); output_printer->print_test_result(result); @@ -1481,7 +1382,8 @@ struct test_case { } // allocate - ggml_backend_buffer_ptr buf(ggml_backend_alloc_ctx_tensors(ctx.get(), backend)); // smart ptr + ggml_backend_buffer_ptr buf( + ggml_backend_alloc_ctx_tensors(ctx.get(), backend)); // smart ptr if (buf == NULL) { printf("failed to allocate tensors\n"); @@ -1498,27 +1400,33 @@ struct test_case { // warmup run ggml_status status = ggml_backend_graph_compute(backend, gf); if (status != GGML_STATUS_SUCCESS) { - fprintf(stderr, "%s: ggml_backend_graph_compute failed. status=%s \n", __func__, ggml_status_to_string(status)); + fprintf(stderr, "%s: ggml_backend_graph_compute failed. status=%s \n", __func__, + ggml_status_to_string(status)); return false; } // determine number of runs int n_runs; - bool is_cpu = ggml_backend_dev_type(ggml_backend_get_device(backend)) == GGML_BACKEND_DEVICE_TYPE_CPU; + bool is_cpu = + ggml_backend_dev_type(ggml_backend_get_device(backend)) == GGML_BACKEND_DEVICE_TYPE_CPU; if (op_flops(out) > 0) { // based on flops const uint64_t GFLOP = 1000 * 1000 * 1000; - const uint64_t target_flops_cpu = 8ULL * GFLOP; + const uint64_t target_flops_cpu = 8ULL * GFLOP; const uint64_t target_flops_gpu = 100ULL * GFLOP; uint64_t target_flops = is_cpu ? target_flops_cpu : target_flops_gpu; - n_runs = (int)std::min(ggml_graph_size(gf) - ggml_graph_n_nodes(gf), target_flops / op_flops(out)) + 1; + n_runs = (int)std::min(ggml_graph_size(gf) - ggml_graph_n_nodes(gf), + target_flops / op_flops(out)) + + 1; } else { // based on memory size const size_t GB = 1ULL << 30; - const size_t target_size_cpu = 8 * GB; + const size_t target_size_cpu = 8 * GB; const size_t target_size_gpu = 32 * GB; size_t target_size = is_cpu ? target_size_cpu : target_size_gpu; - n_runs = (int)std::min(ggml_graph_size(gf) - ggml_graph_n_nodes(gf), target_size / op_size(out)) + 1; + n_runs = (int)std::min(ggml_graph_size(gf) - ggml_graph_n_nodes(gf), + target_size / op_size(out)) + + 1; } // duplicate the op @@ -1553,7 +1461,8 @@ struct test_case { int64_t start_time = ggml_time_us(); ggml_status status = ggml_backend_graph_compute(backend, gf); if (status != GGML_STATUS_SUCCESS) { - fprintf(stderr, "%s: ggml_backend_graph_compute failed. status=%s \n", __func__, ggml_status_to_string(status)); + fprintf(stderr, "%s: ggml_backend_graph_compute failed. status=%s \n", __func__, + ggml_status_to_string(status)); return false; } int64_t end_time = ggml_time_us(); @@ -1561,17 +1470,20 @@ struct test_case { total_time_us += end_time - start_time; total_mem += mem; total_runs += n_runs; - } while (total_time_us < 1000*1000); // run for at least 1 second + } while (total_time_us < 1000 * 1000); // run for at least 1 second // Create test result - double avg_time_us = (double) total_time_us / total_runs; - double calculated_flops = (op_flops(out) > 0) ? (op_flops(out) * total_runs) / (total_time_us / 1e6) : 0.0; + double avg_time_us = (double)total_time_us / total_runs; + double calculated_flops = + (op_flops(out) > 0) ? (op_flops(out) * total_runs) / (total_time_us / 1e6) : 0.0; double calculated_bandwidth = - (op_flops(out) == 0) ? total_mem / (total_time_us / 1e6) / 1024.0 / 1024.0 / 1024.0 : 0.0; + (op_flops(out) == 0) ? total_mem / (total_time_us / 1e6) / 1024.0 / 1024.0 / 1024.0 + : 0.0; size_t calculated_memory_kb = op_size(out) / 1024; - test_result result(ggml_backend_name(backend), current_op_name, vars(), "perf", true, true, "", avg_time_us, - calculated_flops, calculated_bandwidth, calculated_memory_kb, total_runs); + test_result result(ggml_backend_name(backend), current_op_name, vars(), "perf", true, true, + "", avg_time_us, calculated_flops, calculated_bandwidth, + calculated_memory_kb, total_runs); if (output_printer) { output_printer->print_test_result(result); @@ -1580,13 +1492,15 @@ struct test_case { return true; } - bool eval_support(ggml_backend_t backend, const char * op_names_filter, printer * output_printer) { + bool + eval_support(ggml_backend_t backend, const char * op_names_filter, printer * output_printer) { mode = MODE_SUPPORT; static const size_t graph_nodes = 8192; ggml_init_params params = { - /* .mem_size = */ ggml_tensor_overhead()*128 + ggml_graph_overhead_custom(graph_nodes, false), + /* .mem_size = */ ggml_tensor_overhead() * 128 + + ggml_graph_overhead_custom(graph_nodes, false), /* .mem_base = */ NULL, /* .no_alloc = */ true, }; @@ -1596,7 +1510,7 @@ struct test_case { gf = ggml_new_graph_custom(ctx.get(), graph_nodes, false); ggml_tensor * out = build_graph(ctx.get()); - current_op_name = op_desc(out); + current_op_name = op_desc(out); if (!matches_filter(out, op_names_filter)) { return true; @@ -1605,10 +1519,12 @@ struct test_case { bool supported = ggml_backend_supports_op(backend, out); std::string device_desc = ggml_backend_dev_description(ggml_backend_get_device(backend)); - std::string backend_reg_name = ggml_backend_reg_name(ggml_backend_dev_backend_reg(ggml_backend_get_device(backend))); + std::string backend_reg_name = + ggml_backend_reg_name(ggml_backend_dev_backend_reg(ggml_backend_get_device(backend))); - test_result result(ggml_backend_name(backend), current_op_name, vars(), "support", supported, supported, - supported ? "yes" : "no", 0.0, 0.0, 0.0, 0, 0, device_desc, backend_reg_name); + test_result result(ggml_backend_name(backend), current_op_name, vars(), "support", + supported, supported, supported ? "yes" : "no", 0.0, 0.0, 0.0, 0, 0, + device_desc, backend_reg_name); output_printer->print_test_result(result); @@ -1620,7 +1536,8 @@ struct test_case { const std::vector expect = grad_expect(); ggml_init_params params = { - /* .mem_size = */ ggml_tensor_overhead()*128 + 2*ggml_graph_overhead_custom(GGML_DEFAULT_GRAPH_SIZE, true), + /* .mem_size = */ ggml_tensor_overhead() * 128 + + 2 * ggml_graph_overhead_custom(GGML_DEFAULT_GRAPH_SIZE, true), /* .mem_base = */ NULL, /* .no_alloc = */ true, }; @@ -1637,48 +1554,52 @@ struct test_case { } if (out->type != GGML_TYPE_F32) { - output_printer->print_operation(test_operation_info(op_desc(out), vars(), ggml_backend_name(backend), - test_status_t::NOT_SUPPORTED, - out->name + std::string("->type != FP32"))); + output_printer->print_operation(test_operation_info( + op_desc(out), vars(), ggml_backend_name(backend), test_status_t::NOT_SUPPORTED, + out->name + std::string("->type != FP32"))); return true; } // Print operation info first - output_printer->print_operation(test_operation_info(op_desc(out), vars(), ggml_backend_name(backend))); + output_printer->print_operation( + test_operation_info(op_desc(out), vars(), ggml_backend_name(backend))); // check if the backend supports the ops - bool supported = true; - bool any_params = false; + bool supported = true; + bool any_params = false; std::string failure_reason; - for (ggml_tensor * t = ggml_get_first_tensor(ctx.get()); t != NULL; t = ggml_get_next_tensor(ctx.get(), t)) { + for (ggml_tensor * t = ggml_get_first_tensor(ctx.get()); t != NULL; + t = ggml_get_next_tensor(ctx.get(), t)) { if (!ggml_backend_supports_op(backend, t)) { - supported = false; + supported = false; failure_reason = ggml_backend_name(backend); break; } if ((t->flags & GGML_TENSOR_FLAG_PARAM)) { any_params = true; if (t->type != GGML_TYPE_F32) { - supported = false; + supported = false; failure_reason = std::string(t->name) + "->type != FP32"; break; } } } if (!any_params) { - supported = false; + supported = false; failure_reason = op_desc(out); } if (!supported) { - output_printer->print_operation(test_operation_info(op_desc(out), vars(), ggml_backend_name(backend), - test_status_t::NOT_SUPPORTED, failure_reason)); + output_printer->print_operation( + test_operation_info(op_desc(out), vars(), ggml_backend_name(backend), + test_status_t::NOT_SUPPORTED, failure_reason)); return true; } int64_t ngrads = 0; - for (ggml_tensor * t = ggml_get_first_tensor(ctx.get()); t != NULL; t = ggml_get_next_tensor(ctx.get(), t)) { + for (ggml_tensor * t = ggml_get_first_tensor(ctx.get()); t != NULL; + t = ggml_get_next_tensor(ctx.get(), t)) { if (t->flags & GGML_TENSOR_FLAG_PARAM) { ngrads += ggml_nelements(t); } @@ -1690,7 +1611,6 @@ struct test_case { return true; } - if (!ggml_is_scalar(out)) { out = ggml_sum(ctx.get(), out); ggml_set_name(out, "sum_of_out"); @@ -1702,23 +1622,26 @@ struct test_case { ggml_build_backward_expand(ctx.get(), gb, nullptr); if (expect.size() != 1 || expect[0] != 0.0f) { GGML_ASSERT(ggml_graph_n_nodes(gb) > ggml_graph_n_nodes(gf)); - for (ggml_tensor * t = ggml_get_first_tensor(ctx.get()); t != NULL; t = ggml_get_next_tensor(ctx.get(), t)) { - GGML_ASSERT(!(t->flags & GGML_TENSOR_FLAG_PARAM) || ggml_graph_get_grad(gb, t)->op != GGML_OP_NONE); + for (ggml_tensor * t = ggml_get_first_tensor(ctx.get()); t != NULL; + t = ggml_get_next_tensor(ctx.get(), t)) { + GGML_ASSERT(!(t->flags & GGML_TENSOR_FLAG_PARAM) || + ggml_graph_get_grad(gb, t)->op != GGML_OP_NONE); } } - for (ggml_tensor * t = ggml_get_first_tensor(ctx.get()); t != NULL; t = ggml_get_next_tensor(ctx.get(), t)) { + for (ggml_tensor * t = ggml_get_first_tensor(ctx.get()); t != NULL; + t = ggml_get_next_tensor(ctx.get(), t)) { if (!ggml_backend_supports_op(backend, t)) { - output_printer->print_operation(test_operation_info(op_desc(out), vars(), ggml_backend_name(backend), - test_status_t::NOT_SUPPORTED, - ggml_backend_name(backend))); + output_printer->print_operation( + test_operation_info(op_desc(out), vars(), ggml_backend_name(backend), + test_status_t::NOT_SUPPORTED, ggml_backend_name(backend))); supported = false; break; } if ((t->flags & GGML_TENSOR_FLAG_PARAM) && t->type != GGML_TYPE_F32) { - output_printer->print_operation(test_operation_info(op_desc(out), vars(), ggml_backend_name(backend), - test_status_t::NOT_SUPPORTED, - std::string(t->name) + "->type != FP32")); + output_printer->print_operation(test_operation_info( + op_desc(out), vars(), ggml_backend_name(backend), test_status_t::NOT_SUPPORTED, + std::string(t->name) + "->type != FP32")); supported = false; break; } @@ -1728,7 +1651,8 @@ struct test_case { } // allocate - ggml_backend_buffer_ptr buf(ggml_backend_alloc_ctx_tensors(ctx.get(), backend)); // smart ptr + ggml_backend_buffer_ptr buf( + ggml_backend_alloc_ctx_tensors(ctx.get(), backend)); // smart ptr if (buf == NULL) { test_operation_info info(op_desc(out), vars(), ggml_backend_name(backend)); info.set_error("allocation", ""); @@ -1737,21 +1661,24 @@ struct test_case { } initialize_tensors(ctx.get()); // Randomizes all tensors (including gradients). - ggml_graph_reset(gb); // Sets gradients to 1 if loss, 0 otherwise. + ggml_graph_reset(gb); // Sets gradients to 1 if loss, 0 otherwise. ggml_status status = ggml_backend_graph_compute(backend, gf); if (status != GGML_STATUS_SUCCESS) { - fprintf(stderr, "%s: ggml_backend_graph_compute failed. status=%s \n", __func__, ggml_status_to_string(status)); + fprintf(stderr, "%s: ggml_backend_graph_compute failed. status=%s \n", __func__, + ggml_status_to_string(status)); return false; } status = ggml_backend_graph_compute(backend, gb); if (status != GGML_STATUS_SUCCESS) { - fprintf(stderr, "%s: ggml_backend_graph_compute failed. status=%s \n", __func__, ggml_status_to_string(status)); + fprintf(stderr, "%s: ggml_backend_graph_compute failed. status=%s \n", __func__, + ggml_status_to_string(status)); return false; } bool ok = true; - for (struct ggml_tensor * t = ggml_get_first_tensor(ctx.get()); t != nullptr; t = ggml_get_next_tensor(ctx.get(), t)) { + for (struct ggml_tensor * t = ggml_get_first_tensor(ctx.get()); t != nullptr; + t = ggml_get_next_tensor(ctx.get(), t)) { if (!(t->flags & GGML_TENSOR_FLAG_PARAM)) { continue; } @@ -1790,49 +1717,54 @@ struct test_case { const float eps = grad_eps(); for (int64_t i = 0; i < ne; ++i) { - const float xiu = x0[i] + 1.0f*eps; // x, index i, up - const float xiuh = x0[i] + 0.5f*eps; // x, index i, up half - const float xidh = x0[i] - 0.5f*eps; // x, index i, down half - const float xid = x0[i] - 1.0f*eps; // x, index i, down + const float xiu = x0[i] + 1.0f * eps; // x, index i, up + const float xiuh = x0[i] + 0.5f * eps; // x, index i, up half + const float xidh = x0[i] - 0.5f * eps; // x, index i, down half + const float xid = x0[i] - 1.0f * eps; // x, index i, down float fu, fuh, fdh, fd; // output values for xiu, xiuh, xid, xidh - ggml_backend_tensor_set(t, &xiu, i*sizeof(float), sizeof(float)); + ggml_backend_tensor_set(t, &xiu, i * sizeof(float), sizeof(float)); status = ggml_backend_graph_compute(backend, gf); if (status != GGML_STATUS_SUCCESS) { - fprintf(stderr, "%s: ggml_backend_graph_compute failed. status=%s \n", __func__, ggml_status_to_string(status)); + fprintf(stderr, "%s: ggml_backend_graph_compute failed. status=%s \n", __func__, + ggml_status_to_string(status)); return false; } ggml_backend_tensor_get(out, &fu, 0, ggml_nbytes(out)); - ggml_backend_tensor_set(t, &xid, i*sizeof(float), sizeof(float)); + ggml_backend_tensor_set(t, &xid, i * sizeof(float), sizeof(float)); status = ggml_backend_graph_compute(backend, gf); if (status != GGML_STATUS_SUCCESS) { - fprintf(stderr, "%s: ggml_backend_graph_compute failed. status=%s \n", __func__, ggml_status_to_string(status)); + fprintf(stderr, "%s: ggml_backend_graph_compute failed. status=%s \n", __func__, + ggml_status_to_string(status)); return false; } ggml_backend_tensor_get(out, &fd, 0, ggml_nbytes(out)); if (grad_precise()) { - ggml_backend_tensor_set(t, &xiuh, i*sizeof(float), sizeof(float)); + ggml_backend_tensor_set(t, &xiuh, i * sizeof(float), sizeof(float)); status = ggml_backend_graph_compute(backend, gf); if (status != GGML_STATUS_SUCCESS) { - fprintf(stderr, "%s: ggml_backend_graph_compute failed. status=%s \n", __func__, ggml_status_to_string(status)); + fprintf(stderr, "%s: ggml_backend_graph_compute failed. status=%s \n", + __func__, ggml_status_to_string(status)); return false; } ggml_backend_tensor_get(out, &fuh, 0, ggml_nbytes(out)); - ggml_backend_tensor_set(t, &xidh, i*sizeof(float), sizeof(float)); + ggml_backend_tensor_set(t, &xidh, i * sizeof(float), sizeof(float)); status = ggml_backend_graph_compute(backend, gf); if (status != GGML_STATUS_SUCCESS) { - fprintf(stderr, "%s: ggml_backend_graph_compute failed. status=%s \n", __func__, ggml_status_to_string(status)); + fprintf(stderr, "%s: ggml_backend_graph_compute failed. status=%s \n", + __func__, ggml_status_to_string(status)); return false; } ggml_backend_tensor_get(out, &fdh, 0, ggml_nbytes(out)); - gn[i] = (8.0*(double)fuh + (double)fd - (8.0*(double)fdh + (double)fu)) / (6.0*(double)eps); + gn[i] = (8.0 * (double)fuh + (double)fd - (8.0 * (double)fdh + (double)fu)) / + (6.0 * (double)eps); } else { - gn[i] = (fu - fd) / (2.0f*eps); + gn[i] = (fu - fd) / (2.0f * eps); } ggml_backend_tensor_set(t, x0.data(), 0, ggml_nbytes(t)); @@ -1867,58 +1799,10 @@ struct test_case { } }; - // ################################### // ## Section 2: GGML Op Defintions ## // ################################### - -// The following is an example showing the bare minimum for creating a test for a GGML op. - -// GGML_OP_EXAMPLE -struct test_example : public test_case { - // Always define these 2 or variants thereof: - const ggml_type type; // The type of the input tensors. - const std::array ne; // The shape of the input tensors. - // For some ops it's necessary to define multiple types or shapes for the inputs. - // Or they may need additional parameters. - - // Put all parameters needed to fully define the test into one of the VARS_TO_STR macros. - // In most cases these are just the properties of the struct that you defined above. - // This is needed for info prints. - std::string vars() override { - return VARS_TO_STR2(type, ne); - } - - // Define a constructor for the struct. - // In most cases it will be sufficient to have the same arguments as the struct has properties - // and just use initializer lists. - test_example(ggml_type type = GGML_TYPE_F32, - std::array ne = {10, 5, 4, 3}) - : type(type), ne(ne) {} - - // Define how a simple GGML compute graph can be constructed for the new GGML op. - ggml_tensor * build_graph(ggml_context * ctx) override { - // Step 1: create input tensors that don't depend on any other tensors: - ggml_tensor * a = ggml_new_tensor(ctx, type, 4, ne.data()); - ggml_set_name(a, "a"); // Setting names is optional but it's useful for debugging. - - ggml_tensor * b = ggml_new_tensor(ctx, type, 4, ne.data()); - ggml_set_name(b, "b"); - - // Step 2: use the op that you want to test in the GGML compute graph. - ggml_tensor * out = ggml_add(ctx, a, b); // For this example we're just doing a simple addition. - ggml_set_name(out, "out"); - - // Step 3: return the output tensor. - return out; - } - // In order to also check the gradients for your op, add calls like ggml_set_param(a) - // immediately after you create the tensors. - // This is optional and only makes sense if a backward pass has actually been implemented for the new op. -}; - - // GGML_OP_UNARY struct test_unary : public test_case { const ggml_unary_op op; @@ -1926,20 +1810,19 @@ struct test_unary : public test_case { const std::array ne_a; int v; // view (1 : non-contiguous a) - std::string vars() override { - return VARS_TO_STR3(type, ne_a, v); - } + std::string vars() override { return VARS_TO_STR3(type, ne_a, v); } test_unary(ggml_unary_op op, - ggml_type type = GGML_TYPE_F32, - std::array ne_a = {128, 2, 2, 2}, - int v = 0) - : op(op), type(type), ne_a(ne_a), v(v) {} + ggml_type type = GGML_TYPE_F32, + std::array ne_a = {128, 2, 2, 2}, + int v = 0) : + op(op), type(type), ne_a(ne_a), v(v) {} ggml_tensor * build_graph(ggml_context * ctx) override { - const bool grad_supported = op == GGML_UNARY_OP_ABS || op == GGML_UNARY_OP_SGN || op == GGML_UNARY_OP_NEG || - op == GGML_UNARY_OP_STEP || op == GGML_UNARY_OP_RELU || op == GGML_UNARY_OP_SILU || - op == GGML_UNARY_OP_EXPM1 || op == GGML_UNARY_OP_SOFTPLUS; + const bool grad_supported = op == GGML_UNARY_OP_ABS || op == GGML_UNARY_OP_SGN || + op == GGML_UNARY_OP_NEG || op == GGML_UNARY_OP_STEP || + op == GGML_UNARY_OP_RELU || op == GGML_UNARY_OP_SILU || + op == GGML_UNARY_OP_EXPM1 || op == GGML_UNARY_OP_SOFTPLUS; ggml_tensor * a; if (v & 1) { @@ -1954,7 +1837,8 @@ struct test_unary : public test_case { } ggml_set_name(a, "a"); - a = ggml_view_4d(ctx, a, ne_a[0], ne_a[1], ne_a[2], ne_a[3], a->nb[1], a->nb[2], a->nb[3], 0); + a = ggml_view_4d(ctx, a, ne_a[0], ne_a[1], ne_a[2], ne_a[3], a->nb[1], a->nb[2], + a->nb[3], 0); ggml_set_name(a, "view_of_a"); } else { a = ggml_new_tensor(ctx, type, 4, ne_a.data()); @@ -1971,15 +1855,14 @@ struct test_unary : public test_case { } void initialize_tensors(ggml_context * ctx) override { - for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != NULL; t = ggml_get_next_tensor(ctx, t)) { + for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != NULL; + t = ggml_get_next_tensor(ctx, t)) { // test extended range of values to check for NaNs in GELU init_tensor_uniform(t, -150.f, 150.f); } } - float grad_eps() override { - return 15.0f; - } + float grad_eps() override { return 15.0f; } std::vector grad_expect() override { if (op == GGML_UNARY_OP_ABS) { @@ -1993,5081 +1876,639 @@ struct test_unary : public test_case { } return {}; } - }; -// GGML_OP_GLU -struct test_glu : public test_case { - const ggml_glu_op op; +// GGML_OP_ARGMAX +struct test_argmax : public test_case { const ggml_type type; - const std::array ne_a; - int v; // view (1 : non-contiguous a) - bool swapped; + const std::array ne; - std::string vars() override { - return VARS_TO_STR4(type, ne_a, v, swapped); - } + std::string vars() override { return VARS_TO_STR2(type, ne); } - test_glu(ggml_glu_op op, - ggml_type type = GGML_TYPE_F32, - std::array ne_a = {128, 2, 2, 2}, - int v = 0, - bool swapped = false) - : op(op), type(type), ne_a(ne_a), v(v), swapped(swapped) {} + test_argmax(ggml_type type = GGML_TYPE_F32, std::array ne = {10, 100, 1, 1}) : + type(type), ne(ne) {} ggml_tensor * build_graph(ggml_context * ctx) override { - ggml_tensor * a; - if (v & 1) { - auto ne = ne_a; ne[0] *= 3; - a = ggml_new_tensor(ctx, type, 4, ne.data()); - ggml_set_name(a, "a"); - - a = ggml_view_4d(ctx, a, ne_a[0], ne_a[1], ne_a[2], ne_a[3], a->nb[1], a->nb[2], a->nb[3], 0); - ggml_set_name(a, "view_of_a"); - } else { - a = ggml_new_tensor(ctx, type, 4, ne_a.data()); - ggml_set_name(a, "a"); - } + ggml_tensor * a = ggml_new_tensor(ctx, type, 4, ne.data()); + ggml_set_name(a, "a"); - ggml_tensor * out = ggml_glu(ctx, a, op, swapped); + ggml_tensor * out = ggml_argmax(ctx, a); ggml_set_name(out, "out"); return out; } void initialize_tensors(ggml_context * ctx) override { - for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != NULL; t = ggml_get_next_tensor(ctx, t)) { - // test extended range of values to check for NaNs in GELU - init_tensor_uniform(t, -150.f, 150.f); + std::random_device rd; + std::default_random_engine rng(rd()); + for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != NULL; + t = ggml_get_next_tensor(ctx, t)) { + if (t->type == GGML_TYPE_F32) { + // initialize with unique values to avoid ties + for (int64_t r = 0; r < ggml_nrows(t); r++) { + std::vector data(t->ne[0]); + for (int i = 0; i < t->ne[0]; i++) { + data[i] = i; + } + std::shuffle(data.begin(), data.end(), rng); + ggml_backend_tensor_set(t, data.data(), r * t->nb[1], t->ne[0] * sizeof(float)); + } + } else { + init_tensor_uniform(t); + } } } + + double max_nmse_err() override { return 0.0; } }; -struct test_glu_split : public test_case { - const ggml_glu_op op; +// GGML_OP_ADD +// GGML_OP_SUB +// GGML_OP_MUL +// GGML_OP_DIV +struct test_bin_bcast : public test_case { + using op_t = ggml_tensor * (*)(ggml_context *, ggml_tensor *, ggml_tensor *); + op_t op; const ggml_type type; - const std::array ne_a; - int v; // view (1 : non-contiguous a) - - std::string vars() override { - return VARS_TO_STR3(type, ne_a, v) + ",split"; - } - - test_glu_split(ggml_glu_op op, - ggml_type type = GGML_TYPE_F32, - std::array ne_a = {128, 2, 2, 2}, - int v = 0) - : op(op), type(type), ne_a(ne_a), v(v) {} + const std::array ne; + const std::array nr; + int nf; // number of fused ops, nf == 1 -> single op (no fusion) + bool perm1; // permute src1? - ggml_tensor * build_graph(ggml_context * ctx) override { - ggml_tensor * a; - ggml_tensor * b; - if (v & 1) { - auto ne = ne_a; ne[0] *= 3; - a = ggml_new_tensor(ctx, type, 4, ne.data()); - ggml_set_param(a); - ggml_set_name(a, "a"); + bool run_whole_graph() override { return nf > 1; } - a = ggml_view_4d(ctx, a, ne_a[0], ne_a[1], ne_a[2], ne_a[3], a->nb[1], a->nb[2], a->nb[3], 0); - ggml_set_name(a, "view_of_a"); + std::string vars() override { return VARS_TO_STR5(type, ne, nr, nf, perm1); } - b = ggml_new_tensor(ctx, type, 4, ne.data()); - ggml_set_param(b); - ggml_set_name(b, "b"); + size_t op_size(ggml_tensor * t) override { return ggml_nbytes(t) * 3; } - b = ggml_view_4d(ctx, b, ne_a[0], ne_a[1], ne_a[2], ne_a[3], b->nb[1], b->nb[2], b->nb[3], 0); - ggml_set_name(a, "view_of_b"); - } else { - a = ggml_new_tensor(ctx, type, 4, ne_a.data()); - ggml_set_param(a); - ggml_set_name(a, "a"); + test_bin_bcast(op_t op, + ggml_type type = GGML_TYPE_F32, + std::array ne = {10, 10, 1, 1}, + std::array nr = {1, 2, 1, 1}, + int nf = 1, + bool perm1 = false) : + op(op), type(type), ne(ne), nr(nr), nf(nf), perm1(perm1) {} - b = ggml_new_tensor(ctx, type, 4, ne_a.data()); - ggml_set_param(b); - ggml_set_name(b, "b"); + double max_nmse_err(ggml_backend_t backend) override { + // HSA backend converts F16 to BF16, which has lower precision (7-bit mantissa) + if ((type == GGML_TYPE_F16 || type == GGML_TYPE_BF16) && + backend_has_feature(backend, "SUBSTITUTE_FP16_BF16")) { + return 1e-4; // BF16 precision limit } - - ggml_tensor * out = ggml_glu_split(ctx, a, b, op); - ggml_set_name(out, "out"); - - return out; + return 1e-7; } - void initialize_tensors(ggml_context * ctx) override { - for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != NULL; t = ggml_get_next_tensor(ctx, t)) { - // test extended range of values to check for NaNs in GELU - init_tensor_uniform(t, -150.f, 150.f); - } - } -}; + ggml_tensor * build_graph(ggml_context * ctx) override { + GGML_ASSERT(nf <= 16); -struct test_swiglu_oai : public test_case { - const ggml_type type; - const std::array ne_a; - int v; // view (1 : non-contiguous a) - float alpha; - float limit; + ggml_tensor * a = ggml_new_tensor_4d(ctx, type, ne[0] * nr[0], ne[1] * nr[1], ne[2] * nr[2], + ne[3] * nr[3]); + ggml_set_name(a, "a"); - std::string vars() override { - return VARS_TO_STR5(type, ne_a, v, alpha, limit); - } + ggml_tensor * b[16]; + for (int i = 0; i < nf; ++i) { + if (perm1) { + const int p[4] = {1, 2, 0, 3}; // hardcoded for now - test_swiglu_oai(ggml_type type = GGML_TYPE_F32, - std::array ne_a = {128, 2, 2, 2}, - int v = 0, - float alpha = 1.702f, - float limit = 7.0f) - : type(type), ne_a(ne_a), v(v), alpha(alpha), limit(limit) {} + b[i] = ggml_new_tensor_4d(ctx, type, ne[p[0]], ne[p[1]], ne[p[2]], ne[p[3]]); + b[i] = ggml_permute(ctx, b[i], p[0], p[1], p[2], p[3]); + } else { + b[i] = ggml_new_tensor(ctx, type, 4, ne.data()); + } + ggml_set_name(b[i], (std::string("b") + std::to_string(i)).c_str()); + } - ggml_tensor * build_graph(ggml_context * ctx) override { - ggml_tensor * a; - ggml_tensor * b; - if (v & 1) { - auto ne = ne_a; ne[0] *= 3; - a = ggml_new_tensor(ctx, type, 4, ne.data()); + // The backward pass supports broadcasting only for GGML_ADD: + const bool grad_supported = + op == ggml_add && ggml_are_same_shape(a, b[0]) && nf == 1 && !perm1; + if (grad_supported) { ggml_set_param(a); - ggml_set_name(a, "a"); - - a = ggml_view_4d(ctx, a, ne_a[0], ne_a[1], ne_a[2], ne_a[3], a->nb[1], a->nb[2], a->nb[3], 0); - ggml_set_name(a, "view_of_a"); - - b = ggml_new_tensor(ctx, type, 4, ne.data()); - ggml_set_param(b); - ggml_set_name(b, "b"); + ggml_set_param(b[0]); + } - b = ggml_view_4d(ctx, b, ne_a[0], ne_a[1], ne_a[2], ne_a[3], b->nb[1], b->nb[2], b->nb[3], 0); - ggml_set_name(a, "view_of_b"); - } else { - a = ggml_new_tensor(ctx, type, 4, ne_a.data()); - ggml_set_param(a); - ggml_set_name(a, "a"); + ggml_tensor * out = a; - b = ggml_new_tensor(ctx, type, 4, ne_a.data()); - ggml_set_param(b); - ggml_set_name(b, "b"); + for (int i = 0; i < nf; ++i) { + out = op(ctx, out, b[i]); } - ggml_tensor * out = ggml_swiglu_oai(ctx, a, b, alpha, limit); ggml_set_name(out, "out"); return out; } void initialize_tensors(ggml_context * ctx) override { - for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != NULL; t = ggml_get_next_tensor(ctx, t)) { - // test extended range of values to check for NaNs in GELU - init_tensor_uniform(t, -150.f, 150.f); + for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != NULL; + t = ggml_get_next_tensor(ctx, t)) { + if (op == ggml_mul || op == ggml_div) { + // MUL and DIV have numerical issues around zero: + init_tensor_uniform(t, 0.9f, 1.1f); + } else { + init_tensor_uniform(t); + } } } -}; -// GGML_OP_GET_ROWS -struct test_get_rows : public test_case { - const ggml_type type; - const int n; // cols - const int m; // rows - const int r; // rows to get - const int be1; // batch size - const int be2; // batch size - const bool v; // view (non-contiguous src1) + float grad_eps() override { + return 0.1f * (op == ggml_mul ? ne[0] * ne[1] * ne[2] * ne[3] : 1); + } - std::string vars() override { - return VARS_TO_STR7(type, n, m, r, be1, be2, v); - } - - test_get_rows(ggml_type type = GGML_TYPE_F32, int n = 10, int m = 5, int r = 3, int be1 = 1, int be2 = 1, bool v = false) - : type(type), n(n), m(m), r(r), be1(be1), be2(be2), v(v) {} - - ggml_tensor * build_graph(ggml_context * ctx) override { - ggml_tensor * in = ggml_new_tensor_4d(ctx, type, n, m, be1, be2); - ggml_set_name(in, "in"); - - ggml_tensor * rows = ggml_new_tensor_3d(ctx, GGML_TYPE_I32, r, be1, be2); - ggml_set_name(rows, "rows"); - if (v) { - rows = ggml_view_3d(ctx, rows, r/2, be1, be2, rows->nb[1], rows->nb[2], 0); - ggml_set_name(rows, "view_of_rows"); - } - - const bool grad_supported = ggml_is_matrix(in) && ggml_is_vector(rows); - if (grad_supported) { - ggml_set_param(in); - // rows is a constant input -> no gradients - } - - ggml_tensor * out = ggml_get_rows(ctx, in, rows); - ggml_set_name(out, "out"); - - return out; - } - - void initialize_tensors(ggml_context * ctx) override { - for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != NULL; t = ggml_get_next_tensor(ctx, t)) { - if (t->type == GGML_TYPE_I32) { - if (ggml_is_view_op(t->op)) { continue; } - // rows - std::vector data(r*be1*be2); - for (int i = 0; i < r*be1*be2; i++) { - data[i] = rand() % m; - } - ggml_backend_tensor_set(t, data.data(), 0, r * be1 * be2 * sizeof(int)); - } else { - init_tensor_uniform(t); - } - } - } -}; - -// GGML_OP_GET_ROWS_BACK -struct test_get_rows_back : public test_case { - const ggml_type type; - const int n; // cols - const int m; // rows - const int r; // rows to get - const int b; // batch size - const bool v; // view (non-contiguous src1) - - std::string vars() override { - return VARS_TO_STR6(type, n, m, r, b, v); - } - - test_get_rows_back(ggml_type type = GGML_TYPE_F32, int n = 10, int m = 5, int r = 3, int b = 1, bool v = false) - : type(type), n(n), m(m), r(r), b(b), v(v) {} - - ggml_tensor * build_graph(ggml_context * ctx) override { - ggml_tensor * in_forward = ggml_new_tensor_3d(ctx, type, n, m, b); - ggml_set_name(in_forward, "in_forward"); - - ggml_tensor * rows = ggml_new_tensor_2d(ctx, GGML_TYPE_I32, r, b); - ggml_set_name(rows, "rows"); - if (v) { - rows = ggml_view_2d(ctx, rows, r/2, b, rows->nb[1], 0); - ggml_set_name(rows, "view_of_rows"); - } + bool grad_precise() override { return op == ggml_div; } - ggml_tensor * grad = ggml_new_tensor_3d(ctx, type, n, r, b); - ggml_set_name(grad, "grad"); - - ggml_tensor * out = ggml_get_rows_back(ctx, grad, rows, in_forward); - ggml_set_name(out, "out"); - - return out; - } - - void initialize_tensors(ggml_context * ctx) override { - for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != NULL; t = ggml_get_next_tensor(ctx, t)) { - if (t->type == GGML_TYPE_I32) { - if (ggml_is_view_op(t->op)) { continue; } - // rows - std::vector data(r*b); - for (int i = 0; i < r*b; i++) { - data[i] = rand() % m; - } - ggml_backend_tensor_set(t, data.data(), 0, r * b * sizeof(int)); - } else { - init_tensor_uniform(t); - } - } - } + double max_maa_err() override { return op == ggml_add ? 1e-4 : 1e-3; } }; -static void init_set_rows_row_ids(ggml_tensor * t, int num_rows) { - std::random_device rd; - std::default_random_engine rng(rd()); - for (int i2 = 0; i2 < t->ne[2]; i2++) { - for (int i1 = 0; i1 < t->ne[1]; i1++) { - // generate a shuffled subset of row indices - std::vector data(num_rows); - for (int i = 0; i < num_rows; i++) { - data[i] = i; - } - std::shuffle(data.begin(), data.end(), rng); - data.resize(t->ne[0]); - - const size_t offs = i1*t->nb[1] + i2*t->nb[2]; - if (t->type == GGML_TYPE_I32) { - // TODO: Make a template or something - std::vector data_i32(t->ne[0]); - for (int i = 0; i < t->ne[0]; i++) { - data_i32[i] = static_cast(data[i]); - } - ggml_backend_tensor_set(t, data_i32.data(), offs, t->ne[0]*sizeof(int32_t)); - } else { - ggml_backend_tensor_set(t, data.data(), offs, t->ne[0]*sizeof(int64_t)); - } - } - } -} - -// GGML_OP_SET_ROWS -struct test_set_rows : public test_case { - const ggml_type type; - const ggml_type type_idx; - const std::array ne; - const std::array nr23; // broadcast only dims 2 and 3 - const int r; // rows to set - const bool v; // view (non-contiguous src1) +// GGML_OP_MUL_MAT +struct test_mul_mat : public test_case { + const ggml_type type_a; + const ggml_type type_b; + const int64_t m; + const int64_t n; + const int64_t k; + const std::array bs; // dims 3 and 4 + const std::array nr; // repeat in dims 3 and 4 + const std::array per; // permutation of dimensions + const int64_t k_v; // size of k in memory, resulting in a non-contiguous view for k_v > k, no + // view for k_v == 0 + const uint32_t o; // number of outputs std::string vars() override { - return VARS_TO_STR6(type, type_idx, ne, nr23, r, v); - } - - test_set_rows(ggml_type type, - ggml_type type_idx, - std::array ne, - std::array nr23, - int r, bool v = false) - : type(type), type_idx(type_idx), ne(ne), nr23(nr23), r(r), v(v) {} - - ggml_tensor * build_graph(ggml_context * ctx) override { - ggml_tensor * dst = ggml_new_tensor_4d(ctx, type, ne[0], ne[1], ne[2]*nr23[0], ne[3]*nr23[1]); - ggml_set_name(dst, "dst"); - - ggml_tensor * src = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, ne[0], r, ne[2]*nr23[0], ne[3]*nr23[1]); - ggml_set_name(src, "src"); - - ggml_tensor * row_idxs = ggml_new_tensor_3d(ctx, type_idx, r, ne[2], ne[3]); - ggml_set_name(row_idxs, "row_idxs"); - - if (v) { - src = ggml_view_4d(ctx, src, ne[0], r/2, ne[2]*nr23[0], ne[3]*nr23[1], src->nb[1], src->nb[2], src->nb[3], 0); - row_idxs = ggml_view_3d(ctx, row_idxs, r/2, ne[2], ne[3], row_idxs->nb[1], row_idxs->nb[2], 0); - ggml_set_name(row_idxs, "view_of_rows"); - } - - ggml_tensor * out = ggml_set_rows(ctx, dst, src, row_idxs); - ggml_set_name(out, "out"); - - return out; + return VARS_TO_STR10(type_a, type_b, m, n, k, bs, nr, per, k_v, o); } - void initialize_tensors(ggml_context * ctx) override { - for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != NULL; t = ggml_get_next_tensor(ctx, t)) { - if (t->type == GGML_TYPE_I64 || t->type == GGML_TYPE_I32) { - if (ggml_is_view_op(t->op)) { - continue; - } - - init_set_rows_row_ids(t, ne[1]); - } else { - init_tensor_uniform(t); - } - } - } + double max_nmse_err() override { return 5e-4; } - double max_nmse_err() override { - if (type == GGML_TYPE_Q4_0 || type == GGML_TYPE_Q4_1 || type == GGML_TYPE_IQ4_NL || - type == GGML_TYPE_Q5_0 || type == GGML_TYPE_Q5_1 || type == GGML_TYPE_Q8_0) { - // estimate what the max nmse error would be if one quantized value is - // off by one. The test values are distributed in [-1,1], so it'll be - // roughly (2.0 / 2^bits)^2, divided by the mean square value of the reference, - // which is roughly 0.25 times the number of elements. - double err_estimate = 1.0f/8.0f; - if (type == GGML_TYPE_Q5_0 || type == GGML_TYPE_Q5_1) { - err_estimate /= 2.0f; - } - if (type == GGML_TYPE_Q8_0) { - err_estimate /= 8.0f; - } - err_estimate *= err_estimate; - err_estimate /= 0.25f*float(ne[0] * r * ne[2]*nr23[0] * ne[3]*nr23[1]); - return err_estimate; + double max_nmse_err(ggml_backend_t backend) override { + // for blackwell we quantize activations to mxfp4 instead of q8_1 so we add higher tolerance + if (type_a == GGML_TYPE_MXFP4 && backend_has_feature(backend, "BLACKWELL_NATIVE_FP4")) { + return 2e-2; } - return 1e-7; + return max_nmse_err(); } -}; - -// GGML_OP_ROPE + GGML_OP_VIEW + GGML_OP_SET_ROWS -struct test_rope_set_rows : public test_case { - const ggml_type type; - const ggml_type type_idx; - const std::array ne_a; - int mode; - const int n_ctx{512}; - const int n_dims{128}; - std::string vars() override { - return VARS_TO_STR4(type, type_idx, ne_a, mode); - } + int64_t grad_nmax() override { return 20000; } - std::string op_desc(ggml_tensor * t) override { + uint64_t op_flops(ggml_tensor * t) override { GGML_UNUSED(t); - return "ROPE_SET_ROWS"; + return 2 * m * n * k * bs[0] * nr[0] * bs[1] * nr[1]; } - bool run_whole_graph() override { return true; } - - test_rope_set_rows(ggml_type type, - ggml_type type_idx, - std::array ne_a, - int mode) - : type(type), type_idx(type_idx), ne_a(ne_a), mode(mode) {} + test_mul_mat(ggml_type type_a = GGML_TYPE_F32, + ggml_type type_b = GGML_TYPE_F32, + int64_t m = 32, + int64_t n = 32, + int64_t k = 32, + std::array bs = {10, 10}, + std::array nr = {2, 2}, + std::array per = {0, 1, 2, 3}, + int64_t k_v = 0, + uint32_t o = 1) : + type_a(type_a), + type_b(type_b), + m(m), + n(n), + k(k), + bs(bs), + nr(nr), + per(per), + k_v(k_v), + o(o) {} ggml_tensor * build_graph(ggml_context * ctx) override { - ggml_tensor * a = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, ne_a[0], ne_a[1], ne_a[2], 1); - ggml_set_name(a, "a"); - - const bool is_mrope = mode & GGML_ROPE_TYPE_MROPE; - const bool is_vision = mode == GGML_ROPE_TYPE_VISION; - - ggml_tensor * pos; - if (is_mrope || is_vision) { - pos = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, ne_a[2] * 4); - } else { - pos = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, ne_a[2]); - } - ggml_set_name(pos, "pos"); - - float fs = 1.4245f; - float ef = 0.7465f; - float af = 1.4245f; - ggml_tensor * freq = nullptr; - - ggml_tensor * rope = nullptr; - if (is_mrope) { - if (is_vision) { - GGML_ASSERT(n_dims/4 > 0); - int rope_sections[4] = {n_dims/4, n_dims/4, 0, 0}; // Vision-RoPE only use first two dimension for image (x, y) coordinate - rope = ggml_rope_multi(ctx, a, pos, freq, n_dims/2, rope_sections, mode, 0, 10000.0f, fs, ef, af, 1.0f, 1.0f); - } else { - GGML_ASSERT(n_dims/3 > 0); - int rope_sections[4] = {n_dims/3, n_dims/3, n_dims/3, 0}; - rope = ggml_rope_multi(ctx, a, pos, freq, n_dims, rope_sections, mode, 0, 10000.0f, fs, ef, af, 1.0f, 1.0f); - } - } else { - rope = ggml_rope(ctx, a, pos, ne_a[0], mode); - } - - ggml_tensor * view = ggml_view_2d(ctx, rope, ne_a[0] * ne_a[1], ne_a[2], rope->nb[2], 0); - - ggml_tensor * dst = ggml_new_tensor_4d(ctx, type, ne_a[0] * ne_a[1], ne_a[2] * ne_a[3], 1, 1); - ggml_set_name(dst, "dst"); - - ggml_tensor * row_idxs = ggml_new_tensor_3d(ctx, type_idx, ne_a[2], 1, 1); - ggml_set_name(row_idxs, "row_idxs"); + // C^T = A * B^T: (k, m) * (k, n) => (m, n) + ggml_tensor * a; + ggml_tensor * b; - ggml_tensor * out = ggml_set_rows(ctx, dst, view, row_idxs); - ggml_set_name(out, "out"); + const int npermuted = (per[0] != 0) + (per[1] != 1) + (per[2] != 2) + (per[3] != 3); + if (npermuted > 0) { + GGML_ASSERT(npermuted == 2); + GGML_ASSERT(k_v == 0); // not handled + GGML_ASSERT(!ggml_is_quantized(type_a) || per[0] == 0); + GGML_ASSERT(!ggml_is_quantized(type_b) || per[0] == 0); - return out; - } + // Create tensors with the permuted dimensions, then permute them back to the dimensions + // given by m,n,k. + const int64_t ne_a[4] = {k, m, bs[0], bs[1]}; + const int64_t ne_b[4] = {k, n, bs[0] * nr[0], bs[1] * nr[1]}; - void initialize_tensors(ggml_context * ctx) override { - for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != NULL; t = ggml_get_next_tensor(ctx, t)) { - if (strcmp(t->name, "row_idxs") == 0) { - if (ggml_is_view_op(t->op)) { - continue; - } - init_set_rows_row_ids(t, ne_a[2]); - } else if (t->type == GGML_TYPE_I32) { - // pos - const int num_pos_ids = (mode & GGML_ROPE_TYPE_MROPE) ? ne_a[2] * 4 : ne_a[2]; - std::vector data(num_pos_ids); - for (int i = 0; i < num_pos_ids; i++) { - data[i] = rand() % n_ctx; - } - ggml_backend_tensor_set(t, data.data(), 0, num_pos_ids * sizeof(int)); - } else { - if (t->ne[0] == n_dims/2) { - // frequency factors in the range [0.9f, 1.1f] - init_tensor_uniform(t, 0.9f, 1.1f); - } else { - init_tensor_uniform(t); + a = ggml_new_tensor_4d(ctx, type_a, ne_a[per[0]], ne_a[per[1]], ne_a[per[2]], + ne_a[per[3]]); + b = ggml_new_tensor_4d(ctx, type_b, ne_b[per[0]], ne_b[per[1]], ne_b[per[2]], + ne_b[per[3]]); + if (!ggml_is_quantized(type_a)) { + if (bs[1] == 1 && nr[1] == 1) { + ggml_set_param(a); } + ggml_set_param(b); } - } - } -}; - -// GGML_OP_RMS_NORM + GGML_OP_MUL + GGML_OP_ROPE (+ GGML_OP_VIEW + GGML_OP_SET_ROWS) -struct test_rms_norm_mul_rope : public test_case { - const std::array ne; - const float eps; - const bool multi_add; // test a sequence of adds feeding into rms_norm - const bool set_rows; - int mode; - - std::string op_desc(ggml_tensor * t) override { - GGML_UNUSED(t); - return "RMS_NORM_MUL_ROPE"; - } - - bool run_whole_graph() override { return true; } - - std::string vars() override { - return VARS_TO_STR5(ne, eps, multi_add, set_rows, mode); - } - - test_rms_norm_mul_rope(std::array ne, float eps = 1e-6f, bool multi_add = false, - bool set_rows = false, int mode = GGML_ROPE_TYPE_NORMAL) - : ne(ne), eps(eps), multi_add(multi_add), set_rows(set_rows), mode(mode) {} - - ggml_tensor * build_graph(ggml_context * ctx) override { - ggml_tensor * a = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, ne[0], ne[1], ne[2], 1); - ggml_tensor * b = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, ne[0], ne[1], ne[2], 1); - ggml_tensor * c = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, ne[0], ne[1], ne[2], 1); - - if (multi_add) { - a = ggml_add(ctx, ggml_add(ctx, a, b), c); - } - - a = ggml_mul(ctx, ggml_rms_norm(ctx, a, eps), b); - - ggml_tensor * pos = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, ne[2]); - - ggml_tensor * rope = ggml_rope(ctx, a, pos, ne[0], mode); - - ggml_tensor * out; - - if (set_rows) { - ggml_tensor * view = ggml_view_2d(ctx, rope, ne[0] * ne[1], ne[2], rope->nb[2], 0); - - ggml_tensor * dst = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, ne[0] * ne[1], ne[2] * ne[3], 1, 1); - ggml_set_name(dst, "dst"); - - ggml_tensor * row_idxs = ggml_new_tensor_3d(ctx, GGML_TYPE_I64, ne[2], 1, 1); - ggml_set_name(row_idxs, "row_idxs"); + ggml_set_name(a, "a"); + ggml_set_name(b, "b"); - out = ggml_set_rows(ctx, dst, view, row_idxs); - ggml_set_name(out, "out"); + a = ggml_permute(ctx, a, per[0], per[1], per[2], per[3]); + b = ggml_permute(ctx, b, per[0], per[1], per[2], per[3]); + ggml_set_name(a, "a_permuted"); + ggml_set_name(b, "b_permuted"); } else { - out = rope; - } - - return out; - } + const int64_t k_physical = k_v == 0 ? k : k_v; + a = ggml_new_tensor_4d(ctx, type_a, k_physical, m, bs[0], bs[1]); + b = ggml_new_tensor_4d(ctx, type_b, k_physical, n, bs[0] * nr[0], bs[1] * nr[1]); - void initialize_tensors(ggml_context * ctx) override { - for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != NULL; t = ggml_get_next_tensor(ctx, t)) { - if (t->type == GGML_TYPE_I64 || t->type == GGML_TYPE_I32) { - if (ggml_is_view_op(t->op)) { - continue; + if (!ggml_is_quantized(type_a)) { + if (bs[1] == 1 && nr[1] == 1) { + ggml_set_param(a); } + ggml_set_param(b); + } - init_set_rows_row_ids(t, ne[2]); - } else { - init_tensor_uniform(t); + if (k_v != 0) { + GGML_ASSERT(k_v > k); + a = ggml_view_4d(ctx, a, k, m, bs[0], bs[1], a->nb[1], a->nb[2], a->nb[3], 0); + b = ggml_view_4d(ctx, b, k, n, bs[0] * nr[0], bs[1] * nr[1], b->nb[1], b->nb[2], + b->nb[3], 0); } + ggml_set_name(a, "a"); + ggml_set_name(b, "b"); } - } -}; -// GGML_OP_ARGMAX -struct test_argmax : public test_case { - const ggml_type type; - const std::array ne; - - std::string vars() override { - return VARS_TO_STR2(type, ne); - } - - test_argmax(ggml_type type = GGML_TYPE_F32, - std::array ne = {10, 100, 1, 1}) - : type(type), ne(ne) {} - - ggml_tensor * build_graph(ggml_context * ctx) override { - ggml_tensor * a = ggml_new_tensor(ctx, type, 4, ne.data()); - ggml_set_name(a, "a"); - - ggml_tensor * out = ggml_argmax(ctx, a); + ggml_tensor * out = ggml_mul_mat(ctx, a, b); ggml_set_name(out, "out"); + for (uint32_t i = 1; i < o; ++i) { + ggml_tensor * out2 = ggml_mul_mat(ctx, a, b); + ggml_set_name(out2, "out2"); + out = ggml_add(ctx, out, out2); + } return out; } - void initialize_tensors(ggml_context * ctx) override { - std::random_device rd; - std::default_random_engine rng(rd()); - for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != NULL; t = ggml_get_next_tensor(ctx, t)) { - if (t->type == GGML_TYPE_F32) { - // initialize with unique values to avoid ties - for (int64_t r = 0; r < ggml_nrows(t); r++) { - std::vector data(t->ne[0]); - for (int i = 0; i < t->ne[0]; i++) { - data[i] = i; - } - std::shuffle(data.begin(), data.end(), rng); - ggml_backend_tensor_set(t, data.data(), r * t->nb[1], t->ne[0] * sizeof(float)); - } - } else { - init_tensor_uniform(t); - } - } - } + bool run_whole_graph() override { return o > 1; } - double max_nmse_err() override { - return 0.0; + std::string op_desc(ggml_tensor * t) override { + GGML_UNUSED(t); + return ggml_op_name(GGML_OP_MUL_MAT); } }; -// GGML_OP_COUNT_EQUAL -struct test_count_equal : public test_case { - const ggml_type type; - const std::array ne; - - std::string vars() override { - return VARS_TO_STR2(type, ne); - } - - test_count_equal(ggml_type type = GGML_TYPE_F32, - std::array ne = {4, 500, 1, 1}) - : type(type), ne(ne) {} - - ggml_tensor * build_graph(ggml_context * ctx) override { - ggml_tensor * a = ggml_new_tensor(ctx, type, 4, ne.data()); - ggml_set_name(a, "a"); - - ggml_tensor * a_argmax = ggml_argmax(ctx, a); - ggml_set_name(a_argmax, "a_argmax"); - - ggml_tensor * b = ggml_new_tensor(ctx, type, 4, ne.data()); - ggml_set_name(b, "b"); - - ggml_tensor * b_argmax = ggml_argmax(ctx, b); - ggml_set_name(b_argmax, "b_argmax"); - - ggml_tensor * out = ggml_count_equal(ctx, a_argmax, b_argmax); - ggml_set_name(out, "out"); - - return out; - } - - double max_nmse_err() override { - return 0.0; - } - - void initialize_tensors(ggml_context * ctx) override { - std::random_device rd; - std::default_random_engine rng(rd()); - for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != NULL; t = ggml_get_next_tensor(ctx, t)) { - if (t->type == GGML_TYPE_F32) { - // initialize with unique values to avoid ties - for (int64_t r = 0; r < ggml_nrows(t); r++) { - std::vector data(t->ne[0]); - for (int i = 0; i < t->ne[0]; i++) { - data[i] = i; - } - std::shuffle(data.begin(), data.end(), rng); - ggml_backend_tensor_set(t, data.data(), r * t->nb[1], t->ne[0] * sizeof(float)); - } - } else { - init_tensor_uniform(t); - } - } - } -}; - -// GGML_OP_REPEAT -struct test_repeat : public test_case { - const ggml_type type; - const std::array ne; - const std::array nr; - - std::string vars() override { - return VARS_TO_STR3(type, ne, nr); - } - - size_t op_size(ggml_tensor * t) override { - return ggml_nbytes(t) * 2; - } - - test_repeat(ggml_type type = GGML_TYPE_F32, - std::array ne = {10, 5, 4, 3}, - std::array nr = {2, 2, 2, 2}) - : type(type), ne(ne), nr(nr) {} - - ggml_tensor * build_graph(ggml_context * ctx) override { - ggml_tensor * target = ggml_new_tensor_4d(ctx, type, ne[0]*nr[0], ne[1]*nr[1], ne[2]*nr[2], ne[3]*nr[3]); - ggml_set_name(target, "target"); - - ggml_tensor * src = ggml_new_tensor(ctx, type, 4, ne.data()); - ggml_set_param(src); - ggml_set_name(src, "src"); - - ggml_tensor * out = ggml_repeat(ctx, src, target); - ggml_set_name(out, "out"); - - return out; - } -}; - -// GGML_OP_REPEAT_BACK -struct test_repeat_back : public test_case { - const ggml_type type; - const std::array ne; - const std::array nr; - const bool v; // whether src is a noncontiguous view - - std::string vars() override { - return VARS_TO_STR4(type, ne, nr, v); - } - - size_t op_size(ggml_tensor * t) override { - return ggml_nbytes(t) * 2; - } - - test_repeat_back(ggml_type type = GGML_TYPE_F32, - std::array ne = {8, 6, 4, 2}, - std::array nr = {2, 2, 2, 2}, - bool v = false) - : type(type), ne(ne), nr(nr), v(v) {} - - ggml_tensor * build_graph(ggml_context * ctx) override { - ggml_tensor * src = ggml_new_tensor_4d(ctx, type, ne[0]*nr[0], ne[1]*nr[1], ne[2]*nr[2], ne[3]*nr[3]); - ggml_set_name(src, "src"); - - if (v) { - GGML_ASSERT(ne[0] % 2 == 0); - GGML_ASSERT(ne[1] % 2 == 0); - GGML_ASSERT(ne[2] % 2 == 0); - GGML_ASSERT(ne[3] % 2 == 0); - GGML_ASSERT(nr[0] % 2 == 0 || nr[0] == 1); - GGML_ASSERT(nr[1] % 2 == 0 || nr[1] == 1); - GGML_ASSERT(nr[2] % 2 == 0 || nr[2] == 1); - GGML_ASSERT(nr[3] % 2 == 0 || nr[3] == 1); - - const int64_t ne00 = nr[0] == 1 ? src->ne[0] : src->ne[0] / 2; - const int64_t ne01 = nr[1] == 1 ? src->ne[1] : src->ne[1] / 2; - const int64_t ne02 = nr[2] == 1 ? src->ne[2] : src->ne[2] / 2; - const int64_t ne03 = nr[3] == 1 ? src->ne[3] : src->ne[3] / 2; - - src = ggml_view_4d(ctx, src, ne00, ne01, ne02, ne03, src->nb[1], src->nb[2], src->nb[3], 0); - } - - ggml_tensor * target = ggml_new_tensor(ctx, type, 4, ne.data()); - ggml_set_name(target, "target"); - - ggml_tensor * out = ggml_repeat_back(ctx, src, target); - ggml_set_name(out, "out"); - - return out; - } -}; - -// GGML_OP_DUP -struct test_dup : public test_case { - const ggml_type type; - const std::array ne; - const std::array permute; - bool _use_permute; - - std::string vars() override { - std::string v = VARS_TO_STR2(type, ne); - if (_use_permute) v += "," + VAR_TO_STR(permute); - return v; - } - - test_dup(ggml_type type = GGML_TYPE_F32, - std::array ne = {10, 10, 20, 1}, - std::array permute = {0, 0, 0, 0}) - : type(type), ne(ne), permute(permute), - _use_permute(permute[0] + permute[1] + permute[2] + permute[3] > 0) {} - - ggml_tensor * build_graph(ggml_context * ctx) override { - ggml_tensor * src = ggml_new_tensor(ctx, type, 4, ne.data()); - ggml_set_param(src); - ggml_set_name(src, "src"); - - if (_use_permute) { - src = ggml_permute(ctx, src, permute[0], permute[1], permute[2], permute[3]); - ggml_set_name(src, "src_permuted"); - } - - ggml_tensor * out = ggml_dup(ctx, src); - ggml_set_name(out, "out"); - - return out; - } -}; - -// GGML_OP_SET -struct test_set : public test_case { - const ggml_type type_src; - const ggml_type type_dst; - const std::array ne; - const int dim; - const bool inplace; - - std::string vars() override { - return VARS_TO_STR5(type_src, type_dst, ne, dim, inplace); - } - - size_t op_size(ggml_tensor * t) override { - return ggml_nbytes(t) + ggml_nbytes(t->src[0]); - } - - test_set(ggml_type type_src = GGML_TYPE_F32, ggml_type type_dst = GGML_TYPE_F32, - std::array ne = {6, 5, 4, 3}, int dim = 1, bool inplace = false) - : type_src(type_src), type_dst(type_dst), ne(ne), dim(dim), inplace(inplace) {} - - ggml_tensor * build_graph(ggml_context * ctx) override { - ggml_tensor * src = ggml_new_tensor(ctx, type_src, 4, ne.data()); - ggml_set_param(src); - ggml_set_name(src, "src"); - - auto ne_dst = ne; - for (int i = 0; i < dim; ++i) { - ne_dst[i] *= 2; - } - ggml_tensor * dst = ggml_new_tensor(ctx, type_dst, 4, ne_dst.data()); - ggml_set_param(dst); - ggml_set_name(dst, "dst"); - - size_t offset = 0; - for (int i = 0; i < dim; ++i) { - offset += ((ne_dst[i] - ne[i])/2)*dst->nb[i]; - } - ggml_tensor * out; - if (inplace) { - out = ggml_set_inplace(ctx, dst, src, - // The backward pass requires setting a contiguous region: - src->nb[1], src->nb[2], src->nb[3], offset); - } else { - out = ggml_set(ctx, dst, src, - // The backward pass requires setting a contiguous region: - src->nb[1], src->nb[2], src->nb[3], offset); - } - ggml_set_name(out, "out"); - - return out; - } -}; - -// GGML_OP_CPY -struct test_cpy : public test_case { - const ggml_type type_src; - const ggml_type type_dst; - const std::array ne; - const std::array permute_src; - const std::array permute_dst; - bool _src_use_permute; - bool _dst_use_permute; - bool _src_transpose; - - std::string vars() override { - return VARS_TO_STR6(type_src, type_dst, ne, permute_src, permute_dst, _src_transpose); - } - - double max_nmse_err() override { - if (type_src == type_dst) { - return 0.0; - } - if (type_dst == GGML_TYPE_Q4_0 || type_dst == GGML_TYPE_Q4_1 || type_dst == GGML_TYPE_IQ4_NL || - type_dst == GGML_TYPE_Q5_0 || type_dst == GGML_TYPE_Q5_1 || type_dst == GGML_TYPE_Q8_0) { - // estimate what the max nmse error would be if one quantized value is - // off by one. The test values are distributed in [-150,150], so it'll be - // roughly (150*2.0 / 2^bits)^2, divided by the mean square value of the reference, - // which is roughly 0.25*150^2 times the number of elements. - double err_estimate = 1.0f/8.0f * 150.0f; - if (type_dst == GGML_TYPE_IQ4_NL) { - // iq4_nl values are a bit more spread out - err_estimate *= 2.0f; - } - if (type_dst == GGML_TYPE_Q5_0 || type_dst == GGML_TYPE_Q5_1) { - err_estimate /= 2.0f; - } - if (type_dst == GGML_TYPE_Q8_0) { - err_estimate /= 8.0f; - } - err_estimate *= err_estimate; - err_estimate /= (150.0f*150.0f*0.25f)*float(ne[0] * ne[1] * ne[2] * ne[3]); - return err_estimate; - } - return 1e-6; - } - - size_t op_size(ggml_tensor * t) override { - return ggml_nbytes(t) + ggml_nbytes(t->src[0]); - } - - test_cpy(ggml_type type_src = GGML_TYPE_F32, ggml_type type_dst = GGML_TYPE_F32, - std::array ne = {10, 10, 10, 1}, - std::array permute_src = {0, 0, 0, 0}, - std::array permute_dst = {0, 0, 0, 0}, - bool transpose_src = false) - : type_src(type_src), type_dst(type_dst), ne(ne), permute_src(permute_src), permute_dst(permute_dst), - _src_use_permute(permute_src[0] + permute_src[1] + permute_src[2] + permute_src[3] > 0), - _dst_use_permute(permute_dst[0] + permute_dst[1] + permute_dst[2] + permute_dst[3] > 0), - _src_transpose(transpose_src){} - - ggml_tensor * build_graph(ggml_context * ctx) override { - ggml_tensor * src = ggml_new_tensor(ctx, type_src, 4, ne.data()); - ggml_set_param(src); - ggml_set_name(src, "src"); - - if (_src_use_permute) { - src = ggml_permute(ctx, src, permute_src[0], permute_src[1], permute_src[2], permute_src[3]); - ggml_set_name(src, "src_permuted"); - } - - if (_src_transpose) { - src = ggml_transpose(ctx, src); - ggml_set_name(src, "src_transposed"); - } - - ggml_tensor * dst = ggml_new_tensor(ctx, type_dst, 4, src->ne); - ggml_set_name(dst, "dst"); - - if (_dst_use_permute) { - dst = ggml_permute(ctx, dst, permute_dst[0], permute_dst[1], permute_dst[2], permute_dst[3]); - ggml_set_name(dst, "dst_permuted"); - } - - ggml_tensor * out = ggml_cpy(ctx, src, dst); - ggml_set_name(out, "out"); - - return out; - } - - void initialize_tensors(ggml_context * ctx) override { - for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != NULL; t = ggml_get_next_tensor(ctx, t)) { - // test extended range of values to check if casting between f32 and i32 is consistent - init_tensor_uniform(t, -150.f, 150.f); - } - } -}; - -// GGML_OP_CONT -struct test_cont : public test_case { - const ggml_type type; - const std::array ne; - bool use_view_slice; - - std::string vars() override { - return VARS_TO_STR3(type, ne, use_view_slice); - } - - test_cont(ggml_type type = GGML_TYPE_F32, - std::array ne = {10, 10, 10, 1}, - bool use_view_slice = false) - : type(type), ne(ne), use_view_slice(use_view_slice) {} - - ggml_tensor * build_graph(ggml_context * ctx) override { - ggml_tensor * src = ggml_new_tensor(ctx, type, 4, ne.data()); - ggml_set_param(src); - ggml_set_name(src, "src"); - - - ggml_tensor * dst; - if (use_view_slice) { - dst = ggml_view_4d(ctx, src, src->ne[0], 1, src->ne[2], src->ne[3], - src->nb[1], src->nb[2], src->nb[3], src->nb[0] * (src->ne[1] - 1)); - ggml_set_name(dst, "src_view_slice"); - } else { - dst = ggml_transpose(ctx, src); - ggml_set_name(dst, "src_transposed"); - } - - ggml_tensor * out = ggml_cont(ctx, dst); - ggml_set_name(out, "out"); - - return out; - } -}; - -// GGML_OP_ADD -// GGML_OP_SUB -// GGML_OP_MUL -// GGML_OP_DIV -struct test_bin_bcast : public test_case { - using op_t = ggml_tensor * (*) (ggml_context *, ggml_tensor *, ggml_tensor *); - op_t op; - const ggml_type type; - const std::array ne; - const std::array nr; - int nf; // number of fused ops, nf == 1 -> single op (no fusion) - bool perm1; // permute src1? - - bool run_whole_graph() override { return nf > 1; } - - std::string vars() override { - return VARS_TO_STR5(type, ne, nr, nf, perm1); - } - - size_t op_size(ggml_tensor * t) override { - return ggml_nbytes(t) * 3; - } - - test_bin_bcast(op_t op, ggml_type type = GGML_TYPE_F32, - std::array ne = {10, 10, 1, 1}, - std::array nr = {1, 2, 1, 1}, - int nf = 1, - bool perm1 = false) - : op(op), type(type), ne(ne), nr(nr), nf(nf), perm1(perm1) {} - - double max_nmse_err(ggml_backend_t backend) override { - // HSA backend converts F16 to BF16, which has lower precision (7-bit mantissa) - if ((type == GGML_TYPE_F16 || type == GGML_TYPE_BF16) && - backend_has_feature(backend, "SUBSTITUTE_FP16_BF16")) { - return 1e-4; // BF16 precision limit - } - return 1e-7; - } - - ggml_tensor * build_graph(ggml_context * ctx) override { - GGML_ASSERT(nf <= 16); - - ggml_tensor * a = ggml_new_tensor_4d(ctx, type, ne[0]*nr[0], ne[1]*nr[1], ne[2]*nr[2], ne[3]*nr[3]); - ggml_set_name(a, "a"); - - ggml_tensor * b[16]; - for (int i = 0; i < nf; ++i) { - if (perm1) { - const int p[4] = { 1, 2, 0, 3 }; // hardcoded for now - - b[i] = ggml_new_tensor_4d(ctx, type, ne[p[0]], ne[p[1]], ne[p[2]], ne[p[3]]); - b[i] = ggml_permute(ctx, b[i], p[0], p[1], p[2], p[3]); - } else { - b[i] = ggml_new_tensor(ctx, type, 4, ne.data()); - } - ggml_set_name(b[i], (std::string("b") + std::to_string(i)).c_str()); - } - - // The backward pass supports broadcasting only for GGML_ADD: - const bool grad_supported = op == ggml_add && ggml_are_same_shape(a, b[0]) && nf == 1 && !perm1; - if (grad_supported) { - ggml_set_param(a); - ggml_set_param(b[0]); - } - - ggml_tensor * out = a; - - for (int i = 0; i < nf; ++i) { - out = op(ctx, out, b[i]); - } - - ggml_set_name(out, "out"); - - return out; - } - - void initialize_tensors(ggml_context * ctx) override { - for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != NULL; t = ggml_get_next_tensor(ctx, t)) { - if (op == ggml_mul || op == ggml_div) { - // MUL and DIV have numerical issues around zero: - init_tensor_uniform(t, 0.9f, 1.1f); - } else { - init_tensor_uniform(t); - } - } - } - - float grad_eps() override { - return 0.1f * (op == ggml_mul ? ne[0]*ne[1]*ne[2]*ne[3] : 1); - } - - bool grad_precise() override { - return op == ggml_div; - } - - double max_maa_err() override { - return op == ggml_add ? 1e-4 : 1e-3; - } -}; - -// GGML_OP_ADD_ID -struct test_add_id : public test_case { - const ggml_type type_a; - const ggml_type type_b; - const int64_t n_embd; - const int64_t n_experts; - const int64_t n_experts_used; - const int64_t n_token; - - std::string vars() override { - return VARS_TO_STR6(type_a, type_b, n_embd, n_experts, n_experts_used, n_token); - } - - size_t op_size(ggml_tensor * t) override { - return ggml_nbytes(t) + ggml_nbytes(t->src[0]) + ggml_nbytes(t->src[2]); - } - - test_add_id(ggml_type type_a = GGML_TYPE_F32, - ggml_type type_b = GGML_TYPE_F32, - int64_t n_embd = 128, - int64_t n_experts = 16, - int64_t n_experts_used = 8, - int64_t n_token = 10) - : type_a(type_a), type_b(type_b), n_embd(n_embd), - n_experts(n_experts), n_experts_used(n_experts_used), n_token(n_token) {} - - ggml_tensor * build_graph(ggml_context * ctx) override { - ggml_tensor * a = ggml_new_tensor_3d(ctx, type_a, n_embd, n_experts_used, n_token); - ggml_tensor * b = ggml_new_tensor_2d(ctx, type_b, n_embd, n_experts); - ggml_tensor * ids = ggml_new_tensor_2d(ctx, GGML_TYPE_I32, n_experts, n_token); - if (n_experts_used != n_experts) { - ids = ggml_view_2d(ctx, ids, n_experts_used, n_token, ids->nb[1], 0); - ggml_set_name(ids, "view_of_ids"); - } - - ggml_tensor * out = ggml_add_id(ctx, a, b, ids); - ggml_set_name(out, "out"); - return out; - } - - void initialize_tensors(ggml_context * ctx) override { - for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != NULL; t = ggml_get_next_tensor(ctx, t)) { - if (t->type == GGML_TYPE_I32) { - if (ggml_is_view_op(t->op)) { continue; } - std::random_device rd; - std::default_random_engine rng(rd()); - // ids - for (int64_t r = 0; r < ggml_nrows(t); r++) { - std::vector data(t->ne[0]); - for (int i = 0; i < t->ne[0]; i++) { - data[i] = i % n_experts; - } - std::shuffle(data.begin(), data.end(), rng); - ggml_backend_tensor_set(t, data.data(), r * t->nb[1], t->ne[0] * sizeof(int32_t)); - } - } else { - init_tensor_uniform(t); - } - } - } -}; - -// GGML_OP_ADD1 -struct test_add1 : public test_case { - const ggml_type type; - const std::array ne; - - std::string vars() override { - return VARS_TO_STR2(type, ne); - } - - test_add1(ggml_type type = GGML_TYPE_F32, - std::array ne = {10, 5, 4, 3}) - : type(type), ne(ne) {} - - ggml_tensor * build_graph(ggml_context * ctx) override { - ggml_tensor * a = ggml_new_tensor(ctx, type, 4, ne.data()); - ggml_set_param(a); - ggml_set_name(a, "a"); - - ggml_tensor * b = ggml_new_tensor_1d(ctx, type, 1); - // ggml_set_param(b); // TODO: implement - ggml_set_name(b, "b"); - - ggml_tensor * out = ggml_add1(ctx, a, b); - ggml_set_name(out, "out"); - - return out; - } - - float grad_eps() override { - return 0.1f * ne[0]*ne[1]*ne[2]*ne[3]; - } -}; - -// GGML_OP_SCALE -struct test_scale : public test_case { - const ggml_type type; - const std::array ne; - float scale; - float bias; - bool inplace; - - std::string vars() override { - return VARS_TO_STR5(type, ne, scale, bias, inplace); - } - - test_scale(ggml_type type = GGML_TYPE_F32, - std::array ne = {10, 10, 10, 10}, - float scale = 2.0f, - float bias = 0.0f, - bool inplace = false) - : type(type), ne(ne), scale(scale), bias(bias), inplace(inplace) {} - - ggml_tensor * build_graph(ggml_context * ctx) override { - ggml_tensor * a = ggml_new_tensor(ctx, type, 4, ne.data()); - ggml_set_param(a); - ggml_set_name(a, "a"); - - ggml_tensor * out; - if (inplace) { - out = ggml_scale_bias_inplace(ctx, a, scale, bias); - } else { - out = ggml_scale_bias(ctx, a, scale, bias); - } - ggml_set_name(out, "out"); - - return out; - } -}; - -// GGML_OP_SCALE + GGML_UNARY_OP_TANH + GGML_OP_SCALE -struct test_softcap : public test_case { - const ggml_type type; - const std::array ne; - float softcap; - - std::string op_desc(ggml_tensor * t) override { - GGML_UNUSED(t); - return "SOFTCAP"; - } - - bool run_whole_graph() override { return true; } - - std::string vars() override { - return VARS_TO_STR3(type, ne, softcap); - } - - test_softcap(ggml_type type = GGML_TYPE_F32, - std::array ne = {10, 10, 10, 10}, - float softcap = 30.0f) - : type(type), ne(ne), softcap(softcap) {} - - ggml_tensor * build_graph(ggml_context * ctx) override { - ggml_tensor * a = ggml_new_tensor(ctx, type, 4, ne.data()); - - ggml_set_param(a); - ggml_set_name(a, "a"); - - ggml_tensor * out = ggml_scale(ctx, ggml_tanh(ctx, ggml_scale(ctx, a, 1.0f / softcap)), softcap); - ggml_set_name(out, "out"); - - return out; - } -}; - -// GGML_OP_SILU_BACK -struct test_silu_back : public test_case { - const ggml_type type; - const std::array ne; - float eps; - - std::string vars() override { - return VARS_TO_STR3(type, ne, eps); - } - - test_silu_back(ggml_type type = GGML_TYPE_F32, - std::array ne = {64, 5, 4, 3}, - float eps = 1e-6f) - : type(type), ne(ne), eps(eps) {} - - ggml_tensor * build_graph(ggml_context * ctx) override { - ggml_tensor * a = ggml_new_tensor(ctx, type, 4, ne.data()); - ggml_set_name(a, "a"); - - ggml_tensor * grad = ggml_new_tensor(ctx, type, 4, ne.data()); - ggml_set_name(grad, "grad"); - - ggml_tensor * out = ggml_silu_back(ctx, a, grad); - ggml_set_name(out, "out"); - - return out; - } - - bool grad_precise() override { - return true; - } -}; - -// GGML_OP_NORM -struct test_norm : public test_case { - const ggml_type type; - const std::array ne; - const bool v; // whether a is a non-contiguous view - const float eps; - - std::string vars() override { - return VARS_TO_STR4(type, ne, v, eps); - } - - test_norm(ggml_type type = GGML_TYPE_F32, - std::array ne = {64, 5, 4, 3}, - bool v = false, - float eps = 1e-6f) - : type(type), ne(ne), v(v), eps(eps) {} - - ggml_tensor * build_graph(ggml_context * ctx) override { - ggml_tensor * a = ggml_new_tensor(ctx, type, 4, ne.data()); - ggml_set_name(a, "a"); - - if (v) { - a = ggml_view_4d(ctx, a, a->ne[0]/2, a->ne[1]/2, a->ne[2]/2, a->ne[3]/2, a->nb[1], a->nb[2], a->nb[3], 0); - ggml_set_name(a, "view of a"); - } - - ggml_tensor * out = ggml_norm(ctx, a, eps); - ggml_set_name(out, "out"); - - return out; - } -}; - -// GGML_OP_NORM + GGML_OP_MUL + GGML_OP_ADD -struct test_norm_mul_add : public test_case { - const ggml_type type; - const std::array ne; - float eps; - const bool broadcast; - - std::string op_desc(ggml_tensor * t) override { - GGML_UNUSED(t); - return "NORM_MUL_ADD"; - } - - bool run_whole_graph() override { return true; } - - std::string vars() override { - return VARS_TO_STR4(type, ne, eps, broadcast); - } - - test_norm_mul_add(ggml_type type = GGML_TYPE_F32, - std::array ne = {128, 2, 1, 1}, - float eps = 1e-5f, - bool broadcast = false) - : type(type), ne(ne), eps(eps), broadcast(broadcast) {} - - ggml_tensor * build_graph(ggml_context * ctx) override { - std::array broadcast_dims = {ne[0], ne[1] * 2, ne[2] * 2, ne[3] * 2}; - - ggml_tensor * a = ggml_new_tensor(ctx, type, 4, broadcast ? broadcast_dims.data() : ne.data()); - ggml_tensor * w = ggml_new_tensor(ctx, type, 4, ne.data()); - ggml_tensor * b = ggml_new_tensor(ctx, type, 4, ne.data()); - ggml_set_param(a); ggml_set_param(w); ggml_set_param(b); - ggml_set_name(a, "a"); ggml_set_name(w, "w"); ggml_set_name(b, "b"); - - // Use a, w and b early to avoid OP_NONE in graph - a = ggml_add(ctx, ggml_add(ctx, a, w), b); - - ggml_tensor * n = ggml_norm(ctx, a, eps); - ggml_tensor * m = ggml_mul(ctx, n, w); - ggml_tensor * out = ggml_add(ctx, m, b); - ggml_set_name(out, "out"); - return out; - } -}; -// GGML_OP_RMS_NORM -struct test_rms_norm : public test_case { - const ggml_type type; - const std::array ne; - const bool v; // whether a is a non-contiguous view - const float eps; - const bool inplace; // whether to do the operation inplace - - std::string vars() override { - return VARS_TO_STR5(type, ne, v, eps, inplace); - } - - test_rms_norm(ggml_type type = GGML_TYPE_F32, - std::array ne = {64, 5, 4, 3}, - bool v = false, - float eps = 1e-6f, - bool inplace = false) - : type(type), ne(ne), v(v), eps(eps), inplace(inplace) {} - - ggml_tensor * build_graph(ggml_context * ctx) override { - ggml_tensor * a = ggml_new_tensor(ctx, type, 4, ne.data()); - ggml_set_param(a); - ggml_set_name(a, "a"); - - if (v) { - a = ggml_view_4d(ctx, a, a->ne[0]/2, a->ne[1]/2, a->ne[2]/2, a->ne[3]/2, a->nb[1], a->nb[2], a->nb[3], 0); - ggml_set_name(a, "view of a"); - } - - ggml_tensor * out; - if (inplace) { - out = ggml_rms_norm_inplace(ctx, a, eps); - } else { - out = ggml_rms_norm(ctx, a, eps); - } - ggml_set_name(out, "out"); - - return out; - } - - void initialize_tensors(ggml_context * ctx) override { - for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != NULL; t = ggml_get_next_tensor(ctx, t)) { - init_tensor_uniform(t, -10.f, 10.f); - } - } - - float grad_eps() override { - return 1.0f; - } - - bool grad_precise() override { - return true; - } -}; - -// GGML_OP_RMS_NORM_BACK -struct test_rms_norm_back : public test_case { - const ggml_type type; - const std::array ne; - const float eps; - - std::string vars() override { - return VARS_TO_STR3(type, ne, eps); - } - - test_rms_norm_back(ggml_type type = GGML_TYPE_F32, - std::array ne = {64, 5, 4, 3}, - float eps = 1e-6f) - : type(type), ne(ne), eps(eps) {} - - ggml_tensor * build_graph(ggml_context * ctx) override { - ggml_tensor * a = ggml_new_tensor(ctx, type, 4, ne.data()); - ggml_set_name(a, "a"); - - ggml_tensor * b = ggml_new_tensor(ctx, type, 4, ne.data()); - ggml_set_name(b, "b"); - - ggml_tensor * out = ggml_rms_norm_back(ctx, a, b, eps); - ggml_set_name(out, "out"); - - return out; - } - - void initialize_tensors(ggml_context * ctx) override { - for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != NULL; t = ggml_get_next_tensor(ctx, t)) { - init_tensor_uniform(t, -10.f, 10.f); - } - } -}; - -// GGML_OP_RMS_NORM + GGML_OP_MUL + GGML_OP_ADD -struct test_rms_norm_mul_add : public test_case { - const ggml_type type; - const std::array ne; - const float eps; - const bool broadcast; - const bool multi_add; // test a sequence of adds feeding into rms_norm - - std::string op_desc(ggml_tensor * t) override { - GGML_UNUSED(t); - return "RMS_NORM_MUL_ADD"; - } - - bool run_whole_graph() override { return true; } - - std::string vars() override { - return VARS_TO_STR5(type, ne, eps, broadcast, multi_add); - } - - test_rms_norm_mul_add(ggml_type type = GGML_TYPE_F32, - std::array ne = {64, 5, 4, 3}, - float eps = 1e-6f, bool broadcast = false, bool multi_add = false) - : type(type), ne(ne), eps(eps), broadcast(broadcast), multi_add(multi_add) {} - - ggml_tensor * build_graph(ggml_context * ctx) override { - std::array broadcast_dims = {ne[0]*2, ne[1]*3, ne[2]*3, ne[3]*4}; - - ggml_tensor * a = ggml_new_tensor(ctx, type, 4, broadcast ? broadcast_dims.data() : ne.data()); - ggml_tensor * b = ggml_new_tensor(ctx, type, 4, ne.data()); - ggml_tensor * c = ggml_new_tensor(ctx, type, 4, ne.data()); - - ggml_set_param(a); - ggml_set_name(a, "a"); - ggml_set_param(b); - ggml_set_name(b, "b"); - ggml_set_param(c); - ggml_set_name(c, "c"); - - // Use a, b and c early, so we don't end up with an OP_NONE between rms_norm and mul - a = ggml_add(ctx, ggml_add(ctx, a, b), c); - if (multi_add) { - a = ggml_add(ctx, ggml_add(ctx, a, b), c); - } - ggml_tensor * out = ggml_add(ctx, ggml_mul(ctx, ggml_rms_norm(ctx, a, eps), b), c); - ggml_set_name(out, "out"); - - return out; - } - - void initialize_tensors(ggml_context * ctx) override { - for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != NULL; t = ggml_get_next_tensor(ctx, t)) { - init_tensor_uniform(t, -10.f, 10.f); - } - } - - float grad_eps() override { - return 1.0f; - } - - bool grad_precise() override { - return true; - } -}; - -// GGML_OP_ADD + GGML_OP_RMS_NORM (fused operation) -struct test_add_rms_norm : public test_case { - const ggml_type type; - const std::array ne; - const float eps; - const bool broadcast; - - std::string op_desc(ggml_tensor * t) override { - GGML_UNUSED(t); - return "ADD_RMS_NORM"; - } - - bool run_whole_graph() override { return true; } - - std::string vars() override { - return VARS_TO_STR4(type, ne, eps, broadcast); - } - - test_add_rms_norm(ggml_type type = GGML_TYPE_F32, - std::array ne = {64, 5, 4, 3}, - float eps = 1e-6f, bool broadcast = false) - : type(type), ne(ne), eps(eps), broadcast(broadcast) {} - - ggml_tensor * build_graph(ggml_context * ctx) override { - std::array broadcast_dims = {ne[0]*2, ne[1]*3, ne[2]*3, ne[3]*4}; - - ggml_tensor * a = ggml_new_tensor(ctx, type, 4, broadcast ? broadcast_dims.data() : ne.data()); - ggml_tensor * b = ggml_new_tensor(ctx, type, 4, ne.data()); - - ggml_set_param(a); - ggml_set_name(a, "a"); - ggml_set_param(b); - ggml_set_name(b, "b"); - - // ADD operation followed by RMS_NORM - ggml_tensor * add_result = ggml_add(ctx, a, b); - ggml_set_name(add_result, "add_result"); - - ggml_tensor * out = ggml_rms_norm(ctx, add_result, eps); - ggml_set_name(out, "out"); - - return out; - } - - void initialize_tensors(ggml_context * ctx) override { - for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != NULL; t = ggml_get_next_tensor(ctx, t)) { - init_tensor_uniform(t, -10.f, 10.f); - } - } - - float grad_eps() override { - return 1.0f; - } - - bool grad_precise() override { - return true; - } -}; - -// GGML_OP_SSM_CONV -struct test_ssm_conv : public test_case { - const ggml_type type; - const std::array ne_a; - const std::array ne_b; - - std::string vars() override { - return VARS_TO_STR3(type, ne_a, ne_b); - } - - test_ssm_conv(ggml_type type = GGML_TYPE_F32, - std::array ne_a = {10, 10, 10, 1}, - std::array ne_b = {3, 3, 1, 1}) - : type(type), ne_a(ne_a), ne_b(ne_b) {} - - ggml_tensor * build_graph(ggml_context * ctx) override { - ggml_tensor * a = ggml_new_tensor(ctx, type, 4, ne_a.data()); - ggml_tensor * b = ggml_new_tensor(ctx, type, 4, ne_b.data()); - ggml_tensor * out = ggml_ssm_conv(ctx, a, b); - return out; - } -}; - -// GGML_OP_SSM_SCAN -struct test_ssm_scan : public test_case { - const ggml_type type; - - const int64_t d_state; - const int64_t head_dim; - const int64_t n_head; - const int64_t n_group; - const int64_t n_seq_tokens; - const int64_t n_seqs; - - std::string vars() override { - return VARS_TO_STR7(type, d_state, head_dim, n_head, n_group, n_seq_tokens, n_seqs); - } - - test_ssm_scan(ggml_type type = GGML_TYPE_F32, - int64_t d_state = 32, - int64_t head_dim = 1, // non-zero for Mamba-2 - int64_t n_head = 32, - int64_t n_group = 1, - int64_t n_seq_tokens = 32, - int64_t n_seqs = 32) - : type(type), d_state(d_state), head_dim(head_dim), n_head(n_head), n_group(n_group), n_seq_tokens(n_seq_tokens), n_seqs(n_seqs) {} - - ggml_tensor * build_graph(ggml_context * ctx) override { - ggml_tensor * s = ggml_new_tensor_4d(ctx, type, d_state, head_dim, n_head, n_seqs); - ggml_tensor * x = ggml_new_tensor_4d(ctx, type, head_dim, n_head, n_seq_tokens, n_seqs); - ggml_tensor * dt = ggml_new_tensor_3d(ctx, type, n_head, n_seq_tokens, n_seqs); - ggml_tensor * A = ggml_new_tensor_2d(ctx, type, (head_dim > 1) ? 1 : d_state, n_head); - ggml_tensor * B = ggml_new_tensor_4d(ctx, type, d_state, n_group, n_seq_tokens, n_seqs); - ggml_tensor * C = ggml_new_tensor_4d(ctx, type, d_state, n_group, n_seq_tokens, n_seqs); - ggml_tensor * ids = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, n_seqs); - ggml_tensor * out = ggml_ssm_scan(ctx, s, x, dt, A, B, C, ids); - return out; - } - - // similar to test_mul_mat_id - void initialize_tensors(ggml_context * ctx) override { - std::random_device rd; - std::default_random_engine rng(rd()); - for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != NULL; t = ggml_get_next_tensor(ctx, t)) { - if (t->type == GGML_TYPE_I32) { - if (ggml_is_view_op(t->op)) { continue; } - // ids - for (int64_t r = 0; r < ggml_nrows(t); r++) { - std::vector data(t->ne[0]); - for (int i = 0; i < t->ne[0]; i++) { - data[i] = i; - } - std::shuffle(data.begin(), data.end(), rng); - ggml_backend_tensor_set(t, data.data(), r * t->nb[1], t->ne[0] * sizeof(int32_t)); - } - } else { - init_tensor_uniform(t); - } - } - } -}; - -// GGML_OP_RWKV_WKV6 -struct test_rwkv_wkv6 : public test_case { - const ggml_type type; - - const int64_t head_count; - const int64_t head_size; - const int64_t n_seq_tokens; - const int64_t n_seqs; - - std::string vars() override { - return VARS_TO_STR5(type, head_count, head_size, n_seq_tokens, n_seqs); - } - - test_rwkv_wkv6(ggml_type type = GGML_TYPE_F32, - int64_t head_count = 32, int64_t head_size = 64, int64_t n_seq_tokens = 32, int64_t n_seqs = 32) - : type(type), head_count(head_count), head_size(head_size), n_seq_tokens(n_seq_tokens), n_seqs(n_seqs) {} - - ggml_tensor * build_graph(ggml_context * ctx) override { - const int64_t n_tokens = n_seq_tokens * n_seqs; - ggml_tensor * r = ggml_new_tensor(ctx, type, 3, std::vector{ head_size, head_count, n_tokens }.data()); - ggml_tensor * k = ggml_new_tensor(ctx, type, 3, std::vector{ head_size, head_count, n_tokens }.data()); - ggml_tensor * v = ggml_new_tensor(ctx, type, 3, std::vector{ head_size, head_count, n_tokens }.data()); - ggml_tensor * tf = ggml_new_tensor(ctx, type, 2, std::vector{ head_size, head_count }.data()); - ggml_tensor * td = ggml_new_tensor(ctx, type, 3, std::vector{ head_size, head_count, n_tokens }.data()); - ggml_tensor * s = ggml_new_tensor(ctx, type, 2, std::vector{ head_size * head_size * head_count, n_seqs }.data()); - ggml_tensor * out = ggml_rwkv_wkv6(ctx, k, v, r, tf, td, s); - return out; - } -}; - -// GGML_OP_GATED_LINEAR_ATTN -struct test_gla : public test_case { - const ggml_type type; - - const int64_t head_count; - const int64_t head_size; - const int64_t n_seq_tokens; - const int64_t n_seqs; - - std::string vars() override { - return VARS_TO_STR5(type, head_count, head_size, n_seq_tokens, n_seqs); - } - - test_gla(ggml_type type = GGML_TYPE_F32, - int64_t head_count = 32, int64_t head_size = 64, int64_t n_seq_tokens = 32, int64_t n_seqs = 32) - : type(type), head_count(head_count), head_size(head_size), n_seq_tokens(n_seq_tokens), n_seqs(n_seqs) {} - - ggml_tensor * build_graph(ggml_context * ctx) override { - const int64_t n_tokens = n_seq_tokens * n_seqs; - ggml_tensor * q = ggml_new_tensor(ctx, type, 3, std::vector{ head_size, head_count, n_tokens }.data()); - ggml_tensor * k = ggml_new_tensor(ctx, type, 3, std::vector{ head_size, head_count, n_tokens }.data()); - ggml_tensor * v = ggml_new_tensor(ctx, type, 3, std::vector{ head_size, head_count, n_tokens }.data()); - ggml_tensor * g = ggml_new_tensor(ctx, type, 3, std::vector{ head_size, head_count, n_tokens }.data()); - ggml_tensor * s = ggml_new_tensor(ctx, type, 2, std::vector{ head_size * head_size * head_count, n_seqs }.data()); - ggml_tensor * out = ggml_gated_linear_attn(ctx, k, v, q, g, s, pow(head_size, -0.5)); - return out; - } -}; - -// GGML_OP_RWKV_WKV7 -struct test_rwkv_wkv7 : public test_case { - const ggml_type type; - - const int64_t head_count; - const int64_t head_size; - const int64_t n_seq_tokens; - const int64_t n_seqs; - - std::string vars() override { - return VARS_TO_STR5(type, head_count, head_size, n_seq_tokens, n_seqs); - } - - test_rwkv_wkv7(ggml_type type = GGML_TYPE_F32, - int64_t head_count = 32, int64_t head_size = 64, int64_t n_seq_tokens = 32, int64_t n_seqs = 32) - : type(type), head_count(head_count), head_size(head_size), n_seq_tokens(n_seq_tokens), n_seqs(n_seqs) {} - - ggml_tensor * build_graph(ggml_context * ctx) override { - const int64_t n_tokens = n_seq_tokens * n_seqs; - ggml_tensor * r = ggml_new_tensor(ctx, type, 3, std::vector{ head_size, head_count, n_tokens }.data()); - ggml_tensor * w = ggml_new_tensor(ctx, type, 3, std::vector{ head_size, head_count, n_tokens }.data()); - ggml_tensor * k = ggml_new_tensor(ctx, type, 3, std::vector{ head_size, head_count, n_tokens }.data()); - ggml_tensor * v = ggml_new_tensor(ctx, type, 3, std::vector{ head_size, head_count, n_tokens }.data()); - ggml_tensor * a = ggml_new_tensor(ctx, type, 3, std::vector{ head_size, head_count, n_tokens }.data()); - ggml_tensor * b = ggml_new_tensor(ctx, type, 3, std::vector{ head_size, head_count, n_tokens }.data()); - // Outputs may become NaN with long seqlen without these normalization - a = ggml_l2_norm(ctx, a, 1e-7F); - b = ggml_l2_norm(ctx, b, 1e-7F); - ggml_tensor * s = ggml_new_tensor(ctx, type, 2, std::vector{ head_size * head_size * head_count, n_seqs }.data()); - ggml_tensor * out = ggml_rwkv_wkv7(ctx, r, w, k, v, a, b, s); - return out; - } -}; - -// GGML_OP_MUL_MAT -struct test_mul_mat : public test_case { - const ggml_type type_a; - const ggml_type type_b; - const int64_t m; - const int64_t n; - const int64_t k; - const std::array bs; // dims 3 and 4 - const std::array nr; // repeat in dims 3 and 4 - const std::array per; // permutation of dimensions - const int64_t k_v; // size of k in memory, resulting in a non-contiguous view for k_v > k, no view for k_v == 0 - const uint32_t o; // number of outputs - - std::string vars() override { - return VARS_TO_STR10(type_a, type_b, m, n, k, bs, nr, per, k_v, o); - } - - double max_nmse_err() override { - return 5e-4; - } - - double max_nmse_err(ggml_backend_t backend) override { - // for blackwell we quantize activations to mxfp4 instead of q8_1 so we add higher tolerance - if (type_a == GGML_TYPE_MXFP4 && backend_has_feature(backend, "BLACKWELL_NATIVE_FP4")) { - return 2e-2; - } - return max_nmse_err(); - } - - int64_t grad_nmax() override { - return 20000; - } - - uint64_t op_flops(ggml_tensor * t) override { - GGML_UNUSED(t); - return 2 * m * n * k * bs[0] * nr[0] * bs[1] * nr[1]; - } - - test_mul_mat(ggml_type type_a = GGML_TYPE_F32, ggml_type type_b = GGML_TYPE_F32, - int64_t m = 32, int64_t n = 32, int64_t k = 32, - std::array bs = {10, 10}, - std::array nr = {2, 2}, - std::array per = {0, 1, 2, 3}, - int64_t k_v = 0, uint32_t o = 1) - : type_a(type_a), type_b(type_b), m(m), n(n), k(k), bs(bs), nr(nr), per(per), k_v(k_v), o(o) {} - - ggml_tensor * build_graph(ggml_context * ctx) override { - // C^T = A * B^T: (k, m) * (k, n) => (m, n) - ggml_tensor * a; - ggml_tensor * b; - - const int npermuted = (per[0] != 0) + (per[1] != 1) + (per[2] != 2) + (per[3] != 3); - if (npermuted > 0) { - GGML_ASSERT(npermuted == 2); - GGML_ASSERT(k_v == 0); // not handled - GGML_ASSERT(!ggml_is_quantized(type_a) || per[0] == 0); - GGML_ASSERT(!ggml_is_quantized(type_b) || per[0] == 0); - - // Create tensors with the permuted dimensions, then permute them back to the dimensions given by m,n,k. - const int64_t ne_a[4] = {k, m, bs[0], bs[1]}; - const int64_t ne_b[4] = {k, n, bs[0]*nr[0], bs[1]*nr[1]}; - - a = ggml_new_tensor_4d(ctx, type_a, ne_a[per[0]], ne_a[per[1]], ne_a[per[2]], ne_a[per[3]]); - b = ggml_new_tensor_4d(ctx, type_b, ne_b[per[0]], ne_b[per[1]], ne_b[per[2]], ne_b[per[3]]); - if (!ggml_is_quantized(type_a)) { - if (bs[1] == 1 && nr[1] == 1) { - ggml_set_param(a); - } - ggml_set_param(b); - } - ggml_set_name(a, "a"); - ggml_set_name(b, "b"); - - a = ggml_permute(ctx, a, per[0], per[1], per[2], per[3]); - b = ggml_permute(ctx, b, per[0], per[1], per[2], per[3]); - ggml_set_name(a, "a_permuted"); - ggml_set_name(b, "b_permuted"); - } else { - const int64_t k_physical = k_v == 0 ? k : k_v; - a = ggml_new_tensor_4d(ctx, type_a, k_physical, m, bs[0], bs[1]); - b = ggml_new_tensor_4d(ctx, type_b, k_physical, n, bs[0]*nr[0], bs[1]*nr[1]); - - if (!ggml_is_quantized(type_a)) { - if (bs[1] == 1 && nr[1] == 1) { - ggml_set_param(a); - } - ggml_set_param(b); - } - - if (k_v != 0) { - GGML_ASSERT(k_v > k); - a = ggml_view_4d(ctx, a, k, m, bs[0], bs[1], a->nb[1], a->nb[2], a->nb[3], 0); - b = ggml_view_4d(ctx, b, k, n, bs[0]*nr[0], bs[1]*nr[1], b->nb[1], b->nb[2], b->nb[3], 0); - } - ggml_set_name(a, "a"); - ggml_set_name(b, "b"); - } - - ggml_tensor * out = ggml_mul_mat(ctx, a, b); - ggml_set_name(out, "out"); - for (uint32_t i = 1; i < o; ++i) { - ggml_tensor * out2 = ggml_mul_mat(ctx, a, b); - ggml_set_name(out2, "out2"); - out = ggml_add(ctx, out, out2); - } - - return out; - } - - bool run_whole_graph() override { return o > 1; } - - std::string op_desc(ggml_tensor * t) override { - GGML_UNUSED(t); - return ggml_op_name(GGML_OP_MUL_MAT); - } -}; - -static void init_mul_mat_id_tensors(ggml_context * ctx, int n_mats) { - std::random_device rd; - std::default_random_engine rng(rd()); - for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != NULL; t = ggml_get_next_tensor(ctx, t)) { - if (t->type == GGML_TYPE_I32) { - if (ggml_is_view_op(t->op)) { continue; } - // ids - for (int64_t r = 0; r < ggml_nrows(t); r++) { - std::vector data(t->ne[0]); - for (int i = 0; i < t->ne[0]; i++) { - data[i] = i % n_mats; - } - std::shuffle(data.begin(), data.end(), rng); - ggml_backend_tensor_set(t, data.data(), r * t->nb[1], t->ne[0] * sizeof(int32_t)); - } - } else { - init_tensor_uniform(t); - } - } -} - -// GGML_OP_MUL_MAT_ID -struct test_mul_mat_id : public test_case { - const ggml_type type_a; - const ggml_type type_b; - const int n_mats; - const int n_used; - const bool b; // broadcast b matrix - const int64_t m; - const int64_t n; - const int64_t k; - - std::string vars() override { - return VARS_TO_STR8(type_a, type_b, n_mats, n_used, b, m, n, k); - } - - double max_nmse_err() override { - return 5e-4; - } - - double max_nmse_err(ggml_backend_t backend) override { - // for blackwell we quantize activations to mxfp4 instead of q8_1 so we add higher tolerance - if (type_a == GGML_TYPE_MXFP4 && backend_has_feature(backend, "BLACKWELL_NATIVE_FP4")) { - return 2e-2; - } - return max_nmse_err(); - } - - uint64_t op_flops(ggml_tensor * t) override { - GGML_UNUSED(t); - return 2 * m * k * n * n_used; - } - - test_mul_mat_id(ggml_type type_a = GGML_TYPE_F32, ggml_type type_b = GGML_TYPE_F32, - int n_mats = 8, int n_used = 2, bool b = false, - int64_t m = 32, int64_t n = 32, int64_t k = 32) - : type_a(type_a), type_b(type_b), n_mats(n_mats), n_used(n_used), b(b), - m(m), n(n), k(k) { - GGML_ASSERT(n_used <= n_mats); - } - - ggml_tensor * build_graph(ggml_context * ctx) override { - // C^T = A * B^T: (k, m) * (k, n) => (m, n) - ggml_tensor * as = ggml_new_tensor_3d(ctx, type_a, k, m, n_mats); - ggml_set_name(as, "as"); - - ggml_tensor * ids = ggml_new_tensor_2d(ctx, GGML_TYPE_I32, n_mats, n); - ggml_set_name(ids, "ids"); - if (n_used != n_mats) { - ids = ggml_view_2d(ctx, ids, n_used, n, ids->nb[1], 0); - ggml_set_name(ids, "view_of_ids"); - } - - ggml_tensor * b = ggml_new_tensor_3d(ctx, type_b, k, this->b ? 1 : n_used, n); - ggml_set_name(b, "b"); - - ggml_tensor * out = ggml_mul_mat_id(ctx, as, b, ids); - ggml_set_name(out, "out"); - - return out; - } - - void initialize_tensors(ggml_context * ctx) override { - init_mul_mat_id_tensors(ctx, n_mats); - } -}; - -// GGML_OP_MUL_MAT_ID + GGML_OP_ADD or GGML_OP_MUL -struct test_mul_mat_id_fusion : public test_case { - const ggml_type type_a; - const ggml_type type_b; - const int n_mats; - const int n_used; - const bool b; // broadcast b matrix - const int64_t m; - const int64_t n; - const int64_t k; - const uint32_t o; // number of outputs - const bool mul; - - std::string vars() override { - return VARS_TO_STR10(type_a, type_b, n_mats, n_used, b, m, n, k, o, mul); - } - - double max_nmse_err() override { - return 5e-4; - } - - uint64_t op_flops(ggml_tensor * t) override { - GGML_UNUSED(t); - return 2 * m * k * n * n_used; - } - - test_mul_mat_id_fusion(ggml_type type_a = GGML_TYPE_F32, ggml_type type_b = GGML_TYPE_F32, - int n_mats = 8, int n_used = 2, bool b = false, - int64_t m = 32, int64_t n = 32, int64_t k = 32, uint32_t o = 1, bool mul = false) - : type_a(type_a), type_b(type_b), n_mats(n_mats), n_used(n_used), b(b), - m(m), n(n), k(k), o(o), mul(mul) { - GGML_ASSERT(n_used <= n_mats); - } - - ggml_tensor * build_graph(ggml_context * ctx) override { - // C^T = A * B^T: (k, m) * (k, n) => (m, n) - ggml_tensor * as = ggml_new_tensor_3d(ctx, type_a, k, m, n_mats); - ggml_set_name(as, "as"); - - ggml_tensor * ids = ggml_new_tensor_2d(ctx, GGML_TYPE_I32, n_mats, n); - ggml_set_name(ids, "ids"); - if (n_used != n_mats) { - ids = ggml_view_2d(ctx, ids, n_used, n, ids->nb[1], 0); - ggml_set_name(ids, "view_of_ids"); - } - - ggml_tensor * b = ggml_new_tensor_3d(ctx, type_b, k, this->b ? 1 : n_used, n); - ggml_set_name(b, "b"); - - ggml_tensor * out = ggml_mul_mat_id(ctx, as, b, ids); - ggml_set_name(out, "out"); - - for (uint32_t i = 1; i < o; ++i) { - ggml_tensor * a2 = ggml_new_tensor_3d(ctx, type_a, k, m, n_mats); - ggml_tensor * out2 = ggml_mul_mat_id(ctx, a2, b, ids); - ggml_set_name(out2, "out2"); - out = ggml_add(ctx, out, out2); - } - - if (mul) { - std::array ne { 1, out->ne[1], out->ne[2], out->ne[3] }; - ne[0] = 1; - ggml_tensor * m = ggml_new_tensor(ctx, out->type, 4, ne.data()); - out = ggml_mul(ctx, out, m); - } - - return out; - } - - void initialize_tensors(ggml_context * ctx) override { - init_mul_mat_id_tensors(ctx, n_mats); - } - - bool run_whole_graph() override { return true; } - - std::string op_desc(ggml_tensor * t) override { - GGML_UNUSED(t); - return "MUL_MAT_ID_FUSION"; - } -}; - -// GGML_OP_OUT_PROD -struct test_out_prod : public test_case { - const ggml_type type_a; - const ggml_type type_b; - const int64_t m; - const int64_t n; - const int64_t k; - const std::array bs; // dims 3 and 4 - const std::array nr; // repeat in dims 3 and 4 - const bool trans_b; - - std::string vars() override { - return VARS_TO_STR8(type_a, type_b, m, n, k, bs, nr, trans_b); - } - - double max_nmse_err() override { - return 5e-4; - } - - test_out_prod(ggml_type type_a = GGML_TYPE_F32, ggml_type type_b = GGML_TYPE_F32, - int64_t m = 32, int64_t n = 32, int64_t k = 32, - std::array bs = {10, 10}, - std::array nr = {2, 2}, - bool trans_b = false) - : type_a(type_a), type_b(type_b), m(m), n(n), k(k), bs(bs), nr(nr), trans_b(trans_b) {} - - ggml_tensor * build_graph(ggml_context * ctx) override { - ggml_tensor * a = ggml_new_tensor_4d(ctx, type_a, m, k, bs[0], bs[1]); - ggml_set_name(a, "a"); - - ggml_tensor * b; - if (trans_b) { - b = ggml_new_tensor_4d(ctx, type_b, k, n, bs[0]*nr[0], bs[1]*nr[1]); - b = ggml_transpose(ctx, b); - } else { - b = ggml_new_tensor_4d(ctx, type_b, n, k, bs[0]*nr[0], bs[1]*nr[1]); - } - ggml_set_name(b, "b"); - - ggml_tensor * out = ggml_out_prod(ctx, a, b); - ggml_set_name(out, "out"); - - return out; - } -}; - -// GGML_OP_SQR -struct test_sqr : public test_case { - const ggml_type type; - const std::array ne; - - std::string vars() override { - return VARS_TO_STR2(type, ne); - } - - test_sqr(ggml_type type = GGML_TYPE_F32, - std::array ne = {10, 5, 4, 3}) - : type(type), ne(ne) {} - - ggml_tensor * build_graph(ggml_context * ctx) override { - ggml_tensor * a = ggml_new_tensor(ctx, type, 4, ne.data()); - ggml_set_param(a); - ggml_set_name(a, "a"); - - ggml_tensor * out = ggml_sqr(ctx, a); - ggml_set_name(out, "out"); - - return out; - } - - float grad_eps() override { - return 0.1f * 0.25f*ne[0]*ne[1]*ne[2]*ne[3]; // 10% of expected value of sum. - } -}; - -// GGML_OP_SQRT -struct test_sqrt : public test_case { - const ggml_type type; - const std::array ne; - - std::string vars() override { - return VARS_TO_STR2(type, ne); - } - - test_sqrt(ggml_type type = GGML_TYPE_F32, - std::array ne = {10, 3, 3, 2}) - : type(type), ne(ne) {} - - ggml_tensor * build_graph(ggml_context * ctx) override { - ggml_tensor * a = ggml_new_tensor(ctx, type, 4, ne.data()); - ggml_set_param(a); - ggml_set_name(a, "a"); - - ggml_tensor * out = ggml_sqrt(ctx, a); - ggml_set_name(out, "out"); - - return out; - } - - void initialize_tensors(ggml_context * ctx) override { - // fill with positive values - for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != NULL; t = ggml_get_next_tensor(ctx, t)) { - init_tensor_uniform(t, 50.0f, 100.0f); - } - } - - float grad_eps() override { - return 20.0f; - } - - bool grad_precise() override { - return true; - } -}; - -// GGML_OP_LOG -struct test_log : public test_case { - const ggml_type type; - const std::array ne; - - std::string vars() override { - return VARS_TO_STR2(type, ne); - } - - test_log(ggml_type type = GGML_TYPE_F32, - std::array ne = {10, 5, 4, 3}) - : type(type), ne(ne) {} - - ggml_tensor * build_graph(ggml_context * ctx) override { - ggml_tensor * a = ggml_new_tensor(ctx, type, 4, ne.data()); - ggml_set_param(a); - ggml_set_name(a, "a"); - - ggml_tensor * out = ggml_log(ctx, a); - ggml_set_name(out, "out"); - - return out; - } - - void initialize_tensors(ggml_context * ctx) override { - for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != NULL; t = ggml_get_next_tensor(ctx, t)) { - // log(1) == 0, cluster values there to keep the sum low for better precision in the backward pass: - init_tensor_uniform(t, 0.9f, 1.1f); - } - } - - bool grad_precise() override { - return true; - } -}; - -// GGML_OP_SIN -struct test_sin : public test_case { - const ggml_type type; - const std::array ne; - - std::string vars() override { - return VARS_TO_STR2(type, ne); - } - - test_sin(ggml_type type = GGML_TYPE_F32, - std::array ne = {10, 2, 2, 2}) - : type(type), ne(ne) {} - - ggml_tensor * build_graph(ggml_context * ctx) override { - ggml_tensor * a = ggml_new_tensor(ctx, type, 4, ne.data()); - ggml_set_param(a); - ggml_set_name(a, "a"); - - ggml_tensor * out = ggml_sin(ctx, a); - ggml_set_name(out, "out"); - - return out; - } - - void initialize_tensors(ggml_context * ctx) override { - for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != NULL; t = ggml_get_next_tensor(ctx, t)) { - init_tensor_uniform(t, -6.5f, 6.5f); // Covers interval [-2*pi, 2*pi]. - } - } - - double max_maa_err() override { - return 1e-3; - } - - float grad_eps() override { - return 0.2f; - } - - bool grad_precise() override { - return true; - } -}; - -// GGML_OP_COS -struct test_cos : public test_case { - const ggml_type type; - const std::array ne; - - std::string vars() override { - return VARS_TO_STR2(type, ne); - } - - test_cos(ggml_type type = GGML_TYPE_F32, - std::array ne = {10, 2, 2, 2}) - : type(type), ne(ne) {} - - ggml_tensor * build_graph(ggml_context * ctx) override { - ggml_tensor * a = ggml_new_tensor(ctx, type, 4, ne.data()); - ggml_set_param(a); - ggml_set_name(a, "a"); - - ggml_tensor * out = ggml_cos(ctx, a); - ggml_set_name(out, "out"); - - return out; - } - - void initialize_tensors(ggml_context * ctx) override { - for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != NULL; t = ggml_get_next_tensor(ctx, t)) { - init_tensor_uniform(t, -6.5f, 6.5f); // Covers interval [-2*pi, 2*pi]. - } - } - - double max_maa_err() override { - return 1e-3; - } - - float grad_eps() override { - return 0.2f; - } - - bool grad_precise() override { - return true; - } -}; - -// GGML_OP_CLAMP -struct test_clamp : public test_case { - const ggml_type type; - const std::array ne; - float min; - float max; - - std::string vars() override { - return VARS_TO_STR4(type, ne, min, max); - } - - test_clamp(ggml_type type = GGML_TYPE_F32, - std::array ne = {10, 5, 4, 3}, - float min = -0.5f, float max = 0.5f) - : type(type), ne(ne), min(min), max(max) {} - - ggml_tensor * build_graph(ggml_context * ctx) override { - ggml_tensor * a = ggml_new_tensor(ctx, type, 4, ne.data()); - ggml_set_name(a, "a"); - - ggml_tensor * out = ggml_clamp(ctx, a, min, max); - ggml_set_name(out, "out"); - - return out; - } - - float grad_eps() override { - return 1e-2f; - } - - std::vector grad_expect() override { - return {0.0f, 1.0f}; - } -}; - -// GGML_OP_FLOOR -struct test_floor : public test_case { - const ggml_type type; - const std::array ne; - - std::string vars() override { - return VARS_TO_STR2(type, ne); - } - - test_floor(ggml_type type = GGML_TYPE_F32, - std::array ne = {10, 2, 2, 2}) - : type(type), ne(ne) {} - - ggml_tensor * build_graph(ggml_context * ctx) override { - ggml_tensor * a = ggml_new_tensor(ctx, type, 4, ne.data()); - ggml_set_param(a); - ggml_set_name(a, "a"); - - ggml_tensor * out = ggml_floor(ctx, a); - ggml_set_name(out, "out"); - - return out; - } - - void initialize_tensors(ggml_context * ctx) override { - for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != NULL; t = ggml_get_next_tensor(ctx, t)) { - init_tensor_uniform(t, -10.0f, 10.0f); - } - } -}; - -// GGML_OP_CEIL -struct test_ceil : public test_case { - const ggml_type type; - const std::array ne; - - std::string vars() override { - return VARS_TO_STR2(type, ne); - } - - test_ceil(ggml_type type = GGML_TYPE_F32, - std::array ne = {10, 2, 2, 2}) - : type(type), ne(ne) {} - - ggml_tensor * build_graph(ggml_context * ctx) override { - ggml_tensor * a = ggml_new_tensor(ctx, type, 4, ne.data()); - ggml_set_param(a); - ggml_set_name(a, "a"); - - ggml_tensor * out = ggml_ceil(ctx, a); - ggml_set_name(out, "out"); - - return out; - } - - void initialize_tensors(ggml_context * ctx) override { - for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != NULL; t = ggml_get_next_tensor(ctx, t)) { - init_tensor_uniform(t, -10.0f, 10.0f); - } - } -}; - -// GGML_OP_ROUND -struct test_round : public test_case { - const ggml_type type; - const std::array ne; - - std::string vars() override { - return VARS_TO_STR2(type, ne); - } - - test_round(ggml_type type = GGML_TYPE_F32, - std::array ne = {10, 2, 2, 2}) - : type(type), ne(ne) {} - - ggml_tensor * build_graph(ggml_context * ctx) override { - ggml_tensor * a = ggml_new_tensor(ctx, type, 4, ne.data()); - ggml_set_param(a); - ggml_set_name(a, "a"); - - ggml_tensor * out = ggml_round(ctx, a); - ggml_set_name(out, "out"); - - return out; - } - - void initialize_tensors(ggml_context * ctx) override { - for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != NULL; t = ggml_get_next_tensor(ctx, t)) { - init_tensor_uniform(t, -10.0f, 10.0f); - } - } -}; - -// GGML_OP_TRUNC -struct test_trunc : public test_case { - const ggml_type type; - const std::array ne; - - std::string vars() override { - return VARS_TO_STR2(type, ne); - } - - test_trunc(ggml_type type = GGML_TYPE_F32, - std::array ne = {10, 2, 2, 2}) - : type(type), ne(ne) {} - - ggml_tensor * build_graph(ggml_context * ctx) override { - ggml_tensor * a = ggml_new_tensor(ctx, type, 4, ne.data()); - ggml_set_param(a); - ggml_set_name(a, "a"); - - ggml_tensor * out = ggml_trunc(ctx, a); - ggml_set_name(out, "out"); - - return out; - } - - void initialize_tensors(ggml_context * ctx) override { - for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != NULL; t = ggml_get_next_tensor(ctx, t)) { - init_tensor_uniform(t, -10.0f, 10.0f); - } - } -}; - -// GGML_OP_DIAG_MASK_INF -struct test_diag_mask_inf : public test_case { - const ggml_type type; - const std::array ne; - const int n_past; - - std::string vars() override { - return VARS_TO_STR3(type, ne, n_past); - } - - test_diag_mask_inf(ggml_type type = GGML_TYPE_F32, - std::array ne = {10, 10, 3, 2}, - int n_past = 5) - : type(type), ne(ne), n_past(n_past) {} - - ggml_tensor * build_graph(ggml_context * ctx) override { - ggml_tensor * a = ggml_new_tensor(ctx, type, 4, ne.data()); - ggml_set_param(a); - ggml_set_name(a, "a"); - - ggml_tensor * out = ggml_diag_mask_inf(ctx, a, n_past); - ggml_set_name(out, "out"); - - return out; - } -}; - -// GGML_OP_SOFT_MAX -struct test_soft_max : public test_case { - const ggml_type type; - const std::array ne; - const bool mask; - const bool sinks; - const ggml_type m_prec; - const std::array nr23; // broadcast only dims 2 and 3 - const float scale; - const float max_bias; - const bool inplace; - - std::string vars() override { - return VARS_TO_STR9(type, ne, mask, sinks, m_prec, nr23, scale, max_bias, inplace); - } - - // the 1024 test with bias occasionally fails: - // SOFT_MAX(type=f32,ne=[1024,16,1,1],mask=1,scale=1.000000,max_bias=8.000000): [SOFT_MAX] NMSE = 0.000000103 > 0.000000100 FAIL - virtual double max_nmse_err() override { - return 1e-6; - } - - test_soft_max(ggml_type type = GGML_TYPE_F32, - std::array ne = {10, 5, 4, 3}, - bool mask = false, - bool sinks = false, - ggml_type m_prec = GGML_TYPE_F32, - std::array nr23 = {1, 1}, - float scale = 1.0f, - float max_bias = 0.0f, - bool inplace = false) - : type(type), ne(ne), mask(mask), sinks(sinks), m_prec(m_prec), nr23(nr23), scale(scale), max_bias(max_bias), inplace(inplace) {} - - ggml_tensor * build_graph(ggml_context * ctx) override { - ggml_tensor * a = ggml_new_tensor_4d(ctx, type, ne[0], ne[1], ne[2]*nr23[0], ne[3]*nr23[1]); - ggml_set_param(a); - ggml_set_name(a, "a"); - - ggml_tensor * mask = nullptr; - if (this->mask) { - mask = ggml_new_tensor_4d(ctx, m_prec, ne[0], ne[1], ne[2], ne[3]); - ggml_set_name(mask, "mask"); - } - - ggml_tensor * sinks = nullptr; - if (this->sinks) { - sinks = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, ne[2]*nr23[0]); - ggml_set_name(sinks, "sinks"); - } - - ggml_tensor * out; - if (inplace) { - out = ggml_soft_max_ext_inplace(ctx, a, mask, scale, max_bias); - } else { - out = ggml_soft_max_ext(ctx, a, mask, scale, max_bias); - } - ggml_soft_max_add_sinks(out, sinks); - ggml_set_name(out, "out"); - - return out; - } - - bool grad_precise() override { - return true; - } -}; - -// GGML_OP_SOFT_MAX_BACK -struct test_soft_max_back : public test_case { - const ggml_type type; - const std::array ne; - const float scale; - const float max_bias; - - std::string vars() override { - return VARS_TO_STR4(type, ne, scale, max_bias); - } - - test_soft_max_back(ggml_type type = GGML_TYPE_F32, - std::array ne = {10, 5, 4, 3}, - float scale = 1.0f, - float max_bias = 0.0f) - : type(type), ne(ne), scale(scale), max_bias(max_bias) {} - - ggml_tensor * build_graph(ggml_context * ctx) override { - ggml_tensor * a = ggml_new_tensor(ctx, type, 4, ne.data()); - ggml_set_name(a, "a"); - - ggml_tensor * b = ggml_new_tensor(ctx, type, 4, ne.data()); - ggml_set_name(a, "a"); - - ggml_tensor * out = ggml_soft_max_ext_back(ctx, a, b, scale, max_bias); - ggml_set_name(out, "out"); - - return out; - } -}; - -// GGML_OP_ROPE + GGML_OP_ROPE_BACK -struct test_rope : public test_case { - const ggml_type type; - const std::array ne_a; - int n_dims; - int mode; - int n_ctx; // used to generate positions - float fs; // freq_scale - float ef; // ext_factor - float af; // attn_factor - bool ff; - int v; // view (1 : non-contiguous a) - bool forward; - bool inplace; - - std::string vars() override { - // forward can be inferred from the op, does not need to be printed - return VARS_TO_STR11(type, ne_a, n_dims, mode, n_ctx, fs, ef, af, ff, v, inplace); - } - - test_rope(ggml_type type = GGML_TYPE_F32, - std::array ne_a = {10, 5, 3, 1}, - int n_dims = 10, int mode = GGML_ROPE_TYPE_NORMAL, int n_ctx = 512, float fs = 1.0f, - float ef = 0.0f, float af = 0.0f, bool ff = false, int v = 0, bool forward = true, bool inplace = false) - : type(type), ne_a(ne_a), n_dims(n_dims), mode(mode), n_ctx(n_ctx), fs(fs), ef(ef), af(af), ff(ff), v(v), forward(forward), inplace(inplace) {} - - ggml_tensor * build_graph(ggml_context * ctx) override { - ggml_tensor * a; - if (v & 1) { - auto ne = ne_a; ne[0] *= 2; ne[1] *= 4; ne[2] *= 3; - a = ggml_new_tensor(ctx, type, 4, ne.data()); - if (forward) { - ggml_set_param(a); - } - ggml_set_name(a, "a"); - - a = ggml_view_4d(ctx, a, ne_a[0], ne_a[1], ne_a[2], ne_a[3], a->nb[1], a->nb[2], a->nb[3], 0); - ggml_set_name(a, "view_of_a"); - } else { - a = ggml_new_tensor(ctx, type, 4, ne_a.data()); - if (forward) { - ggml_set_param(a); - } - ggml_set_name(a, "a"); - } - - const bool is_mrope = mode & GGML_ROPE_TYPE_MROPE; - const bool is_vision = mode == GGML_ROPE_TYPE_VISION; - - ggml_tensor * pos; - if (is_mrope || is_vision) { - pos = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, ne_a[2] * 4); - } else { - pos = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, ne_a[2]); - } - ggml_set_name(pos, "pos"); - - ggml_tensor * freq = nullptr; - if (ff) { - freq = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, n_dims/2); - ggml_set_name(freq, "freq"); - } - - ggml_tensor * out; - if (is_mrope) { - if (is_vision) { - GGML_ASSERT(n_dims/4 > 0); - int rope_sections[4] = {n_dims/4, n_dims/4, 0, 0}; // Vision-RoPE only use first two dimension for image (x, y) coordinate - if (forward) { - if (inplace) { - out = ggml_rope_multi_inplace(ctx, a, pos, freq, n_dims/2, rope_sections, mode, 0, 10000.0f, fs, ef, af, 1.0f, 1.0f); - } else { - out = ggml_rope_multi(ctx, a, pos, freq, n_dims/2, rope_sections, mode, 0, 10000.0f, fs, ef, af, 1.0f, 1.0f); - } - } else { - out = ggml_rope_multi_back(ctx, a, pos, freq, n_dims/2, rope_sections, mode, 0, 10000.0f, fs, ef, af, 1.0f, 1.0f); - } - } else { - GGML_ASSERT(n_dims/3 > 0); - int rope_sections[4] = {n_dims/3, n_dims/3, n_dims/3, 0}; - if (forward) { - if (inplace) { - out = ggml_rope_multi_inplace(ctx, a, pos, freq, n_dims, rope_sections, mode, 0, 10000.0f, fs, ef, af, 1.0f, 1.0f); - } else { - out = ggml_rope_multi(ctx, a, pos, freq, n_dims, rope_sections, mode, 0, 10000.0f, fs, ef, af, 1.0f, 1.0f); - } - } else { - out = ggml_rope_multi_back(ctx, a, pos, freq, n_dims, rope_sections, mode, 0, 10000.0f, fs, ef, af, 1.0f, 1.0f); - } - } - } else { - if (forward) { - if (inplace) { - out = ggml_rope_ext_inplace(ctx, a, pos, freq, n_dims, mode, 0, 10000.0f, fs, ef, af, 1.0f, 1.0f); - } else { - out = ggml_rope_ext(ctx, a, pos, freq, n_dims, mode, 0, 10000.0f, fs, ef, af, 1.0f, 1.0f); - } - } else { - out = ggml_rope_ext_back(ctx, a, pos, freq, n_dims, mode, 0, 10000.0f, fs, ef, af, 1.0f, 1.0f); - } - - // TODO: add test with a non-contiguous view as input ; this case is needed for build_rope_2d in clip.cpp - } - ggml_set_name(out, "out"); - - return out; - } - - void initialize_tensors(ggml_context * ctx) override { - for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != NULL; t = ggml_get_next_tensor(ctx, t)) { - if (t->type == GGML_TYPE_I32) { - // pos - const int num_pos_ids = (mode & GGML_ROPE_TYPE_MROPE) ? ne_a[2] * 4 : ne_a[2]; - std::vector data(num_pos_ids); - for (int i = 0; i < num_pos_ids; i++) { - data[i] = rand() % n_ctx; - } - ggml_backend_tensor_set(t, data.data(), 0, num_pos_ids * sizeof(int)); - } else { - if (t->ne[0] == n_dims/2) { - // frequency factors in the range [0.9f, 1.1f] - init_tensor_uniform(t, 0.9f, 1.1f); - } else { - init_tensor_uniform(t); - } - } - } - } - - double max_maa_err() override { - return 1e-3; - } - - bool grad_precise() override { - return true; - } -}; - -// GGML_OP_POOL2D -struct test_pool2d : public test_case { - enum ggml_op_pool pool_type; - const ggml_type type_input; - const std::array ne_input; - // kernel size - const int k0; - const int k1; - // stride - const int s0; - const int s1; - // padding - const int p0; - const int p1; - - std::string vars() override { - return VARS_TO_STR9(pool_type, type_input, ne_input, k0, k1, s0, s1, p0, p1); - } - - test_pool2d(ggml_op_pool pool_type = GGML_OP_POOL_AVG, - ggml_type type_input = GGML_TYPE_F32, - std::array ne_input = {10, 10, 3, 1}, // [input_width, input_height, input_channels, 1] - int k0 = 3, int k1 = 3, - int s0 = 1, int s1 = 1, - int p0 = 1, int p1 = 1) - : pool_type(pool_type), type_input(type_input), ne_input(ne_input), k0(k0), k1(k1), s0(s0), s1(s1), p0(p0), p1(p1) {} - - ggml_tensor * build_graph(ggml_context * ctx) override { - ggml_tensor * input = ggml_new_tensor(ctx, type_input, 4, ne_input.data()); - ggml_set_param(input); - ggml_set_name(input, "input"); - - ggml_tensor * out = ggml_pool_2d(ctx, input, pool_type, k0, k1, s0, s1, p0, p1); - ggml_set_name(out, "out"); - - return out; - } -}; - -// GGML_OP_POOL1D -struct test_pool1d : public test_case { - enum ggml_op_pool pool_type; - const ggml_type type_input; - const std::array ne_input; - const int k0; - const int s0; - const int p0; - - std::string vars() override { - return VARS_TO_STR6(pool_type, type_input, ne_input, k0, s0, p0); - } - - test_pool1d(ggml_op_pool pool_type = GGML_OP_POOL_AVG, - ggml_type type_input = GGML_TYPE_F32, - std::array ne_input = {10, 1, 1, 1}, - int k0 = 3, int s0 = 3, int p0 = 0) - : pool_type(pool_type), type_input(type_input), ne_input(ne_input), k0(k0), s0(s0), p0(p0) {} - - ggml_tensor * build_graph(ggml_context * ctx) override { - ggml_tensor * input = ggml_new_tensor(ctx, type_input, 4, ne_input.data()); - ggml_set_param(input); - ggml_set_name(input, "input"); - - ggml_tensor * out = ggml_pool_1d(ctx, input, pool_type, k0, s0, p0); - ggml_set_name(out, "out"); - - return out; - } -}; - -// GGML_OP_CONV_TRANSPOSE_1D -struct test_conv_transpose_1d : public test_case { - const std::array ne_input; - const std::array ne_kernel; - - const int s0; // stride - const int p0; // padding - const int d0; // dilation - - std::string vars() override { - return VARS_TO_STR5(ne_input, ne_kernel, s0, p0, d0); - } - - test_conv_transpose_1d(std::array ne_input = {197, 32, 1, 1}, // [input_width, input_channels, 1 /* assert in cpu kernel*/, 1 (should be batch)] - std::array ne_kernel = {16, 32, 32, 1}, // [kernel_width, output_channels, input_channels, 1 (should be batch)] - int s0 = 1, int p0 = 0, int d0 = 1) - : ne_input(ne_input), ne_kernel(ne_kernel), s0(s0), p0(p0), d0(d0) {} - - ggml_tensor * build_graph(ggml_context * ctx) override { - ggml_tensor * input = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, ne_input.data()); - ggml_set_name(input, "input"); - - ggml_tensor * kernel = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, ne_kernel.data()); - ggml_set_name(kernel, "kernel"); - - ggml_tensor * out = ggml_conv_transpose_1d(ctx, kernel, input, s0, p0, d0); - ggml_set_name(out, "out"); - - return out; - } -}; - -// GGML_OP_CONV_TRANSPOSE_2D -struct test_conv_transpose_2d : public test_case { - const std::array ne_input; - const std::array ne_kernel; - const int stride; - - std::string vars() override { - return VARS_TO_STR3(ne_input, ne_kernel, stride); - } - - double max_nmse_err() override { - return 5e-4; // The default 1e-7 is too small for Vulkan. - } - - test_conv_transpose_2d(std::array ne_input = {10, 10, 3, 1}, // [input_width, input_height, input_channels, 1] - std::array ne_kernel = {3, 3, 3, 1}, // [kernel_width, kernel_height, input_channels, 1] - int stride = 1) - : ne_input(ne_input), ne_kernel(ne_kernel), stride(stride){} - - ggml_tensor * build_graph(ggml_context * ctx) override { - ggml_tensor * input = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, ne_input.data()); - ggml_set_name(input, "input"); - - ggml_tensor * kernel = ggml_new_tensor(ctx, GGML_TYPE_F16, 4, ne_kernel.data()); - ggml_set_name(kernel, "kernel"); - - ggml_tensor * out = ggml_conv_transpose_2d_p0(ctx, kernel, input, stride); - ggml_set_name(out, "out"); - - return out; - } -}; - -// GGML_OP_IM2COL -struct test_im2col : public test_case { - const ggml_type type_input; - const ggml_type type_kernel; - const ggml_type dst_type; - const std::array ne_input; - const std::array ne_kernel; - // stride - const int s0; - const int s1; - // padding - const int p0; - const int p1; - // dilation - const int d0; - const int d1; - // mode - const bool is_2D; - - std::string vars() override { - return VARS_TO_STR12(type_input, type_kernel, dst_type, ne_input, ne_kernel, s0, s1, p0, p1, d0, d1, is_2D); - } - - test_im2col(ggml_type type_input = GGML_TYPE_F32, ggml_type type_kernel = GGML_TYPE_F16, ggml_type dst_type = GGML_TYPE_F32, - std::array ne_input = {10, 10, 3, 1}, // [input_width, input_height, input_channels, 1] - std::array ne_kernel = {3, 3, 3, 1}, // [kernel_width, kernel_height, input_channels, 1] - int s0 = 1, int s1 = 1, - int p0 = 1, int p1 = 1, - int d0 = 1, int d1 = 1, - bool is_2D = true) - : type_input(type_input), type_kernel(type_kernel), dst_type(dst_type), ne_input(ne_input), ne_kernel(ne_kernel), s0(s0), s1(s1), p0(p0), p1(p1), d0(d0), d1(d1), is_2D(is_2D) {} - - ggml_tensor * build_graph(ggml_context * ctx) override { - ggml_tensor * input = ggml_new_tensor(ctx, type_input, 4, ne_input.data()); - ggml_set_param(input); - ggml_set_name(input, "input"); - - ggml_tensor * kernel = ggml_new_tensor(ctx, type_kernel, 4, ne_kernel.data()); - ggml_set_name(kernel, "kernel"); - - ggml_tensor * out = ggml_im2col(ctx, kernel, input, s0, s1, p0, p1, d0, d1, is_2D, dst_type); - ggml_set_name(out, "out"); - - return out; - } -}; - -// GGML_OP_IM2COL_3D -struct test_im2col_3d : public test_case { - const ggml_type type_input; - const ggml_type type_kernel; - const ggml_type dst_type; - const std::array ne_input; - const std::array ne_kernel; - // stride - const int s0; - const int s1; - const int s2; - // padding - const int p0; - const int p1; - const int p2; - // dilation - const int d0; - const int d1; - const int d2; - - const int64_t IC; - const bool v; - - std::string vars() override { - return VARS_TO_STR16(type_input, type_kernel, dst_type, ne_input, ne_kernel, IC, s0, s1, s2, p0, p1, p2, d0, d1, d2, v); - } - - test_im2col_3d(ggml_type type_input = GGML_TYPE_F32, ggml_type type_kernel = GGML_TYPE_F16, ggml_type dst_type = GGML_TYPE_F32, - std::array ne_input = {10, 10, 10, 9}, // [OC*IC, KD, KH, KW] - std::array ne_kernel = {3, 3, 3, 1}, // [N*IC, ID, IH, IW] - int64_t IC = 3, - int s0 = 1, int s1 = 1, int s2 = 1, - int p0 = 1, int p1 = 1, int p2 = 1, - int d0 = 1, int d1 = 1, int d2 = 1, - bool v = false) - : type_input(type_input), type_kernel(type_kernel), dst_type(dst_type), ne_input(ne_input), ne_kernel(ne_kernel), s0(s0), s1(s1), s2(s2), p0(p0), p1(p1), p2(p2), d0(d0), d1(d1), d2(d2), IC(IC), v(v) {} - - ggml_tensor * build_graph(ggml_context * ctx) override { - ggml_tensor * input = ggml_new_tensor(ctx, type_input, 4, ne_input.data()); - ggml_set_param(input); - ggml_set_name(input, "input"); - - if (v) { - input = ggml_view_4d(ctx, input, ne_input[0] - 2, ne_input[1] - 2, ne_input[2] - 2, ne_input[3] - 2, input->nb[1], input->nb[2], input->nb[3], 0); - ggml_set_name(input, "view_of_input"); - } - - ggml_tensor * kernel = ggml_new_tensor(ctx, type_kernel, 4, ne_kernel.data()); - ggml_set_name(kernel, "kernel"); - - ggml_tensor * out = ggml_im2col_3d(ctx, kernel, input, IC, s0, s1, s2, p0, p1, p2, d0, d1, d2, dst_type); - ggml_set_name(out, "out"); - - return out; - } -}; - -// CONV_2D -struct test_conv_2d : public test_case { - const std::array ne_input; - const std::array ne_kernel; - const ggml_type type_kernel; - const int stride0; - const int stride1; - const int padding0; - const int padding1; - const int dilation0; - const int dilation1; - // Whether the inputs are contiguous in the channel dim or the width dim - const bool cwhn; - - // If true, the direct CONV_2D will be used in the graph, otherwise it - // uses ggml_conv_2d: - // * if the program is called with -o CONV_2D_DIRECT_IMPL, the - // CONV_2D graph will be built, while - // * if the program is called with -o CONV_2D_INDIRECT_IMPL, the - // IM2COL -> MUL_MM graph will be built. - - std::string vars() override { - return VARS_TO_STR10(ne_input, ne_kernel, type_kernel, stride0, stride1, padding0, padding1, dilation0, dilation1, cwhn); - } - - double max_nmse_err() override { - return 5e-4; - } - - uint64_t op_flops(ggml_tensor * t) override { - GGML_UNUSED(t); - // Just counting matmul costs: - // KxCRS @ CRSxNPQ = KxNPQ --> KxNPQx(CRS+CRS-1) flops - - // Copied from ggml.c: int64_t ggml_calc_conv_output_size(int64_t ins, int64_t ks, int s, int p, int d) - auto calc_conv_output_size = [](int64_t ins, int64_t ks, int s, int p, int d) -> int64_t { - return (ins + 2 * p - d * (ks - 1) - 1) / s + 1; - }; - - int64_t W = ne_input[0]; - int64_t H = ne_input[1]; - int64_t KW = ne_kernel[0]; - int64_t KH = ne_kernel[1]; - int64_t Cin = ne_kernel[2]; - int64_t Cout = ne_kernel[3]; - int64_t N = ne_input[3]; - int64_t OH = calc_conv_output_size(H, KH, stride0, padding0, dilation0); - int64_t OW = calc_conv_output_size(W, KW, stride0, padding0, dilation0); - - int64_t K = Cout; - int64_t CRS = Cin * KH * KW; - int64_t NPQ = N * OH * OW; - - return K * NPQ * (2 * CRS - 1); - } - - test_conv_2d(std::array ne_input = { 64, 64, 16, 1 }, - std::array ne_kernel = { 3, 3, 1, 16 }, ggml_type type_kernel = GGML_TYPE_F32, int stride0 = 1, - int stride1 = 1, int padding0 = 0, int padding1 = 0, int dilation0 = 1, int dilation1 = 1, bool cwhn = false) : - ne_input(ne_input), - ne_kernel(ne_kernel), - type_kernel(type_kernel), - stride0(stride0), - stride1(stride1), - padding0(padding0), - padding1(padding1), - dilation0(dilation0), - dilation1(dilation1), - cwhn(cwhn) {} - - ggml_tensor * build_graph(ggml_context * ctx) override { - ggml_tensor * input = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, ne_input.data()); - ggml_set_name(input, "input"); - - ggml_tensor * kernel = ggml_new_tensor(ctx, type_kernel, 4, ne_kernel.data()); - ggml_set_name(kernel, "kernel"); - - if (cwhn) { - // change memory layout to channel-most-contiguous (CWHN), - // then permute it back so NE matches the original input - input = ggml_cont(ctx, ggml_permute(ctx, input, 1, 2, 0, 3)); - input = ggml_permute(ctx, input, 2, 0, 1, 3); - kernel = ggml_cont(ctx, ggml_permute(ctx, kernel, 2, 3, 1, 0)); - kernel = ggml_permute(ctx, kernel, 3, 2, 0, 1); - } - - ggml_tensor * out = - ggml_conv_2d_direct(ctx, kernel, input, stride0, stride1, padding0, padding1, dilation0, dilation1); - ggml_set_name(out, "out"); - return out; - } -}; - -// GGML_OP_CONV_2D_DW -struct test_conv_2d_dw : public test_case { - const std::array ne_input; - const std::array ne_kernel; - const int stride; - const int padding; - const int dilation; - const bool cwhn; - - std::string vars() override { - return VARS_TO_STR6(ne_input, ne_kernel, stride, padding, dilation, cwhn); - } - - test_conv_2d_dw(std::array ne_input = {64, 64, 16, 1}, - std::array ne_kernel = {3, 3, 1, 16}, - int stride = 1, int padding = 0, int dilation = 1, bool cwhn = false) - : ne_input(ne_input), ne_kernel(ne_kernel), stride(stride), padding(padding), dilation(dilation), cwhn(cwhn) {} - - ggml_tensor * build_graph(ggml_context * ctx) override { - ggml_tensor * input = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, ne_input.data()); - ggml_set_name(input, "input"); - - ggml_tensor * kernel = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, ne_kernel.data()); - ggml_set_name(kernel, "kernel"); - - if (cwhn) { - // change memory layout to channel-most-contiguous (CWHN), - // then permute it back so NE matches the original input - input = ggml_cont(ctx, ggml_permute(ctx, input, 1, 2, 0, 3)); - input = ggml_permute(ctx, input, 2, 0, 1, 3); - kernel = ggml_cont(ctx, ggml_permute(ctx, kernel, 2, 3, 1, 0)); - kernel = ggml_permute(ctx, kernel, 3, 2, 0, 1); - } - - ggml_tensor * out = ggml_conv_2d_dw_direct( - ctx, kernel, input, - stride, stride, padding, padding, dilation, dilation); - ggml_set_name(out, "out"); - return out; - } -}; - -// GGML_OP_CONV_3D -struct test_conv_3d : public test_case { - // Logical 5D dimensions - const int64_t N, IC, ID, IH, IW; - const int64_t OC, KD, KH, KW; - // Conv params - const int s0, s1, s2; - const int p0, p1, p2; - const int d0, d1, d2; - // Types - const ggml_type type_kernel; - - std::string op_desc(ggml_tensor * t) override { - GGML_UNUSED(t); - return "CONV_3D"; - } - - std::string vars() override { - return VARS_TO_STR11(N, IC, ID, IH, IW, OC, KD, KH, KW, s0, s1) + "," + - VARS_TO_STR8(s2, p0, p1, p2, d0, d1, d2, type_kernel); - } - - double max_nmse_err() override { - return 5e-4; - } - - uint64_t op_flops(ggml_tensor * t) override { - GGML_UNUSED(t); - auto calc_conv_output_size = [](int64_t ins, int64_t ks, int s, int p, int d) -> int64_t { - return (ins + 2 * p - d * (ks - 1) - 1) / s + 1; - }; - const int64_t OD = calc_conv_output_size(ID, KD, s2, p2, d2); - const int64_t OH = calc_conv_output_size(IH, KH, s1, p1, d1); - const int64_t OW = calc_conv_output_size(IW, KW, s0, p0, d0); - - return (uint64_t)N * OC * OD * OH * OW * (2 * IC * KD * KH * KW - 1); - } - - test_conv_3d( - int64_t N, int64_t IC, int64_t ID, int64_t IH, int64_t IW, - int64_t OC, int64_t KD, int64_t KH, int64_t KW, - int s0, int s1, int s2, - int p0, int p1, int p2, - int d0, int d1, int d2, - ggml_type type_kernel - ) : N(N), IC(IC), ID(ID), IH(IH), IW(IW), - OC(OC), KD(KD), KH(KH), KW(KW), - s0(s0), s1(s1), s2(s2), - p0(p0), p1(p1), p2(p2), - d0(d0), d1(d1), d2(d2), - type_kernel(type_kernel) {} - - ggml_tensor * build_graph(ggml_context * ctx) override { - // GGML input tensor is packed as [W, H, D, C*N] - const int64_t ne_input[] = {IW, IH, ID, IC * N}; - ggml_tensor * input = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, ne_input); - ggml_set_name(input, "input"); - - // GGML kernel tensor is packed as [KW, KH, KD, IC*OC] - const int64_t ne_kernel[] = {KW, KH, KD, IC * OC}; - ggml_tensor * kernel = ggml_new_tensor(ctx, type_kernel, 4, ne_kernel); - ggml_set_name(kernel, "kernel"); - - ggml_tensor * out = ggml_conv_3d_direct(ctx, kernel, input, s0, s1, s2, p0, p1, p2, d0, d1, d2, (int)IC, (int)N, (int)OC); - ggml_set_name(out, "out"); - return out; - } -}; - -// GGML_OP_CONCAT -struct test_concat : public test_case { - const ggml_type type; - const std::array ne_a; - const int64_t ne_b_d; - const int dim; - const int v; // view (1 << 0: non-cont a, 1 << 1: non-cont b) - - std::string vars() override { - return VARS_TO_STR5(type, ne_a, ne_b_d, dim, v); - } - - test_concat(ggml_type type = GGML_TYPE_F32, - std::array ne_a = {10, 5, 5, 5}, - int64_t ne_b_d = 5, - int dim = 2, int v = 0) - : type(type), ne_a(ne_a), ne_b_d(ne_b_d), dim(dim), v(v) {} - - ggml_tensor * build_graph(ggml_context * ctx) override { - auto ne_b = ne_a; - ne_b[dim] = ne_b_d; - ggml_tensor * a; - if (v & 1) { - auto ne = ne_a; ne[0] *= 2; ne[1] *= 4; ne[2] *= 3; - a = ggml_new_tensor(ctx, type, 4, ne.data()); - ggml_set_name(a, "a"); - - a = ggml_view_4d(ctx, a, ne_a[0], ne_a[1], ne_a[2], ne_a[3], a->nb[1], a->nb[2], a->nb[3], 0); - ggml_set_name(a, "view_of_a"); - } else { - a = ggml_new_tensor(ctx, type, 4, ne_a.data()); - ggml_set_name(a, "a"); - } - ggml_tensor * b; - if (v & 2) { - auto ne = ne_b; ne[0] *= 3; ne[1] *= 2; ne[2] *= 4; - b = ggml_new_tensor(ctx, type, 4, ne.data()); - ggml_set_name(b, "b"); - - b = ggml_view_4d(ctx, b, ne_b[0], ne_b[1], ne_b[2], ne_b[3], b->nb[1], b->nb[2], b->nb[3], 0); - ggml_set_name(b, "view_of_b"); - } else { - b = ggml_new_tensor(ctx, type, 4, ne_b.data()); - ggml_set_name(b, "b"); - } - - ggml_tensor * out = ggml_concat(ctx, a, b, dim); - ggml_set_name(out, "out"); - - return out; - } -}; - -// GGML_OP_ARGSORT -struct test_argsort : public test_case { - const ggml_type type; - const std::array ne; - ggml_sort_order order; - - std::string vars() override { - return VARS_TO_STR3(type, ne, order); - } - - test_argsort(ggml_type type = GGML_TYPE_F32, - std::array ne = {16, 10, 10, 10}, - ggml_sort_order order = GGML_SORT_ORDER_ASC) - : type(type), ne(ne), order(order) {} - - ggml_tensor * build_graph(ggml_context * ctx) override { - ggml_tensor * a = ggml_new_tensor(ctx, type, 4, ne.data()); - ggml_set_name(a, "a"); - - ggml_tensor * out = ggml_argsort(ctx, a, order); - ggml_set_name(out, "out"); - - return out; - } - - void initialize_tensors(ggml_context * ctx) override { - std::random_device rd; - std::default_random_engine rng(rd()); - for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != NULL; t = ggml_get_next_tensor(ctx, t)) { - if (t->type == GGML_TYPE_I32) { - // indices - std::vector data(ggml_nelements(t)); - for (int i = 0; i < ggml_nelements(t); i++) { - data[i] = rand(); - } - std::shuffle(data.begin(), data.end(), rng); - ggml_backend_tensor_set(t, data.data(), 0, ne[0]*ne[1]*ne[2]*ne[3] * sizeof(int)); - } else if (t->type == GGML_TYPE_F32) { - // initialize with unique values to avoid ties - for (int64_t r = 0; r < ggml_nrows(t); r++) { - std::vector data(t->ne[0]); - for (int i = 0; i < t->ne[0]; i++) { - data[i] = i; - } - std::shuffle(data.begin(), data.end(), rng); - ggml_backend_tensor_set(t, data.data(), r * t->nb[1], t->ne[0] * sizeof(float)); - } - } else { - GGML_ABORT("fatal error"); - } - } - } -}; - -// GGML_OP_TOP_K -struct test_top_k : public test_case { - const ggml_type type; - const std::array ne; - const int k; - const bool ties; - ggml_tensor * input {}; - - std::string vars() override { - return VARS_TO_STR4(type, ne, k, ties); - } - - test_top_k(ggml_type type = GGML_TYPE_F32, - std::array ne = {16, 10, 10, 10}, - int k = 4, bool ties = false) - : type(type), ne(ne), k(k), ties(ties) {} - - double max_err() override { - return 0.0; - } - - // When there are ties, only validate the final result. - // The logic in err can't handle the sentinel tensors. - bool run_whole_graph() override { return ties; } - - double err(const float * a, const float * b, size_t n) override { - // When there are no ties, we expect the exact same set of indices, - // but possibly in a different order. When there are ties, the indices - // can be different but the input values they correspond to should be - // the same. The logic for ties could work for non-ties, but only for - // the output tensor, not for the sentinel tensors. - if (ties) { - std::vector src(ggml_nelements(input)); - - ggml_backend_tensor_get(input, src.data(), 0, ggml_nelements(input) * ggml_type_size(type)); - - double diff = 0.0f; - - GGML_ASSERT(n == (size_t)(ggml_nrows(input) * k)); - int64_t cols = input->ne[0]; - std::vector ia(k); - std::vector ib(k); - std::vector asrc(k); - std::vector bsrc(k); - for (int64_t r = 0; r < ggml_nrows(input); r++) { - // Convert indices for the row back to integer - for (int64_t c = 0; c < k; c++) { - ia[c] = (int32_t)a[r * k + c]; - ib[c] = (int32_t)b[r * k + c]; - } - // The src values for each row should match. - for (int64_t c = 0; c < k; c++) { - asrc[c] = src[r * cols + ia[c]]; - bsrc[c] = src[r * cols + ib[c]]; - } - diff += jdst(asrc.data(), bsrc.data(), k); - // There should be no duplicate indices - std::sort(ia.begin(), ia.end()); - std::sort(ib.begin(), ib.end()); - if (std::adjacent_find(ia.begin(), ia.end()) != ia.end()) { - diff += 1; - } - if (std::adjacent_find(ib.begin(), ib.end()) != ib.end()) { - diff += 1; - } - } - return diff; - } else { - std::vector ia(n); - std::vector ib(n); - - double diff = 0.0f; - - for (size_t i = 0; i < n; i++) { - ia[i] = (int32_t) a[i]; - ib[i] = (int32_t) b[i]; - - // penalize the result if the data is not integer valued - diff += std::fabs(a[i] - ia[i]); - diff += std::fabs(b[i] - ib[i]); - } - - return diff + jdst(ia.data(), ib.data(), n); - } - } - - ggml_tensor * build_graph(ggml_context * ctx) override { - ggml_tensor * a = ggml_new_tensor(ctx, type, 4, ne.data()); - ggml_set_name(a, "a"); - - // Save 'a' for err() - input = a; - - ggml_tensor * out = ggml_top_k(ctx, a, k); - ggml_set_name(out, "out"); - - return out; - } - - void initialize_tensors(ggml_context * ctx) override { - std::random_device rd; - std::default_random_engine rng(rd()); - for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != NULL; t = ggml_get_next_tensor(ctx, t)) { - int tie_denom = std::max(1, std::min(10, k / 2)); - for (int64_t r = 0; r < ggml_nrows(t); r++) { - std::vector data(t->ne[0]); - for (int i = 0; i < t->ne[0]; i++) { - if (ties) { - // integer division to introduce duplicates - data[i] = i / tie_denom; - } else { - data[i] = i; - } - } - std::shuffle(data.begin(), data.end(), rng); - ggml_backend_tensor_set(t, data.data(), r * t->nb[1], t->ne[0] * sizeof(float)); - } - } - } -}; - -enum MoeGatingFunc { - GATING_FUNC_SOFTMAX, - GATING_FUNC_SIGMOID, - GATING_FUNC_SOFTMAX_WEIGHT, -}; - -struct test_topk_moe : public test_case { - const std::array ne; - const int n_expert_used; - const bool with_norm; - const bool bias_probs; - const MoeGatingFunc gating_func; - const float scale_w; - ggml_tensor * weights {}; - ggml_tensor * selected_experts {}; - - test_topk_moe(std::array ne = { 10, 5, 1, 1 }, - int n_expert_used = 1, - bool with_norm = false, - bool bias_probs = false, - MoeGatingFunc gating_func = GATING_FUNC_SOFTMAX, - float scale_w = 0.0f) : - ne(ne), - n_expert_used(n_expert_used), - with_norm(with_norm), - bias_probs(bias_probs), - gating_func(gating_func), - scale_w(scale_w) { - GGML_ASSERT(n_expert_used <= ne[0]); - } - - std::string vars() override { return VARS_TO_STR6(ne, n_expert_used, with_norm, bias_probs, gating_func, scale_w); } - - std::string op_desc(ggml_tensor * t) override { - GGML_UNUSED(t); - return "TOPK_MOE"; - } - - bool run_whole_graph() override { return true; } - - ggml_tensor * build_graph(ggml_context * ctx) override { - const int n_expert = ne[0]; - const int n_tokens = ne[1]; - - ggml_tensor * logits = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, ne.data()); - ggml_tensor * probs = - (gating_func == GATING_FUNC_SOFTMAX) ? ggml_soft_max(ctx, logits) : - (gating_func == GATING_FUNC_SIGMOID) ? ggml_sigmoid(ctx, logits) : logits; - ggml_set_name(probs, "probs"); - - ggml_tensor * selection_probs = probs; - if (bias_probs) { - ggml_tensor * exp_probs_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, ne[0]); - ggml_set_name(exp_probs_b, "exp_probs_b"); - selection_probs = ggml_add(ctx, probs, exp_probs_b); - ggml_set_name(selection_probs, "selection_probs"); - } - - selected_experts = ggml_argsort_top_k(ctx, selection_probs, n_expert_used); // [n_expert_used, n_tokens] - ggml_set_name(selected_experts, "selected_experts"); - - weights = ggml_get_rows(ctx, ggml_reshape_3d(ctx, probs, 1, n_expert, n_tokens), selected_experts); // [1, n_expert_used, n_tokens] - ggml_set_name(weights, "weights"); - - if (gating_func == GATING_FUNC_SOFTMAX_WEIGHT) { - weights = ggml_reshape_2d(ctx, weights, n_expert_used, n_tokens); - weights = ggml_soft_max(ctx, weights); // [n_expert_used, n_tokens] - weights = ggml_reshape_3d(ctx, weights, 1, n_expert_used, n_tokens); - } - - if (with_norm) { - weights = ggml_reshape_2d(ctx, weights, n_expert_used, n_tokens); - ggml_tensor * weights_sum = ggml_sum_rows(ctx, weights); // [1, n_tokens] - ggml_set_name(weights_sum, "weights_sum"); - - weights_sum = ggml_clamp(ctx, weights_sum, 6.103515625e-5, INFINITY); - weights = ggml_div(ctx, weights, weights_sum); // [n_expert_used, n_tokens] - weights = ggml_reshape_3d(ctx, weights, 1, n_expert_used, n_tokens); - } - - if (scale_w) { - weights = ggml_scale(ctx, weights, scale_w); - } - - ggml_set_name(weights, "weights"); - return weights; - } - // Verify two outputs - std::vector fusion_test_nodes() override { return { selected_experts, weights }; } - - // allow output in arbitrary order - double err(const float * a, const float * b, size_t n) override { - std::vector a2(n); - std::vector b2(n); - for (size_t i = 0; i < n; ++i) { - a2[i] = a[i]; - b2[i] = b[i]; - } - std::sort(a2.begin(), a2.end()); - std::sort(b2.begin(), b2.end()); - return nmse(a2.data(), b2.data(), n); - } -}; - -struct test_mul_mat_vec_fusion : public test_case { - const ggml_type type; - const ggml_glu_op glu_op; - const int64_t m; - const int64_t n; - const int64_t k; - const bool use_id; - const int n_mats; - const int n_used; - const bool b; // broadcast b matrix (only for use_id) - const bool with_bias; - const bool with_gate; - std::array batch_dims; - - test_mul_mat_vec_fusion(ggml_type type, ggml_glu_op op, int64_t m, int64_t n, int64_t k, - bool use_id = false, int n_mats = 1, int n_used = 1, bool b = false, bool with_bias = false, bool with_gate = true, - std::array batch_dims = {4, 2}) - : type(type), glu_op(op), m(m), n(n), k(k), use_id(use_id), n_mats(n_mats), n_used(n_used), b(b), with_bias(with_bias), with_gate(with_gate), batch_dims(batch_dims) { - if (use_id) { - GGML_ASSERT(n_used <= n_mats); - } - } - - std::string vars() override { - return VARS_TO_STR12(type, glu_op, m, n, k, use_id, n_mats, n_used, b, with_bias, with_gate, batch_dims); - } - - std::string op_desc(ggml_tensor * t) override { - GGML_UNUSED(t); - return "MUL_MAT_VEC_FUSION"; - } - - bool run_whole_graph() override { return true; } - - ggml_tensor * build_gate(ggml_context * ctx, ggml_tensor * ffn_gate, ggml_tensor * ffn_up) { - ggml_tensor * out = nullptr; - if (with_gate) { - if (glu_op == GGML_GLU_OP_SWIGLU_OAI) { - constexpr float alpha = 1.702f; - constexpr float limit = 7.0f; - out = ggml_swiglu_oai(ctx, ffn_gate, ffn_up, alpha, limit); - } else { - out = ggml_glu_split(ctx, ffn_gate, ffn_up, glu_op); - } - } - return out; - } - - ggml_tensor * build_graph(ggml_context * ctx) override { - if (!use_id) { - const int channels = batch_dims[0]; - const int samples = batch_dims[1]; - std::array ne = { k, m, channels, samples }; - std::array ne0 = { k, n, channels, samples }; - - ggml_tensor * cur = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, ne.data()); - ggml_tensor * gate = with_gate ? ggml_new_tensor(ctx, type, 4, ne0.data()) : nullptr; - ggml_tensor * up = ggml_new_tensor(ctx, type, 4, ne0.data()); - - ggml_tensor * ffn_up = ggml_mul_mat(ctx, up, cur); - if (with_bias) { - std::array bias_ne = { ffn_up->ne[0], 1, channels, samples }; - ggml_tensor * up_bias = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, bias_ne.data()); - ffn_up = ggml_add(ctx, ffn_up, up_bias); - } - - ggml_tensor * ffn_gate = with_gate ? ggml_mul_mat(ctx, gate, cur) : nullptr; - if (with_bias && with_gate) { - std::array bias_ne = { ffn_gate->ne[0], 1, channels, samples }; - ggml_tensor * gate_bias = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, bias_ne.data()); - ffn_gate = ggml_add(ctx, ffn_gate, gate_bias); - } - - ggml_tensor * out = with_gate ? build_gate(ctx, ffn_gate, ffn_up) : ffn_up; - - std::array bias2_ne = { out->ne[0], 1, channels, samples }; - ggml_tensor * bias2 = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, bias2_ne.data()); - out = ggml_add(ctx, out, bias2); - - ggml_set_name(out, "out"); - return out; - } else { - ggml_tensor * gates = ggml_new_tensor_3d(ctx, type, k, n, n_mats); - ggml_tensor * ups = ggml_new_tensor_3d(ctx, type, k, n, n_mats); - ggml_tensor * ids = ggml_new_tensor_2d(ctx, GGML_TYPE_I32, n_mats, m); - - if (n_used != n_mats) { - ids = ggml_view_2d(ctx, ids, n_used, m, ids->nb[1], 0); - } - - ggml_tensor * cur = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, k, this->b ? 1 : n_used, m); - ggml_set_name(cur, "cur"); - - ggml_tensor * ffn_up = ggml_mul_mat_id(ctx, ups, cur, ids); - if (with_bias) { - ggml_tensor * up_bias_param = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, ffn_up->ne[0], n_mats); - ffn_up = ggml_add_id(ctx, ffn_up, up_bias_param, ids); - } - - ggml_tensor * ffn_gate = with_gate? ggml_mul_mat_id(ctx, gates, cur, ids) : nullptr; - if (with_bias && with_gate) { - ggml_tensor * gate_bias_param = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, ffn_gate->ne[0], n_mats); - ffn_gate = ggml_add_id(ctx, ffn_gate, gate_bias_param, ids); - } - - ggml_tensor * out = with_gate ? build_gate(ctx, ffn_gate, ffn_up) : ffn_up; - - std::array scale_ne { 1, out->ne[1], out->ne[2], out->ne[3] }; - ggml_tensor * scale = ggml_new_tensor(ctx, out->type, 4, scale_ne.data()); - out = ggml_mul(ctx, out, scale); - - ggml_set_name(out, "out"); - return out; - } - } - - void initialize_tensors(ggml_context * ctx) override { - if (!use_id) { - for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != NULL; t = ggml_get_next_tensor(ctx, t)) { - init_tensor_uniform(t); - } - } else { - init_mul_mat_id_tensors(ctx, n_mats); - } - } - - double max_nmse_err() override { - return 5e-3; - } -}; - -// GGML_OP_SUM -struct test_sum : public test_case { - const ggml_type type; - const std::array ne; - const std::array permute; - bool _use_permute; - - std::string vars() override { - std::string v = VARS_TO_STR2(type, ne); - if (_use_permute) v += "," + VAR_TO_STR(permute); - return v; - } - - test_sum(ggml_type type = GGML_TYPE_F32, - std::array ne = {10, 5, 4, 3}, - std::array permute = {0, 0, 0, 0}) - : type(type), ne(ne), permute(permute), - _use_permute(permute[0] + permute[1] + permute[2] + permute[3] > 0) {} - - ggml_tensor * build_graph(ggml_context * ctx) override { - ggml_tensor * a = ggml_new_tensor(ctx, type, 4, ne.data()); - ggml_set_param(a); - ggml_set_name(a, "a"); - - if (_use_permute) { - a = ggml_permute(ctx, a, permute[0], permute[1], permute[2], permute[3]); - ggml_set_name(a, "a_permuted"); - } - - ggml_tensor * out = ggml_sum(ctx, a); - ggml_set_name(out, "out"); - - return out; - } - - float grad_eps() override { - return 0.1f * sqrtf(ne[0]*ne[1]*ne[2]*ne[3]); - } - - // Don't center the distribution around zero. Helps to avoid catastrophic cancellation. - void initialize_tensors(ggml_context * ctx) override { - for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != nullptr; t = ggml_get_next_tensor(ctx, t)) { - init_tensor_uniform(t, -0.9f, 1.1f); - } - } -}; - -// GGML_OP_SUM_ROWS -struct test_sum_rows : public test_case { - const ggml_type type; - const std::array ne; - const bool permute; - const bool slice; - - std::string vars() override { - return VARS_TO_STR4(type, ne, permute, slice); - } - - test_sum_rows(ggml_type type = GGML_TYPE_F32, - std::array ne = {10, 5, 4, 3}, - bool permute = false, bool slice = false) - : type(type), ne(ne), permute(permute), slice(slice) {} - - ggml_tensor * build_graph(ggml_context * ctx) override { - ggml_tensor * a = ggml_new_tensor(ctx, type, 4, ne.data()); - ggml_set_param(a); - ggml_set_name(a, "a"); - - if (slice) { - a = ggml_view_4d(ctx, a, - ne[0], ne[1], ne[2] / 2, ne[3] - 1, - a->nb[1], a->nb[2] * 2, a->nb[3], /*offset=*/a->nb[3]); - } - if (permute) { - a = ggml_permute(ctx, a, 0, 2, 3, 1); - } - - ggml_tensor * out = ggml_sum_rows(ctx, a); - ggml_set_name(out, "out"); - - return out; - } -}; - -// GGML_OP_MEAN -struct test_mean : public test_case { - const ggml_type type; - const std::array ne; - - std::string vars() override { - return VARS_TO_STR2(type, ne); - } - - test_mean(ggml_type type = GGML_TYPE_F32, - std::array ne = {10, 5, 4, 3}) - : type(type), ne(ne) {} - - ggml_tensor * build_graph(ggml_context * ctx) override { - ggml_tensor * a = ggml_new_tensor(ctx, type, 4, ne.data()); - ggml_set_param(a); - ggml_set_name(a, "a"); - - ggml_tensor * out = ggml_mean(ctx, a); - ggml_set_name(out, "out"); - - return out; - } - - float grad_eps() override { - return 0.1f * ne[0]*ne[1]*ne[2]*ne[3]; - } - - // Don't center the distribution around zero. Helps to avoid catastrophic cancellation. - void initialize_tensors(ggml_context * ctx) override { - for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != nullptr; t = ggml_get_next_tensor(ctx, t)) { - init_tensor_uniform(t, -0.9f, 1.1f); - } - } -}; - -// GGML_OP_UPSCALE -struct test_upscale : public test_case { - const ggml_type type; - const std::array ne; - const int32_t scale_factor; - const bool transpose; - const ggml_scale_mode mode; - - std::string vars() override { - return VARS_TO_STR5(type, ne, scale_factor, mode, transpose); - } - - test_upscale(ggml_type type = GGML_TYPE_F32, - std::array ne = {512, 512, 3, 1}, - int32_t scale_factor = 2, ggml_scale_mode mode = GGML_SCALE_MODE_NEAREST, bool transpose = false) - : type(type), ne(ne), scale_factor(scale_factor), transpose(transpose), mode(mode) {} - - ggml_tensor * build_graph(ggml_context * ctx) override { - ggml_tensor * a = ggml_new_tensor(ctx, type, 4, ne.data()); - ggml_set_name(a, "a"); - - if (transpose) { - a = ggml_transpose(ctx, a); - ggml_set_name(a, "a_transposed"); - } - - ggml_tensor * out = ggml_upscale(ctx, a, scale_factor, mode); - ggml_set_name(out, "out"); - - return out; - } -}; - -// GGML_OP_UPSCALE (via ggml_interpolate) -struct test_interpolate : public test_case { - const ggml_type type; - const std::array ne; - const std::array ne_tgt; - const ggml_scale_mode mode = GGML_SCALE_MODE_NEAREST; - - std::string vars() override { - return VARS_TO_STR4(type, ne, ne_tgt, mode); - } - - test_interpolate(ggml_type type = GGML_TYPE_F32, - std::array ne = {2, 5, 7, 11}, - std::array ne_tgt = {5, 7, 11, 13}, - ggml_scale_mode mode = GGML_SCALE_MODE_NEAREST) - : type(type), ne(ne), ne_tgt(ne_tgt), mode(mode) {} - - ggml_tensor * build_graph(ggml_context * ctx) override { - ggml_tensor * a = ggml_new_tensor(ctx, type, 4, ne.data()); - ggml_set_name(a, "a"); - - ggml_tensor * out = ggml_interpolate(ctx, a, ne_tgt[0], ne_tgt[1],ne_tgt[2], ne_tgt[3], mode); - ggml_set_name(out, "out"); - - return out; - } -}; - -// GGML_OP_GROUP_NORM -struct test_group_norm : public test_case { - const ggml_type type; - const std::array ne; - const int32_t num_groups; - const float eps; - - std::string vars() override { - return VARS_TO_STR4(type, ne, num_groups, eps); - } - - test_group_norm(ggml_type type = GGML_TYPE_F32, - std::array ne = {64, 64, 320, 1}, - int32_t num_groups = 32, - float eps = 1e-6f) - : type(type), ne(ne), num_groups(num_groups), eps(eps) {} - - ggml_tensor * build_graph(ggml_context * ctx) override { - ggml_tensor * a = ggml_new_tensor(ctx, type, 4, ne.data()); - ggml_set_name(a, "a"); - - ggml_tensor * out = ggml_group_norm(ctx, a, num_groups, eps); - ggml_set_name(out, "out"); - - return out; - } -}; - -// GGML_OP_GROUP_NORM + GGML_OP_MUL + GGML_OP_ADD -struct test_group_norm_mul_add : public test_case { - const ggml_type type; - const std::array ne; - int num_groups; - float eps; - - std::string op_desc(ggml_tensor * t) override { - GGML_UNUSED(t); - return "GROUP_NORM_MUL_ADD"; - } - - bool run_whole_graph() override { return true; } - - std::string vars() override { - return VARS_TO_STR4(type, ne, num_groups, eps); - } - - test_group_norm_mul_add(ggml_type type = GGML_TYPE_F32, - std::array ne = {128, 1, 1, 1}, - int num_groups = 4, - float eps = 1e-5f) - : type(type), ne(ne), num_groups(num_groups), eps(eps) {} - - ggml_tensor * build_graph(ggml_context * ctx) override { - ggml_tensor * a = ggml_new_tensor(ctx, type, 4, ne.data()); - ggml_tensor * w = ggml_new_tensor(ctx, type, 4, ne.data()); - ggml_tensor * b = ggml_new_tensor(ctx, type, 4, ne.data()); - ggml_set_param(a); ggml_set_param(w); ggml_set_param(b); - ggml_set_name(a, "a"); ggml_set_name(w, "w"); ggml_set_name(b, "b"); - ggml_tensor * n = ggml_group_norm(ctx, a, num_groups, eps); - ggml_tensor * m = ggml_mul(ctx, n, w); - ggml_tensor * out = ggml_add(ctx, m, b); - ggml_set_name(out, "out"); - return out; - } -}; - -// GGML_OP_L2_NORM -struct test_l2_norm : public test_case { - const ggml_type type; - const std::array ne; - const float eps; - bool v; - - std::string vars() override { - return VARS_TO_STR4(type, ne, eps, v); - } - - test_l2_norm(ggml_type type = GGML_TYPE_F32, - std::array ne = {64, 64, 320, 1}, - float eps = 1e-12f, - bool v = false) - : type(type), ne(ne), eps(eps), v(v) {} - - ggml_tensor * build_graph(ggml_context * ctx) override { - ggml_tensor * a = ggml_new_tensor(ctx, type, 4, ne.data()); - ggml_set_name(a, "a"); - - if (v) { - a = ggml_view_4d(ctx, a, a->ne[0]/2, a->ne[1]/2, a->ne[2]/2, a->ne[3]/2, a->nb[1], a->nb[2], a->nb[3], 0); - ggml_set_name(a, "view of a"); - } - - ggml_tensor * out = ggml_l2_norm(ctx, a, eps); - ggml_set_name(out, "out"); - - return out; - } -}; - -// GGML_OP_ACC -struct test_acc : public test_case { - const ggml_type type; - const std::array ne_a; - const std::array ne_b; - const int64_t stride_dim; - - std::string vars() override { - return VARS_TO_STR4(type, ne_a, ne_b, stride_dim); - } - - test_acc(ggml_type type = GGML_TYPE_F32, - std::array ne_a = {256, 17, 2, 3}, - std::array ne_b = {256, 16, 2, 3}, - uint64_t stride_dim = -1) - : type(type), ne_a(ne_a), ne_b(ne_b), stride_dim(stride_dim) {} - - ggml_tensor * build_graph(ggml_context * ctx) override { - ggml_tensor * a = ggml_new_tensor(ctx, type, 4, ne_a.data()); - ggml_set_param(a); - ggml_set_name(a, "a"); - - ggml_tensor * b; - if (stride_dim == 1 || stride_dim == 2 || stride_dim == 3) { - // Create a larger tensor and take a view at a non-zero offset. - // This tests that the backend correctly handles b's data offset - std::array ne_b_pad = {ne_b[0], ne_b[1], ne_b[2], ne_b[3]}; - ne_b_pad[stride_dim] += 1; - ggml_tensor * b_pad = ggml_new_tensor(ctx, type, 4, ne_b_pad.data()); - ggml_set_param(b_pad); - ggml_set_name(b_pad, "b_pad"); - // View that skips the first row, so b has a non-zero byte offset - b = ggml_view_4d(ctx, b_pad, - ne_b[0], ne_b[1], ne_b[2], ne_b[3], - b_pad->nb[1], b_pad->nb[2], b_pad->nb[3], - b_pad->nb[1]); - } else { - b = ggml_new_tensor(ctx, type, 4, ne_b.data()); - ggml_set_param(b); - } - ggml_set_name(b, "b"); - - // When ne_b[0] < ne_a[0], a->nb[1] != b->nb[1], so the stride - // parameters to ggml_acc don't match b's natural stride. - ggml_tensor * out = ggml_acc(ctx, a, b, a->nb[1], a->nb[2], a->nb[3], 0); - ggml_set_name(out, "out"); - - return out; - } -}; - -// GGML_OP_PAD -struct test_pad : public test_case { - const ggml_type type; - const std::array ne_a; - const int pad_0; - const int pad_1; - const bool circular; - - std::string vars() override { - return VARS_TO_STR5(type, ne_a, pad_0, pad_1, circular); - } - - test_pad(ggml_type type = GGML_TYPE_F32, - std::array ne_a = {512, 512, 1, 1}, - int pad_0 = 1, int pad_1 = 1, bool circular = false) - : type(type), ne_a(ne_a), pad_0(pad_0), pad_1(pad_1), circular(circular) {} - - ggml_tensor * build_graph(ggml_context * ctx) override { - ggml_tensor * a = ggml_new_tensor(ctx, type, 4, ne_a.data()); - ggml_set_name(a, "a"); - - ggml_tensor * out = circular - ? ggml_pad_circular(ctx, a, pad_0, pad_1, 0, 0) - : ggml_pad(ctx, a, pad_0, pad_1, 0, 0); - ggml_set_name(out, "out"); - - return out; - } -}; - -// GGML_OP_PAD (with extension) -struct test_pad_ext : public test_case { - const ggml_type type; - const std::array ne_a; - const int lp0; - const int rp0; - const int lp1; - const int rp1; - const int lp2; - const int rp2; - const int lp3; - const int rp3; - const int tfrm; // 0 - none, 1 - non-cont, 2 - perm - const bool circular; - - std::string vars() override { - return VARS_TO_STR12(type, ne_a, lp0, rp0, lp1, rp1, lp2, rp2, lp3, rp3, tfrm, circular); - } - - test_pad_ext(ggml_type type = GGML_TYPE_F32, - std::array ne_a = {512, 512, 3, 1}, - int lp0 = 1, int rp0 = 1, int lp1 = 1, int rp1 = 1, - int lp2 = 1, int rp2 = 1, int lp3 = 1, int rp3 = 1, - int tfrm = 0, bool circular = false) - : type(type), ne_a(ne_a), lp0(lp0), rp0(rp0), lp1(lp1), rp1(rp1), lp2(lp2), rp2(rp2), lp3(lp3), rp3(rp3), - tfrm(tfrm), circular(circular) {} - - ggml_tensor * build_graph(ggml_context * ctx) override { - ggml_tensor * a = ggml_new_tensor(ctx, type, 4, ne_a.data()); - ggml_set_name(a, "a"); - - if (tfrm == 1) { - a = ggml_view_4d(ctx, a, (a->ne[0] + 1) / 2, (a->ne[1] + 1) / 2, (a->ne[2] + 1) / 2, (a->ne[3] + 1) / 2, a->nb[1], a->nb[2], a->nb[3], 0); - ggml_set_name(a, "view of a"); - } else if (tfrm == 2) { - a = ggml_permute(ctx, a, 2, 1, 0, 3); - ggml_set_name(a, "permuted a"); - } - - ggml_tensor * out = circular - ? ggml_pad_ext_circular(ctx, a, lp0, rp0, lp1, rp1, lp2, rp2, lp3, rp3) - : ggml_pad_ext (ctx, a, lp0, rp0, lp1, rp1, lp2, rp2, lp3, rp3); - ggml_set_name(out, "out"); - - return out; - } -}; - -// GGML_OP_PAD_REFLECT_1D -struct test_pad_reflect_1d : public test_case { - const ggml_type type; - const std::array ne_a; - const int pad_0; - const int pad_1; - - std::string vars() override { - return VARS_TO_STR4(type, ne_a, pad_0, pad_1); - } - - test_pad_reflect_1d(ggml_type type = GGML_TYPE_F32, - std::array ne_a = {512, 34, 2, 1}, - int pad_0 = 10, int pad_1 = 9) - : type(type), ne_a(ne_a), pad_0(pad_0), pad_1(pad_1) {} - - ggml_tensor * build_graph(ggml_context * ctx) override { - ggml_tensor * a = ggml_new_tensor(ctx, type, 2, ne_a.data()); - ggml_set_name(a, "a"); - - ggml_tensor * out = ggml_pad_reflect_1d(ctx, a, pad_0, pad_1); - ggml_set_name(out, "out"); - - return out; - } -}; - -// GGML_OP_ROLL -struct test_roll : public test_case { - const int shift0; - const int shift1; - const int shift3; - const int shift4; - - std::string vars() override { - return VARS_TO_STR4(shift0, shift1, shift3, shift4); - } - - test_roll(int shift0 = 3, int shift1 = -2, int shift3 = 1, int shift4 = -1) - : shift0(shift0), shift1(shift1), shift3(shift3), shift4(shift4) {} - - ggml_tensor * build_graph(ggml_context * ctx) override { - int64_t ne[4] = {10, 5, 4, 3}; - ggml_tensor * a = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, ne); - ggml_set_name(a, "a"); - - ggml_tensor * out = ggml_roll(ctx, a, shift0, shift1, shift3, shift4); - ggml_set_name(out, "out"); - - return out; - } -}; - -// GGML_OP_ARANGE -struct test_arange : public test_case { - const ggml_type type; - const float start; - const float stop; - const float step; - - std::string vars() override { - return VARS_TO_STR4(type, start, stop, step); - } - - test_arange(ggml_type type = GGML_TYPE_F32, - float start = 0.f, float stop = 10.f, float step = 1.f) - : type(type), start(start), stop(stop), step(step) {} - - ggml_tensor * build_graph(ggml_context * ctx) override { - ggml_tensor * out = ggml_arange(ctx, start, stop, step); - ggml_set_name(out, "out"); - - return out; - } -}; - -// GGML_OP_TIMESTEP_EMBEDDING -struct test_timestep_embedding : public test_case { - const ggml_type type; - const std::array ne_a; - const int dim; - const int max_period; - - std::string vars() override { - return VARS_TO_STR4(type, ne_a, dim, max_period); - } - - test_timestep_embedding(ggml_type type = GGML_TYPE_F32, - std::array ne_a = {2, 1, 1, 1}, - int dim = 320, int max_period=10000) - : type(type), ne_a(ne_a), dim(dim), max_period(max_period) {} - - ggml_tensor * build_graph(ggml_context * ctx) override { - ggml_tensor * a = ggml_new_tensor(ctx, type, 4, ne_a.data()); - ggml_set_name(a, "a"); - - ggml_tensor * out = ggml_timestep_embedding(ctx, a, dim, max_period); - ggml_set_name(out, "out"); - - return out; - } -}; - -// GGML_OP_LEAKY_RELU -struct test_leaky_relu : public test_case { - const ggml_type type; - const std::array ne_a; - const float negative_slope; - - std::string vars() override { - return VARS_TO_STR3(type, ne_a, negative_slope); - } - - test_leaky_relu(ggml_type type = GGML_TYPE_F32, - std::array ne_a = {10, 5, 4, 3}, - float negative_slope = 0.1f) - : type(type), ne_a(ne_a), negative_slope(negative_slope) {} - - ggml_tensor * build_graph(ggml_context * ctx) override { - ggml_tensor * a = ggml_new_tensor(ctx, type, 4, ne_a.data()); - ggml_set_name(a, "a"); - - ggml_tensor * out = ggml_leaky_relu(ctx, a, negative_slope, true); - ggml_set_name(out, "out"); - - return out; - } -}; - -// GGML_OP_FLASH_ATTN_EXT -struct test_flash_attn_ext : public test_case { - const int64_t hsk; // K head size - const int64_t hsv; // V head size - const int64_t nh; // num heads - const std::array nr23; // repeat in dim 2 and 3, tests for grouped-query attention - const int64_t kv; // kv size - const int64_t nb; // batch size - - const bool mask; // use mask - const bool sinks; // use sinks - - const float max_bias; // ALiBi - const float logit_softcap; // Gemma 2 - - const ggml_prec prec; - const ggml_type type_KV; - std::array permute; - - std::string vars() override { - return VARS_TO_STR13(hsk, hsv, nh, nr23, kv, nb, mask, sinks, max_bias, logit_softcap, prec, type_KV, permute); - } - - double max_nmse_err() override { - return 5e-4; - } - - uint64_t op_flops(ggml_tensor * t) override { - GGML_UNUSED(t); - // Just counting matmul costs: - // Q*K^T is nb x hsk x kv, P*V is nb x kv x hsv, per head - return (2 * nh*nr23[0] * nb * (hsk + hsv) * kv)*nr23[1]; - } - - test_flash_attn_ext(int64_t hsk = 128, int64_t hsv = 128, int64_t nh = 32, std::array nr23 = {1, 1}, int64_t kv = 96, int64_t nb = 8, - bool mask = true, bool sinks = false, float max_bias = 0.0f, float logit_softcap = 0.0f, ggml_prec prec = GGML_PREC_F32, - ggml_type type_KV = GGML_TYPE_F16, std::array permute = {0, 1, 2, 3}) - : hsk(hsk), hsv(hsv), nh(nh), nr23(nr23), kv(kv), nb(nb), mask(mask), sinks(sinks), max_bias(max_bias), logit_softcap(logit_softcap), prec(prec), type_KV(type_KV), permute(permute) {} - - ggml_tensor * build_graph(ggml_context * ctx) override { - const int64_t hsk_padded = GGML_PAD(hsk, ggml_blck_size(type_KV)); - const int64_t hsv_padded = GGML_PAD(hsv, ggml_blck_size(type_KV)); - - auto const &create_permuted = [&](ggml_type type, int64_t ne0, int64_t ne1, int64_t ne2, int64_t ne3, bool is_view) -> ggml_tensor * { - int64_t ne[4] = {ne0, ne1, ne2, ne3}; - int64_t ne_perm[4]; - for (int i = 0; i < 4; ++i) { - ne_perm[permute[i]] = ne[i]; - } - ggml_tensor * t; - if (is_view) { - ggml_tensor * t0 = ggml_new_tensor_4d(ctx, type, ne_perm[0], 2*ne_perm[1], ne_perm[2], ne_perm[3]); - t = ggml_view_4d(ctx, t0, ne_perm[0], ne_perm[1], ne_perm[2], ne_perm[3], t0->nb[1], t0->nb[2], t0->nb[3], 0); - } else { - t = ggml_new_tensor_4d(ctx, type, ne_perm[0], ne_perm[1], ne_perm[2], ne_perm[3]); - } - if (permute != std::array{0, 1, 2, 3}) { - t = ggml_permute(ctx, t, permute[0], permute[1], permute[2], permute[3]); - } - return t; - }; - - ggml_tensor * q = create_permuted(GGML_TYPE_F32, hsk_padded, nb, nh*nr23[0], nr23[1], false); - ggml_set_name(q, "q"); - - ggml_tensor * k = create_permuted(type_KV, hsk_padded, kv, nh, nr23[1], true); // the K tensor is usually a view of the K cache - ggml_set_name(k, "k"); - - ggml_tensor * v = nullptr; - if (hsk_padded == 576 && hsv_padded == 512) { - // TODO: this branch should become a separate test case parameter instead of hardcoding this for these head shapes - - // in this branch, the V cache is sub-view of the K cache. this is used by some MLA-based models - // for more info: - // - https://github.com/ggml-org/llama.cpp/pull/13435 - // - https://github.com/ggml-org/llama.cpp/pull/18953#issuecomment-3774948392 - // - https://github.com/ggml-org/llama.cpp/pull/18986 - v = ggml_view_4d(ctx, k, hsv_padded, kv, nh, nr23[1], k->nb[1], k->nb[2], k->nb[3], 0); - } else { - v = create_permuted(type_KV, hsv_padded, kv, nh, nr23[1], true); // the V tensor is usually a view of the V cache - } - ggml_set_name(v, "v"); - - ggml_tensor * m = nullptr; - if (mask) { - m = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, kv, nb, 1, nr23[1]); - ggml_set_name(m, "m"); - } - - ggml_tensor * s = nullptr; - if (sinks) { - s = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, q->ne[2]); - ggml_set_name(s, "s"); - } - - ggml_tensor * out = ggml_flash_attn_ext(ctx, q, k, v, m, 1.0f/sqrtf(hsk), max_bias, logit_softcap); - ggml_flash_attn_ext_add_sinks(out, s); - ggml_flash_attn_ext_set_prec (out, prec); - ggml_set_name(out, "out"); - - return out; - } - - void initialize_tensors(ggml_context * ctx) override { - for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != NULL; t = ggml_get_next_tensor(ctx, t)) { - if (strcmp(t->name, "s") == 0) { - // make the sink values more noticable in order to trigger a test failure when the implementation is wrong - init_tensor_uniform(t, -10.0f, 10.0f); - } else if (strcmp(t->name, "m") == 0) { - init_tensor_kq_mask(t); - } else { - init_tensor_uniform(t); - } - } - } - - bool grad_precise() override { - return true; - } -}; - -// GGML_OP_CROSS_ENTROPY_LOSS -struct test_cross_entropy_loss : public test_case { - const ggml_type type; - const std::array ne; - - std::string vars() override { - return VARS_TO_STR2(type, ne); - } - - test_cross_entropy_loss(ggml_type type = GGML_TYPE_F32, - std::array ne = {10, 5, 4, 3}) - : type(type), ne(ne) {} - - ggml_tensor * build_graph(ggml_context * ctx) override { - ggml_tensor * logits = ggml_new_tensor(ctx, type, 4, ne.data()); - ggml_set_param(logits); - ggml_set_name(logits, "logits"); - - ggml_tensor * labels = ggml_new_tensor(ctx, type, 4, ne.data()); - // The labels are assumed to be constant -> no gradients. - ggml_set_name(labels, "labels"); - - // Ensure labels add up to 1: - labels = ggml_soft_max(ctx, labels); - ggml_set_name(labels, "labels_normalized"); - - ggml_tensor * out = ggml_cross_entropy_loss(ctx, logits, labels); - ggml_set_name(out, "out"); - - return out; - } - - void initialize_tensors(ggml_context * ctx) override { - // For larger abs. diffs between logits softmax is more linear, therefore more precise num. gradients. - for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != NULL; t = ggml_get_next_tensor(ctx, t)) { - init_tensor_uniform(t, -100.0f, 100.0f); - } - } - - float grad_eps() override { - return 1.0f; - } - - bool grad_precise() override { - return true; - } -}; - -// GGML_OP_CROSS_ENTROPY_LOSS_BACK -struct test_cross_entropy_loss_back : public test_case { - const ggml_type type; - const std::array ne; - - std::string vars() override { - return VARS_TO_STR2(type, ne); - } - - test_cross_entropy_loss_back(ggml_type type = GGML_TYPE_F32, - std::array ne = {10, 5, 4, 3}) - : type(type), ne(ne) {} - - ggml_tensor * build_graph(ggml_context * ctx) override { - ggml_tensor * grad = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 1); - ggml_set_name(grad, "grad"); - - ggml_tensor * logits = ggml_new_tensor(ctx, type, 4, ne.data()); - ggml_set_name(logits, "logits"); - - ggml_tensor * labels = ggml_new_tensor(ctx, type, 4, ne.data()); - ggml_set_name(labels, "labels"); - - // Ensure labels add up to 1: - labels = ggml_soft_max(ctx, labels); - ggml_set_name(labels, "labels_normalized"); - - ggml_tensor * out = ggml_cross_entropy_loss_back(ctx, grad, logits, labels); - ggml_set_name(out, "out"); - - return out; - } -}; - -// GGML_OP_OPT_STEP_ADAMW -struct test_opt_step_adamw : public test_case { - const ggml_type type; - const std::array ne; - - std::string vars() override { - return VARS_TO_STR2(type, ne); - } - - test_opt_step_adamw(ggml_type type = GGML_TYPE_F32, - std::array ne = {10, 5, 4, 3}) - : type(type), ne(ne) {} - - ggml_tensor * build_graph(ggml_context * ctx) override { - ggml_tensor * a = ggml_new_tensor_4d(ctx, type, ne[0], ne[1], ne[2], ne[3]); - ggml_set_param(a); // Despite tensor a having gradients the output tensor will not. - ggml_set_name(a, "a"); - - ggml_tensor * grad = ggml_new_tensor_4d(ctx, type, ne[0], ne[1], ne[2], ne[3]); - ggml_set_name(grad, "grad"); - - ggml_tensor * grad_m = ggml_new_tensor_4d(ctx, type, ne[0], ne[1], ne[2], ne[3]); - ggml_set_name(grad_m, "grad_m"); - - ggml_tensor * grad_v = ggml_new_tensor_4d(ctx, type, ne[0], ne[1], ne[2], ne[3]); - ggml_set_name(grad_v, "grad_v"); - - ggml_tensor * adamw_params = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 7); - ggml_set_name(adamw_params, "adamw_params"); - - ggml_tensor * out = ggml_opt_step_adamw(ctx, a, grad, grad_m, grad_v, adamw_params); - ggml_set_name(out, "out"); - - return out; - } - - void initialize_tensors(ggml_context * ctx) override { - for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != NULL; t = ggml_get_next_tensor(ctx, t)) { - init_tensor_uniform(t, 0.0f, 1.0f); // grad_v and adamw_params need non-negative values. - } - } - - bool grad_precise() override { - return true; - } -}; - -// GGML_OP_OPT_STEP_SGD -struct test_opt_step_sgd : public test_case { - const ggml_type type; - const std::array ne; - - std::string vars() override { return VARS_TO_STR2(type, ne); } - - test_opt_step_sgd(ggml_type type = GGML_TYPE_F32, - std::array ne = { 10, 5, 4, 3 }) - : type(type), ne(ne) {} - - ggml_tensor * build_graph(ggml_context * ctx) override { - ggml_tensor * a = ggml_new_tensor_4d(ctx, type, ne[0], ne[1], ne[2], ne[3]); - ggml_set_param(a); // Despite tensor a having gradients the output tensor will not. - ggml_set_name(a, "a"); - - ggml_tensor * grad = ggml_new_tensor_4d(ctx, type, ne[0], ne[1], ne[2], ne[3]); - ggml_set_name(grad, "grad"); - - ggml_tensor * sgd_params = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 2); - ggml_set_name(sgd_params, "sgd_params"); - - ggml_tensor * out = ggml_opt_step_sgd(ctx, a, grad, sgd_params); - - ggml_set_name(out, "out"); - - return out; - } - - void initialize_tensors(ggml_context * ctx) override { - for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != NULL; t = ggml_get_next_tensor(ctx, t)) { - init_tensor_uniform(t, 0.0f, 1.0f); // sgd_params need non-negative values. - } - } - - bool grad_precise() override { - return true; - } -}; - -// GGML_OP_CUMSUM -struct test_cumsum : public test_case { - const ggml_type type; - const std::array ne; - - std::string vars() override { return VARS_TO_STR2(type, ne); } - - test_cumsum(ggml_type type = GGML_TYPE_F32, - std::array ne = { 10, 5, 4, 3 }) - : type(type), ne(ne) {} - - ggml_tensor * build_graph(ggml_context * ctx) override { - ggml_tensor * a = ggml_new_tensor_4d(ctx, type, ne[0], ne[1], ne[2], ne[3]); - ggml_set_param(a); - ggml_set_name(a, "a"); - - ggml_tensor * out = ggml_cumsum(ctx, a); - - ggml_set_name(out, "out"); - - return out; - } - - void initialize_tensors(ggml_context * ctx) override { - for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != NULL; t = ggml_get_next_tensor(ctx, t)) { - init_tensor_uniform(t, -1.0f, 1.0f); - } - } -}; - -// GGML_OP_XIELU -struct test_xielu : public test_case { - const ggml_type type; - const std::array ne; - - std::string vars() override { return VARS_TO_STR2(type, ne); } - - test_xielu(ggml_type type = GGML_TYPE_F32, - std::array ne = { 10, 5, 4, 3 }) - : type(type), ne(ne) {} - - ggml_tensor * build_graph(ggml_context * ctx) override { - ggml_tensor * a = ggml_new_tensor_4d(ctx, type, ne[0], ne[1], ne[2], ne[3]); - ggml_set_param(a); - ggml_set_name(a, "a"); - - float alpha_n = 4.0f; - float alpha_p = 20.0f; - float beta = 0.5f; - float eps = 0.0000001f; - - ggml_tensor * out = ggml_xielu(ctx, a, alpha_n, alpha_p, beta, eps); - - ggml_set_name(out, "out"); - - return out; - } - - void initialize_tensors(ggml_context * ctx) override { - for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != NULL; t = ggml_get_next_tensor(ctx, t)) { - init_tensor_uniform(t, -1.0f, 1.0f); - } - } -}; - -// GGML_OP_TRI -struct test_tri : public test_case { - const ggml_type type; - const std::array ne; - const ggml_tri_type tri_type; - - std::string vars() override { return VARS_TO_STR3(type, ne, tri_type); } - - test_tri(ggml_tri_type tri_type, ggml_type type = GGML_TYPE_F32, - std::array ne = { 10, 10, 4, 3 }) - : type(type), ne(ne), tri_type(tri_type) { - GGML_ASSERT(ne[0] == ne[1]); - } - - ggml_tensor * build_graph(ggml_context * ctx) override { - ggml_tensor * a = ggml_new_tensor_4d(ctx, type, ne[0], ne[1], ne[2], ne[3]); - ggml_set_param(a); - ggml_set_name(a, "a"); - - ggml_tensor * out = ggml_tri(ctx, a, tri_type); - - ggml_set_name(out, "out"); - - return out; - } - - void initialize_tensors(ggml_context * ctx) override { - for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != NULL; t = ggml_get_next_tensor(ctx, t)) { - init_tensor_uniform(t, -1.0f, 1.0f); - } - } -}; - -// GGML_OP_FILL -struct test_fill : public test_case { - const ggml_type type; - const std::array ne; - float c; - - std::string vars() override { return VARS_TO_STR3(type, ne, c); } - - test_fill(float c, ggml_type type = GGML_TYPE_F32, - std::array ne = { 10, 10, 4, 3 }) - : type(type), ne(ne), c(c) {} - - ggml_tensor * build_graph(ggml_context * ctx) override { - ggml_tensor * a = ggml_new_tensor_4d(ctx, type, ne[0], ne[1], ne[2], ne[3]); - ggml_set_param(a); - ggml_set_name(a, "a"); - - ggml_tensor * out = ggml_fill(ctx, a, c); - - ggml_set_name(out, "out"); - - return out; - } -}; - -// GGML_OP_SOLVE_TRI -struct test_solve_tri : public test_case { - const ggml_type type; - const std::array ne_lhs; - const std::array ne_rhs; - - std::string vars() override { return VARS_TO_STR3(type, ne_lhs, ne_rhs); } - - uint64_t op_flops(ggml_tensor * t) override { - GGML_UNUSED(t); - int64_t n = ne_lhs[0]; - int64_t k = ne_rhs[0]; - int64_t batch = ne_lhs[2] * ne_lhs[3]; - // n * (n + 1) / 2 non-zero elements of lhs, 2 flops each, for each col of rhs - return n * (n + 1) * k * batch; - } - - test_solve_tri(ggml_type type = GGML_TYPE_F32, - std::array ne_lhs = { 10, 10, 4, 3 }, - std::array ne_rhs = { 3, 10, 4, 3 } - ) - : type(type), ne_lhs(ne_lhs), ne_rhs(ne_rhs) {} - - ggml_tensor * build_graph(ggml_context * ctx) override { - ggml_tensor * a = ggml_new_tensor_4d(ctx, type, ne_lhs[0], ne_lhs[1], ne_lhs[2], ne_lhs[3]); - ggml_set_param(a); - ggml_set_name(a, "a"); - - ggml_tensor * b = ggml_new_tensor_4d(ctx, type, ne_rhs[0], ne_rhs[1], ne_rhs[2], ne_rhs[3]); - ggml_set_param(b); - ggml_set_name(b, "b"); - - ggml_tensor * out = ggml_solve_tri(ctx, a, b, true, true, false); - ggml_set_name(out, "out"); - - return out; - } - - void initialize_tensors(ggml_context * ctx) override { - for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != NULL; t = ggml_get_next_tensor(ctx, t)) { - if (strcmp(t->name, "a") == 0) { - // note: avoid zeros in the diagonal - init_tensor_tril(t, 0.1, 1.0f); - } else { - init_tensor_uniform(t, -1.0f, 1.0f); +static void init_mul_mat_id_tensors(ggml_context * ctx, int n_mats) { + std::random_device rd; + std::default_random_engine rng(rd()); + for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != NULL; + t = ggml_get_next_tensor(ctx, t)) { + if (t->type == GGML_TYPE_I32) { + if (ggml_is_view_op(t->op)) { + continue; } - } - } -}; - -// GGML_OP_DIAG -struct test_diag : public test_case { - const ggml_type type; - const std::array ne; - - std::string vars() override { return VARS_TO_STR2(type, ne); } - - test_diag(ggml_type type = GGML_TYPE_F32, - std::array ne = { 10, 1, 4, 3 }) - : type(type), ne(ne) {} - - ggml_tensor * build_graph(ggml_context * ctx) override { - GGML_ASSERT(ne[1] == 1); - ggml_tensor * a = ggml_new_tensor_4d(ctx, type, ne[0], ne[1], ne[2], ne[3]); - ggml_set_param(a); - ggml_set_name(a, "a"); - - ggml_tensor * out = ggml_diag(ctx, a); - ggml_set_name(out, "out"); - - return out; - } -}; - - -enum llm_norm_type { - LLM_NORM, - LLM_NORM_RMS, -}; - -struct llama_hparams { - uint32_t n_vocab; - uint32_t n_embd; - uint32_t n_head; - uint32_t n_head_kv; - static constexpr uint32_t n_layer = 1; - uint32_t n_rot; - uint32_t n_embd_head; // dimension of values (d_v) - uint32_t n_ff; - - float f_norm_eps; - float f_norm_rms_eps; - - // cparams - static constexpr uint32_t n_ctx = 512; // user-specified context size - static constexpr uint32_t n_ctx_orig = n_ctx; - - // batch - int32_t n_tokens; - - // llm_build_context - static constexpr int32_t n_kv = 32; // size of KV cache to consider (n_kv <= n_ctx - static constexpr int32_t kv_head = 1; // index of where we store new KV data in the cache - - uint32_t n_embd_gqa() const { // dimension of key embeddings across all k-v heads - return n_embd_head * n_head_kv; - } -}; - -// LLM base class -struct test_llm : public test_case { - llama_hparams hp; - -protected: - test_llm(llama_hparams hp) - : hp(std::move(hp)) { - } - -public: - struct ggml_tensor * llm_build_norm( - struct ggml_context * ctx, - struct ggml_tensor * cur, - struct ggml_tensor * mw, - struct ggml_tensor * mb, - llm_norm_type type) { - switch (type) { - case LLM_NORM: cur = ggml_norm (ctx, cur, hp.f_norm_eps); break; - case LLM_NORM_RMS: cur = ggml_rms_norm(ctx, cur, hp.f_norm_rms_eps); break; - } - cur = ggml_mul(ctx, cur, mw); - if (mb) { - cur = ggml_add(ctx, cur, mb); - } - return cur; - } - - void llm_build_kv_store( - struct ggml_context * ctx, - struct ggml_tensor * k_l, - struct ggml_tensor * v_l, - struct ggml_tensor * k_cur, - struct ggml_tensor * v_cur) { - // compute the transposed [n_tokens, n_embd] V matrix - struct ggml_tensor * v_cur_t = ggml_transpose(ctx, ggml_reshape_2d(ctx, v_cur, hp.n_embd_gqa(), hp.n_tokens)); - - struct ggml_tensor * k_cache_view = ggml_view_1d(ctx, k_l, hp.n_tokens*hp.n_embd_gqa(), - (ggml_row_size(k_l->type, hp.n_embd_gqa()))*hp.kv_head); - - struct ggml_tensor * v_cache_view = ggml_view_2d(ctx, v_l, hp.n_tokens, hp.n_embd_gqa(), - ( hp.n_ctx)*ggml_element_size(v_l), - (hp.kv_head)*ggml_element_size(v_l)); - - // important: storing RoPE-ed version of K in the KV cache! - ggml_cpy(ctx, k_cur, k_cache_view); - ggml_cpy(ctx, v_cur_t, v_cache_view); - } - - struct ggml_tensor * llm_build_kqv( - struct ggml_context * ctx, - struct ggml_tensor * k_l, - struct ggml_tensor * v_l, - struct ggml_tensor * q_cur, - struct ggml_tensor * kq_mask, - float kq_scale) { - struct ggml_tensor * q = ggml_permute(ctx, q_cur, 0, 2, 1, 3); - - struct ggml_tensor * k = - ggml_view_3d(ctx, k_l, - hp.n_embd_head, hp.n_kv, hp.n_head_kv, - ggml_row_size(k_l->type, hp.n_embd_gqa()), - ggml_row_size(k_l->type, hp.n_embd_head), - 0); - - struct ggml_tensor * kq = ggml_mul_mat(ctx, k, q); - - kq = ggml_soft_max_ext(ctx, kq, kq_mask, kq_scale, 0.0f); - - // split cached v into n_head heads - struct ggml_tensor * v = - ggml_view_3d(ctx, v_l, - hp.n_kv, hp.n_embd_head, hp.n_head_kv, - ggml_element_size(v_l)*hp.n_ctx, - ggml_element_size(v_l)*hp.n_ctx*hp.n_embd_head, - 0); - - struct ggml_tensor * kqv = ggml_mul_mat(ctx, v, kq); - - struct ggml_tensor * kqv_merged = ggml_permute(ctx, kqv, 0, 2, 1, 3); - - struct ggml_tensor * cur = ggml_cont_2d(ctx, kqv_merged, hp.n_embd_head*hp.n_head, hp.n_tokens); - - struct ggml_tensor * wo = ggml_new_tensor_2d(ctx, GGML_TYPE_Q4_0, hp.n_embd, hp.n_embd); - cur = ggml_mul_mat(ctx, wo, cur); - - return cur; - } - - void initialize_tensors(ggml_context * ctx) override { - for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != NULL; t = ggml_get_next_tensor(ctx, t)) { - if (t->type == GGML_TYPE_I32) { - // pos - std::vector data(hp.n_tokens); - for (int i = 0; i < hp.n_tokens; i++) { - data[i] = rand() % hp.n_ctx; + // ids + for (int64_t r = 0; r < ggml_nrows(t); r++) { + std::vector data(t->ne[0]); + for (int i = 0; i < t->ne[0]; i++) { + data[i] = i % n_mats; } - ggml_backend_tensor_set(t, data.data(), 0, hp.n_tokens * sizeof(int)); - } else { - init_tensor_uniform(t); + std::shuffle(data.begin(), data.end(), rng); + ggml_backend_tensor_set(t, data.data(), r * t->nb[1], t->ne[0] * sizeof(int32_t)); } + } else { + init_tensor_uniform(t); } } -}; - -// Llama -struct test_llama : public test_llm { - static constexpr float freq_base = 10000.0f; - static constexpr float freq_scale = 1.0f; - static constexpr float ext_factor = 0.0f; - static constexpr float attn_factor = 1.0f; - static constexpr float beta_fast = 32.0f; - static constexpr float beta_slow = 1.0f; - bool fused; +} - std::string op_desc(ggml_tensor * t) override { - GGML_UNUSED(t); - return "LLAMA"; - } +// GGML_OP_SOFT_MAX +struct test_soft_max : public test_case { + const ggml_type type; + const std::array ne; + const bool mask; + const bool sinks; + const ggml_type m_prec; + const std::array nr23; // broadcast only dims 2 and 3 + const float scale; + const float max_bias; + const bool inplace; std::string vars() override { - auto n_tokens = hp.n_tokens; - return VARS_TO_STR1(n_tokens); + return VARS_TO_STR9(type, ne, mask, sinks, m_prec, nr23, scale, max_bias, inplace); } - double max_nmse_err() override { - return 2e-3; - } + // the 1024 test with bias occasionally fails: + // SOFT_MAX(type=f32,ne=[1024,16,1,1],mask=1,scale=1.000000,max_bias=8.000000): [SOFT_MAX] NMSE + // = 0.000000103 > 0.000000100 FAIL + virtual double max_nmse_err() override { return 1e-6; } - bool run_whole_graph() override { return fused; } - - test_llama(int n_tokens = 1, bool fused = false) - : test_llm({ - /*n_vocab =*/ 32000, - /*n_embd =*/ 3200, - /*n_head =*/ 32, - /*n_head_kv =*/ 32, - /*n_rot =*/ 100, - /*n_embd_head =*/ 100, - /*n_ff =*/ 8640, - /*f_norm_eps =*/ 0.f, - /*f_norm_rms_eps =*/ 1e-5f, - /*n_tokens =*/ n_tokens, - }) - , fused(fused) - { - } + test_soft_max(ggml_type type = GGML_TYPE_F32, + std::array ne = {10, 5, 4, 3}, + bool mask = false, + bool sinks = false, + ggml_type m_prec = GGML_TYPE_F32, + std::array nr23 = {1, 1}, + float scale = 1.0f, + float max_bias = 0.0f, + bool inplace = false) : + type(type), + ne(ne), + mask(mask), + sinks(sinks), + m_prec(m_prec), + nr23(nr23), + scale(scale), + max_bias(max_bias), + inplace(inplace) {} ggml_tensor * build_graph(ggml_context * ctx) override { - struct ggml_tensor * cur; - struct ggml_tensor * inpL; - - inpL = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, hp.n_embd, hp.n_tokens); - - // inp_pos - contains the positions - struct ggml_tensor * inp_pos = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, hp.n_tokens); - - // KQ_mask (mask for 1 head, it will be broadcasted to all heads) - struct ggml_tensor * KQ_mask = ggml_new_tensor_3d(ctx, GGML_TYPE_F16, hp.n_kv, hp.n_tokens, 1); - - ggml_tensor * k_l = ggml_new_tensor_1d(ctx, GGML_TYPE_F16, 1638400); - ggml_tensor * v_l = ggml_new_tensor_1d(ctx, GGML_TYPE_F16, 1638400); - - for (uint32_t il = 0; il < hp.n_layer; ++il) { - struct ggml_tensor * inpSA = inpL; - - // norm - ggml_tensor * attn_norm = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, hp.n_embd); - cur = llm_build_norm(ctx, inpL, attn_norm, nullptr, LLM_NORM_RMS); - - // self-attention - { - ggml_tensor * wq = ggml_new_tensor_2d(ctx, GGML_TYPE_Q4_0, hp.n_embd, hp.n_embd); - ggml_tensor * wk = ggml_new_tensor_2d(ctx, GGML_TYPE_Q4_0, hp.n_embd, hp.n_embd_gqa()); - ggml_tensor * wv = ggml_new_tensor_2d(ctx, GGML_TYPE_Q4_0, hp.n_embd, hp.n_embd_gqa()); - - // compute Q and K and RoPE them - struct ggml_tensor * Qcur = ggml_mul_mat(ctx, wq, cur); - struct ggml_tensor * Kcur = ggml_mul_mat(ctx, wk, cur); - struct ggml_tensor * Vcur = ggml_mul_mat(ctx, wv, cur); - - Qcur = ggml_rope_ext( - ctx, ggml_reshape_3d(ctx, Qcur, hp.n_embd_head, hp.n_head, hp.n_tokens), inp_pos, nullptr, - hp.n_rot, 0, hp.n_ctx_orig, freq_base, freq_scale, - ext_factor, attn_factor, beta_fast, beta_slow - ); - - Kcur = ggml_rope_ext( - ctx, ggml_reshape_3d(ctx, Kcur, hp.n_embd_head, hp.n_head_kv, hp.n_tokens), inp_pos, nullptr, - hp.n_rot, 0, hp.n_ctx_orig, freq_base, freq_scale, - ext_factor, attn_factor, beta_fast, beta_slow - ); + ggml_tensor * a = + ggml_new_tensor_4d(ctx, type, ne[0], ne[1], ne[2] * nr23[0], ne[3] * nr23[1]); + ggml_set_param(a); + ggml_set_name(a, "a"); - llm_build_kv_store(ctx, k_l, v_l, Kcur, Vcur); + ggml_tensor * mask = nullptr; + if (this->mask) { + mask = ggml_new_tensor_4d(ctx, m_prec, ne[0], ne[1], ne[2], ne[3]); + ggml_set_name(mask, "mask"); + } - cur = llm_build_kqv(ctx, k_l, v_l, Qcur, KQ_mask, 1.0f/sqrtf(float(hp.n_embd_head))); - } + ggml_tensor * sinks = nullptr; + if (this->sinks) { + sinks = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, ne[2] * nr23[0]); + ggml_set_name(sinks, "sinks"); + } - struct ggml_tensor * ffn_inp = ggml_add(ctx, cur, inpSA); + ggml_tensor * out; + if (inplace) { + out = ggml_soft_max_ext_inplace(ctx, a, mask, scale, max_bias); + } else { + out = ggml_soft_max_ext(ctx, a, mask, scale, max_bias); + } + ggml_soft_max_add_sinks(out, sinks); + ggml_set_name(out, "out"); - // feed-forward network - ggml_tensor * ffn_norm = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, hp.n_embd); - cur = llm_build_norm(ctx, ffn_inp, ffn_norm, nullptr, LLM_NORM_RMS); + return out; + } - ggml_tensor * ffn_gate = ggml_new_tensor_2d(ctx, GGML_TYPE_Q4_0, hp.n_embd, hp.n_ff); - ggml_tensor * ffn_down = ggml_new_tensor_2d(ctx, GGML_TYPE_Q4_0, hp.n_ff, hp.n_embd); - ggml_tensor * ffn_up = ggml_new_tensor_2d(ctx, GGML_TYPE_Q4_0, hp.n_embd, hp.n_ff); - struct ggml_tensor * tmp = ggml_mul_mat(ctx, ffn_up, cur); - cur = ggml_mul_mat(ctx, ffn_gate, cur); - cur = ggml_silu(ctx, cur); - cur = ggml_mul(ctx, cur, tmp); - cur = ggml_mul_mat(ctx, ffn_down, cur); + bool grad_precise() override { return true; } +}; - cur = ggml_add(ctx, cur, ffn_inp); +// GGML_OP_POOL2D +struct test_pool2d : public test_case { + enum ggml_op_pool pool_type; + const ggml_type type_input; + const std::array ne_input; + // kernel size + const int k0; + const int k1; + // stride + const int s0; + const int s1; + // padding + const int p0; + const int p1; - // input for next layer - inpL = cur; - } + std::string vars() override { + return VARS_TO_STR9(pool_type, type_input, ne_input, k0, k1, s0, s1, p0, p1); + } - cur = inpL; + test_pool2d(ggml_op_pool pool_type = GGML_OP_POOL_AVG, + ggml_type type_input = GGML_TYPE_F32, + std::array ne_input = + {10, 10, 3, 1}, // [input_width, input_height, input_channels, 1] + int k0 = 3, + int k1 = 3, + int s0 = 1, + int s1 = 1, + int p0 = 1, + int p1 = 1) : + pool_type(pool_type), + type_input(type_input), + ne_input(ne_input), + k0(k0), + k1(k1), + s0(s0), + s1(s1), + p0(p0), + p1(p1) {} - ggml_tensor * output_norm = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, hp.n_embd); - cur = llm_build_norm(ctx, cur, output_norm, nullptr, LLM_NORM_RMS); + ggml_tensor * build_graph(ggml_context * ctx) override { + ggml_tensor * input = ggml_new_tensor(ctx, type_input, 4, ne_input.data()); + ggml_set_param(input); + ggml_set_name(input, "input"); - // lm_head - ggml_tensor * output = ggml_new_tensor_2d(ctx, GGML_TYPE_Q4_0, hp.n_embd, hp.n_vocab); - cur = ggml_mul_mat(ctx, output, cur); + ggml_tensor * out = ggml_pool_2d(ctx, input, pool_type, k0, k1, s0, s1, p0, p1); + ggml_set_name(out, "out"); - return cur; + return out; } }; -// Falcon -struct test_falcon : public test_llm { - static constexpr float freq_base = 10000.0f; - static constexpr float freq_scale = 1.0f; - static constexpr float ext_factor = 0.0f; - static constexpr float attn_factor = 1.0f; - static constexpr float beta_fast = 32.0f; - static constexpr float beta_slow = 1.0f; +// CONV_2D +struct test_conv_2d : public test_case { + const std::array ne_input; + const std::array ne_kernel; + const ggml_type type_kernel; + const int stride0; + const int stride1; + const int padding0; + const int padding1; + const int dilation0; + const int dilation1; + // Whether the inputs are contiguous in the channel dim or the width dim + const bool cwhn; - std::string op_desc(ggml_tensor * t) override { - GGML_UNUSED(t); - return "FALCON"; - } + // If true, the direct CONV_2D will be used in the graph, otherwise it + // uses ggml_conv_2d: + // * if the program is called with -o CONV_2D_DIRECT_IMPL, the + // CONV_2D graph will be built, while + // * if the program is called with -o CONV_2D_INDIRECT_IMPL, the + // IM2COL -> MUL_MM graph will be built. std::string vars() override { - auto n_tokens = hp.n_tokens; - return VARS_TO_STR1(n_tokens); + return VARS_TO_STR10(ne_input, ne_kernel, type_kernel, stride0, stride1, padding0, padding1, + dilation0, dilation1, cwhn); } - double max_nmse_err() override { - return 2e-3; - } - - test_falcon(int n_tokens = 1) - : test_llm({ - /*n_vocab =*/ 32000, - /*n_embd =*/ 3200, - /*n_head =*/ 50, - /*n_head_kv =*/ 1, - /*n_rot =*/ 64, - /*n_embd_head =*/ 64, - /*n_ff =*/ 8640, - /*f_norm_eps =*/ 1e-5f, - /*f_norm_rms_eps =*/ 0.f, - /*n_tokens =*/ n_tokens, - }) { - } - - ggml_tensor * build_graph(ggml_context * ctx) override { - struct ggml_tensor * cur; - struct ggml_tensor * inpL; + double max_nmse_err() override { return 5e-4; } - inpL = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, hp.n_embd, hp.n_tokens); + uint64_t op_flops(ggml_tensor * t) override { + GGML_UNUSED(t); + // Just counting matmul costs: + // KxCRS @ CRSxNPQ = KxNPQ --> KxNPQx(CRS+CRS-1) flops - // inp_pos - contains the positions - struct ggml_tensor * inp_pos = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, hp.n_tokens); + // Copied from ggml.c: int64_t ggml_calc_conv_output_size(int64_t ins, int64_t ks, int s, + // int p, int d) + auto calc_conv_output_size = [](int64_t ins, int64_t ks, int s, int p, int d) -> int64_t { + return (ins + 2 * p - d * (ks - 1) - 1) / s + 1; + }; - // KQ_mask (mask for 1 head, it will be broadcasted to all heads) - struct ggml_tensor * KQ_mask = ggml_new_tensor_3d(ctx, GGML_TYPE_F16, hp.n_kv, hp.n_tokens, 1); + int64_t W = ne_input[0]; + int64_t H = ne_input[1]; + int64_t KW = ne_kernel[0]; + int64_t KH = ne_kernel[1]; + int64_t Cin = ne_kernel[2]; + int64_t Cout = ne_kernel[3]; + int64_t N = ne_input[3]; + int64_t OH = calc_conv_output_size(H, KH, stride0, padding0, dilation0); + int64_t OW = calc_conv_output_size(W, KW, stride0, padding0, dilation0); - ggml_tensor * k_l = ggml_new_tensor_1d(ctx, GGML_TYPE_F16, 1638400); - ggml_tensor * v_l = ggml_new_tensor_1d(ctx, GGML_TYPE_F16, 1638400); + int64_t K = Cout; + int64_t CRS = Cin * KH * KW; + int64_t NPQ = N * OH * OW; - for (uint32_t il = 0; il < hp.n_layer; ++il) { - // norm - ggml_tensor * attn_norm_w = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, hp.n_embd); - ggml_tensor * attn_norm_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, hp.n_embd); - ggml_tensor * attn_norm = llm_build_norm(ctx, inpL, attn_norm_w, attn_norm_b, LLM_NORM); + return K * NPQ * (2 * CRS - 1); + } - // self-attention - { - cur = attn_norm; + test_conv_2d(std::array ne_input = {64, 64, 16, 1}, + std::array ne_kernel = {3, 3, 1, 16}, + ggml_type type_kernel = GGML_TYPE_F32, + int stride0 = 1, + int stride1 = 1, + int padding0 = 0, + int padding1 = 0, + int dilation0 = 1, + int dilation1 = 1, + bool cwhn = false) : + ne_input(ne_input), + ne_kernel(ne_kernel), + type_kernel(type_kernel), + stride0(stride0), + stride1(stride1), + padding0(padding0), + padding1(padding1), + dilation0(dilation0), + dilation1(dilation1), + cwhn(cwhn) {} - ggml_tensor * wqkv = ggml_new_tensor_2d(ctx, GGML_TYPE_Q4_0, hp.n_embd, hp.n_embd + 2*hp.n_embd_gqa()); + ggml_tensor * build_graph(ggml_context * ctx) override { + ggml_tensor * input = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, ne_input.data()); + ggml_set_name(input, "input"); - cur = ggml_mul_mat(ctx, wqkv, cur); + ggml_tensor * kernel = ggml_new_tensor(ctx, type_kernel, 4, ne_kernel.data()); + ggml_set_name(kernel, "kernel"); - struct ggml_tensor * Qcur = ggml_cont(ctx, ggml_view_2d(ctx, cur, hp.n_embd, hp.n_tokens, cur->nb[1], 0*sizeof(float)*(hp.n_embd))); - struct ggml_tensor * Kcur = ggml_cont(ctx, ggml_view_2d(ctx, cur, hp.n_embd_gqa(), hp.n_tokens, cur->nb[1], 1*sizeof(float)*(hp.n_embd))); - struct ggml_tensor * Vcur = ggml_cont(ctx, ggml_view_2d(ctx, cur, hp.n_embd_gqa(), hp.n_tokens, cur->nb[1], 1*sizeof(float)*(hp.n_embd + hp.n_embd_gqa()))); + if (cwhn) { + // change memory layout to channel-most-contiguous (CWHN), + // then permute it back so NE matches the original input + input = ggml_cont(ctx, ggml_permute(ctx, input, 1, 2, 0, 3)); + input = ggml_permute(ctx, input, 2, 0, 1, 3); + kernel = ggml_cont(ctx, ggml_permute(ctx, kernel, 2, 3, 1, 0)); + kernel = ggml_permute(ctx, kernel, 3, 2, 0, 1); + } - Qcur = ggml_reshape_3d(ctx, Qcur, hp.n_embd_head, hp.n_head, hp.n_tokens); - Kcur = ggml_reshape_3d(ctx, Kcur, hp.n_embd_head, hp.n_head_kv, hp.n_tokens); + ggml_tensor * out = ggml_conv_2d_direct(ctx, kernel, input, stride0, stride1, padding0, + padding1, dilation0, dilation1); + ggml_set_name(out, "out"); + return out; + } +}; - // using mode = 2 for neox mode - Qcur = ggml_rope_ext( - ctx, Qcur, inp_pos, nullptr, hp.n_rot, 2, hp.n_ctx_orig, - freq_base, freq_scale, ext_factor, attn_factor, beta_fast, beta_slow - ); +// GGML_OP_CROSS_ENTROPY_LOSS +struct test_cross_entropy_loss : public test_case { + const ggml_type type; + const std::array ne; - Kcur = ggml_rope_ext( - ctx, Kcur, inp_pos, nullptr, hp.n_rot, 2, hp.n_ctx_orig, - freq_base, freq_scale, ext_factor, attn_factor, beta_fast, beta_slow - ); + std::string vars() override { return VARS_TO_STR2(type, ne); } - llm_build_kv_store(ctx, k_l, v_l, Kcur, Vcur); + test_cross_entropy_loss(ggml_type type = GGML_TYPE_F32, + std::array ne = {10, 5, 4, 3}) : + type(type), ne(ne) {} - cur = llm_build_kqv(ctx, k_l, v_l, Qcur, KQ_mask, 1.0f/sqrtf(float(hp.n_embd_head))); - } + ggml_tensor * build_graph(ggml_context * ctx) override { + ggml_tensor * logits = ggml_new_tensor(ctx, type, 4, ne.data()); + ggml_set_param(logits); + ggml_set_name(logits, "logits"); - struct ggml_tensor * ffn_inp = cur; + ggml_tensor * labels = ggml_new_tensor(ctx, type, 4, ne.data()); + // The labels are assumed to be constant -> no gradients. + ggml_set_name(labels, "labels"); - // feed forward - { - ggml_tensor * ffn_up = ggml_new_tensor_2d(ctx, GGML_TYPE_Q4_0, hp.n_embd, hp.n_ff); - ggml_tensor * ffn_down = ggml_new_tensor_2d(ctx, GGML_TYPE_Q4_0, hp.n_ff, hp.n_embd); - cur = attn_norm; - cur = ggml_mul_mat(ctx, ffn_up, cur); - cur = ggml_gelu(ctx, cur); - cur = ggml_mul_mat(ctx, ffn_down, cur); - } + // Ensure labels add up to 1: + labels = ggml_soft_max(ctx, labels); + ggml_set_name(labels, "labels_normalized"); - cur = ggml_add(ctx, cur, ffn_inp); + ggml_tensor * out = ggml_cross_entropy_loss(ctx, logits, labels); + ggml_set_name(out, "out"); - cur = ggml_add(ctx, cur, inpL); + return out; + } - // input for next layer - inpL = cur; + void initialize_tensors(ggml_context * ctx) override { + // For larger abs. diffs between logits softmax is more linear, therefore more precise num. + // gradients. + for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != NULL; + t = ggml_get_next_tensor(ctx, t)) { + init_tensor_uniform(t, -100.0f, 100.0f); } + } - cur = inpL; - - ggml_tensor * output_norm = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, hp.n_embd); - ggml_tensor * output_norm_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, hp.n_embd); - cur = llm_build_norm(ctx, cur, output_norm, output_norm_b, LLM_NORM); - - // lm_head - ggml_tensor * output = ggml_new_tensor_2d(ctx, GGML_TYPE_Q8_0, hp.n_embd, hp.n_vocab); - cur = ggml_mul_mat(ctx, output, cur); + float grad_eps() override { return 1.0f; } - return cur; - } + bool grad_precise() override { return true; } }; - -// ########################################### // ## Section 3: GGML Op Test Instantiation ## // ########################################### -static const ggml_type all_types[] = { - GGML_TYPE_F32, GGML_TYPE_F16, GGML_TYPE_BF16, - GGML_TYPE_Q4_0, GGML_TYPE_Q4_1, - GGML_TYPE_Q5_0, GGML_TYPE_Q5_1, - GGML_TYPE_Q8_0, - GGML_TYPE_MXFP4, - GGML_TYPE_Q2_K, GGML_TYPE_Q3_K, - GGML_TYPE_Q4_K, GGML_TYPE_Q5_K, - GGML_TYPE_Q6_K, - // GGML_TYPE_TQ1_0, GGML_TYPE_TQ2_0, // TODO: implement for all backends - GGML_TYPE_IQ2_XXS, GGML_TYPE_IQ2_XS, GGML_TYPE_IQ2_S, - GGML_TYPE_IQ3_XXS, GGML_TYPE_IQ1_S, GGML_TYPE_IQ1_M, - GGML_TYPE_IQ4_NL, GGML_TYPE_IQ3_S, GGML_TYPE_IQ4_XS, -}; - -static const ggml_type base_types[] = { - GGML_TYPE_F32, GGML_TYPE_F16, - GGML_TYPE_Q8_0, // for I8MM tests - GGML_TYPE_Q4_0, - GGML_TYPE_Q4_1, // for I8MM tests - GGML_TYPE_Q4_K, - GGML_TYPE_MXFP4, // TODO: or "other" - GGML_TYPE_IQ2_XXS -}; - -static const ggml_type other_types[] = { - GGML_TYPE_Q4_1, - GGML_TYPE_Q5_0, GGML_TYPE_Q5_1, - GGML_TYPE_Q8_0, - GGML_TYPE_Q2_K, GGML_TYPE_Q3_K, - GGML_TYPE_Q5_K, - GGML_TYPE_Q6_K, - // GGML_TYPE_TQ1_0, GGML_TYPE_TQ2_0, // TODO: implement for all backends - GGML_TYPE_IQ2_XS, GGML_TYPE_IQ2_S, - GGML_TYPE_IQ3_XXS, GGML_TYPE_IQ1_S, GGML_TYPE_IQ1_M, - GGML_TYPE_IQ4_NL, GGML_TYPE_IQ3_S, GGML_TYPE_IQ4_XS, - GGML_TYPE_BF16, -}; #ifdef _MSC_VER // Workaround long compile time with msvc #pragma optimize("", off) #endif -// Test cases for evaluation: should try to cover edge cases while using small input sizes to keep the runtime low +// Test cases for evaluation: should try to cover edge cases while using small input sizes to keep +// the runtime low static std::vector> make_test_cases_eval() { std::vector> test_cases; std::default_random_engine rng(0); - // MNIST layer tests (FP32) - // Layer 1: input [784] x weights [784, 500] -> [500, 500] - test_cases.emplace_back(new test_mul_mat(GGML_TYPE_F32, GGML_TYPE_F32, 500, 500, 784, {1, 1}, {1, 1})); - // Layer 1: add bias [500, 500] + [500, 1] -> [500, 500] (broadcast) - test_cases.emplace_back(new test_bin_bcast(ggml_add, GGML_TYPE_F32, {500, 1, 1, 1}, {1, 500, 1, 1})); - // Layer 1: ReLU activation [500] -> [500] - test_cases.emplace_back(new test_unary(GGML_UNARY_OP_RELU, GGML_TYPE_F32, {500, 1, 1, 1})); - // Layer 2: hidden [500] x weights [500, 10] -> [10, 500] - test_cases.emplace_back(new test_mul_mat(GGML_TYPE_F32, GGML_TYPE_F32, 10, 500, 500, {1, 1}, {1, 1})); - // Layer 2: add bias [10, 500] + [10, 1] -> [10, 500] (broadcast) - test_cases.emplace_back(new test_bin_bcast(ggml_add, GGML_TYPE_F32, {10, 1, 1, 1}, {1, 500, 1, 1})); - // Layer 2: argmax [10, 500] -> [500, 1] + // MNIST-MLP layer tests (FP32, batch=500) + // FC1: images [784, 500] x fc1_weight [784, 500] -> [500, 500] + test_cases.emplace_back( + new test_mul_mat(GGML_TYPE_F32, GGML_TYPE_F32, 500, 500, 784, {1, 1}, {1, 1})); + // FC1 bias: [500, 500] + fc1_bias [500] -> [500, 500] (broadcast) + test_cases.emplace_back( + new test_bin_bcast(ggml_add, GGML_TYPE_F32, {500, 1, 1, 1}, {1, 500, 1, 1})); + // FC1 ReLU: [500, 500] -> [500, 500] + test_cases.emplace_back(new test_unary(GGML_UNARY_OP_RELU, GGML_TYPE_F32, {500, 500, 1, 1})); + // FC2: fc1_out [500, 500] x fc2_weight [500, 10] -> [10, 500] + test_cases.emplace_back( + new test_mul_mat(GGML_TYPE_F32, GGML_TYPE_F32, 10, 500, 500, {1, 1}, {1, 1})); + // FC2 bias: [10, 500] + fc2_bias [10] -> [10, 500] (broadcast) + test_cases.emplace_back( + new test_bin_bcast(ggml_add, GGML_TYPE_F32, {10, 1, 1, 1}, {1, 500, 1, 1})); + // Argmax of logits [10, 500] -> [500] test_cases.emplace_back(new test_argmax(GGML_TYPE_F32, {10, 500, 1, 1})); - // Cross entropy loss: [10, 500] x [10, 500] -> [1, 1] + // Cross entropy loss on logits [10, 500] test_cases.emplace_back(new test_cross_entropy_loss(GGML_TYPE_F32, {10, 500, 1, 1})); - // Cross entropy loss: [500, 1] x [500, 1] -> [1, 1] - test_cases.emplace_back(new test_cross_entropy_loss(GGML_TYPE_F32, {500, 1, 1, 1})); - - // Soft max: [10, 500] x [10, 500] -> [1, 1] + // Softmax on logits [10, 500] test_cases.emplace_back(new test_soft_max(GGML_TYPE_F32, {10, 500, 1, 1})); - // Soft max: [500, 1] x [500, 1] -> [1, 1] - test_cases.emplace_back(new test_soft_max(GGML_TYPE_F32, {500, 1, 1, 1})); + + // MNIST-CNN layer tests (FP32, batch=500, NCB=8) + // Conv1: images [28, 28, 1, 500] x conv1_kernel [3, 3, 1, 8], stride=1, pad=1 -> [28, 28, 8, + // 500] + test_cases.emplace_back( + new test_conv_2d({28, 28, 1, 500}, {3, 3, 1, 8}, GGML_TYPE_F32, 1, 1, 1, 1, 1, 1)); + // Conv1 bias + ReLU: [28, 28, 8, 500] + test_cases.emplace_back(new test_unary(GGML_UNARY_OP_RELU, GGML_TYPE_F32, {28, 28, 8, 500})); + // MaxPool1: [28, 28, 8, 500] with 2x2 kernel, stride=2, pad=0 -> [14, 14, 8, 500] + test_cases.emplace_back( + new test_pool2d(GGML_OP_POOL_MAX, GGML_TYPE_F32, {28, 28, 8, 500}, 2, 2, 2, 2, 0, 0)); + // MaxPool1 with padding: [28, 28, 8, 500] with 3x3 kernel, stride=2, pad=1 -> [14, 14, 8, 500] + test_cases.emplace_back( + new test_pool2d(GGML_OP_POOL_MAX, GGML_TYPE_F32, {28, 28, 8, 500}, 3, 3, 2, 2, 1, 1)); + // MaxPool1 F16 input: [28, 28, 8, 500] with 2x2 kernel, stride=2, pad=0 -> [14, 14, 8, 500] + test_cases.emplace_back( + new test_pool2d(GGML_OP_POOL_MAX, GGML_TYPE_F16, {28, 28, 8, 500}, 2, 2, 2, 2, 0, 0)); + // Conv2: [14, 14, 8, 500] x conv2_kernel [3, 3, 8, 16], stride=1, pad=1 -> [14, 14, 16, 500] + test_cases.emplace_back( + new test_conv_2d({14, 14, 8, 500}, {3, 3, 8, 16}, GGML_TYPE_F32, 1, 1, 1, 1, 1, 1)); + // Conv2 bias + ReLU: [14, 14, 16, 500] + test_cases.emplace_back(new test_unary(GGML_UNARY_OP_RELU, GGML_TYPE_F32, {14, 14, 16, 500})); + // MaxPool2: [14, 14, 16, 500] with 2x2 kernel, stride=2, pad=0 -> [7, 7, 16, 500] + test_cases.emplace_back( + new test_pool2d(GGML_OP_POOL_MAX, GGML_TYPE_F32, {14, 14, 16, 500}, 2, 2, 2, 2, 0, 0)); + // Dense: flattened [784, 500] x dense_weight [784, 10] -> [10, 500] (784 = 7*7*16) + test_cases.emplace_back( + new test_mul_mat(GGML_TYPE_F32, GGML_TYPE_F32, 10, 500, 784, {1, 1}, {1, 1})); + // Dense bias: [10, 500] + dense_bias [10] -> [10, 500] (broadcast) + test_cases.emplace_back( + new test_bin_bcast(ggml_add, GGML_TYPE_F32, {10, 1, 1, 1}, {1, 500, 1, 1})); + // Argmax of CNN logits [10, 500] -> [500] + test_cases.emplace_back(new test_argmax(GGML_TYPE_F32, {10, 500, 1, 1})); return test_cases; } @@ -7075,262 +2516,45 @@ static std::vector> make_test_cases_eval() { #pragma optimize("", on) #endif -// Test cases for performance evaluation: should be representative of real-world use cases +// Test cases for performance evaluation: should be representative of MNIST workloads static std::vector> make_test_cases_perf() { std::vector> test_cases; - // Conv2d: K=CRS=NPQ=4096 matmul performance - uint32_t iwh_idx = 0; - uint32_t kwh_idx = 1; - uint32_t Cout_idx = 2; - uint32_t Cin_idx = 3; - uint32_t B_idx = 4; - std::vector> cases = { - //{IWH, KWH, Cout, Cin, B} - // K=CRS=NPQ=4096 conv2d matmul performance - {19, 4, 4096, 256, 16}, - // K=128, CRS=128, NPQ=4096 - { 19, 4, 128, 8, 16}, - // K=130, CRS=128, NPQ=4096 - { 19, 4, 130, 8, 16}, - // Edge case: K x CRS is small - { 19, 2, 4, 4, 16}, - // A ConvNet's first layer - { 224, 3, 8, 3, 1 }, - // A ConvNet's first layer with 2x2 convolution, and 1 channel - { 224, 2, 8, 1, 1 }, - // A ConvNet's first layer with 2x2 convolution, and 1 channel, several images in the batch - { 224, 2, 8, 1, 8 }, - // A middle layer of a ConvNet - { 58, 3, 64, 32, 1 }, - // A middle layer of a ConvNet, several images in the batch - { 58, 3, 64, 32, 8 }, - // A deep layer of a ConvNet, several images in the batch - { 16, 3, 512, 128, 8 }, - // High resolution output (large NPQ) - {1536, 3, 64, 32, 1 }, - }; - - for (auto kernel_type : {GGML_TYPE_F32, GGML_TYPE_F16}) { - for (auto act_case : cases) { - // Direct CONV_2D - test_cases.emplace_back(new test_conv_2d( - { act_case[iwh_idx], act_case[iwh_idx], act_case[Cin_idx], act_case[B_idx] }, - { act_case[kwh_idx], act_case[kwh_idx], act_case[Cin_idx], act_case[Cout_idx] }, - kernel_type, 1, 1, 0, 0, 1, 1, false)); - } - } - - test_cases.emplace_back(new test_bin_bcast(ggml_add, GGML_TYPE_F32, {4096, 1, 1, 1}, {1, 1, 1, 1})); - test_cases.emplace_back(new test_bin_bcast(ggml_add, GGML_TYPE_F32, {4096, 1, 1, 1}, {1, 512, 1, 1})); - - test_cases.emplace_back(new test_cpy(GGML_TYPE_F32, GGML_TYPE_F16, {512, 3072, 1, 1})); - test_cases.emplace_back(new test_cpy(GGML_TYPE_F32, GGML_TYPE_F32, {8192, 512, 2, 1}, {0, 2, 1, 3})); - test_cases.emplace_back(new test_cpy(GGML_TYPE_F32, GGML_TYPE_F32, {3072, 512, 2, 1}, {0, 2, 1, 3})); - test_cases.emplace_back(new test_cpy(GGML_TYPE_F32, GGML_TYPE_Q4_0, {8192, 512, 2, 1})); - test_cases.emplace_back(new test_cpy(GGML_TYPE_Q4_0, GGML_TYPE_F32, {8192, 512, 2, 1})); - - test_cases.emplace_back(new test_cpy(GGML_TYPE_F32, GGML_TYPE_F32, {768*1024, 256, 1, 1}, {1, 0, 2, 3}, {0, 0, 0, 0})); - test_cases.emplace_back(new test_cpy(GGML_TYPE_F16, GGML_TYPE_F16, {768*1024, 256, 1, 1}, {1, 0, 2, 3}, {0, 0, 0, 0})); - test_cases.emplace_back(new test_cpy(GGML_TYPE_F16, GGML_TYPE_F16, {768, 1024, 256, 1}, {1, 0, 2, 3}, {0, 0, 0, 0})); - test_cases.emplace_back(new test_cpy(GGML_TYPE_BF16, GGML_TYPE_BF16, {768, 1024, 256, 1}, {1, 0, 2, 3}, {0, 0, 0, 0})); - - test_cases.emplace_back(new test_cpy(GGML_TYPE_F32, GGML_TYPE_F32, {768*1024, 256, 1, 1}, {0, 0, 0, 0}, {0, 0, 0, 0}, true)); - test_cases.emplace_back(new test_cpy(GGML_TYPE_F32, GGML_TYPE_F32, {768, 1024, 256, 1}, {0, 0, 0, 0}, {0, 0, 0, 0}, true)); - test_cases.emplace_back(new test_cpy(GGML_TYPE_F16, GGML_TYPE_F16, {768*1024, 256, 1, 1}, {0, 0, 0, 0}, {0, 0, 0, 0}, true)); - test_cases.emplace_back(new test_cpy(GGML_TYPE_F16, GGML_TYPE_F16, {768, 1024, 256, 1}, {0, 0, 0, 0}, {0, 0, 0, 0}, true)); - test_cases.emplace_back(new test_cpy(GGML_TYPE_BF16, GGML_TYPE_BF16, {768, 1024, 256, 1}, {0, 0, 0, 0}, {0, 0, 0, 0}, true)); - - - test_cases.emplace_back(new test_soft_max(GGML_TYPE_F32, {4096, 4096, 5, 1}, false, false, GGML_TYPE_F32, {1, 1}, 1.0f, 0.0f)); - test_cases.emplace_back(new test_soft_max(GGML_TYPE_F32, {12888, 256, 5, 1}, false, false, GGML_TYPE_F32, {1, 1}, 1.0f, 0.0f)); - test_cases.emplace_back(new test_soft_max(GGML_TYPE_F32, {77, 4096, 5, 1}, false, false, GGML_TYPE_F32, {1, 1}, 1.0f, 0.0f)); - test_cases.emplace_back(new test_soft_max(GGML_TYPE_F32, {1024, 1024, 10, 1}, false, false, GGML_TYPE_F32, {1, 1}, 1.0f, 0.0f)); - test_cases.emplace_back(new test_soft_max(GGML_TYPE_F32, {77, 1024, 10, 1}, false, false, GGML_TYPE_F32, {1, 1}, 1.0f, 0.0f)); - test_cases.emplace_back(new test_soft_max(GGML_TYPE_F32, {256, 256, 20, 1}, false, false, GGML_TYPE_F32, {1, 1}, 1.0f, 0.0f)); - test_cases.emplace_back(new test_soft_max(GGML_TYPE_F32, {64, 64, 20, 1}, false, false, GGML_TYPE_F32, {1, 1}, 1.0f, 0.0f)); - test_cases.emplace_back(new test_soft_max(GGML_TYPE_F32, {77, 64, 20, 1}, false, false, GGML_TYPE_F32, {1, 1}, 1.0f, 0.0f)); - - test_cases.emplace_back(new test_argmax(GGML_TYPE_F32, {32, 10, 1, 1})); - test_cases.emplace_back(new test_argmax(GGML_TYPE_F32, {1024, 10, 1, 1})); - test_cases.emplace_back(new test_argmax(GGML_TYPE_F32, {32000, 512, 1, 1})); - - test_cases.emplace_back(new test_pad_reflect_1d(GGML_TYPE_F32, {512, 34, 2, 1})); - test_cases.emplace_back(new test_pad_reflect_1d(GGML_TYPE_F32, {3000, 80, 1, 1})); - test_cases.emplace_back(new test_pad_reflect_1d(GGML_TYPE_F32, {3000, 80, 4, 1})); - test_cases.emplace_back(new test_pad_reflect_1d(GGML_TYPE_F32, {3000, 384, 1, 1})); - test_cases.emplace_back(new test_pad_reflect_1d(GGML_TYPE_F32, {3000, 384, 4, 1})); - - test_cases.emplace_back(new test_mul_mat(GGML_TYPE_F16, GGML_TYPE_F32, 16416, 1, 128, {8, 1}, {4, 1}, {0, 2, 1, 3})); - test_cases.emplace_back(new test_mul_mat(GGML_TYPE_F16, GGML_TYPE_F32, 128, 1, 16416, {8, 1}, {4, 1}, {0, 1, 2, 3}, 2*16416)); - - test_cases.emplace_back(new test_solve_tri(GGML_TYPE_F32, { 64, 64, 4, 4 }, { 32, 64, 4, 4 })); - test_cases.emplace_back(new test_solve_tri(GGML_TYPE_F32, { 128, 128, 4, 2 }, { 32, 128, 4, 2 })); - // qwen3next with CHUNK_SIZE 64 - test_cases.emplace_back(new test_solve_tri(GGML_TYPE_F32, { 64, 64, 8, 32 }, { 64, 64, 8, 32 })); - // qwen3next with CHUNK_SIZE 128 - test_cases.emplace_back(new test_solve_tri(GGML_TYPE_F32, { 128, 128, 4, 32 }, { 128, 128, 4, 32 })); - test_cases.emplace_back(new test_solve_tri(GGML_TYPE_F32, { 256, 256, 4, 2 }, { 128, 256, 4, 2 })); - - test_cases.emplace_back(new test_tri(GGML_TRI_TYPE_LOWER, GGML_TYPE_F32, { 256, 256, 4, 4 })); - test_cases.emplace_back(new test_tri(GGML_TRI_TYPE_UPPER_DIAG, GGML_TYPE_F32, { 1024, 1024, 8, 4 })); - - test_cases.emplace_back(new test_cumsum(GGML_TYPE_F32, { 128, 128, 4, 4 })); - test_cases.emplace_back(new test_cumsum(GGML_TYPE_F32, { 2048, 16, 5, 4 })); - test_cases.emplace_back(new test_cumsum(GGML_TYPE_F32, { 20000, 10, 4, 1 })); - - for (int bs : {1, 2, 3, 4, 5, 8, 512}) { - for (ggml_type type_a : all_types) { - for (ggml_type type_b : {GGML_TYPE_F32}) { - test_cases.emplace_back(new test_mul_mat(type_a, type_b, 4096, bs, 14336, {1, 1}, {1, 1})); - } - } - } - - // qwen3-30b-a3b - for (int bs : {1, 4, 8, 32, 64, 128, 256, 512}) { - for (ggml_type type_a : {GGML_TYPE_F32, GGML_TYPE_F16, GGML_TYPE_Q4_0, GGML_TYPE_Q8_0, GGML_TYPE_Q4_K, GGML_TYPE_Q6_K, GGML_TYPE_IQ2_XS}) { - for (ggml_type type_b : {GGML_TYPE_F32}) { - test_cases.emplace_back(new test_mul_mat_id(type_a, type_b, 128, 8, false, 768, bs, 2048)); - test_cases.emplace_back(new test_mul_mat_id_fusion(type_a, type_b, 128, 8, false, 768, bs, 2048, 1)); - } - } - } - - for (int bs : {1, 4, 8, 32, 64, 128, 256, 512}) { - for (ggml_type type_a : {GGML_TYPE_F32, GGML_TYPE_F16, GGML_TYPE_Q4_0, GGML_TYPE_Q8_0, GGML_TYPE_Q4_K, GGML_TYPE_Q6_K, GGML_TYPE_IQ2_XS}) { - for (ggml_type type_b : {GGML_TYPE_F32}) { - test_cases.emplace_back(new test_mul_mat_id(type_a, type_b, 32, 4, false, 1792, bs, 2048)); - test_cases.emplace_back(new test_mul_mat_id_fusion(type_a, type_b, 32, 4, false, 1792, bs, 2048, 1)); - } - } - } - - - // gpt-oss-20b - for (int bs : {1, 4, 8, 512}) { - for (ggml_type type_a : {GGML_TYPE_MXFP4}) { - for (ggml_type type_b : {GGML_TYPE_F32}) { - test_cases.emplace_back(new test_mul_mat_id(type_a, type_b, 32, 4, false, 2880, bs, 2880)); - test_cases.emplace_back(new test_mul_mat_id_fusion(type_a, type_b, 32, 4, false, 2880, bs, 2880, 1)); - } - } - } - - for (int K : {3, 5}) { - for (int IC : {256, 2560}) { - for (int IW_IH : {32, 64, 256}) { - if (IC == 2560 && IW_IH == 256) { - // too big - continue; - } - test_cases.emplace_back(new test_im2col(GGML_TYPE_F32, GGML_TYPE_F16, GGML_TYPE_F32, {IW_IH, IW_IH, IC, 1}, {K, K, IC, 1}, 1, 1, 1, 1, 1, 1, true)); - } - } - } - - // Qwen3-VL-8B https://github.com/ggml-org/llama.cpp/issues/17012 - test_cases.emplace_back(new test_flash_attn_ext(72, 72, 16, {1, 1}, 5776, 5776, false, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16)); - - test_cases.emplace_back(new test_flash_attn_ext(64, 64, 8, {8, 1}, 7680, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16)); - test_cases.emplace_back(new test_flash_attn_ext(64, 64, 8, {8, 1}, 7680, 4, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16)); - - for (int kv : { 4096, 8192, 16384, }) { - for (int hs : { 64, 128, }) { - for (int nr : { 1, 4, }) { - test_cases.emplace_back(new test_flash_attn_ext(hs, hs, 8, {nr, 1}, kv, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16)); - } - } - } - - for (int col : {8192, 16384, 32768, 65536, 131072, 262144, 524288}) { - for (int rows : {1, 4, 16}){ - test_cases.emplace_back(new test_soft_max(GGML_TYPE_F32, {col, rows, 1, 1}, false, false, GGML_TYPE_F32, {1, 1}, 1.0f, 0.0f)); - } - } - - test_cases.emplace_back(new test_conv_2d_dw({512, 512, 256, 1}, {3, 3, 1, 256}, 1, 1, 1, false)); - test_cases.emplace_back(new test_conv_2d_dw({512, 512, 256, 1}, {3, 3, 1, 256}, 1, 1, 1, true)); - - test_cases.emplace_back(new test_conv_transpose_2d({256, 256, 256, 1}, {3, 3, 16, 256}, 1)); - test_cases.emplace_back(new test_conv_transpose_2d({16, 16, 16, 1}, {3, 3, 8, 16}, 1)); - test_cases.emplace_back(new test_conv_transpose_2d({10, 10, 9, 1}, {3, 3, 1, 9}, 2)); - - test_cases.emplace_back(new test_mean(GGML_TYPE_F32, {256, 256, 3, 1})); - - - for (int n_token : {1, 512}) { - test_cases.emplace_back(new test_add_id(GGML_TYPE_F32, GGML_TYPE_F32, 2880, 128, 4, n_token)); - test_cases.emplace_back(new test_add_id(GGML_TYPE_F32, GGML_TYPE_F32, 2880, 32, 4, n_token)); - } - - for (bool fw : {true, false}) { // fw == forward - for (ggml_type type : {GGML_TYPE_F32, GGML_TYPE_F16}) { - for (bool ff : {false, true}) { // freq_factors - for (float v : { 0, 1 }) { - test_cases.emplace_back(new test_rope(type, {128, 32, 512, 1}, 128, GGML_ROPE_TYPE_NORMAL, 512, 1.0f, 0.0f, 1.0f, ff, v, fw)); // llama 7B - test_cases.emplace_back(new test_rope(type, {128, 64, 512, 1}, 128, GGML_ROPE_TYPE_NORMAL, 512, 1.0f, 0.0f, 1.0f, ff, v, fw)); // llama 65B - test_cases.emplace_back(new test_rope(type, { 80, 32, 512, 1}, 20, GGML_ROPE_TYPE_NEOX, 512, 1.0f, 0.0f, 1.0f, ff, v, fw)); // neox (stablelm) - test_cases.emplace_back(new test_rope(type, { 64, 8, 512, 1}, 64, GGML_ROPE_TYPE_NEOX, 512, 1.0f, 0.0f, 1.0f, ff, v, fw)); // neox (falcon 40B) - test_cases.emplace_back(new test_rope(type, {128, 12, 512, 1}, 128, GGML_ROPE_TYPE_MROPE, 512, 1.0f, 0.0f, 1.0f, ff, v, fw)); // rope_multi,m-rope (qwen2vl 2B) - test_cases.emplace_back(new test_rope(type, {128, 12, 512, 1}, 128, GGML_ROPE_TYPE_IMROPE, 512, 1.0f, 0.0f, 1.0f, ff, v, fw)); // rope_multi,imrope (qwen3vl 2B) - test_cases.emplace_back(new test_rope(type, { 80, 16, 2, 1}, 80, GGML_ROPE_TYPE_VISION, 512, 1.0f, 0.0f, 1.0f, ff, v, fw)); // rope_multi,m-rope (qwen2vl ViT) - } - } - } - } - - std::vector> reduce_rows_cases = { - { 8192, 1, 1, 1 }, - { 8192, 8192, 1, 1 }, - { 128, 8192, 1, 1 }, - }; - - for (auto it: reduce_rows_cases){ - test_cases.emplace_back(new test_mean(GGML_TYPE_F32, it)); - test_cases.emplace_back(new test_sum_rows(GGML_TYPE_F32, it)); - test_cases.emplace_back(new test_sum(GGML_TYPE_F32, it)); - } - - test_cases.emplace_back(new test_argsort(GGML_TYPE_F32, {65000, 16, 1, 1})); - test_cases.emplace_back(new test_argsort(GGML_TYPE_F32, {200000, 1, 1, 1})); - test_cases.emplace_back(new test_argsort(GGML_TYPE_F32, {200000, 16, 1, 1})); - - test_cases.emplace_back(new test_top_k(GGML_TYPE_F32, {2, 1, 1, 1}, 1)); - for (auto k : {1, 10, 40, 400}) { - for (auto nrows : {1, 16}) { - for (auto cols : {k, 1000, 65000, 200000}) { - test_cases.emplace_back(new test_top_k(GGML_TYPE_F32, {cols, nrows, 1, 1}, k)); - } - } - } - - for (auto nrows : {1, 4, 8, 16}) { - for (auto cols : {128, 1024, 4096, 8192, 16384, 32768, 65536, 131072, 200000, 2000000}) { - test_cases.emplace_back(new test_cumsum(GGML_TYPE_F32, {cols, nrows, 1, 1})); - } - } - - // Examples from granite-4.0-h-1b/ggml-model-Q8_0.gguf - test_cases.emplace_back(new test_ssm_conv(GGML_TYPE_F32, {515, 3328, 1, 1}, {4, 3328, 1, 1})); // prefill - test_cases.emplace_back(new test_ssm_conv(GGML_TYPE_F32, {4, 3328, 1, 1}, {4, 3328, 1, 1})); // generate - test_cases.emplace_back(new test_ssm_scan(GGML_TYPE_F32, 128, 64, 48, 1, 512, 1)); // prefill - test_cases.emplace_back(new test_ssm_scan(GGML_TYPE_F32, 128, 64, 48, 1, 1, 1)); // generate + // MNIST-MLP: FC1 [784, 500] x [784, 500] + test_cases.emplace_back( + new test_mul_mat(GGML_TYPE_F32, GGML_TYPE_F32, 500, 500, 784, {1, 1}, {1, 1})); + // MNIST-MLP: FC2 [500, 500] x [500, 10] + test_cases.emplace_back( + new test_mul_mat(GGML_TYPE_F32, GGML_TYPE_F32, 10, 500, 500, {1, 1}, {1, 1})); + // MNIST-MLP: bias broadcast add [500, 500] + [500] + test_cases.emplace_back( + new test_bin_bcast(ggml_add, GGML_TYPE_F32, {500, 1, 1, 1}, {1, 500, 1, 1})); + // MNIST-MLP: softmax on logits [10, 500] + test_cases.emplace_back(new test_soft_max(GGML_TYPE_F32, {10, 500, 1, 1}, false, false, + GGML_TYPE_F32, {1, 1}, 1.0f, 0.0f)); + // MNIST-MLP: argmax of logits [10, 500] + test_cases.emplace_back(new test_argmax(GGML_TYPE_F32, {10, 500, 1, 1})); - // acc - test_cases.emplace_back(new test_acc(GGML_TYPE_F32, {256, 17, 1, 1}, {256, 16, 1, 1}, -1)); - test_cases.emplace_back(new test_acc(GGML_TYPE_F32, {256, 17, 2, 3}, {256, 16, 2, 3}, -1)); - test_cases.emplace_back(new test_acc(GGML_TYPE_F32, {256, 17, 2, 3}, {128, 16, 2, 3}, -1)); - test_cases.emplace_back(new test_acc(GGML_TYPE_F32, {256, 17, 2, 3}, {256, 16, 2, 3}, 1)); - test_cases.emplace_back(new test_acc(GGML_TYPE_F32, {256, 17, 2, 3}, {128, 16, 2, 3}, 2)); - test_cases.emplace_back(new test_acc(GGML_TYPE_F32, {256, 17, 2, 3}, {64, 16, 2, 3}, 3)); + // MNIST-CNN: Conv1 [28, 28, 1, 500] x [3, 3, 1, 8], stride=1, pad=1 + test_cases.emplace_back( + new test_conv_2d({28, 28, 1, 500}, {3, 3, 1, 8}, GGML_TYPE_F32, 1, 1, 1, 1, 1, 1, false)); + // MNIST-CNN: Conv2 [14, 14, 8, 500] x [3, 3, 8, 16], stride=1, pad=1 + test_cases.emplace_back( + new test_conv_2d({14, 14, 8, 500}, {3, 3, 8, 16}, GGML_TYPE_F32, 1, 1, 1, 1, 1, 1, false)); + // MNIST-CNN: Dense [784, 500] x [784, 10] (784 = 7*7*16) + test_cases.emplace_back( + new test_mul_mat(GGML_TYPE_F32, GGML_TYPE_F32, 10, 500, 784, {1, 1}, {1, 1})); return test_cases; } -static bool test_backend(ggml_backend_t backend, test_mode mode, const char * op_names_filter, const char * params_filter, +static bool test_backend(ggml_backend_t backend, + test_mode mode, + const char * op_names_filter, + const char * params_filter, printer * output_printer) { - auto filter_test_cases = [](std::vector> & test_cases, const char * params_filter) { + auto filter_test_cases = [](std::vector> & test_cases, + const char * params_filter) { if (params_filter == nullptr) { return; } @@ -7360,16 +2584,18 @@ static bool test_backend(ggml_backend_t backend, test_mode mode, const char * op // Use reference implementation on the CPU backend for comparison using ggml_backend_cpu_set_use_ref_t = void (*)(ggml_backend_t, bool); auto * reg = ggml_backend_dev_backend_reg(ggml_backend_get_device(backend_cpu)); - auto * set_use_ref = (ggml_backend_cpu_set_use_ref_t) ggml_backend_reg_get_proc_address(reg, "ggml_backend_cpu_set_use_ref"); + auto * set_use_ref = (ggml_backend_cpu_set_use_ref_t)ggml_backend_reg_get_proc_address( + reg, "ggml_backend_cpu_set_use_ref"); if (set_use_ref) { set_use_ref(backend_cpu, true); } size_t n_ok = 0; - size_t tests_run = 0; + size_t tests_run = 0; std::vector failed_tests; for (auto & test : test_cases) { - test_status_t status = test->eval(backend, backend_cpu, op_names_filter, output_printer); + test_status_t status = + test->eval(backend, backend_cpu, op_names_filter, output_printer); if (status == test_status_t::SKIPPED || status == test_status_t::NOT_SUPPORTED) { continue; } @@ -7416,12 +2642,11 @@ static bool test_backend(ggml_backend_t backend, test_mode mode, const char * op filter_test_cases(test_cases, params_filter); // Filter out fusion cases - test_cases.erase( - std::remove_if(test_cases.begin(), test_cases.end(), [](const std::unique_ptr & tc) { - return tc->run_whole_graph(); - }), - test_cases.end() - ); + test_cases.erase(std::remove_if(test_cases.begin(), test_cases.end(), + [](const std::unique_ptr & tc) { + return tc->run_whole_graph(); + }), + test_cases.end()); for (auto & test : test_cases) { test->eval_support(backend, op_names_filter, output_printer); @@ -7455,12 +2680,8 @@ static void show_test_coverage() { std::set all_ops; for (int i = 1; i < GGML_OP_COUNT; i++) { auto op = (enum ggml_op)i; - if (op == GGML_OP_VIEW || - op == GGML_OP_RESHAPE || - op == GGML_OP_PERMUTE || - op == GGML_OP_TRANSPOSE || - op == GGML_OP_CONT || - op == GGML_OP_GLU || + if (op == GGML_OP_VIEW || op == GGML_OP_RESHAPE || op == GGML_OP_PERMUTE || + op == GGML_OP_TRANSPOSE || op == GGML_OP_CONT || op == GGML_OP_GLU || op == GGML_OP_UNARY) { continue; } @@ -7475,16 +2696,14 @@ static void show_test_coverage() { auto test_cases = make_test_cases_eval(); // Filter out fusion cases test_cases.erase( - std::remove_if(test_cases.begin(), test_cases.end(), [](const std::unique_ptr & tc) { - return tc->run_whole_graph(); - }), - test_cases.end() - ); + std::remove_if(test_cases.begin(), test_cases.end(), + [](const std::unique_ptr & tc) { return tc->run_whole_graph(); }), + test_cases.end()); std::set tested_ops; ggml_init_params params = { - /* .mem_size = */ ggml_tensor_overhead()*128 + ggml_graph_overhead(), + /* .mem_size = */ ggml_tensor_overhead() * 128 + ggml_graph_overhead(), /* .mem_base = */ NULL, /* .no_alloc = */ true, }; @@ -7533,14 +2752,18 @@ static void show_test_coverage() { } static void usage(char ** argv) { - printf("Usage: %s [mode] [-o ] [-b ] [-p ] [--output ] [--list-ops] [--show-coverage]\n", argv[0]); + printf("Usage: %s [mode] [-o ] [-b ] [-p ] [--output " + "] [--list-ops] [--show-coverage]\n", + argv[0]); printf(" valid modes:\n"); printf(" - test (default, compare with CPU backend for correctness)\n"); - printf(" - grad (compare gradients from backpropagation with method of finite differences)\n"); + printf(" - grad (compare gradients from backpropagation with method of finite " + "differences)\n"); printf(" - perf (performance evaluation)\n"); printf(" - support (probe backend operation support)\n"); printf(" op names for -o are as given by ggml_op_desc() (e.g. ADD, MUL_MAT, etc),\n"); - printf(" optionally including the full test case string (e.g. \"ADD(type=f16,ne=[1,1,8,1],nr=[1,1,1,1],nf=1)\")\n"); + printf(" optionally including the full test case string (e.g. " + "\"ADD(type=f16,ne=[1,1,8,1],nr=[1,1,1,1],nf=1)\")\n"); printf(" --output specifies output format (default: console, options: console, sql, csv)\n"); printf(" --list-ops lists all available GGML operations\n"); printf(" --show-coverage shows test coverage\n"); @@ -7622,15 +2845,17 @@ int main(int argc, char ** argv) { ggml_backend_dev_t dev = ggml_backend_dev_get(i); if (backend_filter != NULL && strcmp(backend_filter, ggml_backend_dev_name(dev)) != 0) { - output_printer->print_backend_init( - backend_init_info(i, ggml_backend_dev_count(), ggml_backend_dev_name(dev), true, "Skipping")); + output_printer->print_backend_init(backend_init_info( + i, ggml_backend_dev_count(), ggml_backend_dev_name(dev), true, "Skipping")); n_ok++; continue; } - if (backend_filter == NULL && ggml_backend_dev_type(dev) == GGML_BACKEND_DEVICE_TYPE_CPU && mode != MODE_GRAD) { - output_printer->print_backend_init(backend_init_info( - i, ggml_backend_dev_count(), ggml_backend_dev_name(dev), true, "Skipping CPU backend")); + if (backend_filter == NULL && ggml_backend_dev_type(dev) == GGML_BACKEND_DEVICE_TYPE_CPU && + mode != MODE_GRAD) { + output_printer->print_backend_init(backend_init_info(i, ggml_backend_dev_count(), + ggml_backend_dev_name(dev), true, + "Skipping CPU backend")); n_ok++; continue; } @@ -7639,25 +2864,27 @@ int main(int argc, char ** argv) { GGML_ASSERT(backend != NULL); ggml_backend_reg_t reg = ggml_backend_dev_backend_reg(dev); - auto ggml_backend_set_n_threads_fn = (ggml_backend_set_n_threads_t) ggml_backend_reg_get_proc_address(reg, "ggml_backend_set_n_threads"); + auto ggml_backend_set_n_threads_fn = + (ggml_backend_set_n_threads_t)ggml_backend_reg_get_proc_address( + reg, "ggml_backend_set_n_threads"); if (ggml_backend_set_n_threads_fn) { // TODO: better value for n_threads ggml_backend_set_n_threads_fn(backend, N_THREADS); } - size_t free, total; // NOLINT + size_t free, total; // NOLINT ggml_backend_dev_memory(dev, &free, &total); - output_printer->print_backend_init(backend_init_info(i, ggml_backend_dev_count(), ggml_backend_dev_name(dev), - false, "", ggml_backend_dev_description(dev), - total / 1024 / 1024, free / 1024 / 1024, true)); + output_printer->print_backend_init(backend_init_info( + i, ggml_backend_dev_count(), ggml_backend_dev_name(dev), false, "", + ggml_backend_dev_description(dev), total / 1024 / 1024, free / 1024 / 1024, true)); bool ok = test_backend(backend, mode, op_names_filter, params_filter, output_printer.get()); if (ok) { n_ok++; } - output_printer->print_backend_status( - backend_status_info(ggml_backend_name(backend), ok ? test_status_t::OK : test_status_t::FAIL)); + output_printer->print_backend_status(backend_status_info( + ggml_backend_name(backend), ok ? test_status_t::OK : test_status_t::FAIL)); ggml_backend_free(backend); }