Update blocking pass to support slm (#838)

chencha3 · web-flow · commit ae9fe176b562 · 2024-08-21T08:38:12.000-05:00
Accesses (load/store) to SLM use different instructions and has different size constraints as compared to block load/store operations for global memory.
diff --git a/include/imex/Conversion/XeTileToXeGPU/XeTileToXeGPUConversion.h b/include/imex/Conversion/XeTileToXeGPU/XeTileToXeGPUConversion.h
@@ -159,12 +159,14 @@ class XeOneToNConversion : public XeConversionPattern<TileUsageAnalysis> {
     // (convertedTypes.size() == 1) we will reuse the current value. Otherwise,
     // it has one-to-n mapping, and the new value should be an
     // UnrealizedConversionCastOp.
-    for (auto &value : remappedValues) {
+    for (size_t i = 0; i < remappedValues.size(); i++) {
+      auto value = remappedValues[i];
       auto castOp = value.getDefiningOp<mlir::UnrealizedConversionCastOp>();
-      if (castOp && castOp.getInputs().size() > 1)
+      auto valueTy = value.getType();
+      if (castOp && valueTy == op->getOperand(i).getType())
         convertedValues.push_back(castOp.getInputs());
       else
-        convertedValues.push_back(value);
+        convertedValues.push_back(remappedValues[i]);
     }
 
     auto sourceOp = llvm::dyn_cast<SourceOp>(op);
diff --git a/include/imex/Dialect/XeTile/IR/XeTileOps.h b/include/imex/Dialect/XeTile/IR/XeTileOps.h
@@ -14,7 +14,6 @@
 
 #ifndef _XETILE_OPS_H_INCLUDED_
 #define _XETILE_OPS_H_INCLUDED_
-
 #include <mlir/Dialect/Vector/IR/VectorOps.h>
 #include <mlir/IR/BuiltinTypeInterfaces.h>
 #include <mlir/IR/BuiltinTypes.h>
diff --git a/include/imex/Utils/XeCommon.h b/include/imex/Utils/XeCommon.h
@@ -268,8 +268,11 @@ class PropagateAnalysis {
 
       auto *op = getDefineOrParentOp(value);
 
-      // stop when meet a function.
-      if (!op || llvm::isa<mlir::FunctionOpInterface>(op))
+      // stop when meet a function or ops, e.g., arith.truncf.
+      // since their source and results could have different bitwidth,
+      // in which case the block size cannot be propagated.
+      if (!op || llvm::isa<mlir::FunctionOpInterface>(op) ||
+          llvm::isa<mlir::CastOpInterface>(op))
         continue;
 
       OpAttrMap[value] = attr;
diff --git a/lib/Conversion/XeTileToXeGPU/XeTileOpConversion.cpp b/lib/Conversion/XeTileToXeGPU/XeTileOpConversion.cpp
@@ -637,7 +637,10 @@ struct SgLoadTileOpPattern : public XeOneToNConversion<xetile::LoadTileOp> {
       bool isPowerOf2 = (width & (width - 1)) == 0;
       return isPowerOf2 & (width < 32) & (width > 1);
     };
-    if (isForDPASB(op) && factor > 1)
+    // vnni can only be applied when the blockSZ[0] >= factor
+    // for shape, e.g., 1xN, vnni cannot be applied, since no
+    // vnni transform available)
+    if (isForDPASB(op) && factor > 1 && blockSZ[0] >= factor)
       vnniAttr = mlir::UnitAttr::get(ctx);
 
     mlir::DenseI64ArrayAttr transposeAttr;
@@ -661,6 +664,12 @@ struct SgLoadTileOpPattern : public XeOneToNConversion<xetile::LoadTileOp> {
           .notifyMatchFailure(op, "Unsupported order");
     }
 
+    // vnni and transpose are not available for SLM memory scope.
+    if (tileTy.getMemoryScopeAsInt() == 3) {
+      vnniAttr = nullptr;
+      transposeBitWidthAttr = nullptr;
+    }
+
     rewriter.setInsertionPoint(op);
     llvm::SmallVector<::mlir::Value> xegpuOps;
     for (auto src : sources) {
diff --git a/lib/Dialect/XeTile/Transforms/Blocking.cpp b/lib/Dialect/XeTile/Transforms/Blocking.cpp
@@ -14,7 +14,6 @@
 /// such that each pieces can be handled by a hardware instruction.
 ///
 //===----------------------------------------------------------------------===//
-
 #include <mlir/Conversion/LLVMCommon/TypeConverter.h>
 #include <mlir/Dialect/Arith/IR/Arith.h>
 #include <mlir/Dialect/Func/IR/FuncOps.h>
@@ -779,66 +778,83 @@ struct InitTileOpPattern
           op, "Skipped InitTileOp because the result tile is not rank 2.\n");
 
     auto innerBlocks = tileTy.getInnerBlocks();
+    auto memorySpace = op.getSourceMemorySpaceAsInt();
 
     // skip it if innerBlocks has been set by user or compiler.
     if (innerBlocks)
       return mlir::failure();
 
     auto elemTy = tileTy.getElementType();
     int elementSize = elemTy.getIntOrFloatBitWidth();
-    if (isForPrefetch(op)) {
-      innerBlocks = mlir::DenseI64ArrayAttr::get(
-          getContext(), getInnerBlockSizes<Prefetch>(
-                            op.getOperation(), elemTy, tileTy.getShape()[0],
-                            tileTy.getShape()[1], this->uArchInterface));
-    } else if (isForLoad(op)) {
-
-      // Set transpose and vnni
-      bool vnni = false;
-      bool transpose = false;
-
-      auto order = tileTy.getOrder();
-      if (order[0] == 0 && order[1] == 1)
-        transpose = true;
-
-      for (auto user : getEffectiveUsers(op)) {
-        if (auto loadTileOp = llvm::dyn_cast<xetile::LoadTileOp>(user)) {
-          if (isForDPASB(loadTileOp) && elementSize < 32) {
-            vnni = true;
-            break;
+
+    if (memorySpace == 3) {                    // for shared memory
+      const unsigned int lscConstraints = 512; // 512 bytes constraint by lsc
+      const unsigned int subgroupSize = 16;
+      auto shape = tileTy.getShape();
+      int64_t innerBlockSizes[2];
+      // prefer to use gather loads with 16 simd lanes
+      innerBlockSizes[0] = shape[0] % subgroupSize == 0 ? 16 : 1;
+      innerBlockSizes[1] =
+          (lscConstraints * 8) / (elementSize * innerBlockSizes[0]);
+      innerBlockSizes[1] =
+          std::min<int64_t>(innerBlockSizes[1], tileTy.getShape()[1]);
+      innerBlocks = mlir::DenseI64ArrayAttr::get(getContext(), innerBlockSizes);
+    } else { // for global memory
+      if (isForPrefetch(op)) {
+        innerBlocks = mlir::DenseI64ArrayAttr::get(
+            getContext(), getInnerBlockSizes<Prefetch>(
+                              op.getOperation(), elemTy, tileTy.getShape()[0],
+                              tileTy.getShape()[1], this->uArchInterface));
+      } else if (isForLoad(op)) {
+
+        // Set transpose and vnni
+        bool vnni = false;
+        bool transpose = false;
+
+        auto order = tileTy.getOrder();
+        if (order[0] == 0 && order[1] == 1)
+          transpose = true;
+
+        for (auto user : getEffectiveUsers(op)) {
+          if (auto loadTileOp = llvm::dyn_cast<xetile::LoadTileOp>(user)) {
+            if (isForDPASB(loadTileOp) && elementSize < 32) {
+              vnni = true;
+              break;
+            }
           }
         }
-      }
 
-      if (vnni && transpose && elementSize < 32) {
-        int factor = 32 / elementSize;
-        vnni = false;
-        llvm::SmallVector<int64_t, 2> innerBlock = getInnerBlockSizes<Load>(
-            op.getOperation(), mlir::FloatType::getF32(getContext()),
-            tileTy.getShape()[1], (tileTy.getShape()[0]) / factor,
-            this->uArchInterface, vnni, transpose);
-        std::swap(innerBlock[0], innerBlock[1]);
-        innerBlock[0] *= factor;
-        innerBlocks = mlir::DenseI64ArrayAttr::get(getContext(), innerBlock);
-
-      } else if (transpose && elementSize < 32) {
-        return rewriter.notifyMatchFailure(op, "Invalid transpose.");
-      } else {
+        if (vnni && transpose && elementSize < 32) {
+          int factor = 32 / elementSize;
+          vnni = false;
+          llvm::SmallVector<int64_t, 2> innerBlock = getInnerBlockSizes<Load>(
+              op.getOperation(), mlir::FloatType::getF32(getContext()),
+              tileTy.getShape()[1], (tileTy.getShape()[0]) / factor,
+              this->uArchInterface, vnni, transpose);
+          std::swap(innerBlock[0], innerBlock[1]);
+          innerBlock[0] *= factor;
+          innerBlocks = mlir::DenseI64ArrayAttr::get(getContext(), innerBlock);
+
+        } else if (transpose && elementSize < 32) {
+          return rewriter.notifyMatchFailure(op, "Invalid transpose.");
+        } else {
+          innerBlocks = mlir::DenseI64ArrayAttr::get(
+              getContext(),
+              getInnerBlockSizes<Load>(
+                  op.getOperation(), elemTy, tileTy.getShape()[0],
+                  tileTy.getShape()[1], this->uArchInterface, vnni, transpose));
+        }
+      } else if (isForStore(op)) {
         innerBlocks = mlir::DenseI64ArrayAttr::get(
-            getContext(),
-            getInnerBlockSizes<Load>(op.getOperation(), elemTy,
-                                     tileTy.getShape()[0], tileTy.getShape()[1],
-                                     this->uArchInterface, vnni, transpose));
+            getContext(), getInnerBlockSizes<Store>(
+                              op.getOperation(), elemTy, tileTy.getShape()[0],
+                              tileTy.getShape()[1], this->uArchInterface));
+      } else {
+        return rewriter.notifyMatchFailure(
+            op,
+            "The tile is used for multiple purpose. The init-duplicate pass "
+            "should be run first to resolve this issue.");
       }
-    } else if (isForStore(op)) {
-      innerBlocks = mlir::DenseI64ArrayAttr::get(
-          getContext(), getInnerBlockSizes<Store>(
-                            op.getOperation(), elemTy, tileTy.getShape()[0],
-                            tileTy.getShape()[1], this->uArchInterface));
-    } else {
-      return rewriter.notifyMatchFailure(
-          op, "The tile is used for multiple purpose. The init-duplicate pass "
-              "should be run first to resolve this issue.");
     }
 
     if (innerBlocks.empty()) {
diff --git a/lib/Utils/XeArch.cpp b/lib/Utils/XeArch.cpp
@@ -303,6 +303,10 @@ mlir::LogicalResult XeuArchInterface::isLegalLoad2dOp(mlir::Operation *op) {
   if (auto loadOp = llvm::dyn_cast<mlir::xegpu::LoadNdOp>(op)) {
     auto tdescTy = loadOp.getTensorDescType();
 
+    // TODO: need more thinking on SLM
+    if (tdescTy.getMemoryScope() == mlir::xegpu::MemoryScope::SLM)
+      return mlir::success();
+
     int elementSize = loadOp.getTensorDescType().getElementTypeBitWidth();
 
     LoadStore2DConfig loadParams;
@@ -342,6 +346,10 @@ mlir::LogicalResult XeuArchInterface::isLegalStore2dOp(mlir::Operation *op) {
     auto tdescTy = storeOp.getTensorDescType();
     int elementSize = tdescTy.getElementTypeBitWidth();
 
+    // TODO: need more thinking on SLM
+    if (tdescTy.getMemoryScope() == mlir::xegpu::MemoryScope::SLM)
+      return mlir::success();
+
     LoadStore2DConfig storeParams;
     bool vnni = false;
     bool transpose = false;
diff --git a/test/Conversion/XeTileToXeGPU/sg_gemm_1k_1k_1k_f16_f32_slm.mlir b/test/Conversion/XeTileToXeGPU/sg_gemm_1k_1k_1k_f16_f32_slm.mlir
diff --git a/test/Conversion/XeTileToXeGPU/sg_mixed_scf.mlir b/test/Conversion/XeTileToXeGPU/sg_mixed_scf.mlir
diff --git a/test/Dialect/XeTile/Transforms/sg_gemm_1k_1k_1k_f16_f32_slm.mlir b/test/Dialect/XeTile/Transforms/sg_gemm_1k_1k_1k_f16_f32_slm.mlir
