[tuner]:expose the helper function for level name

Signed-off-by: Bangtian Liu <[email protected]>
iree-org · Dec 7, 2024 · a080908 · a080908
1 parent 0edeb8e
commit a080908
Show file tree

Hide file tree

Showing 4 changed files with 15 additions and 6 deletions.
diff --git a/compiler/bindings/python/IREECompilerDialectsModule.cpp b/compiler/bindings/python/IREECompilerDialectsModule.cpp
@@ -364,7 +364,7 @@ PYBIND11_MODULE(_ireeCompilerDialects, m) {
 
             size_t len = mlirArrayAttrGetNumElements(workgroupAttr);
             std::vector<int64_t> workgroup(len);
-            for (size_t i = 0, e = len; i < e; ++i) {
+            for (size_t i = 0; i < len; ++i) {
               MlirAttribute attr = mlirArrayAttrGetElement(workgroupAttr, i);
               workgroup[i] = mlirIntegerAttrGetValueInt(attr);
             }
@@ -381,7 +381,7 @@ PYBIND11_MODULE(_ireeCompilerDialects, m) {
 
             size_t len = mlirArrayAttrGetNumElements(reductionAttr);
             std::vector<int64_t> reduction(len);
-            for (size_t i = 0, e = len; i < e; ++i) {
+            for (size_t i = 0; i < len; ++i) {
               MlirAttribute attr = mlirArrayAttrGetElement(reductionAttr, i);
               reduction[i] = mlirIntegerAttrGetValueInt(attr);
             }

diff --git a/compiler/src/iree/compiler/API/Internal/IREEGPUDialectCAPI.cpp b/compiler/src/iree/compiler/API/Internal/IREEGPUDialectCAPI.cpp
@@ -223,12 +223,17 @@ ireeGPUTileSizes ireeGPULoweringConfigAttrGetTileSizes(MlirAttribute attr) {
           unwrap(attr))
           .getAttributes();
 
-  constexpr mlir::StringLiteral workgroupName = "workgroup";
+  llvm::StringRef workgroupName =
+      mlir::iree_compiler::IREE::GPU::getTilingLevelName(
+          mlir::iree_compiler::IREE::GPU::TilingLevel::Workgroup);
+
   if (auto workgroupArray = dict.getAs<mlir::ArrayAttr>(workgroupName)) {
     tilesizes.workgroupAttr = wrap(workgroupArray);
   }
 
-  constexpr mlir::StringLiteral reductionName = "reduction";
+  llvm::StringRef reductionName =
+      mlir::iree_compiler::IREE::GPU::getTilingLevelName(
+          mlir::iree_compiler::IREE::GPU::TilingLevel::Reduction);
   if (auto reductionArray = dict.getAs<mlir::ArrayAttr>(reductionName)) {
     tilesizes.reductionAttr = wrap(reductionArray);
   }

diff --git a/compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUAttrs.cpp b/compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUAttrs.cpp
@@ -1104,7 +1104,7 @@ constexpr StringLiteral kThreadLevelName = "thread";
 constexpr StringLiteral kSubgroupLevelName = "subgroup";
 constexpr StringLiteral kLaneLevelName = "lane";
 
-static StringRef getTilingLevelName(GPU::TilingLevel level) {
+StringRef getTilingLevelName(GPU::TilingLevel level) {
   switch (level) {
   case GPU::TilingLevel::Workgroup:
     return kWorkgroupLevelName;
@@ -1120,7 +1120,7 @@ static StringRef getTilingLevelName(GPU::TilingLevel level) {
     return kLaneLevelName;
   }
   assert(false && "Unknown tiling level");
-  return StringAttr();
+  return StringRef();
 }
 
 static SmallVector<int64_t> getIntegerVector(ArrayAttr array) {

diff --git a/compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUAttrs.h b/compiler/src/iree/compiler/Codegen/Dialect/GPU/IR/IREEGPUAttrs.h
@@ -68,6 +68,10 @@ MMASingleSubgroupLayout getSingleSubgroupLayout(VirtualMMAIntrinsic intrinsic,
 MMASingleSubgroupLayout getSingleSubgroupLayout(MmaInterfaceAttr mmaKind,
                                                 MMAFragment fragment);
 
+/// Returns the name of the tilling `level`, as used in the `lowering_config`
+/// attribute.
+StringRef getTilingLevelName(GPU::TilingLevel level);
+
 } // namespace mlir::iree_compiler::IREE::GPU
 
 // clang-format off