Repository: tenstorrent/tt-mlir
Pull request #3196: Add init_short and uninit ops for ttkernel tilize and untilize
Author: login: akshaybaviskar
Labels: name: community
Created: 2025-05-04T06:15:34Z | Merged: 2025-05-09T01:09:10Z

### Ticket
Link to Github Issue

### Problem description
Add init_short and uninit ops for ttkernel tilize and untilize.

### What's changed
Add tilize_init_short,  tilize_uninit, untilize_init_short,  untilize_uninit ops to ttkernel ops.
Also add lowering to emitc for the above ops.

### Checklist
- [ ] New/Existing tests provide coverage for changes



Diff:
diff --git a/include/ttmlir/Dialect/TTKernel/IR/TTKernelOps.td b/include/ttmlir/Dialect/TTKernel/IR/TTKernelOps.td
index a4ab3f40225..149d74d6b27 100644
--- a/include/ttmlir/Dialect/TTKernel/IR/TTKernelOps.td
+++ b/include/ttmlir/Dialect/TTKernel/IR/TTKernelOps.td
@@ -380,10 +380,10 @@ def TTKernel_CBWaitFrontOp : TTKernel_Op<"cb_wait_front"> {
 // TTKernel Tile operations
 //===----------------------------------------------------------------------===//
 
-def TTKernel_TilizeInitOp : TTKernel_Op<"tilize_init"> {
-    let summary = "TilizeInitOp call.";
+def TTKernel_TilizeBlockOp : TTKernel_Op<"tilize_block"> {
+    let summary = "TilizeBlockOp call.";
     let description = [{
-      TilizeInitOp operation
+      TilizeBlockOp operation
     }];
 
     let arguments = (ins TTKernel_CB:$cbIn, I32:$numTiles, TTKernel_CB:$cbOut);
@@ -391,26 +391,35 @@ def TTKernel_TilizeInitOp : TTKernel_Op<"tilize_init"> {
     let hasVerifier = 1;
 }
 
-def TTKernel_UntilizeInitOp : TTKernel_Op<"untilize_init"> {
-    let summary = "UntilizeInitOp call.";
+def TTKernel_TilizeInitOp : TTKernel_Op<"tilize_init"> {
+    let summary = "TilizeInitOp call.";
     let description = [{
-      UntilizeInitOp operation
+      Initialize the tilize operation. To be called once at beginning of a
+      kernel.
     }];
 
-    let arguments = (ins TTKernel_CB:$cbIn, TTKernel_CB:$cbOut);
+    let arguments = (ins TTKernel_CB:$cbIn, I32:$numTiles, TTKernel_CB:$cbOut);
 
     let hasVerifier = 1;
 }
 
-def TTKernel_TilizeBlockOp : TTKernel_Op<"tilize_block"> {
-    let summary = "TilizeBlockOp call.";
+def TTKernel_TilizeInitShortOp : TTKernel_Op<"tilize_init_short"> {
+    let summary = "TilizeInitShortOp call.";
+     let description = [{
+      Re-initialize for the tilize operation. This can be called after a full
+      init.
+     }];
+
+    let arguments = (ins TTKernel_CB:$cbIn, I32:$numiles, TTKernel_CB:$cbOut);
+}
+
+def TTKernel_TilizeUninitOp : TTKernel_Op<"tilize_uninit"> {
+    let summary = "TilizeUninitOp call.";
     let description = [{
-      TilizeBlockOp operation
+      Uninitialize tilize operation before re-initializing for another operation.
     }];
 
-    let arguments = (ins TTKernel_CB:$cbIn, I32:$numTiles, TTKernel_CB:$cbOut);
-
-    let hasVerifier = 1;
+    let arguments = (ins TTKernel_CB:$cbI, TTKernel_CB:$cbOut);
 }
 
 def TTKernel_UntilizeBlockOp : TTKernel_Op<"untilize_block"> {
@@ -424,6 +433,37 @@ def TTKernel_UntilizeBlockOp : TTKernel_Op<"untilize_block"> {
     let hasVerifier = 1;
 }
 
+def TTKernel_UntilizeInitOp : TTKernel_Op<"untilize_init"> {
+    let summary = "UntilizeInitOp call.";
+    let description = [{
+      Init function for untilize operations, to be used at the beginning of the
+      kernel.
+    }];
+
+    let arguments = (ins TTKernel_CB:$cbIn, TTKernel_CB:$cbOut);
+
+    let hasVerifier = 1;
+}
+
+def TTKernel_UntilizeInitShortOp : TTKernel_Op<"untilize_init_short"> {
+    let summary = "UntilizeInitShortOp call.";
+    let description = [{
+      Re-initialize for the tilize operation. This can be called after a full
+      init.
+    }];
+
+    let arguments = (ins TTKernel_CB:$cbIn);
+}
+
+def TTKernel_UntilizeUninitOp : TTKernel_Op<"untilize_uninit"> {
+    let summary = "UntilizeUninitOp call.";
+    let description = [{
+      Uninitialize untilize operation, to allow initializing another operation.
+    }];
+
+    let arguments = (ins TTKernel_CB:$cbIn);
+}
+
 //===----------------------------------------------------------------------===//
 // TTKernel NOC operations
 //===----------------------------------------------------------------------===//
diff --git a/lib/Conversion/TTKernelToEmitC/TTKernelToEmitC.cpp b/lib/Conversion/TTKernelToEmitC/TTKernelToEmitC.cpp
index 6ec91e93339..f143816aee0 100644
--- a/lib/Conversion/TTKernelToEmitC/TTKernelToEmitC.cpp
+++ b/lib/Conversion/TTKernelToEmitC/TTKernelToEmitC.cpp
@@ -570,7 +570,11 @@ class ConvertTTKernelToEmitCPass
         TTKernelToEmitCOpaqueRewriter<ttkernel::CBReserveBackOp>,
         TTKernelToEmitCOpaqueRewriter<ttkernel::CBWaitFrontOp>,
         TTKernelToEmitCOpaqueRewriter<ttkernel::TilizeInitOp>,
+        TTKernelToEmitCOpaqueRewriter<ttkernel::TilizeInitShortOp>,
+        TTKernelToEmitCOpaqueRewriter<ttkernel::TilizeUninitOp>,
         TTKernelToEmitCOpaqueRewriter<ttkernel::UntilizeInitOp>,
+        TTKernelToEmitCOpaqueRewriter<ttkernel::UntilizeInitShortOp>,
+        TTKernelToEmitCOpaqueRewriter<ttkernel::UntilizeUninitOp>,
         TTKernelToEmitCOpaqueRewriter<ttkernel::TilizeBlockOp>,
         TTKernelToEmitCOpaqueRewriter<ttkernel::UntilizeBlockOp>,
         TTKernelToEmitCOpaqueRewriter<ttkernel::BinaryOpInitCommonOp>,
diff --git a/test/ttmlir/Conversion/TTKernelToEmitC/ttkernel.mlir b/test/ttmlir/Conversion/TTKernelToEmitC/ttkernel.mlir
index 727a9b85f92..9860695c37a 100644
--- a/test/ttmlir/Conversion/TTKernelToEmitC/ttkernel.mlir
+++ b/test/ttmlir/Conversion/TTKernelToEmitC/ttkernel.mlir
@@ -333,6 +333,26 @@ module {
       return
     }
 
+    // CHECK-LABEL: func @tilize_init_short
+    func.func @tilize_init_short(%in_cb: !cb0_scalar, %out_cb: !cb1_tiles) -> () attributes {ttkernel.thread = #ttkernel.thread<compute>} {
+      // CHECK: %[[IN_CB:.*]] = "emitc.constant"{{.+}}<"::tt::CB">
+      // CHECK: %[[OUT_CB:.*]] = "emitc.constant"{{.+}}<"::tt::CB">
+      // CHECK: %[[NUM_TILES:.*]] = "emitc.constant"
+      %num_tiles = arith.constant 3 : i32
+      // CHECK: emitc.call_opaque "tilize_init_short"(%[[IN_CB]], %[[NUM_TILES]], %[[OUT_CB]])
+      "ttkernel.tilize_init_short"(%in_cb, %num_tiles, %out_cb) : (!cb0_scalar, i32, !cb1_tiles) -> ()
+      return
+    }
+
+    // CHECK-LABEL: func @tilize_uninit
+    func.func @tilize_uninit(%in_cb: !cb0_scalar, %out_cb: !cb1_tiles) -> () attributes {ttkernel.thread = #ttkernel.thread<compute>} {
+      // CHECK: %[[IN_CB:.*]] = "emitc.constant"{{.+}}<"::tt::CB">
+      // CHECK: %[[OUT_CB:.*]] = "emitc.constant"{{.+}}<"::tt::CB">
+      // CHECK: emitc.call_opaque "tilize_uninit"(%[[IN_CB]], %[[OUT_CB]])
+      "ttkernel.tilize_uninit"(%in_cb, %out_cb) : (!cb0_scalar, !cb1_tiles) -> ()
+      return
+    }
+
     // CHECK-LABEL: func @untilize_init
     func.func @untilize_init(%in_cb: !cb0_tiles, %out_cb: !cb1_scalar) -> () attributes {ttkernel.thread = #ttkernel.thread<compute>} {
       // CHECK: %[[IN_CB:.*]] = "emitc.constant"{{.+}}<"::tt::CB">
@@ -342,6 +362,22 @@ module {
       return
     }
 
+    // CHECK-LABEL: func @untilize_init_short
+    func.func @untilize_init_short(%in_cb: !cb0_tiles) -> () attributes {ttkernel.thread = #ttkernel.thread<compute>} {
+      // CHECK: %[[IN_CB:.*]] = "emitc.constant"{{.+}}<"::tt::CB">
+      // CHECK: emitc.call_opaque "untilize_init_short"(%[[IN_CB]])
+      "ttkernel.untilize_init_short"(%in_cb) : (!cb0_tiles) -> ()
+      return
+    }
+
+    // CHECK-LABEL: func @untilize_uninit
+    func.func @untilize_uninit(%in_cb: !cb0_tiles) -> () attributes {ttkernel.thread = #ttkernel.thread<compute>} {
+      // CHECK: %[[IN_CB:.*]] = "emitc.constant"{{.+}}<"::tt::CB">
+      // CHECK: emitc.call_opaque "untilize_uninit"(%[[IN_CB]])
+      "ttkernel.untilize_uninit"(%in_cb) : (!cb0_tiles) -> ()
+      return
+    }
+
     // CHECK-LABEL: func @tilize_block
     func.func @tilize_block(%in_cb: !cb0_scalar, %out_cb: !cb1_tiles) -> () attributes {ttkernel.thread = #ttkernel.thread<compute>} {
       // CHECK: %[[IN_CB:.*]] = "emitc.constant"{{.+}}<"::tt::CB">
diff --git a/test/ttmlir/Dialect/TTKernel/ops.mlir b/test/ttmlir/Dialect/TTKernel/ops.mlir
new file mode 100644
index 00000000000..68a865fef9c
--- /dev/null
+++ b/test/ttmlir/Dialect/TTKernel/ops.mlir
@@ -0,0 +1,26 @@
+// RUN: ttmlir-opt %s | FileCheck %s
+
+#l1_ = #tt.memory_space<l1>
+// CHECK-LABEL: func @test_untilize_init_short
+func.func @test_untilize_init_short(%tilized_cb : !ttkernel.cb<cb_in0, 0, memref<4x4x!tt.tile<32x32, f32>, #l1_>, 4096, 1>) -> () {
+  "ttkernel.untilize_init_short"(%tilized_cb) : (!ttkernel.cb<cb_in0, 0, memref<4x4x!tt.tile<32x32, f32>, #l1_>, 4096, 1>) -> ()
+  return
+}
+
+// CHECK-LABEL: func @test_tilize_init_short
+func.func @test_tilize_init_short(%tilized_cb : !ttkernel.cb<cb_in0, 0, memref<4x4x!tt.tile<32x32, f32>, #l1_>, 4096, 1>, %untilized_cb : !ttkernel.cb<cb_in0, 0, memref<128x128xf32, #l1_>, 512, 1>, %num_tiles : i32) -> () {
+  "ttkernel.tilize_init_short"(%untilized_cb, %num_tiles, %tilized_cb) : (!ttkernel.cb<cb_in0, 0, memref<128x128xf32, #l1_>, 512, 1>, i32, !ttkernel.cb<cb_in0, 0, memref<4x4x!tt.tile<32x32, f32>, #l1_>, 4096, 1>) -> ()
+  return
+}
+
+// CHECK-LABEL: func.func @test_tilize_uninit
+func.func @test_tilize_uninit(%tilized_cb : !ttkernel.cb<cb_in0, 0, memref<4x4x!tt.tile<32x32, f32>, #l1_>, 4096, 1>, %untilized_cb : !ttkernel.cb<cb_in0, 0, memref<128x128xf32, #l1_>, 512, 1>) -> () {
+  "ttkernel.tilize_uninit"(%untilized_cb, %tilized_cb) : (!ttkernel.cb<cb_in0, 0, memref<128x128xf32, #l1_>, 512, 1>, !ttkernel.cb<cb_in0, 0, memref<4x4x!tt.tile<32x32, f32>, #l1_>, 4096, 1>) -> ()
+  return
+}
+
+// CHECK-LABEL: func.func @test_untilize_uninit
+func.func @test_untilize_uninit(%tilized_cb : !ttkernel.cb<cb_in0, 0, memref<4x4x!tt.tile<32x32, f32>, #l1_>, 4096, 1>) -> () {
+  "ttkernel.untilize_uninit"(%tilized_cb) : (!ttkernel.cb<cb_in0, 0, memref<4x4x!tt.tile<32x32, f32>, #l1_>, 4096, 1>) -> ()
+  return
+}


Reviews:
user: login: jdesousa-TT
state: CHANGES_REQUESTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-05T15:42:10Z
user: login: jdesousa-TT
body: Looks good to me, thanks! 
state: APPROVED
author_association: CONTRIBUTOR
submitted_at: 2025-05-08T17:48:18Z
user: login: nsmithtt
state: APPROVED
author_association: CONTRIBUTOR
submitted_at: 2025-05-08T18:24:59Z

Review comments:
diff_hunk: @@ -359,10 +359,52 @@ def TTKernel_CBWaitFrontOp : TTKernel_Op<"cb_wait_front"> {
 // TTKernel Tile operations
 //===----------------------------------------------------------------------===//
 
+def TTKernel_TilizeBlockOp : TTKernel_Op<"tilize_block"> {
+    let summary = "TilizeBlockOp call.";
+    let description = [{
+      TilizeBlockOp operation
+    }];
+
+    let arguments = (ins TTKernel_CB:$cbIn, I32:$numTiles, TTKernel_CB:$cbOut);
+
+    let hasVerifier = 1;
+}
+
 def TTKernel_TilizeInitOp : TTKernel_Op<"tilize_init"> {
     let summary = "TilizeInitOp call.";
     let description = [{
-      TilizeInitOp operation
+      Initialize the tilize operation. To be called once at beginning of a
+      kernel.
+    }];
+
+    let arguments = (ins TTKernel_CB:$cbIn, I32:$numTiles, TTKernel_CB:$cbOut);
+
+    let hasVerifier = 1;
+}
+
+def TTKernel_TilizeInitShortOp : TTKernel_Op<"tilize_init_short"> {
+    let summary = "TilizeInitShortOp call.";
+     let description = [{
+      Re-initialize for the tilize operation. This can be called after a full
+      init.
+     }];
+ 
+    let arguments = (ins TTKernel_CB:$cbIn, I32:$numiles, TTKernel_CB:$cbOut);
+}
+
+def TTKernel_TilizeUninitOp : TTKernel_Op<"tilize_uninit"> {
+    let summary = "TilizeUninitOp call.";
+    let description = [{
+      Uninitialize tilize operation before re-initializing for another operation.
+    }];
+
+    let arguments = (ins TTKernel_CB:$cbI, I32:$numTiles, TTKernel_CB:$cbOut);
path: include/ttmlir/Dialect/TTKernel/IR/TTKernelOps.td
user: login: jdesousa-TT
body: `tilize_uninit` only takes input/output cbs as arguments. See `tt_metal/include/compute_kernel_api/tilize.h`
created_at: 2025-05-05T15:38:02Z
author_association: CONTRIBUTOR
diff_hunk: @@ -373,34 +415,32 @@ def TTKernel_TilizeInitOp : TTKernel_Op<"tilize_init"> {
 def TTKernel_UntilizeInitOp : TTKernel_Op<"untilize_init"> {
     let summary = "UntilizeInitOp call.";
     let description = [{
-      UntilizeInitOp operation
+      Init function for untilize operations, to be used at the beginning of the
+      kernel.
     }];
 
     let arguments = (ins TTKernel_CB:$cbIn, TTKernel_CB:$cbOut);
 
     let hasVerifier = 1;
 }
 
-def TTKernel_TilizeBlockOp : TTKernel_Op<"tilize_block"> {
-    let summary = "TilizeBlockOp call.";
+def TTKernel_UntilizeInitShortOp : TTKernel_Op<"untilize_init_short"> {
+    let summary = "UntilizeInitShortOp call.";
     let description = [{
-      TilizeBlockOp operation
+      Re-initialize for the tilize operation. This can be called after a full
+      init.
     }];
 
     let arguments = (ins TTKernel_CB:$cbIn, I32:$numTiles, TTKernel_CB:$cbOut);
path: include/ttmlir/Dialect/TTKernel/IR/TTKernelOps.td
user: login: jdesousa-TT
body: `untilize_init_short` only takes input cb as the argument. See `tt_metal/include/compute_kernel_api/untilize.h`
created_at: 2025-05-05T15:40:29Z
author_association: CONTRIBUTOR
diff_hunk: @@ -373,34 +415,32 @@ def TTKernel_TilizeInitOp : TTKernel_Op<"tilize_init"> {
 def TTKernel_UntilizeInitOp : TTKernel_Op<"untilize_init"> {
     let summary = "UntilizeInitOp call.";
     let description = [{
-      UntilizeInitOp operation
+      Init function for untilize operations, to be used at the beginning of the
+      kernel.
     }];
 
     let arguments = (ins TTKernel_CB:$cbIn, TTKernel_CB:$cbOut);
 
     let hasVerifier = 1;
 }
 
-def TTKernel_TilizeBlockOp : TTKernel_Op<"tilize_block"> {
-    let summary = "TilizeBlockOp call.";
+def TTKernel_UntilizeInitShortOp : TTKernel_Op<"untilize_init_short"> {
+    let summary = "UntilizeInitShortOp call.";
     let description = [{
-      TilizeBlockOp operation
+      Re-initialize for the tilize operation. This can be called after a full
+      init.
     }];
 
     let arguments = (ins TTKernel_CB:$cbIn, I32:$numTiles, TTKernel_CB:$cbOut);
-
-    let hasVerifier = 1;
 }
 
-def TTKernel_UntilizeBlockOp : TTKernel_Op<"untilize_block"> {
-    let summary = "UntilizeBlockOp call.";
+def TTKernel_UntilizeUninitOp : TTKernel_Op<"untilize_uninit"> {
+    let summary = "UntilizeUninitOp call.";
     let description = [{
-      UntilizeBlockOp operation
+      Uninitialize untilize operation, to allow initializing another operation.
     }];
 
     let arguments = (ins TTKernel_CB:$cbIn, I32:$numTiles, TTKernel_CB:$cbOut);
path: include/ttmlir/Dialect/TTKernel/IR/TTKernelOps.td
user: login: jdesousa-TT
body: `uintilize_uninit` only takes the input cb, same as above.
created_at: 2025-05-05T15:41:49Z
author_association: CONTRIBUTOR

PR comments:
user: login: codecov-commenter
created_at: 2025-05-05T13:43:31Z
body: ## [Codecov](https://app.codecov.io/gh/tenstorrent/tt-mlir/pull/3196?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=tenstorrent) Report
:white_check_mark: All modified and coverable lines are covered by tests.
:white_check_mark: Project coverage is 79.87%. Comparing base ([`40082a6`](https://app.codecov.io/gh/tenstorrent/tt-mlir/commit/40082a648f7963a959a784c7c69130f124a1a2fb?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=tenstorrent)) to head ([`0d0b403`](https://app.codecov.io/gh/tenstorrent/tt-mlir/commit/0d0b4035568e237bb3c5cdb42bae557555e255c7?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=tenstorrent)).
:warning: Report is 2024 commits behind head on main.

<details><summary>Additional details and impacted files</summary>



```diff
@@            Coverage Diff             @@
##             main    #3196      +/-   ##
==========================================
- Coverage   79.92%   79.87%   -0.06%     
==========================================
  Files         200      200              
  Lines       26473    26473              
==========================================
- Hits        21159    21144      -15     
- Misses       5314     5329      +15     
```
</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/tenstorrent/tt-mlir/pull/3196?dropdown=coverage&src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=tenstorrent).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=tenstorrent).
author_association: NONE
user: login: akshaybaviskar
created_at: 2025-05-07T12:06:50Z
body: Thanks for the feedback @jdesousa-TT , addressed all the comments. Also, added the lowering to emitC for the newly added ops. Please re-review.
author_association: CONTRIBUTOR
user: login: akshaybaviskar
created_at: 2025-05-08T16:25:53Z
body: Hi @nsmithtt  @jdesousa-TT, can you please re-review this PR?
Thanks.
author_association: CONTRIBUTOR
user: login: nsmithtt
created_at: 2025-05-08T18:24:34Z
body: LGTM too! Currently failing CI due to syntax errors, please run the following and then update the PR:
```bash
pre-commit run --all-files
```
author_association: CONTRIBUTOR
user: login: akshaybaviskar
created_at: 2025-05-08T20:50:39Z
body: > LGTM too! Currently failing CI due to syntax errors, please run the following and then update the PR:
> 
> ```shell
> pre-commit run --all-files
> ```

Thanks, done.
author_association: CONTRIBUTOR