relax
relax copied to clipboard
[DISCUSS] Naming of variables during transformation
We have lots of implicit agreement on variable naming, for example:
-
lv
means local variable, whilegv
means global variable -
sh
means shape variable -
storage
means storage,alloc
andtensor
both mean allocated tensor - ...
Those naming conventions often are introduced by different transformation pass and the reflected property are often orthogonal, for example, a variable can be a local variable also a shape variable. Using those conventions together would make the final program a mess. For example:
Orignal program:
@tvm.script.ir_module
class Module:
@relax.function
def foo(x: Tensor[_, "float32"]) -> Tensor[_, _]:
# block 0
with relax.dataflow():
sh = relax.call_packed("vm.builtin.shape_of", x)
x0 = relax.match_shape(sh, (n, m))
y = relax.call_dps((n, (m * 2)), "test.vm.tile", x)
relax.output(y)
return y
After lowering:
@tvm.script.ir_module
class Module:
@tir.prim_func
def shape_func(heap_1: tir.handle) -> None:
# function attr dict
tir.func_attr({"global_symbol": "shape_func"})
H_1 = tir.match_buffer(heap_1, [tir.int64(4)], dtype="int64")
# body
tir.store(H_1.data, 2, tir.load("int64", H_1.data, 0) * (tir.load("int64", H_1.data, 1) * tir.int64(2)) * tir.int64(4), True)
@tir.prim_func
def shape_func1(heap_1: tir.handle) -> None:
# function attr dict
tir.func_attr({"global_symbol": "shape_func1"})
H_1 = tir.match_buffer(heap_1, [tir.int64(4)], dtype="int64")
# body
tir.store(H_1.data, 0, tir.load("int64", H_1.data, 0), True)
tir.store(H_1.data, 3, tir.load("int64", H_1.data, 1) * tir.int64(2), True)
@relax.function
def foo(x: Tensor[_, "float32"]) -> Tensor[_, _]:
# block 0
shape_heap: Tensor[(4,), "int64"] = relax.call_packed("vm.builtin.alloc_shape_heap", (4,))
# block 1
sh = relax.call_packed("vm.builtin.shape_of", x)
gv = relax.vm.builtin.store_shape(sh, shape_heap, indices=[0, 1], attrs_type_key="relax.attrs.ShapeHeapAttrs")
_ = shape_func(shape_heap)
sh1 = relax.vm.builtin.load_shape(shape_heap, indices=[2], attrs_type_key="relax.attrs.ShapeHeapAttrs")
storage = relax.vm.builtin.alloc_storage(sh1, device_type=1, dtype="float32", attrs_type_key="relax.attrs.AllocStorageAttrs")
_1 = shape_func1(shape_heap)
sh11 = relax.vm.builtin.load_shape(shape_heap, indices=[0, 3], attrs_type_key="relax.attrs.ShapeHeapAttrs")
tensor = relax.vm.builtin.alloc_tensor(storage, sh11, offset=0, dtype="float32", attrs_type_key="relax.attrs.AllocTensorAttrs")
alloc = tensor
_2 = relax.call_packed("test.vm.tile", x, alloc)
y = alloc
# block 2
return y
We might want to have a unified naming convention instead of creating one by one in each transformation to improve readability.
summary of offline discussion during open source meeting:
- We should definitely standardize the naming convention across passes at least. It would be good to keep some semantic hint in the name if possible, although this might need some further discussion to agree on a convention.
- We should have a single NameTable that lives across the entire compilation pipeline, so that passes can query it without causing duplicate names (currently a new NameTable is created for every BlockBuilder, so each pass has a new NameTable).
- People prefer the unique names to start with index 0, e.g. "x" -> "x0" (the first time gettin a unique name for "x").
- There could still be potential for edge case unique name generation causing something like "x1" -> "x11". We should try to identify the causes and prevent this if possible.