[Unity][IR] Purity Tracking #14394

slyubomirsky · 2023-03-24T03:38:43Z

In this PR, I am beginning to implement the tracking of function purity as part of the StructInfo system. This will allow the compiler to enforce that no impure function (one that can possibly have visible side effects) can be called in a DataflowBlock. Tracking this requires noting which operators are pure or impure (I am presently doing this with an operator attribute called FPurity, a simple boolean), though dealing with calls to other Relax functions means that this information must also be tracked via the StructInfo system.

Additionally, it is difficult to infer the purity of a function in the general case (when there are calls to other Relax functions), so this change does require users to annotate impure functions using a new field on functions, is_pure (in TVMScript, this can be done using R.is_pure() or R.is_impure()). Since most Relax functions are likely to be pure and purity is the default assumption, this will hopefully not be a large imposition on users. We can consider inferring purity in the easier cases, since those are likely to be common.

Note that PackedFuncs are conservatively treated as impure. However, in situations where they are needed inside a dataflow block, a call to a PackedFunc that is, in reality, pure can be done via the new operator call_pure_packed or the existing operator call_dps_packed (it is assumed that any PackedFunc used with it will be pure). (Similarly, a new operator invoke_pure_closure is introduced as a counterpart to invoke_closure for dealing with closure objects, though this really should only come up with the LambdaLifting pass.)

As an "escape hatch" to the purity system, one can use the attribute relax.force_pure, which indicates to the compiler to treat the entire function as pure even if it contains an impure call. Additionally, even though PackedFuncs are normally treated as impure, a user can use call_pure_packed or call_dps_packed to call PackedFuncs in dataflow blocks when appropriate. These can be used to deal with the following situations:

A function does side effects but only on a value that will not be exposed anywhere else or on a new value that will be returned. Even though the individual actions are "impure," the overall function fulfills the definition of being pure. relax.force_pure would be useful here.
A PackedFunc is, in reality, pure. call_pure_packed or call_dps_packed are useful in this situation.

Changes include:

Enforcing that impure functions are not used in DataflowBlocks in the well-formed check.
Enforcing that functions that are not labeled impure do not contain impure calls (unless relax.force_pure is set).
Implementing the call_pure operator

Still to be done:

Label purity for all operators
Address certain tricky bugs

(Also address the design concerns below)

The process of implementing purity tracking has revealed a few issues that may require some further design discussion.

Labeling purity or impurity

Using function attributes to label purity/impurity seems very messy. It might be worth making this part of the Function node's AST to avoid having to wrangle attributes in many places in the codebase.

The treatment of `call_pure`

call_pure presents a dilemma because many passes look for calls to certain operators, but with call_pure, these would be "wrapped" like so: Call(Op("relay.call_pure"), [inner_op, arg1, ..., argn], attrs, sinfo_args). In the posted PR, many passes needed to be revised to look for instances of call_pure so that they could be treated the same as calls to ordinary operators. It might cause less disruption to passes to turn call_pure into a specialized AST node that literally wraps a call node. This way visitors or mutators could have an easy default case and the lower-level code generation passes could simply ignore the call_pure node and deal with the wrapped call. As painful as introducing a new AST node would be, it would likely be more maintainable than having every pass have to have a special case for the call_pure operator.

The more restricted call_pure_packed/invoke_pure_closure are less likely to come up in passes and so are less likely to need special handling.

Staging

Related to the above issue, having purity tracking in the well-formedness analysis means that even low-level passes like VMLowerBuiltin, which replaces some Relax operators with PackedFuncs (treated as impure by default, as they are in principle dangerous black boxes), have to insert calls to call_pure to avoid triggering purity errors. Perhaps it might be useful to ignore purity by the time we reach lower stages of compilation, much as the ToNonDataflow pass is used in the VM build process to eliminate DataflowBlocks.

Update: Based on the TVM Unity Community Meeting on Mar. 28, 2023, I've provisionally adopted a variant of this approach thanks to a suggestion by @tqchen. Namely, there is a pass called RemovePurityChecking that simply adds the ForcePure attribute to all pure functions and unwraps all invocations of call_pure_packed, call_pure_dps_packed, and invoke_pure_closure. This allows lower-level passes to remain unchanged and not to have to consider purity at all. It is inserted after ToNonDataflow in the build() function in vm_build.py.

Update 2: Thanks to suggestions by @tqchen, I've made is_pure and force_pure into fields on functions rather than attributes and added syntactic sugar in TVMScript to make these easy to set.

Update 3: Per @tqchen's suggestion force_pure will be kept as an attribute instead. This is because it acts more as a direction to the compiler rather than an inherent property of a function.

tvm-bot · 2023-03-24T03:38:46Z

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}

tvm-bot · 2023-03-24T03:38:47Z

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}

tvm-bot · 2023-03-24T03:38:47Z

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}

tests/python/relax/test_analysis_contains_impure_call.py

src/relax/transform/legalize_ops.cc

src/relax/backend/vm/vm_builtin_lower.cc

src/relax/op/op.cc

tests/python/relax/test_analysis_contains_impure_call.py

python/tvm/script/parser/relax/parser.py

slyubomirsky · 2023-03-27T03:05:18Z

This will probably keep failing tests because lots of them depend on syntactic matches and this has messed up a lot of it. I will start going through them, but some have proven very difficult to debug.

slyubomirsky · 2023-03-29T02:31:41Z

tests/python/relax/test_tvmscript_parser.py

+            # slight hack: normally, we would prefer to use True, but the func attrs, when printed,
+            # will have it as 1, so it would fail roundtripping otherwise
+            R.func_attr({"ForcePure": 1})


@yongwww, would you happen to know how I could avoid this (print the attribute as True instead of 1)? It's not super-pressing, but it would help for roundtripping

the related printing should be https://github.com/apache/tvm/blob/unity/src/script/printer/relax/function.cc#L54-L58. Feel free to leave a todo for me

slyubomirsky · 2023-03-29T02:38:45Z

More general question (maybe one to leave to a community meeting): Where should we communicate that we expect certain passes to be run before others? We already have ToNonDataflow and this PR adds RemovePurityChecking. I also saw that a line in VMShapeLower indicates that it expects LambdaLifting to have already been used.

slyubomirsky · 2023-03-29T03:21:52Z

src/script/printer/relax/call.cc

+Optional<ExprDoc> PrintAssertOp(const relax::Call& n, const ObjectPath& n_p, const IRDocsifier& d) {
+  static const Op& assert_op = Op::Get("relax.assert_op");
+  if (!n->op.same_as(assert_op)) {
+    return NullOpt;
+  }
+  ICHECK(n->args.size() >= 2);
+  // special handling: it is important to indicate that the format string (second argument)
+  // is the _format_ string, or else roundtripping will fail
+  // (the format string will be interpreted as an argument and there will be a new default format
+  // string given)
+  Array<ExprDoc> args;
+  args.push_back(d->AsDoc<ExprDoc>(n->args[0], n_p->Attr("args")->ArrayIndex(0)));
+  ExprDoc second_arg = d->AsDoc<ExprDoc>(n->args[1], n_p->Attr("args")->ArrayIndex(1));
+  for (size_t i = 2; i < n->args.size(); i++) {
+    args.push_back(d->AsDoc<ExprDoc>(n->args[i], n_p->Attr("args")->ArrayIndex(i)));
+  }
+  return Relax(d, "assert_op")->Call(args, {"format"}, {second_arg});
+}
+
+Optional<ExprDoc> PrintRelaxPrint(const relax::Call& n, const ObjectPath& n_p,
+                                  const IRDocsifier& d) {
+  static const Op& print_op = Op::Get("relax.print");
+  if (!n->op.same_as(print_op)) {
+    return NullOpt;
+  }
+  ICHECK(n->args.size() >= 1);
+  // special handling: it is important to indicate that the format string (first argument)
+  // is the _format_ string, or else roundtripping will fail
+  // (the format string will be interpreted as an argument and there will be a new default format
+  // string given)
+  ExprDoc first_arg = d->AsDoc<ExprDoc>(n->args[0], n_p->Attr("args")->ArrayIndex(0));
+  Array<ExprDoc> args;
+  for (size_t i = 1; i < n->args.size(); i++) {
+    args.push_back(d->AsDoc<ExprDoc>(n->args[i], n_p->Attr("args")->ArrayIndex(i)));
+  }
+  return Relax(d, "print")->Call(args, {"format"}, {first_arg});
+}


One would think this isn't related to these changes at all, but in the process of testing roundtripping for call_pure, I discovered roundtripping bugs for print and assert_op (both having to do with their format strings).

tests/python/relax/test_transform_legalize_ops_manipulate.py

tests/python/relax/test_codegen_dnnl.py

slyubomirsky · 2023-03-30T03:25:58Z

src/relax/transform/run_codegen.cc

+          // preserve the purity: if the func was originally pure, wrap call_pure
+          bool purity = GetStructInfoAs<FuncStructInfoNode>(gvar)->purity;
+          auto ret = create_call_dps_packed(new_func, func->ret_struct_info);
+          if (purity) {
+            return WrapCallPure(ret);
+          }
+          return ret;


This is another case where being clear about staging might make some of this reasoning unnecessary. We do codegen before doing the VM build in test_codegen_dnnl.py, so I wasn't sure if I should have eliminated dataflow blocks. The fusion passes seem to rely on having dataflow blocks, so I can't trivially get rid of purity checking and call_pure before fusion (unless we also disable checking purity from the well-formed check--we could consider having a third, internal-only annotation, that disables all purity checks, even inside dataflow blocks).

python/tvm/relax/analysis/analysis.py

python/tvm/script/parser/relax/parser.py

yongwww · 2023-04-01T05:05:38Z

src/relax/op/op.cc

@@ -177,7 +227,9 @@ TVM_REGISTER_OP("relax.call_builtin_with_ctx")
    .set_num_inputs(4)


just noticed this should be 2, it is not related to your change...

Good observation

src/relax/op/op.cc

yongwww · 2023-04-01T05:26:36Z

src/relax/op/tensor/manipulate.cc

@@ -1357,7 +1370,8 @@ TVM_REGISTER_OP("relax.cumsum")
    .set_attrs_type<CumsumAttrs>()
    .set_num_inputs(1)
    .add_argument("data", "Tensor", "The input tensor.")
-    .set_attr<FInferStructInfo>("FInferStructInfo", InferStructInfoCumsum);
+    .set_attr<FInferStructInfo>("FInferStructInfo", InferStructInfoCumsum)
+    .set_attr<Bool>("FPurity", Bool(true));


will we run into issue if FPurity is not specified explicitly?

We should specify it explicitly. I have an assert in the well-formedness checker that it's defined.

slyubomirsky · 2023-04-03T18:04:58Z

Tests failing due to the rebase (have to update some more ops/passes), will address shortly.

…ping/unwrapping behavior

…Params

slyubomirsky · 2023-05-18T01:10:36Z

Updated to reflect the fix from #14864, thanks to @jinhongyii 😄

slyubomirsky added status: WIP branch: unity labels Mar 24, 2023

slyubomirsky marked this pull request as draft March 24, 2023 03:39

slyubomirsky commented Mar 24, 2023

View reviewed changes

tests/python/relax/test_analysis_contains_impure_call.py Show resolved Hide resolved

slyubomirsky commented Mar 24, 2023

View reviewed changes

src/relax/transform/legalize_ops.cc Outdated Show resolved Hide resolved

slyubomirsky commented Mar 24, 2023

View reviewed changes

src/relax/backend/vm/vm_builtin_lower.cc Outdated Show resolved Hide resolved

slyubomirsky force-pushed the purity-tracking branch from c36741f to 1b23bd6 Compare March 24, 2023 03:49

slyubomirsky commented Mar 26, 2023

View reviewed changes

src/relax/op/op.cc Show resolved Hide resolved

slyubomirsky commented Mar 26, 2023

View reviewed changes

tests/python/relax/test_analysis_contains_impure_call.py Outdated Show resolved Hide resolved

slyubomirsky commented Mar 26, 2023

View reviewed changes

python/tvm/script/parser/relax/parser.py Outdated Show resolved Hide resolved

slyubomirsky force-pushed the purity-tracking branch from f246b75 to 1953c74 Compare March 26, 2023 22:55

slyubomirsky changed the title ~~[Unity][IR][WIP] Purity Tracking~~ [Unity][IR] Purity Tracking Mar 26, 2023

slyubomirsky removed the status: WIP label Mar 26, 2023

slyubomirsky marked this pull request as ready for review March 26, 2023 22:56

slyubomirsky commented Mar 29, 2023

View reviewed changes

tests/python/relax/test_transform_legalize_ops_manipulate.py Outdated Show resolved Hide resolved

slyubomirsky commented Mar 29, 2023

View reviewed changes

tests/python/relax/test_codegen_dnnl.py Outdated Show resolved Hide resolved

slyubomirsky commented Mar 30, 2023

View reviewed changes

yongwww reviewed Apr 1, 2023

View reviewed changes

tqchen force-pushed the unity branch 2 times, most recently from a425bc7 to 5c8b7af Compare April 1, 2023 20:00

slyubomirsky force-pushed the purity-tracking branch from 54747d7 to 68968c9 Compare April 3, 2023 16:55

slyubomirsky force-pushed the purity-tracking branch from 68968c9 to 1e34266 Compare April 3, 2023 18:07

slyubomirsky added 22 commits May 17, 2023 20:43

Correct rebase errror

382e467

Address dynamic_strided_slice

419a021

Linting: Remove unused import

4c9d28f

Fix incorrect purity annotations

51a5d42

Fix incorrect comment in well_formed

b91878e

Add explanatory comment in unusual test case

1387180

Use call_pure_packed in pipeline test

cd8675a

Update transform_fuse_ops test

f26eedb

Use ForcePure for the CUDA graph rewrite

697919d

Update comment

e00ba5c

Factor out search for purity annotation

fdb3370

Remove call_dps_pure_packed (call_dps_packed is pure) and remove wrap…

0fa4843

…ping/unwrapping behavior

Remove unused var

446d863

Correct rebase mistake

587ba5f

Use ForcePure for tests of low-level codegen

ba78ec4

lint

85d5f62

Make LazyTransformParams compatible with purity tracking

f1982f5

Make is_pure and force_pure fields in Function instead of attrs

d83bfbb

Lint

3ba972f

Unused imports

3616ec0

Use an attribute (relax.force_pure) to control forcing purity

e786f92

Indicate that RemovePurityChecking is also required for LazyTransform…

c920bf5

…Params

slyubomirsky force-pushed the purity-tracking branch from d1651ce to c920bf5 Compare May 18, 2023 01:10

tqchen approved these changes May 18, 2023

View reviewed changes

slyubomirsky merged commit 7aff6bf into apache:unity May 18, 2023

slyubomirsky mentioned this pull request May 18, 2023

[Unity][Draft][WIP] Relax language specification #14148

Closed

masahi mentioned this pull request May 19, 2023

[Unity][CUTLASS] Fix for purity tracking #14891

Merged

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[Unity][IR] Purity Tracking #14394

[Unity][IR] Purity Tracking #14394

slyubomirsky commented Mar 24, 2023 •

edited

Loading

tvm-bot commented Mar 24, 2023

tvm-bot commented Mar 24, 2023

tvm-bot commented Mar 24, 2023

slyubomirsky commented Mar 27, 2023

slyubomirsky Mar 29, 2023 •

edited

Loading

yongwww Apr 7, 2023 •

edited

Loading

slyubomirsky commented Mar 29, 2023

slyubomirsky Mar 29, 2023

slyubomirsky Mar 30, 2023 •

edited

Loading

yongwww Apr 1, 2023

slyubomirsky Apr 3, 2023

yongwww Apr 1, 2023

slyubomirsky Apr 3, 2023

slyubomirsky commented Apr 3, 2023

slyubomirsky commented May 18, 2023

		@@ -177,7 +227,9 @@ TVM_REGISTER_OP("relax.call_builtin_with_ctx")
		.set_num_inputs(4)

[Unity][IR] Purity Tracking #14394

[Unity][IR] Purity Tracking #14394

Conversation

slyubomirsky commented Mar 24, 2023 • edited Loading

Labeling purity or impurity

The treatment of call_pure

Staging

tvm-bot commented Mar 24, 2023

tvm-bot commented Mar 24, 2023

tvm-bot commented Mar 24, 2023

slyubomirsky commented Mar 27, 2023

slyubomirsky Mar 29, 2023 • edited Loading

Choose a reason for hiding this comment

yongwww Apr 7, 2023 • edited Loading

Choose a reason for hiding this comment

slyubomirsky commented Mar 29, 2023

slyubomirsky Mar 29, 2023

Choose a reason for hiding this comment

slyubomirsky Mar 30, 2023 • edited Loading

Choose a reason for hiding this comment

yongwww Apr 1, 2023

Choose a reason for hiding this comment

slyubomirsky Apr 3, 2023

Choose a reason for hiding this comment

yongwww Apr 1, 2023

Choose a reason for hiding this comment

slyubomirsky Apr 3, 2023

Choose a reason for hiding this comment

slyubomirsky commented Apr 3, 2023

slyubomirsky commented May 18, 2023

slyubomirsky commented Mar 24, 2023 •

edited

Loading

The treatment of `call_pure`

slyubomirsky Mar 29, 2023 •

edited

Loading

yongwww Apr 7, 2023 •

edited

Loading

slyubomirsky Mar 30, 2023 •

edited

Loading