AffineLoopFusion: Prevent fusion of multi-out-edge producer loops #259

dcaballe · 2019-11-22T19:25:28Z

In #162 I introduced a bug that
incorrectly allowed fusion of producer loops with multiple outgoing
edges. This commit fixes that problem. It also introduces a new flag to
disable sibling loop fusion so that we can test producer-consumer fusion
in isolation.

tensorflow#162 introduced a bug that incorrectly allowed fusion of producer loops with multiple outgoing edges. This commit fixes that problem. It also introduces a new flag to disable sibling loop fusion so that we can test producer-consumer fusion in isolation.

bondhugula · 2019-11-25T03:52:19Z

lib/Transforms/LoopFusion.cpp

+         "Expected only outgoing edges from memref in srcId");
+
+  // Return false if 'srcNode' has more than one output edge on 'memref'.
+  if (memrefNumOutEdges != 1)


bondhugula · 2019-11-25T03:53:37Z

lib/Transforms/LoopFusion.cpp

-    fuseProducerConsumerNodes(/*maxSrcUserCount=*/1);
-    fuseSiblingNodes();
+    if (clSiblingFusion) {
+      fuseProducerConsumerNodes(/*maxSrcUserCount=*/1);


Why is the call to fuseProducerConsumerNodes here under clSiblingFusion?

I assumed that if sibling fusion was disabled, only one fuseProducerConsumerNodes pass was needed and I left the one with the highest maxSrcUserCount. I now see comment in line 1472 and I'm not sure if a single pass will address all the producer-consumer cases. Could a first pass expose new producer-consumer cases that are not handled by the own pass?

Yes, please put call the "fuseSiblingNodes()" under the flag "clSiblingFusion", otherwise I think users will get confused. If you need to disable producer-consumer fusion, then it should be under a separate flag.

However, we may not need the flag at all, if you can construct a test case which is illegal for sibling fusion.

// producer loop
for ...
store %v, %a[%i]

// Consumer loop 1
for ...
load %a[%i]
store %v, %b[%i]

// Consumer loop 2
for ...
load %a[%i]
load %b[%i]

The dependence on '%b' between the two consumer nodes should prevent sibling fusion

bondhugula · 2019-11-25T04:02:09Z

test/Transforms/loop-fusion.mlir

+// multiple outgoing edges. Sibling loop fusion is disabled to properly test
+// producer-consumer fusion in isolation.
+// CHECK-LABEL: func @should_not_fuse_multi_outgoing_edge_store_producer
+func @should_not_fuse_multi_outgoing_edge_store_producer() {


Producer consumer fusion by itself is expected to fuse everything here. You shouldn't need sibling fusion to be turned on for fusion here. As such, the desired output even with -sibling-fusion=false would be the fused nest. While the fused nest could be seen as a sibling fusion followed by producer consumer fusion (which is probably what you had in mind?), it can also be viewed just as a sequence of two producer consumer fusions.

Disabling sibling fusion would be expected to stop fusion in other cases like these.

for %i = ... %v = affine.load %A[%i] affine.store %v, %B[%i] for %j = ... %v = affine.load %A[%i] affine.store %v, %C[%i]

There is no producer consumer relationship or reuse, but "sibling fusion" (being for input/RAR reuse) will fuse these.

Yes, that makes sense to me. A mature producer-consumer fusion algorithm should fuse loops below. However, what I'm trying to test here is that the current producer-consumer algorithm is not fusing a producer with more than one outgoing edge because that is not supported currently (note the three NOSIBLING checks). If I didn't disable sibling fusion, producer-consumer algorithm wouldn't fuse the loops but sibling fusion would and I couldn't be sure who did it.

Can you try a unit test as follows. This test would not allow sibling fusion of the two consumer loop nests because of the dependence on '%b'.

// producer loop
for ...
store %v, %a[%i]

// Consumer loop 1
for ...
load %a[%i]
store %v, %b[%i]

// Consumer loop 2
for ...
load %a[%i]
load %b[%i]

Thanks! I’ve been playing with the suggested test case:

func @should_not_fuse_multi_outgoing_edge_store_producer2(%a : memref<1xf32>, %b : memref<1xf32>) { %cst = constant 0.000000e+00 : f32 affine.for %arg0 = 0 to 1 { affine.store %cst, %a[%arg0] : memref<1xf32> } affine.for %arg0 = 0 to 1 { %0 = affine.load %a[%arg0] : memref<1xf32> affine.store %cst, %b[%arg0] : memref<1xf32> } affine.for %arg0 = 0 to 1 { %0 = affine.load %a[%arg0] : memref<1xf32> %1 = affine.load %b[%arg0] : memref<1xf32> } return }

Producer-consumer fusion fuses the three loops. It fuses 2nd and 3rd first, then 1st becomes a producer with a single outgoing edge and it’s fused later so we cannot use this test.
However, I noticed that if I make ‘%a’ and external memref in my original test case, I can test exactly what I want and sibling fusion doesn’t kick in. I think we can proceed this way and remove clSiblingFusion for now, although I kind of like that flag. It's useful for debugging/independent testing.

dcaballe · 2019-11-26T01:24:38Z

Thanks for the comments, @bondhugula!
I can't add reviewers to MLIR PRs probably because I don't have permission. Maybe @andydavis1 is also interested?

andydavis1 · 2019-11-26T01:47:27Z

Hi Diego, thanks for contributing! I can take a look at this...

dcaballe

Modified original test and removed clSiblingFusion.
I also added a new test (and fixes related to it) that it asserted now. It’s a producer with multiple outgoing edges where there is a single store with a single outgoing edge.

Please, let me know if you have any other comments.
Thanks!

dcaballe · 2019-11-27T19:26:37Z

test/Transforms/loop-fusion.mlir

+// multiple outgoing edges. Sibling loop fusion is disabled to properly test
+// producer-consumer fusion in isolation.
+// CHECK-LABEL: func @should_not_fuse_multi_outgoing_edge_store_producer
+func @should_not_fuse_multi_outgoing_edge_store_producer() {


Thanks! I’ve been playing with the suggested test case:

func @should_not_fuse_multi_outgoing_edge_store_producer2(%a : memref<1xf32>, %b : memref<1xf32>) { %cst = constant 0.000000e+00 : f32 affine.for %arg0 = 0 to 1 { affine.store %cst, %a[%arg0] : memref<1xf32> } affine.for %arg0 = 0 to 1 { %0 = affine.load %a[%arg0] : memref<1xf32> affine.store %cst, %b[%arg0] : memref<1xf32> } affine.for %arg0 = 0 to 1 { %0 = affine.load %a[%arg0] : memref<1xf32> %1 = affine.load %b[%arg0] : memref<1xf32> } return }

Producer-consumer fusion fuses the three loops. It fuses 2nd and 3rd first, then 1st becomes a producer with a single outgoing edge and it’s fused later so we cannot use this test.
However, I noticed that if I make ‘%a’ and external memref in my original test case, I can test exactly what I want and sibling fusion doesn’t kick in. I think we can proceed this way and remove clSiblingFusion for now, although I kind of like that flag. It's useful for debugging/independent testing.

bondhugula · 2019-11-27T19:56:08Z

Modified original test and removed clSiblingFusion.
I also added a new test (and fixes related to it) that it asserted now. It’s a producer with multiple outgoing edges where there is a single store with a single outgoing edge.

Please, let me know if you have any other comments.
Thanks!

I'm a bit confused here. Why should the fusion be prevented here in cases like these where it's valid (referring to the test case you added) - where you have a single producer and multiple consumers? I recall this was what you had addressed in your earlier commit. Are you introducing the limitation back? :-)

dcaballe · 2019-11-27T20:31:44Z

I'm a bit confused here. Why should the fusion be prevented here in cases like these where it's valid (referring to the test case you added) - where you have a single producer and multiple consumers? I recall this was what you had addressed in your earlier commit. Are you introducing the limitation back? :-)

My previous commit introduced support for multi-store producers with a single outgoing edge within the function. That is still working fine. The test I added is intact: https://github.com/tensorflow/mlir/pull/259/files/b8e7f532a846e8a4f76b04473db21e9ffb7d15f7#diff-ea810dda41c98feaab3d122f33648fb5R2327. However, as part of that change, I asserted on that producers reaching canFuseSrcWhichWritesToLiveOut should have only one outgoing edge (https://github.com/tensorflow/mlir/pull/162/files#diff-5b47f3417e354b0773b092edc5d8740cR1014). That is not true because a single store may have multiple outgoing edges and loads can also have outgoing edges. That is the piece I’m reverting and adding test for it, which is a bit complicated due to the producer-consumer fusion and sibling fusion interaction.

Does it make sense now? :)

bondhugula · 2019-11-27T20:35:13Z

I'm a bit confused here. Why should the fusion be prevented here in cases like these where it's valid (referring to the test case you added) - where you have a single producer and multiple consumers? I recall this was what you had addressed in your earlier commit. Are you introducing the limitation back? :-)

My previous commit introduced support for multi-store producers with a single outgoing edge within the function. That is still working fine. The test I added is intact: https://github.com/tensorflow/mlir/pull/259/files/b8e7f532a846e8a4f76b04473db21e9ffb7d15f7#diff-ea810dda41c98feaab3d122f33648fb5R2327. However, as part of that change, I asserted on that producers reaching canFuseSrcWhichWritesToLiveOut should have only one outgoing edge (https://github.com/tensorflow/mlir/pull/162/files#diff-5b47f3417e354b0773b092edc5d8740cR1014). That is not true because a single store may have multiple outgoing edges and loads can also have outgoing edges. That is the piece I’m reverting and adding test for it, which is a bit complicated due to the producer-consumer fusion and sibling fusion interaction.

Does it make sense now? :)

I see - sounds good then.

tensorflow/mlir#162 introduced a bug that incorrectly allowed fusion of producer loops with multiple outgoing edges. This commit fixes that problem. It also introduces a new flag to disable sibling loop fusion so that we can test producer-consumer fusion in isolation. Closes #259 COPYBARA_INTEGRATE_REVIEW=tensorflow/mlir#259 from dcaballe:dcaballe/fix_multi_out_edge_producer_fusion 578d5661705fd5c56c555832d5e0528df88c5282 PiperOrigin-RevId: 283531105 Change-Id: I3a6173463ea20bd35555c24fa451bfbf2dfac098

tensorflow/mlir#162 introduced a bug that incorrectly allowed fusion of producer loops with multiple outgoing edges. This commit fixes that problem. It also introduces a new flag to disable sibling loop fusion so that we can test producer-consumer fusion in isolation. Closes tensorflow/mlir#259 COPYBARA_INTEGRATE_REVIEW=tensorflow/mlir#259 from dcaballe:dcaballe/fix_multi_out_edge_producer_fusion 578d5661705fd5c56c555832d5e0528df88c5282 PiperOrigin-RevId: 283531105

googlebot added the cla: yes label Nov 22, 2019

bondhugula reviewed Nov 25, 2019

View reviewed changes

Address initial feedback

2c01828

dcaballe added 4 commits November 26, 2019 21:32

Fix lit test check

f158781

Disable sibling fusion only with clSiblingFusion flag

dcb1ee6

Remove clSiblingFusion flag

bc2770e

Remove clSiglingFusion, add new lit test

578d566

dcaballe commented Nov 27, 2019

View reviewed changes

andydavis1 approved these changes Nov 27, 2019

View reviewed changes

ftynse added the kokoro:run label Dec 3, 2019

kokoro-team removed the kokoro:run label Dec 3, 2019

ftynse added the ready to pull label Dec 3, 2019

mlir-copybara-bot closed this in 69607ea Dec 3, 2019

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

AffineLoopFusion: Prevent fusion of multi-out-edge producer loops #259

AffineLoopFusion: Prevent fusion of multi-out-edge producer loops #259

dcaballe commented Nov 22, 2019

bondhugula Nov 25, 2019

bondhugula Nov 25, 2019

dcaballe Nov 26, 2019

andydavis1 Nov 27, 2019

andydavis1 Nov 27, 2019

bondhugula Nov 25, 2019 •

edited

dcaballe Nov 26, 2019

andydavis1 Nov 27, 2019

dcaballe Nov 27, 2019

dcaballe commented Nov 26, 2019

andydavis1 commented Nov 26, 2019

dcaballe left a comment

dcaballe Nov 27, 2019

bondhugula commented Nov 27, 2019

dcaballe commented Nov 27, 2019

bondhugula commented Nov 27, 2019

AffineLoopFusion: Prevent fusion of multi-out-edge producer loops #259

AffineLoopFusion: Prevent fusion of multi-out-edge producer loops #259

Conversation

dcaballe commented Nov 22, 2019

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

bondhugula Nov 25, 2019 • edited

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

dcaballe commented Nov 26, 2019

andydavis1 commented Nov 26, 2019

dcaballe left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

bondhugula commented Nov 27, 2019

dcaballe commented Nov 27, 2019

bondhugula commented Nov 27, 2019

bondhugula Nov 25, 2019 •

edited