TVM TIR

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# 2021-09-08 18:19
import tvm
from tvm import te
import numpy as np

N = 10

dtype = "float32"
A = te.placeholder((N,), name="A")
B = te.placeholder((N,), name="B")


def test_ir_builder(A, B, C):
    ib = tvm.tir.ir_builder.create()
    Aptr = ib.buffer_ptr(A)
    Bptr = ib.buffer_ptr(B)
    Cptr = ib.buffer_ptr(C)
    with ib.for_range(0, N, name="k") as i:
        Cptr[i] = Aptr[i] + Bptr[i]
    return ib.get()


C = te.extern(
    A.shape,
    [A, B],
    lambda ins, outs: test_ir_builder(ins[0], ins[1], outs[0]),
    name="add",
    dtype=dtype,
)
s = te.create_schedule(C.op)

fadd = tvm.build(s, [A, B, C], "llvm")

dev = tvm.cpu(0)
a = tvm.nd.array(np.ones(N).astype(A.dtype), dev)
b = tvm.nd.array(np.ones(N).astype(B.dtype), dev)
c = tvm.nd.array(np.zeros(N, dtype=C.dtype), dev)
fadd(a, b, c)

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

import tvm
from tvm import te
import numpy as np


def test_ir_builder(C):
    ib = tvm.tir.ir_builder.create()
    Cptr = ib.buffer_ptr(C)
    with ib.for_range(0, 10, name="i") as i:
        with ib.for_range(0, 10, name="j") as j:
            Cptr[i * 10 + j] = tvm.tir.call_extern("int32", "get_value", i * 3 + j * 1)
    return ib.get()


C = te.extern(
    (100,),
    [],
    lambda ins, outs: test_ir_builder(outs[0]),
    name="test",
    dtype="int32",
)

s = te.create_schedule(C.op)
fadd = tvm.build(s, [C], "c")
print(fadd.get_source())