Format python files using Black

Author: stijnh

python/kernel_launcher/reader.py

@@ -163,7 +163,9 @@ def _parse_scalar_argument(entry):
     return np.frombuffer(bytes(data), dtype=dtype)[0]
 
 
-def _parse_array_file(file_name: str, data_dir: str, expect_hash: str, dtype, validate: bool):
+def _parse_array_file(
+    file_name: str, data_dir: str, expect_hash: str, dtype, validate: bool
+):
     file_path = os.path.join(data_dir, file_name)
 
     if file_name.endswith(".gz"):
@@ -190,15 +192,21 @@ def _parse_array_argument(entry: dict, data_dir: str, validate_checksum: bool):
         dtype = _type_name_to_dtype(type_name[:-1])
 
     if dtype is None:
-        logger.warning(f"unknown type \"{type_name}\", falling back to byte array")
+        logger.warning(f'unknown type "{type_name}", falling back to byte array')
         dtype = np.byte
 
-    arg = _parse_array_file(entry["file"], data_dir, entry.get("file_hash"),
-                            dtype, validate_checksum)
+    arg = _parse_array_file(
+        entry["file"], data_dir, entry.get("file_hash"), dtype, validate_checksum
+    )
 
     if "reference_file" in entry:
-        answer = _parse_array_file(entry["reference_file"], data_dir, entry.get("reference_hash"),
-                                   dtype, validate_checksum)
+        answer = _parse_array_file(
+            entry["reference_file"],
+            data_dir,
+            entry.get("reference_hash"),
+            dtype,
+            validate_checksum,
+        )
     else:
         answer = None
 
@@ -230,7 +238,15 @@ def __init__(self, obj, data_dir: str, validate_checksum: bool = True):
             self.args.append(arg)
             self.answers.append(answer)
 
-    def _tune_options(self, working_dir=None, lang="cupy", compiler_options=None, defines=None, device=0, **kwargs):
+    def _tune_options(
+        self,
+        working_dir=None,
+        lang="cupy",
+        compiler_options=None,
+        defines=None,
+        device=0,
+        **kwargs,
+    ):
         if working_dir is None:
             working_dir = os.getcwd()
 
@@ -277,20 +293,21 @@ def grid_size(config):
         restrictions = [e.resolve(**context) for e in self.space.restrictions]
 
         options = dict(
-                kernel_name=self.kernel.generate_name(),
-                kernel_source=self.kernel.generate_source(working_dir),
-                arguments=self.args,
-                problem_size=grid_size,
-                restrictions=lambda config: all(f(config) for f in restrictions),
-                defines=all_defines,
-                compiler_options=compiler_options,
-                block_size_names=block_size_names,
-                grid_div_x=[],
-                grid_div_y=[],
-                grid_div_z=[],
-                lang=lang,
-                device=device,
-                **kwargs)
+            kernel_name=self.kernel.generate_name(),
+            kernel_source=self.kernel.generate_source(working_dir),
+            arguments=self.args,
+            problem_size=grid_size,
+            restrictions=lambda config: all(f(config) for f in restrictions),
+            defines=all_defines,
+            compiler_options=compiler_options,
+            block_size_names=block_size_names,
+            grid_div_x=[],
+            grid_div_y=[],
+            grid_div_z=[],
+            lang=lang,
+            device=device,
+            **kwargs,
+        )
 
         os.chdir(working_dir)
         return extra_params, options
@@ -351,16 +368,25 @@ def tune(self, params=None, **kwargs):
             strategy = "brute_force" if total_configs < 100 else "bayes_opt"
 
         return kernel_tuner.tune_kernel(
-                tune_params=params,
-                strategy=strategy,
-                answer=answer,
-                verify=verify,
-                **options)
+            tune_params=params,
+            strategy=strategy,
+            answer=answer,
+            verify=verify,
+            **options,
+        )
 
 
 def _fancy_verify(answers, outputs, *, atol=None):
-    INTEGRAL_DTYPES = [np.int8, np.int16, np.int32, np.int64,
-                       np.uint8, np.uint16, np.uint32, np.uint64]
+    INTEGRAL_DTYPES = [
+        np.int8,
+        np.int16,
+        np.int32,
+        np.int64,
+        np.uint8,
+        np.uint16,
+        np.uint32,
+        np.uint64,
+    ]
     FLOATING_DTYPES = [np.float16, np.float32, np.float64]
     PRINT_TOP_VALUES = 25
     DEFAULT_ATOL = 1e-8
@@ -378,7 +404,9 @@ def _fancy_verify(answers, outputs, *, atol=None):
             continue
 
         if output.dtype != expected.dtype or output.shape != expected.shape:
-            raise RuntimeError(f"arrays data type or shape do not match: {output} and {expected}")
+            raise RuntimeError(
+                f"arrays data type or shape do not match: {output} and {expected}"
+            )
 
         if output.dtype in INTEGRAL_DTYPES:
             matches = output == expected
@@ -401,14 +429,18 @@ def _fancy_verify(answers, outputs, *, atol=None):
         # indices = indices[np.argsort(errors[indices], kind="stable")][::-1]
 
         percentage = nerrors / len(output) * 100
-        print(f"argument {index + 1} fails validation: {nerrors} incorrect values" +
-              f"({percentage:.5}%)")
+        print(
+            f"argument {index + 1} fails validation: {nerrors} incorrect values"
+            + f"({percentage:.5}%)"
+        )
 
         errors = np.abs(output - expected)
 
         for index in indices[:PRINT_TOP_VALUES]:
-            print(f" * at index {index}: {output[index]} != {expected[index]} " +
-                  f"(error: {errors[index]})")
+            print(
+                f" * at index {index}: {output[index]} != {expected[index]} "
+                + f"(error: {errors[index]})"
+            )
 
         if nerrors > PRINT_TOP_VALUES:
             print(f" * ({nerrors - PRINT_TOP_VALUES} more entries have been omitted)")
@@ -583,22 +615,27 @@ def resolve(self, problem_size, **kwargs):
 
 class DeviceAttributeExpr(Expr):
     # Map cuda.h names to cupy names
-    NAME_MAPPING = dict([
-        ('MAX_THREADS_PER_BLOCK', 'MaxThreadsPerBlock'),
-        ('MAX_BLOCK_DIM_X', 'MaxBlockDimX'),
-        ('MAX_BLOCK_DIM_Y', 'MaxBlockDimY'),
-        ('MAX_BLOCK_DIM_Z', 'MaxBlockDimZ'),
-        ('MAX_GRID_DIM_X', 'MaxGridDimX'),
-        ('MAX_GRID_DIM_Y', 'MaxGridDimY'),
-        ('MAX_GRID_DIM_Z', 'MaxGridDimZ'),
-        ('MAX_SHARED_MEMORY_PER_BLOCK', 'MaxSharedMemoryPerBlock'),
-        ('WARP_SIZE', 'WarpSize'),
-        ('MAX_REGISTERS_PER_BLOCK', 'MaxRegistersPerBlock'),
-        ('MULTIPROCESSOR_COUNT', 'MultiProcessorCount'),
-        ('MAX_THREADS_PER_MULTIPROCESSOR', 'MaxThreadsPerMultiProcessor'),
-        ('MAX_SHARED_MEMORY_PER_MULTIPROCESSOR', 'MaxSharedMemoryPerMultiprocessor'),
-        ('MAX_REGISTERS_PER_MULTIPROCESSOR', 'MaxRegistersPerMultiprocessor'),
-    ])
+    NAME_MAPPING = dict(
+        [
+            ("MAX_THREADS_PER_BLOCK", "MaxThreadsPerBlock"),
+            ("MAX_BLOCK_DIM_X", "MaxBlockDimX"),
+            ("MAX_BLOCK_DIM_Y", "MaxBlockDimY"),
+            ("MAX_BLOCK_DIM_Z", "MaxBlockDimZ"),
+            ("MAX_GRID_DIM_X", "MaxGridDimX"),
+            ("MAX_GRID_DIM_Y", "MaxGridDimY"),
+            ("MAX_GRID_DIM_Z", "MaxGridDimZ"),
+            ("MAX_SHARED_MEMORY_PER_BLOCK", "MaxSharedMemoryPerBlock"),
+            ("WARP_SIZE", "WarpSize"),
+            ("MAX_REGISTERS_PER_BLOCK", "MaxRegistersPerBlock"),
+            ("MULTIPROCESSOR_COUNT", "MultiProcessorCount"),
+            ("MAX_THREADS_PER_MULTIPROCESSOR", "MaxThreadsPerMultiProcessor"),
+            (
+                "MAX_SHARED_MEMORY_PER_MULTIPROCESSOR",
+                "MaxSharedMemoryPerMultiprocessor",
+            ),
+            ("MAX_REGISTERS_PER_MULTIPROCESSOR", "MaxRegistersPerMultiprocessor"),
+        ]
+    )
 
     def __init__(self, name):
         self.name = name
@@ -630,8 +667,10 @@ def evaluate(self, config):
         index = self.condition.evaluate(config)
 
         if not is_int_like(index) or index < 0 or index >= len(self.options):
-            raise RuntimeError("expression must yield an integer in " +
-                               f"range 0..{len(self.options)}: {self}")
+            raise RuntimeError(
+                "expression must yield an integer in "
+                + f"range 0..{len(self.options)}: {self}"
+            )
 
         return self.options[int(index)].evaluate(config)
 
@@ -641,8 +680,7 @@ def visit_children(self, fun):
 
 def _parse_expr(entry) -> Expr:
     # literal int, str or float becomes ValueExpr.
-    if isinstance(entry, (int, str, float)) or \
-            entry is None:
+    if isinstance(entry, (int, str, float)) or entry is None:
         return ValueExpr(entry)
 
     # Otherwise it must be an operator expression

python/kernel_launcher/wisdom.py

@@ -17,21 +17,39 @@
 WISDOM_OBJECTIVE = "time"
 
 
-def write_wisdom_for_problem(path: str, problem: TuningProblem, results: list,
-                             env: dict, **kwargs):
-    """Write the results of ``TuningProblem.tune`` to a wisdom file. This function calls ``write_wisdom``
+def write_wisdom_for_problem(
+    path: str, problem: TuningProblem, results: list, env: dict, **kwargs
+):
+    """Write the results of ``TuningProblem.tune`` to a wisdom file. This function just wraps ``write_wisdom``
 
     :param path: Directory were wisdom files are stored. Alternatively, this can be a file name ending with ``.wisdom``.
     :param problem: The ``TuningProblem``.
     :param results: The results returned by ``kernel_tuner.tune_kernel``.
     :param env: The environment returned by ``kernel_tuner.tune_kernel``.
     :param kwargs: Additional keyword arguments passed to `write_wisdom`.
     """
-    return write_wisdom(path, problem.key, problem.space.params, problem.problem_size, results, env, **kwargs)
-
-
-def write_wisdom(path: str, key: str, params: dict, problem_size: list, results: list, env: dict, *,
-                 max_results=5, merge_existing_results=False):
+    return write_wisdom(
+        path,
+        problem.key,
+        problem.space.params,
+        problem.problem_size,
+        results,
+        env,
+        **kwargs,
+    )
+
+
+def write_wisdom(
+    path: str,
+    key: str,
+    params: dict,
+    problem_size: list,
+    results: list,
+    env: dict,
+    *,
+    max_results=5,
+    merge_existing_results=False,
+):
     """Write the results of ``kernel_tuner.tune_kernel`` to a wisdom file.
 
     :param path: Directory were wisdom files are stored. Alternatively, this can be a file name ending with ``.wisdom``.
@@ -42,7 +60,8 @@ def write_wisdom(path: str, key: str, params: dict, problem_size: list, results:
     :param env: The environment returned by ``kernel_tuner.tune_kernel``
     :param max_results: Only the top ``max_results`` results are written in the wisdom file.
     :param merge_existing_results: If ``True``, existing results in the wisdom file for the same problem size and
-                                   environment are merged with the provided ``results``.
+                                   environment are merged with the provided ``results``. If ``False``, existing
+                                   results for the same problem size and environment are overwritten.
 
     """
     device_name = env["device_name"]
@@ -60,9 +79,11 @@ def write_wisdom(path: str, key: str, params: dict, problem_size: list, results:
 
             for line, record in _parse_wisdom(handle):
                 # Skip lines that have a matching problem_size and device_name
-                if not record or \
-                        record["problem_size"] != problem_size or \
-                        record["environment"].get("device_name") != device_name:
+                if (
+                    not record
+                    or record["problem_size"] != problem_size
+                    or record["environment"].get("device_name") != device_name
+                ):
                     lines.append(line)
                 elif merge_existing_results:
                     index = tuple(record["config"])
@@ -105,7 +126,9 @@ def write_wisdom(path: str, key: str, params: dict, problem_size: list, results:
         handle.writelines(line + "\n" for line in lines)
 
 
-def read_wisdom(path: str, key: str = None, params: dict = None, *, error_if_missing: bool = True) -> list:
+def read_wisdom(
+    path: str, key: str = None, params: dict = None, *, error_if_missing: bool = True
+) -> list:
     """
     Read the results of a wisdom file.
 
@@ -168,26 +191,31 @@ def _check_header(line: str, key: str, params: dict):
     if key is not None:
         if data.get("key") != key:
             print(data)
-            raise RuntimeError(f"invalid key in wisdom file: {key} != {data.get('key')}")
+            raise RuntimeError(
+                f"invalid key in wisdom file: {key} != {data.get('key')}"
+            )
 
     keys = data.get("tunable_parameters", [])
     if params is not None:
         if set(keys) != set(params.keys()):
-            raise RuntimeError("invalid tunable parameters in wisdom file: " +
-                               f"{list(params.keys())} != {keys}")
+            raise RuntimeError(
+                "invalid tunable parameters in wisdom file: "
+                + f"{list(params.keys())} != {keys}"
+            )
 
     return keys
 
 
 def _create_header(key: str, param_keys: list) -> str:
-    """Header of wisdom file (ie, the first line).
-    """
-    return json.dumps({
-        "version": WISDOM_VERSION,
-        "objective": WISDOM_OBJECTIVE,
-        "tunable_parameters": list(param_keys),
-        "key": key,
-    })
+    """Header of wisdom file (ie, the first line)."""
+    return json.dumps(
+        {
+            "version": WISDOM_VERSION,
+            "objective": WISDOM_OBJECTIVE,
+            "tunable_parameters": list(param_keys),
+            "key": key,
+        }
+    )
 
 
 def _is_valid_config(config):
@@ -230,7 +258,9 @@ def _wisdom_file(path, key):
         filename = re.sub("[^0-9a-zA-Z_.-]", "_", key) + ".wisdom"
         return os.path.join(path, filename)
     else:
-        raise ValueError(f"path must be a directory or a file ending with .wisdom: {path}")
+        raise ValueError(
+            f"path must be a directory or a file ending with .wisdom: {path}"
+        )
 
 
 def _build_environment(env):
@@ -243,13 +273,15 @@ def _build_environment(env):
     # Kernel tuner related
     try:
         import kernel_tuner
+
         env["kernel_tuner_version"] = kernel_tuner.__version__
     except AttributeError as e:
         logger.warning(f"ignore error: kernel_tuner.__version__ is not available: {e}")
 
     # CUDA related
     try:
         import pycuda.driver
+
         env["cuda_driver_version"] = pycuda.driver.get_driver_version()
 
         major, minor, patch = pycuda.driver.get_version()