From f15e704d33d234cf9a63d8a6665c256f41783e49 Mon Sep 17 00:00:00 2001
From: Jacob DeSousa <jdesousa@tenstorrent.com>
Date: Mon, 24 Feb 2025 20:42:26 -0500
Subject: [PATCH] Add tilize/untilize ttir generic region ops (#2274)

### Ticket
#1948

### What's changed
Add `ttir.tile_tilize_block` and `ttir.tile_untilize_block` ops to the
TTIR generic region ops.
Add respective verifiers to match input/output element types.

### Checklist
- [x] New/Existing tests provide coverage for changes
---
 .../Dialect/TTIR/IR/TTIRGenericRegionOps.td   | 34 +++++++++++++++++++
 lib/Dialect/TTIR/IR/TTIRGenericRegionOps.cpp  | 34 +++++++++++++++++++
 2 files changed, 68 insertions(+)

diff --git a/include/ttmlir/Dialect/TTIR/IR/TTIRGenericRegionOps.td b/include/ttmlir/Dialect/TTIR/IR/TTIRGenericRegionOps.td
index ca06fb5724..f3b9757556 100644
--- a/include/ttmlir/Dialect/TTIR/IR/TTIRGenericRegionOps.td
+++ b/include/ttmlir/Dialect/TTIR/IR/TTIRGenericRegionOps.td
@@ -136,4 +136,38 @@ def TTIR_TileMatmulBlockOp : TTIR_GenericRegionOp<"tile_matmul_block",
     let hasVerifier = 1;
 }
 
+def TTIR_TileTilizeBlockOp : TTIR_GenericRegionOp<"tile_tilize_block",
+  [DestinationStyleOpInterface, MemoryEffects<[MemRead, MemWrite]>]> {
+    let summary = "TTIR Tile Tilize Block Op";
+    let description = [{
+        The `tile_tilize_block` operation tilizes the input row major memref block and outputs the memref containing the tilized data.
+    }];
+
+    let arguments = (ins AnyNon0RankedMemRef:$input,
+                         AnyNon0RankedMemRef:$output);
+
+    let extraClassDeclaration = [{
+      MutableOperandRange getDpsInitsMutable() { return getOutputMutable(); }
+    }];
+
+    let hasVerifier = 1;
+}
+
+def TTIR_TileUntilizeBlockOp : TTIR_GenericRegionOp<"tile_untilize_block",
+  [DestinationStyleOpInterface, MemoryEffects<[MemRead, MemWrite]>]> {
+    let summary = "TTIR Tile Untilize Block Op";
+    let description = [{
+        The `tile_untilize_block` operation untilizes the input tilized memref block and outputs the memref contianing the row major data.
+    }];
+
+    let arguments = (ins AnyNon0RankedMemRef:$input,
+                         AnyNon0RankedMemRef:$output);
+
+    let extraClassDeclaration = [{
+      MutableOperandRange getDpsInitsMutable() { return getOutputMutable(); }
+    }];
+
+    let hasVerifier = 1;
+}
+
 #endif // TTMLIR_TTMLIR_DIALECT_TTIR_TTIRGENERICREGIONOPS_TD
diff --git a/lib/Dialect/TTIR/IR/TTIRGenericRegionOps.cpp b/lib/Dialect/TTIR/IR/TTIRGenericRegionOps.cpp
index bd15ddccec..9a86dd8c5f 100644
--- a/lib/Dialect/TTIR/IR/TTIRGenericRegionOps.cpp
+++ b/lib/Dialect/TTIR/IR/TTIRGenericRegionOps.cpp
@@ -22,3 +22,37 @@ ::mlir::LogicalResult mlir::tt::ttir::TileMatmulBlockOp::verify() {
 
   return success();
 }
+
+// TileTilizeBlockOp verification
+::mlir::LogicalResult mlir::tt::ttir::TileTilizeBlockOp::verify() {
+
+  if (llvm::isa<mlir::tt::TileType>(getInput().getType().getElementType())) {
+    return emitOpError(
+        "MemRef operand to TileTilizeBlock must not have tt.tile "
+        "element type");
+  }
+
+  if (!llvm::isa<mlir::tt::TileType>(getOutput().getType().getElementType())) {
+    return emitOpError("MemRef result of TileTilizeBlock must have tt.tile "
+                       "element type");
+  }
+
+  return success();
+}
+
+// TileUntilizeBlockOp verification
+::mlir::LogicalResult mlir::tt::ttir::TileUntilizeBlockOp::verify() {
+
+  if (!llvm::isa<mlir::tt::TileType>(getInput().getType().getElementType())) {
+    return emitOpError("MemRef operand to TileUntilizeBlock must have tt.tile "
+                       "element type");
+  }
+
+  if (llvm::isa<mlir::tt::TileType>(getOutput().getType().getElementType())) {
+    return emitOpError(
+        "MemRef result of TileUntilizeBlock must not have tt.tile "
+        "element type");
+  }
+
+  return success();
+}