[mlir] Mode for explicitly controlling the fusion kind

[sumesh13]

Summary

New mode option that allows for either running the default fusion kind that happens today or doing either of producer-consumer or sibling fusion. This will also be helpful to minimize the compile-time of the fusion tests.

[bondhugula]

Thanks for adding this. Please use enum valued cl option to support this.

mlir/include/mlir/Transforms/Passes.h

82

Missing documentation for new arg.

mlir/include/mlir/Transforms/Passes.td

139

Nit: space after comma.

139

Please avoid using hard coded values like 0, 1, ... as pass options. It's not readable when used in run commands, textual pass pipelines, etc. Please instead use enum valued command-line options. See test-legalize-mode for reference..

140

Producer- -> producer->

mlir/lib/Transforms/LoopFusion.cpp

1409–1414

All functions are supposed to have doc comments.

[bondhugula]

mlir/include/mlir/Transforms/Passes.h

85–86

Why is this a bool? Don't you have three modes? If this is a bug, the testing isn't adequate.

mlir/lib/Transforms/LoopFusion.cpp

67

How has unsigned become a bool here?!

97–98

Triple /// comments.

mlir/test/Transforms/loop-fusion-4.mlir

1–3

Use enum valued command line options. Eg:

mlir-opt -h | grep test-legalize-mode -A 3
  --test-legalize-mode=<value>                          - The legalization mode to use with the test driver
    =analysis                                           -   Perform an analysis conversion
    =full                                               -   Perform a full conversion
    =partial                                            -   Perform a partial conversion

[dcaballe]

Thanks, Sumesh! A couple of comments...

mlir/lib/Transforms/LoopFusion.cpp

1414

Probably good to rename this now to something more meaningful.

1415

IIRC, maxSrcUserCount = 1 was used with the purpose of enabling some sibling fusion scenarios. I think this should be std::numeric_limits<unsigned>::max() instead?

[sumesh13]

Address code review comments

• Input option as enum
• Producer Fusion with max src count

[sumesh13]

mlir/lib/Transforms/LoopFusion.cpp

1414

Makes sense. renamed as runGreedyFusion

[bondhugula]

mlir/include/mlir/Transforms/Passes.td

142

Not immediately clear how the "greedy" aspect is different from the other two. Rephrase to:
Perform greedy (both producer-consumer and sibling) fusion.?

144–147

Nit: leave a space after commas.

mlir/test/Transforms/loop-fusion-4.mlir

1–2

You don't need to repeat affine-fusion again for the option: just -affine-loop-fusion="mode=producer" should be good enough I think.

7–8

Nit: Use a hyphen: SIBLING-MAXIMAL-...

[sumesh13]

Address code review comments

[bondhugula]

LGTM - thanks.

mlir/test/Transforms/loop-fusion-4.mlir

37–42

Nit: hyphen between PRODUCER and CONSUMER.

[dcaballe]

Thanks!

mlir/include/mlir/Transforms/Passes.h

 #include "llvm/Support/Debug.h"
 #include <limits>
 
 namespace mlir {
 
 class AffineForOp;
 class GreedyRewriteConfig;
 
+/// Fusion mode to attempt. The default mode `Greedy` does both
+/// producer-consumer and sibling fusion.
+enum FusionMode { Greedy, ProducerConsumer, Sibling };
+
 //===----------------------------------------------------------------------===//
 // Passes
 //===----------------------------------------------------------------------===//
 
 /// Creates an instance of the BufferDeallocation pass to free all allocated
 /// buffers.
 std::unique_ptr<Pass> createBufferDeallocationPass();
 

 
 /// Creates an instance of the Canonicalizer pass with the specified config.
 std::unique_ptr<Pass>
 createCanonicalizerPass(const GreedyRewriteConfig &config);
 
 /// Creates a pass to perform common sub expression elimination.
 std::unique_ptr<Pass> createCSEPass();
 
-/// Creates a loop fusion pass which fuses loops. Buffers of size less than or
-/// equal to `localBufSizeThreshold` are promoted to memory space
-/// `fastMemorySpace'.
+/// Creates a loop fusion pass which fuses loops according to type of fusion
+/// specified in `fusionMode`. Buffers of size less than or equal to
+/// `localBufSizeThreshold` are promoted to memory space `fastMemorySpace`.
 std::unique_ptr<OperationPass<FuncOp>>
 createLoopFusionPass(unsigned fastMemorySpace = 0,
                      uint64_t localBufSizeThreshold = 0,
-                     bool maximalFusion = false);
+                     bool maximalFusion = false,
+                     enum FusionMode fusionMode = FusionMode::Greedy);
 
 /// Creates a loop invariant code motion pass that hoists loop invariant
 /// instructions out of the loop.
 std::unique_ptr<Pass> createLoopInvariantCodeMotionPass();
 
 /// Creates a pass to pipeline explicit movement of data across levels of the
 /// memory hierarchy.
 std::unique_ptr<OperationPass<FuncOp>> createPipelineDataTransferPass();

mlir/include/mlir/Transforms/Passes.td

let options = [
     Option<"fastMemorySpace", "fusion-fast-mem-space", "unsigned",
            /*default=*/"0",
            "Faster memory space number to promote fusion buffers to">,
     Option<"localBufSizeThreshold", "fusion-local-buf-threshold", "uint64_t",
            /*default=*/"0", "Threshold size (KiB) for promoting local buffers "
                             "to fast memory space">,
     Option<"maximalFusion", "fusion-maximal", "bool", /*default=*/"false",
            "Enables maximal loop fusion">,
+    Option<"affineFusionMode", "mode", "enum FusionMode",
+           "mlir::FusionMode::Greedy", "fusion mode to attempt",
+           "llvm::cl::values(clEnumValN(mlir::FusionMode::Greedy,"
+           " \"greedy\", \"Perform greedy (both producer-consumer and sibling)  fusion\"), "
+           "clEnumValN( mlir::FusionMode::ProducerConsumer, "
+           "\"producer\", \"Perform only producer-consumer fusion\"), "
+           "clEnumValN( mlir::FusionMode::Sibling, "
+           "\"sibling\", \"Perform only sibling fusion\"))">,
     ];
   let dependentDialects = ["memref::MemRefDialect"];
 }
 
 def AffinePipelineDataTransfer
     : FunctionPass<"affine-pipeline-data-transfer"> {
   let summary = "Pipeline non-blocking data transfers between explicitly "
                 "managed levels of the memory hierarchy";
   let description = [{

mlir/lib/Transforms/LoopFusion.cpp

 // TODO: Support fusion of source loop nests which write to multiple
 // memrefs, where each memref can have multiple users (if profitable).
 // TODO: Extend this pass to check for fusion preventing dependences,
 // and add support for more general loop fusion algorithms.
 
 struct LoopFusion : public AffineLoopFusionBase<LoopFusion> {
   LoopFusion() = default;
   LoopFusion(unsigned fastMemorySpace, uint64_t localBufSizeThresholdBytes,
-             bool maximalFusion) {
+             bool maximalFusion, enum FusionMode affineFusionMode) {
     this->fastMemorySpace = fastMemorySpace;
     this->localBufSizeThreshold = localBufSizeThresholdBytes / 1024;
     this->maximalFusion = maximalFusion;
+    this->affineFusionMode = affineFusionMode;
   }
 
   void runOnFunction() override;
 };
 
 } // end anonymous namespace
 
 std::unique_ptr<OperationPass<FuncOp>>
 mlir::createLoopFusionPass(unsigned fastMemorySpace,
-                           uint64_t localBufSizeThreshold, bool maximalFusion) {
+                           uint64_t localBufSizeThreshold, bool maximalFusion,
+                           enum FusionMode affineFusionMode) {
   return std::make_unique<LoopFusion>(fastMemorySpace, localBufSizeThreshold,
-                                      maximalFusion);
+                                      maximalFusion, affineFusionMode);
 }
 
 namespace {
 
 // LoopNestStateCollector walks loop nests and collects load and store
 // operations, and whether or not a region holding op other than ForOp and IfOp
 // was encountered in the loop nest.
 struct LoopNestStateCollector {
opToWalk->walk([&](Operation *op) {
       else if (isa<AffineReadOpInterface>(op))
         loadOpInsts.push_back(op);
       else if (isa<AffineWriteOpInterface>(op))
         storeOpInsts.push_back(op);
     });
   }
 };
 
 // MemRefDependenceGraph is a graph data structure where graph nodes are
 // top-level operations in a FuncOp which contain load/store ops, and edges
 // are memref dependences between the nodes.
 // TODO: Add a more flexible dependence graph representation.
 // TODO: Add a depth parameter to dependence graph construction.
 struct MemRefDependenceGraph {
 public:
   // Node represents a node in the graph. A Node is either an entire loop nest
   // rooted at the top level which contains loads/stores, or a top level
   // load/store.
void init() {
     // TODO: Add a priority queue for prioritizing nodes by different
     // metrics (e.g. arithmetic intensity/flops-to-bytes ratio).
     worklist.clear();
     for (auto &idAndNode : mdg->nodes) {
       const Node &node = idAndNode.second;
       worklist.push_back(node.id);
     }
   }
+  /// Run only sibling fusion on the `mdg`.
+  void runSiblingFusionOnly() {
+    fuseSiblingNodes();
+    eraseUnusedMemRefAllocations();
+  }
+
+  /// Run only producer/consumer fusion on the `mdg`.
+  void runProducerConsumerFusionOnly() {
+    fuseProducerConsumerNodes(
+        /*maxSrcUserCount=*/std::numeric_limits<unsigned>::max());
+    eraseUnusedMemRefAllocations();
+  }
 
   // Run the GreedyFusion pass.
   // *) First pass through the nodes fuses single-use producer nodes into their
   //    unique consumer.
   // *) Second pass fuses sibling nodes which share no dependence edges.
   // *) Third pass fuses any remaining producer nodes into their users.
-  void run() {
+  void runGreedyFusion() {
     // TODO: Run this repeatedly until a fixed-point is reached.
     fuseProducerConsumerNodes(/*maxSrcUserCount=*/1);
     fuseSiblingNodes();
     fuseProducerConsumerNodes(
         /*maxSrcUserCount=*/std::numeric_limits<unsigned>::max());
     eraseUnusedMemRefAllocations();
   }
 
   void fuseProducerConsumerNodes(unsigned maxSrcUserCount) {
if (!g.init(getFunction()))
     return;
 
   Optional<unsigned> fastMemorySpaceOpt;
   if (fastMemorySpace.hasValue())
     fastMemorySpaceOpt = fastMemorySpace;
   unsigned localBufSizeThresholdBytes = localBufSizeThreshold * 1024;
   GreedyFusion fusion(&g, localBufSizeThresholdBytes, fastMemorySpaceOpt,
                       maximalFusion, computeToleranceThreshold);
-  fusion.run();
+
+  if (affineFusionMode == FusionMode::ProducerConsumer)
+    fusion.runProducerConsumerFusionOnly();
+  else if (affineFusionMode == FusionMode::Sibling)
+    fusion.runSiblingFusionOnly();
+  else
+    fusion.runGreedyFusion();
 }

mlir/lib/Transforms/PassDetail.h

 //===- PassDetail.h - Transforms Pass class details -------------*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef TRANSFORMS_PASSDETAIL_H_
 #define TRANSFORMS_PASSDETAIL_H_
 
 #include "mlir/Pass/Pass.h"
+#include "mlir/Transforms/Passes.h"
 
 namespace mlir {
 class AffineDialect;
 
 // Forward declaration from Dialect.h
 template <typename ConcreteDialect>
 void registerDialect(DialectRegistry &registry);
 

mlir/test/Transforms/loop-fusion-4.mlir

-// RUN: mlir-opt -allow-unregistered-dialect %s -affine-loop-fusion -split-input-file | FileCheck %s
-// RUN: mlir-opt -allow-unregistered-dialect %s -affine-loop-fusion="fusion-maximal" -split-input-file | FileCheck %s --check-prefix=MAXIMAL
+// RUN: mlir-opt -allow-unregistered-dialect %s -affine-loop-fusion="mode=producer" -split-input-file | FileCheck %s --check-prefix=PRODUCER-CONSUMER
+// RUN: mlir-opt -allow-unregistered-dialect %s -affine-loop-fusion="fusion-maximal mode=sibling" -split-input-file | FileCheck %s --check-prefix=SIBLING-MAXIMAL
 
-// Part I of fusion tests in  mlir/test/Transforms/loop-fusion.mlir. 
+// Part I of fusion tests in  mlir/test/Transforms/loop-fusion.mlir.
 // Part II of fusion tests in mlir/test/Transforms/loop-fusion-2.mlir
 // Part III of fusion tests in mlir/test/Transforms/loop-fusion-3.mlir
 
-// -----
-
-func @reduce_add_non_maximal_f32_f32(%arg0: memref<64x64xf32, 1>, %arg1 : memref<1x64xf32, 1>, %arg2 : memref<1x64xf32, 1>) {
-    %cst_0 = constant 0.000000e+00 : f32
-    %cst_1 = constant 1.000000e+00 : f32
-    affine.for %arg3 = 0 to 1 {
-      affine.for %arg4 = 0 to 64 {
-        %accum = affine.for %arg5 = 0 to 64 iter_args (%prevAccum = %cst_0) -> f32 {
-          %4 = affine.load %arg0[%arg5, %arg4] : memref<64x64xf32, 1>
-          %5 = addf %prevAccum, %4 : f32
-          affine.yield %5 : f32
-        }
-        %accum_dbl = addf %accum, %accum : f32
-        affine.store %accum_dbl, %arg1[%arg3, %arg4] : memref<1x64xf32, 1>
-      }
-    }
-    affine.for %arg3 = 0 to 1 {
-      affine.for %arg4 = 0 to 64 {
-        // Following loop  trip count does not match the corresponding source trip count.
-        %accum = affine.for %arg5 = 0 to 32 iter_args (%prevAccum = %cst_1) -> f32 {
-          %4 = affine.load %arg0[%arg5, %arg4] : memref<64x64xf32, 1>
-          %5 = mulf %prevAccum, %4 : f32
-          affine.yield %5 : f32
-        }
-        %accum_sqr = mulf %accum, %accum : f32
-        affine.store %accum_sqr, %arg2[%arg3, %arg4] : memref<1x64xf32, 1>
-      }
-    }
-    return
-}
-// Test checks the loop structure is preserved after sibling fusion
-// since the destination loop and source loop trip counts do not
-// match.
-// MAXIMAL-LABEL:   func @reduce_add_non_maximal_f32_f32(
-// MAXIMAL:        %[[cst_0:.*]] = constant 0.000000e+00 : f32
-// MAXIMAL-NEXT:        %[[cst_1:.*]] = constant 1.000000e+00 : f32
-// MAXIMAL-NEXT:           affine.for %[[idx_0:.*]]= 0 to 1 {
-// MAXIMAL-NEXT:             affine.for %[[idx_1:.*]] = 0 to 64 {
-// MAXIMAL-NEXT:               %[[result_1:.*]] = affine.for %[[idx_2:.*]] = 0 to 32 iter_args(%[[iter_0:.*]] = %[[cst_1]]) -> (f32) {
-// MAXIMAL-NEXT:                 %[[result_0:.*]] = affine.for %[[idx_3:.*]] = 0 to 64 iter_args(%[[iter_1:.*]] = %[[cst_0]]) -> (f32) {
-
 // Expects fusion of producer into consumer at depth 4 and subsequent removal of
 // source loop.
-// CHECK-LABEL: func @unflatten4d
+// PRODUCER-CONSUMER-LABEL: func @unflatten4d
 func @unflatten4d(%arg1: memref<7x8x9x10xf32>) {
   %m = memref.alloc() : memref<5040xf32>
   %cf7 = constant 7.0 : f32
 
   affine.for %i0 = 0 to 7 {
     affine.for %i1 = 0 to 8 {
       affine.for %i2 = 0 to 9 {
         affine.for %i3 = 0 to 10 {
affine.for %i1 = 0 to 8 {
           affine.store %v0, %arg1[%i0, %i1, %i2, %i3] : memref<7x8x9x10xf32>
         }
       }
     }
   }
   return
 }
 
-// CHECK:        affine.for
-// CHECK-NEXT:     affine.for
-// CHECK-NEXT:       affine.for
-// CHECK-NEXT:         affine.for
-// CHECK-NOT:    affine.for
-// CHECK: return
+// PRODUCER-CONSUMER:        affine.for
+// PRODUCER-CONSUMER-NEXT:     affine.for
+// PRODUCER-CONSUMER-NEXT:       affine.for
+// PRODUCER-CONSUMER-NEXT:         affine.for
+// PRODUCER-CONSUMER-NOT:    affine.for
+// PRODUCER-CONSUMER: return
 
 // -----
 
 // Expects fusion of producer into consumer at depth 2 and subsequent removal of
 // source loop.
-// CHECK-LABEL: func @unflatten2d_with_transpose
+// PRODUCER-CONSUMER-LABEL: func @unflatten2d_with_transpose
 func @unflatten2d_with_transpose(%arg1: memref<8x7xf32>) {
   %m = memref.alloc() : memref<56xf32>
   %cf7 = constant 7.0 : f32
 
   affine.for %i0 = 0 to 7 {
     affine.for %i1 = 0 to 8 {
       affine.store %cf7, %m[8 * %i0 + %i1] : memref<56xf32>
     }
   }
   affine.for %i0 = 0 to 8 {
     affine.for %i1 = 0 to 7 {
       %v0 = affine.load %m[%i0 + 8 * %i1] : memref<56xf32>
       affine.store %v0, %arg1[%i0, %i1] : memref<8x7xf32>
     }
   }
   return
 }
 
-// CHECK:        affine.for
-// CHECK-NEXT:     affine.for
-// CHECK-NOT:    affine.for
-// CHECK: return
+// PRODUCER-CONSUMER:        affine.for
+// PRODUCER-CONSUMER-NEXT:     affine.for
+// PRODUCER-CONSUMER-NOT:    affine.for
+// PRODUCER-CONSUMER: return
+
+// -----
+
+// SIBLING-MAXIMAL-LABEL:   func @reduce_add_non_maximal_f32_f32(
+func @reduce_add_non_maximal_f32_f32(%arg0: memref<64x64xf32, 1>, %arg1 : memref<1x64xf32, 1>, %arg2 : memref<1x64xf32, 1>) {
+    %cst_0 = constant 0.000000e+00 : f32
+    %cst_1 = constant 1.000000e+00 : f32
+    affine.for %arg3 = 0 to 1 {
+      affine.for %arg4 = 0 to 64 {
+        %accum = affine.for %arg5 = 0 to 64 iter_args (%prevAccum = %cst_0) -> f32 {
+          %4 = affine.load %arg0[%arg5, %arg4] : memref<64x64xf32, 1>
+          %5 = addf %prevAccum, %4 : f32
+          affine.yield %5 : f32
+        }
+        %accum_dbl = addf %accum, %accum : f32
+        affine.store %accum_dbl, %arg1[%arg3, %arg4] : memref<1x64xf32, 1>
+      }
+    }
+    affine.for %arg3 = 0 to 1 {
+      affine.for %arg4 = 0 to 64 {
+        // Following loop  trip count does not match the corresponding source trip count.
+        %accum = affine.for %arg5 = 0 to 32 iter_args (%prevAccum = %cst_1) -> f32 {
+          %4 = affine.load %arg0[%arg5, %arg4] : memref<64x64xf32, 1>
+          %5 = mulf %prevAccum, %4 : f32
+          affine.yield %5 : f32
+        }
+        %accum_sqr = mulf %accum, %accum : f32
+        affine.store %accum_sqr, %arg2[%arg3, %arg4] : memref<1x64xf32, 1>
+      }
+    }
+    return
+}
+// Test checks the loop structure is preserved after sibling fusion
+// since the destination loop and source loop trip counts do not
+// match.
+// SIBLING-MAXIMAL:        %[[cst_0:.*]] = constant 0.000000e+00 : f32
+// SIBLING-MAXIMAL-NEXT:        %[[cst_1:.*]] = constant 1.000000e+00 : f32
+// SIBLING-MAXIMAL-NEXT:           affine.for %[[idx_0:.*]]= 0 to 1 {
+// SIBLING-MAXIMAL-NEXT:             affine.for %[[idx_1:.*]] = 0 to 64 {
+// SIBLING-MAXIMAL-NEXT:               %[[result_1:.*]] = affine.for %[[idx_2:.*]] = 0 to 32 iter_args(%[[iter_0:.*]] = %[[cst_1]]) -> (f32) {
+// SIBLING-MAXIMAL-NEXT:                 %[[result_0:.*]] = affine.for %[[idx_3:.*]] = 0 to 64 iter_args(%[[iter_1:.*]] = %[[cst_0]]) -> (f32) {
 No newline at end of file