[Discuss] Aliasing

[slyubomirsky]

The current implementation of Relax allows tensors to be aliased.

from __future__ import annotations
import os

import numpy as np
import tvm
import tvm.script
from tvm import relax
from tvm.script import relax as R

@tvm.register_func("test.increment")
def increment(inp):
    inp[:] = np.array([inp.numpy()[0] + np.array(1, dtype="int32")], dtype="int32")

@tvm.script.ir_module
class Alias:
    @R.function
    def main(x: Tensor((1,), "int32")):
        y = x
        p1 = R.call_packed("test.increment", y, type_args=(Void))
        p2 = relax.print(x, y, format="x: {}, y: {}")
        return p1

mod = Alias                                                                              
target = tvm.target.Target("llvm", host="llvm")
ex = relax.vm.build(mod, target)
vm = relax.VirtualMachine(ex, tvm.cpu())

ret = vm["main"](tvm.nd.array([1]))                                                    

When executed, this prints x: [2], y: [2], suggesting that they point to the same tensor. This is not altogether unreasonable. However, if we permit tensor values to alias, then it might make sense to provide a way to copy them in cases when users want to keep the original value as well as a mutated one. As it stands, there is a builtin VM operator that does copying but it's apparently only used for the implementation of if-else expressions (note: we should test if that breaks aliasing, actually) -- perhaps we should expose it.

Any thoughts on aliasing? It's perhaps reasonable for users to expect values to alias (see the August 9 community discussion), so we may want to make sure that it's possible to make copies and check if values are aliased.

[slyubomirsky]

For the record, if-else expressions do not break aliasing. This experiment has the same result as the one above:

@tvm.script.ir_module
class TestIfAlias:
    @R.function
    def branch(cond: Tensor((), "bool"), x: Tensor((1,), "int32")) -> Tensor((1,), "int32"):
        if cond:
            ret = x
    else:
            ret = x
        return ret

    @R.function
    def main(cond: Tensor((), "bool"), x: Tensor((1,), "int32")) -> Object:
        y = branch(cond, x)
    p1 = R.call_packed("test.increment", y, type_args=(Void))
    p2 = relax.print(x, y, format="x: {}, y: {}")
        return p2

mod =  TestIfAlias
target = tvm.target.Target("llvm", host="llvm")
ex = relax.vm.build(mod, target)
vm = relax.VirtualMachine(ex, tvm.cpu())

ret = vm["main"](1, tvm.nd.array([1])) # prints "x: [2], y: [2]"

[sunggg]

@slyubomirsky Thanks for bringing this up! This is an interesting problem. I wonder what kind of sticky situations we would face if we keep the current behavior. For example, what do we expect if we alias tensors within data flow block outside?

 @R.function
 def main(x: Tensor((1,), "int32")):
       with relax.dataflow():
             y = relax.add(x,x) 
             z = relax.add(y,x)
       alias = y
       # play around with alias like in-place update

[slyubomirsky]

If y is a DataflowVar, then it would not be visible outside the dataflow block. I wonder if there is some way to mess with the safety of dataflow blocks this way, yes.

[sunggg]

Also curious if alias would complicate liveness analysis for memory footprint.

[slyubomirsky]

It adds some complexity, yes. You would have to think about the actual memory regions. A region of memory is live if there is any live reference to it, so you would have to track which references in fact point to the same region.

[sunggg]

Right. Might be good to list out the trade-offs for having this alias support. For now, imho, aliasing seems more like a UX feature rather than must-have feature (please correct me if I'm wrong). Don't have a clear opinion yet, but have a doubt if aliasing is worth supporting despite the extra complication that it would bring. It might be okay starting from conservative, restricted approach (no alias) and gradually introduce new features like alias as we observe the real community needs.

[slyubomirsky]

Aliasing exists in the current implementation, the question is whether we want the semantics of the language to permit it or not. I would argue that it is actually the less complex option because preventing aliasing would require a lot of copying of NDArrays and asking packed funcs very, very nicely not to break the abstraction.

[sunggg]

Right, it already exists haha I guess I might subconsciously want to avoid alias if I can due to the PTSD from my limited experience with may-alias relationship - it made my optimization behave more conservative than it was supposed to!

preventing aliasing would require a lot of copying of NDArrays and asking packed funcs very, very nicely not to break the abstraction.

Would you give me an example?

[slyubomirsky]

Here are some choices we might make: If we don't want to permit aliasing, then we would have an assignment like y = x in the above example perform a copy. Then y and x are separate tensors. However, a PackedFuncs make things trickier unless we always want to copy everything all the time (which incurs some overhead). Suppose we have a PackedFunc that simply takes in an NDArray x and returns that same NDArray. The result of the PackedFunc aliases its argument. If we want to prevent aliasing in the language, we would have to check whether the result is an alias for something else and make a copy. However, an even more evil example might be a PackedFunc that returns a pointer to a global NDArray it keeps. Every time we call this PackedFunc, the result aliases the results of past calls. We would have no way to automatically detect this, so we may just ask users very nicely to not do things like that in PackedFuncs (on pain of the program breaking). Again, always copying tensors returned by PackedFunc calls would solve this problem, but it could introduce lots of unnecessary overhead. We could potentially be smart about this if we rely on user annotations to tell us when results will not alias (putting a burden on users).

my_array = tvm.nd.array([...])

@tvm.register_global("my_evil_func")
def my_evil_func():
    return my_array

Those are not the only choices we could make, just one example. Another choice would be having some kind of COW layer on top of NDArrays in the language, but that would introduce overhead at runtime. I am sure there might be more options, too. What makes things very complicated is exposing PackedFuncs in the language, since they can do anything they want and users are certain to take advantage of that flexibility. Given that, I would actually advocate for permitting aliasing and being up-front about it, while also providing ways to make copies when needed and check if values are aliases.

[sunggg]

@slyubomirsky Really appreciate the detailed explanation :)
Alias is much more common than my previous naive understanding - wasn't relating alias to function calls. I'm convinced! In summary, I agree with the following direction:

1. Support alias and provide a way to make copies
2. Address potential corner cases
3. (Optional) Take user annotation for better analysis e.g., TIR seems to have attr for noalias: https://github.com/tlc-pack/relax/blob/relax/include/tvm/tir/function.h#L324

[slyubomirsky]

Oh boy, here's a tricky case that we might want to think about: Suppose we have a function that captures a tensor from an outer scope. If tensor values can be aliased, then that might provide a way to mess with dataflow blocks

Example:

@R.function
def outer(x: Tensor((1,), "int32"):
    @R.function
    def inner():
        with R.dataflow():
            dx = x
            R.output(dx)
        return dx
    w = relax.print(inner())
    y = R.call_packed("increment", x, type_args=(Void))
    z = relax.print(inner())
    return z

With the argument [1], this prints:

[1]
[2]

Even though inner() is a pure function, its output can be affected by mutating the x that is captured. :zany_face:

This is, I guess, not so crazy (Racket and other LISP-like languages also do closure capturing by reference). It just makes it a little tricky to think about closure capturing if there are external references remaining. Even a function that appears to be "pure" (not mutating any values) may not be idempotent.

This is another area where we can make a choice: Capture by value or capture by reference (or have some way to choose, like C++ lambdas offer).

[slyubomirsky]

Note: In-place mutation can also be done to scalars! This may actually be a reason to distinguish between pure ints/bools/etc and 0-rank tensors, but I would really rather we not distinguish them.

@tvm.register_func("test.inc_scalar")
def inc_scalar(s):
    t = tvm.nd.array(int(s.numpy()) + 1)
    t.copyto(s)

@tvm.script.ir_module
class IncrementScalar:
    @R.function
    def main(x: Tensor((), "int32")):
        y = R.call_packed("test.inc_scalar", x, type_args=(Void))
        return x # x will be incremented in-place

[tqchen]

aliasing indeed create complication(need for alias analysis or do nothing). On the otherhand, they are somewhat needed mainly because of the abstraction and interfacing with python's data structure (which themselves alias).

From practical pov, most of the optimizations focused on the dataflow block and we do not need to worry about aliasing there. The capture example is indeed interesting, similar property holds there, since as we entering the dataflow block, the boundary (the vars that goes into it and implicit variables that are captured by closure) remains immutable during the dataflow block.

[slyubomirsky]

Agreed that this is less of a problem in dataflow blocks. We should adequately warn users of the potential for aliasing and provide facilities to be able to deal with aliased values when they are writing non-dataflow code.