From e923065eaab5e2d4c5b34cb43d89b3c0fd09f9ca Mon Sep 17 00:00:00 2001 From: Quinn Dawkins Date: Sat, 21 Sep 2024 23:53:06 -0400 Subject: [PATCH] [Codegen][GPU] Use I64ArrayAttr for tile sizes for simpler printing 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. --- .../Dialect/GPU/TargetUtils/ConfigUtils.cpp | 12 ++++---- .../test/ROCDL/config_tile_and_fuse.mlir | 28 +++++++++---------- 2 files changed, 20 insertions(+), 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 index cee4a5a49cde..9b6180b7d4d8 100644 --- 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 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 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 index e7b0ba4ec68d..93736c704412 100644 --- 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 -// 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 -// 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 // 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 // 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 // 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] // -----