Custom Torch Op to Custom QONNX op export

[jurevreca12]

Hello,
I am trying to figure out how to add a custom op in torch and use the expor_qonnx function convert it to a custom QONNX op (CustomOp class). I defined the torch custom op as such:

import torch
import math

@torch.library.custom_op("chisel4ml::fft", mutates_args=())
def fft_c4ml(x: torch.Tensor, win_fn: torch.Tensor, fft_size: int, num_frames: int) -> torch.Tensor:
    return torch.fft.fft(x * win_fn,
                         norm="backward"
    ).real

# Use register_fake to add a ``FakeTensor`` kernel for the operator
@fft_c4ml.register_fake
def _(x, win_fn, fft_size, num_frames):
    assert len(x.shape) == 3
    batch_x, num_frames_x, frame_length_x = x.shape
    assert num_frames_x == num_frames
    assert frame_length_x == fft_size
    assert len(win_fn).shape == 1
    assert len(win_fn) == fft_size
    output_length = int(2**(math.log2(fft_size) + 3))
    return torch.empty((batch_x, num_frames, output_length))

The custom op I defined as such (template):

class FFT_c4ml(CustomOp):

    def get_nodeattr_types(self):
        return {
            "fft_size": (128,),
            "num_frames": (8,),
            "win_fn": (np.array([]),),
        }

    def make_shape_compatible_op(self, model):
        # For Quant the output shape should be the same as the input shape.
        # Get the output shape from the input
        out_shape = model.get_tensor_shape(self.onnx_node.input[0])

        # implement tensor with correct shape
        values = np.random.randn(*out_shape).astype(np.float32)
        return helper.make_node(
            "Constant",
            inputs=[],
            outputs=[self.onnx_node.output[0]],
            value=helper.make_tensor(
                name="const_tensor",
                data_type=TensorProto.FLOAT,
                dims=values.shape,
                vals=values.flatten().astype(float),
            ),
            name=self.onnx_node.name,
        )

    def get_integer_datatype(self, model):
        return DataType["FLOAT32"]

    def get_scaled_integer_datatype(self, model):
        return DataType["FLOAT32"]

    def get_output_dtype(self, model):
        return DataType["FLOAT32"]

    def infer_node_datatype(self, model):
        node = self.onnx_node
        model.set_tensor_datatype(node.output[0], DataType["FLOAT32"])


    def execute_node(self, context, graph):
        node = self.onnx_node
        inp_tensor = context[node.input[0]]
        win_fn = self.get_nodeattr('win_fn')
        results = []
        for x in inp_tensor:
            res = np.fft.fft(x * win_fn, 
                             norm="backward", 
                             axis=-1).real
            results.append(np.expand_dims(res, axis=-1).astype(np.float32))
        ret = np.reshape(
            np.array(results),
            [len(inp_tensor), self.cfg.num_frames, self.cfg.fft_size, 1],
        )
        context[node.output[0]] = ret


    def verify_node(self):
        pass

I feel like I am pretty close to success, I just need to "tie" these things together. I am searching through the code and I see that ONNXBaseManager class has a register_custom_fns class, however the "Quant" ops seem to use a class based on torch.autograf.Function (brevitas/src/brevitas/export/onnx/qonnx/function.py).

Does anyone have experience with this? What do I need to do here?

[jurevreca12]

I managed to make it work. See QONNX PR: fastmachinelearning/qonnx#144