Add IntelDPASLayout to Gluon (#5273)

This PR fixes #5159

Author: HBN-MichalSzy

python/src/gluon_ir.cc

@@ -5,6 +5,7 @@
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/Types.h"
 #include "third_party/amd/include/Dialect/TritonAMDGPU/IR/Dialect.h"
+#include "third_party/intel/include/Dialect/TritonIntelGPU/IR/Dialect.h"
 #include "triton/Analysis/Utility.h"
 #include "triton/Dialect/Gluon/IR/Dialect.h"
 #include "triton/Dialect/TritonGPU/IR/Attributes.h"
@@ -102,12 +103,15 @@ struct GluonLayouts {
   py::handle AMDMFMALayout;
   py::handle AMDWMMALayout;
   py::handle PaddedSharedLayout;
+  py::handle IntelDPASLayout;
 
   GluonLayouts() {
     auto layouts =
         py::module::import("triton.experimental.gluon.language._layouts");
     auto amdLayouts =
         py::module::import("triton.experimental.gluon.language.amd._layouts");
+    auto intelLayouts =
+        py::module::import("triton.experimental.gluon.language.intel._layouts");
     AutoLayout = py::object(layouts.attr("AutoLayout")).release();
     BlockedLayout = py::object(layouts.attr("BlockedLayout")).release();
     SliceLayout = py::object(layouts.attr("SliceLayout")).release();
@@ -125,6 +129,8 @@ struct GluonLayouts {
     AMDWMMALayout = py::object(amdLayouts.attr("AMDWMMALayout")).release();
     PaddedSharedLayout =
         py::object(layouts.attr("PaddedSharedLayout")).release();
+    IntelDPASLayout =
+        py::object(intelLayouts.attr("IntelDPASLayout")).release();
 
     auto core = py::module::import("triton.language.core");
   }
@@ -247,6 +253,12 @@ py::object layoutToGluon(Attribute layout) {
     return layouts.PaddedSharedLayout(intervalPaddingPairs,
                                       ll.getBases().lookup(kOffset),
                                       ll.getBases().lookup(kBlock), shape);
+  } else if (auto intelDpas = dyn_cast<ttg::intel::DpasEncodingAttr>(layout)) {
+    return layouts.IntelDPASLayout(
+        intelDpas.getRepeatCount(), intelDpas.getSystolicDepth(),
+        intelDpas.getExecutionSize(), intelDpas.getOpsPerChannel(),
+        toStdVector(intelDpas.getWarpsPerCTA()),
+        toStdVector(intelDpas.getRepCluster()), intelDpas.getThreadsPerWarp());
   }
 
   throw py::value_error("Unhandled encoding encountered");
@@ -385,6 +397,17 @@ void init_gluon_ir(py::module &&m) {
              return ttg::AMDWmmaEncodingAttr::get(
                  ctx, version, transposed, warpsPerCta, ctaLayout, instrShape);
            })
+      .def("get_intel_dpas_layout",
+           [](GluonOpBuilder &self, unsigned repeatCount,
+              unsigned systolicDepth, unsigned executionSize,
+              unsigned opsPerChannel, std::vector<unsigned> &warpsPerCTA,
+              std::vector<unsigned> &repCluster,
+              unsigned threadsPerWarp) -> Attribute {
+             auto ctx = self.getContext();
+             return ttg::intel::DpasEncodingAttr::get(
+                 ctx, repeatCount, systolicDepth, executionSize, opsPerChannel,
+                 warpsPerCTA, repCluster, threadsPerWarp);
+           })
       .def("get_padded_shared_layout",
            [](GluonOpBuilder &self, std::vector<unsigned> &intervals,
               std::vector<unsigned> &paddings,

python/test/gluon/test_lowerings.py

@@ -125,6 +125,8 @@ def _reduce_layouts():
         ttgl.amd.AMDMFMALayout(version=4, instr_shape=[32, 32, 16], transposed=True, warps_per_cta=[1, 4]),
         ttgl.amd.AMDWMMALayout(version=1, transposed=True, warps_per_cta=[1, 4]),
         ttgl.amd.AMDWMMALayout(version=2, transposed=True, warps_per_cta=[1, 4]),
+        ttgl.intel.IntelDPASLayout(repeatCount=8, systolic_depth=8, execution_size=8, ops_per_chan=1,
+                                   warps_per_cta=[4, 1], rep_cluster=[1, 1], threads_per_warp=32),
         ttgl.DotOperandLayout(
             parent=ttgl.NVMMADistributedLayout(version=[2, 0], warps_per_cta=[2, 4], ctas_per_cga=[1, 1],
                                                cta_split_num=[1, 1], cta_order=[0, 1], instr_shape=[16, 8]),

python/triton/experimental/gluon/language/__init__.py

@@ -125,4 +125,5 @@
 
 from . import nvidia
 from . import amd
+from . import intel
 from . import extra

python/triton/experimental/gluon/language/intel/__init__.py

@@ -0,0 +1,3 @@
+from ._layouts import IntelDPASLayout
+
+__all__ = ["IntelDPASLayout"]

python/triton/experimental/gluon/language/intel/_layouts.py

@@ -0,0 +1,86 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import List, Optional
+from triton.language.core import _unwrap_if_constexpr
+
+from triton.experimental.gluon.language._layouts import DistributedLayout
+
+__all__ = [
+    "IntelDPASLayout",
+]
+
+
+@dataclass(frozen=True)
+class IntelDPASLayout(DistributedLayout):
+    """
+    Represents a layout for Intel DPAS (dot product accumulator) operations.
+
+    Args:
+        repeatCount (int): Number of repeats for the operation.
+        systolic_depth (int): Systolic array depth.
+        execution_size (int): Execution size.
+        ops_per_chan (int): Operations per channel.
+        warps_per_cta (List[int]): Warp layout in the block.
+        rep_cluster (List[int]): Cluster repetition configuration.
+        threads_per_warp (int): Number of threads per warp.
+    """
+
+    repeatCount: int
+    systolic_depth: int
+    execution_size: int
+    ops_per_chan: int
+    warps_per_cta: List[int]
+    rep_cluster: List[int]
+    threads_per_warp: int
+    cta_order: Optional[List[int]] = None
+
+    def __post_init__(self):
+        super().__setattr__("repeatCount", _unwrap_if_constexpr(self.repeatCount))
+        super().__setattr__("systolic_depth", _unwrap_if_constexpr(self.systolic_depth))
+        super().__setattr__("execution_size", _unwrap_if_constexpr(self.execution_size))
+        super().__setattr__("ops_per_chan", _unwrap_if_constexpr(self.ops_per_chan))
+        super().__setattr__("warps_per_cta", _unwrap_if_constexpr(self.warps_per_cta))
+        super().__setattr__("rep_cluster", _unwrap_if_constexpr(self.rep_cluster))
+        super().__setattr__("threads_per_warp", _unwrap_if_constexpr(self.threads_per_warp))
+        # Compute cta_order as reversed range of warps_per_cta length, if not provided
+        super().__setattr__("cta_order", list(reversed(range(len(self.warps_per_cta)))))
+
+        self.verify()
+
+    def _to_ir(self, builder):
+        # TODO: Replace with actual Intel DPAS IR builder method
+        return builder.get_intel_dpas_layout(
+            self.repeatCount,
+            self.systolic_depth,
+            self.execution_size,
+            self.ops_per_chan,
+            self.warps_per_cta,
+            self.rep_cluster,
+            self.threads_per_warp,
+        )
+
+    def mangle(self) -> str:
+
+        def stringify(x):
+            if x is None:
+                return ""
+            return "_".join(map(str, x))
+
+        return f"IntelDPAS_{self.repeatCount}_{self.systolic_depth}_{self.execution_size}_{self.ops_per_chan}_{stringify(self.warps_per_cta)}_{stringify(self.rep_cluster)}_{self.threads_per_warp}_IntelDPAS"
+
+    def verify(self):
+        # TODO Do we need verify?
+        return
+
+    def __hash__(self):
+        return hash((
+            self.repeatCount,
+            self.systolic_depth,
+            self.execution_size,
+            self.ops_per_chan,
+            tuple(self.warps_per_cta),
+            tuple(self.rep_cluster),
+            self.threads_per_warp,
+            tuple(self.cta_order),
+        ))