(
self,
node: torch.fx.Node,
xnn_graph: XNNGraph,
vals_to_ids: Dict[torch.fx.Node, int],
debug_handle: int,
)
| 31 | super().__init__(*args) |
| 32 | |
| 33 | def define_node( |
| 34 | self, |
| 35 | node: torch.fx.Node, |
| 36 | xnn_graph: XNNGraph, |
| 37 | vals_to_ids: Dict[torch.fx.Node, int], |
| 38 | debug_handle: int, |
| 39 | ) -> None: |
| 40 | |
| 41 | # input |
| 42 | input_node = get_input_node(node, 0) |
| 43 | input_quant_params = QuantParams.from_inputs(input_node, self._exported_program) |
| 44 | self.define_tensor( |
| 45 | input_node, |
| 46 | xnn_graph, |
| 47 | vals_to_ids, |
| 48 | quant_params=input_quant_params, |
| 49 | ) |
| 50 | input_id = vals_to_ids[input_node] |
| 51 | |
| 52 | # filter |
| 53 | weight_node = get_input_node(node, 1) |
| 54 | weight_quant_params = QuantParams.from_weights( |
| 55 | weight_node, self._exported_program |
| 56 | ) |
| 57 | self.define_tensor( |
| 58 | weight_node, |
| 59 | xnn_graph, |
| 60 | vals_to_ids, |
| 61 | quant_params=weight_quant_params, |
| 62 | ) |
| 63 | filter_id = vals_to_ids[weight_node] |
| 64 | |
| 65 | # bias |
| 66 | if len(node.args) > 2: |
| 67 | bias_node = get_input_node(node, 2) |
| 68 | bias_quant_params = QuantParams.from_bias( |
| 69 | bias_node, weight_quant_params, input_quant_params |
| 70 | ) |
| 71 | # For dynamic quantization, there are no kernels with fp16 bias |
| 72 | # So we need to force the fp16 bias to fp32 |
| 73 | force_fp32 = False |
| 74 | if input_quant_params is not None and input_quant_params.is_dynamic: |
| 75 | force_fp32 = True |
| 76 | |
| 77 | self.define_tensor( |
| 78 | get_input_node(node, 2), |
| 79 | xnn_graph, |
| 80 | vals_to_ids, |
| 81 | quant_params=bias_quant_params, |
| 82 | force_fp32=force_fp32, |
| 83 | ) |
| 84 | bias_id = vals_to_ids[bias_node] |
| 85 | else: |
| 86 | bias_id = XNN_INVALID_VALUE_ID |
| 87 | |
| 88 | # output |
| 89 | output_min_max = FuseActivationPass.get_fused_activation(node) |
| 90 | output_quant_params = QuantParams.from_outputs(node) |
nothing calls this directly
no test coverage detected