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

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=https://github.com/tensorflow/mlir/pull/259 from dcaballe:dcaballe/fix_multi_out_edge_producer_fusion 578d5661705fd5c56c555832d5e0528df88c5282
PiperOrigin-RevId: 283531105

1 files changed, 53 insertions, 0 deletions

diff --git a/mlir/test/Transforms/loop-fusion.mlir b/mlir/test/Transforms/loop-fusion.mlir
index 7431eade896..339cc31f549 100644
--- a/mlir/test/Transforms/loop-fusion.mlir
+++ b/mlir/test/Transforms/loop-fusion.mlir
@@ -2388,6 +2388,59 @@ func @mul_add_0(%arg0: memref<3x4xf32>, %arg1: memref<4x3xf32>, %arg2: memref<3x
   // CHECK-NEXT:   }
   // CHECK-NEXT: }
   // CHECK-NEXT: return
+  return
+}
+
+// -----
+
+// Verify that 'fuseProducerConsumerNodes' doesn't fuse a producer loop with
+// a store that has multiple outgoing edges. Sibling loop fusion should not fuse
+// any of these loops due to dependencies on external memref '%a'.
 
+// CHECK-LABEL: func @should_not_fuse_multi_outgoing_edge_store_producer1
+func @should_not_fuse_multi_outgoing_edge_store_producer1(%a : 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.for %arg0 = 0 to 1 {
+    %0 = affine.load %a[%arg0] : memref<1xf32>
+  }
+  // CHECK: affine.for %{{.*}} = 0 to 1
+  // CHECK: affine.for %{{.*}} = 0 to 1
+  // CHECK: affine.for %{{.*}} = 0 to 1
+  return
+}
+
+// -----
+
+// Verify that 'fuseProducerConsumerNodes' fuses a producer loop that: 1) has
+// multiple outgoing edges, 2) producer store has a single outgoing edge.
+// Sibling loop fusion should not fuse any of these loops due to
+// dependencies on external memrefs '%a' and '%b'.
+
+// CHECK-LABEL: func @should_fuse_producer_with_multi_outgoing_edges
+func @should_fuse_producer_with_multi_outgoing_edges(%a : memref<1xf32>, %b : memref<1xf32>) {
+  %cst = constant 0.000000e+00 : f32
+  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 {
+    affine.store %cst, %a[%arg0] : memref<1xf32>
+    %1 = affine.load %b[%arg0] : memref<1xf32>
+  }
+  // CHECK: affine.for %{{.*}} = 0 to 1
+  // CHECK-NEXT: affine.load %[[A:.*]][{{.*}}]
+  // CHECK-NEXT: affine.store %{{.*}}, %[[B:.*]][{{.*}}]
+  // CHECK-NEXT: affine.store %{{.*}}, %[[A]]
+  // CHECK-NEXT: affine.load %[[B]]
+  // CHECK-NOT: affine.for %{{.*}}
   return
 }