From 0edeb8e2e3232f85d6ca8d7201bada1f023d974c Mon Sep 17 00:00:00 2001
From: Bangtian Liu <liubangtian@gmail.com>
Date: Thu, 5 Dec 2024 23:12:47 -0800
Subject: [PATCH] [tuner]: use attribute for workgroup and reduction array

Signed-off-by: Bangtian Liu <liubangtian@gmail.com>
---
 .../c/iree/compiler/dialects/iree_gpu.h       |  6 +--
 .../python/IREECompilerDialectsModule.cpp     | 50 +++++++++++++------
 .../API/Internal/IREEGPUDialectCAPI.cpp       | 25 ++++------
 3 files changed, 47 insertions(+), 34 deletions(-)

diff --git a/compiler/bindings/c/iree/compiler/dialects/iree_gpu.h b/compiler/bindings/c/iree/compiler/dialects/iree_gpu.h
index 44a018f8a8e6..4b7836128142 100644
--- a/compiler/bindings/c/iree/compiler/dialects/iree_gpu.h
+++ b/compiler/bindings/c/iree/compiler/dialects/iree_gpu.h
@@ -94,10 +94,8 @@ MLIR_CAPI_EXPORTED MlirAttribute
 ireeGPULoweringConfigAttrGetAttributes(MlirAttribute attr);
 
 struct ireeGPUTileSizes {
-  const int64_t *workgroupTileSizes;
-  size_t numWorkgroupTileSizes;
-  const int64_t *reductionTileSizes;
-  size_t numReductionTileSizes;
+  MlirAttribute workgroupAttr;
+  MlirAttribute reductionAttr;
 };
 
 MLIR_CAPI_EXPORTED ireeGPUTileSizes
diff --git a/compiler/bindings/python/IREECompilerDialectsModule.cpp b/compiler/bindings/python/IREECompilerDialectsModule.cpp
index 89c8d413adfc..b764558760d0 100644
--- a/compiler/bindings/python/IREECompilerDialectsModule.cpp
+++ b/compiler/bindings/python/IREECompilerDialectsModule.cpp
@@ -353,22 +353,40 @@ PYBIND11_MODULE(_ireeCompilerDialects, m) {
           "Gets an #iree_gpu.lowering_config from parameters.")
       .def_property_readonly("attributes",
                              ireeGPULoweringConfigAttrGetAttributes)
-      .def_property_readonly("workgroup_tile_sizes",
-                             [](MlirAttribute self) -> std::vector<int64_t> {
-                               auto tilesizes =
-                                   ireeGPULoweringConfigAttrGetTileSizes(self);
-                               return {tilesizes.workgroupTileSizes,
-                                       tilesizes.workgroupTileSizes +
-                                           tilesizes.numWorkgroupTileSizes};
-                             })
-      .def_property_readonly("reduction_tile_sizes",
-                             [](MlirAttribute self) -> std::vector<int64_t> {
-                               auto tilesizes =
-                                   ireeGPULoweringConfigAttrGetTileSizes(self);
-                               return {tilesizes.reductionTileSizes,
-                                       tilesizes.reductionTileSizes +
-                                           tilesizes.numReductionTileSizes};
-                             })
+      .def_property_readonly(
+          "workgroup_tile_sizes",
+          [](MlirAttribute self) -> std::vector<int64_t> {
+            auto tilesizes = ireeGPULoweringConfigAttrGetTileSizes(self);
+            MlirAttribute workgroupAttr = tilesizes.workgroupAttr;
+            if (mlirAttributeIsNull(workgroupAttr)) {
+              return {};
+            }
+
+            size_t len = mlirArrayAttrGetNumElements(workgroupAttr);
+            std::vector<int64_t> workgroup(len);
+            for (size_t i = 0, e = len; i < e; ++i) {
+              MlirAttribute attr = mlirArrayAttrGetElement(workgroupAttr, i);
+              workgroup[i] = mlirIntegerAttrGetValueInt(attr);
+            }
+            return workgroup;
+          })
+      .def_property_readonly(
+          "reduction_tile_sizes",
+          [](MlirAttribute self) -> std::vector<int64_t> {
+            auto tilesizes = ireeGPULoweringConfigAttrGetTileSizes(self);
+            MlirAttribute reductionAttr = tilesizes.reductionAttr;
+            if (mlirAttributeIsNull(reductionAttr)) {
+              return {};
+            }
+
+            size_t len = mlirArrayAttrGetNumElements(reductionAttr);
+            std::vector<int64_t> reduction(len);
+            for (size_t i = 0, e = len; i < e; ++i) {
+              MlirAttribute attr = mlirArrayAttrGetElement(reductionAttr, i);
+              reduction[i] = mlirIntegerAttrGetValueInt(attr);
+            }
+            return reduction;
+          })
       .def_property_readonly(
           "subgroup_count_mn",
           [](MlirAttribute self) -> py::tuple {
diff --git a/compiler/src/iree/compiler/API/Internal/IREEGPUDialectCAPI.cpp b/compiler/src/iree/compiler/API/Internal/IREEGPUDialectCAPI.cpp
index 1ec55ba94590..58434ee337ff 100644
--- a/compiler/src/iree/compiler/API/Internal/IREEGPUDialectCAPI.cpp
+++ b/compiler/src/iree/compiler/API/Internal/IREEGPUDialectCAPI.cpp
@@ -218,23 +218,20 @@ MlirAttribute ireeGPULoweringConfigAttrGetAttributes(MlirAttribute attr) {
 ireeGPUTileSizes ireeGPULoweringConfigAttrGetTileSizes(MlirAttribute attr) {
   assert(ireeAttributeIsAGPULoweringConfigAttr(attr));
   ireeGPUTileSizes tilesizes = {};
-  auto loweringConfigAttr =
+  mlir::DictionaryAttr dict =
       llvm::cast<mlir::iree_compiler::IREE::GPU::LoweringConfigAttr>(
-          unwrap(attr));
-
-  llvm::SmallVector<int64_t> workgroups =
-      loweringConfigAttr.getWorkgroupTileSizes();
-  tilesizes.workgroupTileSizes = workgroups.data();
-  tilesizes.numWorkgroupTileSizes = workgroups.size();
+          unwrap(attr))
+          .getAttributes();
 
-  llvm::SmallVector<int64_t> reductions =
-      loweringConfigAttr.getStaticTilingLevelSizes(
-          llvm::to_underlying(
-              mlir::iree_compiler::IREE::GPU::TilingLevel::Reduction),
-          nullptr);
-  tilesizes.reductionTileSizes = reductions.data();
-  tilesizes.numReductionTileSizes = reductions.size();
+  constexpr mlir::StringLiteral workgroupName = "workgroup";
+  if (auto workgroupArray = dict.getAs<mlir::ArrayAttr>(workgroupName)) {
+    tilesizes.workgroupAttr = wrap(workgroupArray);
+  }
 
+  constexpr mlir::StringLiteral reductionName = "reduction";
+  if (auto reductionArray = dict.getAs<mlir::ArrayAttr>(reductionName)) {
+    tilesizes.reductionAttr = wrap(reductionArray);
+  }
   return tilesizes;
 }