[Relax][Transform] Preserve param names in LiftTransformParams by Lunderberg · Pull Request #16594 · apache/tvm

Lunderberg · 2024-02-16T19:28:43Z

The relax.transform.LiftTransformParams pass splits apart a relax function, extracting the steps that could be performed at compile-time. Prior to this commit, the transformed parameters were named param0, param1, and so on.

This commit updates the LiftTransformParams pass to preserve any human-readable parameter names. The parameter names for the updated function are taken from the original parameter names, if no transformation is performed, or from the internal variable binding, if a transformation is applied. This implementation uses LambdaLift internally, relying on the changes made in
#16306.

The `relax.transform.LiftTransformParams` pass splits apart a relax function, extracting the steps that could be performed at compile-time. Prior to this commit, the transformed parameters were named `param0`, `param1`, and so on. This commit updates the `LiftTransformParams` pass to preserve any human-readable parameter names. The parameter names for the updated function are taken from the original parameter names, if no transformation is performed, or from the internal variable binding, if a transformation is applied. This implementation uses `LambdaLift` internally, relying on the changes made in apache#16306.

slyubomirsky

This is a substantial reorganization of LiftTransformParams and it's nicely modular. Thanks for the changes!

slyubomirsky · 2024-02-21T22:50:27Z

+  // A post-proc utility to reproduce the previous behavior of
+  // LiftTransformParams.


IMO (total nitpick) this is not phrasing we should really have in the codebase itself (as opposed to in the PR description) since someone reading the code will not know what the "previous" version was. We could probably just leave out the bit about the "previous behavior."

No, I completely agree with you there. I should have kept references to "previous behavior" in the commit message instead, where there's a clear before/after. Updated the comment to describe this as the third step in LiftTransformParams, and to include the names of the other two steps.

slyubomirsky · 2024-02-21T22:52:47Z

+  // LiftTransformParams.
+  //
+  // 1. Partition each function into a compile-time and run-time
+  // lambda functions.


Suggested change

// lambda functions.

// lambda function.

Typo

Thank you, and fixed!

slyubomirsky · 2024-02-21T22:59:01Z

-  std::unordered_map<Var, Expr, ObjectPtrHash, ObjectPtrEqual> param_remap_;
-  // The plan of lifting the transform params
-  LiftTransformParamsInfoPlan lift_plan_;
+// Adapted from https://stackoverflow.com/a/2072890


Thank you. At some point, this may be useful as a member function of tvm::String, similar to how python has str.endswith and str.startswith.

slyubomirsky · 2024-02-21T23:00:50Z

+
+  /*! \brief Bindings that can be lifted out into a pre-processing
+   *
+   * - All bindings in `liftable_bindings` are suitable for use in a


Is assume liftable_bindings is an old name that should be updated?

Thank you, and that is exactly correct. I replaced all references, but didn't grep the comments for spots that needed to be updated. Fixed.

slyubomirsky · 2024-02-21T23:30:33Z

    // Cond 5. Do not lift declarations of external functions
-    if (binding->value.as<relax::ExternFuncNode>()) {
-      can_lift = false;
+    if (value.as<relax::ExternFuncNode>()) {


I'm curious when we ever have an ExternFunc as the RHS of a binding.

It happened more often prior to #15900, but rarely happens now.

At some point, I'd like to allow a relax function to specify when tvm::runtime::Registry::Get is invoked. It's legal for a R.ExternFunc to appear within the IRModule::functions table, on the RHS of a binding, or inline for a relax::CallNode::op. Currently, all three of those are resolved when the relax::VirtualMachine is instantiated, but they could be used to represent binding at instantiation-time, at function start-time, and at function execution-time. In that case, the ExternFunc would become more common as the RHS of a binding.

I guess ExternFuncs are permitted to be used in that way per the spec as first-class values, so this is good to handle.

slyubomirsky · 2024-02-21T23:32:58Z

  }

  Expr VisitExpr_(const VarNode* var) override {
-    Var new_var = GetRef<Var>(var);


Was there any bug you encountered due to not doing this visit before? If so, that would be a good regression test to throw in.

Hmm, this was definitely required during initial development, and was intended to handle cases where a replacement in ExprMutator::var_remap_ occurs outside of the planned replacements in the ReplacementPlan.replace_usage_. This can happen if the removal of trivial bindings triggers InferStructInfo in a downstream operation. However, that doesn't reproduce the bug, nor does that bug occur in anything on my local dev branch.

I'm removing this line from the PR, and can make a separate PR (with test case!) if/when the bug resurfaces.

I don't think it's necessary to remove, especially since it seems like a correct change. I'm glad that there wasn't a bug that you encountered because of it.

Double-checking, I think the reason this line is no longer necessary is due to the ExprMutator::VisitExpr_(new_var.get()); at the end of the function. It may have been necessary during some of the development steps, but is no longer necessary. Thinking through the reasoning:

The var_remap_ would apply if VisitVarDef

Canonicalization updates a variable in VisitVarDef if the plan_.replace_binding is populated.

From (1) and (2), anything that would be populated in var_remap_ would also be populated in plan_.replace_bindings

Whenever we populate plan_.replace_binding, we also populate plan_.replace_usage.

From (3) and (4), anything that would be replaced by var_remap_ could instead be replaced by plan_.replace_usage.

So, I think it isn't necessary to have two calls to ExprMutator::VisitExpr_(const Varnode*) in the implementation.

Lunderberg mentioned this pull request Feb 16, 2024

[Transform] Check for zero-param operators in LiftTransformParams #16595

Merged

slyubomirsky approved these changes Feb 21, 2024

View reviewed changes

Update based on review comments

a297f3f

Lunderberg merged commit e715814 into apache:main Feb 23, 2024

Lunderberg deleted the relax_preserve_names_in_lift_transform_params branch February 23, 2024 14:36

ysh329 mentioned this pull request Apr 21, 2024

[Release] v0.16.0 Release Candidate Notes #16911

Closed

		// A post-proc utility to reproduce the previous behavior of
		// LiftTransformParams.

Conversation

Lunderberg commented Feb 16, 2024

Uh oh!

slyubomirsky left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants