Quantization Lowering Support

[WilliamZhang20]

In my file here, I train a model using Quantization-Aware-Training (QAT).

I then export MLIR at this file in TOSA. Note that inside that file, the attempt to export in TOSA fails as there are no TOSA ops, we still have torch.aten ops.

However, when running:

torch-mlir-opt qat_model_torch_dialect.mlir     --torch-unpack-quant-tensor     --convert-torch-to-tosa     --canonicalize     --symbol-dce     -o final_tosa
_model.mlir

I get:

<unknown>:0: error: unable to schedule pass 'UnpackQuantTensor' on a PassManager intended to run on 'builtin.module'!

Any reason for this?
I am also unsure how to use the newest TorchAO pt2e QAT with MLIR integration.

[sahas3]

I added the following after https://github.com/WilliamZhang20/ECE298A-TPU/blob/main/test/tpu/train_qat_model.py#L90

    quantized_model = torch.export.export(quantized_model, example_inputs)

    module = fx.export_and_import(
        quantized_model , output_type="linalg_on_tensors",
    )

and was able to get MLIR in linalg+other dialect. Trying to get TOSA fails with failed to legalize operation 'torch.aten._make_per_tensor_quantized_tensor' that was explicitly marked illegal. @sjarus / @justin-ngo-arm can probably confirm but I think quantized models are not well supported in the TOSA path.

Regarding the torch-unpack-quant-tensor pass it was added quite sometime back in https://github.com/llvm/torch-mlir/pull/2402/files -- it doesn't appear to be part of any pass pipelines that are still operational now in the new torch.export path. Looking at the changes, it seems to provide a way to lower quant.matmul_rhs_group_quant -- I am not sure if this op is even produced in the PT2 export QAT path anymore. I at least didn't see this op in the IR for the model you shared.