precommit formatting

Signed-off-by: geomin12 <[email protected]>

linalg_ops/attention/generate_e2e_attention_tests.py

@@ -50,12 +50,14 @@ class ShapesId(enum.Enum):
     MEDIUM = "medium"
     LARGE = "large"
 
+
 # Enumerates ways to construct MLIR tensor types.
 @enum.unique
 class Dynamicity(enum.Enum):
     MIXED = "mixed"  # Randomly mix '?' and values. Example: tensor<?x4xf32>.
     STATIC = "static"  # Use fixed values everywhere. Example: tensor<4x6xf32>.
 
+
 # batch: Batch dimension
 # m: M dimension of first and second matmul
 # n: N dimension of second matmul
@@ -90,6 +92,7 @@ def get_test_shapes(shapes_id: ShapesId):
 
     raise ValueError(shapes_id)
 
+
 # A shape dimension value, i.e. a size value that could appear in a MLIR type
 # such as 'tensor<?x4xf32>'. None means a dynamic size, similar to '?' in MLIR.
 @dataclasses.dataclass
@@ -118,6 +121,7 @@ def int_or_question_mark(s: DimSize):
 def int_or_DYN(s: DimSize):
     return s.value or "DYN"
 
+
 # Determines the shape of input and kernel tensors.
 @dataclasses.dataclass
 class TestInputTensorShapes:
@@ -133,7 +137,6 @@ class TestInputTensorShapes:
 # converts from the runtime shape dimensions in TestShape and given dynamicity to
 # the set of shapes to be used in a test function's input tensors.
 def generate_shapes_and_scale(shape: TestShapeAndScale, dynamicity: Dynamicity):
-
     m = shape_dim(shape.m, dynamicity)
     k2 = shape_dim(shape.k2, dynamicity)
     if dynamicity == Dynamicity.MIXED:
@@ -262,18 +265,23 @@ def generate_function(
     )
     if dynamicity == Dynamicity.MIXED:
         func_definition = func_definition + (
-             f"  %c1 = arith.constant 1 : index\n"
-             f"  %m_in = tensor.dim %query, %c1 : {query_tensor_type}\n"
-             f"  %k2_in = tensor.dim %key, %c1 : {key_tensor_type}\n"
-             f"  %m = util.assume.int %m_in<udiv = 16> : index\n"
-             f"  %k2 = util.assume.int %k2_in<udiv = 16> : index\n"
-             f"  %query_mixed = flow.tensor.tie_shape %query : {query_tensor_type}""{%m}\n"
-             f"  %key_mixed = flow.tensor.tie_shape %key : {key_tensor_type}""{%k2}\n"
-             f"  %value_mixed = flow.tensor.tie_shape %value : {value_tensor_type}""{%k2}\n"
-             f"  %result0 = tensor.empty(%m): {result_tensor_type}\n"
+            f"  %c1 = arith.constant 1 : index\n"
+            f"  %m_in = tensor.dim %query, %c1 : {query_tensor_type}\n"
+            f"  %k2_in = tensor.dim %key, %c1 : {key_tensor_type}\n"
+            f"  %m = util.assume.int %m_in<udiv = 16> : index\n"
+            f"  %k2 = util.assume.int %k2_in<udiv = 16> : index\n"
+            f"  %query_mixed = flow.tensor.tie_shape %query : {query_tensor_type}"
+            "{%m}\n"
+            f"  %key_mixed = flow.tensor.tie_shape %key : {key_tensor_type}"
+            "{%k2}\n"
+            f"  %value_mixed = flow.tensor.tie_shape %value : {value_tensor_type}"
+            "{%k2}\n"
+            f"  %result0 = tensor.empty(%m): {result_tensor_type}\n"
         )
     else:
-        func_definition = func_definition + ( f"  %result0 = tensor.empty(): {result_tensor_type}\n")
+        func_definition = func_definition + (
+            f"  %result0 = tensor.empty(): {result_tensor_type}\n"
+        )
 
     func_definition = func_definition + (
         f"  %scale_f16 = arith.truncf %scale : {F32} to {F16}\n"

litert_models/utils.py

@@ -87,10 +87,7 @@ def compile_mlir_with_iree(
     compile_args.extend(compile_flags)
     compile_args.extend(["-o", iree_module_path])
     compile_cmd = subprocess.list2cmdline(compile_args)
-    logger.info(
-        f"Launching compile command:\n"  #
-        f"  cd {cwd} && {compile_cmd}"
-    )
+    logger.info(f"Launching compile command:\n" f"  cd {cwd} && {compile_cmd}")  #
     ret = subprocess.run(compile_cmd, shell=True, capture_output=True)
     if ret.returncode != 0:
         logger.error(f"Compilation of '{iree_module_path}' failed")

onnx_models/conftest.py

@@ -251,10 +251,7 @@ def compile_mlir_with_iree(
     compile_args.extend(compile_flags)
     compile_args.extend(["-o", iree_module_path.relative_to(cwd)])
     compile_cmd = subprocess.list2cmdline(compile_args)
-    logger.info(
-        f"Launching compile command:\n"  #
-        f"  cd {cwd} && {compile_cmd}"
-    )
+    logger.info(f"Launching compile command:\n" f"  cd {cwd} && {compile_cmd}")  #
     ret = subprocess.run(compile_cmd, shell=True, capture_output=True, cwd=cwd)
     if ret.returncode != 0:
         logger.error(f"Compilation of '{iree_module_path}' failed")
@@ -271,10 +268,7 @@ def run_iree_module(iree_module_path: Path, run_flags: list[str]):
     run_args = ["iree-run-module", f"--module={iree_module_path.relative_to(cwd)}"]
     run_args.extend(run_flags)
     run_cmd = subprocess.list2cmdline(run_args)
-    logger.info(
-        f"Launching run command:\n"  #
-        f"  cd {cwd} && {run_cmd}"
-    )
+    logger.info(f"Launching run command:\n" f"  cd {cwd} && {run_cmd}")  #
     ret = subprocess.run(run_cmd, shell=True, capture_output=True, cwd=cwd)
     if ret.returncode != 0:
         logger.error(f"Run of '{iree_module_path}' failed")

onnx_models/utils.py

@@ -169,10 +169,7 @@ def import_onnx_model_to_mlir(onnx_path: Path) -> Path:
         str(imported_mlir_path),
     ]
     import_cmd = subprocess.list2cmdline(import_args)
-    logger.info(
-        f"Running import command:\n"  #
-        f"  {import_cmd}"
-    )
+    logger.info(f"Running import command:\n" f"  {import_cmd}")  #
     ret = subprocess.run(import_cmd, shell=True, capture_output=True)
     if ret.returncode != 0:
         logger.error(f"Import of '{onnx_path.name}' failed!")

onnx_ops/conftest.py

@@ -329,8 +329,7 @@ def runtest(self):
     def test_compile(self):
         cwd = self.test_cwd
         logging.getLogger().info(
-            f"Launching compile command:\n"  #
-            f"cd {cwd} && {self.compile_cmd}"
+            f"Launching compile command:\n" f"cd {cwd} && {self.compile_cmd}"  #
         )
         proc = subprocess.run(
             self.compile_cmd, shell=True, capture_output=True, cwd=cwd
@@ -346,8 +345,7 @@ def test_compile(self):
     def test_run(self):
         cwd = self.test_cwd
         logging.getLogger().info(
-            f"Launching run command:\n"  #
-            f"cd {cwd} && {self.run_cmd}"
+            f"Launching run command:\n" f"cd {cwd} && {self.run_cmd}"  #
         )
         proc = subprocess.run(self.run_cmd, shell=True, capture_output=True, cwd=cwd)
         if proc.returncode != 0:

sharktank_models/benchmarks/conftest.py

@@ -18,53 +18,78 @@
 
 logger = logging.getLogger(__name__)
 
+
 def pytest_sessionstart(session):
-    with open("job_summary.md", "a") as job_summary, open("job_summary.json", "w+") as content:
+    with open("job_summary.md", "a") as job_summary, open(
+        "job_summary.json", "w+"
+    ) as content:
         print(f"{sku.upper()} Complete Benchmark Summary:\n", file=job_summary)
         json.dump({}, content)
-        
+
     logger.info("Pytest benchmark test session is starting")
-
+
+
 def pytest_sessionfinish(session, exitstatus):
     markdown_data = {
-        "time_summary": ["Model name", "Submodel name", "Current time (ms)", "Expected/golden time (ms)"],
-        "dispatch_summary": ["Model name", "Submodel name", "Current dispatch count", "Expected/golden dispatch count"],
-        "size_summary": ["Model name", "Submodel name", "Current binary size (bytes)", "Expected/golden binary size (bytes)"]
+        "time_summary": [
+            "Model name",
+            "Submodel name",
+            "Current time (ms)",
+            "Expected/golden time (ms)",
+        ],
+        "dispatch_summary": [
+            "Model name",
+            "Submodel name",
+            "Current dispatch count",
+            "Expected/golden dispatch count",
+        ],
+        "size_summary": [
+            "Model name",
+            "Submodel name",
+            "Current binary size (bytes)",
+            "Expected/golden binary size (bytes)",
+        ],
     }
-
-    with open("job_summary.md", "a") as job_summary, open("job_summary.json", "r") as content:
+
+    with open("job_summary.md", "a") as job_summary, open(
+        "job_summary.json", "r"
+    ) as content:
         summary_data = json.loads(content.read())
         for key, value in markdown_data.items():
             if key in summary_data:
                 table_data = tabulate.tabulate(
                     summary_data.get(key), headers=value, tablefmt="pipe"
                 )
                 print("\n" + table_data, file=job_summary)
-    
+
     logger.info("Pytest benchmark test session has finished")
 
+
 def pytest_collect_file(parent, file_path):
-    if file_path.suffix == ".json" and "job_summary" not in file_path.name and "benchmarks" in str(THIS_DIR):
+    if (
+        file_path.suffix == ".json"
+        and "job_summary" not in file_path.name
+        and "benchmarks" in str(THIS_DIR)
+    ):
         return SharkTankModelBenchmarkTests.from_parent(parent, path=file_path)
 
-@dataclass(frozen = True)
+
+@dataclass(frozen=True)
 class BenchmarkTestSpec:
     model_name: str
     benchmark_file_name: str
-
+
+
 class SharkTankModelBenchmarkTests(pytest.File):
-
     def collect(self):
         path = str(self.path).split("/")
         benchmark_file_name = path[-1].replace(".json", "")
         model_name = path[-2]
-        
+
         item_name = f"{model_name} :: {benchmark_file_name}"
-        
+
         spec = BenchmarkTestSpec(
-            model_name = model_name,
-            benchmark_file_name = benchmark_file_name
+            model_name=model_name, benchmark_file_name=benchmark_file_name
         )
-        
+
         yield ModelBenchmarkRunItem.from_parent(self, name=item_name, spec=spec)
-

sharktank_models/benchmarks/model_benchmark_run.py

@@ -30,6 +30,8 @@
 """
 Helper methods
 """
+
+
 # Converts a list of inputs into compiler friendly input arguments
 def get_input_list(input_list):
     return [f"--input={entry}" for entry in input_list]
@@ -85,13 +87,14 @@ def e2e_iree_benchmark_module_args(modules, file_suffix):
 
 
 class ModelBenchmarkRunItem(pytest.Item):
-
     def __init__(self, spec, **kwargs):
         super().__init__(**kwargs)
         self.spec = spec
         self.model_name = self.spec.model_name
         self.benchmark_file_name = self.spec.benchmark_file_name
-        SUBMODEL_FILE_PATH = THIS_DIR / f"{self.model_name}/{self.benchmark_file_name}.json"
+        SUBMODEL_FILE_PATH = (
+            THIS_DIR / f"{self.model_name}/{self.benchmark_file_name}.json"
+        )
         split_file_name = self.benchmark_file_name.split("_")
         self.submodel_name = "_".join(split_file_name[:-1])
         type_of_backend = split_file_name[-1]
@@ -164,9 +167,13 @@ def runtest(self):
 
         # If there are modules for an e2e pipeline test, reset exec_args and directory_compile variables to custom variables
         if self.modules:
-            all_modules_found, exec_args = e2e_iree_benchmark_module_args(self.modules, self.file_suffix)
+            all_modules_found, exec_args = e2e_iree_benchmark_module_args(
+                self.modules, self.file_suffix
+            )
             if not all_modules_found:
-                pytest.skip(f"Modules needed for {self.model_name} :: {self.submodel_name} not found, unable to run benchmark tests. Skipping...")
+                pytest.skip(
+                    f"Modules needed for {self.model_name} :: {self.submodel_name} not found, unable to run benchmark tests. Skipping..."
+                )
             vmfb_file_path = f"{vmfb_dir}/{self.compiled_file_name}.vmfb"
 
         exec_args += (
@@ -176,9 +183,11 @@ def runtest(self):
             + get_input_list(self.inputs)
             + self.benchmark_flags
         )
-        
+
         if not Path(vmfb_file_path).is_file():
-            pytest.skip(f"Vmfb file for {self.model_name} :: {self.submodel_name} was not found. Unable to run benchmark tests, skipping...")
+            pytest.skip(
+                f"Vmfb file for {self.model_name} :: {self.submodel_name} was not found. Unable to run benchmark tests, skipping..."
+            )
 
         # run iree benchmark command
         ret_value, output = iree_benchmark_module(
@@ -198,10 +207,17 @@ def runtest(self):
         # golden time check
         if self.golden_time:
             # Writing to time summary
-            mean_time_row = [self.model_name, self.submodel_name, str(benchmark_mean_time), self.golden_time]
+            mean_time_row = [
+                self.model_name,
+                self.submodel_name,
+                str(benchmark_mean_time),
+                self.golden_time,
+            ]
             with open("job_summary.json", "r+") as job_summary:
                 file_data = json.loads(job_summary.read())
-                file_data["time_summary"] = file_data.get("time_summary", []) + [mean_time_row]
+                file_data["time_summary"] = file_data.get("time_summary", []) + [
+                    mean_time_row
+                ]
                 job_summary.seek(0)
                 json.dump(file_data, job_summary)
 
@@ -226,10 +242,17 @@ def runtest(self):
                 comp_stats["stream-aggregate"]["execution"]["dispatch-count"]
             )
 
-            dispatch_count_row = [self.model_name, self.submodel_name, dispatch_count, self.golden_dispatch]
+            dispatch_count_row = [
+                self.model_name,
+                self.submodel_name,
+                dispatch_count,
+                self.golden_dispatch,
+            ]
             with open("job_summary.json", "r+") as job_summary:
                 file_data = json.loads(job_summary.read())
-                file_data["dispatch_summary"] = file_data.get("dispatch_summary", []) + [mean_time_row]
+                file_data["dispatch_summary"] = file_data.get(
+                    "dispatch_summary", []
+                ) + [mean_time_row]
                 job_summary.seek(0)
                 json.dump(file_data, job_summary)
 
@@ -250,13 +273,20 @@ def runtest(self):
             module_path = f"{directory_compile}/model.{self.file_suffix}.vmfb"
             binary_size = Path(module_path).stat().st_size
 
-            binary_size_row = [self.model_name, self.submodel_name, binary_size, self.golden_size]
+            binary_size_row = [
+                self.model_name,
+                self.submodel_name,
+                binary_size,
+                self.golden_size,
+            ]
             with open("job_summary.json", "r+") as job_summary:
                 file_data = json.loads(job_summary.read())
-                file_data["size_summary"] = file_data.get("size_summary", []) + [mean_time_row]
+                file_data["size_summary"] = file_data.get("size_summary", []) + [
+                    mean_time_row
+                ]
                 job_summary.seek(0)
                 json.dump(file_data, job_summary)
-                
+
             logger.info(
                 (
                     f"{self.model_name} {self.submodel_name} binary size: {binary_size} bytes"
@@ -269,9 +299,9 @@ def runtest(self):
                 self.golden_size,
                 f"{self.model_name} {self.submodel_name} binary size should not get bigger",
             )
-            
+
     def repr_failure(self, excinfo):
         return super().repr_failure(excinfo)
-        
+
     def reportinfo(self):
         return self.path, 0, f"usecase: {self.name}"