Improves the blocking logic and fixing Issue #413 (#722)

Author: chencha3

include/imex/Dialect/XeTile/Transforms/Blocking.h

@@ -74,9 +74,8 @@ class XeTileConversion : public imex::XeConversionPattern<AnalysisT> {
         mlir::DenseI64ArrayAttr::get(src.getContext(), innerBlocks));
   }
 
-  xetile::TilePackOp addPackOp(mlir::Value src,
-                               llvm::ArrayRef<int64_t> targetBlkSizes,
-                               OpPatternRewriter &rewriter) const {
+  mlir::Value addPackOp(mlir::Value src, llvm::ArrayRef<int64_t> targetBlkSizes,
+                        OpPatternRewriter &rewriter) const {
     auto srcTy = src.getType().dyn_cast<mlir::VectorType>();
     assert(srcTy && targetBlkSizes.size() == 2);
     auto shape = srcTy.getShape();

include/imex/Utils/XeCommon.h

@@ -198,7 +198,7 @@ class TileUsageAnalysis {
 // optimizations.
 class PropagateAnalysis {
 private:
-  llvm::DenseMap<mlir::Operation *, mlir::DenseI64ArrayAttr> OpAttrMap;
+  llvm::DenseMap<mlir::Value, mlir::DenseI64ArrayAttr> OpAttrMap;
 
 public:
   PropagateAnalysis(mlir::Operation *op) {
@@ -214,11 +214,25 @@ class PropagateAnalysis {
     });
   }
 
-  bool maybeUpdated(mlir::Operation *op) { return OpAttrMap.count(op); }
+  bool maybeUpdated(mlir::Operation *op) const {
+    assert(op->getNumResults() == 1);
+    auto v = op->getResult(0);
+    return OpAttrMap.count(v);
+  }
+
+  mlir::DenseI64ArrayAttr getValue(mlir::Value value) const {
+    auto it = OpAttrMap.find(value);
+    if (it != OpAttrMap.end())
+      return it->second;
+    return {};
+  }
 
-  mlir::DenseI64ArrayAttr getValue(mlir::Operation *op) {
-    if (OpAttrMap.count(op))
-      return OpAttrMap[op];
+  mlir::DenseI64ArrayAttr getValue(mlir::Operation *op) const {
+    assert(op->getNumResults() == 1);
+    auto v = op->getResult(0);
+    auto it = OpAttrMap.find(v);
+    if (it != OpAttrMap.end())
+      return it->second;
     return {};
   }
 
@@ -256,9 +270,9 @@ class PropagateAnalysis {
 
       // stop when meet a function.
       if (!op || llvm::isa<mlir::FunctionOpInterface>(op))
-        return;
+        continue;
 
-      OpAttrMap[op] = attr;
+      OpAttrMap[value] = attr;
 
       if (auto forOp = llvm::dyn_cast<mlir::scf::ForOp>(op)) {
         auto opr = getOperandForArg(forOp, value);
@@ -447,6 +461,10 @@ class XeConversionPattern : public mlir::RewritePattern {
     return {};
   }
 
+  mlir::DenseI64ArrayAttr getValue(mlir::Value value) const {
+    return llvm::cast<PropagateAnalysis>(analysis).getValue(value);
+  }
+
   template <typename = typename std::enable_if<
                 std::is_same_v<AnalysisT, TileUsageAnalysis>>>
   bool isForDPASA(imex::xetile::LoadTileOp op) const {

lib/Dialect/XeTile/Transforms/BlockAligning.cpp

@@ -99,23 +99,29 @@ struct SCFForOpPattern
   ::mlir::LogicalResult
   matchAndRewrite(mlir::scf::ForOp op, OpAdaptor adaptor,
                   OpPatternRewriter &rewriter) const override {
+    // we need to update the SCF ForOp if the types of its init arg values
+    // do not match the types of the region iter args, or the init arg value
+    // is defined by a TilePackOp. Otherwise we can skip the op.
+    bool changed = false;
     llvm::SmallVector<mlir::Value> newInitArgs;
-    llvm::SmallVector<mlir::DenseI64ArrayAttr> oldBlockSize;
-    llvm::SmallVector<mlir::DenseI64ArrayAttr> newBlockSize;
-    for (auto arg : adaptor.getInitArgs()) {
+    llvm::SmallVector<mlir::DenseI64ArrayAttr> oldBlockSizes;
+    llvm::SmallVector<mlir::DenseI64ArrayAttr> newBlockSizes;
+    for (auto [i, arg] : llvm::enumerate(adaptor.getInitArgs())) {
+      auto blockArg = op.getRegionIterArg(i);
       auto defOp = arg.getDefiningOp<xetile::TilePackOp>();
-      if (auto blockSize = getValue(defOp)) {
-        newBlockSize.push_back(blockSize);
-        oldBlockSize.push_back(defOp.getInnerBlocksAttr());
-        auto repackOp = addUnpackAndPackOps(arg, blockSize, rewriter);
-        newInitArgs.push_back(repackOp);
-      } else {
-        oldBlockSize.push_back({});
-        newBlockSize.push_back({});
-        newInitArgs.push_back(arg);
-      }
+      auto oldSize = defOp ? defOp.getInnerBlocksAttr() : DenseI64ArrayAttr();
+      auto newSize = defOp ? getValue(blockArg) : DenseI64ArrayAttr();
+      auto newArg =
+          defOp && newSize ? addUnpackAndPackOps(arg, newSize, rewriter) : arg;
+      oldBlockSizes.push_back(oldSize);
+      newBlockSizes.push_back(newSize);
+      newInitArgs.push_back(newArg);
+      changed |= (newArg.getType() != blockArg.getType());
     }
 
+    if (!changed)
+      return mlir::failure();
+
     auto newOp = rewriter.create<mlir::scf::ForOp>(
         op.getLoc(), adaptor.getLowerBound(), adaptor.getUpperBound(),
         adaptor.getStep(), newInitArgs);
@@ -124,15 +130,19 @@ struct SCFForOpPattern
     mlir::Block *newBlock = newOp.getBody();
     llvm::SmallVector<mlir::Value> newArguments;
     auto numCtrlOprs = newOp.getNumInductionVars();
+    // remove the terminator of the new block
+    if (newBlock->mightHaveTerminator())
+      rewriter.eraseOp(newBlock->getTerminator());
+
     // add UnpackOp and PackOp pairs to the block arguments
     // if the corresponding init arg is repacked, such that
     // the old unpack op using it in the body will be folded
     for (auto [i, arg] : llvm::enumerate(newBlock->getArguments())) {
-      if (i < numCtrlOprs || !oldBlockSize[i - numCtrlOprs]) {
+      if (i < numCtrlOprs || !oldBlockSizes[i - numCtrlOprs]) {
         newArguments.push_back(arg);
       } else {
         auto repackOp =
-            addUnpackAndPackOps(arg, oldBlockSize[i - numCtrlOprs], rewriter);
+            addUnpackAndPackOps(arg, oldBlockSizes[i - numCtrlOprs], rewriter);
         newArguments.push_back(repackOp);
       }
     }
@@ -145,9 +155,9 @@ struct SCFForOpPattern
     mlir::OpBuilder::InsertionGuard g(rewriter);
     rewriter.startOpModification(yieldOp);
     for (auto [i, v] : llvm::enumerate(yieldOp.getResults())) {
-      if (newBlockSize[i]) {
+      if (newBlockSizes[i]) {
         rewriter.setInsertionPointAfter(v.getDefiningOp());
-        auto repack = addUnpackAndPackOps(v, newBlockSize[i], rewriter);
+        auto repack = addUnpackAndPackOps(v, newBlockSizes[i], rewriter);
         yieldOp->setOperand(i, repack);
       }
     }
@@ -157,8 +167,8 @@ struct SCFForOpPattern
     rewriter.setInsertionPointAfter(op);
     llvm::SmallVector<mlir::Value> newValues;
     for (auto [i, result] : llvm::enumerate(newOp->getResults())) {
-      if (newInitArgs[i].getDefiningOp<xetile::TilePackOp>()) {
-        auto unpack = addUnpackAndPackOps(result, oldBlockSize[i], rewriter);
+      if (oldBlockSizes[i]) {
+        auto unpack = addUnpackAndPackOps(result, oldBlockSizes[i], rewriter);
         newValues.push_back(unpack);
       } else {
         newValues.push_back(result);
@@ -247,6 +257,9 @@ struct UpdateTileOffsetOpPattern
   ::mlir::LogicalResult
   matchAndRewrite(xetile::UpdateTileOffsetOp op, OpAdaptor adaptor,
                   OpPatternRewriter &rewriter) const override {
+    if (adaptor.getTile().getType() == op.getResult().getType())
+      return mlir::failure();
+
     rewriter.replaceOpWithNewOp<xetile::UpdateTileOffsetOp>(
         op, adaptor.getTile().getType(), adaptor.getTile(),
         adaptor.getOffsetX(), adaptor.getOffsetY());
@@ -291,9 +304,9 @@ class XeTileBlockAligningPass : public imex::impl::XeTileBlockAligningBase<
     // Use TopDown traversal order, and only look at existing ops
     // to simpliy the code logic and speedup the pass
     mlir::GreedyRewriteConfig config;
+    config.enableRegionSimplification = false;
     config.useTopDownTraversal = true;
-    config.maxIterations = 2;
-    config.strictMode = GreedyRewriteStrictness::ExistingOps;
+    config.strictMode = GreedyRewriteStrictness::ExistingAndNewOps;
     if (failed(
             applyPatternsAndFoldGreedily(mod, std::move(patterns), config))) {
       return signalPassFailure();