Function run_program

test/opt/test_tensor_cores.py:23–26

(prg:UOp, bufs:list[Buffer])

AMX = "AMX" in DEV.arch

def run_program(prg:UOp, bufs:list[Buffer]):
  buf_uops = [UOp.new_buffer(b.device, b.size, b.dtype) for b in bufs]
  for u,b in zip(buf_uops, bufs): buffers[u] = b
  run_linear(UOp(Ops.LINEAR, src=(prg.call(*buf_uops),)))

def helper_tc_ensure_uops_and_opts_count(N: int, M:int, K:int, dtype_in:DType, dtype_out:DType, axis:int=0, tc_select:int=-1, tc_opt:int=0,
                                         ensure_triggered:bool=True):