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

[Adam Siemieniuk]

================
@@ -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(VectorType originalType,
+                                               LayoutAttr layout);
+
+// Helper to get the distributed vector type for a given vector type according
+// to a given LayoutAttr.
----------------
adam-smnk wrote:

nit: I think the descriptions are reversed as this one matches the other function

https://github.com/llvm/llvm-project/pull/135271