Module: TinyNNMetal

Defined in:
lib/toy/ffi/tinynn_metal.rb

Class Method Summary collapse

Class Method Details

.download_matmul(sess, tensor, m, n) ⇒ Object 



337
338
339
340
341
342
343
344
345
346
347
348
349
350
 
# File 'lib/toy/ffi/tinynn_metal.rb', line 337

def self.download_matmul(sess, tensor, m, n)
  TinyNNMetal.tnn_download(sess, tensor)
  out = Mat.new(m, n)
  i = 0
  while i < m
    j = 0
    while j < n
      out.flat[i * n + j] = TinyNNMetal.tnn_scratch_get(sess, j * m + i)
      j = j + 1
    end
    i = i + 1
  end
  out
end

.download_row_major(sess, dl_handle, rows, cols) ⇒ Object 



319
320
321
322
323
324
325
326
327
328
329
 
# File 'lib/toy/ffi/tinynn_metal.rb', line 319

def self.download_row_major(sess, dl_handle, rows, cols)
  TinyNNMetal.tnn_download(sess, dl_handle)
  out = Mat.new(rows, cols)
  n = rows * cols
  i = 0
  while i < n
    out.flat[i] = TinyNNMetal.tnn_scratch_get(sess, i)
    i = i + 1
  end
  out
end

.download_to_mat(sess, dl_handle, rows, cols) ⇒ Object 



331
332
333
334
335
 
# File 'lib/toy/ffi/tinynn_metal.rb', line 331

def self.download_to_mat(sess, dl_handle, rows, cols)
  out = Mat.new(rows, cols)
  TinyNNMetal.tnn_download_to_f64_array(sess, dl_handle, out.flat, rows * cols)
  out
end

.stage_row_major_and_upload(sess, target, m) ⇒ Object 



356
357
358
 
# File 'lib/toy/ffi/tinynn_metal.rb', line 356

def self.stage_row_major_and_upload(sess, target, m)
  TinyNNMetal.tnn_upload_from_float_array(sess, target, m.flat, m.nrows * m.ncols)
end

.stage_transposed_and_upload(sess, target, b) ⇒ Object 



352
353
354
 
# File 'lib/toy/ffi/tinynn_metal.rb', line 352

def self.stage_transposed_and_upload(sess, target, b)
  TinyNNMetal.tnn_upload_transposed_f64(sess, target, b.flat, b.nrows, b.ncols)
end

.upload_int_array(sess, tensor, indices) ⇒ Object 



300
301
302
 
# File 'lib/toy/ffi/tinynn_metal.rb', line 300

def self.upload_int_array(sess, tensor, indices)
  TinyNNMetal.tnn_upload_from_int_array(sess, tensor, indices, indices.length)
end

.upload_row_major(sess, tensor, mat) ⇒ Object

Ruby-level convenience helpers — these are pure-Ruby wrappers over the FFI surface. The mirror generator only rewrites the per-method ‘TinyNN.` qualifiers, not the helper-method bodies themselves; callers in the SmolLM2 mirror reference `TinyNNMetal.upload_int_array` etc., so we define them here. Bodies are identical to lib/toy/ffi/tinynn.rb.



296
297
298
 
# File 'lib/toy/ffi/tinynn_metal.rb', line 296

def self.upload_row_major(sess, tensor, mat)
  TinyNNMetal.tnn_upload_from_float_array(sess, tensor, mat.flat, mat.nrows * mat.ncols)
end

.upload_transposed(sess, tensor, mat) ⇒ Object 



304
305
306
307
308
309
310
311
312
313
314
315
316
317
 
# File 'lib/toy/ffi/tinynn_metal.rb', line 304

def self.upload_transposed(sess, tensor, mat)
  br = mat.nrows
  bc = mat.ncols
  i = 0
  while i < br
    j = 0
    while j < bc
      TinyNNMetal.tnn_scratch_set(sess, j * br + i, mat.flat[i * bc + j])
      j = j + 1
    end
    i = i + 1
  end
  TinyNNMetal.tnn_upload(sess, tensor)
end