add option to print layout results

Author: charithaintc

mlir/include/mlir/Dialect/XeGPU/Transforms/Passes.td

@@ -29,7 +29,7 @@ def XeGPUSubgroupDistribute : Pass<"xegpu-subgroup-distribute"> {
                            "vector::VectorDialect"];
 }
 
-def XeGPULayoutPropagate : Pass<"xegpu-propagate-layout"> {
+def XeGPUPropagateLayout : Pass<"xegpu-propagate-layout"> {
   let summary = "Propagate and assign XeGPU layout information";
   let description = [{
     This pass propagates the XeGPU layout information accross ops. Starting
@@ -40,6 +40,10 @@ def XeGPULayoutPropagate : Pass<"xegpu-propagate-layout"> {
   }];
   let dependentDialects = ["memref::MemRefDialect", "xegpu::XeGPUDialect",
                            "vector::VectorDialect"];
+  let options = [Option<
+    "printOnly", "print-analysis-only", "bool",
+    /*default=*/"false",
+    "Print the result of layout propagation analysis and exit.">];
 }
 
 def XeGPUWgToSgDistribute : Pass<"xegpu-wg-to-sg-distribute"> {

mlir/lib/Dialect/XeGPU/Transforms/CMakeLists.txt

@@ -4,7 +4,7 @@ add_mlir_dialect_library(MLIRXeGPUTransforms
   XeGPUSubgroupDistribute.cpp
   XeGPUUnroll.cpp
   XeGPUWgToSgDistribute.cpp
-  XeGPULayoutPropagate.cpp
+  XeGPUPropagateLayout.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/XeGPU

mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutPropagate.cpp renamed to mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp

@@ -1,4 +1,4 @@
-//===- XeGPULayoutPropagate.cpp - XeGPU Layout Propagation ------*- C++ -*-===//
+//===- XeGPUPropagateLayout.cpp - XeGPU Layout Propagation ------*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
@@ -38,7 +38,7 @@
 
 namespace mlir {
 namespace xegpu {
-#define GEN_PASS_DEF_XEGPULAYOUTPROPAGATE
+#define GEN_PASS_DEF_XEGPUPROPAGATELAYOUT
 #include "mlir/Dialect/XeGPU/Transforms/Passes.h.inc"
 } // namespace xegpu
 } // namespace mlir
@@ -622,8 +622,7 @@ LayoutInfo RunLayoutInfoPropagation::getLayoutInfo(Value val) {
 }
 
 // Print the analysis result for debugging purposes.
-[[maybe_unused]] void
-RunLayoutInfoPropagation::printAnalysisResult(llvm::raw_ostream &os) {
+void RunLayoutInfoPropagation::printAnalysisResult(llvm::raw_ostream &os) {
   auto printFunctionResult = [&](FunctionOpInterface funcOp) {
     os << "function: " << funcOp.getName() << ":\n";
     // Function arguments
@@ -828,15 +827,25 @@ static LogicalResult updateFunctionOpInterface(mlir::OpBuilder &builder,
 }
 
 namespace {
-struct XeGPULayoutPropagatePass final
-    : public xegpu::impl::XeGPULayoutPropagateBase<XeGPULayoutPropagatePass> {
+struct XeGPUPropagateLayoutPass final
+    : public xegpu::impl::XeGPUPropagateLayoutBase<XeGPUPropagateLayoutPass> {
+  XeGPUPropagateLayoutPass() = default;
+  XeGPUPropagateLayoutPass(const XeGPUPropagateLayoutPass &other) = default;
+  XeGPUPropagateLayoutPass(xegpu::XeGPUPropagateLayoutOptions options)
+      : XeGPUPropagateLayoutBase(options) {}
   void runOnOperation() override;
 };
 
 } // namespace
 
-void XeGPULayoutPropagatePass::runOnOperation() {
+void XeGPUPropagateLayoutPass::runOnOperation() {
   auto &analysis = getAnalysis<RunLayoutInfoPropagation>();
+  // Print the analysis result and exit. (for debugging purposes)
+  if (printOnly) {
+    auto &os = llvm::outs();
+    analysis.printAnalysisResult(os);
+    return;
+  }
   // Helper to convert LayoutInfo to xegpu::LayoutAttr.
   auto getXeGPULayoutForValue = [&](Value val) -> xegpu::LayoutAttr {
     LayoutInfo layout = analysis.getLayoutInfo(val);