From 12650a383f1faf8f2800ecb5a7377fa2a57d1ed5 Mon Sep 17 00:00:00 2001 From: "Gusthinna Waduge, Charitha Saumya" Date: Tue, 7 Nov 2023 18:29:17 +0000 Subject: [PATCH] Updates to XeTile definition. - Remove init_coop_tile op and update tests - Some fixes to op assembly format --- include/imex/Dialect/XeTile/IR/XeTileOps.td | 50 +++------------------ test/Dialect/XeTile/IR/ops.mlir | 28 ++---------- 2 files changed, 11 insertions(+), 67 deletions(-) diff --git a/include/imex/Dialect/XeTile/IR/XeTileOps.td b/include/imex/Dialect/XeTile/IR/XeTileOps.td index 87c2d47b4..e106b60ac 100644 --- a/include/imex/Dialect/XeTile/IR/XeTileOps.td +++ b/include/imex/Dialect/XeTile/IR/XeTileOps.td @@ -206,40 +206,6 @@ def XeTile_InitTileOp : XeTile_Op<"init_tile", [Pure, AttrSizedOperandSegments]> } -def XeTile_InitCoopTileOp : XeTile_Op<"init_coop_tile", [Pure]> { - let summary = "splits XeTile into multiple smaller tiles for cooprative work among subgroups"; - let description = [{ - The "init_coop_tile" operation splits an XeTile among multiple subgroups within the same - workgroup. This enables cooperative work among the subgroups such as cooperative prefetch, - load A or B matrices. Subgroups within the same workgroup share the memory accesses. - - This operation has the following arguments: - * source : this is the XeTile that will be split among subgroups - * coop_id : ID assigned to each subgroup - * coop_size : size assigned to each subgroup - - Example 1: - - ```mlir - %1 = xetile.init_coop_tile %src, 0, 64 - : tile<64x64xf32>, index, index -> tile<8x8xf32> - ``` - }]; - - let arguments = (ins - XeTile: $source, - Index: $coop_id, - Index: $coop_size); - let results = (outs XeTile: $result); - - let assemblyFormat = [{ - $source `,` ` ` $coop_id `,` ` ` $coop_size attr-dict - `:` qualified(type($source)) `,` qualified(type($coop_id)) `,` qualified(type($coop_size)) - `->` qualified(type($result)) - }]; - -} - def XeTile_LoadTileOp : XeTile_Op<"load_tile", []> { let summary = "Loads a tile into a register region"; let description = [{ @@ -339,19 +305,17 @@ def XeTile_StoreTileOp : XeTile_Op<"store_tile", []> { def XeTile_PrefetchTileOp : XeTile_Op<"prefetch_tile", []> { let summary = "preftech tiles into cache cooperatively"; let description = [{ - "preftech_tile" operation can be used to prefetch smaller tiles into the cache. - A larger tile can be split into smaller tiles using the init_coop_tile operation and - then the resulting smaller tiles can be passed into this operation for prefetching. + "preftech_tile" operation can be used to smaller tiles into the cache. This operation takes following arguments: * tile : tile to prefetch into the cache Example 1: ```mlir - %1 = xetile.init_coop_tile %src, %c0, %c1 - : tile<64x64xf32>, index, index -> tile<8x8xf32> + %1 = memref.alloc() : memref<1024x1024xf16> + %2 = xetile.init_tile %src[0, 0] : !xetile.tile<128x128xf16> - xetile.prefetch_tile %1 : (tile<8x8xf32>) + xetile.prefetch_tile %2 : tile<8x8xf32> ``` @@ -360,7 +324,7 @@ def XeTile_PrefetchTileOp : XeTile_Op<"prefetch_tile", []> { let arguments = (ins XeTile:$tile); let assemblyFormat = [{ - $tile attr-dict `:` `(` qualified(type($tile)) `)` + $tile attr-dict `:` qualified(type($tile)) }]; } @@ -409,8 +373,8 @@ def XeTile_TileMMAOp : XeTile_Op<"tile_mma", [Pure]> { let results = (outs XeTile_2DOr4DVector: $output); // let assemblyFormat = [{ - // $a`,` ` `$b (`,` ` `$c^)? attr-dict `:` `(`qualified(type($a))`,` ` `qualified(type($b)) - // (`,` ` `qualified(type($c))^)?`)` `->` qualified(type($output)) + // $a`,` ` `$b (`,` ` `$c^)? attr-dict `:` qualified(type($a))`,` ` `qualified(type($b)) + // (`,` ` `qualified(type($c))^)? `->` qualified(type($output)) // }]; let hasCustomAssemblyFormat = true; diff --git a/test/Dialect/XeTile/IR/ops.mlir b/test/Dialect/XeTile/IR/ops.mlir index cef7228e2..c34623519 100644 --- a/test/Dialect/XeTile/IR/ops.mlir +++ b/test/Dialect/XeTile/IR/ops.mlir @@ -111,19 +111,6 @@ func.func @test_init_tile_using_addr(%src: i64, %dim0_size : index, %dim1_size : return } -// CHECK-LABEL: func @test_init_coop_tile({{.*}}) { -func.func @test_init_coop_tile(%src: !xetile.tile<64x64xf16>) { - %c0 = arith.constant 0 : index - %c1 = arith.constant 64 : index - - // CHECK: xetile.init_coop_tile - // CHECK-SAME: !xetile.tile<64x64xf16>, index, index -> !xetile.tile<8x8xf16> - %1 = xetile.init_coop_tile %src, %c0, %c1 - : !xetile.tile<64x64xf16>, index, index -> !xetile.tile<8x8xf16> - - return -} - // CHECK-LABEL: func @test_load_tile({{.*}}) { func.func @test_load_tile(%src: !xetile.tile<64x32xf16>, %src1 : !xetile.tile<128x128xf16, #xe_map>) { @@ -180,19 +167,12 @@ func.func @test_store_tile(%value1 : vector<64x32xf16>, return } -// CHECK-LABEL: func @test_coop_prefetch_tile({{.*}}) { -func.func @test_coop_prefetch_tile(%src: !xetile.tile<64x64xf16>) { - %c0 = arith.constant 0 : index - %c1 = arith.constant 64 : index - - // CHECK: xetile.init_coop_tile - // CHECK-SAME: !xetile.tile<64x64xf16>, index, index -> !xetile.tile<8x8xf16> - %1 = xetile.init_coop_tile %src, %c0, %c1 - : !xetile.tile<64x64xf16>, index, index -> !xetile.tile<8x8xf16> +// CHECK-LABEL: func @test_prefetch_tile({{.*}}) { +func.func @test_prefetch_tile(%src: !xetile.tile<64x64xf16>) { // CHECK: xetile.prefetch_tile - // CHECK-SAME: (!xetile.tile<8x8xf16>) - xetile.prefetch_tile %1 : (!xetile.tile<8x8xf16>) + // CHECK-SAME: !xetile.tile<64x64xf16> + xetile.prefetch_tile %src : !xetile.tile<64x64xf16> return }