From 2a95b1e5c9f599604d114189059aa72580fd11bf Mon Sep 17 00:00:00 2001
From: Milan Topalovic <163355844+mtopalovicTT@users.noreply.github.com>
Date: Fri, 6 Sep 2024 21:43:48 +0200
Subject: [PATCH] Adding transpose_twice test (#606)

* Fixing layouting for output and op rewrite

* Adding rewrite

* Addressing comments
---
 .../Dialect/TTNN/transpose/transpose_twice.mlir     | 13 +++++++++++++
 1 file changed, 13 insertions(+)
 create mode 100644 test/ttmlir/Dialect/TTNN/transpose/transpose_twice.mlir

diff --git a/test/ttmlir/Dialect/TTNN/transpose/transpose_twice.mlir b/test/ttmlir/Dialect/TTNN/transpose/transpose_twice.mlir
new file mode 100644
index 000000000..f18e0e5cb
--- /dev/null
+++ b/test/ttmlir/Dialect/TTNN/transpose/transpose_twice.mlir
@@ -0,0 +1,13 @@
+// RUN: ttmlir-opt --ttir-to-ttnn-backend-pipeline %s | FileCheck %s
+#any_device_tile = #tt.operand_constraint<dram|l1|tile|any_device_tile>
+module attributes {} {
+  func.func @forward(%arg0: tensor<1x16x32x64xbf16>) -> tensor<1x32x64x16xbf16> {
+    %0 = tensor.empty() : tensor<1x64x32x16xbf16>
+    // CHECK: %[[C:.*]] = "ttnn.transpose"[[C:.*]]
+    %1 = "ttir.transpose"(%arg0, %0) <{dim0 = -3 : si32, dim1 = -1 : si32, operand_constraints = [#any_device_tile, #any_device_tile]}> : (tensor<1x16x32x64xbf16>, tensor<1x64x32x16xbf16>) -> tensor<1x64x32x16xbf16>
+    %2 = tensor.empty() : tensor<1x32x64x16xbf16>
+    // CHECK: %[[C:.*]] = "ttnn.transpose
+    %3 = "ttir.transpose"(%1, %2) <{dim0 = -3 : si32, dim1 = -2 : si32, operand_constraints = [#any_device_tile, #any_device_tile]}> : (tensor<1x64x32x16xbf16>, tensor<1x32x64x16xbf16>) -> tensor<1x32x64x16xbf16>
+    return %3 : tensor<1x32x64x16xbf16>
+  }
+}