Convert directory fbcode/aitemplate to use the Ruff Formatter (#1030)

Summary:
Pull Request resolved: #1030

Converts the directory specified to use the Ruff formatter in pyfmt

ruff_dog

If this diff causes merge conflicts when rebasing, please run
`hg status -n -0 --change . -I '**/*.{py,pyi}' | xargs -0 arc pyfmt`
on your diff, and amend any changes before rebasing onto latest.
That should help reduce or eliminate any merge conflicts.

allow-large-files

Reviewed By: amyreese

Differential Revision: D64264218

examples/01_resnet-50/benchmark_ait.py

@@ -45,7 +45,6 @@ def mark_output(y):
 
 
 def compile_module(model_name, batch_size, **kwargs):
-
     if model_name != "resnet50":
         raise NotImplementedError
 

examples/01_resnet-50/weight_utils.py

@@ -17,7 +17,6 @@
 Only tested on resnet50
 """
 
-
 import pickle
 import re
 

examples/02_detectron2/demo.py

@@ -15,6 +15,7 @@
 """
 A main inference script for rcnn models
 """
+
 import glob
 import os
 

examples/02_detectron2/modeling/roi_heads/roi_heads.py

@@ -59,7 +59,6 @@ def get_shape(self, x):
         return shape
 
     def forward(self, features: Dict[str, Tensor], rois: Tensor, proposals: Tensor):
-
         box_features = [features[f] for f in self.in_features]
         roi_feat = self.box_head(box_features, rois)
         detections = self.box_predictor(roi_feat, proposals)

examples/02_detectron2/predictor/builtin_meta.py

@@ -25,7 +25,6 @@
 COCO model (with correct class names and colors).
 """
 
-
 # All coco categories, together with their nice-looking visualization colors
 # It's from https://github.com/cocodataset/panopticapi/blob/master/panoptic_coco_categories.json
 COCO_CATEGORIES = [

examples/04_vit/weight_utils.py

@@ -12,8 +12,8 @@
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
 #
-"""script for converting vit model from timm to ait
-"""
+"""script for converting vit model from timm to ait"""
+
 import pickle
 
 import click

examples/05_stable_diffusion/scripts/download_pipeline.py

@@ -36,7 +36,6 @@
     help="Pipeline files local directory.",
 )
 def download_pipeline_files(model_name, token, save_directory) -> None:
-
     StableDiffusionPipeline.from_pretrained(
         model_name,
         revision="fp16",

fx2ait/fx2ait/acc_tracer/acc_normalizer.py

@@ -283,7 +283,10 @@ def move_kwargs_to_acc_out_ty(
 
     for kwarg_replacement_tuple in normalization_info.kwargs_to_move_to_acc_out_ty:
         if len(kwarg_replacement_tuple) == 2:
-            orig_kwarg_name, tmd_field_name, move_to_qparams = *kwarg_replacement_tuple, False  # type: ignore[misc]
+            orig_kwarg_name, tmd_field_name, move_to_qparams = (
+                *kwarg_replacement_tuple,
+                False,
+            )  # type: ignore[misc]
         else:
             assert len(kwarg_replacement_tuple) == 3
             orig_kwarg_name, tmd_field_name, move_to_qparams = kwarg_replacement_tuple  # type: ignore[misc]
@@ -331,9 +334,7 @@ def get_normalized_kwargs(
                 new_kwargs[new_kwarg_name] = node.args[i]
             else:
                 # Verify the arg we're trying to normalize was optional.
-                assert (
-                    is_optional
-                ), f"Cannot normalize {orig_kwargs_names} to {new_kwarg_name} for {node.name}"
+                assert is_optional, f"Cannot normalize {orig_kwargs_names} to {new_kwarg_name} for {node.name}"
         else:
             new_kwargs[new_kwarg_name] = node.kwargs[orig_kwargs_name]
 

fx2ait/fx2ait/acc_tracer/acc_tracer.py

@@ -462,7 +462,10 @@ def __init__(self, orig):
                 for k, v in orig.__dict__.items():
                     if k == "_modules":
                         for mod_k, mod_v in v.items():
-                            if getattr(mod_v, "_base_class_origin", type(mod_v)) in leaf_module_list:  # type: ignore[operator]
+                            if (
+                                getattr(mod_v, "_base_class_origin", type(mod_v))
+                                in leaf_module_list
+                            ):  # type: ignore[operator]
                                 _LOGGER.info(
                                     f"Skip rewriting leaf module {type(mod_v)}"
                                 )

fx2ait/fx2ait/converters/ait_converters.py

@@ -920,8 +920,8 @@ def acc_ops_conv_transpose2d(
         # Grouped conv doesn't currently work on AIT CUDA, manually map
         groups = kwargs["groups"]
         assert (
-            w_last_dim * groups
-        ) % 8 == 0, f"cutlass needs weight output channel={w_last_dim*groups} is not divisble by 8! This restriction may be not valid in newer version"
+            (w_last_dim * groups) % 8 == 0
+        ), f"cutlass needs weight output channel={w_last_dim*groups} is not divisble by 8! This restriction may be not valid in newer version"
 
         group_size = input_val.shape()[3]._attrs["values"][0] // groups
         w_group_size = weight.shape()[0]._attrs["values"][0] // groups
@@ -1767,7 +1767,7 @@ def acc_ops_to_dtype(
     input_val = kwargs["input"]
 
     def _get_cast_to_dtype_from_kwargs(
-        kwargs: Dict[str, Argument]
+        kwargs: Dict[str, Argument],
     ) -> Optional[torch.dtype]:
         torch_dtype_to_ait_dtype_str = {
             torch.float: "float32",

fx2ait/fx2ait/extension.py

@@ -27,7 +27,6 @@ def is_oss_ait_model():
 
 
 def _get_extension_path(lib_name):
-
     lib_dir = os.path.dirname(__file__)
 
     loader_details = (

fx2ait/fx2ait/tools/ait_subgraph_rewriter.py

@@ -348,7 +348,6 @@ def _replace_pattern(
     replacement: Union[Callable, GraphModule],
     match_filters: List[Callable[["InternalMatch", Graph, Graph], bool]] = None,  # type: ignore[name-defined]
 ) -> List[ReplacedPatterns]:
-
     if match_filters is None:
         match_filters = []
 
@@ -392,7 +391,6 @@ def _replace_pattern(
 
     match_and_replacements = []
     for match in _matches:
-
         # Build connecting between replacement graph's input and original graph input producer node
 
         # Initialize `val_map` with mappings from placeholder nodes in
@@ -458,7 +456,6 @@ def get_replacement_nodes(curr_node: Node):
                 and node.op != "output"
                 and node.target != torch.ops.aten.sym_size
             ):
-
                 gn = match.nodes_map[node]
                 gm.graph.erase_node(gn)
         match_and_replacements.append(

python/aitemplate/backend/__init__.py

@@ -15,6 +15,7 @@
 """
 Backend for AITemplate.
 """
+
 from aitemplate.backend import (  # noqa
     backend_spec,
     builder,

python/aitemplate/backend/codegen.py

@@ -651,9 +651,7 @@ def _codegen_output_tensor(self, tensor: Tensor) -> None:
         elif external_tensor is not None:
             # Special view cases for outputs; we can hit this case if the output
             # is a view of a constant, input, or another output.
-            assert (
-                is_view
-            ), f"orig_tensor is not None, but node {name} is not marked as a view! Node: {tensor}"
+            assert is_view, f"orig_tensor is not None, but node {name} is not marked as a view! Node: {tensor}"
             self.set_inputs.append(
                 check_not_null(tensor, output_idx, skip_if_lower_bound_is_zero=True)
             )

python/aitemplate/backend/common/concatenate_common.py

@@ -15,6 +15,7 @@
 """
 backend concatenate function common templates.
 """
+
 from copy import deepcopy
 from typing import List
 

python/aitemplate/backend/common/split_common.py

@@ -15,6 +15,7 @@
 """
 Backend-agnostic function templates for split.
 """
+
 import jinja2
 
 FUNC_DECL_TEMPLATE = jinja2.Template(

python/aitemplate/backend/common/tensor/permute0213_common.py

@@ -23,6 +23,7 @@
 in the blockIdx.z for the direct kernel launch. The input and output
 pointers are shifted accordingly in the kernel code.
 """
+
 from typing import Any, Dict
 
 import jinja2

python/aitemplate/backend/common/tensor/permute021_common.py

@@ -21,6 +21,7 @@
 i.e. Output[d0, ..., dn-3, dn-1, dn-2] = Input[d0, ..., dn-3, dn-2, dn-1]
 
 """
+
 from typing import Any, Dict
 
 import jinja2

python/aitemplate/backend/common/tensor/permute102_common.py

@@ -36,6 +36,7 @@
 starting from 17 items, the approach #1 corresponds to the same data
 movement, just through the SMEM and with more index computation.
 """
+
 from typing import Any, Dict
 
 import jinja2

python/aitemplate/backend/common/tensor/permute210_common.py

@@ -26,6 +26,7 @@
 The 4 for thread blocks indicates each thread is responsible of 4 elements.
 We use TILE_SIZE = 32 for the time being.
 """
+
 from typing import Any, Dict
 
 import jinja2

python/aitemplate/backend/common/tensor/slice_common.py

@@ -15,6 +15,7 @@
 """
 Slice backend common implementation.
 """
+
 import jinja2
 
 

python/aitemplate/backend/common/tensor/slice_reshape_scatter_common.py

@@ -15,6 +15,7 @@
 """
 Slice reshape backend common implementation.
 """
+
 import functools
 
 import jinja2

python/aitemplate/backend/common/vision_ops/efficient_nms_kernel.py

@@ -15,6 +15,7 @@
 """
 efficient_nms function gpu kernel.
 """
+
 import jinja2
 
 kernel = jinja2.Template(

python/aitemplate/backend/common/vision_ops/nms_kernel.py

@@ -15,6 +15,7 @@
 """
 nms kernel template.
 """
+
 import jinja2
 
 KERNEL_TEMPLATE = jinja2.Template(

python/aitemplate/backend/cuda/__init__.py

@@ -16,6 +16,7 @@
 """
 CUDA backend codegen functions.
 """
+
 from aitemplate.backend.cuda import (
     builder_cmake,
     cuda_common,

python/aitemplate/backend/cuda/attention/__init__.py

@@ -15,6 +15,7 @@
 """
 cuda flash_attention module init
 """
+
 from aitemplate.backend.cuda.attention import flash_attention, mem_eff_attention
 
 __all__ = ["flash_attention", "mem_eff_attention"]

python/aitemplate/backend/cuda/attention/flash_attention.py

@@ -15,6 +15,7 @@
 """
 attention kernel codegen for CUDA.
 """
+
 from typing import Any, Dict
 
 import jinja2

python/aitemplate/backend/cuda/attention/mem_eff_attention.py

@@ -15,6 +15,7 @@
 """
 Attention kernel codegen for CUDA.
 """
+
 from typing import Any, Dict
 
 import jinja2

python/aitemplate/backend/cuda/b2b_bmm/classic_b2b_bmm.py

@@ -15,6 +15,7 @@
 """
 classic_b2b_bmm kernel codegen for CUDA.
 """
+
 from typing import Any, Dict
 
 import jinja2

python/aitemplate/backend/cuda/b2b_bmm/fmha_style_b2b_bmm.py

@@ -15,6 +15,7 @@
 """
 fmha_style_b2b_bmm kernel codegen for CUDA.
 """
+
 from typing import Any, Dict
 
 import jinja2

python/aitemplate/backend/cuda/b2b_bmm/grouped_classic_b2b_bmm.py

@@ -15,6 +15,7 @@
 """
 classic_b2b_bmm kernel codegen for CUDA.
 """
+
 from typing import Any, Dict
 
 import jinja2

python/aitemplate/backend/cuda/b2b_bmm/grouped_fmha_style_b2b_bmm.py

@@ -15,6 +15,7 @@
 """
 grouped_fmha_style_b2b_bmm kernel codegen for CUDA.
 """
+
 from typing import Any, Dict
 
 import jinja2

python/aitemplate/backend/cuda/common/__init__.py

@@ -16,4 +16,5 @@
 """
 CUDA Common module init
 """
+
 from aitemplate.backend.cuda.common.dummy_op import *

python/aitemplate/backend/cuda/conv2d/__init__.py

@@ -16,6 +16,7 @@
 """
 cuda conv2d module init
 """
+
 from aitemplate.backend.cuda.conv2d import (
     conv2d,
     conv2d_bias,

python/aitemplate/backend/cuda/conv2d/common.py

@@ -15,6 +15,7 @@
 """
 common template for conv2d
 """
+
 import re
 
 from collections import OrderedDict

python/aitemplate/backend/cuda/conv2d/conv2d.py

@@ -15,6 +15,7 @@
 """
 Codegen for conv2d.
 """
+
 from aitemplate.backend import registry
 from aitemplate.backend.cuda.conv2d import common
 

python/aitemplate/backend/cuda/conv2d/conv2d_bias.py

@@ -15,6 +15,7 @@
 """
 conv2d bias codegen
 """
+
 from aitemplate.backend import registry
 from aitemplate.backend.cuda.conv2d import common, common_conv2d_bias_activation as cba
 

python/aitemplate/backend/cuda/conv2d/conv2d_bias_add.py

@@ -15,6 +15,7 @@
 """
 conv2d bias add codegen
 """
+
 from aitemplate.backend import registry
 from aitemplate.backend.cuda.conv2d import (
     common,

python/aitemplate/backend/cuda/conv2d/conv2d_bias_add_hardswish.py

@@ -15,6 +15,7 @@
 """
 conv2d bias add hardswish codegen
 """
+
 from aitemplate.backend import registry
 from aitemplate.backend.cuda.conv2d import (
     common,

python/aitemplate/backend/cuda/conv2d/conv2d_bias_add_relu.py

@@ -15,6 +15,7 @@
 """
 conv2d bias add relu codegen
 """
+
 from aitemplate.backend import registry
 from aitemplate.backend.cuda.conv2d import (
     common,

python/aitemplate/backend/cuda/conv2d/conv2d_bias_hardswish.py

@@ -15,6 +15,7 @@
 """
 conv2d bias hardswish codegen
 """
+
 from aitemplate.backend import registry
 from aitemplate.backend.cuda.conv2d import common, common_conv2d_bias_activation as cba
 

python/aitemplate/backend/cuda/conv2d/conv2d_bias_relu.py

@@ -15,6 +15,7 @@
 """
 conv2d bias relu codegen
 """
+
 from aitemplate.backend import registry
 from aitemplate.backend.cuda.conv2d import common, common_conv2d_bias_activation as cba