Fix lowering logic for low-precision types in XeTileToXeGPU (#817)

intel · Jul 30, 2024 · 4e86e50 · 4e86e50
1 parent a4d993d
commit 4e86e50
Showing 1 changed file with 6 additions and 1 deletion.
diff --git a/lib/Conversion/XeTileToXeGPU/XeTileOpConversion.cpp b/lib/Conversion/XeTileToXeGPU/XeTileOpConversion.cpp
@@ -610,6 +610,11 @@ struct SgLoadTileOpPattern : public XeOneToNConversion<xetile::LoadTileOp> {
     // TODO: move these two into architecture abstracture in future.
     const int SIMD_WIDTH_IN_BITS = 32;
     int factor = SIMD_WIDTH_IN_BITS / elemTy.getIntOrFloatBitWidth();
+    // TODO: use uArch for this?
+    auto isLowPrecision = [](unsigned int width) -> bool {
+      bool isPowerOf2 = (width & (width - 1)) == 0;
+      return isPowerOf2 & (width < 32) & (width > 1);
+    };
     if (isForDPASB(op) && factor > 1)
       vnniAttr = mlir::UnitAttr::get(ctx);
 
@@ -621,7 +626,7 @@ struct SgLoadTileOpPattern : public XeOneToNConversion<xetile::LoadTileOp> {
       auto elemWidth = elemTy.getIntOrFloatBitWidth();
       if (elemWidth == 32) {
         transposeAttr = rewriter.getDenseI64ArrayAttr({1, 0});
-      } else if (elemWidth == 16 && vnniAttr) {
+      } else if (isLowPrecision(elemWidth) && vnniAttr) {
         transposeAttr = rewriter.getDenseI64ArrayAttr({1, 0});
         transposeBitWidthAttr = rewriter.getI32IntegerAttr(32);
         vnniAttr = nullptr;