[Mlir-commits] [mlir] [uArch][XeGPU] Add XeGPU uArch definition. (PR #153706)

[Artem Kroviakov]

================
@@ -0,0 +1,172 @@
+#include "mlir/Dialect/XeGPU/uArch/IntelGpuXe2.h"
+#include "mlir/IR/BuiltinTypes.h"
+#include "llvm/Support/DebugLog.h"
+#include <algorithm>
+#include <vector>
+
+#define DEBUG_TYPE "xegpu-uarch"
+
+using namespace mlir::xegpu::uArch;
+using namespace mlir::xegpu::uArch::Xe2Plus;
+
+namespace mlir {
+namespace xegpu {
+namespace uArch {
+namespace Xe2Plus {
+
+std::vector<std::pair<uint32_t, uint32_t>>
+DPASInstruction::getSupportedShapes(mlir::Type dataType,
+                                    MMAOpndEnum matrixType) {
+  auto combineVectors = [](const std::vector<uint32_t> &a,
+                           const std::vector<uint32_t> &b)
+      -> std::vector<std::pair<uint32_t, uint32_t>> {
+    std::vector<std::pair<uint32_t, uint32_t>> result;
+    for (unsigned x : a) {
+      for (unsigned y : b) {
+        result.emplace_back(x, y);
+      }
+    }
+    return result;
+  };
+
+  auto M = getSupportedM(dataType);
+  auto K = getSupportedK(dataType);
+  auto N = getSupportedN(dataType);
+  std::vector<std::pair<unsigned, unsigned>> resultMatrix;
+
+  switch (matrixType) {
+  case MMAOpndEnum::MatrixA:
+    resultMatrix = combineVectors(M, K);
+    break;
+  case MMAOpndEnum::MatrixB:
+    resultMatrix = combineVectors(K, N);
+    break;
+  case MMAOpndEnum::MatrixC:
+    resultMatrix = combineVectors(M, N);
+    break;
+  case MMAOpndEnum::MatrixD:
+    resultMatrix = combineVectors(M, N);
+    break;
+  }
+  return resultMatrix;
+}
+
+std::vector<mlir::Type>
+DPASInstruction::getSupportedTypes(MLIRContext &context,
----------------
akroviakov wrote:

 What is the motivating use case behind this method?

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