[Unity] Implement relax.Function.bind_symbolic_vars

[Lunderberg]

If a function has dynamic shape parameters, it can be useful to replace them with static parameters (e.g. when producing several models within the same family). This commit introduces a utility function relax.Function.bind_symbolic_vars, which allows symbolic variables to be replaced with static values.

This is a related to the parameter binding done in relax.transform.BindParam, but does not require the bound parameter to be fully static data array.

[tvm-bot]

Thanks for contributing to TVM! Please refer to the contributing guidelines https://tvm.apache.org/docs/contribute/ for useful information and tips. Please request code reviews from Reviewers by @-ing them in a comment.

• cc @quic-sanirudh _{See #10317 for details}

_{Generated by tvm-bot}

[sunggg]

Thanks, @Lunderberg! We might have implemented the similar feature haha
Here is my PR #15246 under review.

[Lunderberg]

Ooh, thank you for pointing it out! That would explain why I didn't see the functionality anywhere in the unity branch itself.

Looking through it, I think the main differences are:

• This PR accepts either tir.Var (less ambiguous, easier to generate programmatically), or strings (more ambiguous, much easier to hand-write), where [Unity][Pass] BindSymVar pass #15246 only accepts strings.

• This PR operates on a relax function, where [Unity][Pass] BindSymVar pass #15246 operates on an IRModule.

• This PR would ignore functions that are defined within a local function definition, where [Unity][Pass] BindSymVar pass #15246 would update their symbolic variables.

• This PR would replace any TIR usage supported by tir::ExprMutator, where [Unity][Pass] BindSymVar pass #15246 checks for specific TIR nodes.

• This PR raises an error if no symbolic variable is known by that name, where [Unity][Pass] BindSymVar pass #15246 would silently ignore it.

Thoughts on making a merged PR with the best of both implementations:

• Providing methods to update either a IRModule transform or a relax Function would provide the most flexibility.
• Implementing a IRModule transform on top of a relax Function update would be require constructing fewer intermediates than implementing a relax Function update on top of an IRModule transform.

Any preference on which of the two PRs should be updated overall? Since both provide similar functionality, I don't have a huge preference.

[sunggg]

Thank you for the great summary!
In fact, I like your PR in the sense that it offers broader supports for TIR and it accepts both strings and tir.Var.
Since I think a pass would be convenient when we are dealing with IRModule with multiple functions inside, how about creating a function pass that wraps your utility function? I agree that this would be easier path to merge our PRs.

[Lunderberg]

@sunggg How does the IRModule pass in the latest commit look?

I pulled in the unit tests and interface from your PR, but made the following changes while integrating the two.

• Swapped order of binding_map and func_name, to allow a default value of func_name = None. If no function name is specified, all functions within the module with be updated.
• Pulled in the error checking from the relax.Function.bind_symbolic_vars for the IRModule pass. If a single function is updated, each value in the binding_map must be used in that function. If multiple functions are updated, each values in the binding_map must occur as a replacement in at least one function.
• Allows replacement either by TIR variable or by name. If multiple functions are updated, a replacement by name may result in updates to a different TIR variable in each function, so long as it uniquely maps to a TIR variable within a given function. (e.g. Replacing "batch_size" in all functions in a module.)

The error checking was why I ended up going with a module pass instead of just wrapping the utility with a function pass. Since the string names would be very convenient to hand-write, they're also the greatest risk of a user typo, which I wanted to guard against. Having a function pass would require either disabling the "all replacements are used", or restricting the updates to a single function.

[sunggg]

Overall, LGTM. Thank you for merging with my PR and making it better!
One comment that is not directly related to this PR.

[sunggg]

I'm surprised to see this case passing the well_formed check. I think this should be filtered before any pass is applied, because this seems incorrect to me.

[Lunderberg]

On the C++ side, both TIR variables and Relax variables have reference identity, and duplicate names can exist within the same function. (Technically, the variables only have a name_hint, not a name, which I believe is to avoid the assumption of unique names. (Though, double parenthetically, some of the python APIs expose name_hint as name.)) TVMScript has a tighter restriction, where variables tracked by the parser may not repeat within the same scope. In this test, I define the variables with the same TIR variable name outside of TVMScript, to produce the ambiguous test case anyways.

As a bit of fun, the function can be inspected, and even though TVMScript modifies the variable names to avoid printing duplicate variables, the underlying objects still have duplicate names.

M = tir.Var("M", "int64")
N1 = tir.Var("N", "int64")
N2 = tir.Var("N", "int64")

dtype = "float32"


@R.function
def func(
    A: R.Tensor([M, N2], dtype), B: R.Tensor([N1, N2], dtype)
) -> R.Tensor([M, N1], dtype):
    A: R.Tensor([M, N2], dtype) = A
    B: R.Tensor([N1, N2], dtype) = B
    out = R.linear(A, B)
    return out


# Prints  B: R.Tensor(("N", "N_1"), dtype="float32")
func.show()
# But the true name can still be extracted
assert func.params[1].struct_info.shape[0].name == "N"
assert func.params[1].struct_info.shape[1].name == "N"
# And the function is well-formed, despite the duplicate variable
# names.
assert tvm.relax.analysis.well_formed(tvm.IRModule.from_expr(func))

[sunggg]

Thanks for the explanation, and yeah I see why it currently passes the well-formed checker. But I'm wondering if this is what we want or what it is supposed to be. I feel like this duplicated name would be very confusing if we are debugging things, especially when we print out IRs. Is there any necessity for this feature?

[Lunderberg]

Hmm. It is really useful if you're merging parts of TIR/Relax that were generated independently. Those are guaranteed to have different tir.Var and relax.Var objects, but can very easily have duplicate names. I think the only place we de-duplicate the names themselves is in an IRModule, and even that gets kind of hairy with both String and GlobalVar acting as unrelated unique keys.

I think that re-use of the same name should be well-formed, but there should be a utility pass (potentially called several times during lowering) that de-duplicates the name within a function. That way, it provides the better debugging experience, but doesn't require renaming variables when constructing TIR. It also avoids promising that variable names will remain stable, or that they should be relied upon, since that isn't a stable assumption at the moment.

[sunggg]

Ah, my apology. I missed the part that it already shows renamed varname when printing functions in your example. Just realized by trying by my own. My concern was printing the same name for different variables, which is turned out to be a false alarm. So, I agree with you, this seems okay. Thank you for clarifying!

[Lunderberg]

@sunggg Also, I've added you as a co-author for this PR.

[slyubomirsky]

I wonder if this is worth changing to avoid needing a special case here. Namely, we can use PrimExpr arguments for strided_slice (the reason it's implemented this way is probably because we didn't have PrimExpr at the time we added this op).

[Lunderberg]

That's true, and I'd like to update strided_slice at some point as well. The comments in the strided_slice implementation state that it should be updated when relax.PrimValue is more widely supported, so that is the long-term goal. However, I think that's a larger change, as it would be a change to the Relax IR, and not just a transformation between Relax functions, and should probably be left for a later PR.

[slyubomirsky]

When does this happen? I think by definition, the symbolic variables are defined at their first appearance in the signature (going left to right)

[Lunderberg]

The symbolic variables are defined in their first appearance in the signature, but only if that appearance is sufficient to provide a definition. If the symbolic variable can be matched to a specific expression in the caller (e.g. ), then it is bound. If the symbolic variable cannot be matched to a specific expression in the caller, then it is a usage-site instead. So, R.Tensor(shape=[16]) binding to R.Tensor(shape=[sym_var]) implies a binding of {sym_var: 16}, but R.Tensor(shape=[16]) binding to R.Tensor(shape=[sym_var1 * sym_var2]) cannot provide any definitions.

So, this edge case occurs in cases like the one shown below, where the first occurrence is a usage, and the definition doesn't occur until a later parameter.

def func(A: R.Tensor(shape=[sym_var1*sym_var2]), B: R.Tensor(shape=[sym_var1, sym_var2])):
    #                       ^^^^^^^^                                ^^^^^^^^
    #                       Usage of sym_var1                       Definition of sym_var1
    ...

[slyubomirsky]

Aha, I hadn't thought of that case. Very good.

[slyubomirsky]

Do we use this pattern elsewhere? I'm not personally very fond of it because it could create problems for adding more modes later, but I'm not against adding it if we do this already.

[Lunderberg]

The pattern of GT/LT enums is used a couple of places, most commonly in the OpPatternKind when checking if an operation is as permissible or more permissible than required for a transformation. That said, I agree that this isn't the most readable, and it probably would be better to have separate visitors to collect defined variables and to collect free variables.

[Lunderberg]

Making two separate visitors ended up with too much duplicate code, so I instead made VisitMode explicitly be a bitflag, with kProvideDefinition and kRequireDefinition. I think this cleans up the implementation, allowing the behavior to be specified with more granularity, and avoids the use of >= and <= for enums.

[slyubomirsky]

These seem like solid changes and the implementation is straightforward. Thank you very much! Please have a look at my comments if you have the chance.