How to query for memory descriptors for weights_layer_md and weights_iter_md in RNN LSTM from primitive descriptor

[lacak-sk]

Hello there are examples in C++ that query for memory descriptors for weights_layer and weights_iter in RNN LSTM primitive, that are:

• lstm_pd.weights_desc()
• lstm_pd.weights_iter_desc()
  (https://uxlfoundation.github.io/oneDNN/page_lstm_example_cpp.html#doxid-lstm-example-cpp)

But how to do it in plain C implementation? There is C API function: dnnl_primitive_desc_query_md(primitive_desc, what, index), where argument "what" is of dnnl_query_t type.
But dnnl_query_t type has only some options there, one of them is dnnl_query_weights_md but I do not see there another option which will distinguish weights_layer and weights_iter.

Can I use?:

• dnnl_primitive_desc_query_md(primitive_desc, dnnl_query_exec_arg_md, DNNL_ARG_WEIGHTS_LAYER)
• dnnl_primitive_desc_query_md(primitive_desc, dnnl_query_exec_arg_md, DNNL_ARG_WEIGHTS_ITER)
  or:
• dnnl_primitive_desc_query_md(primitive_desc, dnnl_query_weights_md, DNNL_ARG_WEIGHTS_LAYER)
• dnnl_primitive_desc_query_md(primitive_desc, dnnl_query_weights_md, DNNL_ARG_WEIGHTS_ITER)

And will be memory descriptors for:

• DNNL_ARG_WEIGHTS_LAYER and DNNL_ARG_DIFF_WEIGHTS_LAYER the same?
• DNNL_ARG_WEIGHTS_ITER and DNNL_ARG_DIFF_WEIGHTS_ITER the same?

Thanks

[shu1chen]

Hello @lacak-sk, great questions!

But how to do it in plain C implementation?

You're right that dnnl_query_weights_md is a generic query and doesn't distinguish between weights_layer and weights_iter. It's typically used for primitives with a single weights argument (e.g., convolution).
For RNN primitives like LSTM, which have multiple weights tensors, you should use dnnl_query_exec_arg_md with the appropriate argument index:

const_dnnl_memory_desc_t weights_layer_md = 
    dnnl_primitive_desc_query_md(lstm_pd, dnnl_query_exec_arg_md, DNNL_ARG_WEIGHTS_LAYER);

const_dnnl_memory_desc_t weights_iter_md = 
    dnnl_primitive_desc_query_md(lstm_pd, dnnl_query_exec_arg_md, DNNL_ARG_WEIGHTS_ITER);

And will be memory descriptors for:
DNNL_ARG_WEIGHTS_LAYER and DNNL_ARG_DIFF_WEIGHTS_LAYER the same?
DNNL_ARG_WEIGHTS_ITER and DNNL_ARG_DIFF_WEIGHTS_ITER the same?

In general, yes, they often have the same shape and layout. However, this is not guaranteed. Forward weights may use data types like f32, bf16, or quantized formats, while gradient (diff) weights typically use accumulation-friendly types like f32 for numerical stability.
So, it's best to query them separately to ensure correctness:

const_dnnl_memory_desc_t diff_weights_layer_md = 
    dnnl_primitive_desc_query_md(lstm_bwd_pd, dnnl_query_exec_arg_md, DNNL_ARG_DIFF_WEIGHTS_LAYER);
    
const_dnnl_memory_desc_t diff_weights_iter_md = 
    dnnl_primitive_desc_query_md(lstm_bwd_pd, dnnl_query_exec_arg_md, DNNL_ARG_DIFF_WEIGHTS_ITER);