[BE]: Apply ruff PERF403 to use dict comprehensions more often (pytorch#149257)

Fixes #ISSUE_NUMBER

Pull Request resolved: pytorch#149257
Approved by: https://github.com/jansel

Author: Skylion007

.github/scripts/trymerge.py

@@ -819,10 +819,9 @@ def _get_reviews(self) -> list[tuple[str, str]]:
                     cursor=info["reviews"]["pageInfo"]["startCursor"],
                 )
                 info = rc["data"]["repository"]["pullRequest"]
-        reviews = {}
-        for author, state in self._reviews:
-            if state != "COMMENTED":
-                reviews[author] = state
+        reviews = {
+            author: state for author, state in self._reviews if state != "COMMENTED"
+        }
         return list(reviews.items())
 
     def get_approved_by(self) -> list[str]:

benchmarks/operator_benchmark/benchmark_core.py

@@ -296,8 +296,7 @@ def split(s):
             (key.strip(), value.strip())
             for key, value in map(lambda str: str.split(":"), key_vals)  # noqa: C417
         ]  # ['M: (32, 16)', 'ZPB: 2'] -> [('M', '(32, 16)'), ('ZPB', '2')]
-        for key, value in key_vals:
-            out[key] = value
+        out.update(key_vals)
 
         return out
 

pyproject.toml

@@ -72,7 +72,6 @@ ignore = [
     # these ignores are from ruff PERF; please fix!
     "PERF203",
     "PERF401",
-    "PERF403",
     # these ignores are from PYI; please fix!
     "PYI024",
     "PYI036",

test/inductor/test_custom_post_grad_passes.py

@@ -219,9 +219,7 @@ def merge_mm_shared_rhs(graph: fx.Graph):
             for m in matmuls:
                 rhs_vals[m.args[1]].add(m)
 
-            order = {}
-            for idx, n in enumerate(graph.nodes):
-                order[n] = idx
+            order = {n: idx for idx, n in enumerate(graph.nodes)}
 
             for rhs, matmuls in rhs_vals.items():
                 if len(matmuls) == 1:

test/test_fx.py

@@ -2324,9 +2324,7 @@ def test_deepcopy_recursion_depth(self):
 
         copied_graph = copy.deepcopy(g)
 
-        val_map = {}
-        for orig_node, new_node in zip(g.nodes, copied_graph.nodes):
-            val_map[orig_node] = new_node
+        val_map = dict(zip(g.nodes, copied_graph.nodes))
 
         for orig_node, new_node in zip(g.nodes, copied_graph.nodes):
             orig_users = set(orig_node.users.keys())

test/test_jit.py

@@ -1761,8 +1761,7 @@ def doit(x, y):
         for node in g.nodes():
             n_ = g2.createClone(node, lambda x: g_to_g2[x])
             g2.appendNode(n_)
-            for o, no in zip(node.outputs(), n_.outputs()):
-                g_to_g2[o] = no
+            g_to_g2.update(zip(node.outputs(), n_.outputs()))
 
         for node in g.outputs():
             g2.registerOutput(g_to_g2[node])

test/test_linalg.py

@@ -91,9 +91,7 @@ def tunableop_matmul(device, dtype):
 
 def get_tunableop_validators():
     assert len(torch.cuda.tunable.get_validators()) > 0
-    validators = {}
-    for key, value in torch.cuda.tunable.get_validators():
-        validators[key] = value
+    validators = dict(torch.cuda.tunable.get_validators())
     return validators
 
 class TestLinalg(TestCase):

torch/_export/serde/serialize.py

@@ -1564,9 +1564,7 @@ def serialize(self, exported_program: ep.ExportedProgram) -> _SerializedProgram:
         # TODO: Directly serialize exported_program.constants once
         # CustomClassHolders get stored in the ExportedProgram rather than in
         # the graph
-        constants: dict[str, Any] = {}
-        for n, c in gm_serializer.custom_objs.items():
-            constants[n] = c
+        constants: dict[str, Any] = gm_serializer.custom_objs.copy()
         for n, t in exported_program.constants.items():
             assert n not in constants
             constants[n] = t

torch/_functorch/partitioners.py

@@ -580,9 +580,7 @@ def reordering_to_mimic_autograd_engine(gm: fx.GraphModule) -> fx.GraphModule:
     for node in gm.graph.find_nodes(op="placeholder"):
         env[node] = new_graph.node_copy(node, lambda x: env[x])
 
-    order = {}
-    for idx, node in enumerate(gm.graph.nodes):
-        order[node] = idx
+    order = {node: idx for idx, node in enumerate(gm.graph.nodes)}
 
     def insert_node_in_graph(node):
         cur_nodes = [node]

torch/_functorch/top_operators_github_usage.py

@@ -625,8 +625,5 @@ def get_nn_functional_top_list():
     return top_nn_functional_
 
 
-usage_count = {}
-for k, v in get_nn_functional_top_list():
-    usage_count[k] = v
-for k, v in top_torch:
-    usage_count[k] = v
+usage_count = dict(get_nn_functional_top_list())
+usage_count.update(top_torch)

torch/_inductor/autotune_process.py

@@ -398,12 +398,9 @@ def benchmark(
         assert self.processes is not None, "Tuning process pool is not initialized"
         assert self.executor is not None
 
-        results = {}
-
         # Use a ThreadExecutorPool to spread the work across the subprocesses and
         # to grab subprocesses as soon as they're free.
-        for choice, result in zip(choices, self.executor.map(self.target, choices)):
-            results[choice] = result
+        results = dict(zip(choices, self.executor.map(self.target, choices)))
 
         return results
 

torch/_inductor/memory.py

@@ -267,9 +267,9 @@ class BufferInfo:
 
     # get the execution step of each node, this will be used to determine
     # the end_step of buffers
-    node_to_step: dict[BaseSchedulerNode, int] = dict()
-    for step, node in enumerate(nodes):
-        node_to_step[node] = step
+    node_to_step: dict[BaseSchedulerNode, int] = {
+        node: step for step, node in enumerate(nodes)
+    }
 
     # get buffers' size and liveliness information
     buf_info_list: list[BufferInfo] = []

torch/_inductor/runtime/triton_heuristics.py

@@ -154,8 +154,7 @@ def _dump_launch_params(args, kwargs, launcher, kernel_name, grid):
         else:
             call_kwargs[k] = v
     if not triton_version_uses_attrs_dict():
-        for k, v in launcher.config.kwargs.items():
-            call_kwargs[k] = v
+        call_kwargs.update(launcher.config.kwargs)
     call_kwargs["num_warps"] = launcher.config.num_warps
     call_kwargs["num_stages"] = launcher.config.num_stages
     args_str = [*call_args]

torch/_inductor/utils.py

@@ -340,8 +340,7 @@ def _type_of(key: Optional[torch.dtype]) -> str:
         "uint64": "u64",
     }
     # reinterpret can create triton type
-    for v in list(tys.values()):
-        tys[v] = v
+    tys.update({v: v for v in list(tys.values())})
     return key if isinstance(key, str) else f"*{tys[dtype_str]}"
 
 
@@ -635,9 +634,7 @@ def get_kernel_metadata(
             single_graph = inductor_nodes[0].graph
             # create a map of idx -> node and cache it
             if not hasattr(single_graph, "_inductor_kernel_metadata_node_to_idx_map"):
-                node_to_idx_map = {}
-                for idx, n in enumerate(single_graph.nodes):
-                    node_to_idx_map[n] = idx
+                node_to_idx_map = {n: idx for idx, n in enumerate(single_graph.nodes)}
                 single_graph._inductor_kernel_metadata_node_to_idx_map = node_to_idx_map  # type: ignore[attr-defined]
             inductor_nodes.sort(
                 key=lambda n: single_graph._inductor_kernel_metadata_node_to_idx_map[n]  # type: ignore[attr-defined]

torch/ao/ns/_numeric_suite.py

@@ -368,9 +368,7 @@ def prepare_model_with_stubs(
         "quantization_api._numeric_suite.prepare_model_with_stubs"
     )
 
-    float_module_children = {}
-    for name, mod in float_module.named_children():
-        float_module_children[name] = mod
+    float_module_children = dict(float_module.named_children())
 
     reassign = {}
     for name, mod in q_module.named_children():

torch/ao/quantization/_correct_bias.py

@@ -119,10 +119,11 @@ def bias_correction(
         float_model, quantized_model, _supported_modules, MeanShadowLogger
     )
 
-    uncorrected_modules = {}
-    for name, submodule in quantized_model.named_modules():
-        if type(submodule) in target_modules:
-            uncorrected_modules[name] = submodule
+    uncorrected_modules = {
+        name: submodule
+        for name, submodule in quantized_model.named_modules()
+        if type(submodule) in target_modules
+    }
 
     for uncorrected_module in uncorrected_modules:
         quantized_submodule = get_module(quantized_model, uncorrected_module)

torch/ao/quantization/fx/qconfig_mapping_utils.py

@@ -376,10 +376,8 @@ def _get_flattened_qconfig_dict(
     flattened: dict[Union[Callable, str], QConfigAny] = {
         "": qconfig_mapping.global_qconfig
     }
-    for obj, qconfig in qconfig_mapping.object_type_qconfigs.items():
-        flattened[obj] = qconfig
-    for obj, qconfig in qconfig_mapping.module_name_qconfigs.items():
-        flattened[obj] = qconfig
+    flattened.update(qconfig_mapping.object_type_qconfigs)
+    flattened.update(qconfig_mapping.module_name_qconfigs)  # type: ignore[arg-type]
     return flattened
 
 

torch/distributed/checkpoint/state_dict.py

@@ -596,8 +596,7 @@ def _load_model_state_dict(
             )
         elif info.full_state_dict:
             _distribute_state_dict(state_dict, local_state_dict, device=devices.pop())
-        for fqn, local_state in local_state_dict.items():
-            state_dict[fqn] = local_state
+        state_dict.update(local_state_dict)
 
     with info.fsdp_context():
         return cast(

torch/distributed/fsdp/_optim_utils.py

@@ -314,11 +314,9 @@ def _unflatten_communicated_optim_state(
             unflat_state_param[state_name] = optim_state
 
         # Add zero-dimension tensor state: take the target rank's value
-        for state_name, zero_dim_tensor in sorted_items(zero_dim_tensor_state):
-            unflat_state_param[state_name] = zero_dim_tensor
+        unflat_state_param.update(sorted_items(zero_dim_tensor_state))
         # Add non-tensor state: take the target rank's value
-        for state_name, non_tensor in sorted_items(non_tensor_state):
-            unflat_state_param[state_name] = non_tensor
+        unflat_state_param.update(sorted_items(non_tensor_state))
         unflat_param_state.append(unflat_state_param)
     return unflat_param_state
 
@@ -1827,11 +1825,12 @@ def _convert_state_with_flat_params(
             )
             if to_save:
                 assert len(unflat_state) == len(optim_state_key.unflat_param_names)
-                for unflat_param_name, unflat_param_state in zip(
-                    optim_state_key.unflat_param_names,
-                    unflat_state,
-                ):
-                    fsdp_osd_state[unflat_param_name] = unflat_param_state
+                fsdp_osd_state.update(
+                    zip(
+                        optim_state_key.unflat_param_names,
+                        unflat_state,
+                    )
+                )
         elif to_save:
             assert len(optim_state_key.unflat_param_names) == 1
             unflat_param_name = optim_state_key.unflat_param_names[0]

torch/distributed/tensor/_ops/_common_rules.py

@@ -265,7 +265,6 @@ def pointwise_rule(op_schema: OpSchema, linearity: bool = False) -> OutputShardi
     # check if we replace the all inputs dim char with singleton dimension,
     # if we replace all inputs, we also need to replace the output dimension.
     for output_dim_idx in range(len(out_dimchars)):
-        out_dimchar = out_dimchars[output_dim_idx]
         if singleton_counter[output_dim_idx] == len(input_specs):
             out_dimchars = _replace_char_in_str(out_dimchars, "1", output_dim_idx)
 
@@ -274,12 +273,10 @@ def pointwise_rule(op_schema: OpSchema, linearity: bool = False) -> OutputShardi
     enforce_sharding: dict[str, int] = {}
     if _is_inplace_op(op_schema.op):
         # inplace op should keep the input sharding it writes to
-        for out_dimchar, mesh_dim in zip(out_dimchars, input_specs[0].dim_map):
-            enforce_sharding[out_dimchar] = mesh_dim
+        enforce_sharding.update(zip(out_dimchars, input_specs[0].dim_map))
     elif _is_out_variant_op(op_schema.op):
         out_spec = cast(DTensorSpec, op_schema.kwargs_schema["out"])
-        for out_dimchar, mesh_dim in zip(out_dimchars, out_spec.dim_map):
-            enforce_sharding[out_dimchar] = mesh_dim
+        enforce_sharding.update(zip(out_dimchars, out_spec.dim_map))
 
     return einop_rule(
         fmt,

torch/onnx/_internal/_exporter_legacy.py

@@ -596,13 +596,14 @@ def export(self) -> _onnx_program.ONNXProgram:
             # not valid.
             # Concrete data is expected to be filled for those initializers later during `ONNXProgram.save`.
             if self.options.fake_context is not None:
-                initializers_with_real_tensors: dict[str, torch.Tensor] = {}
-                for (
-                    initializer_name,
-                    initializer,
-                ) in onnxscript_graph.initializers.items():
-                    if not isinstance(initializer, torch._subclasses.FakeTensor):
-                        initializers_with_real_tensors[initializer_name] = initializer
+                initializers_with_real_tensors: dict[str, torch.Tensor] = {
+                    initializer_name: initializer
+                    for (
+                        initializer_name,
+                        initializer,
+                    ) in onnxscript_graph.initializers.items()
+                    if not isinstance(initializer, torch._subclasses.FakeTensor)
+                }
                 onnxscript_graph.initializers = initializers_with_real_tensors
 
             # Export TorchScript graph to ONNX ModelProto.

torch/testing/_internal/common_cuda.py

@@ -217,8 +217,7 @@ def wrapper(f):
 
         @functools.wraps(f)
         def wrapped(*args, **kwargs):
-            for k, v in zip(arg_names, args):
-                kwargs[k] = v
+            kwargs.update(zip(arg_names, args))
             cond = torch.cuda.is_tf32_supported()
             if 'device' in kwargs:
                 cond = cond and (torch.device(kwargs['device']).type == 'cuda')

torch/testing/_internal/common_mkldnn.py

@@ -60,8 +60,7 @@ def wrapper(f):
 
         @functools.wraps(f)
         def wrapped(*args, **kwargs):
-            for k, v in zip(arg_names, args):
-                kwargs[k] = v
+            kwargs.update(zip(arg_names, args))
             cond = bf32_is_not_fp32()
             if "device" in kwargs:
                 cond = cond and (torch.device(kwargs["device"]).type == "cpu")