( # noqa: C901
self,
inputs_repsets: TensorRepSetList,
outputs_repsets: TensorRepSetList,
op_node: torch.fx.Node,
texture_limits: ImageExtents,
)
| 1286 | """ |
| 1287 | |
| 1288 | def __init__( # noqa: C901 |
| 1289 | self, |
| 1290 | inputs_repsets: TensorRepSetList, |
| 1291 | outputs_repsets: TensorRepSetList, |
| 1292 | op_node: torch.fx.Node, |
| 1293 | texture_limits: ImageExtents, |
| 1294 | ): |
| 1295 | self.op_node = op_node |
| 1296 | |
| 1297 | # inputs_repset_list is received from the operator registration. If a different |
| 1298 | # repset is defined for each input tensor, then assume that the input tensor |
| 1299 | # representations do not need to be synchronized. |
| 1300 | if len(inputs_repsets) > 1: |
| 1301 | self.sync_args_repr = False |
| 1302 | # Otherwise, default to True |
| 1303 | else: |
| 1304 | self.sync_args_repr = True |
| 1305 | |
| 1306 | # outputs_repset_list is received from the operator registration. If a different |
| 1307 | # repset is defined for each output tensor, then assume that the output tensor |
| 1308 | # representations do not need to be synchronized. |
| 1309 | if len(outputs_repsets) > 1: |
| 1310 | self.sync_outs_repr = False |
| 1311 | else: |
| 1312 | self.sync_outs_repr = True |
| 1313 | |
| 1314 | # Try to determine the index of the "primary" argument, i.e. the first non |
| 1315 | # constant tensor argument. For the vast majority of operators with tensor |
| 1316 | # arguments, this will be the first argument. |
| 1317 | self.primary_arg_idx: Optional[int] = None |
| 1318 | for i, arg_node in enumerate(self.op_node.args): |
| 1319 | arg_node_repset = inputs_repsets[i] |
| 1320 | if not is_tensor_arg_node(arg_node): |
| 1321 | continue |
| 1322 | if arg_node_repset is None: |
| 1323 | continue |
| 1324 | if arg_node_repset.is_empty(): |
| 1325 | continue |
| 1326 | |
| 1327 | self.primary_arg_idx = i |
| 1328 | break |
| 1329 | |
| 1330 | # If the repset of the primary input and the primary output are the same, then |
| 1331 | # assume they need to be the same. |
| 1332 | self.sync_primary_io_repr = self.primary_arg_idx is not None |
| 1333 | if self.primary_arg_idx is not None: |
| 1334 | if inputs_repsets[self.primary_arg_idx] != outputs_repsets[0]: |
| 1335 | self.sync_primary_io_repr = False |
| 1336 | |
| 1337 | # Now, go through the arguments of the operator and create a filtered repset |
| 1338 | # for each based on the actual tensor value. |
| 1339 | args_repset_list = TensorRepSetList([]) |
| 1340 | common_arg_repset = ANY_STORAGE_INCL_PACKED_INT8 |
| 1341 | for i, arg_node in enumerate(op_node.args): |
| 1342 | arg_repset = inputs_repsets[i] |
| 1343 | |
| 1344 | # Use ANY_STORAGE_INCL_PACKED_INT8 for non-tensor nodes so they don't cause the op |
| 1345 | # repsets to appear empty |
nothing calls this directly
no test coverage detected