Lecture 24: Machine Learning Compiler and Deployment

!pip install --pre  mlc-ai-nightly-cu121 mlc-llm-nightly-cu121 -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()

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

from mlc_llm import MLCEngine

# Create engine
model = "HF://mlc-ai/Llama-3-8B-Instruct-q4f16_1-MLC"
engine = MLCEngine(model)

# Run chat completion in OpenAI API.
for response in engine.chat.completions.create(
    messages=[{"role": "user", "content": "Tell me about cmu?"}],
    model=model,
    stream=True,
):
    for choice in response.choices:
        print(choice.delta.content, end="", flush=True)
print("\n")