#14063: Fix test by replace wrong default arguments in kernel code

tenstorrent · Dec 4, 2024 · 0f73762 · 0f73762
1 parent 98be31d
commit 0f73762
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 4 deletions.
diff --git a/tests/tt_eager/python_api_testing/unit_testing/misc/test_rotary_embedding.py b/tests/tt_eager/python_api_testing/unit_testing/misc/test_rotary_embedding.py
@@ -100,7 +100,6 @@ def test_rotary_embedding_prefill(W, Z, Y, X, cache_size, in_sharded, out_sharde
     assert p
 
 
-@skip_for_blackhole("Mismatching on Blackhole, see #12349")
 @pytest.mark.parametrize(
     "W, Z, Y, X",
     ([1, 1, 32, 64], [1, 71, 32, 64], [1, 1, 64, 64], [1, 71, 64, 64], [1, 32, 32, 64], [1, 2, 32, 64]),
@@ -246,7 +245,6 @@ def test_rotary_embedding_prefill_fp32(
     assert p
 
 
-@skip_for_blackhole("Mismatching on Blackhole, see #12349")
 @pytest.mark.skipif(is_grayskull(), reason="GS does not support fp32")
 @pytest.mark.parametrize("W, Z, Y, X", [(1, 1, 32, 64)])
 @pytest.mark.parametrize("cache_size", [2048])

diff --git a/...ons/experimental/transformer/rotary_embedding/device/kernels/compute/rotary_embedding.cpp b/...ons/experimental/transformer/rotary_embedding/device/kernels/compute/rotary_embedding.cpp
@@ -51,7 +51,7 @@ ALWI void UNTILIZE_TILES(uint32_t in0_cb, uint32_t out_cb, uint32_t num_tiles) {
 }
 
 ALWI void TILIZE_ROWS(uint32_t in0_cb, uint32_t sync_cb, uint32_t out_cb, uint32_t num_tiles) {
-    tilize_init_short(in0_cb, num_tiles);
+    tilize_init_short(in0_cb, num_tiles, out_cb);
     cb_wait_front(in0_cb, num_tiles);
     cb_wait_front(sync_cb, num_tiles);
     cb_reserve_back(out_cb, num_tiles);
@@ -61,7 +61,7 @@ ALWI void TILIZE_ROWS(uint32_t in0_cb, uint32_t sync_cb, uint32_t out_cb, uint32
     // Pop shared cbs after tilize
     cb_pop_front(in0_cb, num_tiles);
     cb_pop_front(sync_cb, num_tiles);
-    tilize_uninit(in0_cb);
+    tilize_uninit(in0_cb, out_cb);
 }
 
 namespace NAMESPACE {