Adding custom quantized C++ op

username: rvijayc
cooked: <p>Hi,</p>
<p>I am following the example in the following file that creates “reference quantized representations” of different corresponding floating point ops.</p><aside class="onebox githubblob" data-onebox-src="https://github.com/pytorch/pytorch/blob/main/torch/ao/quantization/pt2e/representation/rewrite.py">
  <header class="source">

      <a href="https://github.com/pytorch/pytorch/blob/main/torch/ao/quantization/pt2e/representation/rewrite.py" target="_blank" rel="noopener nofollow ugc">github.com</a>
  </header>

  <article class="onebox-body">
    <h4><a href="https://github.com/pytorch/pytorch/blob/main/torch/ao/quantization/pt2e/representation/rewrite.py" target="_blank" rel="noopener nofollow ugc">pytorch/pytorch/blob/main/torch/ao/quantization/pt2e/representation/rewrite.py</a></h4>


      <pre><code class="lang-py">import torch
from torch.fx import GraphModule
from ..utils import (
    get_aten_graph_module,
    remove_tensor_overload_for_qdq_ops,
    _replace_literals_with_new_placeholders,
    _replace_literals_with_existing_placeholders,
)
from torch.ao.quantization.fx._decomposed import quantized_decomposed_lib  # noqa: F401
from torch.fx.subgraph_rewriter import replace_pattern
from torch._higher_order_ops.out_dtype import out_dtype
from typing import Optional, Callable, Tuple, Any
from dataclasses import dataclass

from functools import partial

__all__ = [
    "reference_representation_rewrite",
]

</code></pre>



  This file has been truncated. <a href="https://github.com/pytorch/pytorch/blob/main/torch/ao/quantization/pt2e/representation/rewrite.py" target="_blank" rel="noopener nofollow ugc">show original</a>

  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>

<p>From what I understand, it seems like it is trying to match “dq → float_op → q” and replace it with some reference quantized op.</p>
<p>What I would like to do is to provide my own C++ implementation for the following function in <code>rewrite.py</code>.</p>
<pre><code class="lang-python">def _reference_quantized_linear(
    x_i8, x_scale, x_zero_point, x_quant_min, x_quant_max,
    weight_i8, weight_scale, weight_zero_point, weight_quant_min, weight_quant_max,
    bias_fp32,
    out_scale, out_zero_point, out_quant_min, out_quant_max
):
</code></pre>
<p>The question I grappling with is how to setup the signature of the function. There seems to two choices:</p>
<p>Use <code>torch::Tensor</code> for all arguments (including scalar arguments like scales and zero points). For example:</p>
<pre><code class="lang-cpp">    m.def("fc_relu_s8u8s8(Tensor x_i8, Tensor x_scale, Tensor xzp, Tensor xmin, Tensor xmax, Tensor w_i8, Tensor w_scale, Tensor wzp, Tensor wmin, Tensor wmax, Tensor bfp32, Tensor y_scale, Tensor yzp, Tensor ymin, Tensor ymax) -&gt; Tensor");
</code></pre>
<p>If I use the above representation, I can replace the pattern <code>_qdq_quantized_linear</code> with <code>fc_relu_s8u8s8</code> in the graph, but when I run the graph, it fails because it cannot cast scalar values (like <code>x_scale</code>, <code>x_zp</code> etc., to Tensors when calling <code>fc_relu_s8u8s8</code>).</p>
<p>Or, I could use scalar arguments for scales and zero points as below:</p>
<pre><code class="lang-cpp">    m.def("fc_prelu_s8u8s8_raw(Tensor x_i8, float xs, int xzp, int xmin, int xmax, Tensor w_i8, double ws, int wzp, int wmin,int wmax, Tensor bfp32, doublw ys, int yzp, int ymin, int ymax) -&gt; Tensor");
</code></pre>
<p>If I use this approach the Subgraph rewriter complains about “dead code” (which I guess is related to the scalar arguments).</p>
<p>Is there a way around this, or is there some reference code available that shows how to do such a thing? I am stuck between a rock and a hard place.</p>
<p>Appreciate your help or insights on this. Please let me know if I can provide more information. I am using the recent torch from github.</p>
<pre><code class="lang-auto">In [5]: torch.version.__version__
Out[5]: '2.2.0a0+gitb4ce501'
</code></pre>
<p>Thanks,<br>
Vijay.</p>

username: rvijayc
cooked: <p>OK. Found the problem. I defined my op using the actual types for each parameter:</p>
<pre><code class="lang-cpp">m.def("fc_prelu_s8u8s8_raw(Tensor x_i8, float xs, int xzp, int xmin, int xmax, Tensor w_i8, double ws, int wzp, int wmin,int wmax, Tensor bfp32, double ys, int yzp, int ymin, int ymax) -&gt; Tensor");
</code></pre>
<p>Then, I followed the instructions in this document <a href="https://docs.google.com/document/d/1_W62p8WJOQQUzPsJYa7s701JXt0qf2OfLub2sbkHOaU/edit#heading=h.3rgxk3v387wl" rel="noopener nofollow ugc">The C++ Custom Operators Manual - Google Docs</a> to bind the OP to torch ( as <code>torch.op.&lt;my_lib&gt;.fc_prelu_s8u8s8</code> ). This doc should really be a part of the official PyTorch docs!</p>
<p>Then, instead of <code>get_aten_graph_module</code> which internally calls <code>capture_pre_autograd_graph</code> (which I am guessing is equivalent to <code>torch.export.export</code> that exports a FX graph by passing real inputs to the module, I used <code>torch.fx.symbolic_trace</code> to convert the wrapper function that wraps <code>torch.ops.&lt;my_lib&gt;.fc_prelu_s8u8s8</code>) into a FX graph. With this approach, I don’t see any warning about literals or dead code and I can able to rewrite the <code>quant -&gt; linear -&gt; dequant</code> combination with my new fixed point Op.</p>
<pre><code class="lang-python">        # replacement = get_aten_graph_module(replacement, example_inputs)  # type: ignore[arg-type, assignment]
        replacement = torch.fx.symbolic_trace(replacement)
</code></pre>