[MLIR][Affine] Extend/generalize MDG to properly add edges between non-affine ops (#125451)

Drop arbitrary checks and hacks from affine fusion MDG construction and
handle all ops using memory read/write effects. This has been a long
pending change and it now makes affine fusion more powerful in the
presence of non-affine ops and does not limit fusion in parts of the
block where it is feasible simply because of non-affine ops elsewhere or
intervening non-affine users.

Populate memref read and write ops in non-affine region holding ops and
non-affine ops at the top level of the Block properly; add the
appropriate edges to MDG. Use memory read-write effects and drop
assumptions and special handling of ops due to historic reasons.

Update MDG to drop unnecessary "unhandled region" hack. This hack is no
longer needed with the update to fully and properly construct the MDG.

MDG edges now capture dependences between nodes completely. Drop
non-affine users check. With the MDG generalization to properly include
edges
between non-affine nodes/operations, the non-affine users on path check
in fusion is no longer needed. Add more test cases to exercise MDG
generalization.

Drop unnecessary failure when encountering side-effect-free affine.if
ops.

Improve documentation on MDG.

Author: bondhugula

mlir/include/mlir/Dialect/Affine/Analysis/Utils.h

@@ -37,19 +37,28 @@ struct MemRefAccess;
 // was encountered in the loop nest.
 struct LoopNestStateCollector {
   SmallVector<AffineForOp, 4> forOps;
+  // Affine loads.
   SmallVector<Operation *, 4> loadOpInsts;
+  // Affine stores.
   SmallVector<Operation *, 4> storeOpInsts;
-  bool hasNonAffineRegionOp = false;
+  // Non-affine loads.
+  SmallVector<Operation *, 4> memrefLoads;
+  // Non-affine stores.
+  SmallVector<Operation *, 4> memrefStores;
+  // Free operations.
+  SmallVector<Operation *, 4> memrefFrees;
 
   // Collects load and store operations, and whether or not a region holding op
   // other than ForOp and IfOp was encountered in the loop nest.
   void collect(Operation *opToWalk);
 };
 
 // MemRefDependenceGraph is a graph data structure where graph nodes are
-// top-level operations in a `Block` which contain load/store ops, and edges
-// are memref dependences between the nodes.
-// TODO: Add a more flexible dependence graph representation.
+// top-level operations in a `Block` and edges are memref dependences or SSA
+// dependences (on memrefs) between the nodes. Nodes are created for all
+// top-level operations except in certain cases (see `init` method). Edges are
+// created between nodes with a dependence (see `Edge` documentation). Edges
+// aren't created from/to nodes that have no memory effects.
 struct MemRefDependenceGraph {
 public:
   // Node represents a node in the graph. A Node is either an entire loop nest
@@ -60,10 +69,18 @@ struct MemRefDependenceGraph {
     unsigned id;
     // The top-level statement which is (or contains) a load/store.
     Operation *op;
-    // List of load operations.
+    // List of affine loads.
     SmallVector<Operation *, 4> loads;
-    // List of store op insts.
+    // List of non-affine loads.
+    SmallVector<Operation *, 4> memrefLoads;
+    // List of affine store ops.
     SmallVector<Operation *, 4> stores;
+    // List of non-affine stores.
+    SmallVector<Operation *, 4> memrefStores;
+    // List of free operations.
+    SmallVector<Operation *, 4> memrefFrees;
+    // Set of private memrefs used in this node.
+    DenseSet<Value> privateMemrefs;
 
     Node(unsigned id, Operation *op) : id(id), op(op) {}
 
@@ -73,6 +90,13 @@ struct MemRefDependenceGraph {
     // Returns the store op count for 'memref'.
     unsigned getStoreOpCount(Value memref) const;
 
+    /// Returns true if there exists an operation with a write memory effect to
+    /// `memref` in this node.
+    unsigned hasStore(Value memref) const;
+
+    // Returns true if the node has a free op on `memref`.
+    unsigned hasFree(Value memref) const;
+
     // Returns all store ops in 'storeOps' which access 'memref'.
     void getStoreOpsForMemref(Value memref,
                               SmallVectorImpl<Operation *> *storeOps) const;
@@ -86,7 +110,16 @@ struct MemRefDependenceGraph {
     void getLoadAndStoreMemrefSet(DenseSet<Value> *loadAndStoreMemrefSet) const;
   };
 
-  // Edge represents a data dependence between nodes in the graph.
+  // Edge represents a data dependence between nodes in the graph. It can either
+  // be a memory dependence or an SSA dependence. In the former case, it
+  // corresponds to a pair of memory accesses to the same memref or aliasing
+  // memrefs where at least one of them has a write or free memory effect. The
+  // memory accesses need not be affine load/store operations. Operations are
+  // checked for read/write effects and edges may be added conservatively. Edges
+  // are not created to/from nodes that have no memory effect. An exception to
+  // this are SSA dependences between operations that define memrefs (like
+  // alloc's, view-like ops) and their memory-effecting users that are enclosed
+  // in loops.
   struct Edge {
     // The id of the node at the other end of the edge.
     // If this edge is stored in Edge = Node.inEdges[i], then
@@ -182,9 +215,12 @@ struct MemRefDependenceGraph {
   // of sibling node 'sibId' into node 'dstId'.
   void updateEdges(unsigned sibId, unsigned dstId);
 
-  // Adds ops in 'loads' and 'stores' to node at 'id'.
-  void addToNode(unsigned id, const SmallVectorImpl<Operation *> &loads,
-                 const SmallVectorImpl<Operation *> &stores);
+  // Adds the specified ops to lists of node at 'id'.
+  void addToNode(unsigned id, ArrayRef<Operation *> loads,
+                 ArrayRef<Operation *> stores,
+                 ArrayRef<Operation *> memrefLoads,
+                 ArrayRef<Operation *> memrefStores,
+                 ArrayRef<Operation *> memrefFrees);
 
   void clearNodeLoadAndStores(unsigned id);