Add support for 1D batched-GEMV broadcasting

Author: affifboudaoud

dace/frontend/python/replacements/linalg.py

@@ -86,6 +86,34 @@ def _matmult(visitor: ProgramVisitor, sdfg: SDFG, state: SDFGState, op1: str, op
 
         output_shape = (1, )
 
+    elif len(arr1.shape) == 1 and len(arr2.shape) > 2:  # vector @ batched matrix (e.g., [k] @ [b, k, n])
+
+        res = symbolic.equal(arr1.shape[0], arr2.shape[-2])
+        if res is None:
+            warnings.warn(
+                f'Length of vector {arr1.shape[0]} and second-last dimension of tensor {arr2.shape[-2]} '
+                f'may not match', UserWarning)
+        elif not res:
+            raise SyntaxError(f"Length of vector {arr1.shape[0]} must match "
+                              f"second-last dimension of tensor {arr2.shape[-2]}")
+
+        # Output has all batch dimensions plus the last dimension of arr2
+        output_shape = arr2.shape[:-2] + (arr2.shape[-1], )
+
+    elif len(arr1.shape) > 2 and len(arr2.shape) == 1:  # batched matrix @ vector (e.g., [b, m, k] @ [k])
+
+        res = symbolic.equal(arr1.shape[-1], arr2.shape[0])
+        if res is None:
+            warnings.warn(
+                f'Last dimension of tensor {arr1.shape[-1]} and length of vector {arr2.shape[0]} '
+                f'may not match', UserWarning)
+        elif not res:
+            raise SyntaxError(f"Last dimension of tensor {arr1.shape[-1]} must match "
+                              f"length of vector {arr2.shape[0]}")
+
+        # Output has all batch dimensions plus the second-last dimension of arr1
+        output_shape = arr1.shape[:-1]
+
     else:  # Dunno what this is, bail
 
         raise SyntaxError("Cannot multiply arrays with shapes: {} and {}".format(arr1.shape, arr2.shape))

dace/libraries/blas/blas_helpers.py

@@ -110,9 +110,23 @@ def get_gemm_opts(a_strides, b_strides, c_strides) -> Dict[str, Any]:
     #       |    |      |
     #     use these 3 to detect correct option
 
-    sAM, sAK = a_strides[-2:]
-    sBK, sBN = b_strides[-2:]
-    sCM, sCN = c_strides[-2:]
+    # Handle 1D inputs by treating them as column/row vectors
+    # [k] -> treat as [k, 1] with stride [1, k] for column vector
+    if len(a_strides) == 1:
+        sAM, sAK = a_strides[0], 1  # Treat as column vector [k, 1]
+    else:
+        sAM, sAK = a_strides[-2:]
+
+    # Treat as row vector [1, k] -> transposed to [k, 1]
+    if len(b_strides) == 1:
+        sBK, sBN = 1, b_strides[0]
+    else:
+        sBK, sBN = b_strides[-2:]
+
+    if len(c_strides) == 1:
+        sCM, sCN = c_strides[0], 1
+    else:
+        sCM, sCN = c_strides[-2:]
 
     opts = {
         'mkm': {