feat: add maxpool_2d implementation using linalg

Author: xTayEx

frontend/Python/graph/type.py

@@ -35,7 +35,7 @@ class TensorDType(Enum):
     - Bool: str
         Represents the boolean data type.
     """
-    
+
     Int8 = "int8"
     Int32 = "int32"
     Int64 = "int64"
@@ -47,7 +47,7 @@ class TensorDType(Enum):
 
 class TensorMeta:
     """
-    Store tensor metadata, including shape and data type, while overlooking raw 
+    Store tensor metadata, including shape and data type, while overlooking raw
     data.
 
     Attributes:
@@ -58,7 +58,7 @@ class TensorMeta:
 
     Methods:
     - __init__(shape: tuple, dtype: str) -> None:
-        Initializes a new instance of the TensorMeta class with the specified 
+        Initializes a new instance of the TensorMeta class with the specified
         shape and data type.
 
     Example:
@@ -79,6 +79,7 @@ def __init__(self, shape, dtype) -> None:
         self.shape = shape
         self.dtype = dtype
 
+
 class DeviceType(Enum):
     """
     Enumeration class representing different types of devices.
@@ -91,6 +92,7 @@ class DeviceType(Enum):
     Each attribute represents a specific device type and is associated with a
     string value.
     """
-    CPU = 'cpu'
-    GPU = 'gpu'
-    UNKNOW = 'unknow'
+
+    CPU = "cpu"
+    GPU = "gpu"
+    UNKNOWN = "unknow"

frontend/Python/ops/linalg.py

@@ -2560,13 +2560,10 @@ def convolution2d_op(
     dilations = node.args[5]
 
     input_val = symbol_table.get((str(input_), 0))
-    input_shape = list(ir.RankedTensorType(input_val.type).shape)
     filter_val = symbol_table.get((str(filter_), 0))
-    filter_shape = list(ir.RankedTensorType(filter_val.type).shape)
     dtype = node.tensor_meta["dtype"]
     result_element_type = mlir_element_type_get(dtype)
     out_shape = node.tensor_meta["shape"]
-    result_tensor_type = ir.RankedTensorType.get(out_shape, result_element_type)
     strides_attr = ir._denseI64ArrayAttr(strides, None)
     dilations_attr = ir._denseI64ArrayAttr(dilations, None)
     conv2d_result = tensor.EmptyOp(out_shape, result_element_type)
@@ -2582,7 +2579,6 @@ def convolution2d_op(
     if len(node._parents) > 2:
         bias_tensor = symbol_table.get((str(bias), 0))
         init = tensor.EmptyOp(out_shape, result_element_type)
-        print(f"{bias_tensor}, {init}, {out_shape}")
         broadcasted = linalg.broadcast(
             bias_tensor, outs=[init], dimensions=[0, 2, 3]
         )
@@ -2592,6 +2588,40 @@ def convolution2d_op(
     return op_to_return
 
 
+def maxpool2d_op(
+    node: Conv2dOp, symbol_table: Dict[Tuple[str, int], ir.Operation]
+):
+    # print(node.kwargs, node.args)
+    input_ = node.args[0]
+    kernel_size = node.args[1]
+    strides = node.args[2]
+    dtype = node.tensor_meta["dtype"]
+    result_element_type = mlir_element_type_get(dtype)
+    result_shape = node.tensor_meta["shape"]
+
+    input_value = symbol_table.get((str(input_), 0))
+    kernel_size_value = tensor.EmptyOp(kernel_size, result_element_type)
+
+    if len(node.args) > 3:
+        dilations = node.args[4]
+    else:
+        dilations = [1, 1]
+
+    strides_attr = ir._denseI64ArrayAttr(strides, None)
+    dilations_attr = ir._denseI64ArrayAttr(dilations, None)
+
+    result = tensor.EmptyOp(result_shape, result_element_type)
+    op = linalg.pooling_nchw_max(
+        input_value,
+        kernel_size_value,
+        outs=[result],
+        strides=strides_attr,
+        dilations=dilations_attr,
+    )
+
+    return op
+
+
 ops_registry = {
     "MatmulOp": matmul_op,
     "TransposeMatmulFusedOp": matmul_transpose_b_op,
@@ -2635,4 +2665,5 @@ def convolution2d_op(
     "CopyOp": copy_op,
     "SliceScatterOp": slice_scatter_op,
     "Conv2dOp": convolution2d_op,
+    "MaxPool2dOp": maxpool2d_op,
 }