Skip to content

Commit 7e3f0d0

Browse files
committed
Able to raise gemv
1 parent 37dd847 commit 7e3f0d0

File tree

1 file changed

+20
-0
lines changed

1 file changed

+20
-0
lines changed

generic_solver/kernel_library.mlir

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -170,6 +170,26 @@ module {
170170
kernel.yield %result : tensor<i32>
171171
}
172172

173+
// General Matrix-Vector Multiply (GEMV)
174+
kernel.defn @gemv_simple(%A: tensor<?x?xf64>, %x: tensor<?xf64>, %y: tensor<?xf64>) -> tensor<?xf64> {
175+
// Simple matrix-vector multiplication: y += A * x
176+
%result = linalg.generic {
177+
indexing_maps = [
178+
affine_map<(d0, d1) -> (d1, d0)>, // Matrix A[d0, d1]
179+
affine_map<(d0, d1) -> (d0)>, // Vector x[d1]
180+
affine_map<(d0, d1) -> (d1)> // Vector y[d0]
181+
],
182+
iterator_types = ["parallel", "reduction"]
183+
} ins(%A, %x : tensor<?x?xf64>, tensor<?xf64>)
184+
outs(%y : tensor<?xf64>) {
185+
^bb0(%a: f64, %x_val: f64, %y_val: f64):
186+
%product = arith.mulf %a, %x_val : f64
187+
%result = arith.addf %y_val, %product : f64
188+
linalg.yield %result : f64
189+
} -> tensor<?xf64>
190+
kernel.yield %result : tensor<?xf64>
191+
}
192+
173193
// Index of minimum absolute value operation definition with linalg.generic representation
174194
kernel.defn @iamin_linalg(%X: tensor<?xf32>, %init: tensor<i32>) -> tensor<i32> {
175195
// Implementation using linalg.generic

0 commit comments

Comments
 (0)