Updates to XeTile definition.

- Remove init_coop_tile op and update tests - Some fixes to op assembly format
intel · Nov 7, 2023 · 12650a3 · 12650a3
1 parent 5e32af3
commit 12650a3
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 67 deletions.
diff --git 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
@@ -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
 }