TypeError: Cannot update immutable parameter Z_DRAM

[dinghongsong]

Describe the bug

When I run the codes in contributed/matmul.py, I encounter the following error:

TypeError: Cannot update immutable parameter Z_DRAM

This is due to using Z as a parameter to matmul, which is decorated with nki.baremetal().

Is there any way to fix this error?

Expected Behavior

The Optimized matrix multiplication kernel works correctly.

Current Behavior

Cannot update immutable parameter Z_DRAM
when invoke
nl.store(Z_DRAM[m_start:m_end, n_start:n_end], value=Z_SBUF[m1])

Reproduction Steps

python contributed/matmul.py

Regression Issue

• Select this option if this issue appears to be a regression.

Possible Solution

No response

Additional Information/Context

No response

neuronx-cc version used

NeuronX Compiler version 2.16.372.0+4a9b2326

Framework(s) and their versions used (JAX, PyTorch, etc..)

Python version 3.10.12 HWM version 2.16.0.372+4a9b2326 NumPy version 1.25.2

[aws-liuayu]

Thanks @dinghongsong for filing the issue. We will take a look and get back to you.

[aws-zhehongb]

Looks like we need to change the function singature:

@nki.jit
def nki_matmul_basic_(lhsT, rhs):

we are not passing the result as parameter anymore. Check here

[ggumen]

you can read about that error message and how to fix it here: https://awsdocs-neuron.readthedocs-hosted.com/en/latest/general/nki/api/nki.errors.html#err-cannot-update-immutable-parameter

[dinghongsong]

Hi @ggumen @aws-zhehongb @aws-liuayu,

Thank you for your response. This issue has been resolved.
However, I noticed that the implementation of NKI matrix multiplication is still slower than torch.matmul(A.T, B).
The related discussion link is as follows.

#45

I'm not sure if I'm missing something.