[Unity] Extend RemoveAllUnused to support relax::Expr

[Lunderberg]

Prior to this commit, the RemoveAllUnused utility could only apply to relax::Function instances. This commit extends the allowed usage to apply to any relax::Expr that contains variable bindings, such as relax::SeqExpr.

[jinhongyii]

This PR has a bug when handling a use-def chain which ends at an impure call in a binding block.
Suppose the code is like

b = 2*a
c = impure_call(b)
return None

Then b is an unused var (by your definition), but not an impure call, so it will be removed, which breaks the program.

[jinhongyii]

btw, since we have an assumption that binding block is impure, why do we try to remove unused in binding block?

[Lunderberg]

Then b is an unused var (by your definition), but not an impure call, so it will be removed, which breaks the program.

Thank you, and I can reproduce the error with the test case below.

def test_binding_block_keep_pure_func_used_only_for_impure():
    """Remove unused dataflow bindings

    Removal of unused bindings may not remove side effects.  Unused
    bindings whose value is an impure operation (e.g. `R.call_packed`)
    may not be removed, nor may any of their inputs.
    """

    @R.function(private=True)
    def before(x: R.Tensor((32, 32), "int32")):
        y = x * R.const(2)
        z = R.call_packed(
            "function_maybe_with_side_effects", y, sinfo_args=(R.Tensor((32, 32), "int32"))
        )
        return R.tuple()

    expected = before

    after = remove_all_unused(before)
    tvm.ir.assert_structural_equal(expected, after)

For the purpose of removing unused bindings, an argument to an impure function is equivalent to an output, and should be retained. Thankfully, this should be a pretty straightforward change to make.

btw, since we have an assumption that binding block is impure, why do we try to remove unused in binding block?

I'd phrase it slightly differently: A BindingBlock has no assumptions about the purity of its bound values, where a DataflowBlock may assume that all bound values are pure. That is, we don't need to assume that every step of a BindingBlock is impure, but we also cannot assume that any step is pure. Simplifications that require bound values to be pure are still legal within a BindingBlock, so long as the purity is validated before making a change. The advantage of the DataflowBlock is that it lets us skip that validation step, because the contents are required to be pure.

In context of this PR, this check is encoded here. A binding is removable if it is unused, and if it's computation is pure, where purity can be established either from being within a DataflowBlock, or from checking the purity explicitly. The bug is that the set of unused_vars_ is determined by flowing backwards from function outputs, where it should be determined by flowing backwards from function outputs and arguments to impure functions.

[Lunderberg]

@jinhongyii I've opened #16036, which updates the variable usage tracking to check for arguments to impure functions.