Update PropagateLayout pass (#837)

- added support for ExtractOp
- rename it to apply-vnni-transform

Author: chencha3

include/imex/Transforms/Passes.h

@@ -31,6 +31,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> createVnniTransformationPass();
 
 #define GEN_PASS_DECL
 #include "imex/Transforms/Passes.h.inc"

include/imex/Transforms/Passes.td

@@ -146,5 +146,14 @@ def RemoveSingleElemVector : Pass<"imex-remove-single-elem-vector"> {
   ];
 }
 
+def VnniTransformation : Pass<"imex-xegpu-apply-vnni-transformation"> {
+  let summary = "apply vnni transformation for to B operand of dpas instructions if necessary.";
+  let constructor = "imex::createVnniTransformationPass()";
+
+  let dependentDialects = [
+    "::mlir::vector::VectorDialect"
+  ];
+}
+
 
 #endif // _IMEX_TRANSFORMS_PASSES_TD_INCLUDED_

lib/Conversion/XeGPUToVC/XeGPUToVC.cpp

@@ -140,28 +140,27 @@ static Value adjustBasePointer(ConversionPatternRewriter &rewriter,
   auto strides = mlir::getStridesAndOffset(memType).first;
   int64_t i = memType.getRank() - 1;
 
-  auto computeBase =
-      [&](Value base) {
-        for (; i >= 0; --i) {
-          unsigned stride =
-              strides[i] * memType.getElementType().getIntOrFloatBitWidth() / 8;
-          auto factor = createIndexConstant(stride);
-          auto offset = offsets.pop_back_val();
-          Value offsetVal;
-
-          if (offset.is<Value>()) {
-            offsetVal = offset.get<Value>();
-          } else {
-            offsetVal = createIndexConstant(
-                llvm::cast<IntegerAttr>(offset.get<Attribute>()).getInt());
-          }
-          auto linearOffset =
-              rewriter.create<arith::MulIOp>(loc, offsetVal, factor);
-          base = rewriter.create<arith::AddIOp>(loc, base, linearOffset);
-        }
+  auto computeBase = [&](Value base) {
+    for (; i >= 0; --i) {
+      unsigned stride =
+          strides[i] * memType.getElementType().getIntOrFloatBitWidth() / 8;
+      auto factor = createIndexConstant(stride);
+      auto offset = offsets.pop_back_val();
+      Value offsetVal;
+
+      if (offset.is<Value>()) {
+        offsetVal = offset.get<Value>();
+      } else {
+        offsetVal = createIndexConstant(
+            llvm::cast<IntegerAttr>(offset.get<Attribute>()).getInt());
+      }
+      auto linearOffset =
+          rewriter.create<arith::MulIOp>(loc, offsetVal, factor);
+      base = rewriter.create<arith::AddIOp>(loc, base, linearOffset);
+    }
 
-        return base;
-      };
+    return base;
+  };
 
   if (tileRank == 2 && memType.getRank() > 2) {
     // base address of plane for 2d: base addr of memref + offsets (starting

lib/Dialect/XeTile/IR/XeTileDialect.cpp

@@ -120,8 +120,10 @@ mlir::LogicalResult XeTileAttr::verify(
 
   if (order != mlir::DenseI32ArrayAttr() && order.size() != 2)
     emitError() << "expect integer array of size 2 for order";
-  if (inner_blocks != mlir::DenseI64ArrayAttr() && (inner_blocks.size() > 0 && inner_blocks.size() != 2))
-    emitError() << "expect integer array of size 2 for non empty inner_blocks attribute";
+  if (inner_blocks != mlir::DenseI64ArrayAttr() &&
+      (inner_blocks.size() > 0 && inner_blocks.size() != 2))
+    emitError() << "expect integer array of size 2 for non empty inner_blocks "
+                   "attribute";
   return mlir::success();
 }
 

lib/Transforms/CMakeLists.txt

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