Skip to content

Instantly share code, notes, and snippets.

@michel-steuwer

michel-steuwer/naive_mm.py

Created September 11, 2020 15:52
Show Gist options
  • Save michel-steuwer/8816c7ff13e2d7f0f903f062e9085c94 to your computer and use it in GitHub Desktop.
Save michel-steuwer/8816c7ff13e2d7f0f903f062e9085c94 to your computer and use it in GitHub Desktop.
Download ZIP
TVM Matrix Multiplication Example
Raw
naive_mm.py
# Naive algorithm
k = tvm.reduce_axis((0, K), 'k')
A = tvm.placeholder((M, K), name='A')
B = tvm.placeholder((K, N), name='B')
C = tvm.compute((M, N),lambda x, y: tvm.sum(A[x, k] * B[k, y], axis=k),name='C')
# Default schedule
s = tvm.create_schedule(C.op)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment