[Affine-fusion] Fix a bug in mod detection

[vinayaka-bandishti]

Summary

Fix a bug in detecting unknown ids as mods of known ids that was
preventing certain fusions.

While at this, fix the function signature of detectAsMod function to
have output as the last argument.

[bondhugula]

LGTM - thanks.

[bondhugula]

mlir/lib/Analysis/FlatLinearValueConstraints.cpp

234–236

Missing update to the doc comment -- this one is pretty important. It's unclear what offset and num are otherwise.

324–327

Please use the ternary operator.

[vinayaka-bandishti]

Address comments

mlir/lib/Analysis/FlatLinearValueConstraints.cpp

 // `var_r = var_n mod divisor`.
 //
 // Additionally, in a special case of the above constaints where `q_expr` is an
 // variable itself that is not yet known (say `var_q`), it can be written as a
 // floordiv in the following way:
 //
 // `var_q = var_n floordiv divisor`.
 //
+// First 'num' dimensional variables starting at 'offset' are
+// derived/to-be-derived in terms of the remaining variables. The remaining
+// variables are assigned trivial affine expressions in `memo`. For example,
+// memo is initilized as follows for a `cst` with 5 dims, when offset=2, num=2:
+// memo ==>  d0  d1  .   .   d2 ...
+// cst  ==>  c0  c1  c2  c3  c4 ...
+//
 // Returns true if the above mod or floordiv are detected, updating 'memo' with
 // these new expressions. Returns false otherwise.
 static bool detectAsMod(const FlatLinearConstraints &cst, unsigned pos,
-                        int64_t lbConst, int64_t ubConst,
-                        SmallVectorImpl<AffineExpr> &memo,
-                        MLIRContext *context) {
+                        unsigned offset, unsigned num, int64_t lbConst,
+                        int64_t ubConst, MLIRContext *context,
+                        SmallVectorImpl<AffineExpr> &memo) {
   assert(pos < cst.getNumVars() && "invalid position");
 
   // Check if a divisor satisfying the condition `0 <= var_r <= divisor - 1` can
   // be determined.
   if (lbConst != 0 || ubConst < 1)
     return false;
   int64_t divisor = ubConst + 1;
 
for (unsigned curEquality = 0, numEqualities = cst.getNumEqualities();
     // TODO: Handle AffineSymbolExpr as well. There is no reason to restrict it
     // to dims themselves.
     auto dimExpr = dividendExpr.dyn_cast<AffineDimExpr>();
     if (!dimExpr)
       continue;
 
     // Express `var_r` as `var_n % divisor` and store the expression in `memo`.
     if (quotientCount >= 1) {
-      auto ub = cst.getConstantBound64(BoundType::UB, dimExpr.getPosition());
+      // Find the column corresponding to `dimExpr`. `num` columns starting at
+      // `offset` correspond to previously unknown variables. The column
+      // corresponding to the trivially known `dimExpr` can be on either side
+      // of these.
+      unsigned dimExprPos = dimExpr.getPosition();
+      unsigned dimExprCol = dimExprPos < offset ? dimExprPos : dimExprPos + num;
+      auto ub = cst.getConstantBound64(BoundType::UB, dimExprCol);
       // If `var_n` has an upperbound that is less than the divisor, mod can be
       // eliminated altogether.
       if (ub && *ub < divisor)
         memo[pos] = dimExpr;
       else
         memo[pos] = dimExpr % divisor;
       // If a unique quotient `var_q` was seen, it can be expressed as
       // `var_n floordiv divisor`.
       if (quotientCount == 1 && !memo[quotientPosition])
         memo[quotientPosition] = dimExpr.floorDiv(divisor) * quotientSign;
 
for (unsigned pos = 0; pos < getNumVars(); pos++) {
         if (*lbConst == *ubConst) {
           memo[pos] = getAffineConstantExpr(*lbConst, context);
           changed = true;
           continue;
         }
 
         // Detect a variable as modulo of another variable w.r.t a
         // constant.
-        if (detectAsMod(*this, pos, *lbConst, *ubConst, memo, context)) {
+        if (detectAsMod(*this, pos, offset, num, *lbConst, *ubConst, context,
+                        memo)) {
           changed = true;
           continue;
         }
       }
 
       // Detect a variable as a floordiv of an affine function of other
       // variables (divisor is a positive constant).
       if (detectAsFloorDiv(*this, pos, context, memo)) {

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

affine.for %i1 = 0 to 10 {
     %v0 = affine.load %m[%i1] : memref<10xf32>
   }
   // PRODUCER-CONSUMER:      affine.for %{{.*}} = 0 to 10 {
   // PRODUCER-CONSUMER-NEXT:   affine.store %{{.*}}, %{{.*}}[0] : memref<1xf32>
   // PRODUCER-CONSUMER-NEXT:   affine.load %{{.*}}[0] : memref<1xf32>
   // PRODUCER-CONSUMER-NEXT: }
   return
 }
+
+// -----
+
+// PRODUCER-CONSUMER-LABEL: @fuse_higher_dim_nest_into_lower_dim_nest
+func.func @fuse_higher_dim_nest_into_lower_dim_nest() {
+  %A = memref.alloc() : memref<8x12x128x64xf32>
+  %B = memref.alloc() : memref<8x128x12x64xf32>
+  affine.for %arg205 = 0 to 8 {
+    affine.for %arg206 = 0 to 128 {
+      affine.for %arg207 = 0 to 12 {
+        affine.for %arg208 = 0 to 64 {
+          %a = affine.load %A[%arg205, %arg207, %arg206, %arg208] : memref<8x12x128x64xf32>
+          affine.store %a, %B[%arg205, %arg206, %arg207, %arg208] : memref<8x128x12x64xf32>
+        }
+      }
+    }
+  }
+  %C = memref.alloc() : memref<8x128x768xf16>
+  affine.for %arg205 = 0 to 8 {
+    affine.for %arg206 = 0 to 128 {
+      affine.for %arg207 = 0 to 768 {
+        %b = affine.load %B[%arg205, %arg206, %arg207 floordiv 64, %arg207 mod 64] : memref<8x128x12x64xf32>
+        %c = arith.truncf %b : f32 to f16
+        affine.store %c, %C[%arg205, %arg206, %arg207] : memref<8x128x768xf16>
+      }
+    }
+  }
+
+  // Check that fusion happens into the innermost loop of the consumer.
+  // PRODUCER-CONSUMER:      affine.for
+  // PRODUCER-CONSUMER-NEXT:   affine.for %{{.*}} = 0 to 128
+  // PRODUCER-CONSUMER-NEXT:     affine.for %{{.*}} = 0 to 768
+  // PRODUCER-CONSUMER-NOT:  affine.for
+  // PRODUCER-CONSUMER:      return
+  return
+}