import tvm
from tvm import te

import sys


# compute
k = te.reduce_axis((0, 512), "k")
A = te.placeholder((256, 512), name="A")
B = te.placeholder((512, 128), name="B")
C = te.compute((256, 128), lambda x, y: te.sum(A[x, k] * B[k, y], axis=k), name="C")


# schedule
s = te.create_schedule(C.op)


# lower and build

# target = "cuda"
stmt = tvm.lower(s,[A, B, C], simple_mode=True)



# func = tvm.build(s, [A, B, C])
# print(func)