[Relay][Training][Pass] Factor out first-order AD to a module pass #7677

altanh · 2021-03-17T03:38:47Z

Since there is negligible code sharing between the first-order and higher-order AD in Relay, I've factored out the first-order AD pass to a separate file. Additionally, I've made it a proper IRModule pass, and fixed a few spots where adding tuples wasn't being lifted correctly. I tried my best to add some diagnostics.

To complement AD, I've also added a pass called ConcretizeLike that transforms the *_like operators generated during AD to the explicit-shape equivalent (e.g. zeros_like(x, y) -> zeros(x, y.shape)), which removes some forward dependencies and could enable more opportunity for operator fusion. Note that this is the same thing as LikeZapp from @t-vi's blogpost with perhaps a bit more error checking.

cc @t-vi @MarisaKirisame @tqchen @jroesch

MarisaKirisame · 2021-03-17T03:48:57Z

src/relay/transforms/concretize_like.cc

+      auto attrs = make_object<InitOpAttrs>();
+      auto cshape =
+          MakeConstantTensor(DataType::Int(32), {static_cast<int64_t>(shape.size())}, shape);
+      attrs->shape = shape;


Suggested change

auto attrs = make_object<InitOpAttrs>();

auto cshape =

MakeConstantTensor(DataType::Int(32), {static_cast<int64_t>(shape.size())}, shape);

attrs->shape = shape;

auto attrs = make_object<InitOpAttrs>();

attrs->shape = shape;

auto cshape =

MakeConstantTensor(DataType::Int(32), {static_cast<int64_t>(shape.size())}, shape);

MarisaKirisame · 2021-03-17T03:50:35Z

src/relay/transforms/concretize_like.cc

+      if (const auto* imm = dim.as<IntImmNode>()) {
+        cshape.push_back(Integer(GetRef<IntImm>(imm)));
+        continue;
+      }
+      return post;


Suggested change

if (const auto* imm = dim.as<IntImmNode>()) {

cshape.push_back(Integer(GetRef<IntImm>(imm)));

continue;

}

return post;

if (const auto* imm = dim.as<IntImmNode>()) {

cshape.push_back(Integer(GetRef<IntImm>(imm)));

} else {

return post;

}

MarisaKirisame · 2021-03-17T03:51:24Z

src/relay/transforms/concretize_like.cc

+    if (call_node->args.size() == 2) {
+      return concrete_map_.at(op_ref)(node_map[data_pat_][0], cshape, like_ty->dtype);
+    }
+    return concrete_map_.at(op_ref)(Expr(), cshape, like_ty->dtype);


why empty Expr()?

Yeah maybe this is too much of a hack, I'm just using it as a placeholder since the unary matches won't have a corresponding data Expr node. I'll rework this tmrw

It might be better to refer to SimplifyExpr pass to separate unary and binary ops. Then maybe we could have a base struct to put the sharable logic.

tqchen · 2021-03-17T13:41:08Z

cc @comaniac

comaniac

Also cc @yzhliu

comaniac · 2021-03-17T16:56:03Z

src/relay/transforms/concretize_like.cc

+    for (const auto& pr : concrete_map_) {
+      if (!op_pat_.defined()) {
+        op_pat_ = IsExpr(pr.first);
+      } else {
+        op_pat_ = op_pat_ || IsExpr(pr.first);
+      }
+    }


Could you ellaborate what this loop does? It seems to me that it unions all the patterns that match ops in concrete_map_, but I didn't find op_pat_ being used else where.

The construction of concrete_map_ could be static, so we should be able to move it out of the constructor.

thanks for the catch, this was left over from before

true, I'm not sure what the code style is for defining static variables like this but I'll try something

just to follow up on the idea of making these static, I'm mainly concerned about how the static lifetime interacts with the Operator registry, do you know how that works?

I don't think it would be an issue. If you search static const Op& op in the code base, you'll find lots of use cases.

comaniac · 2021-03-17T16:56:33Z

src/relay/transforms/concretize_like.cc

+    data_pat_ = IsWildcard();
+    like_pat_ = IsWildcard();
+    unary_like_pat_ = (IsOp("zeros_like") || IsOp("ones_like"))({like_pat_});
+    binary_like_pat_ = (IsOp("reshape_like") || IsOp("collapse_sum_like") ||
+                        IsOp("broadcast_to_like"))({data_pat_, like_pat_});


These patterns can also be defined statically.

Now we have two places to specify the supported *like ops. I feel we should define a list of "unary like" and "binary like" ops once and used them both in concrete_map_ and here. As a result, we could have the logic similar to L61-67 to construct the pattern.

good points, thanks

comaniac · 2021-03-17T16:59:32Z

src/relay/transforms/concretize_like.cc

+    // skip trying to support this for now (ironic, as I was the one who added the feature)
+    if (const auto* attrs = call_node->attrs.as<ReshapeLikeAttrs>()) {
+      if (attrs->lhs_begin != 0 || attrs->rhs_begin != 0 || attrs->lhs_end.defined() ||
+          attrs->rhs_end.defined()) {
+        return post;
+      }
+    }


This is too ad-hoc. It means we may not concretize *like ops in certain situations. Instead of hacking the unified callabck function, we should maintain this logic in the op specific function. I can think of two ways to achieve this:

Put the logic in the beginning of concrete_map_ functions and return the same op if not applicable.

Construct another checker map that include checker functions of each op, and invoke the corresponding checker function here.

yeah I agree, thanks

comaniac · 2021-03-17T17:04:36Z

src/relay/transforms/concretize_like.cc

+    CHECK(like->checked_type_.defined())
+        << "ConcretizeLike requires checked types to be populated, please run type inference";


We already have L88-90 so this check seems useless.

comaniac · 2021-03-17T17:09:05Z

src/relay/transforms/concretize_like.cc

+    if (call_node->args.size() == 2) {
+      return concrete_map_.at(op_ref)(node_map[data_pat_][0], cshape, like_ty->dtype);
+    }
+    return concrete_map_.at(op_ref)(Expr(), cshape, like_ty->dtype);


It might be better to refer to SimplifyExpr pass to separate unary and binary ops. Then maybe we could have a base struct to put the sharable logic.

comaniac · 2021-03-17T17:13:30Z

tests/python/relay/test_pass_concretize_like.py

+    data = relay.var("data", shape=(2, 3, 4), dtype="float32")
+    shape_like = relay.var("shape_like", shape=(6, 2, 2), dtype="float32")
+    f = relay.Function([data, shape_like], relay.reshape_like(data, shape_like))
+    f_expected = relay.Function([data, shape_like], relay.reshape(data, (6, 2, 2)))


This exposes a question of the current implementation: after the ConcretizeLike pass we expect unused arguments. If we don't need shape_like tensor anymore, we should remove it from the function.

this is a good point that I was wondering about while I wrote the tests. However this might be complicated, as removing an parameter means we'll need to remove the argument at each callsite (and then maybe re-apply the optimization until fixed point). I imagine we won't want to remove the parameter from a global function as then the user will have to ensure they don't pass the removed argument.

It should be fine to have another internal pass being called after concretizing like ops to check and mutate the function signatures.

comaniac · 2021-03-17T17:14:12Z

tests/python/relay/test_pass_concretize_like.py

+
+    mod = tvm.IRModule.from_expr(f)
+    mod_concrete = relay.transform.ConcretizeLike()(mod)
+    assert tvm.ir.structural_equal(mod_concrete["main"], f_expected)


Add if __name__ == "__main__":

altanh · 2021-03-17T18:10:15Z

@comaniac thanks for the helpful feedback, it seems there's some design decisions that would be worth combining with existing approaches. I'll spend some time this week to incorporate the feedback, but in the meantime, I think I'll remove the ConcretizeLike pass from this PR and send a follow up (since the AD refactor can stand alone). How does that sound?

comaniac · 2021-03-17T18:42:25Z

@comaniac thanks for the helpful feedback, it seems there's some design decisions that would be worth combining with existing approaches. I'll spend some time this week to incorporate the feedback, but in the meantime, I think I'll remove the ConcretizeLike pass from this PR and send a follow up (since the AD refactor can stand alone). How does that sound?

Sounds good to me and I'm good with the first-order AD changes in general.

tqchen · 2021-03-18T13:09:27Z

@MarisaKirisame please take another look and https://tvm.apache.org/docs/contribute/code_review.html#approve-and-request-changes-explicitly

tqchen · 2021-03-18T14:04:56Z

Thanks @altanh @comaniac @MarisaKirisame

…pache#7677)

altanh added 5 commits March 15, 2021 16:08

first order AD as a separate pass

ebf056b

ConcretizeLike pass

4b74bd6

some refactoring, add unit tests for ConcretizeLike

bb4ec4c

revert changes unrelated to this PR

ee72271

missed one

4a01859

MarisaKirisame reviewed Mar 17, 2021

View reviewed changes

tqchen self-assigned this Mar 17, 2021

tqchen added the status: need review label Mar 17, 2021

comaniac requested changes Mar 17, 2021

View reviewed changes

delete ConcretizeLike

775086b

comaniac approved these changes Mar 17, 2021

View reviewed changes

tqchen changed the title ~~[Relay][Training][Pass] Factor out first-order AD to a module pass, and add ConcretizeLike pass~~ [Relay][Training][Pass] Factor out first-order AD to a module pass Mar 17, 2021

MarisaKirisame approved these changes Mar 18, 2021

View reviewed changes

tqchen merged commit 45442ed into apache:main Mar 18, 2021

tqchen added status: accepted and removed status: need review labels Mar 18, 2021

trevor-m pushed a commit to trevor-m/tvm that referenced this pull request May 6, 2021

[Relay][Training][Pass] Factor out first-order AD to a module pass (a…

2fd5c04

…pache#7677)

trevor-m pushed a commit to neo-ai/tvm that referenced this pull request May 11, 2021

[Relay][Training][Pass] Factor out first-order AD to a module pass (a…

85a98ce

…pache#7677)

junrushao mentioned this pull request Nov 1, 2021

Apache TVM v0.8 Release Note Candidate #9416

Closed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[Relay][Training][Pass] Factor out first-order AD to a module pass #7677

[Relay][Training][Pass] Factor out first-order AD to a module pass #7677

altanh commented Mar 17, 2021

MarisaKirisame Mar 17, 2021

MarisaKirisame Mar 17, 2021

MarisaKirisame Mar 17, 2021

altanh Mar 17, 2021

comaniac Mar 17, 2021

tqchen commented Mar 17, 2021

comaniac left a comment

comaniac Mar 17, 2021

altanh Mar 17, 2021

altanh Mar 17, 2021

comaniac Mar 17, 2021

comaniac Mar 17, 2021

altanh Mar 17, 2021

comaniac Mar 17, 2021

altanh Mar 17, 2021

comaniac Mar 17, 2021

comaniac Mar 17, 2021

comaniac Mar 17, 2021

altanh Mar 17, 2021

comaniac Mar 17, 2021

comaniac Mar 17, 2021

altanh commented Mar 17, 2021

comaniac commented Mar 17, 2021

tqchen commented Mar 18, 2021

tqchen commented Mar 18, 2021

		CHECK(like->checked_type_.defined())
		<< "ConcretizeLike requires checked types to be populated, please run type inference";

[Relay][Training][Pass] Factor out first-order AD to a module pass #7677

[Relay][Training][Pass] Factor out first-order AD to a module pass #7677

Conversation

altanh commented Mar 17, 2021

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

tqchen commented Mar 17, 2021

comaniac left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

altanh commented Mar 17, 2021

comaniac commented Mar 17, 2021

tqchen commented Mar 18, 2021

tqchen commented Mar 18, 2021