Addressed comments

tenstorrent · Nov 25, 2024 · 03d62a5 · 03d62a5
1 parent 850686a
commit 03d62a5
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 4 deletions.
diff --git a/lib/Dialect/TTIR/Transforms/Reshape.cpp b/lib/Dialect/TTIR/Transforms/Reshape.cpp
@@ -13,7 +13,7 @@ namespace mlir::tt::ttir {
 #include "ttmlir/Dialect/TTIR/Transforms/Passes.h.inc"
 
 //===----------------------------------------------------------------------===//
-// Constant as fill pass
+// Reshape folding pass
 //===----------------------------------------------------------------------===//
 
 class TTIRReshapeFoldingRewriter : public OpRewritePattern<ReshapeOp> {

diff --git a/test/ttmlir/Dialect/TTIR/reshape/reshape_test.mlir b/test/ttmlir/Dialect/TTIR/reshape/reshape_test.mlir
@@ -1,7 +1,8 @@
 // RUN: ttmlir-opt --ttir-reshape-fold %s| FileCheck %s
-#any_device_tile = #tt.operand_constraint<dram|l1|tile|none|interleaved|single_bank|height_sharded|width_sharded|block_sharded|any_layout|any_device_tile>
-module @jit_ravel attributes {mhlo.num_partitions = 1 : i32, mhlo.num_replicas = 1 : i32} {
-  func.func public @main(%arg0: tensor<1xi32>) -> (tensor<1xi32> {jax.result_info = ""}) {
+#any_device_tile = #tt.operand_constraint<dram|l1|tile|any_device_tile>
+// Tests if we fold when translating from "ttir.reshape" which is called on the two same shapes.
+module @reshape_test {
+  func.func @main(%arg0: tensor<1xi32>) -> (tensor<1xi32> {jax.result_info = ""}) {
     %0 = tensor.empty() : tensor<1xi32>
     %1 = "ttir.reshape"(%arg0, %0) <{operand_constraints = [#any_device_tile, #any_device_tile], shape = [1 : i32]}> : (tensor<1xi32>, tensor<1xi32>) -> tensor<1xi32>
     // CHECK: return %arg0 : tensor<1xi32>