Better torch integration (#66)

* Initial implementation of the nn.Model patching mechanism

* Start the support for native pytorch functions

* Testing NPU Module in a challenging scenario

* Remove model saving

* Convert module directly from pytorch

* Further additions

* Support reshape using sequence or multiple args

* Add l2 and gelu

* Add additional bindings

* Support both gelu operators

* Add additional methods

* Add adaptive pool and flatten

* Implemented avg pooling

* Adding max pooling

* Add common operations

* Fix CI

* Implementation of batchnorm

* Different sdpa tests

Author: alessandropalla

include/intel_npu_acceleration_library/nn_factory.h

@@ -122,6 +122,75 @@ class ModelFactory : public intel_npu_acceleration_library::OVInferenceModel {
         return conv.get();
     }
 
+    /**
+     * @brief Create a new average pooling operation
+     * @param input pooling input
+     * @param strides pooling strides
+     * @param pads_begin pooling padding begin
+     * @param pads_ends pooling padding end
+     * @param kernel pooling kernel
+     * @param exclude_pad exclude padding from the average calculation
+     * @param rounding_type rounding type
+     * @param auto_pad padding type
+     * @return ov::op::Op*
+     */
+    ov::op::Op* average_pooling(ov::op::Op* input, std::vector<size_t> strides, std::vector<size_t> pads_begin,
+                                std::vector<size_t> pads_ends, std::vector<size_t> kernel, bool exclude_pad = false,
+                                ov::op::RoundingType rounding_type = ov::op::RoundingType::FLOOR,
+                                ov::op::PadType auto_pad = ov::op::PadType::EXPLICIT) {
+        auto pool = std::make_shared<ov::opset1::AvgPool>(input->output(0), ov::Strides(strides), pads_begin, pads_ends,
+                                                          kernel, exclude_pad, rounding_type, auto_pad);
+        operations.push_back(pool);
+        return pool.get();
+    }
+
+    /**
+     * @brief Create a new adaptive average pooling operation
+     * @param input pooling input
+     * @param output_shape output shape
+     * @return ov::op::Op*
+     */
+    ov::op::Op* adaptive_average_pool(ov::op::Op* input, ov::op::Op* output_shape) {
+        auto pool = std::make_shared<ov::opset8::AdaptiveAvgPool>(input->output(0), output_shape->output(0));
+        operations.push_back(pool);
+        return pool.get();
+    }
+
+    /**
+     * @brief Create a new max pooling operation
+     * @param input pooling input
+     * @param strides pooling strides
+     * @param pads_begin pooling padding begin
+     * @param pads_ends pooling padding end
+     * @param kernel pooling kernel
+     * @param exclude_pad exclude padding from the max calculation
+     * @param rounding_type rounding type
+     * @param auto_pad padding type
+     * @return ov::op::Op*
+     */
+    ov::op::Op* max_pooling(ov::op::Op* input, std::vector<size_t> strides, std::vector<size_t> pads_begin,
+                            std::vector<size_t> pads_ends, std::vector<size_t> kernel,
+                            ov::op::RoundingType rounding_type = ov::op::RoundingType::FLOOR,
+                            ov::op::PadType auto_pad = ov::op::PadType::EXPLICIT) {
+        auto pool = std::make_shared<ov::opset1::MaxPool>(input->output(0), ov::Strides(strides), pads_begin, pads_ends,
+                                                          kernel, rounding_type, auto_pad);
+        operations.push_back(pool);
+        return pool.get();
+    }
+
+    /**
+     * @brief Create a new adaptive max pooling operation
+     * @param input pooling input
+     * @param output_shape output shape
+     * @return ov::op::Op*
+     */
+    ov::op::Op* adaptive_max_pool(ov::op::Op* input, ov::op::Op* output_shape) {
+        auto pool = std::make_shared<ov::opset8::AdaptiveMaxPool>(input->output(0), output_shape->output(0),
+                                                                  ov::element::i64);
+        operations.push_back(pool);
+        return pool.get();
+    }
+
     /**
      * @brief Create a new gather operation
      *
@@ -374,8 +443,8 @@ class ModelFactory : public intel_npu_acceleration_library::OVInferenceModel {
      * @param input operation's input node
      * @return ov::op::Op*
      */
-    ov::op::Op* gelu(ov::op::Op* input) {
-        auto gelu = std::make_shared<ov::opset7::Gelu>(input->output(0), ov::op::GeluApproximationMode::TANH);
+    ov::op::Op* gelu(ov::op::Op* input, ov::op::GeluApproximationMode mode) {
+        auto gelu = std::make_shared<ov::opset7::Gelu>(input->output(0), mode);
         operations.push_back(gelu);
         return gelu.get();
     }
@@ -695,11 +764,19 @@ class ModelFactory : public intel_npu_acceleration_library::OVInferenceModel {
      */
     ov::op::Op* scaled_dot_product_attention(ov::op::Op* query, ov::op::Op* key, ov::op::Op* value,
                                              ov::op::Op* attn_mask, bool is_causal) {
-        auto sdpa = std::make_shared<ov::opset13::ScaledDotProductAttention>(
-                query->output(0), key->output(0), value->output(0), attn_mask->output(0), is_causal);
-
-        operations.push_back(sdpa);
-        return sdpa.get();
+        if (attn_mask == nullptr) {
+            auto sdpa = std::make_shared<ov::opset13::ScaledDotProductAttention>(query->output(0), key->output(0),
+                                                                                 value->output(0), is_causal);
+
+            operations.push_back(sdpa);
+            return sdpa.get();
+        } else {
+            auto sdpa = std::make_shared<ov::opset13::ScaledDotProductAttention>(
+                    query->output(0), key->output(0), value->output(0), attn_mask->output(0), is_causal);
+
+            operations.push_back(sdpa);
+            return sdpa.get();
+        }
     }
 
     /**

intel_npu_acceleration_library/backend/bindings.py

@@ -180,6 +180,39 @@ def init_network_factory(lib: ctypes.CDLL):
     ]
     lib.convolution.restype = handler
 
+    lib.avg_pooling.argtypes = [
+        handler,
+        handler,
+        ctypes.c_int,
+        c_u32_array,
+        ctypes.c_int,
+        c_u32_array,
+        ctypes.c_int,
+        c_u32_array,
+        ctypes.c_int,
+        c_u32_array,
+        ctypes.c_bool,
+        ctypes.c_int,
+        ctypes.c_int,
+    ]
+    lib.avg_pooling.restype = handler
+
+    lib.max_pooling.argtypes = [
+        handler,
+        handler,
+        ctypes.c_int,
+        c_u32_array,
+        ctypes.c_int,
+        c_u32_array,
+        ctypes.c_int,
+        c_u32_array,
+        ctypes.c_int,
+        c_u32_array,
+        ctypes.c_int,
+        ctypes.c_int,
+    ]
+    lib.max_pooling.restype = handler
+
     for op in get_supported_ops():
         fn = getattr(lib, op.name)
         fn.argtypes = [handler] * (op.inputs + 1) + list(op.parameters)

intel_npu_acceleration_library/backend/factory.py

@@ -13,6 +13,7 @@
 import numpy.typing as npt
 import numpy as np
 import ctypes
+import torch
 
 
 F = TypeVar("F", bound=Callable[..., Any])
@@ -165,12 +166,12 @@ def to(self, tensor: ctypes._Pointer, dtype: npt.DTypeLike) -> ctypes._Pointer:
     @return_tensor
     def constant(
         self,
-        data: Union[np.array, Sequence[int], Sequence[float], int, float],
+        data: Union[np.array, Sequence[int], Sequence[float], int, float, torch.Tensor],
     ) -> ctypes._Pointer:
         """Generate a model input constant.
 
         Args:
-            data (Union[np.array, Sequence[int], Sequence[float], int, float]): constant data
+            data (Union[np.array, Sequence[int], Sequence[float], int, float, torch.Tensor]): constant data
 
         Returns:
             ctypes._Pointer: an instance to a constant object
@@ -185,6 +186,8 @@ def constant(
             data = np.array([data], dtype=np.int64)
         elif isinstance(data, float):
             data = np.array([data], dtype=np.float32)
+        elif isinstance(data, torch.Tensor):
+            data = data.detach().numpy()
 
         dst = data.ctypes.data_as(ctypes.c_void_p)
         shape_ptr = np.array(data.shape, dtype=np.uint32)
@@ -395,6 +398,139 @@ def normL2(
         axis_node = self.constant(axis).node  # type: ignore
         return backend_lib.normL2(self._mm, input_node, axis_node, eps)
 
+    @return_tensor
+    def avg_pooling(
+        self,
+        input: ctypes._Pointer,
+        kernel_size: Union[int, Sequence[int]],
+        strides: Optional[Union[int, Sequence[int]]] = None,
+        padding: int = 0,
+        ceil_mode: bool = False,
+        count_include_pad: bool = True,
+        divisor_override: Optional[int] = None,
+        n_spatial_dims: int = 2,
+    ) -> ctypes._Pointer:
+        """Generate an average pooling layer.
+
+        Args:
+            input (ctypes._Pointer): layer input node
+            kernel_size (Sequence[int]): kernel size
+            strides (Sequence[int]): strides
+            padding (int): padding
+            ceil_mode (bool): ceil mode
+            count_include_pad (bool): count include pad
+            divisor_override (int): divisor override
+            n_spatial_dims (int): number of spatial dimensions
+
+        Raises:
+            NotImplementedError: divisor_override is not supported
+
+        Returns:
+            ctypes._Pointer: output node
+        """
+        if isinstance(kernel_size, int):
+            kernel_size = [kernel_size] * n_spatial_dims
+
+        if strides is None:
+            strides = kernel_size
+        elif isinstance(strides, int):
+            strides = [strides] * n_spatial_dims
+
+        if isinstance(padding, int):
+            padding_begins = [padding] * n_spatial_dims
+            padding_ends = [padding] * n_spatial_dims
+        else:
+            padding_begins = list(padding)
+            padding_ends = list(padding)
+
+        strides_ptr = np.array(strides, dtype=np.uint32)
+        padding_begins_ptr = np.array(padding_begins, dtype=np.uint32)
+        padding_ends_ptr = np.array(padding_ends, dtype=np.uint32)
+        kernel_size_ptr = np.array(kernel_size, dtype=np.uint32)
+
+        rounding_type = 1 if ceil_mode else 0
+        auto_pad = 0  # Hardcoded to explicit padding
+
+        if divisor_override:
+            raise NotImplementedError("divisor_override is not supported")
+
+        return backend_lib.avg_pooling(
+            self._mm,
+            input,
+            strides_ptr.size,
+            strides_ptr,
+            padding_begins_ptr.size,
+            padding_begins_ptr,
+            padding_ends_ptr.size,
+            padding_ends_ptr,
+            kernel_size_ptr.size,
+            kernel_size_ptr,
+            not count_include_pad,  # exclude_pad
+            rounding_type,  # rounding_type
+            auto_pad,  # auto_pad
+        )
+
+    @return_tensor
+    def max_pooling(
+        self,
+        input: ctypes._Pointer,
+        kernel_size: Union[int, Sequence[int]],
+        strides: Optional[Union[int, Sequence[int]]] = None,
+        padding: int = 0,
+        ceil_mode: bool = False,
+        n_spatial_dims: int = 2,
+    ) -> ctypes._Pointer:
+        """Generate an average pooling layer.
+
+        Args:
+            input (ctypes._Pointer): layer input node
+            kernel_size (Sequence[int]): kernel size
+            strides (Sequence[int]): strides
+            padding (int): padding
+            ceil_mode (bool): ceil mode
+            n_spatial_dims (int): number of spatial dimensions
+
+        Returns:
+            ctypes._Pointer: output node
+        """
+        if isinstance(kernel_size, int):
+            kernel_size = [kernel_size] * n_spatial_dims
+
+        if strides is None:
+            strides = kernel_size
+        elif isinstance(strides, int):
+            strides = [strides] * n_spatial_dims
+
+        if isinstance(padding, int):
+            padding_begins = [padding] * n_spatial_dims
+            padding_ends = [padding] * n_spatial_dims
+        else:
+            padding_begins = list(padding)
+            padding_ends = list(padding)
+
+        strides_ptr = np.array(strides, dtype=np.uint32)
+        padding_begins_ptr = np.array(padding_begins, dtype=np.uint32)
+        padding_ends_ptr = np.array(padding_ends, dtype=np.uint32)
+        kernel_size_ptr = np.array(kernel_size, dtype=np.uint32)
+
+        rounding_type = 1 if ceil_mode else 0
+        auto_pad = 0  # Hardcoded to explicit padding
+
+        return backend_lib.max_pooling(
+            self._mm,
+            input,
+            strides_ptr.size,
+            strides_ptr,
+            padding_begins_ptr.size,
+            padding_begins_ptr,
+            padding_ends_ptr.size,
+            padding_ends_ptr,
+            kernel_size_ptr.size,
+            kernel_size_ptr,
+            rounding_type,  # rounding_type
+            auto_pad,  # auto_pad
+        )
+
     def get_output_tensor_shape(self):
         """Get output tensor shape.
 

intel_npu_acceleration_library/backend/ops.py

@@ -52,6 +52,7 @@ def get_supported_ops() -> List[SupportedOp]:
         SupportedOp(name="floor_act", inputs=1),
         SupportedOp(name="grn", inputs=1, parameters=[ctypes.c_float]),
         SupportedOp(name="gelu", inputs=1),
+        SupportedOp(name="gelu_erf", inputs=1),
         SupportedOp(name="log_act", inputs=1),
         SupportedOp(name="negative", inputs=1),
         SupportedOp(name="relu", inputs=1),
@@ -79,6 +80,11 @@ def get_supported_ops() -> List[SupportedOp]:
             inputs=4,
             parameters=[ctypes.c_bool],
         ),
+        SupportedOp(
+            name="scaled_dot_product_attention_simple",
+            inputs=3,
+            parameters=[ctypes.c_bool],
+        ),
         SupportedOp(
             name="normL2",
             inputs=2,
@@ -93,5 +99,7 @@ def get_supported_ops() -> List[SupportedOp]:
         SupportedOp(name="transpose", inputs=2),
         SupportedOp(name="squeeze", inputs=1),
         SupportedOp(name="unsqueeze", inputs=2),
+        SupportedOp(name="adaptive_avg_pool", inputs=2),
+        SupportedOp(name="adaptive_max_pool", inputs=2),
     ]
     return supported_ops