relax icon indicating copy to clipboard operation
relax copied to clipboard

[DISCUSS] Naming of variables during transformation

Open ZihengJiang opened this issue 3 years ago • 1 comments

We have lots of implicit agreement on variable naming, for example:

  • lv means local variable, while gv means global variable
  • sh means shape variable
  • storage means storage, alloc and tensor 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.

ZihengJiang avatar Nov 18 '21 01:11 ZihengJiang

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.

altanh avatar Nov 19 '21 19:11 altanh