[Relay] QLinearMatMul allows 1D weight_scale, weight_zero_point inputs

[yuanfz98]

Hello,

This PR is to fix #9908. As _qnn.op.dense can already accept 1D vectors (

tvm/python/tvm/relay/qnn/op/qnn.py

Line 552 in 6eb4ed8

and optionally the kernel_zero_point will be 1-D vectors instead of scalars.

), we need to enable this feature in QLinearMatMul.

In #9908, the expected behaviour is tested by following code:

import onnx
import torch
from tvm import relay
import tvm 
import onnxruntime
 
from tvm.testing import assert_allclose

onnx_model = onnx.load("dense_h32_w32_c3_sNone_pNone_kNone.pt_quant.onnx")
onnx.checker.check_model(onnx_model)

ort_session = onnxruntime.InferenceSession("dense_h32_w32_c3_sNone_pNone_kNone.pt_quant.onnx")
x = torch.randn(1, 3, 32, 32, requires_grad=True) 

def to_numpy(tensor):
    return tensor.detach().cpu().numpy() if tensor.requires_grad else tensor.cpu().numpy()
 
# compute ONNX Runtime output prediction
ort_inputs = {
    ort_session.get_inputs()[0].name: to_numpy(x)
    }
ort_outs = ort_session.run(None, ort_inputs)

model, params = relay.frontend.from_onnx(onnx_model)

with tvm.transform.PassContext(opt_level=1):
    executor = relay.build_module.create_executor(
        "graph", model, tvm.cpu(0), "llvm", params=params
    ).evaluate()

tvm_output = executor(tvm.nd.array(x.detach())).numpy()

assert_allclose(tvm_output, ort_outs[0])
print(tvm_output, ort_outs[0])

Output:

[[-0.539178   -0.3405335  -0.04540446  0.488098   -0.5505291  -0.54485357
  -0.56188023  0.22134677  0.09648448  0.19864453 -0.27810234 -0.13621339
  -0.3235068   0.18161786 -0.5278269   0.46539575 -0.5278269  -0.4483691
  -0.09648448  0.01702667]] [[-0.539178   -0.3405335  -0.04540446  0.488098   -0.5505291  -0.54485357
  -0.56188023  0.22134677  0.09648448  0.19864453 -0.27810234 -0.13621339
  -0.3235068   0.18161786 -0.5278269   0.46539575 -0.5278269  -0.4483691
  -0.09648448  0.01702667]]

Thank reviewers for your time and have a nice day !

cc @alnah005 @cconvey

[AndrewZhaoLuo]

Thanks for the contribution, LGTM mostly, just one set of comments comment

[AndrewZhaoLuo]

In this codepath, you will return None

[AndrewZhaoLuo]

This changes the behavior of ensure_scalar_shape -- now we can return non-scalar elements, I would instead in QLinearMatMul's try_resolve_to_const simply, wrap ensure_scalar_shape calls with a check if the number of elements is 1.

[AndrewZhaoLuo]

echoing above, instead of adding a new flag to ensure_scalar_shape that breaks invariants, take the check you wrote and push it out here.

[AndrewZhaoLuo]

Also, for each file, we use black as our python formatter. Make sure to run black on all python files or else you will fail our linter.