ttnn.rms_norm fails with larger tensors as inputs

[nemanjagrujic]

ttnn.rms_norm operation fails with larger tensors (and tensors with larger ranks) as inputs. ttnn.rms_norm works for small and simple tensors as inputs such as [224, 256], [1, 256], but fails when more complicated tensors are provided such as [8, 224, 160], [8, 224, 160] with error:

RuntimeError: TT_FATAL @ ../tt_eager/tt_dnn/op_library/layernorm/layernorm_op.cpp:49: gamma.value().get_legacy_shape()[-2] == TILE_HEIGHT

Pytorch implementation works with such inputs. Also documentation at https://tenstorrent.github.io/tt-metal/latest/ttnn/ttnn/ttnn/rms_norm.html does not mention such limitations and allows ranks of tensors up to 4.

Problem is observed on both GS and WH cards.

To Reproduce
Steps to reproduce the behavior:

1. Checkout branch ngrujic/op_bug_unit_tests (soon to be merged into main).
2. Run unit test test_rmsnorm.py using this command:
   pytest tests/ttnn/python_api_testing/non_working_unit_tests/grayskull/test_rmsnorm.py

Expected behavior
There are multiple test cases presented in the unit test, which are failing with error:

"TT_FATAL @ ../tt_eager/tt_dnn/op_library/layernorm/layernorm_op.cpp:49: gamma.value().get_legacy_shape()[-2] == TILE_HEIGHT
backtrace:
 --- /home/ubuntu/tt-metal/ttnn/ttnn/_ttnn.so(+0x16a206) [0x7f002d471206]
 --- /home/ubuntu/tt-metal/build/lib/libtt_eager.so(+0x639e3c) [0x7f00303d0e3c]
 --- /home/ubuntu/tt-metal/build/lib/libtt_eager.so(_ZNSt17_Function_handlerIFSt7variantIJSt10shared_ptrIN2tt8tt_metal7ProgramEESt17reference_wrapperIS4_EEERKNS3_9operation15DeviceOperationISt6vectorINS3_6TensorESaISC_EEEERKSE_RKSB_ISt8optionalIKSC_ESaISM_EERSE_RKSB_ISK_ISC_ESaISS_EEEZNS9_6detail20run_device_operationISE_EET_SK_IS6_INS3_12CommandQueueEEERKNSA_IS10_EESJ_SQ_SW_EUlSH_SJ_SQ_SR_SW_E0_E9_M_invokeERKSt9_Any_dataSH_SJ_SQ_SR_SW_+0x2d) [0x7f00303d0f8d]
 --- std::vector<tt::tt_metal::Tensor, std::allocator<tt::tt_metal::Tensor> > tt::tt_metal::operation::detail::run_device_operation<std::vector<tt::tt_metal::Tensor, std::allocator<tt::tt_metal::Tensor> > >(std::optional<std::reference_wrapper<tt::tt_metal::CommandQueue> >, tt::tt_metal::operation::DeviceOperation<std::vector<tt::tt_metal::Tensor, std::allocator<tt::tt_metal::Tensor> > > const&, std::vector<tt::tt_metal::Tensor, std::allocator<tt::tt_metal::Tensor> > const&, std::vector<std::optional<tt::tt_metal::Tensor const>, std::allocator<std::optional<tt::tt_metal::Tensor const> > > const&, std::vector<std::optional<tt::tt_metal::Tensor>, std::allocator<std::optional<tt::tt_metal::Tensor> > > const&)
 --- std::vector<tt::tt_metal::Tensor, std::allocator<tt::tt_metal::Tensor> > tt::tt_metal::operation::run<std::vector<tt::tt_metal::Tensor, std::allocator<tt::tt_metal::Tensor> > >(tt::tt_metal::operation::DeviceOperation<std::vector<tt::tt_metal::Tensor, std::allocator<tt::tt_metal::Tensor> > > const&, std::vector<tt::tt_metal::Tensor, std::allocator<tt::tt_metal::Tensor> > const&, std::vector<std::optional<tt::tt_metal::Tensor const>, std::allocator<std::optional<tt::tt_metal::Tensor const> > > const&, std::vector<std::optional<tt::tt_metal::Tensor>, std::allocator<std::optional<tt::tt_metal::Tensor> > > const&)
 --- /home/ubuntu/tt-metal/ttnn/ttnn/_ttnn.so(+0x1f7d47) [0x7f002d4fed47]
 --- /home/ubuntu/tt-metal/ttnn/ttnn/_ttnn.so(_ZNSt17_Function_handlerIFSt6vectorIN2tt8tt_metal6TensorESaIS3_EERKS5_RKS0_ISt8optionalIKS3_ESaISA_EERKS0_IS8_IS3_ESaISF_EEEZNKS1_10operations7primary14make_layernormILNS2_13LayerNormTypeE1EEclERS9_fSA_SA_RKNS2_12MemoryConfigERKSt7variantIJNS2_29LayerNormDefaultProgramConfigENS2_38LayerNormShardedMultiCoreProgramConfigEEES8_IKSU_IJNS2_28GrayskullComputeKernelConfigENS2_27WormholeComputeKernelConfigEEEEEUlS7_SE_SJ_E_E9_M_invokeERKSt9_Any_dataS7_SE_SJ_+0x25) [0x7f002d4ff1f5]
 --- /home/ubuntu/tt-metal/build/lib/libtt_eager.so(+0x62bdad) [0x7f00303c2dad]
 --- /home/ubuntu/tt-metal/build/lib/libtt_eager.so(+0x626126) [0x7f00303bd126]
 --- tt::tt_metal::Device::push_work(std::shared_ptr<std::function<void ()> >, bool)
 --- tt::tt_metal::operation::launch_op(std::function<std::vector<tt::tt_metal::Tensor, std::allocator<tt::tt_metal::Tensor> > (std::vector<tt::tt_metal::Tensor, std::allocator<tt::tt_metal::Tensor> > const&, std::vector<std::optional<tt::tt_metal::Tensor const>, std::allocator<std::optional<tt::tt_metal::Tensor const> > > const&, std::vector<std::optional<tt::tt_metal::Tensor>, std::allocator<std::optional<tt::tt_metal::Tensor> > > const&)>&&, std::vector<tt::tt_metal::Tensor, std::allocator<tt::tt_metal::Tensor> >, std::vector<tt::tt_metal::Tensor, std::allocator<tt::tt_metal::Tensor> >&, std::vector<std::optional<tt::tt_metal::Tensor const>, std::allocator<std::optional<tt::tt_metal::Tensor const> > >, std::vector<std::optional<tt::tt_metal::Tensor>, std::allocator<std::optional<tt::tt_metal::Tensor> > >, bool)
 --- /home/ubuntu/tt-metal/ttnn/ttnn/_ttnn.so(+0x21e241) [0x7f002d525241]
 --- /home/ubuntu/tt-metal/ttnn/ttnn/_ttnn.so(+0x23fce8) [0x7f002d546ce8]
 --- /home/ubuntu/tt-metal/ttnn/ttnn/_ttnn.so(_ZNSt17_Function_handlerIFSt6vectorIN2tt8tt_metal6TensorESaIS3_EERKS5_RKS0_ISt8optionalIKS3_ESaISA_EERKS0_IS8_IS3_ESaISF_EEEZNK4ttnn10decorators11operation_tILi0ENSL_10operations13normalization7RMSNormEEclIJRS9_ST_fRKS8_INS2_12MemoryConfigEEEEEDaDpOT_EUlS7_SE_SJ_E_E9_M_invokeERKSt9_Any_dataS7_SE_SJ_+0x25) [0x7f002d547705]
 --- /home/ubuntu/tt-metal/build/lib/libtt_eager.so(+0x62bdad) [0x7f00303c2dad]
 --- /home/ubuntu/tt-metal/build/lib/libtt_eager.so(+0x626126) [0x7f00303bd126]
 --- tt::tt_metal::Device::push_work(std::shared_ptr<std::function<void ()> >, bool)
 --- tt::tt_metal::operation::launch_op(std::function<std::vector<tt::tt_metal::Tensor, std::allocator<tt::tt_metal::Tensor> > (std::vector<tt::tt_metal::Tensor, std::allocator<tt::tt_metal::Tensor> > const&, std::vector<std::optional<tt::tt_metal::Tensor const>, std::allocator<std::optional<tt::tt_metal::Tensor const> > > const&, std::vector<std::optional<tt::tt_metal::Tensor>, std::allocator<std::optional<tt::tt_metal::Tensor> > > const&)>&&, std::vector<tt::tt_metal::Tensor, std::allocator<tt::tt_metal::Tensor> >, std::vector<tt::tt_metal::Tensor, std::allocator<tt::tt_metal::Tensor> >&, std::vector<std::optional<tt::tt_metal::Tensor const>, std::allocator<std::optional<tt::tt_metal::Tensor const> > >, std::vector<std::optional<tt::tt_metal::Tensor>, std::allocator<std::optional<tt::tt_metal::Tensor> > >, bool)
 --- /home/ubuntu/tt-metal/ttnn/ttnn/_ttnn.so(+0x21d8fe) [0x7f002d5248fe]
 --- /home/ubuntu/tt-metal/ttnn/ttnn/_ttnn.so(+0x1a2422) [0x7f002d4a9422]
 --- /home/ubuntu/tt-metal/python_env/bin/python3(PyCFunction_Call+0x59) [0x5d5499]
 --- /home/ubuntu/tt-metal/python_env/bin/python3(_PyObject_MakeTpCall+0x296) [0x5d6066]
 --- /home/ubuntu/tt-metal/python_env/bin/python3() [0x4e2288]
 --- /home/ubuntu/tt-metal/python_env/bin/python3(PyObject_Call+0x62) [0x5d4c12]
 --- /home/ubuntu/tt-metal/python_env/bin/python3() [0x57a7bc]
 --- /home/ubuntu/tt-metal/python_env/bin/python3(PyObject_Call+0x25e) [0x5d4e0e]
 --- /home/ubuntu/tt-metal/python_env/bin/python3(_PyEval_EvalFrameDefault+0x1f26) [0x548a66]
 --- /home/ubuntu/tt-metal/python_env/bin/python3(_PyEval_EvalCodeWithName+0x26a) [0x54552a]
 --- /home/ubuntu/tt-metal/python_env/bin/python3(_PyFunction_Vectorcall+0x393) [0x5d5a23]
 --- /home/ubuntu/tt-metal/python_env/bin/python3(PyObject_Call+0x62) [0x5d4c12]
 --- /home/ubuntu/tt-metal/python_env/bin/python3(_PyEval_EvalFrameDefault+0x1f26) [0x548a66]
 --- /home/ubuntu/tt-metal/python_env/bin/python3(_PyEval_EvalCodeWithName+0x26a) [0x54552a]
 --- /home/ubuntu/tt-metal/python_env/bin/python3(_PyFunction_Vectorcall+0x393) [0x5d5a23]
 --- /home/ubuntu/tt-metal/python_env/bin/python3(PyObject_Call+0x62) [0x5d4c12]
 --- /home/ubuntu/tt-metal/python_env/bin/python3(_PyEval_EvalFrameDefault+0x1f26) [0x548a66]
 --- /home/ubuntu/tt-metal/python_env/bin/python3(_PyEval_EvalCodeWithName+0x26a) [0x54552a]
 --- /home/ubuntu/tt-metal/python_env/bin/python3(_PyFunction_Vectorcall+0x393) [0x5d5a23]
 --- /home/ubuntu/tt-metal/python_env/bin/python3() [0x57a4af]
 --- /home/ubuntu/tt-metal/python_env/bin/python3(_PyObject_MakeTpCall+0x296) [0x5d6066]
 --- /home/ubuntu/tt-metal/python_env/bin/python3(_PyEval_EvalFrameDefault+0x690a) [0x54d44a]
 --- /home/ubuntu/tt-metal/python_env/bin/python3(_PyEval_EvalCodeWithName+0x26a) [0x54552a]
 --- /home/ubuntu/tt-metal/python_env/bin/python3(_PyFunction_Vectorcall+0x393) [0x5d5a23]
 --- /home/ubuntu/tt-metal/python_env/bin/python3(PyObject_Call+0x62) [0x5d4c12]
 --- /home/ubuntu/tt-metal/python_env/bin/python3(_PyEval_EvalFrameDefault+0x1f26) [0x548a66]
 --- /home/ubuntu/tt-metal/python_env/bin/python3(_PyEval_EvalCodeWithName+0x26a) [0x54552a]
 --- /home/ubuntu/tt-metal/python_env/bin/python3(_PyFunction_Vectorcall+0x393) [0x5d5a23]
 --- /home/ubuntu/tt-metal/python_env/bin/python3(PyObject_Call+0x62) [0x5d4c12]
 --- /home/ubuntu/tt-metal/python_env/bin/python3(_PyEval_EvalFrameDefault+0x1f26) [0x548a66]
 --- /home/ubuntu/tt-metal/python_env/bin/python3(_PyEval_EvalCodeWithName+0x26a) [0x54552a]
 --- /home/ubuntu/tt-metal/python_env/bin/python3(_PyFunction_Vectorcall+0x393) [0x5d5a23]
 --- /home/ubuntu/tt-metal/python_env/bin/python3(PyObject_Call+0x62) [0x5d4c12]
 --- /home/ubuntu/tt-metal/python_env/bin/python3(_PyEval_EvalFrameDefault+0x1f26) [0x548a66]
 --- /home/ubuntu/tt-metal/python_env/bin/python3(_PyEval_EvalCodeWithName+0x26a) [0x54552a]
 --- /home/ubuntu/tt-metal/python_env/bin/python3(_PyFunction_Vectorcall+0x393) [0x5d5a23]
 --- /home/ubuntu/tt-metal/python_env/bin/python3(_PyEval_EvalFrameDefault+0x725) [0x547265]
 --- /home/ubuntu/tt-metal/python_env/bin/python3(_PyEval_EvalCodeWithName+0x26a) [0x54552a]
 --- /home/ubuntu/tt-metal/python_env/bin/python3(_PyFunction_Vectorcall+0x393) [0x5d5a23]
 --- /home/ubuntu/tt-metal/python_env/bin/python3(_PyEval_EvalFrameDefault+0x725) [0x547265]
 --- /home/ubuntu/tt-metal/python_env/bin/python3(_PyEval_EvalCodeWithName+0x26a) [0x54552a]
 --- /home/ubuntu/tt-metal/python_env/bin/python3(_PyFunction_Vectorcall+0x393) [0x5d5a23]
 --- /home/ubuntu/tt-metal/python_env/bin/python3(_PyEval_EvalFrameDefault+0x1876) [0x5483b6]
 --- /home/ubuntu/tt-metal/python_env/bin/python3(_PyEval_EvalCodeWithName+0x26a) [0x54552a]
 --- /home/ubuntu/tt-metal/python_env/bin/python3(_PyFunction_Vectorcall+0x393) [0x5d5a23]
 --- /home/ubuntu/tt-metal/python_env/bin/python3(PyObject_Call+0x62) [0x5d4c12]
 --- /home/ubuntu/tt-metal/python_env/bin/python3(_PyEval_EvalFrameDefault+0x1f26) [0x548a66]
 --- /home/ubuntu/tt-metal/python_env/bin/python3(_PyEval_EvalCodeWithName+0x26a) [0x54552a]
 --- /home/ubuntu/tt-metal/python_env/bin/python3(_PyFunction_Vectorcall+0x393) [0x5d5a23]
 --- /home/ubuntu/tt-metal/python_env/bin/python3(PyObject_Call+0x62) [0x5d4c12]
 --- /home/ubuntu/tt-metal/python_env/bin/python3(_PyEval_EvalFrameDefault+0x1f26) [0x548a66]
 --- /home/ubuntu/tt-metal/python_env/bin/python3(_PyEval_EvalCodeWithName+0x26a) [0x54552a]
 --- /home/ubuntu/tt-metal/python_env/bin/python3(_PyFunction_Vectorcall+0x393) [0x5d5a23]
 --- /home/ubuntu/tt-metal/python_env/bin/python3(_PyEval_EvalFrameDefault+0x5d69) [0x54c8a9]
 --- /home/ubuntu/tt-metal/python_env/bin/python3() [0x4e1a5e]
 --- /home/ubuntu/tt-metal/python_env/bin/python3(_PyEval_EvalFrameDefault+0x5d69) [0x54c8a9]
 --- /home/ubuntu/tt-metal/python_env/bin/python3(_PyEval_EvalCodeWithName+0x26a) [0x54552a]
 --- /home/ubuntu/tt-metal/python_env/bin/python3(_PyFunction_Vectorcall+0x393) [0x5d5a23]
 --- /home/ubuntu/tt-metal/python_env/bin/python3() [0x57a4af]
 --- /home/ubuntu/tt-metal/python_env/bin/python3(_PyObject_MakeTpCall+0x296) [0x5d6066]
 --- /home/ubuntu/tt-metal/python_env/bin/python3(_PyEval_EvalFrameDefault+0x690a) [0x54d44a]
 --- /home/ubuntu/tt-metal/python_env/bin/python3(_PyFunction_Vectorcall+0x1b6) [0x5d5846]
 --- /home/ubuntu/tt-metal/python_env/bin/python3(_PyEval_EvalFrameDefault+0x907) [0x547447]
 --- /home/ubuntu/tt-metal/python_env/bin/python3(_PyFunction_Vectorcall+0x1b6) [0x5d5846]
 --- /home/ubuntu/tt-metal/python_env/bin/python3(PyObject_Call+0x62) [0x5d4c12]
 --- /home/ubuntu/tt-metal/python_env/bin/python3(_PyEval_EvalFrameDefault+0x1f26) [0x548a66]
 --- /home/ubuntu/tt-metal/python_env/bin/python3(_PyEval_EvalCodeWithName+0x26a) [0x54552a]
 --- /home/ubuntu/tt-metal/python_env/bin/python3(_PyFunction_Vectorcall+0x393) [0x5d5a23]
 --- /home/ubuntu/tt-metal/python_env/bin/python3(_PyEval_EvalFrameDefault+0x5d69) [0x54c8a9]
 --- /home/ubuntu/tt-metal/python_env/bin/python3() [0x4e1a5e]
 --- /home/ubuntu/tt-metal/python_env/bin/python3(_PyEval_EvalFrameDefault+0x5d69) [0x54c8a9]
 --- /home/ubuntu/tt-metal/python_env/bin/python3(_PyEval_EvalCodeWithName+0x26a) [0x54552a]
 --- /home/ubuntu/tt-metal/python_env/bin/python3(_PyFunction_Vectorcall+0x393) [0x5d5a23]
 --- /home/ubuntu/tt-metal/python_env/bin/python3() [0x57a4af]
 --- /home/ubuntu/tt-metal/python_env/bin/python3(PyObject_Call+0x25e) [0x5d4e0e]
 --- /home/ubuntu/tt-metal/python_env/bin/python3(_PyEval_EvalFrameDefault+0x1f26) [0x548a66]
 --- /home/ubuntu/tt-metal/python_env/bin/python3(_PyEval_EvalCodeWithName+0x26a) [0x54552a]
 --- /home/ubuntu/tt-metal/python_env/bin/python3(_PyFunction_Vectorcall+0x393) [0x5d5a23]
 --- /home/ubuntu/tt-metal/python_env/bin/python3(_PyEval_EvalFrameDefault+0x725) [0x547265]

Getting Additional info for the operation under test and its behavior
To get additional information and results for different combinations of input shapes, types, layouts and memory configs for which this operation was tested you can also run locally sweeps and check the results. To do this you should:

1. Run non working sweep by using pytest tests/ttnn/python_api_testing/sweep_tests/run_sweep_test.py --input-path tests/ttnn/python_api_testing/sweep_tests/test_configs/ci_sweep_tests_broken/grayskull/ttnn_rmsnorm_test.yaml --input-method cli --cli-input results_ttnn_rmsnorm_broken
2. After the run is completed all test sweeps results should be available inside specified output directory (in this case ./results_ttnn_rmsnorm_broken).