[Mlir-commits] [mlir] [mlir][Affine] Expand affine.[de]linearize_index without affine maps (PR #116703)

[llvmlistbot at llvm.org]

================
@@ -58,13 +147,44 @@ struct LowerLinearizeIndexOps final : OpRewritePattern<AffineLinearizeIndexOp> {
       return success();
     }
 
-    SmallVector<OpFoldResult> multiIndex =
-        getAsOpFoldResult(op.getMultiIndex());
-    OpFoldResult linearIndex =
-        linearizeIndex(rewriter, op.getLoc(), multiIndex, op.getMixedBasis());
-    Value linearIndexValue =
-        getValueOrCreateConstantIntOp(rewriter, op.getLoc(), linearIndex);
-    rewriter.replaceOp(op, linearIndexValue);
+    Location loc = op.getLoc();
+    ValueRange multiIndex = op.getMultiIndex();
+    size_t numIndexes = multiIndex.size();
+    ArrayRef<int64_t> staticBasis = op.getStaticBasis();
+    if (numIndexes == staticBasis.size())
+      staticBasis = staticBasis.drop_front();
+
+    SmallVector<Value> strides =
+        computeStrides(loc, rewriter, op.getDynamicBasis(), staticBasis);
+    SmallVector<std::pair<Value, int64_t>> scaledValues;
+    scaledValues.reserve(numIndexes);
+
+    // Note: strides doesn't contain a value for the final element (stride 1)
+    // and everything else lines up. We use the "mutable" accessor so we can get
+    // our hands on an `OpOperand&` for the loop invariant counting function.
+    for (auto [stride, idxOp] :
+         llvm::zip_equal(strides, llvm::drop_end(op.getMultiIndexMutable()))) {
+      Value scaledIdx =
+          rewriter.create<arith::MulIOp>(loc, idxOp.get(), stride);
+      int64_t numHoistableLoops = numEnclosingInvariantLoops(idxOp);
+      scaledValues.emplace_back(scaledIdx, numHoistableLoops);
+    }
+    scaledValues.emplace_back(
+        multiIndex.back(),
+        numEnclosingInvariantLoops(op.getMultiIndexMutable()[numIndexes - 1]));
----------------
MaheshRavishankar wrote:

Ok I see this from  the commit message 
```
In addition, the lowering of affine.linearize_index now sorts the operands by loop-independence, allowing an increased amount of loop-invariant code motion after lowering.
```

I think I understand what this does... that is interesting. I wonder if we can do this separately as an optimizations instead of linking it to linearization lowering. Take a sequence of `arith.mul` operations and reorder them to allow for better hoisting. 

https://github.com/llvm/llvm-project/pull/116703