Merge pull request #52 from flash-algo/add-gemv-pytorch-kernel

[FEATURE SUPPORT] add gemv pytorch kernel

Author: LoserCheems

README.md

@@ -22,7 +22,7 @@ The following common BLAS kernels have been implemented in multiple frameworks.
 | [scal](./docs/scal.md) | scale vector | $y = \alpha y$ | $n$ | $2n$ | [✅](./kernel_course/python_ops/scal.py) | [✅](./kernel_course/pytorch_ops/scal.py) | [✅](./kernel_course/triton_ops/scal.py) | ❌ | [✅](./tests/test_scal.py) |
 | [axpby](./docs/axpby.md) | update vector| $y = \alpha x + \beta y$ | $3n$ | $3n$ | [✅](./kernel_course/python_ops/axpby.py) | [✅](./kernel_course/pytorch_ops/axpby.py) | [✅](./kernel_course/triton_ops/axpby.py) | ❌ | [✅](./tests/test_axpby.py) |
 | [dot](./docs/dot.md) | dot product | $z = x^\top y$ | $2n$ | $2n$ | [✅](./kernel_course/python_ops/dot.py) | [✅](./kernel_course/pytorch_ops/dot.py) | [✅](./kernel_course/triton_ops/dot.py) | ❌ | [✅](./tests/test_dot.py) |
-| [gemv](./docs/gemv.md) | general matrix-vector multiply | $y = \alpha A x + \beta y$ | $2mn$ | $mn + n + 2m$ | [✅](./kernel_course/python_ops/gemv.py) | ❌ | ❌ | ❌ | ❌ |
+| [gemv](./docs/gemv.md) | general matrix-vector multiply | $y = \alpha A x + \beta y$ | $2mn$ | $mn + n + 2m$ | [✅](./kernel_course/python_ops/gemv.py) | [✅](./kernel_course/pytorch_ops/gemv.py) | ❌ | ❌ | ❌ |
 | geru | general rank-1 update | $A = A + \alpha x y^\top$ | $2mn$ | $2mn + m + n$ | ❌ | ❌ | ❌ | ❌ | ❌ |
 | gemm | general matrix-matrix multiply | $C = \alpha A B + \beta C$ | $2mnk$ | $mk + nk + 2mn$ | ❌ | ❌ | ❌ | ❌ | ❌ |
 

kernel_course/python_ops/gemv.py

@@ -18,7 +18,7 @@ def gemv(
         y (torch.Tensor): Vector tensor to be updated.
         alpha (float): Scaling factor for the product of `A` and `x`.
         beta (float): Scaling factor for `y`.
-    
+
     Returns:
         torch.Tensor: The updated tensor `y`.
     """

kernel_course/pytorch_ops/gemv.py

@@ -0,0 +1,28 @@
+import torch
+
+
+def gemv(
+    A: torch.Tensor,
+    x: torch.Tensor,
+    y: torch.Tensor,
+    alpha: float,
+    beta: float,
+) -> torch.Tensor:
+    """
+    Updates tensor `y` by adding the product of matrix `A` and vector `x`
+    scaled by `alpha`, and `y` scaled by `beta` using PyTorch operations.
+
+    Args:
+        A (torch.Tensor): Matrix tensor.
+        x (torch.Tensor): Vector tensor to be multiplied with `A`.
+        y (torch.Tensor): Vector tensor to be updated.
+        alpha (float): Scaling factor for the product of `A` and `x`.
+        beta (float): Scaling factor for `y`.
+
+    Returns:
+        torch.Tensor: The updated tensor `y`.
+    """
+
+    y = torch.add(torch.mul(torch.matmul(A, x), alpha), torch.mul(y, beta))
+
+    return y