Skip to content

Commit 8d083ed

Browse files
mcremon-metafacebook-github-bot
mcremon-meta
authored and
facebook-github-bot
committed
Introduce strongly typed quant/dequant ops (#14268)
Summary: Pull Request resolved: #14268 As titled. Reviewed By: hsharma35 Differential Revision: D82183474
1 parent c67afa3 commit 8d083ed

12 files changed

+582
-20
lines changed

‎backends/cadence/aot/functions.yaml‎

Lines changed: 48 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -184,12 +184,60 @@
184184
- arg_meta: null
185185
kernel_name: impl::generic::quantize_per_tensor_out
186186

187+
- func: cadence::quantize_per_tensor_asym8s.out(Tensor input, float scale, int zero_point, int quant_min, int quant_max, ScalarType dtype, *, Tensor(a!) out) -> Tensor(a!)
188+
variants: function
189+
kernels:
190+
- arg_meta: null
191+
kernel_name: impl::generic::quantize_per_tensor_asym8s_out
192+
193+
- func: cadence::quantize_per_tensor_asym8u.out(Tensor input, float scale, int zero_point, int quant_min, int quant_max, ScalarType dtype, *, Tensor(a!) out) -> Tensor(a!)
194+
variants: function
195+
kernels:
196+
- arg_meta: null
197+
kernel_name: impl::generic::quantize_per_tensor_asym8u_out
198+
199+
- func: cadence::quantize_per_tensor_asym16s.out(Tensor input, float scale, int zero_point, int quant_min, int quant_max, ScalarType dtype, *, Tensor(a!) out) -> Tensor(a!)
200+
variants: function
201+
kernels:
202+
- arg_meta: null
203+
kernel_name: impl::generic::quantize_per_tensor_asym16s_out
204+
205+
- func: cadence::quantize_per_tensor_asym16u.out(Tensor input, float scale, int zero_point, int quant_min, int quant_max, ScalarType dtype, *, Tensor(a!) out) -> Tensor(a!)
206+
variants: function
207+
kernels:
208+
- arg_meta: null
209+
kernel_name: impl::generic::quantize_per_tensor_asym16u_out
210+
187211
- func: cadence::dequantize_per_tensor.out(Tensor input, float scale, int zero_point, int quant_min, int quant_max, ScalarType dtype, *, Tensor(a!) out) -> Tensor(a!)
188212
variants: function
189213
kernels:
190214
- arg_meta: null
191215
kernel_name: impl::generic::dequantize_per_tensor_out
192216

217+
- func: cadence::dequantize_per_tensor_asym8s.out(Tensor input, float scale, int zero_point, int quant_min, int quant_max, ScalarType dtype, *, Tensor(a!) out) -> Tensor(a!)
218+
variants: function
219+
kernels:
220+
- arg_meta: null
221+
kernel_name: impl::generic::dequantize_per_tensor_asym8s_out
222+
223+
- func: cadence::dequantize_per_tensor_asym8u.out(Tensor input, float scale, int zero_point, int quant_min, int quant_max, ScalarType dtype, *, Tensor(a!) out) -> Tensor(a!)
224+
variants: function
225+
kernels:
226+
- arg_meta: null
227+
kernel_name: impl::generic::dequantize_per_tensor_asym8u_out
228+
229+
- func: cadence::dequantize_per_tensor_asym16s.out(Tensor input, float scale, int zero_point, int quant_min, int quant_max, ScalarType dtype, *, Tensor(a!) out) -> Tensor(a!)
230+
variants: function
231+
kernels:
232+
- arg_meta: null
233+
kernel_name: impl::generic::dequantize_per_tensor_asym16s_out
234+
235+
- func: cadence::dequantize_per_tensor_asym16u.out(Tensor input, float scale, int zero_point, int quant_min, int quant_max, ScalarType dtype, *, Tensor(a!) out) -> Tensor(a!)
236+
variants: function
237+
kernels:
238+
- arg_meta: null
239+
kernel_name: impl::generic::dequantize_per_tensor_asym16u_out
240+
193241
- func: cadence::quantized_conv_nchw.out(Tensor input, Tensor weight, Tensor bias, int[] stride, SymInt[] padding, int[] dilation, int groups, int input_zero_point, Tensor weight_zero_point, Tensor bias_scale, float out_scale, int out_zero_point, Tensor out_multiplier, Tensor out_shift, *, Tensor(a!) out) -> Tensor(a!)
194242
kernels:
195243
- arg_meta: null

‎backends/cadence/aot/functions_hifi.yaml‎

Lines changed: 49 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -284,12 +284,61 @@
284284
- arg_meta: null
285285
kernel_name: impl::HiFi::quantize_per_tensor_out
286286

287+
- func: cadence::quantize_per_tensor_asym8s.out(Tensor input, float scale, int zero_point, int quant_min, int quant_max, ScalarType dtype, *, Tensor(a!) out) -> Tensor(a!)
288+
variants: function
289+
kernels:
290+
- arg_meta: null
291+
kernel_name: impl::HiFi::quantize_per_tensor_asym8s_out
292+
293+
- func: cadence::quantize_per_tensor_asym8u.out(Tensor input, float scale, int zero_point, int quant_min, int quant_max, ScalarType dtype, *, Tensor(a!) out) -> Tensor(a!)
294+
variants: function
295+
kernels:
296+
- arg_meta: null
297+
kernel_name: impl::HiFi::quantize_per_tensor_asym8u_out
298+
299+
- func: cadence::quantize_per_tensor_asym16s.out(Tensor input, float scale, int zero_point, int quant_min, int quant_max, ScalarType dtype, *, Tensor(a!) out) -> Tensor(a!)
300+
variants: function
301+
kernels:
302+
- arg_meta: null
303+
kernel_name: impl::HiFi::quantize_per_tensor_asym16s_out
304+
305+
- func: cadence::quantize_per_tensor_asym16u.out(Tensor input, float scale, int zero_point, int quant_min, int quant_max, ScalarType dtype, *, Tensor(a!) out) -> Tensor(a!)
306+
variants: function
307+
kernels:
308+
- arg_meta: null
309+
kernel_name: impl::HiFi::quantize_per_tensor_asym16s_out
310+
311+
287312
- func: cadence::dequantize_per_tensor.out(Tensor input, float scale, int zero_point, int quant_min, int quant_max, ScalarType dtype, *, Tensor(a!) out) -> Tensor(a!)
288313
variants: function
289314
kernels:
290315
- arg_meta: null
291316
kernel_name: impl::HiFi::dequantize_per_tensor_out
292317

318+
- func: cadence::dequantize_per_tensor_asym8s.out(Tensor input, float scale, int zero_point, int quant_min, int quant_max, ScalarType dtype, *, Tensor(a!) out) -> Tensor(a!)
319+
variants: function
320+
kernels:
321+
- arg_meta: null
322+
kernel_name: impl::HiFi::dequantize_per_tensor_asym8s_out
323+
324+
- func: cadence::dequantize_per_tensor_asym8u.out(Tensor input, float scale, int zero_point, int quant_min, int quant_max, ScalarType dtype, *, Tensor(a!) out) -> Tensor(a!)
325+
variants: function
326+
kernels:
327+
- arg_meta: null
328+
kernel_name: impl::HiFi::dequantize_per_tensor_asym8u_out
329+
330+
- func: cadence::dequantize_per_tensor_asym16s.out(Tensor input, float scale, int zero_point, int quant_min, int quant_max, ScalarType dtype, *, Tensor(a!) out) -> Tensor(a!)
331+
variants: function
332+
kernels:
333+
- arg_meta: null
334+
kernel_name: impl::HiFi::dequantize_per_tensor_asym16s_out
335+
336+
- func: cadence::dequantize_per_tensor_asym16u.out(Tensor input, float scale, int zero_point, int quant_min, int quant_max, ScalarType dtype, *, Tensor(a!) out) -> Tensor(a!)
337+
variants: function
338+
kernels:
339+
- arg_meta: null
340+
kernel_name: impl::HiFi::dequantize_per_tensor_asym16u_out
341+
293342
- func: cadence::quantized_conv_nchw.out(Tensor input, Tensor weight, Tensor bias, int[] stride, SymInt[] padding, int[] dilation, int groups, int input_zero_point, Tensor weight_zero_point, Tensor bias_scale, float out_scale, int out_zero_point, Tensor out_multiplier, Tensor out_shift, *, Tensor(a!) out) -> Tensor(a!)
294343
kernels:
295344
- arg_meta: null

‎backends/cadence/aot/ops_registrations.py‎

Lines changed: 148 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -28,12 +28,64 @@
2828
"quantize_per_tensor.out(Tensor input, float scale, int zero_point, int quant_min, int quant_max, ScalarType dtype, *, Tensor(a!) out) -> Tensor(a!)"
2929
)
3030

31+
lib.define(
32+
"quantize_per_tensor_asym8s(Tensor input, float scale, int zero_point, int quant_min, int quant_max, ScalarType dtype) -> (Tensor Z)"
33+
)
34+
lib.define(
35+
"quantize_per_tensor_asym8s.out(Tensor input, float scale, int zero_point, int quant_min, int quant_max, ScalarType dtype, *, Tensor(a!) out) -> Tensor(a!)"
36+
)
37+
38+
lib.define(
39+
"quantize_per_tensor_asym8u(Tensor input, float scale, int zero_point, int quant_min, int quant_max, ScalarType dtype) -> (Tensor Z)"
40+
)
41+
lib.define(
42+
"quantize_per_tensor_asym8u.out(Tensor input, float scale, int zero_point, int quant_min, int quant_max, ScalarType dtype, *, Tensor(a!) out) -> Tensor(a!)"
43+
)
44+
45+
lib.define(
46+
"quantize_per_tensor_asym16s(Tensor input, float scale, int zero_point, int quant_min, int quant_max, ScalarType dtype) -> (Tensor Z)"
47+
)
48+
lib.define(
49+
"quantize_per_tensor_asym16s.out(Tensor input, float scale, int zero_point, int quant_min, int quant_max, ScalarType dtype, *, Tensor(a!) out) -> Tensor(a!)"
50+
)
51+
52+
lib.define(
53+
"quantize_per_tensor_asym16u(Tensor input, float scale, int zero_point, int quant_min, int quant_max, ScalarType dtype) -> (Tensor Z)"
54+
)
55+
lib.define(
56+
"quantize_per_tensor_asym16u.out(Tensor input, float scale, int zero_point, int quant_min, int quant_max, ScalarType dtype, *, Tensor(a!) out) -> Tensor(a!)"
57+
)
58+
3159
lib.define(
3260
"dequantize_per_tensor(Tensor input, float scale, int zero_point, int quant_min, int quant_max, ScalarType dtype) -> (Tensor Z)"
3361
)
3462
lib.define(
3563
"dequantize_per_tensor.out(Tensor input, float scale, int zero_point, int quant_min, int quant_max, ScalarType dtype, *, Tensor(a!) out) -> Tensor(a!)"
3664
)
65+
lib.define(
66+
"dequantize_per_tensor_asym8s(Tensor input, float scale, int zero_point, int quant_min, int quant_max, ScalarType dtype) -> (Tensor Z)"
67+
)
68+
lib.define(
69+
"dequantize_per_tensor_asym8s.out(Tensor input, float scale, int zero_point, int quant_min, int quant_max, ScalarType dtype, *, Tensor(a!) out) -> Tensor(a!)"
70+
)
71+
lib.define(
72+
"dequantize_per_tensor_asym8u(Tensor input, float scale, int zero_point, int quant_min, int quant_max, ScalarType dtype) -> (Tensor Z)"
73+
)
74+
lib.define(
75+
"dequantize_per_tensor_asym8u.out(Tensor input, float scale, int zero_point, int quant_min, int quant_max, ScalarType dtype, *, Tensor(a!) out) -> Tensor(a!)"
76+
)
77+
lib.define(
78+
"dequantize_per_tensor_asym16s(Tensor input, float scale, int zero_point, int quant_min, int quant_max, ScalarType dtype) -> (Tensor Z)"
79+
)
80+
lib.define(
81+
"dequantize_per_tensor_asym16s.out(Tensor input, float scale, int zero_point, int quant_min, int quant_max, ScalarType dtype, *, Tensor(a!) out) -> Tensor(a!)"
82+
)
83+
lib.define(
84+
"dequantize_per_tensor_asym16u(Tensor input, float scale, int zero_point, int quant_min, int quant_max, ScalarType dtype) -> (Tensor Z)"
85+
)
86+
lib.define(
87+
"dequantize_per_tensor_asym16u.out(Tensor input, float scale, int zero_point, int quant_min, int quant_max, ScalarType dtype, *, Tensor(a!) out) -> Tensor(a!)"
88+
)
3789

3890
lib.define(
3991
"quantized_layer_norm(Tensor X, Tensor X_scale, Tensor X_zero_point, int[] normalized_shape, Tensor weight, Tensor bias, float eps, float output_scale, int output_zero_point) -> (Tensor Y)"
@@ -554,6 +606,54 @@ def quantize_per_tensor_meta(
554606
return input.new_empty(input.size(), dtype=dtype)
555607

556608

609+
@register_fake("cadence::quantize_per_tensor_asym8s")
610+
def quantize_per_tensor_asym8s_meta(
611+
input: torch.Tensor,
612+
scale: float,
613+
zero_point: int,
614+
quant_min: int,
615+
quant_max: int,
616+
dtype: torch.dtype,
617+
) -> torch.Tensor:
618+
return input.new_empty(input.size(), dtype=dtype)
619+
620+
621+
@register_fake("cadence::quantize_per_tensor_asym8u")
622+
def quantize_per_tensor_asym8u_meta(
623+
input: torch.Tensor,
624+
scale: float,
625+
zero_point: int,
626+
quant_min: int,
627+
quant_max: int,
628+
dtype: torch.dtype,
629+
) -> torch.Tensor:
630+
return input.new_empty(input.size(), dtype=dtype)
631+
632+
633+
@register_fake("cadence::quantize_per_tensor_asym16s")
634+
def quantize_per_tensor_asym16s_meta(
635+
input: torch.Tensor,
636+
scale: float,
637+
zero_point: int,
638+
quant_min: int,
639+
quant_max: int,
640+
dtype: torch.dtype,
641+
) -> torch.Tensor:
642+
return input.new_empty(input.size(), dtype=dtype)
643+
644+
645+
@register_fake("cadence::quantize_per_tensor_asym16u")
646+
def quantize_per_tensor_asym16u_meta(
647+
input: torch.Tensor,
648+
scale: float,
649+
zero_point: int,
650+
quant_min: int,
651+
quant_max: int,
652+
dtype: torch.dtype,
653+
) -> torch.Tensor:
654+
return input.new_empty(input.size(), dtype=dtype)
655+
656+
557657
@register_fake("cadence::dequantize_per_tensor")
558658
def dequantize_per_tensor_meta(
559659
input: torch.Tensor,
@@ -566,6 +666,54 @@ def dequantize_per_tensor_meta(
566666
return input.new_empty(input.size(), dtype=torch.float)
567667

568668

669+
@register_fake("cadence::dequantize_per_tensor_asym8s")
670+
def dequantize_per_tensor_asym8s_meta(
671+
input: torch.Tensor,
672+
scale: float,
673+
zero_point: int,
674+
quant_min: int,
675+
quant_max: int,
676+
dtype: torch.dtype,
677+
) -> torch.Tensor:
678+
return input.new_empty(input.size(), dtype=torch.float)
679+
680+
681+
@register_fake("cadence::dequantize_per_tensor_asym8u")
682+
def dequantize_per_tensor_asym8u_meta(
683+
input: torch.Tensor,
684+
scale: float,
685+
zero_point: int,
686+
quant_min: int,
687+
quant_max: int,
688+
dtype: torch.dtype,
689+
) -> torch.Tensor:
690+
return input.new_empty(input.size(), dtype=torch.float)
691+
692+
693+
@register_fake("cadence::dequantize_per_tensor_asym16s")
694+
def dequantize_per_tensor_asym16s_meta(
695+
input: torch.Tensor,
696+
scale: float,
697+
zero_point: int,
698+
quant_min: int,
699+
quant_max: int,
700+
dtype: torch.dtype,
701+
) -> torch.Tensor:
702+
return input.new_empty(input.size(), dtype=torch.float)
703+
704+
705+
@register_fake("cadence::dequantize_per_tensor_asym16u")
706+
def dequantize_per_tensor_asym16u_meta(
707+
input: torch.Tensor,
708+
scale: float,
709+
zero_point: int,
710+
quant_min: int,
711+
quant_max: int,
712+
dtype: torch.dtype,
713+
) -> torch.Tensor:
714+
return input.new_empty(input.size(), dtype=torch.float)
715+
716+
569717
@register_fake("cadence::quantized_add")
570718
def quantized_add_meta(
571719
X: torch.Tensor,

‎backends/cadence/aot/type_dispatch.py‎

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,7 @@ class OpConfig:
2727
base_name: str
2828
type_dispatch_suffixes: dict[tuple[torch.dtype, ...], str]
2929
weight_arg_idx: Optional[int] = None
30+
is_quant_op: bool = False
3031
variant: str = "per_tensor"
3132

3233

@@ -100,6 +101,27 @@ class CompileTimeTypeDispatchPass(ExportPass):
100101
},
101102
variant="default",
102103
),
104+
exir_ops.edge.cadence.quantize_per_tensor.default: OpConfig(
105+
"quantize_per_tensor",
106+
type_dispatch_suffixes={
107+
(torch.int8,): "asym8s",
108+
(torch.uint8,): "asym8u",
109+
(torch.int16,): "asym16s",
110+
(torch.uint16,): "asym16s",
111+
},
112+
variant="default",
113+
is_quant_op=True,
114+
),
115+
exir_ops.edge.cadence.dequantize_per_tensor.default: OpConfig(
116+
"dequantize_per_tensor",
117+
type_dispatch_suffixes={
118+
(torch.int8,): "asym8s",
119+
(torch.uint8,): "asym8u",
120+
(torch.int16,): "asym16s",
121+
(torch.uint16,): "asym16s",
122+
},
123+
variant="default",
124+
),
103125
}
104126

105127
def call_operator(
@@ -120,6 +142,8 @@ def call_operator(
120142
if config.weight_arg_idx is not None:
121143
weight_dtype = args[config.weight_arg_idx].to_tensor().dtype
122144
dtype_key = (input_dtype, weight_dtype)
145+
elif config.is_quant_op:
146+
dtype_key = (args[5],)
123147
else:
124148
dtype_key = (input_dtype,)
125149

0 commit comments

Comments
 (0)