add aligning transform into blocking pass (#716)

add aligning transform into blocking pass

Author: chencha3

include/imex/Conversion/XeTileToXeGPU/XeTileToXeGPU.h

@@ -38,7 +38,8 @@ class XeGPUTypeConverter;
 
 /// Populate the given list with patterns rewrite XeTile Ops
 void populateXeTileToXeGPUConversionPatterns(XeGPUTypeConverter &converter,
-                                             mlir::RewritePatternSet &patterns);
+                                             mlir::RewritePatternSet &patterns,
+                                             imex::TileUsageAnalysis &analysis);
 
 /// Create a pass to convert the XeTile dialect to the XeGPU dialect.
 std::unique_ptr<mlir::OperationPass<mlir::gpu::GPUModuleOp>>

include/imex/Conversion/XeTileToXeGPU/XeTileToXeGPUConversion.h

@@ -38,8 +38,7 @@ namespace imex {
 
 class XeGPUTypeConverter : public imex::XeTypeConverter {
 public:
-  XeGPUTypeConverter(mlir::MLIRContext &context,
-                     TileUsageAnalysis *analysis = nullptr);
+  XeGPUTypeConverter(mlir::MLIRContext &context);
 
   std::optional<mlir::LogicalResult>
   convertTileType(xetile::TileType tileTy,
@@ -112,13 +111,15 @@ class XeGPUOneToNPatterRewriter : public mlir::PatternRewriter,
 };
 
 template <typename SourceOp>
-class SgXeTileToXeGPUConversion : public XeConversionPattern {
+class SgXeTileToXeGPUConversion
+    : public XeConversionPattern<TileUsageAnalysis> {
 public:
   SgXeTileToXeGPUConversion(mlir::MLIRContext *context,
                             XeGPUTypeConverter &typeConverter,
+                            TileUsageAnalysis &analysis,
                             mlir::PatternBenefit benefit = 1)
-      : XeConversionPattern(typeConverter, SourceOp::getOperationName(),
-                            benefit, context) {}
+      : XeConversionPattern(typeConverter, analysis,
+                            SourceOp::getOperationName(), benefit, context) {}
 
   using RangeT = llvm::ArrayRef<mlir::ValueRange>;
   using OpAdaptor = typename SourceOp::template GenericAdaptor<RangeT>;

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

@@ -1,4 +1,4 @@
-//===- Blocking.h ----- Blocking Pass -------*- C++ -*-===//
+//===- Blocking.h ----------- Blocking Pass ---------*- C++ -*------------===//
 //
 // Copyright 2024 Intel Corporation
 // Part of the IMEX Project, under the Apache License v2.0 with LLVM Exceptions.
@@ -32,16 +32,17 @@
 
 namespace imex {
 
-template <typename SourceOp>
-class XeTileConversion : public imex::XeConversionPattern {
+template <typename SourceOp, typename AnalysisT>
+class XeTileConversion : public imex::XeConversionPattern<AnalysisT> {
 public:
   using OpAdaptor = typename SourceOp::Adaptor;
   using OpPatternRewriter = typename mlir::PatternRewriter;
 
   XeTileConversion(mlir::MLIRContext *context, XeTypeConverter &typeConverter,
-                   mlir::PatternBenefit benefit = 1)
-      : XeConversionPattern(typeConverter, SourceOp::getOperationName(),
-                            benefit, context) {}
+                   AnalysisT &analysis, mlir::PatternBenefit benefit = 1)
+      : XeConversionPattern<AnalysisT>(typeConverter, analysis,
+                                       SourceOp::getOperationName(), benefit,
+                                       context) {}
 
   mlir::LogicalResult
   matchAndRewrite(mlir::Operation *op,
@@ -70,7 +71,7 @@ class XeTileConversion : public imex::XeConversionPattern {
     auto unpackTy = mlir::VectorType::get(unpackShape, srcTy.getElementType());
     return rewriter.create<xetile::TileUnpackOp>(
         src.getLoc(), unpackTy, src,
-        mlir::DenseI64ArrayAttr::get(getContext(), innerBlocks));
+        mlir::DenseI64ArrayAttr::get(src.getContext(), innerBlocks));
   }
 
   xetile::TilePackOp addPackOp(mlir::Value src,
@@ -86,7 +87,7 @@ class XeTileConversion : public imex::XeConversionPattern {
     auto packTy = mlir::VectorType::get(packShape, srcTy.getElementType());
     auto packOp = rewriter.create<xetile::TilePackOp>(
         src.getLoc(), packTy, src,
-        mlir::DenseI64ArrayAttr::get(getContext(), targetBlkSizes));
+        mlir::DenseI64ArrayAttr::get(src.getContext(), targetBlkSizes));
     return packOp;
   }
 
@@ -98,16 +99,16 @@ class XeTileConversion : public imex::XeConversionPattern {
   }
 };
 
-template <template <typename> class TraitType>
-class XeTileTraitConversion : public imex::XeConversionPattern {
+template <template <typename> class TraitType, typename AnalysisT>
+class XeTileTraitConversion : public imex::XeConversionPattern<AnalysisT> {
 public:
   XeTileTraitConversion(mlir::MLIRContext *context,
-                        XeTypeConverter &typeConverter,
+                        XeTypeConverter &typeConverter, AnalysisT &analysis,
                         mlir::PatternBenefit benefit = 1)
-      : XeConversionPattern(typeConverter, mlir::Pattern::MatchTraitOpTypeTag(),
-                            mlir::TypeID::get<TraitType>(), benefit, context) {}
+      : XeConversionPattern<AnalysisT>(
+            typeConverter, analysis, mlir::Pattern::MatchTraitOpTypeTag(),
+            mlir::TypeID::get<TraitType>(), benefit, context) {}
 };
-
 } // namespace imex
 
 #endif

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

@@ -39,6 +39,7 @@ std::unique_ptr<mlir::Pass> createXeTileInitDuplicatePass();
 
 std::unique_ptr<mlir::Pass>
 createXeTileBlockingPass(const std::string &device = "pvc");
+std::unique_ptr<mlir::Pass> createXeTileBlockAligningPass();
 
 ///
 void populateXeTileInitDuplicatePatterns(imex::XeTypeConverter &converter,

include/imex/Dialect/XeTile/Transforms/Passes.td

@@ -60,4 +60,25 @@ def XeTileBlocking : Pass<"xetile-blocking", "::mlir::gpu::GPUModuleOp">{
  ];
 }
 
+
+// TODO: [block-aligning] remove the following code when upstreaming the pass.
+// The pass is not supposed to be exposed to users. Temporary keep in case we
+// need debug it for down stream development.
+def XeTileBlockAligning: Pass <"xetile-block-aligning", "::mlir::gpu::GPUModuleOp"> {
+  let summary = "optimize the performance for mma.";
+
+  let description = [{
+    This transform is to optimize performance by aligning the block size among operators
+    to reduce in-register data movements. Currently, it mainly focues on the alignment
+    between load and MMA operators.
+  }];
+
+  let constructor = "imex::createXeTileBlockAligningPass()";
+  let dependentDialects = ["imex::xetile::XeTileDialect",
+                           "mlir::arith::ArithDialect",
+                           "mlir::gpu::GPUDialect",
+                           "mlir::memref::MemRefDialect",
+                           "mlir::vector::VectorDialect"];
+}
+
 #endif // _XeTile_PASSES_TD_INCLUDED_

include/imex/Utils/XeCommon.h

@@ -101,18 +101,6 @@ class TileUsageAnalysis {
     }); // walk on LoadTileOp
   };
 
-  uint getUsage(imex::xetile::InitTileOp op) {
-    if (Usage.count(op))
-      return Usage[op];
-    return UsageType::None;
-  }
-
-  uint getUsage(imex::xetile::LoadTileOp op) {
-    if (Usage.count(op))
-      return Usage[op];
-    return UsageType::None;
-  }
-
   bool isForDPASA(imex::xetile::LoadTileOp op) {
     if (Usage.count(op)) {
       return Usage[op] & UsageType::DPAS_A;
@@ -200,14 +188,93 @@ class TileUsageAnalysis {
   llvm::DenseMap<mlir::Operation *, uint> Usage;
 };
 
+// This analysis is used to propagate the inner block size of an operator
+// to its uses or users. Current implementation is to propagate the MMA
+// size used by an MMA operator to the definition (InitTileOp) for its operands.
+// TODO: This analysis can be extended to propagate the block size for other ops
+// such that it can be used as a general analysis for other block size
+// optimizations.
+class PropagateAnalysis {
+private:
+  llvm::DenseMap<mlir::Operation *, mlir::DenseI64ArrayAttr> OpAttrMap;
+
+public:
+  PropagateAnalysis(mlir::Operation *op) {
+    op->walk<mlir::WalkOrder::PostOrder>([&](xetile::TileMMAOp op) {
+      mlir::Operation *operation = op.getOperation();
+      for (auto value : operation->getOperands()) {
+        auto packOp = value.getDefiningOp<xetile::TilePackOp>();
+        if (packOp) {
+          auto blkSZ = packOp.getInnerBlocksAttr();
+          propagate(value, blkSZ);
+        }
+      }
+    });
+  }
+
+  bool maybeUpdated(mlir::Operation *op) { return OpAttrMap.count(op); }
+
+  mlir::DenseI64ArrayAttr getValue(mlir::Operation *op) {
+    if (OpAttrMap.count(op))
+      return OpAttrMap[op];
+    return {};
+  }
+
+private:
+  mlir::Operation *getDefineOrParentOp(mlir::Value value) {
+    if (llvm::isa<mlir::OpResult>(value))
+      return value.getDefiningOp();
+    if (auto arg = llvm::dyn_cast_or_null<mlir::BlockArgument>(value))
+      return arg.getOwner()->getParentOp();
+    return nullptr;
+  };
+
+  mlir::Value getOperandForArg(mlir::scf::ForOp &forOp, mlir::Value &value) {
+    auto arg = llvm::dyn_cast<mlir::BlockArgument>(value);
+    if (arg && arg.getArgNumber() >= forOp.getNumInductionVars()) {
+      auto &iterOperand = *forOp.getTiedLoopInit(arg);
+      auto numCtrlOperands = forOp.getNumControlOperands();
+      auto operandIdx = iterOperand.getOperandNumber();
+      return forOp.getInitArgs()[operandIdx - numCtrlOperands];
+    }
+    return mlir::Value();
+  };
+
+  void propagate(mlir::Value start, mlir::DenseI64ArrayAttr attr) {
+    llvm::SmallVector<mlir::Value> queue;
+    if (bool(start))
+      queue.push_back(start);
+
+    while (queue.size()) {
+      auto value = queue.pop_back_val();
+      if (!bool(value))
+        continue;
+
+      auto *op = getDefineOrParentOp(value);
+
+      // stop when meet a function.
+      if (!op || llvm::isa<mlir::FunctionOpInterface>(op))
+        return;
+
+      OpAttrMap[op] = attr;
+
+      if (auto forOp = llvm::dyn_cast<mlir::scf::ForOp>(op)) {
+        auto opr = getOperandForArg(forOp, value);
+        if (bool(opr))
+          queue.push_back(opr);
+      } else if (op->getNumOperands() == 1) {
+        queue.push_back(op->getOperand(0));
+      }
+    }
+  }
+};
+
 class XeTypeConverter : public mlir::OneToNTypeConverter {
 public:
-  friend class XeConversionPattern;
+  // friend class XeConversionPattern;
   using mlir::OneToNTypeConverter::convertType;
 
-  XeTypeConverter(mlir::MLIRContext &context,
-                  TileUsageAnalysis *analysis = nullptr)
-      : context(context), usageAnalysis(analysis) {
+  XeTypeConverter(mlir::MLIRContext &context) : context(context) {
     addConversion([&](xetile::TileType tileTy,
                       llvm::SmallVectorImpl<mlir::Type> &resultTypes)
                       -> std::optional<mlir::LogicalResult> {
@@ -235,72 +302,95 @@ class XeTypeConverter : public mlir::OneToNTypeConverter {
 
 private:
   mlir::MLIRContext &context;
-
-protected:
-  TileUsageAnalysis *usageAnalysis;
 };
 
 // A simple mlir::RewritePattern wrapper with methods for accessing UsageType
+template <typename AnalysisT>
 class XeConversionPattern : public mlir::RewritePattern {
 public:
   using mlir::RewritePattern::RewritePattern;
 
   template <typename... Args>
-  XeConversionPattern(imex::XeTypeConverter &typeConverter, Args &&...args)
+  XeConversionPattern(imex::XeTypeConverter &typeConverter, AnalysisT &analysis,
+                      Args &&...args)
       : mlir::RewritePattern(std::forward<Args>(args)...),
-        typeConverter(typeConverter) {}
+        typeConverter(typeConverter), analysis(analysis) {}
 
   virtual mlir::LogicalResult
   matchAndRewrite(mlir::Operation *op,
                   mlir::PatternRewriter &rewriter) const override {
     llvm_unreachable("must override matchAndRewrite or a rewrite method");
   };
 
+  imex::XeTypeConverter &getTypeConverter() const { return typeConverter; }
+
+  template <typename ConverterTy>
+  std::enable_if_t<std::is_base_of<mlir::TypeConverter, ConverterTy>::value,
+                   ConverterTy &>
+  getTypeConverter() const {
+    return static_cast<ConverterTy &>(typeConverter);
+  }
+
+protected:
+  imex::XeTypeConverter &typeConverter;
+  AnalysisT &analysis;
+
+  template <typename = typename std::enable_if<
+                std::is_same_v<AnalysisT, PropagateAnalysis>>>
+  mlir::DenseI64ArrayAttr getValue(mlir::Operation *op) const {
+    if (op)
+      return llvm::cast<PropagateAnalysis>(analysis).getValue(op);
+    return {};
+  }
+
+  template <typename = typename std::enable_if<
+                std::is_same_v<AnalysisT, TileUsageAnalysis>>>
   bool isForDPASA(imex::xetile::LoadTileOp op) const {
-    return typeConverter.usageAnalysis->isForDPASA(op);
+    return llvm::cast<TileUsageAnalysis>(analysis).isForDPASA(op);
   }
 
+  template <typename = typename std::enable_if<
+                std::is_same_v<AnalysisT, TileUsageAnalysis>>>
   bool isForDPASB(imex::xetile::LoadTileOp op) const {
-    return typeConverter.usageAnalysis->isForDPASB(op);
+    return llvm::cast<TileUsageAnalysis>(analysis).isForDPASB(op);
   }
 
+  template <typename = typename std::enable_if<
+                std::is_same_v<AnalysisT, TileUsageAnalysis>>>
   bool isForDPASC(imex::xetile::LoadTileOp op) const {
-    return typeConverter.usageAnalysis->isForDPASC(op);
+    return llvm::cast<TileUsageAnalysis>(analysis).isForDPASC(op);
   }
 
+  template <typename = typename std::enable_if<
+                std::is_same_v<AnalysisT, TileUsageAnalysis>>>
   bool isForLoad(imex::xetile::InitTileOp op) const {
-    return typeConverter.usageAnalysis->isForLoad(op);
+    return llvm::cast<TileUsageAnalysis>(analysis).isForLoad(op);
   }
 
+  template <typename = typename std::enable_if<
+                std::is_same_v<AnalysisT, TileUsageAnalysis>>>
   bool isForStore(imex::xetile::InitTileOp op) const {
-    return typeConverter.usageAnalysis->isForStore(op);
+    return llvm::cast<TileUsageAnalysis>(analysis).isForStore(op);
   }
 
+  template <typename = typename std::enable_if<
+                std::is_same_v<AnalysisT, TileUsageAnalysis>>>
   bool isForPrefetch(imex::xetile::InitTileOp op) const {
-    return typeConverter.usageAnalysis->isForPrefetch(op);
+    return llvm::cast<TileUsageAnalysis>(analysis).isForPrefetch(op);
   }
 
+  template <typename = typename std::enable_if<
+                std::is_same_v<AnalysisT, TileUsageAnalysis>>>
   bool isForLoadAndPrefetch(imex::xetile::InitTileOp op) const {
-    return typeConverter.usageAnalysis->isForLoadAndPrefetch(op);
+    return llvm::cast<TileUsageAnalysis>(analysis).isForLoadAndPrefetch(op);
   }
 
+  template <typename = typename std::enable_if<
+                std::is_same_v<AnalysisT, TileUsageAnalysis>>>
   bool isForLoadAndStore(imex::xetile::InitTileOp op) const {
-    return typeConverter.usageAnalysis->isForLoadAndStore(op);
-  }
-
-  imex::XeTypeConverter &getTypeConverter() const { return typeConverter; }
-
-  template <typename ConverterTy>
-  std::enable_if_t<std::is_base_of<mlir::TypeConverter, ConverterTy>::value,
-                   ConverterTy &>
-  getTypeConverter() const {
-    return static_cast<ConverterTy &>(typeConverter);
+    return llvm::cast<TileUsageAnalysis>(analysis).isForLoadAndStore(op);
   }
-
-protected:
-  imex::XeTypeConverter &typeConverter;
 };
-
 } // namespace imex
 
 #endif

lib/Conversion/XeTileToXeGPU/ArithOpConversion.cpp

@@ -81,8 +81,10 @@ bool isLegalArithOp(mlir::Operation *op) {
 }
 
 void populateArithOpConversionPatterns(imex::XeGPUTypeConverter &converter,
-                                       mlir::RewritePatternSet &patterns) {
-  patterns.add<SgArithConstantOpPattern>(patterns.getContext(), converter);
+                                       mlir::RewritePatternSet &patterns,
+                                       TileUsageAnalysis &analysis) {
+  patterns.add<SgArithConstantOpPattern>(patterns.getContext(), converter,
+                                         analysis);
 }
 
 } // namespace imex