[Codegen][GPU] Use I64ArrayAttr for tile sizes for simpler printing (#…

…18575) The extra `: index` type qualifiers on the tile size attributes from using an index array attribute was just distracting. Use I64ArrayAttr instead which has simpler printing.
iree-org · Sep 22, 2024 · 5a6bd8d · 5a6bd8d
1 parent 9ee061d
commit 5a6bd8d
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 20 deletions.
diff --git a/compiler/src/iree/compiler/Codegen/Dialect/GPU/TargetUtils/ConfigUtils.cpp b/compiler/src/iree/compiler/Codegen/Dialect/GPU/TargetUtils/ConfigUtils.cpp
@@ -193,11 +193,11 @@ LogicalResult setMatmulLoweringConfig(IREE::GPU::TargetAttr target,
   SmallVector<NamedAttribute, 1> attrs;
   Builder b(context);
   attrs.emplace_back(StringAttr::get(context, "workgroup"),
-                     b.getIndexArrayAttr(workgroupTileSizes));
+                     b.getI64ArrayAttr(workgroupTileSizes));
   attrs.emplace_back(StringAttr::get(context, "reduction"),
-                     b.getIndexArrayAttr(reductionTileSizes));
+                     b.getI64ArrayAttr(reductionTileSizes));
   attrs.emplace_back(StringAttr::get(context, "subgroup"),
-                     b.getIndexArrayAttr(subgroupTileSizes));
+                     b.getI64ArrayAttr(subgroupTileSizes));
   attrs.emplace_back(StringAttr::get(context, "mma_kind"), mmaKind);
   auto configDict = DictionaryAttr::get(context, attrs);
   auto loweringConfig = IREE::GPU::LoweringConfigAttr::get(context, configDict);
@@ -434,10 +434,10 @@ LogicalResult setTileAndFuseLoweringConfig(IREE::GPU::TargetAttr target,
   SmallVector<NamedAttribute, 1> attrs;
   Builder b(context);
   attrs.emplace_back(StringAttr::get(context, "workgroup"),
-                     b.getIndexArrayAttr(workgroupTileSizes));
+                     b.getI64ArrayAttr(workgroupTileSizes));
 
   attrs.emplace_back(StringAttr::get(context, "thread"),
-                     b.getIndexArrayAttr(threadTileSizes));
+                     b.getI64ArrayAttr(threadTileSizes));
 
   // Heuristic value chosen to limit maximum vector sizes when tiling below.
   const unsigned maxVectorSize = 32;
@@ -467,7 +467,7 @@ LogicalResult setTileAndFuseLoweringConfig(IREE::GPU::TargetAttr target,
   }
   if (llvm::any_of(loopTileSizes, [](int64_t s) { return s != 0; })) {
     attrs.emplace_back(StringAttr::get(context, "reduction"),
-                       b.getIndexArrayAttr(loopTileSizes));
+                       b.getI64ArrayAttr(loopTileSizes));
   }
 
   auto configDict = DictionaryAttr::get(context, attrs);

diff --git a/compiler/src/iree/compiler/Codegen/LLVMGPU/test/ROCDL/config_tile_and_fuse.mlir b/compiler/src/iree/compiler/Codegen/LLVMGPU/test/ROCDL/config_tile_and_fuse.mlir
@@ -35,9 +35,9 @@ func.func @expanded_matmul_transpose_b(%lhs: tensor<2x64x2048xf16>, %rhs: tensor
 
 //       CHECK:   linalg.generic {{.*}}lowering_config = #iree_gpu.lowering_config
 //  CHECK-SAME:     mma_kind = #iree_gpu.mma_layout<MFMA_F32_16x16x16_F16>
-//  CHECK-SAME:     reduction = [0 : index, 0 : index, 0 : index, 0 : index, 4 : index]
-//  CHECK-SAME:     subgroup = [0 : index, 0 : index, 4 : index, 1 : index, 0 : index]
-//  CHECK-SAME:     workgroup = [1 : index, 1 : index, 64 : index, 64 : index, 0 : index]
+//  CHECK-SAME:     reduction = [0, 0, 0, 0, 4]
+//  CHECK-SAME:     subgroup = [0, 0, 4, 1, 0]
+//  CHECK-SAME:     workgroup = [1, 1, 64, 64, 0]
 
 // -----
 
@@ -59,9 +59,9 @@ func.func @mfma_matmul_1024x1024x1024(%lhs: tensor<1024x1024xf16>, %rhs: tensor<
 
 //       CHECK:   linalg.matmul {{.*}}lowering_config = #iree_gpu.lowering_config
 //  CHECK-SAME:     mma_kind = #iree_gpu.mma_layout<MFMA_F32_16x16x16_F16>
-//  CHECK-SAME:     reduction = [0 : index, 0 : index, 2 : index]
-//  CHECK-SAME:     subgroup = [4 : index, 4 : index, 0 : index]
-//  CHECK-SAME:     workgroup = [128 : index, 128 : index, 0 : index]
+//  CHECK-SAME:     reduction = [0, 0, 2]
+//  CHECK-SAME:     subgroup = [4, 4, 0]
+//  CHECK-SAME:     workgroup = [128, 128, 0]
 
 // -----
 
@@ -79,9 +79,9 @@ module {
 // CHECK-LABEL: func.func @conv_nhwc
 //  CHECK-SAME:   #iree_codegen.translation_info<LLVMGPUTileAndFuse workgroup_size = [64, 1, 1] subgroup_size = 64>
 //       CHECK:   linalg.conv_2d_nhwc_hwcf {{.*}} lowering_config = #iree_gpu.lowering_config
-//  CHECK-SAME:     reduction = [0 : index, 0 : index, 0 : index, 0 : index, 1 : index, 3 : index, 4 : index]
-//  CHECK-SAME:     thread = [1 : index, 1 : index, 1 : index, 1 : index, 0 : index, 0 : index, 0 : index]
-//  CHECK-SAME:     workgroup = [1 : index, 1 : index, 1 : index, 64 : index, 0 : index, 0 : index, 0 : index]
+//  CHECK-SAME:     reduction = [0, 0, 0, 0, 1, 3, 4]
+//  CHECK-SAME:     thread = [1, 1, 1, 1, 0, 0, 0]
+//  CHECK-SAME:     workgroup = [1, 1, 1, 64, 0, 0, 0]
 
 // -----
 
@@ -100,9 +100,9 @@ module {
 // CHECK-LABEL: func.func @matmul_dynamic_dim
 //  CHECK-SAME:   #iree_codegen.translation_info<LLVMGPUTileAndFuse workgroup_size = [64, 1, 1] subgroup_size = 64>
 //       CHECK:   linalg.matmul {{.*}}lowering_config = #iree_gpu.lowering_config
-//  CHECK-SAME:     reduction = [0 : index, 0 : index, 4 : index]
-//  CHECK-SAME:     thread = [1 : index, 1 : index, 0 : index]
-//  CHECK-SAME:     workgroup = [1 : index, 64 : index, 0 : index]
+//  CHECK-SAME:     reduction = [0, 0, 4]
+//  CHECK-SAME:     thread = [1, 1, 0]
+//  CHECK-SAME:     workgroup = [1, 64, 0]
 
 // -----
 
@@ -120,8 +120,8 @@ module {
 // CHECK-LABEL: func.func @elementwise_dynamic_dim
 //  CHECK-SAME:   #iree_codegen.translation_info<LLVMGPUTileAndFuse workgroup_size = [64, 1, 1] subgroup_size = 64>
 //       CHECK:   linalg.add {{.*}}lowering_config = #iree_gpu.lowering_config
-//  CHECK-SAME:     thread = [1 : index, 1 : index]
-//  CHECK-SAME:     workgroup = [1 : index, 64 : index]
+//  CHECK-SAME:     thread = [1, 1]
+//  CHECK-SAME:     workgroup = [1, 64]
 
 // -----