[mlir][xegpu] SIMT distribution patterns for XeGPU CreateNdTdesc, LoadNd, StoreNd and Dpas Ops. (llvm#135271)

This PR adds the SIMT distribution patterns for create_nd_tdesc, load_nd, store_nd and dpas XeGPU ops.

Author: charithaintc

mlir/include/mlir/Dialect/XeGPU/IR/XeGPUTypes.td

@@ -189,11 +189,6 @@ def XeGPU_TensorDesc: XeGPUTypeDef<"TensorDesc", "tensor_desc",
         return scatter_attr.getChunkSize().getInt();
       return 1;
     }
-
-    // This returns a vector type that represents the fragment of data owned by
-    // a work item in SIMT mode if this tensor descriptor is used in a XeGPU
-    // load/store operation.
-    FailureOr<VectorType> getDistributedVectorType();
   }];
 
   let hasCustomAssemblyFormat = true;

mlir/include/mlir/Dialect/XeGPU/Transforms/Transforms.h

@@ -16,6 +16,8 @@ namespace xegpu {
 
 /// Appends patterns for folding aliasing ops into XeGPU ops into `patterns`.
 void populateXeGPUFoldAliasOpsPatterns(RewritePatternSet &patterns);
+/// Appends patterns for XeGPU SIMT distribution into `patterns`.
+void populateXeGPUSubgroupDistributePatterns(RewritePatternSet &patterns);
 
 } // namespace xegpu
 } // namespace mlir

mlir/include/mlir/Dialect/XeGPU/Utils/XeGPUUtils.h

@@ -0,0 +1,57 @@
+//===- XeGPUUtils.h - Vector Utilities --------------------------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_DIALECT_XEGPU_UTILS_XEGPUUTILS_H_
+#define MLIR_DIALECT_XEGPU_UTILS_XEGPUUTILS_H_
+
+#include "mlir/IR/BuiltinTypes.h"
+namespace mlir {
+
+class VectorType;
+namespace xegpu {
+class LayoutAttr;
+class TensorDescType;
+} // namespace xegpu
+
+namespace xegpu {
+
+/// If tensor descriptor has a layout attribute it is used in SIMT mode.
+/// In this mode, the distributed vector shape is determined as follows:
+/// Definitions:
+///        lane_data_size = lane_data[0] × lane_data[1]
+///        subgroup_size = lane_layout[0] × lane_layout[1]
+///        distribution_unit_size = subgroup_size × lane_data_size
+///
+/// Case 1: Regular loads/stores.
+/// The following conditions must be met:
+///        * tensor_desc[0] == lane_layout[0]
+/// Distributed vector is a 1D vector with shape:
+///        [chunk_size]
+///
+/// Case 2: Block loads/stores
+/// Additional definitions:
+///        tensor_size = tensor_desc[0] * .. * tensor_desc[r-1] * array_length
+///        n_distribution_units = tensor_size / distribution_unit_size
+///        fragment_size = n_distribution_units * lane_data_size
+/// Given above definitions, the following conditions must be met:
+///        * tensor_desc[0] % (lane_layout[0] × lane_data[0]) == 0
+///        * tensor_desc[1] % (lane_layout[1] × lane_data[1]) == 0
+/// Distributed vector is a 1D vector with shape:
+///        [fragment_size]
+FailureOr<VectorType> getDistributedVectorType(xegpu::TensorDescType tdescTy);
+
+/// Helper to get the distributed vector type for a given vector type according
+/// to a given LayoutAttr.
+FailureOr<VectorType> getDistributedVectorType(VectorType originalType,
+                                               LayoutAttr layout);
+
+} // namespace xegpu
+
+} // namespace mlir
+
+#endif // MLIR_DIALECT_XEGPU_UTILS_XEGPUUTILS_H_

mlir/lib/Dialect/XeGPU/CMakeLists.txt

@@ -1,2 +1,3 @@
 add_subdirectory(IR)
 add_subdirectory(Transforms)
+add_subdirectory(Utils)

mlir/lib/Dialect/XeGPU/IR/XeGPUDialect.cpp

@@ -376,74 +376,6 @@ LogicalResult TensorDescType::verify(
   return success();
 }
 
-// If tensor descriptor has a layout attribute it is used in SIMT mode.
-// In this mode, the distributed vector shape is determined as follows:
-// Definitions:
-//        lane_data_size = lane_data[0] × lane_data[1]
-//        subgroup_size = lane_layout[0] × lane_layout[1]
-//        distribution_unit_size = subgroup_size × lane_data_size
-// ---------------------------------------------------------------------
-// Case 1: Regular loads/stores.
-// ---------------------------------------------------------------------
-// The following conditions must be met:
-//        * tensor_desc[0] == lane_layout[0]
-// Distributed vector is a 1D vector with shape:
-//        [chunk_size]
-// ---------------------------------------------------------------------
-// Case 2: Block loads/stores
-// ---------------------------------------------------------------------
-// Additional definitions:
-//        tensor_size = tensor_desc[0] * .. * tensor_desc[r-1] * array_length
-//        n_distribution_units = tensor_size / distribution_unit_size
-//        fragment_size = n_distribution_units * lane_data_size
-// Given above definitions, the following conditions must be met:
-//        * tensor_desc[0] % (lane_layout[0] × lane_data[0]) == 0
-//        * tensor_desc[1] % (lane_layout[1] × lane_data[1]) == 0
-// Distributed vector is a 1D vector with shape:
-//        [fragment_size]
-FailureOr<VectorType> TensorDescType::getDistributedVectorType() {
-  auto layout = llvm::dyn_cast_if_present<LayoutAttr>(getLayout());
-  // It only works for subgroup level layout, which only has lane_layout
-  // and lane_data, and is to distribute a SIMD code into SIMT code.
-  if (!layout || !layout.isSgLayout())
-    return failure();
-
-  SmallVector<int64_t> laneData(layout.getLaneData().asArrayRef());
-  SmallVector<int64_t> laneLayout(layout.getLaneLayout().asArrayRef());
-  auto tdescShape = getShape();
-
-  // compute sgSize by multiply elements of laneLayout
-  // e.g. for 2D layout, sgSize = laneLayout[0] * laneLayout[1]
-  // e.g. for 1D layout, sgSize = laneLayout[0]
-  auto sgSize = std::accumulate(laneLayout.begin(), laneLayout.end(), 1,
-                                std::multiplies<int64_t>());
-
-  // Case 1: regular loads/stores
-  auto scatterAttr = getEncodingAsScatterTensorDescAttr();
-  if (scatterAttr) {
-    auto chunkSize = scatterAttr.getChunkSize().getInt();
-    // Verify if the first dimension of the tensor descriptor shape is
-    // distributable.
-    assert(tdescShape[0] == laneLayout[0] &&
-           "tensor descriptor shape is not distributable");
-    return VectorType::get({chunkSize}, getElementType());
-  }
-
-  // Case 2: block loads/stores
-  // Check if the tensor descriptor shape is distributable.
-  int64_t tensorSize = 1;
-  for (auto [tdescDim, laneDim, laneDataDim] :
-       llvm::zip_equal(tdescShape, laneLayout, laneData)) {
-    assert((tdescDim % (laneDim * laneDataDim) == 0) &&
-           "tensor descriptor shape is not distributable");
-    tensorSize *= tdescDim;
-  }
-  // tensorSize must be adjusted for array_length.
-  tensorSize *= getArrayLength();
-
-  return VectorType::get({tensorSize / sgSize}, getElementType());
-}
-
 } // namespace xegpu
 } // namespace mlir
 

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

@@ -16,4 +16,7 @@ add_mlir_dialect_library(MLIRXeGPUTransforms
   MLIRPass
   MLIRTransforms
   MLIRGPUDialect
+  MLIRXeGPUUtils
+  MLIRGPUUtils
+  MLIRVectorTransforms
 )