Optimization pass for fusing XeGPU load_nd with vector transpose. (#845)

Author: charithaintc

include/imex/Transforms/Passes.h

@@ -15,6 +15,7 @@
 #define IMEX_TRANSFORMS_PASSES_H_
 
 #include "mlir/Pass/Pass.h"
+#include <memory>
 
 namespace imex {
 //===----------------------------------------------------------------------===//
@@ -31,6 +32,7 @@ std::unique_ptr<mlir::Pass> createRemoveTemporariesPass();
 std::unique_ptr<mlir::Pass> createVectorLinearizePass();
 std::unique_ptr<mlir::Pass> createPropagatePackedLayoutPass();
 std::unique_ptr<mlir::Pass> createRemoveSingleElemVectorPass();
+std::unique_ptr<mlir::Pass> createOptimizeTransposePass();
 std::unique_ptr<mlir::Pass> createVnniTransformationPass();
 
 #define GEN_PASS_DECL

include/imex/Transforms/Passes.td

@@ -155,5 +155,18 @@ def VnniTransformation : Pass<"imex-xegpu-apply-vnni-transformation"> {
   ];
 }
 
+def OptimizeTranspose : Pass<"imex-xegpu-optimize-transpose"> {
+  let summary = "Eliminate in-register vector transpose by fusing with load.";
+  let constructor = "imex::createOptimizeTransposePass()";
+  let dependentDialects = [
+    "::mlir::xegpu::XeGPUDialect"
+  ];
+  let options = [
+     Option<"device", "device", "std::string",
+            /*default=*/"\"pvc\"",
+            "gpu platform architecture where these ops are running">
+ ];
+}
+
 
 #endif // _IMEX_TRANSFORMS_PASSES_TD_INCLUDED_

lib/Transforms/CMakeLists.txt

@@ -11,6 +11,7 @@ add_mlir_library(IMEXTransforms
   SetSPIRVCapabilities.cpp
   VectorLinearize.cpp
   VnniTransformation.cpp
+  OptimizeTranspose.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${PROJECT_SOURCE_DIR}/imex/Transforms