loop_tool_tui_demo.py

import loop_tool as lt
import numpy as np
import timeit


s = lt.SymbolGenerator()

a_np = np.random.randn(128, 128)
b_np = np.random.randn(128, 128)
a = lt.Tensor(a_np).to(s.m, s.k)
b = lt.Tensor(b_np).to(s.k, s.n)

c = (a * b).sum(s.k)
c.resolve()
c_np = np.dot(a_np, b_np)

assert np.allclose(c_np, c.numpy(), atol=0.001, rtol=0.001)
print("pass.")

lt.ui(c, 'test.c')

iters = 300
def bench_lt():
  def f():
    c = (a * b).sum(s.k)
    c.resolve()
  return timeit.timeit(f, number=iters)

def bench_np():
  def f():
    c_np = a_np @ b_np
  return timeit.timeit(f, number=iters)

print(f"ran {iters} iterations of numpy in {bench_np():.2f} seconds")
print(f"ran {iters} iterations of your code in {bench_lt():.2f} seconds!")