Implement 2D pooling operation with max and average pooling options

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) |
 

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`                                   |

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

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

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

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:

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("/")