FFI 容器

FFI 容器#

import tvm.ffi as tvm_ffi

数组#

a = tvm_ffi.convert([1, 2, 3])
assert isinstance(a, tvm_ffi.Array)
assert len(a) == 3
assert a[-1] == 3
a_slice = a[-3:-1]
assert (a_slice[0], a_slice[1]) == (1, 2)

map#

from tvm import te
a = te.var("a")
b = te.var("b")
amap = tvm_ffi.convert({a: 2, b: 3})
assert a in amap
assert len(amap) == 2
dd = dict(amap.items())
assert a in dd
assert b in dd
assert a + 1 not in amap
assert {x for x in amap} == {a, b}
assert set(amap.keys()) == {a, b}
assert set(amap.values()) == {2, 3}

数组映射#

a = tvm_ffi.convert([[1, 2, 3], {"A": 5, "B": 6}])
assert isinstance(a, tvm_ffi.Array)
assert len(a) == 2
assert isinstance(a[0], tvm_ffi.Array)
assert isinstance(a[1], tvm_ffi.Map)
assert tuple(a[0]) == (1, 2, 3)
assert a[1]["A"] == 5
assert a[1]["B"] == 6

a = tvm_ffi.convert([1, 2, 3])
assert dir(a)

a = te.var("a")
b = te.var("b")
amap = tvm_ffi.convert({a: 2, b: 3})
assert dir(amap)

a = te.var("a")
b = te.var("b")
amap = tvm_ffi.convert({a: 2, b: 3})
assert isinstance(amap, tvm.ffi.Map)

整数映射#

amap = tvm_ffi.convert({3: 2, 4: 3})
assert 3 in amap
assert len(amap) == 2
dd = dict(amap.items())
assert 3 in dd
assert 4 in dd
assert 5 not in amap
assert tuple(amap.items()) == ((3, 2), (4, 3))
assert tuple(amap.keys()) == (3, 4)
assert tuple(amap.values()) == (2, 3)

字符串映射#

data = []
for i in reversed(range(10)):
    data.append((f"a{i}", i))
amap = tvm_ffi.convert({k: v for k, v in data})
assert tuple(amap.items()) == tuple(data)
for k, v in data:
    assert k in amap
    assert amap[k] == v
    assert amap.get(k) == v

assert tuple(k for k in amap) == tuple(k for k, _ in data)

KeyError#

import pytest
amap = tvm_ffi.convert({3: 2, 4: 3})
with pytest.raises(KeyError):
    amap[5]

repr#

a = tvm_ffi.convert([1, 2, 3])
assert str(a) == "[1, 2, 3]"
amap = tvm_ffi.convert({3: 2, 4: 3})
assert str(amap) == "{3: 2, 4: 3}"

smap = tvm_ffi.convert({"a": 1, "b": 2})
assert str(smap) == "{'a': 1, 'b': 2}"

序列化#

import tvm
a = tvm_ffi.convert([1, 2, 3.5, True])
json_str = tvm.ir.save_json(a)
a_loaded = tvm.ir.load_json(json_str)
assert a_loaded[1] == 2
assert a_loaded[2] == 3.5
assert a_loaded[3] == True
assert isinstance(a_loaded[3], bool)

a = te.var("a")
b = te.var("b")
amap = tvm_ffi.convert({a: 2, b: 3})
json_str = tvm.ir.save_json(amap)
amap = tvm.ir.load_json(json_str)
assert len(amap) == 2
dd = {kv[0].name: kv[1] for kv in amap.items()}
assert dd == {"a": 2, "b": 3}

in 容器#

arr = tvm_ffi.convert(["a", "b", "c"])
assert "a" in arr
assert tvm.tir.StringImm("a") in arr
assert "d" not in arr

ndarray 容器#

x = tvm.nd.array([1, 2, 3])
arr = tvm_ffi.convert([x, x])
assert arr[0].same_as(x)
assert arr[1].same_as(x)
assert isinstance(arr[0], tvm.nd.NDArray)

ReturnsVariant#

func = tvm.get_global_func("testing.ReturnsVariant")
res_even = func(42)
assert isinstance(res_even, tvm.tir.IntImm)
assert res_even == 21

res_odd = func(17)
assert res_odd == "argument was odd"

AcceptsVariant#

func = tvm.get_global_func("testing.AcceptsVariant")

assert func("string arg") == "ffi.String"
assert func(17) == "ir.IntImm"

func = tvm.get_global_func("testing.AcceptsVariant")
float_arg = tvm.tir.FloatImm("float32", 0.5)

with pytest.raises(Exception):
    func(float_arg)
目录