[AutoDiff] Automatically determine AdStack's size #2438

xumingkuan · 2021-06-18T13:28:14Z

Related issue = #656

This PR uses the control-flow graph to compute the necessary size for each Autodiff stack.

~~In case there may be a loop in which there is a stack push, a limit constexpr int kMaxAdStackSize = 32; is set to prevent infinite loops.~~

We use the Bellman-Ford algorithm to compute the sizes. When there is a positive loop (#pushes > #pops in a loop) for an Autodiff stack, we cannot determine the size, and a default value CompileConfig::default_ad_stack_size = 32 is used.

xumingkuan · 2021-06-18T13:29:57Z

/format

taichi/backends/metal/codegen_metal.cpp

yuanming-hu

I haven't carefully checked ControlFlowGraph::determine_ad_stack_size, but I believe this is great progress and a very useful feature in AutoDiff :-)

At a high level, I think we should add a few test cases using the CHI IR builder. A subtle error in this inference pass may result in undefined behavior (since stack overflows) instead of a laud error message. It would also make sense to add a runtime check of stack overflowing in debug mode.

taichi/transforms/determine_ad_stack_size.cpp

taichi/backends/cc/codegen_cc.cpp

taichi/ir/control_flow_graph.cpp

taichi/ir/control_flow_graph.h

taichi/ir/statements.h

taichi/ir/transforms.h

xumingkuan · 2021-06-19T16:14:09Z

TODOs:

Add C++ tests
Add documentation in transforms.h
Final verify pass?

Co-authored-by: Ye Kuang <k-ye@users.noreply.github.com>

Co-authored-by: Yuanming Hu <yuanming-hu@users.noreply.github.com>

xumingkuan · 2021-06-22T07:11:07Z

verify -> verify_ir_structure
verify_before_codegen

taichi/ir/control_flow_graph.cpp

xumingkuan · 2021-06-22T08:44:50Z

verify -> verify_ir_structure
verify_before_codegen

I think it's better to put this refactoring in another PR. WDYT?

yuanming-hu · 2021-06-22T09:18:08Z

verify -> verify_ir_structure
verify_before_codegen

I think it's better to put this refactoring in another PR. WDYT?

Sounds good!

taichi/backends/cc/codegen_cc.cpp

k-ye · 2021-06-24T06:00:15Z

taichi/ir/control_flow_graph.cpp

+    for (auto &stack : oversized_stacks) {
+      oversized_stacks_name.push_back(stack->name());
+    }
+    TI_WARN(


How serious it is if the necessary AD stack size overflows max_ad_stack_size? If this would directly result in wrong results, we should better report an error and stop immediately?

This is not that serious -- IMHO in most cases, it's the control-flow graph that unable to determine the max capacity. For example:

s = stack() for i in range(10): s.push(i)

The control-flow graph does not have the range 10, and thus cannot determine the capacity.

From another perspective, the current codebase uses a fixed capacity of 16, and it works fine.

On the other hand, I don't see any warnings in the current tests, so maybe even the above case (a loop with #pushes > #pops) doesn't exist, and the control-flow graph is able to determine all maximum capacities. Whether the above case exists depends on the usage of AD-stack in auto_diff.cpp.

Hmm, the code implementation looks like it is able to figure out the necessary stack size, which overflows max_ad_stack_size. But the interpretation is "unable to determine the max capacity". Is it possible to distinguish these two cases, or maybe I'm misunderstanding something ? 🤣

Oh I was using the terms "size" and "capacity" interchangeably... It should be the necessary stack size, not the max capacity.

If auto_diff.cpp is bug-free, then #pushes should be always equal to #pops?

I still don't see why the fact that an undetermined stack would not lead to a bad result. IIUC, this is serious, but there are cases where this stack size just cannot be determined statically? (Like the range(10) example you gave).

I don't see any warnings in the current tests

Were you referring to the python tests, or the new CPP tests? If the former, i think that's because the output are not shown in pytest by default..

Thanks for the explanation. Let's add this comment to the code?

Now I think maybe it's better to use the Bellman-Ford algorithm -- when we are "able to figure out all stack sizes statically", we are guaranteed to figure them out even if a stack needs a large size (and the algorithm's running time is approximately the same); when we are unable to figure out at least one stack size, the algorithm will run slower but we will be sure that we are unable to figure out the stack size statically (instead of a warning about possible overflow) (So, in this case, I think the message should not be a TI_WARN -- maybe a TI_INFO or TI_DEBUG). WDYT?

Now I think maybe it's better to use the Bellman-Ford algorithm

Sounds great!

So, in this case, I think the message should not be a TI_WARN -- maybe a TI_INFO or TI_DEBUG. WDYT?

IIUC, we cannot determine the stack size as soon as the kernel has a loop in it? If so, yeah, maybe it's better to make this TI_DEBUG to reduce the noise..

IIUC, we cannot determine the stack size as soon as the kernel has a loop in it? If so, yeah, maybe it's better to make this TI_DEBUG to reduce the noise..

Right...

@ljcc0930 will help review Bellman-Ford XD

taichi/ir/control_flow_graph.cpp

tests/cpp/transforms/determine_ad_stack_size_test.cpp

k-ye · 2021-06-24T06:26:45Z

taichi/ir/control_flow_graph.cpp

+    for (auto &stack : oversized_stacks) {
+      oversized_stacks_name.push_back(stack->name());
+    }
+    TI_WARN(


Hmm, the code implementation looks like it is able to figure out the necessary stack size, which overflows max_ad_stack_size. But the interpretation is "unable to determine the max capacity". Is it possible to distinguish these two cases, or maybe I'm misunderstanding something ? 🤣

Co-authored-by: Ye Kuang <k-ye@users.noreply.github.com>

tests/cpp/transforms/determine_ad_stack_size_test.cpp

k-ye · 2021-06-26T11:31:38Z

taichi/ir/control_flow_graph.cpp

+    for (auto &stack : oversized_stacks) {
+      oversized_stacks_name.push_back(stack->name());
+    }
+    TI_WARN(


https://stackoverflow.com/questions/14405063/how-can-i-see-normal-print-output-created-during-pytest-run ? (You may have to tweak ti to make this work 🤣 )

k-ye · 2021-06-26T11:34:02Z

taichi/ir/control_flow_graph.cpp

+    for (auto &stack : oversized_stacks) {
+      oversized_stacks_name.push_back(stack->name());
+    }
+    TI_WARN(


OK I see, but I thought oversized_stacks means a different thing? I.e., "I am able to figure out the stack size statically, and it definitely overflows the configured max". So capping it to max_ad_stack_size will lead to wrong result?

k-ye

Great thanks!

ljcc0930

LGTM!

xumingkuan · 2021-06-29T05:31:33Z

Hi @ljcc0930 , could you please re-review 6652b8d ? I wasn't really using Bellman-Ford before this commit. Thanks in advance!

taichi/ir/control_flow_graph.cpp

ljcc0930

LGTM!

[AutoDiff] Automatically determine AdStack's size

2b92337

xumingkuan mentioned this pull request Jun 18, 2021

Advanced optimization #656

Closed

18 tasks

taichi-gardener and others added 4 commits June 18, 2021 13:31

Auto Format

65e8ec4

revert auto format

1cda1cd

revert auto format

78e2486

update comment

6020799

xumingkuan commented Jun 18, 2021

View reviewed changes

taichi/backends/metal/codegen_metal.cpp Show resolved Hide resolved

fix format

8625399

xumingkuan requested review from k-ye and yuanming-hu June 18, 2021 13:39

oops

244b249

yuanming-hu reviewed Jun 19, 2021

View reviewed changes

taichi/transforms/determine_ad_stack_size.cpp Outdated Show resolved Hide resolved

taichi/backends/cc/codegen_cc.cpp Show resolved Hide resolved

taichi/ir/control_flow_graph.cpp Outdated Show resolved Hide resolved

k-ye reviewed Jun 19, 2021

View reviewed changes

taichi/ir/control_flow_graph.h Outdated Show resolved Hide resolved

k-ye reviewed Jun 19, 2021

View reviewed changes

taichi/ir/statements.h Outdated Show resolved Hide resolved

taichi/ir/transforms.h Show resolved Hide resolved

xumingkuan marked this pull request as draft June 19, 2021 16:12

xumingkuan and others added 5 commits June 22, 2021 13:32

Update taichi/ir/statements.h

1ae751f

Co-authored-by: Ye Kuang <k-ye@users.noreply.github.com>

Update taichi/ir/control_flow_graph.cpp

d01283c

Co-authored-by: Yuanming Hu <yuanming-hu@users.noreply.github.com>

Apply review

b2f31b2

Add a documentation

411c332

Add a basic C++ test

9b8cfc5

yuanming-hu reviewed Jun 22, 2021

View reviewed changes

taichi/ir/control_flow_graph.cpp Outdated Show resolved Hide resolved

taichi/ir/control_flow_graph.cpp Outdated Show resolved Hide resolved

taichi/ir/control_flow_graph.cpp Outdated Show resolved Hide resolved

taichi/ir/control_flow_graph.cpp Outdated Show resolved Hide resolved

xumingkuan added 2 commits June 22, 2021 16:20

Add 3 more tests

e04dfe7

Add comments

005a560

xumingkuan marked this pull request as ready for review June 22, 2021 08:46

xumingkuan requested a review from k-ye June 22, 2021 08:46

k-ye reviewed Jun 24, 2021

View reviewed changes

xumingkuan and others added 2 commits June 24, 2021 17:21

Update taichi/ir/control_flow_graph.cpp

6d86c71

Co-authored-by: Ye Kuang <k-ye@users.noreply.github.com>

Apply review, use parameterized tests, fix typo and code format

243764f

xumingkuan requested a review from k-ye June 26, 2021 09:14

k-ye reviewed Jun 26, 2021

View reviewed changes

k-ye approved these changes Jun 28, 2021

View reviewed changes

k-ye requested a review from ljcc0930 June 28, 2021 06:54

ljcc0930 approved these changes Jun 28, 2021

View reviewed changes

xumingkuan added 2 commits June 29, 2021 13:03

Apply review

838bf5d

Use the Bellman-Ford algorithm

6652b8d

xumingkuan requested a review from ljcc0930 June 29, 2021 05:30

xumingkuan commented Jun 29, 2021

View reviewed changes

taichi/ir/control_flow_graph.cpp Outdated Show resolved Hide resolved

xumingkuan added 2 commits June 29, 2021 15:43

Update taichi/ir/control_flow_graph.cpp

39c8d0d

Run Bellman-Ford on each stack separately, fix a bug, add a test

838e969

ljcc0930 approved these changes Jun 29, 2021

View reviewed changes

k-ye approved these changes Jun 29, 2021

View reviewed changes

k-ye merged commit 7fe1cf2 into taichi-dev:master Jun 29, 2021

squarefk mentioned this pull request Jul 1, 2021

[release] v0.7.25 #2482

Merged

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[AutoDiff] Automatically determine AdStack's size #2438

[AutoDiff] Automatically determine AdStack's size #2438

xumingkuan commented Jun 18, 2021 •

edited

Loading

xumingkuan commented Jun 18, 2021

yuanming-hu left a comment

xumingkuan commented Jun 19, 2021 •

edited

Loading

xumingkuan commented Jun 22, 2021

xumingkuan commented Jun 22, 2021

yuanming-hu commented Jun 22, 2021

k-ye Jun 24, 2021

xumingkuan Jun 24, 2021

k-ye Jun 24, 2021

xumingkuan Jun 24, 2021

k-ye Jun 26, 2021

k-ye Jun 28, 2021

xumingkuan Jun 28, 2021 •

edited

Loading

k-ye Jun 28, 2021

xumingkuan Jun 28, 2021

k-ye Jun 28, 2021

k-ye Jun 24, 2021

k-ye Jun 26, 2021

k-ye Jun 26, 2021

k-ye left a comment

ljcc0930 left a comment

xumingkuan commented Jun 29, 2021

ljcc0930 left a comment

[AutoDiff] Automatically determine AdStack's size #2438

[AutoDiff] Automatically determine AdStack's size #2438

Conversation

xumingkuan commented Jun 18, 2021 • edited Loading

xumingkuan commented Jun 18, 2021

yuanming-hu left a comment

Choose a reason for hiding this comment

xumingkuan commented Jun 19, 2021 • edited Loading

xumingkuan commented Jun 22, 2021

xumingkuan commented Jun 22, 2021

yuanming-hu commented Jun 22, 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

Choose a reason for hiding this comment

xumingkuan Jun 28, 2021 • edited Loading

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

k-ye left a comment

Choose a reason for hiding this comment

ljcc0930 left a comment

Choose a reason for hiding this comment

xumingkuan commented Jun 29, 2021

ljcc0930 left a comment

Choose a reason for hiding this comment

xumingkuan commented Jun 18, 2021 •

edited

Loading

xumingkuan commented Jun 19, 2021 •

edited

Loading

xumingkuan Jun 28, 2021 •

edited

Loading