Defines an tensor value into the XNNGraph Args: tensor: EdgeIR Tensor that is being defined into xnn_graph xnn_graph: XNNGraph object for serializing into flatbuffer vals_to_ids: dictionary mapping edge_graph values(node targets) to
( # noqa: C901
self,
tensor: torch.fx.Node,
xnn_graph: XNNGraph,
vals_to_ids: Dict[torch.fx.Node, int],
convert_to_nhwc: bool = False,
swap_in_out_for_weights: bool = False,
quant_params: Optional[QuantParams] = None,
force_fp32: bool = False,
groups: int = 1,
)
| 374 | assert quant_params.is_qc4w, "Only 4b group quantization is supported" |
| 375 | |
| 376 | def define_tensor( # noqa: C901 |
| 377 | self, |
| 378 | tensor: torch.fx.Node, |
| 379 | xnn_graph: XNNGraph, |
| 380 | vals_to_ids: Dict[torch.fx.Node, int], |
| 381 | convert_to_nhwc: bool = False, |
| 382 | swap_in_out_for_weights: bool = False, |
| 383 | quant_params: Optional[QuantParams] = None, |
| 384 | force_fp32: bool = False, |
| 385 | groups: int = 1, |
| 386 | ) -> None: |
| 387 | """ |
| 388 | Defines an tensor value into the XNNGraph |
| 389 | |
| 390 | Args: |
| 391 | tensor: EdgeIR Tensor that is being defined into xnn_graph |
| 392 | xnn_graph: XNNGraph object for serializing into flatbuffer |
| 393 | vals_to_ids: dictionary mapping edge_graph values(node targets) to |
| 394 | their corresponding ids in XNNGraph |
| 395 | convert_to_nhwc: bool to indicate whether tensor shape should be permuted to |
| 396 | reflect the nhwc memory format. |
| 397 | swap_in_out_for_weights: bool to indicate whether tensor shape should be |
| 398 | permuted and reshape from (inc, oc/groups, height, width) to (oc, inc/groups, height, width) |
| 399 | , which should be used for depthwise/transpose convolution |
| 400 | weights. This is only valid for tensors which hold |
| 401 | constant data. If used along with convert_to_nhwc, this |
| 402 | swap will happen before converting to nhwc. |
| 403 | quant_params: Quantization meta data for this tensor, None if it is not quantized |
| 404 | force_fp32: forces tensor to be serialize as fp32, used for bias of dynamically quantized ops |
| 405 | groups: number of groups for swap_in_out_for_weights |
| 406 | """ |
| 407 | |
| 408 | assert ( |
| 409 | swap_in_out_for_weights or groups == 1 |
| 410 | ), "groups is option for swap_in_out_for_weights" |
| 411 | |
| 412 | if tensor in vals_to_ids: |
| 413 | return |
| 414 | |
| 415 | if quant_params is not None: |
| 416 | if quant_params.q_input in vals_to_ids: |
| 417 | vals_to_ids[tensor] = vals_to_ids[quant_params.q_input] |
| 418 | return |
| 419 | # Tag added by ChannelsLastTaggedReshapePass |
| 420 | convert_to_nhwc |= tensor.meta.get( |
| 421 | ChannelsLastTaggedReshapePass.XNN_NHWC_NODE, False |
| 422 | ) |
| 423 | |
| 424 | # Get new xnn id for tensor value |
| 425 | ext_id, id_out, flag = self.gen_ids_and_flags(tensor, xnn_graph, quant_params) |
| 426 | dims = get_shape(tensor) |
| 427 | dims = [1] if len(dims) == 0 else dims |
| 428 | |
| 429 | # check for per_channel_group quantization |
| 430 | if quant_params and quant_params.per_channel_group: |
| 431 | self._check_per_channel_group_params(quant_params, dims) |
| 432 | |
| 433 | # constant values serialize data |
no test coverage detected