ttnn::mean failing thru TTRT but not in EmitC

[svuckovicTT]

For some reason, a simple ttnn::mean test is failing thru TTRT but not thru EmitC path.

Test:

// RUN: ttmlir-opt --ttir-to-ttnn-backend-pipeline="system-desc-path=%system_desc_path%" %s > %t.mlir
// RUN: ttmlir-translate --ttnn-to-flatbuffer %t.mlir > %t.ttnn

func.func @forward(%arg0: tensor<512x1024xbf16>) -> tensor<512x32xbf16> {
  %0 = tensor.empty() : tensor<512x32xbf16>
  %1 = "ttir.mean"(%arg0, %0) <{dim_arg = [-1: i32], keep_dim = true}> : (tensor<512x1024xbf16>, tensor<512x32xbf16>) -> tensor<512x32xbf16>
  return %1 : tensor<512x32xbf16>
}

Repro:

git checkout svuckovic/repro-1942
# build...
llvm-lit -sv test/ttmlir/Silicon/TTNN/repro_mean.mlir
ttrt run build/test/ttmlir/Silicon/TTNN/Output/repro_mean.mlir.tmp.ttnn

My assumption is it has to do with the added memory config in runtime/lib/ttnn/operations/reduction/reduction.cpp::runReductionOp.

[kmabeeTT]

Thanks Sasha, do you have the signature handy for how this test was failing for you?

[svuckovicTT]

Knew I forgot something haha!

Yes, this is the error msg:

                 Always |    FATAL | Input output tensor size mismatch in memcpy: 512 * 2 != 16384 * 2
LOG_ASSERT @ /localdev/svuckovic/_workspace/repos/tt-mlir/runtime/lib/ttnn/runtime.cpp:370: srcTensor.volume() * srcTensor.element_size() == dstTensor.volume() * dstTensor.element_size()
backtrace:
 --- tt::runtime::ttnn::memcpy(tt::runtime::Tensor, tt::runtime::Tensor)
 --- tt::runtime::memcpy(tt::runtime::Tensor, tt::runtime::Tensor)
 --- /opt/ttmlir-toolchain/venv/lib/python3.10/site-packages/ttrt/runtime/_C.cpython-310-x86_64-linux-gnu.so(+0xfdbfe) [0x7f0ad37dcbfe]
 --- /opt/ttmlir-toolchain/venv/lib/python3.10/site-packages/ttrt/runtime/_C.cpython-310-x86_64-linux-gnu.so(+0x11cdfa) [0x7f0ad37fbdfa]
 --- /opt/ttmlir-toolchain/venv/bin/python(+0x15cb2e) [0x5603d815ab2e]
 --- /opt/ttmlir-toolchain/venv/bin/python(_PyObject_MakeTpCall+0x25b) [0x5603d81512db]
 --- /opt/ttmlir-toolchain/venv/bin/python(_PyEval_EvalFrameDefault+0x6b17) [0x5603d8149d27]
 --- /opt/ttmlir-toolchain/venv/bin/python(_PyFunction_Vectorcall+0x7c) [0x5603d815b42c]
 --- /opt/ttmlir-toolchain/venv/bin/python(_PyEval_EvalFrameDefault+0x6bc) [0x5603d81438cc]
 --- /opt/ttmlir-toolchain/venv/bin/python(_PyFunction_Vectorcall+0x7c) [0x5603d815b42c]
 --- /opt/ttmlir-toolchain/venv/bin/python(_PyEval_EvalFrameDefault+0x8ab) [0x5603d8143abb]
 --- /opt/ttmlir-toolchain/venv/bin/python(_PyObject_FastCallDictTstate+0xc4) [0x5603d8150474]
 --- /opt/ttmlir-toolchain/venv/bin/python(_PyObject_Call_Prepend+0x5c) [0x5603d81662bc]
 --- /opt/ttmlir-toolchain/venv/bin/python(+0x2826d0) [0x5603d82806d0]
 --- /opt/ttmlir-toolchain/venv/bin/python(_PyObject_MakeTpCall+0x25b) [0x5603d81512db]
 --- /opt/ttmlir-toolchain/venv/bin/python(_PyEval_EvalFrameDefault+0x64d6) [0x5603d81496e6]
 --- /opt/ttmlir-toolchain/venv/bin/python(_PyFunction_Vectorcall+0x7c) [0x5603d815b42c]
 --- /opt/ttmlir-toolchain/venv/bin/python(_PyEval_EvalFrameDefault+0x6bc) [0x5603d81438cc]
 --- /opt/ttmlir-toolchain/venv/bin/python(+0x142016) [0x5603d8140016]
 --- /opt/ttmlir-toolchain/venv/bin/python(PyEval_EvalCode+0x86) [0x5603d82358b6]
 --- /opt/ttmlir-toolchain/venv/bin/python(+0x262918) [0x5603d8260918]
 --- /opt/ttmlir-toolchain/venv/bin/python(+0x25c1db) [0x5603d825a1db]
 --- /opt/ttmlir-toolchain/venv/bin/python(+0x262665) [0x5603d8260665]
 --- /opt/ttmlir-toolchain/venv/bin/python(_PyRun_SimpleFileObject+0x1a8) [0x5603d825fb48]
 --- /opt/ttmlir-toolchain/venv/bin/python(_PyRun_AnyFileObject+0x43) [0x5603d825f793]
 --- /opt/ttmlir-toolchain/venv/bin/python(Py_RunMain+0x2be) [0x5603d82522ce]
 --- /opt/ttmlir-toolchain/venv/bin/python(Py_BytesMain+0x2d) [0x5603d822870d]
 --- /lib/x86_64-linux-gnu/libc.so.6(+0x29d90) [0x7f0b162bad90]
 --- /lib/x86_64-linux-gnu/libc.so.6(__libc_start_main+0x80) [0x7f0b162bae40]
 --- /opt/ttmlir-toolchain/venv/bin/python(_start+0x25) [0x5603d8228605]

[svuckovicTT]

The diff of 512 vs 16384 is 32, so I'm assuming it's got to do with row_major vs tiled.