Adding a MatMulInteger-like contrib/custom op

[tarekziade]

Hello,
Disclaimer: I am very new to the code base, so I don't understand everything.

We've recently added Transformers.js + onnxruntime (WASM) into Firefox as our inference runtime, and it works great, but we want to improve its speed in particular when running against quantized modules using int8 because it's a good tradeoff in terms of download size for our uses cases. Right now it's slower than fp32 and way slower than the native runtime.

In Firefox Translations, which uses another WASM runtime (Bergamot), we were able to get a significant performance boost by plugging in the WASM a JS import hook to call https://github.com/mozilla/gemmology that is vendored into Firefox.

My goal is to try the same thing with onnxruntime, by writing a new contrib op that will replace MatMulInteger in a WASM environment to call back our gemmology lib inside Firefox and see how that can help.

The first step I took was to take a model that we use, and rewrite its graph. See for instance our fine-tuned mobilebert here https://huggingface.co/Mozilla/mobilebert-uncased-finetuned-LoRA-intent-classifier/tree/main/onnx

My script to rewrite is here https://gist.github.com/tarekziade/08404a890a3cf56878042892f6ce58aa
and the resulting graph looks ok -- I am not 100% sure I did the right thing there but the result in Netron looks correct.

From there I added a first version of custom op here that mimics MatMulInteger for now tarekziade@bcfe7cb (planning to do the WASM hook as a second step)

The op seems to be correctly registered in the ops list, but I get the following error

>>> import onnxruntime as ort; sess_opt = ort.SessionOptions() ; sess_opt.log_severity_level = 0 ; ort_sess = ort.InferenceSession('model_quantized.firefox.onnx', sess_opt);
onnxruntime.capi.onnxruntime_pybind11_state.NotImplemented: [ONNXRuntimeError] : 9 : NOT_IMPLEMENTED : Could not find an implementation for FirefoxMatMulInteger8(1) node with name '/distilbert/transformer/layer.0/attention/q_lin/MatMul_quant

Suggesting that the schema part is not well implemented. I noticed that the call to node.GetExecutionProviderType() in the loop that looks at the graph in session_state.cc, returns an empty value, suggesting that my contrib op is not registered for CPUExecutionProvider.

But this is what I have done by adding

class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, FirefoxMatMulInteger8);

in onnxruntime/contrib_ops/cpu/cpu_contrib_kernels.cc

I am not entirely sure how to solve this, any help would be appreciated, thanks!

[fs-eire]

you need to also add it by appending a new line of BuildKernelCreateInfo in function RegisterCpuContribKernels in onnxruntime/contrib_ops/cpu/cpu_contrib_kernels.cc.

[tarekziade]

Thanks so much, that worked as expected, I was able to replace MatMulInteger with my own class and the rewritten graph works.

I pushed a demo here https://github.com/tarekziade/onnx-custom-op and will now see how to plug Javascript as an extern.

@fs-eire is this something that you have already done in onnxruntime ?

I am planning to expose the JS API as an extern "C" and then call that function from the ops.

[tarekziade]

I want to send the matmul operation to our C++ backend using a JS Call like this:

#ifdef __EMSCRIPTEN__
#include <emscripten/emscripten.h>
// Declare the JavaScript function as an external import
extern "C" {
    // Will be linked to the JavaScript function
   void MatMulOperation(
    CBLAS_TRANSPOSE TransA,
    CBLAS_TRANSPOSE TransB,
    size_t BatchSize,
    MLAS_THREADPOOL* ThreadPool);
}
#else
// Fallback local C++ implementation of the function
void MatMulOperation(
    CBLAS_TRANSPOSE TransA,
    CBLAS_TRANSPOSE TransB,
    size_t BatchSize,
    MLAS_THREADPOOL* ThreadPool) {
    MlasGemmBatch(TransA, TransA, batch_size, ctx->GetOperatorThreadPool());
}
#endif

[fs-eire]

I think you can use the EM_ASM macro to put inline JavaScript code in C++ source code. Not sure if that works.

[tarekziade]

The problem with EM_ASM is that it adds some overhead.

I was able to plug everything together but it's hackish, see tarekziade@a7c1d2b#diff-8e46f75c9f18f62a5a6bdc2f790da0e090b89185ac0c3aa55fba7f7ab76567bb

I declare in the C++ side of things the function with import_name:

extern "C" void __attribute__((import_name("int8MultiplyAndAddBias")))
  int8MultiplyAndAddBias(
      size_t ARowSize, 
      size_t BColSize, 
      size_t AColSizeBRowSize, 
      bool AIsSigned, 
      bool BIsSigned, 
      MLAS_GEMM_QUANT_DATA_PARAMS* gemm_data
  );
#endif

But one problem I faced with this was that the cmake builder will fail because it can't find that export.

So, by adding: "SHELL:-s ERROR_ON_UNDEFINED_SYMBOLS=0" it compiles. But in that case, cmake add its own stub function:

/** @type {function(...*):?} */ function _int8MultiplyAndAddBias() {
 abort("missing function: int8MultiplyAndAddBias");
}

_int8MultiplyAndAddBias.stub = true;

I could not find a way to tell cmake not to generate this stub. If it was not around, as long as _int8MultiplyAndAddBias is declared in our JS environment, things will work.

As a workaround, my patch extends pre.js and implements a custom Module["instantiateWasm"] in order to inject into the imports my function.

Do you know if there's a way to avoid the stub to be generated? or to make cmake happy about not having _int8MultiplyAndAddBias around?

Thanks

[tarekziade]

I've resolved the issue by moving my specific imports to a separate import module. In that case cmake does not generate the stub, but I still need to override Module["instantiateWasm"] to do this

[tarekziade]

https://github.com/tarekziade/onnxruntime/blob/0104484d18dae22fb9ff316a086a0634472711ef/onnxruntime/wasm/pre.js#L87-L134