Lecture 24: Machine Learning Compiler and Deployment

!pip install --pre  mlc-ai-nightly-cu118 mlc-chat-nightly-cu118 -f https://mlc.ai/wheels

import tvm
from tvm.ir.module import IRModule
from tvm.script import tir as T
import numpy as np

def lnumpy_add(a, b, c):
    for i in range(128):
        c[i] = a[i] + b[i]

from tvm import te

A = te.placeholder(shape=(128,), dtype="float32", name="A")
B = te.placeholder(shape=(128,), dtype="float32", name="B")
C = te.compute((128,), lambda i: A[i] + B[i], name="C")
func = te.create_prim_func([A, B, C])
func = func.with_attr("global_symbol", "main")
ir_module = IRModule({"main": func})

ir_module.show()

rt_mod = tvm.build(ir_module, target="llvm")  # The module for CPU backends.
print(type(rt_mod))

func = rt_mod["main"]

func

a = tvm.nd.array(np.arange(128, dtype="float32"))

b = tvm.nd.array(np.ones(128, dtype="float32"))

c = tvm.nd.empty((128,), dtype="float32")

 c

func(a, b, c)

print(a)
print(b)
print(c)

sch = tvm.tir.Schedule(ir_module)

sch.mod.show()

blockC = sch.get_block("C")
i, = sch.get_loops(blockC)

i0, i1 = sch.split(i, factors=[None, 8])

sch.mod.show()

sch.reorder(i1, i0)

sch.mod.show()

sch.reorder(i0, i1)

sch.mod.show()

sch.parallel(i0)
sch.vectorize(i1)

sch.mod.show()

sch = tvm.tir.Schedule(ir_module)
print(type(sch))

ir_module.show()

sch = tvm.tir.Schedule(ir_module)
block_C = sch.get_block("C")
i, = sch.get_loops(block_C)
i0, i1 = sch.split(i, [None, 8])
sch.reorder(i1, i0)
sch.parallel(i1)
sch.mod.show()

# Get block by its name
block_c = sch.get_block("C")
# Get loops surronding the block
(i,) = sch.get_loops(block_c)
# Tile the loop nesting.
i_0, i_1, i_2 = sch.split(i, factors=[None, 4, 4])
sch.mod.show()

sch.reorder(i_1, i_0, i_2)
sch.mod.show()

sch.vectorize(i_2)
sch.mod.show()

M = 1024
K = 1024
N = 1024

# The default tensor type in tvm
dtype = "float32"

target = "llvm"
dev = tvm.device(target, 0)

# Algorithm
k = te.reduce_axis((0, K), "k")
A = te.placeholder((M, K), name="A")
B = te.placeholder((K, N), name="B")
C = te.compute((M, N), lambda m, n: te.sum(A[m, k] * B[k, n], axis=k), name="C")

# Default schedule
func = te.create_prim_func([A, B, C])
func = func.with_attr("global_symbol", "main")
ir_module = IRModule({"main": func})
ir_module.show()

func = tvm.build(ir_module, target="llvm")  # The module for CPU backends.
a_np = np.random.rand(M, K).astype(dtype)
b_np = np.random.rand(K, N).astype(dtype)
a = tvm.nd.array(a_np, dev)
b = tvm.nd.array(b_np, dev)
c = tvm.nd.array(np.zeros((M, N), dtype=dtype), dev)
func(a, b, c)

evaluator = func.time_evaluator("main", dev, number=3)
print("Baseline time cost %g sec" % evaluator(a, b, c).mean)

c.numpy() - a_np @b_np

ir_module.show()

def transform(sch, tile_m, tile_n):
    block_C = sch.get_block("C")
    m, n, k = sch.get_loops(block_C)
    mo, mi = sch.split(m, [None, tile_m])
    no, ni = sch.split(n, [None, tile_n])
    sch.reorder(mo, no, k, mi, ni)
    return sch
sch = tvm.tir.Schedule(ir_module)
sch = transform(sch, 8, 8)
sch.mod.show()

sch = tvm.tir.Schedule(ir_module)
sch = transform(sch, 32, 32)
sch.mod.show()
mod = tvm.build(sch.mod, target="llvm")
new_eval = mod.time_evaluator("main", number=3, dev=tvm.cpu())
print("Transformed time cost %g sec" % new_eval(a, b, c).mean)

!git lfs install

!mkdir -p dist/prebuilt
!git clone https://github.com/mlc-ai/binary-mlc-llm-libs.git dist/prebuilt/lib

!cd dist/prebuilt && git clone https://huggingface.co/mlc-ai/mlc-chat-Llama-2-7b-chat-hf-q4f16_1

from mlc_chat import ChatModule
from mlc_chat.callback import StreamToStdout

cm = ChatModule(model="Llama-2-7b-chat-hf-q4f16_1")

output = cm.generate(
    prompt="When was Python released?",
    progress_callback=StreamToStdout(callback_interval=2),
)

prompt = input("Prompt: ")
output = cm.generate(prompt=prompt, progress_callback=StreamToStdout(callback_interval=2))

print(cm.stats())