kernel working and tested

ggml-org · Nov 20, 2024 · d782d29 · d782d29
1 parent 9cf977d
commit d782d29
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 5 deletions.
diff --git a/src/ggml-metal/ggml-metal.m b/src/ggml-metal/ggml-metal.m
@@ -3485,7 +3485,8 @@ static void ggml_metal_encode_node(
                 [encoder setComputePipelineState:pipeline];
                 [encoder setBuffer:id_src0 offset:offs_src0        atIndex:0];
                 [encoder setBuffer:id_dst  offset:offs_dst         atIndex:1];
-                [encoder setBytes:&nb01    length:sizeof( int64_t) atIndex:2];
+                [encoder setBytes:&ne00    length:sizeof( int64_t) atIndex:2];
+                [encoder setBytes:&nb01    length:sizeof(uint64_t) atIndex:3];
 
                 [encoder dispatchThreadgroups:MTLSizeMake(nrows, 1, 1) threadsPerThreadgroup:MTLSizeMake(1, 1, 1)];
             } break;

diff --git a/src/ggml-metal/ggml-metal.metal b/src/ggml-metal/ggml-metal.metal
@@ -1347,15 +1347,16 @@ kernel void kernel_ssm_scan_f32(
 kernel void kernel_argmax(
         device   const void * x,
         device      int32_t * dst,
-        constant    int64_t & nb01,
+        constant    int64_t & ncols,
+        constant   uint64_t & nb01,
         uint tgpig[[threadgroup_position_in_grid]]) {
     device const float * x_row = (device const float *) ((device const char *) x + tgpig * nb01);
 
-    dst[tpitg] = 0;
+    dst[tgpig] = 0;
 
     for (int i = 0; i < ncols; i++) {
-        if (x_row[i] > x_row[dst[tpitg]]) {
-            dst[tpitg] = i;
+        if (x_row[i] > x_row[dst[tgpig]]) {
+            dst[tgpig] = i;
         }
     }
 }