Initial import: grid-bot — grid trading bot for BTC-USDT on Cifra Markets
This commit is contained in:
Kolp
@@ -0,0 +1,33 @@
|
||||
import argparse
|
||||
import logging
|
||||
import os
|
||||
|
||||
from torch._inductor.autotune_process import TuningProcess
|
||||
from torch._inductor.compile_worker.utils import _async_compile_initializer
|
||||
|
||||
|
||||
log = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--parent", type=int)
|
||||
parser.add_argument("--read-fd", type=int)
|
||||
parser.add_argument("--write-fd", type=int)
|
||||
args = parser.parse_args()
|
||||
read_pipe = os.fdopen(args.read_fd, "rb")
|
||||
write_pipe = os.fdopen(args.write_fd, "wb")
|
||||
|
||||
try:
|
||||
# Ensures the subprocess exits if the parent crashes:
|
||||
_async_compile_initializer(args.parent)
|
||||
TuningProcess.process_main(read_pipe, write_pipe)
|
||||
except Exception:
|
||||
log.exception("Uncaught exception in autotune subprocess")
|
||||
finally:
|
||||
read_pipe.close()
|
||||
write_pipe.close()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,462 @@
|
||||
# mypy: allow-untyped-defs
|
||||
from __future__ import annotations
|
||||
|
||||
import dataclasses
|
||||
import io
|
||||
import logging
|
||||
import os
|
||||
from typing import Any, IO, Literal, Optional, TYPE_CHECKING, Union
|
||||
|
||||
import torch.fx
|
||||
|
||||
from .standalone_compile import CompiledArtifact # noqa: TC001
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from torch._inductor.utils import InputType
|
||||
from torch.export import ExportedProgram
|
||||
from torch.export.pt2_archive._package import AOTICompiledModel
|
||||
from torch.export.pt2_archive._package_weights import Weights
|
||||
from torch.types import FileLike
|
||||
|
||||
__all__ = [
|
||||
"compile",
|
||||
"list_mode_options",
|
||||
"list_options",
|
||||
"cudagraph_mark_step_begin",
|
||||
"standalone_compile",
|
||||
]
|
||||
|
||||
|
||||
log = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def compile(
|
||||
gm: torch.fx.GraphModule,
|
||||
example_inputs: list[InputType],
|
||||
options: dict[str, Any] | None = None,
|
||||
):
|
||||
"""
|
||||
Compile a given FX graph with TorchInductor. This allows compiling
|
||||
FX graphs captured without using TorchDynamo.
|
||||
|
||||
Args:
|
||||
gm: The FX graph to compile.
|
||||
example_inputs: List of tensor inputs.
|
||||
options: Optional dict of config options. See `torch._inductor.config`.
|
||||
|
||||
Returns:
|
||||
Callable with same behavior as gm but faster.
|
||||
"""
|
||||
from .compile_fx import compile_fx
|
||||
|
||||
return compile_fx(gm, example_inputs, config_patches=options)
|
||||
|
||||
|
||||
def aoti_compile_and_package(
|
||||
exported_program: ExportedProgram,
|
||||
_deprecated_unused_args=None,
|
||||
_deprecated_unused_kwargs=None,
|
||||
*,
|
||||
package_path: FileLike | None = None,
|
||||
inductor_configs: dict[str, Any] | None = None,
|
||||
) -> str:
|
||||
"""
|
||||
Compiles the exported program with AOTInductor, and packages it into a .pt2
|
||||
artifact specified by the input package_path. To load the package, you can
|
||||
call ``torch._inductor.aoti_load_package(package_path)``.
|
||||
|
||||
An example usage is as follows:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
ep = torch.export.export(M(), ...)
|
||||
aoti_file = torch._inductor.aoti_compile_and_package(
|
||||
ep, package_path="my_package.pt2"
|
||||
)
|
||||
compiled_model = torch._inductor.aoti_load_package("my_package.pt2")
|
||||
|
||||
To compile and save multiple models into a single ``.pt2`` artifact, you can do
|
||||
the following:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
ep1 = torch.export.export(M1(), ...)
|
||||
aoti_file1 = torch._inductor.aot_compile(
|
||||
ep1, ..., options={"aot_inductor.package": True}
|
||||
)
|
||||
ep2 = torch.export.export(M2(), ...)
|
||||
aoti_file2 = torch._inductor.aot_compile(
|
||||
ep2, ..., options={"aot_inductor.package": True}
|
||||
)
|
||||
|
||||
from torch._inductor.package import package_aoti, load_package
|
||||
|
||||
package_aoti("my_package.pt2", {"model1": aoti_file1, "model2": aoti_file2})
|
||||
|
||||
compiled_model1 = load_package("my_package.pt2", "model1")
|
||||
compiled_model2 = load_package("my_package.pt2", "model2")
|
||||
|
||||
Args:
|
||||
exported_program: An exported program created through a call from torch.export
|
||||
package_path: Optional specified path to the generated .pt2 artifact.
|
||||
inductor_configs: Optional dictionary of configs to control inductor.
|
||||
|
||||
Returns:
|
||||
Path to the generated artifact
|
||||
"""
|
||||
from torch.export import ExportedProgram
|
||||
|
||||
from .debug import aot_inductor_minifier_wrapper
|
||||
|
||||
if not isinstance(exported_program, ExportedProgram):
|
||||
raise ValueError("Only ExportedProgram is supported")
|
||||
|
||||
if exported_program.example_inputs is None:
|
||||
raise RuntimeError(
|
||||
"exported_program.example_inputs is required to be set in order "
|
||||
"for AOTInductor compilation."
|
||||
)
|
||||
|
||||
if _deprecated_unused_args is not None or _deprecated_unused_kwargs is not None:
|
||||
log.warning(
|
||||
"You no longer need to specify args/kwargs to aoti_compile_and_package "
|
||||
"as we can get this information from exported_program.example_inputs."
|
||||
)
|
||||
|
||||
assert (
|
||||
package_path is None
|
||||
or (
|
||||
isinstance(package_path, (io.IOBase, IO))
|
||||
and package_path.writable()
|
||||
and package_path.seekable()
|
||||
)
|
||||
or (
|
||||
isinstance(package_path, (str, os.PathLike))
|
||||
and os.fspath(package_path).endswith(".pt2")
|
||||
)
|
||||
), (
|
||||
f"Expect package path to be a file ending in .pt2, is None, or is a buffer. Instead got {package_path}"
|
||||
)
|
||||
|
||||
inductor_configs = inductor_configs or {}
|
||||
inductor_configs["aot_inductor.package"] = True
|
||||
|
||||
if inductor_configs.get("aot_inductor.output_path"):
|
||||
raise RuntimeError(
|
||||
"Please pass in a package path to aot_inductor_compile() instead "
|
||||
"of setting the aot_inductor.output_path config."
|
||||
)
|
||||
|
||||
# a wrapper around aoti_compile_and_package_inner.
|
||||
return aot_inductor_minifier_wrapper(
|
||||
_aoti_compile_and_package_inner,
|
||||
exported_program,
|
||||
package_path=package_path,
|
||||
inductor_configs=inductor_configs,
|
||||
)
|
||||
|
||||
|
||||
def _aoti_compile_and_package_inner(
|
||||
gm: torch.nn.Module,
|
||||
# flat_example_inputs: List[Any],
|
||||
args: tuple[Any],
|
||||
kwargs: dict[str, Any] | None = None,
|
||||
*,
|
||||
load_and_run: bool = False,
|
||||
check_accuracy: str | None = None,
|
||||
package_path: str | io.BytesIO | None = None,
|
||||
inductor_configs: dict[str, Any] | None = None,
|
||||
):
|
||||
"""
|
||||
See docstring for aoti_compile_and_package.
|
||||
|
||||
If `load_and_run` is True, this function will load the compiled model and run it.
|
||||
This is for the minifier to check the correctness of the compiled model.
|
||||
|
||||
If `check_accuracy` is set, this function will check the accuracy of the compiled
|
||||
model against gm. kwargs must be None if check_accuracy is set.
|
||||
"strict_accuracy" means "we will minify any time we see anything that
|
||||
diverges", whereas "accuracy" is more conservative, and will only minify if there
|
||||
is a meaningful fp64 divergence
|
||||
"""
|
||||
|
||||
if check_accuracy:
|
||||
assert kwargs is None or len(kwargs) == 0, (
|
||||
"when checking for accuracy, the inputs must have been flattened and kwargs is None"
|
||||
)
|
||||
|
||||
from .package import package_aoti
|
||||
|
||||
assert isinstance(gm, torch.fx.GraphModule)
|
||||
|
||||
kwargs = kwargs or {}
|
||||
|
||||
aoti_files = aot_compile(gm, args, kwargs, options=inductor_configs)
|
||||
assert isinstance(aoti_files, list)
|
||||
|
||||
if package_path is None:
|
||||
path = [
|
||||
os.path.splitext(file)[0]
|
||||
for file in aoti_files
|
||||
if isinstance(file, str) and os.path.splitext(file)[1] == ".so"
|
||||
]
|
||||
if len(path) == 0:
|
||||
path = [
|
||||
os.path.splitext(file)[0]
|
||||
for file in aoti_files
|
||||
if isinstance(file, str) and os.path.splitext(file)[1] == ".cpp"
|
||||
]
|
||||
package_path = path[0] + ".pt2"
|
||||
|
||||
res = package_aoti(package_path, aoti_files)
|
||||
assert res == package_path
|
||||
|
||||
if load_and_run or check_accuracy:
|
||||
compiled_model = aoti_load_package(package_path)
|
||||
if check_accuracy:
|
||||
from torch._dynamo.debug_utils import AccuracyError, same_two_models
|
||||
|
||||
# This might look inverted but it's not. strict_accuracy means "we will
|
||||
# minify any time we see anything that diverges", whereas accuracy is more
|
||||
# conservative, and will only minify if there is a meaningful fp64
|
||||
# divergence
|
||||
not_strict_accuracy = check_accuracy == "accuracy"
|
||||
if not same_two_models(
|
||||
gm,
|
||||
compiled_model, # type: ignore[arg-type]
|
||||
args,
|
||||
only_fwd=True,
|
||||
require_fp64=not_strict_accuracy,
|
||||
ignore_non_fp=not_strict_accuracy,
|
||||
):
|
||||
raise AccuracyError("Bad accuracy detected")
|
||||
else:
|
||||
compiled_model(*args, **kwargs)
|
||||
|
||||
return package_path
|
||||
|
||||
|
||||
def aoti_load_package(
|
||||
path: FileLike, run_single_threaded: bool = False, device_index: int = -1
|
||||
) -> AOTICompiledModel:
|
||||
"""
|
||||
Loads the model from the PT2 package.
|
||||
|
||||
If multiple models were packaged into the PT2, this will load the default
|
||||
model. To load a specific model, you can directly call the load API
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
from torch._inductor.package import load_package
|
||||
|
||||
compiled_model1 = load_package("my_package.pt2", "model1")
|
||||
compiled_model2 = load_package("my_package.pt2", "model2")
|
||||
|
||||
Args:
|
||||
path: Path to the .pt2 package
|
||||
run_single_threaded (bool): Whether the model should be run without
|
||||
thread synchronization logic. This is useful to avoid conflicts with
|
||||
CUDAGraphs.
|
||||
device_index (int): The index of the device to which the PT2 package is
|
||||
to be loaded. By default, `device_index=-1` is used, which corresponds
|
||||
to the device `cuda` when using CUDA. Passing `device_index=1` would
|
||||
load the package to `cuda:1`, for example.
|
||||
"""
|
||||
from torch._inductor.package import load_package
|
||||
|
||||
return load_package(
|
||||
path, run_single_threaded=run_single_threaded, device_index=device_index
|
||||
)
|
||||
|
||||
|
||||
def aot_compile(
|
||||
gm: torch.fx.GraphModule,
|
||||
args: tuple[Any, ...],
|
||||
kwargs: dict[str, Any] | None = None,
|
||||
*,
|
||||
options: dict[str, Any] | None = None,
|
||||
) -> str | list[str | Weights] | torch.fx.GraphModule:
|
||||
"""
|
||||
Ahead-of-time compile a given FX graph with TorchInductor into a shared library.
|
||||
|
||||
Args:
|
||||
gm: The FX graph to compile.
|
||||
args: Example arguments
|
||||
kwargs: Example keyword arguments
|
||||
options: Optional dict of config options. See `torch._inductor.config`.
|
||||
|
||||
Returns:
|
||||
Path to the generated shared library, or a list of files generated by
|
||||
AOTI if aot_inductor.package=True.
|
||||
TODO: make it return a list by default
|
||||
"""
|
||||
from .compile_fx import _aoti_flatten_inputs, compile_fx_aot
|
||||
|
||||
if hasattr(gm, "_guards_fn"):
|
||||
# Do not compile the guards function, since it may contain checks
|
||||
# that are not currently supported by AOTI. In particular, non-Tensor
|
||||
# arguments are converted to None and will fail specialization checks.
|
||||
node = next(iter(gm.graph.find_nodes(op="call_module", target="_guards_fn")))
|
||||
gm.graph.erase_node(node)
|
||||
delattr(gm, "_guards_fn")
|
||||
gm.recompile()
|
||||
|
||||
flat_example_inputs, options = _aoti_flatten_inputs(
|
||||
gm, args, kwargs, options=options
|
||||
)
|
||||
from torch._export.utils import _compiling_state_context
|
||||
|
||||
with _compiling_state_context():
|
||||
return compile_fx_aot(
|
||||
gm,
|
||||
flat_example_inputs, # type: ignore[arg-type]
|
||||
config_patches=options,
|
||||
)
|
||||
|
||||
|
||||
lite_mode_options = {
|
||||
# Fallback by default unless users explicitly annotated with
|
||||
# regional inductor compile.
|
||||
"fallback_by_default": True,
|
||||
"selective_decompose": True,
|
||||
# Disable reorder optimizations
|
||||
"reorder_for_peak_memory": False,
|
||||
"reorder_for_compute_comm_overlap": False,
|
||||
"triton.reorder_for_reducing_graph_partitions": False,
|
||||
# Disable pre-, joint-, post-grad passes
|
||||
"use_pre_grad_passes": False,
|
||||
"use_joint_graph_passes": False,
|
||||
"use_post_grad_passes": False,
|
||||
# Disable dead code elimination (dce) and buffer reuse
|
||||
"use_dce": False,
|
||||
"allow_buffer_reuse": False,
|
||||
}
|
||||
|
||||
|
||||
def list_mode_options(
|
||||
mode: str | None = None, dynamic: bool | None = None
|
||||
) -> dict[str, Any]:
|
||||
r"""Returns a dictionary describing the optimizations that each of the available
|
||||
modes passed to `torch.compile()` performs.
|
||||
|
||||
Args:
|
||||
mode (str, optional): The mode to return the optimizations for.
|
||||
If None, returns optimizations for all modes
|
||||
dynamic (bool, optional): Whether dynamic shape is enabled.
|
||||
|
||||
Example::
|
||||
>>> torch._inductor.list_mode_options()
|
||||
"""
|
||||
|
||||
mode_options: dict[str, dict[str, bool]] = {
|
||||
"default": {},
|
||||
# lite backend for opt-in optimizations
|
||||
"lite": lite_mode_options,
|
||||
# enable cudagraphs
|
||||
"reduce-overhead": {
|
||||
"triton.cudagraphs": True,
|
||||
},
|
||||
# enable max-autotune
|
||||
"max-autotune-no-cudagraphs": {
|
||||
"max_autotune": True,
|
||||
"coordinate_descent_tuning": True,
|
||||
},
|
||||
# enable max-autotune
|
||||
# enable cudagraphs
|
||||
"max-autotune": {
|
||||
"max_autotune": True,
|
||||
"triton.cudagraphs": True,
|
||||
"coordinate_descent_tuning": True,
|
||||
},
|
||||
}
|
||||
try:
|
||||
return mode_options[mode] if mode else mode_options
|
||||
except KeyError as e:
|
||||
raise RuntimeError(
|
||||
f"Unrecognized mode={mode}, should be one of: {', '.join(mode_options.keys())}"
|
||||
) from e
|
||||
|
||||
|
||||
def list_options() -> list[str]:
|
||||
r"""Returns a dictionary describing the optimizations and debug configurations
|
||||
that are available to `torch.compile()`.
|
||||
|
||||
The options are documented in `torch._inductor.config`.
|
||||
|
||||
Example::
|
||||
|
||||
>>> torch._inductor.list_options()
|
||||
"""
|
||||
|
||||
from torch._inductor import config
|
||||
|
||||
current_config: dict[str, Any] = config.get_config_copy()
|
||||
|
||||
return list(current_config.keys())
|
||||
|
||||
|
||||
def cudagraph_mark_step_begin():
|
||||
"Indicates that a new iteration of inference or training is about to begin."
|
||||
from .cudagraph_trees import mark_step_begin
|
||||
|
||||
mark_step_begin()
|
||||
|
||||
|
||||
def standalone_compile(
|
||||
gm: torch.fx.GraphModule,
|
||||
example_inputs: list[InputType],
|
||||
*,
|
||||
dynamic_shapes: Literal[
|
||||
"from_example_inputs", "from_tracing_context", "from_graph"
|
||||
] = "from_graph",
|
||||
options: dict[str, Any] | None = None,
|
||||
aot: bool = False, # AOT mode, which uses BundledAOTAutogradCache
|
||||
donate_graph_module: bool = False,
|
||||
) -> CompiledArtifact:
|
||||
"""
|
||||
Precompilation API for inductor.
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
compiled_artifact = torch._inductor.standalone_compile(gm, args)
|
||||
compiled_artifact.save(path=path, format="binary")
|
||||
|
||||
# Later on a new process
|
||||
loaded = torch._inductor.CompiledArtifact.load(path=path, format="binary")
|
||||
compiled_out = loaded(*args)
|
||||
|
||||
Args:
|
||||
gm: Graph Module
|
||||
example_inputs: Inputs for the graph module
|
||||
dynamic_shapes: If "from_graph" (default), we will use the dynamic
|
||||
shapes in the passed-in graph module.
|
||||
If "from_tracing_context", we use the dynamic shape info in the
|
||||
ambient tracing context.
|
||||
If "from_example_inputs", we will specialize the graph on the
|
||||
example_inputs.
|
||||
options: Inductor compilation options
|
||||
donate_graph_module: If True, standalone_compile takes ownership of
|
||||
the graph module and may mutate it, avoiding an internal deepcopy.
|
||||
Defaults to False for backwards compatibility.
|
||||
|
||||
Returns:
|
||||
CompiledArtifact that can be saved to disk or invoked directly.
|
||||
"""
|
||||
from .standalone_compile import standalone_compile
|
||||
|
||||
options = options if options else {}
|
||||
return standalone_compile(
|
||||
gm,
|
||||
example_inputs,
|
||||
dynamic_shapes=dynamic_shapes,
|
||||
options=options,
|
||||
aot=aot,
|
||||
donate_graph_module=donate_graph_module,
|
||||
)
|
||||
|
||||
|
||||
@dataclasses.dataclass
|
||||
class _CudagraphAnnotation:
|
||||
fwd: bool | None
|
||||
bwd: bool | None
|
||||
@@ -0,0 +1,215 @@
|
||||
import logging
|
||||
from dataclasses import dataclass
|
||||
|
||||
import torch
|
||||
|
||||
|
||||
log = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class DeviceInfo:
|
||||
"""
|
||||
Theoretical Numbers from data sheet. If two numbers are given, Tensor/Matrix Core vs not,
|
||||
then the higher number is reported. Sparsity is not considered.
|
||||
|
||||
|
||||
Bandwidth numbers are tricky, because there are platform differences that may not show up in the profiler trace.
|
||||
For example,
|
||||
"""
|
||||
|
||||
tops: dict[torch.dtype | str, float]
|
||||
dram_bw_gbs: float
|
||||
dram_gb: float
|
||||
|
||||
|
||||
# Indexing is based on `torch.cuda.get_device_name()`
|
||||
# TODO investigate profiler support for tf32 and allow device to report correct number when it's turned on.
|
||||
_device_mapping: dict[str, DeviceInfo] = {
|
||||
# Source:
|
||||
# @lint-ignore https://www.nvidia.com/en-us/data-center/h100/
|
||||
"NVIDIA H100": DeviceInfo(
|
||||
tops={
|
||||
torch.float64: 67.0,
|
||||
torch.float32: 67.5,
|
||||
"torch.tf32": 156.0,
|
||||
torch.bfloat16: 1979.0,
|
||||
torch.float16: 1979.0,
|
||||
torch.float8_e8m0fnu: 3958.0,
|
||||
torch.float8_e8m0fnu: 3958.0,
|
||||
torch.float8_e4m3fnuz: 3958.0,
|
||||
torch.float8_e5m2: 3958.0,
|
||||
torch.float8_e5m2fnuz: 3958.0,
|
||||
torch.float8_e8m0fnu: 3958.0,
|
||||
torch.int8: 3958.0,
|
||||
},
|
||||
dram_bw_gbs=3350,
|
||||
dram_gb=80,
|
||||
),
|
||||
# Source:
|
||||
# @lint-ignore https://www.nvidia.com/content/dam/en-zz/Solutions/Data-Center/a100/pdf/
|
||||
# nvidia-a100-datasheet-us-nvidia-1758950-r4-web.pdf
|
||||
"NVIDIA A100": DeviceInfo(
|
||||
tops={
|
||||
torch.float64: 19.5,
|
||||
torch.float32: 19.5,
|
||||
torch.bfloat16: 312.5,
|
||||
torch.float16: 312.5,
|
||||
# Not in datasheet: float8
|
||||
torch.int8: 624.0,
|
||||
"torch.tf32": 156.0,
|
||||
},
|
||||
dram_bw_gbs=2039.0,
|
||||
dram_gb=80.0,
|
||||
),
|
||||
# Source:
|
||||
# @lint-ignore https://resources.nvidia.com/en-us-gpu-resources/l4-tensor-datasheet
|
||||
"NVIDIA L4": DeviceInfo(
|
||||
tops={
|
||||
# This is a guess, not in datasheet
|
||||
torch.float64: 15.1,
|
||||
torch.float32: 30.3,
|
||||
"torch.tf32": 120.0,
|
||||
torch.bfloat16: 242.0,
|
||||
torch.float16: 242.0,
|
||||
torch.float8_e8m0fnu: 485.0,
|
||||
torch.float8_e8m0fnu: 485.0,
|
||||
torch.float8_e4m3fnuz: 485.0,
|
||||
torch.float8_e5m2: 485.0,
|
||||
torch.float8_e5m2fnuz: 485.0,
|
||||
torch.float8_e8m0fnu: 485.0,
|
||||
torch.int8: 485.0,
|
||||
},
|
||||
dram_bw_gbs=3350,
|
||||
dram_gb=24,
|
||||
),
|
||||
# Source:
|
||||
# @lint-ignore https://www.amd.com/content/dam/amd/en/documents\
|
||||
# /instinct-tech-docs/product-briefs/amd-instinct-mi350x-gpu-brochure.pdf
|
||||
"AMD MI350X": DeviceInfo(
|
||||
tops={
|
||||
torch.float64: 72.1,
|
||||
torch.float32: 144.2,
|
||||
# not specified, fall back to float32 numbers
|
||||
"torch.tf32": 144.2,
|
||||
torch.bfloat16: 2309.6,
|
||||
torch.float16: 2309.6,
|
||||
torch.float8_e8m0fnu: 4614.0,
|
||||
torch.float8_e8m0fnu: 4614.0,
|
||||
torch.float8_e4m3fnuz: 4614.0,
|
||||
torch.float8_e5m2: 4614.0,
|
||||
torch.float8_e5m2fnuz: 4614.0,
|
||||
torch.float8_e8m0fnu: 4614.0,
|
||||
torch.int8: 4614.0,
|
||||
},
|
||||
dram_bw_gbs=8000.0,
|
||||
dram_gb=288.0,
|
||||
),
|
||||
# Source:
|
||||
# @lint-ignore https://www.amd.com/content/dam/amd/en/documents\
|
||||
# /instinct-tech-docs/data-sheets/amd-instinct-mi300a-data-sheet.pdf
|
||||
"AMD MI300A": DeviceInfo(
|
||||
tops={
|
||||
torch.float64: 122.6,
|
||||
torch.float32: 122.6,
|
||||
"torch.tf32": 490.3,
|
||||
torch.bfloat16: 980.6,
|
||||
torch.float16: 980.6,
|
||||
torch.float8_e8m0fnu: 1961.2,
|
||||
torch.float8_e8m0fnu: 1961.2,
|
||||
torch.float8_e4m3fnuz: 1961.2,
|
||||
torch.float8_e5m2: 1961.2,
|
||||
torch.float8_e5m2fnuz: 1961.2,
|
||||
torch.float8_e8m0fnu: 1961.2,
|
||||
torch.int8: 1961.2,
|
||||
},
|
||||
dram_bw_gbs=5300.0,
|
||||
dram_gb=128.0,
|
||||
),
|
||||
# Source:
|
||||
# @lint-ignore https://www.amd.com/content/dam/amd/en/documents/\
|
||||
# instinct-tech-docs/data-sheets/amd-instinct-mi300x-data-sheet.pdf
|
||||
"AMD MI300X": DeviceInfo(
|
||||
tops={
|
||||
torch.float64: 163.4,
|
||||
torch.float32: 163.4,
|
||||
"torch.tf32": 653.7,
|
||||
torch.bfloat16: 1307.4,
|
||||
torch.float16: 1307.4,
|
||||
torch.float8_e8m0fnu: 2614.9,
|
||||
torch.float8_e8m0fnu: 2614.9,
|
||||
torch.float8_e4m3fnuz: 2614.9,
|
||||
torch.float8_e5m2: 2614.9,
|
||||
torch.float8_e5m2fnuz: 2614.9,
|
||||
torch.float8_e8m0fnu: 2614.9,
|
||||
torch.int8: 2614.9,
|
||||
},
|
||||
dram_bw_gbs=5300.0,
|
||||
dram_gb=192.0,
|
||||
),
|
||||
# Source:
|
||||
# @lint-ignore https://www.amd.com/content/dam/amd/\
|
||||
# en/documents/instinct-business-docs/product-briefs/instinct-mi210-brochure.pdf
|
||||
"AMD MI210X": DeviceInfo(
|
||||
tops={
|
||||
torch.float64: 45.3,
|
||||
torch.float32: 45.3,
|
||||
# not specified, fall back to float32 numbers
|
||||
"torch.tf32": 45.3,
|
||||
torch.bfloat16: 181.0,
|
||||
torch.float16: 181.0,
|
||||
# not specified, fall back to float16 numbers
|
||||
torch.float8_e8m0fnu: 181.0,
|
||||
torch.float8_e8m0fnu: 181.0,
|
||||
torch.float8_e4m3fnuz: 181.0,
|
||||
torch.float8_e5m2: 181.0,
|
||||
torch.float8_e5m2fnuz: 181.0,
|
||||
torch.float8_e8m0fnu: 181.0,
|
||||
torch.int8: 181.0,
|
||||
},
|
||||
# pcie4.0x16
|
||||
dram_bw_gbs=1600.0,
|
||||
dram_gb=64.0,
|
||||
),
|
||||
}
|
||||
_device_mapping["AMD INSTINCT MI350X"] = _device_mapping["AMD MI350X"]
|
||||
_device_mapping["AMD INSTINCT MI300X"] = _device_mapping["AMD MI300X"]
|
||||
_device_mapping["AMD INSTINCT MI210X"] = _device_mapping["AMD MI210X"]
|
||||
|
||||
|
||||
def lookup_device_info(name: str) -> DeviceInfo | None:
|
||||
"""
|
||||
Problem: when diffing profiles between amd and nvidia, we don't have access to the device information
|
||||
of the other one. Also, since the analysis is static, we should be able to do it on another device unrelated
|
||||
to the recorded device. Therefore, _device_mapping statically contains the information for lots of devices.
|
||||
If one is missing, please run DeviceInfo.get_device_info() and add it to _device_mapping.
|
||||
name (str): name of the device to lookup. Should map onto torch.cuda.get_device_name().
|
||||
"""
|
||||
return _device_mapping.get(name)
|
||||
|
||||
|
||||
def datasheet_tops(dtype: torch.dtype, is_tf32: bool = False) -> float | None:
|
||||
"""
|
||||
Get the theoretical TFLOPS of the device for a given dtype. This can throw an exception if the device
|
||||
is not in the datasheet list above.
|
||||
"""
|
||||
name: str | None = torch.cuda.get_device_name()
|
||||
if name is None:
|
||||
log.info("No device found, returning None")
|
||||
return None
|
||||
device_info = lookup_device_info(name)
|
||||
if device_info is None:
|
||||
log_str = f"Device {name} not in datasheet, returning None"
|
||||
log.info(log_str)
|
||||
return None
|
||||
if dtype not in device_info.tops:
|
||||
log.info(
|
||||
"Device %s does not have a datasheet entry for %s, returning None",
|
||||
name,
|
||||
dtype,
|
||||
)
|
||||
return None
|
||||
|
||||
return device_info.tops[
|
||||
"torch.tf32" if dtype == torch.float32 and is_tf32 else dtype
|
||||
]
|
||||
+823
@@ -0,0 +1,823 @@
|
||||
import json
|
||||
import logging
|
||||
import math
|
||||
from collections import defaultdict
|
||||
from collections.abc import Callable
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Optional
|
||||
|
||||
import torch
|
||||
from torch._inductor.analysis.device_info import DeviceInfo, lookup_device_info
|
||||
from torch._inductor.utils import tabulate_2d, zip_dicts
|
||||
from torch.utils import _pytree as pytree
|
||||
from torch.utils._ordered_set import OrderedSet
|
||||
from torch.utils.flop_counter import flop_registry
|
||||
|
||||
|
||||
log = logging.getLogger(__name__)
|
||||
|
||||
|
||||
ATEN_PREFIX = "aten::"
|
||||
|
||||
|
||||
@dataclass
|
||||
class ProfileEvent:
|
||||
category: str
|
||||
key: str
|
||||
self_device_time_ms: float
|
||||
# the benchmark is run multiple times and we average the count across all the
|
||||
# runs. It should be an integer but define a float just in case.
|
||||
count: float
|
||||
|
||||
|
||||
# adapters convert the json trace into a format that works with flops_counter
|
||||
ArgsType = tuple[tuple[Any, ...], dict[Any, Any]]
|
||||
AdapterType = Callable[[tuple[Any, ...], tuple[Any, ...]], ArgsType]
|
||||
adapters_map: dict[str, AdapterType] = {}
|
||||
|
||||
|
||||
def parse_list(lst: str) -> list[int]:
|
||||
lst = lst.replace("[", "").replace("]", "")
|
||||
substrings = lst.split(",")
|
||||
|
||||
return [int(substring.strip()) for substring in substrings]
|
||||
|
||||
|
||||
def register_adapter(
|
||||
aten: str | list[str],
|
||||
) -> Callable[
|
||||
[AdapterType],
|
||||
AdapterType,
|
||||
]:
|
||||
def decorator(func: AdapterType) -> AdapterType:
|
||||
# pyrefly: ignore [unknown-name]
|
||||
global _adapters_map
|
||||
|
||||
if isinstance(aten, str):
|
||||
adapters_map[aten] = func
|
||||
else:
|
||||
for at in aten:
|
||||
adapters_map[at] = func
|
||||
return func
|
||||
|
||||
return decorator
|
||||
|
||||
|
||||
@register_adapter(["_slow_conv2d_forward"])
|
||||
def _slow_conv2d_adapter(
|
||||
shapes: tuple[Any, ...], concrete: tuple[Any, ...]
|
||||
) -> tuple[tuple[Any], dict[Any, Any]]:
|
||||
tmp = list(shapes)
|
||||
tmp.append(False)
|
||||
tmp2 = list(concrete)
|
||||
if len(tmp2) < 5:
|
||||
raise ParseException("slow conv2d has less than 5 concrete inputs")
|
||||
tmp2[3] = tmp2[4]
|
||||
return conv_adapter(tuple(tmp), tuple(tmp2))
|
||||
|
||||
|
||||
@register_adapter(
|
||||
["convolution", "_convolution", "cudnn_convolution", "convolution_overrideable"]
|
||||
)
|
||||
def conv_adapter(
|
||||
shapes: tuple[Any, ...], concrete: tuple[Any, ...]
|
||||
) -> tuple[tuple[Any], dict[Any, Any]]:
|
||||
tmp = list(shapes)
|
||||
if len(tmp) == 4:
|
||||
transposed = False
|
||||
elif len(tmp) > 6:
|
||||
transposed = bool(tmp[6])
|
||||
tmp[6] = transposed
|
||||
else:
|
||||
raise ParseException(f"Convolution has the wrong number of inputs: {len(tmp)}")
|
||||
|
||||
kwargs: dict[Any, Any] = {}
|
||||
if not transposed:
|
||||
# calculate output shape if not transposed.
|
||||
def conv_out_dims(x: int, kernel: int, stride: int) -> int:
|
||||
return (x - kernel) // stride + 1
|
||||
|
||||
stride = parse_list(concrete[3])
|
||||
inp = shapes[0]
|
||||
w = shapes[1]
|
||||
out_x_y = [conv_out_dims(*args) for args in zip(inp[2:], w[2:], stride)]
|
||||
out = [inp[0], w[0]] + out_x_y # we only need the xy values
|
||||
kwargs["out_val"] = out
|
||||
|
||||
return tuple(tmp), kwargs
|
||||
|
||||
|
||||
def default_adapter(
|
||||
shapes: tuple[Any], concrete: tuple[Any]
|
||||
) -> tuple[tuple[Any], dict[Any, Any]]:
|
||||
return shapes, {}
|
||||
|
||||
|
||||
@register_adapter("addmm")
|
||||
def addmm_adapter(
|
||||
shapes: tuple[Any], concrete: tuple[Any]
|
||||
) -> tuple[tuple[Any], dict[Any, Any]]:
|
||||
tmp = list(shapes)[:3]
|
||||
return tuple(tmp), {}
|
||||
|
||||
|
||||
@register_adapter("bmm")
|
||||
def bmm_adapter(
|
||||
shapes: tuple[Any], concrete: tuple[Any]
|
||||
) -> tuple[tuple[Any], dict[Any, Any]]:
|
||||
tmp = list(shapes)
|
||||
return tuple(tmp[:2]), {}
|
||||
|
||||
|
||||
@register_adapter("baddbmm")
|
||||
def baddbmm_adapter(
|
||||
shapes: tuple[Any], concrete: tuple[Any]
|
||||
) -> tuple[tuple[Any], dict[Any, Any]]:
|
||||
tmp = list(shapes)[:3]
|
||||
return tuple(tmp), {}
|
||||
|
||||
|
||||
@register_adapter("mm")
|
||||
def mm_adapter(
|
||||
shapes: tuple[Any], concrete: tuple[Any]
|
||||
) -> tuple[tuple[Any], dict[Any, Any]]:
|
||||
return shapes, {}
|
||||
|
||||
|
||||
def _parse_kernel_name(name: str) -> str | None:
|
||||
"""
|
||||
parse the name of the kernel from the event name.
|
||||
"""
|
||||
if name.startswith(ATEN_PREFIX):
|
||||
return name[len(ATEN_PREFIX) :]
|
||||
elif "conv" in name:
|
||||
return "convolution"
|
||||
elif "addmm" in name:
|
||||
return "addmm"
|
||||
elif "bmm" in name:
|
||||
return "bmm"
|
||||
elif "baddbmm" in name:
|
||||
return "baddbmm"
|
||||
elif "_mm" in name:
|
||||
return "mm"
|
||||
else:
|
||||
return None
|
||||
|
||||
|
||||
def _calculate_flops(event: dict[str, Any]) -> int:
|
||||
"""
|
||||
This function has to parse the kernel name, which is error prone. There doesn't seem to be another solution that
|
||||
will support all the different backends that can generate kernels, so make sure to update this function when new
|
||||
ops and backends are desired.
|
||||
"""
|
||||
name = event["name"]
|
||||
if "kernel_flop" in event["args"] and event["args"]["kernel_flop"] != 0:
|
||||
return event["args"]["kernel_flop"]
|
||||
op_name = _parse_kernel_name(name)
|
||||
if op_name is None:
|
||||
return 0
|
||||
|
||||
op_obj = getattr(torch.ops.aten, op_name, None)
|
||||
if op_obj is None or op_obj not in flop_registry:
|
||||
return 0
|
||||
|
||||
flop_function = flop_registry[op_obj]
|
||||
|
||||
if "Input Dims" not in event["args"] or "Concrete Inputs" not in event["args"]:
|
||||
return 0
|
||||
input_shapes = event["args"]["Input Dims"]
|
||||
concrete = event["args"]["Concrete Inputs"]
|
||||
if op_name in adapters_map:
|
||||
try:
|
||||
args, kwargs = adapters_map[op_name](input_shapes, concrete)
|
||||
except ParseException as e:
|
||||
msg = f"Failed to parse {op_name} with {e}"
|
||||
log.warning(msg)
|
||||
return 0
|
||||
else:
|
||||
try:
|
||||
args, kwargs = default_adapter(input_shapes, concrete)
|
||||
except ParseException as e:
|
||||
msg = f"Failed to parse {op_name} with {e}"
|
||||
log.warning(msg)
|
||||
return 0
|
||||
return flop_function(*args, **kwargs)
|
||||
|
||||
|
||||
def _get_size_from_string(type_string: str) -> int:
|
||||
if not hasattr(torch, type_string):
|
||||
return 1
|
||||
else:
|
||||
return getattr(torch, type_string).itemsize
|
||||
|
||||
|
||||
def _default_estimate_gb(event: dict[str, Any]) -> float:
|
||||
sizes_and_types = zip(event["args"]["Input Dims"], event["args"]["Input type"])
|
||||
bw = 0
|
||||
for size, typ in sizes_and_types:
|
||||
isize = _get_size_from_string(typ)
|
||||
bw += isize * math.prod(pytree.tree_flatten(size)[0])
|
||||
return bw / 1e9
|
||||
|
||||
|
||||
def _estimate_gb(event: dict[str, Any]) -> float:
|
||||
"""
|
||||
Our best effort to estimate the gb, should be refactored soon with MemoryCounter.
|
||||
"""
|
||||
name = event["name"]
|
||||
if "kernel_num_gb" in event["args"] and event["args"]["kernel_num_gb"] != 0:
|
||||
return event["args"]["kernel_num_gb"]
|
||||
if "Input type" not in event["args"] or "Input Dims" not in event["args"]:
|
||||
return 0
|
||||
op_name = _parse_kernel_name(name)
|
||||
if op_name is None:
|
||||
return _default_estimate_gb(event)
|
||||
|
||||
op_obj = getattr(torch.ops.aten, op_name, None)
|
||||
if op_obj is None:
|
||||
return _default_estimate_gb(event)
|
||||
|
||||
if "Input Dims" not in event["args"] or "Concrete Inputs" not in event["args"]:
|
||||
return _default_estimate_gb(event)
|
||||
input_shapes = event["args"]["Input Dims"]
|
||||
|
||||
# NOTE these will be refactored into a similar object to FlopCounter soon
|
||||
def mm_formula(M: int, N: int, K: int, size: int) -> int:
|
||||
return 2 * (M * K + N * K + M * N) * size
|
||||
|
||||
if op_name == "addmm":
|
||||
add_in_size = math.prod(pytree.tree_flatten(input_shapes[0])[0])
|
||||
add_type_size = _get_size_from_string(event["args"]["Input type"][0])
|
||||
M = input_shapes[1][0]
|
||||
N = input_shapes[1][1]
|
||||
assert input_shapes[1][1] == input_shapes[2][0]
|
||||
K = input_shapes[2][1]
|
||||
mul_type_size = _get_size_from_string(event["args"]["Input type"][1])
|
||||
return (mm_formula(M, N, K, mul_type_size) + add_in_size * add_type_size) / 1e9
|
||||
elif op_name == "mm":
|
||||
M = input_shapes[0][0]
|
||||
N = input_shapes[0][1]
|
||||
assert input_shapes[0][1] == input_shapes[1][0]
|
||||
K = input_shapes[1][1]
|
||||
type_size = _get_size_from_string(event["args"]["Input type"][0])
|
||||
return mm_formula(M, N, K, type_size) / 1e9
|
||||
elif op_name == "baddbmm":
|
||||
add_in_size = math.prod(pytree.tree_flatten(input_shapes[0])[0])
|
||||
add_type_size = _get_size_from_string(event["args"]["Input type"][0])
|
||||
B = input_shapes[0][0]
|
||||
M = input_shapes[1][1]
|
||||
N = input_shapes[1][2]
|
||||
K = input_shapes[2][2]
|
||||
mul_type_size = _get_size_from_string(event["args"]["Input type"][1])
|
||||
return (
|
||||
B * mm_formula(M, N, K, mul_type_size) + add_in_size * add_type_size
|
||||
) / 1e9
|
||||
elif op_name == "bmm":
|
||||
add_in_size = math.prod(pytree.tree_flatten(input_shapes[0])[0])
|
||||
add_type_size = _get_size_from_string(event["args"]["Input type"][0])
|
||||
B = input_shapes[0][0]
|
||||
M = input_shapes[0][1]
|
||||
N = input_shapes[0][2]
|
||||
K = input_shapes[1][2]
|
||||
mul_type_size = _get_size_from_string(event["args"]["Input type"][1])
|
||||
return (
|
||||
B * mm_formula(M, N, K, mul_type_size) + add_in_size * add_type_size
|
||||
) / 1e9
|
||||
elif op_name in [
|
||||
"convolution",
|
||||
"_convolution",
|
||||
"cudnn_convolution",
|
||||
"_slow_conv2d_forward",
|
||||
]:
|
||||
concrete = event["args"]["Concrete Inputs"]
|
||||
|
||||
def conv_out_dim(x: int, kernel: int, stride: int) -> int:
|
||||
return (x - kernel) // stride + 1
|
||||
|
||||
stride = parse_list(
|
||||
concrete[3] if op_name != "_slow_conv2d_forward" else concrete[4]
|
||||
)
|
||||
inp = input_shapes[0]
|
||||
w = input_shapes[1]
|
||||
out_x_y = [conv_out_dim(*args) for args in zip(inp[2:], w[2:], stride)]
|
||||
out = [inp[0], w[0]] + out_x_y
|
||||
# each output element reads in * w * w chunk
|
||||
input_reads = out[0] * out[1] * out[2] * out[3] * inp[1] * w[2] * w[3]
|
||||
# Assume weights are in cache, so only read once
|
||||
weight_reads = w[0] * w[1] * w[2] * w[3]
|
||||
return (input_reads + weight_reads) / 1e9
|
||||
|
||||
return _default_estimate_gb(event)
|
||||
|
||||
|
||||
def _create_extern_mapping(
|
||||
data: dict[str, Any],
|
||||
) -> defaultdict[int, list[dict[str, Any]]]:
|
||||
"""
|
||||
compute a mapping from external ids to non kernels, which contain the information we need to estimate flops etc
|
||||
"""
|
||||
extern_mapping: defaultdict[int, list[dict[str, Any]]] = defaultdict(list)
|
||||
for event in data["traceEvents"]:
|
||||
if (
|
||||
"args" not in event
|
||||
or "External id" not in event["args"]
|
||||
or event["cat"] != "cpu_op"
|
||||
):
|
||||
continue
|
||||
if len(extern_mapping[event["args"]["External id"]]) > 0:
|
||||
raise ParseException("duplicate external id in event")
|
||||
extern_mapping[event["args"]["External id"]].append(event)
|
||||
return extern_mapping
|
||||
|
||||
|
||||
def _augment_trace_helper(data: dict[str, Any]) -> dict[str, Any]:
|
||||
extern_mapping = _create_extern_mapping(data)
|
||||
|
||||
for event in data["traceEvents"]:
|
||||
if "cat" not in event or event["cat"] != "kernel":
|
||||
continue
|
||||
if "args" not in event:
|
||||
raise ParseException(f"kernel has no args: {event}")
|
||||
if "External id" not in event["args"]:
|
||||
event_str = f"kernel has no External id: {event}"
|
||||
log.info(event_str)
|
||||
continue
|
||||
|
||||
external_op = extern_mapping[event["args"]["External id"]][0]
|
||||
flops = _calculate_flops(external_op)
|
||||
if flops == 0:
|
||||
flops = _calculate_flops(event)
|
||||
external_op["args"]["kernel_flop"] = flops
|
||||
external_op["args"]["kernel_num_gb"] = _estimate_gb(external_op)
|
||||
event["args"]["kernel_flop"] = external_op["args"]["kernel_flop"]
|
||||
event["args"]["kernel_num_gb"] = external_op["args"]["kernel_num_gb"]
|
||||
return data
|
||||
|
||||
|
||||
_dtype_map = {
|
||||
"float": torch.float,
|
||||
"float32": torch.float,
|
||||
"int": torch.int,
|
||||
"int8": torch.int8,
|
||||
"int16": torch.int16,
|
||||
"int32": torch.int,
|
||||
"long": torch.long,
|
||||
"long int": torch.long,
|
||||
"bfloat16": torch.bfloat16,
|
||||
"float16": torch.float16,
|
||||
"float64": torch.double,
|
||||
}
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class KernelStats:
|
||||
flops: int
|
||||
bw: float
|
||||
latency: float # us
|
||||
achieved_flops: float
|
||||
achieved_bandwidth: float
|
||||
|
||||
|
||||
KernelNameMap = defaultdict[str, OrderedSet[KernelStats]]
|
||||
|
||||
|
||||
@dataclass(frozen=False)
|
||||
class Device:
|
||||
name: str
|
||||
index: int
|
||||
info: DeviceInfo | None
|
||||
stats: KernelNameMap
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return f"Device({self.name}, {self.index}): {self.info}"
|
||||
|
||||
|
||||
DeviceMap = dict[int, Device]
|
||||
Table = tuple[list[str], dict[str, list[str]]]
|
||||
|
||||
|
||||
class JsonProfile:
|
||||
_devices: DeviceMap
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
path: str,
|
||||
benchmark_name: str | None = None,
|
||||
dtype: torch.dtype | str | None = None,
|
||||
):
|
||||
"""
|
||||
Convenience class for running common operations on chrome/perfetto json traces.
|
||||
"""
|
||||
self.path = path
|
||||
with open(path) as f:
|
||||
self.data = json.load(f)
|
||||
self.events = self.data["traceEvents"]
|
||||
self.benchmark_name = benchmark_name
|
||||
if dtype is None:
|
||||
self.dtype = None
|
||||
elif isinstance(dtype, torch.dtype):
|
||||
# pyrefly: ignore [bad-assignment]
|
||||
self.dtype = dtype
|
||||
else:
|
||||
# pyrefly: ignore [bad-assignment]
|
||||
self.dtype = _dtype_map.get(dtype)
|
||||
self._create_devices()
|
||||
|
||||
def convert_dtype(self, event: dict[str, Any]) -> torch.dtype | None:
|
||||
"""
|
||||
Each op has a list of dtypes for each input arg. We need to convert these into a single dtype for flop estimation.
|
||||
Issues:
|
||||
- converting the strings to concrete torch.dtypes
|
||||
- What if we have float32, float, float16 all in the inputs? Our choice is to use the largest buffer dtype.
|
||||
"""
|
||||
|
||||
if (
|
||||
"Input Dims" not in event["args"]
|
||||
or "Input type" not in event["args"]
|
||||
or "Concrete Inputs" not in event["args"]
|
||||
):
|
||||
if "bfloat16" in event["name"]:
|
||||
return torch.bfloat16
|
||||
elif "float16" in event["name"]:
|
||||
return torch.float16
|
||||
else:
|
||||
return None
|
||||
|
||||
input_sizes = event["args"]["Input Dims"]
|
||||
input_types = event["args"]["Input type"]
|
||||
concrete_inputs = event["args"]["Concrete Inputs"]
|
||||
assert len(input_sizes) == len(input_types)
|
||||
assert len(input_types) == len(concrete_inputs)
|
||||
|
||||
if len(input_sizes) == 0:
|
||||
raise RuntimeError("Empty input_sizes and input_types")
|
||||
|
||||
biggest_size = 0
|
||||
biggest_index = 0
|
||||
for i in range(len(input_sizes)):
|
||||
if concrete_inputs[i] != "":
|
||||
# concrete inputs are usually small tensors, so we can just skip
|
||||
continue
|
||||
my_size = input_sizes[i]
|
||||
total_size = sum(parse_list(my_size))
|
||||
if total_size > biggest_size:
|
||||
biggest_size = total_size
|
||||
biggest_index = i
|
||||
ret_type = input_types[biggest_index]
|
||||
if ret_type in _dtype_map:
|
||||
return _dtype_map[ret_type]
|
||||
raise RuntimeError(f"Unknown type: {ret_type}. Please add to _dtype_map.")
|
||||
|
||||
def _create_devices(self) -> None:
|
||||
self._devices = {}
|
||||
for dev in self.data["deviceProperties"]:
|
||||
name = dev["name"]
|
||||
device_info = lookup_device_info(name)
|
||||
|
||||
if device_info is None:
|
||||
log.info(
|
||||
"Unsupported device in profile: %s, please consider contributing to _device_mapping.",
|
||||
name,
|
||||
)
|
||||
self._devices[dev["id"]] = Device(
|
||||
name, dev["id"], device_info, defaultdict(OrderedSet)
|
||||
)
|
||||
|
||||
def calculate_flops(self, event: dict[str, Any]) -> int:
|
||||
return _calculate_flops(event)
|
||||
|
||||
def estimate_gb(self, event: dict[str, Any]) -> float:
|
||||
return _estimate_gb(event)
|
||||
|
||||
def augment_trace(self) -> None:
|
||||
self.data = _augment_trace_helper(self.data)
|
||||
|
||||
def _compute_stats(self) -> None:
|
||||
"""populates the name -> stats map"""
|
||||
for event in self.events:
|
||||
if "cat" not in event or "args" not in event or event["cat"] != "kernel":
|
||||
continue
|
||||
if "device" not in event["args"]:
|
||||
continue
|
||||
dev_tmp = event["args"]["device"]
|
||||
if dev_tmp not in self._devices:
|
||||
continue
|
||||
dev = self._devices[event["args"]["device"]]
|
||||
|
||||
dur = event["dur"] # us
|
||||
if "kernel_flop" in event["args"]:
|
||||
assert dur != 0
|
||||
# 1,000,000us/s * flop / us
|
||||
op_flops = event["args"]["kernel_flop"] / (dur / 1e6)
|
||||
else:
|
||||
op_flops = 0
|
||||
|
||||
if "kernel_num_gb" in event["args"]:
|
||||
assert dur != 0
|
||||
# 1,000,000us/s * gb = gb/s
|
||||
op_gbps = event["args"]["kernel_num_gb"] / (dur / 1e6)
|
||||
else:
|
||||
op_gbps = 0
|
||||
|
||||
if dev.info is not None:
|
||||
dtype = self.convert_dtype(event) or self.dtype
|
||||
if dtype is None:
|
||||
raise RuntimeError(
|
||||
"dtype is not found on tensor and default dtype is not set"
|
||||
)
|
||||
achieved_flops = 100 * op_flops / (1e12 * dev.info.tops[dtype])
|
||||
achieved_bandwidth = 100 * op_gbps / dev.info.dram_bw_gbs
|
||||
else:
|
||||
achieved_flops = 0
|
||||
achieved_bandwidth = 0
|
||||
|
||||
if "name" not in event["args"]:
|
||||
continue
|
||||
dev.stats[event["name"]].add(
|
||||
KernelStats(
|
||||
flops=op_flops,
|
||||
bw=op_gbps,
|
||||
latency=dur,
|
||||
achieved_bandwidth=achieved_bandwidth,
|
||||
achieved_flops=achieved_flops,
|
||||
)
|
||||
)
|
||||
|
||||
def _create_single_table(self, dev: Device) -> Table:
|
||||
"""Create a table with the devices mapped to indices."""
|
||||
headers = [
|
||||
"Kernel Name",
|
||||
"Kernel Count",
|
||||
"FLOPS",
|
||||
"Kernel Reads (GB)",
|
||||
"Dur (us)",
|
||||
"Achieved FLOPS %",
|
||||
"Achieved Bandwidth %",
|
||||
]
|
||||
rows: dict[str, list[str]] = {}
|
||||
|
||||
def safe_div_format(x: float, y: float) -> str:
|
||||
if y == 0:
|
||||
return "0.0"
|
||||
return f"{x / y:.4f}"
|
||||
|
||||
for kernel_name, stats_set in dev.stats.items():
|
||||
ker_count = 0
|
||||
flops = 0
|
||||
flops_count = 0
|
||||
achieved_flops = 0.0
|
||||
bw = 0.0
|
||||
bw_count = 0
|
||||
achieved_bandwidth = 0.0
|
||||
latency = 0.0
|
||||
for stats in stats_set:
|
||||
if stats.flops != 0:
|
||||
flops += stats.flops
|
||||
achieved_flops += stats.achieved_flops
|
||||
flops_count += 1
|
||||
if stats.bw != 0:
|
||||
bw += stats.bw
|
||||
achieved_bandwidth += stats.achieved_bandwidth
|
||||
bw_count += 1
|
||||
latency += stats.latency
|
||||
ker_count += 1
|
||||
assert ker_count != 0
|
||||
rows[kernel_name] = [
|
||||
str(ker_count),
|
||||
safe_div_format(flops, flops_count),
|
||||
safe_div_format(bw, bw_count),
|
||||
safe_div_format(latency, ker_count),
|
||||
safe_div_format(achieved_flops, flops_count),
|
||||
safe_div_format(achieved_bandwidth, bw_count),
|
||||
]
|
||||
|
||||
return headers, rows
|
||||
|
||||
def _create_tables(self, devs: DeviceMap) -> dict[int, Table]:
|
||||
return {idx: self._create_single_table(dev) for idx, dev in devs.items()}
|
||||
|
||||
def _combine_tables(
|
||||
self, table1: Table, table1_name: str, table2: Table, table2_name: str
|
||||
) -> Table:
|
||||
new_headers = (
|
||||
["Kernel Name"]
|
||||
+ [f"{table1_name} {head}" for head in table1[0][1:]]
|
||||
+ [f"{table2_name} {head}" for head in table2[0][1:]]
|
||||
)
|
||||
t1_length = len(table1[0][1:])
|
||||
t2_length = len(table2[0][1:])
|
||||
new_rows = {}
|
||||
|
||||
for key, row1, row2 in zip_dicts(
|
||||
table1[1],
|
||||
table2[1],
|
||||
d1_default=["Empty"] * t1_length,
|
||||
d2_default=["Empty"] * t2_length,
|
||||
):
|
||||
assert row1 is not None
|
||||
assert row2 is not None
|
||||
new_rows[key] = row1 + row2
|
||||
return new_headers, new_rows
|
||||
|
||||
def report(
|
||||
self, other: Optional["JsonProfile"] = None, name_limit: int = 40
|
||||
) -> str:
|
||||
def create_ret(
|
||||
table_headers: list[str], table_rows: dict[str, list[str]]
|
||||
) -> str:
|
||||
table_flattened = [
|
||||
[kernel_name[:name_limit], *kernel_vals]
|
||||
for kernel_name, kernel_vals in table_rows.items()
|
||||
]
|
||||
return tabulate_2d(table_flattened, headers=table_headers)
|
||||
|
||||
if other is not None:
|
||||
self._compute_stats()
|
||||
other._compute_stats()
|
||||
|
||||
self_tables = self._create_tables(self._devices)
|
||||
other_tables = self._create_tables(other._devices)
|
||||
|
||||
self_name = (
|
||||
self.benchmark_name if self.benchmark_name is not None else "Table 1"
|
||||
)
|
||||
other_name = (
|
||||
other.benchmark_name if other.benchmark_name is not None else "Table 2"
|
||||
)
|
||||
|
||||
ret = []
|
||||
assert self._devices.keys() == other._devices.keys()
|
||||
for device_idx, t1, t2 in zip_dicts(
|
||||
self_tables, other_tables, d1_default=None, d2_default=None
|
||||
):
|
||||
assert t1 is not None
|
||||
assert t2 is not None
|
||||
table_headers, table_rows = self._combine_tables(
|
||||
t1, self_name, t2, other_name
|
||||
)
|
||||
tab_string = create_ret(table_headers, table_rows)
|
||||
# pyrefly: ignore [bad-argument-type]
|
||||
ret.append(f"{self._devices[device_idx]}:\n{tab_string}")
|
||||
return "\n".join(ret)
|
||||
self._compute_stats()
|
||||
|
||||
self_tables = self._create_tables(self._devices)
|
||||
|
||||
ret = []
|
||||
for idx, table in self_tables.items():
|
||||
table_headers, table_rows = table
|
||||
tab_string = create_ret(table_headers, table_rows)
|
||||
# pyrefly: ignore [bad-argument-type]
|
||||
ret.append(f"{self._devices[idx]}:\n{tab_string}")
|
||||
return "\n".join(ret)
|
||||
|
||||
def dump(self, out: str) -> None:
|
||||
with open(out, "w") as f:
|
||||
json.dump(self.data, f)
|
||||
|
||||
def combine_with(self, other: "JsonProfile") -> "JsonProfile":
|
||||
"""
|
||||
Combine this profile with another profile by merging their trace events.
|
||||
Returns a new JsonProfile object with combined data.
|
||||
"""
|
||||
# Create a new combined data structure
|
||||
combined_data = {
|
||||
"traceEvents": self.data["traceEvents"] + other.data["traceEvents"],
|
||||
"deviceProperties": self.data.get("deviceProperties", []),
|
||||
}
|
||||
|
||||
# Merge device properties, avoiding duplicates
|
||||
other_device_props = other.data.get("deviceProperties", [])
|
||||
existing_device_ids = OrderedSet(
|
||||
[dev["id"] for dev in combined_data["deviceProperties"]]
|
||||
)
|
||||
|
||||
for device_prop in other_device_props:
|
||||
if device_prop["id"] not in existing_device_ids:
|
||||
combined_data["deviceProperties"].append(device_prop)
|
||||
|
||||
# Copy any other top-level properties from the first profile
|
||||
for key, value in self.data.items():
|
||||
if key not in combined_data:
|
||||
combined_data[key] = value
|
||||
|
||||
import os
|
||||
|
||||
# Create a temporary file to write the combined data
|
||||
import tempfile
|
||||
|
||||
with tempfile.NamedTemporaryFile(
|
||||
mode="w", suffix=".json", delete=False
|
||||
) as tmp_file:
|
||||
json.dump(combined_data, tmp_file)
|
||||
tmp_path = tmp_file.name
|
||||
|
||||
try:
|
||||
# Create new JsonProfile from the combined data
|
||||
combined_profile = JsonProfile(
|
||||
tmp_path,
|
||||
benchmark_name=f"{self.benchmark_name or 'Profile1'}_+_{other.benchmark_name or 'Profile2'}",
|
||||
dtype=self.dtype or other.dtype,
|
||||
)
|
||||
return combined_profile
|
||||
finally:
|
||||
# Clean up temporary file
|
||||
os.unlink(tmp_path)
|
||||
|
||||
|
||||
class ParseException(RuntimeError):
|
||||
pass
|
||||
|
||||
|
||||
def main() -> None:
|
||||
"""
|
||||
Main function for the profile analysis script.
|
||||
"""
|
||||
import argparse
|
||||
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument(
|
||||
"--diff",
|
||||
nargs=5,
|
||||
metavar=(
|
||||
"input_file1",
|
||||
"name1",
|
||||
"input_file2",
|
||||
"name2",
|
||||
"dtype",
|
||||
),
|
||||
help="Two json traces to compare with, specified as <file1> <name1> <file2> <name2> <dtype>",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--name_limit",
|
||||
type=int,
|
||||
help="the maximum name size in the final report",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--augment_trace",
|
||||
"-a",
|
||||
nargs=3,
|
||||
metavar=("input_file", "output_file", "dtype"),
|
||||
help="Augment a trace with inductor meta information. Provide input and output file paths.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--analysis",
|
||||
nargs=2,
|
||||
metavar=("input_file", "dtype"),
|
||||
help="Run analysis on a single trace, specified as <file> <dtype>",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--combine",
|
||||
nargs="+",
|
||||
metavar=("input_files", "output_file"),
|
||||
help="Combine multiple profiles into a single profile by merging trace events. Specify as <input_file1> \
|
||||
<input_file2> [input_file3 ...] <output_file>. The last argument is the output file, all preceding arguments are \
|
||||
input files to combine.",
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
if args.diff:
|
||||
p1 = JsonProfile(args.diff[0], args.diff[1], dtype=args.diff[4])
|
||||
p1.augment_trace()
|
||||
p2 = JsonProfile(args.diff[2], args.diff[3], dtype=args.diff[4])
|
||||
p2.augment_trace()
|
||||
if args.name_limit:
|
||||
print(p1.report(p2, name_limit=args.name_limit))
|
||||
else:
|
||||
print(p1.report(p2))
|
||||
if args.analysis:
|
||||
p1 = JsonProfile(
|
||||
args.analysis[0],
|
||||
dtype=args.analysis[1],
|
||||
)
|
||||
p1.augment_trace()
|
||||
if args.name_limit:
|
||||
print(p1.report(name_limit=args.name_limit))
|
||||
else:
|
||||
print(p1.report())
|
||||
if args.augment_trace:
|
||||
p = JsonProfile(args.augment_trace[0], dtype=args.augment_trace[2])
|
||||
p.augment_trace()
|
||||
p.dump(args.augment_trace[1])
|
||||
if args.combine:
|
||||
input_files = args.combine[:-1] # All arguments except the last one
|
||||
output_file = args.combine[-1] # Last argument is the output file
|
||||
|
||||
if len(input_files) < 2:
|
||||
print("Error: At least 2 input files are required for combining")
|
||||
return
|
||||
|
||||
# Load the first profile
|
||||
combined = JsonProfile(input_files[0], dtype=None)
|
||||
|
||||
# Iteratively combine with all other profiles
|
||||
for input_file in input_files[1:]:
|
||||
profile = JsonProfile(input_file, dtype=None)
|
||||
combined = combined.combine_with(profile)
|
||||
|
||||
combined.dump(output_file)
|
||||
print(f"Successfully combined {', '.join(input_files)} into {output_file}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
+166
@@ -0,0 +1,166 @@
|
||||
import dataclasses
|
||||
import itertools
|
||||
from typing import Any, TYPE_CHECKING
|
||||
|
||||
import sympy
|
||||
|
||||
import torch
|
||||
from torch._inductor import config
|
||||
from torch._inductor.dtype_propagation import DtypePropagationOpsHandler
|
||||
from torch._inductor.index_propagation import SymPyOps, TypedExpr
|
||||
|
||||
from .ops_handler import DefaultHandler
|
||||
from .virtualized import StoreMode, V
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from torch._inductor.scheduler import SchedulerNode
|
||||
|
||||
|
||||
def construct_symbol(count: int, dtype: torch.dtype) -> sympy.Symbol:
|
||||
return sympy.Symbol(f"unknown_{count}")
|
||||
|
||||
|
||||
class PreservesZeros(SymPyOps, DefaultHandler):
|
||||
"""
|
||||
For prologue kernels where the loads are masked, does the final store of this kernel preserve
|
||||
the zeros.
|
||||
"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.count = itertools.count(0)
|
||||
self.store_preserves_zeros: bool | None = None
|
||||
self.dtype_prop = DtypePropagationOpsHandler()
|
||||
|
||||
def load(self, name: str, index: sympy.Expr) -> TypedExpr:
|
||||
# In prologue fusion, all loads get broadcasted
|
||||
dtype = self.dtype_prop.load(name, index)
|
||||
return TypedExpr(
|
||||
sympy.Float(0) if dtype.is_floating_point else sympy.Integer(0), dtype
|
||||
)
|
||||
|
||||
def store(
|
||||
self, name: str, index: sympy.Expr, value: TypedExpr, mode: "StoreMode" = None
|
||||
) -> None:
|
||||
assert isinstance(self, PreservesZeros)
|
||||
# should only have a single store in prologue
|
||||
assert self.store_preserves_zeros is None
|
||||
self.store_preserves_zeros = value.is_constant() and value.expr == 0
|
||||
|
||||
def indirect_indexing(self, *args: Any, **kwargs: Any) -> sympy.Expr:
|
||||
return construct_symbol(next(self.count), torch.int32)
|
||||
|
||||
def _default(self, name: str, args: tuple[Any, ...], kwargs: dict[str, Any]) -> Any:
|
||||
from torch._inductor.codegen.common import OpDecompositions
|
||||
|
||||
if hasattr(OpDecompositions, name):
|
||||
return getattr(OpDecompositions, name)(*args, **kwargs).value
|
||||
|
||||
dtype = getattr(self.dtype_prop, name)(*args, **kwargs)
|
||||
return TypedExpr(construct_symbol(next(self.count), dtype), dtype)
|
||||
|
||||
|
||||
def prologue_preserves_zero_mask(prologue: "SchedulerNode") -> bool:
|
||||
"""
|
||||
Does this prologue preserve zero masks
|
||||
"""
|
||||
preserves_zeros = PreservesZeros()
|
||||
with V.set_ops_handler(preserves_zeros):
|
||||
prologue._body(*prologue.get_ranges())
|
||||
|
||||
store_preserves_zeros = preserves_zeros.store_preserves_zeros
|
||||
assert isinstance(store_preserves_zeros, bool)
|
||||
|
||||
return store_preserves_zeros
|
||||
|
||||
|
||||
@dataclasses.dataclass
|
||||
class DTypeContainer:
|
||||
dtype: torch.dtype
|
||||
is_scalar: bool = False
|
||||
|
||||
|
||||
class RecordLowPrecisionOps(DefaultHandler):
|
||||
def __init__(self, disallow_fp32_ops: bool = False) -> None:
|
||||
self.disallow_fp32_ops = disallow_fp32_ops
|
||||
self.low_precision_numeric_op = False
|
||||
self.dtype_prop = DtypePropagationOpsHandler()
|
||||
self.non_numeric_ops = (
|
||||
"to_dtype",
|
||||
"constant",
|
||||
"where",
|
||||
)
|
||||
|
||||
def load(self, name: str, index: sympy.Expr) -> DTypeContainer:
|
||||
return DTypeContainer(self.dtype_prop.load(name, index))
|
||||
|
||||
@staticmethod
|
||||
def store(
|
||||
name: str, index: sympy.Expr, value: TypedExpr, mode: "StoreMode" = None
|
||||
) -> None:
|
||||
pass
|
||||
|
||||
def check_bounds(
|
||||
self, expr: sympy.Expr, size: sympy.Expr, lower: bool, upper: bool
|
||||
) -> None:
|
||||
pass
|
||||
|
||||
@staticmethod
|
||||
# pyrefly: ignore [bad-override]
|
||||
def indirect_indexing(*args: Any, **kwargs: Any) -> sympy.Expr:
|
||||
return sympy.S.Zero
|
||||
|
||||
def _default(self, name: str, args: tuple[Any, ...], kwargs: dict[str, Any]) -> Any:
|
||||
out_dtype = getattr(self.dtype_prop, name)(*args, **kwargs)
|
||||
out = DTypeContainer(out_dtype, is_scalar=(name == "constant"))
|
||||
if name == "constant":
|
||||
return DTypeContainer(torch.float, is_scalar=True)
|
||||
|
||||
uses_low_prec = any(
|
||||
isinstance(dtype_cont, DTypeContainer)
|
||||
and dtype_cont.dtype is not None
|
||||
and low_prec_float(dtype_cont.dtype)
|
||||
for dtype_cont in itertools.chain((out,), args, kwargs.values())
|
||||
)
|
||||
|
||||
if uses_low_prec and name not in self.non_numeric_ops:
|
||||
self.low_precision_numeric_op = True
|
||||
|
||||
if (
|
||||
self.disallow_fp32_ops
|
||||
and out.dtype in (torch.float32, torch.float64)
|
||||
and not out.is_scalar
|
||||
):
|
||||
self.low_precision_numeric_op = True
|
||||
|
||||
return out
|
||||
|
||||
|
||||
def low_prec_float(dtype: torch.dtype) -> bool:
|
||||
return dtype.is_floating_point and dtype.itemsize < 4
|
||||
|
||||
|
||||
def can_codegen_without_upcasts(
|
||||
prologue: "SchedulerNode",
|
||||
disallow_fp32_ops: bool = False,
|
||||
) -> bool:
|
||||
"""
|
||||
Can this prologue be run without `upcast_to_fp32` while preserving numerics.
|
||||
|
||||
This is only true if the node only contains dtype conversions, indexing, and other non-arithmetic operators.
|
||||
|
||||
If disallow_fp32_ops is True, then we also disallow ops that are explicitly computed in fp32 or fp64.
|
||||
"""
|
||||
if prologue.get_operation_names() <= V.graph.low_precision_codegen_ops:
|
||||
return True
|
||||
|
||||
low_prec_analysis = RecordLowPrecisionOps(disallow_fp32_ops)
|
||||
|
||||
# Need to turn off upcasting to do analysis of whether we can turn it off
|
||||
with (
|
||||
config.patch("triton.codegen_upcast_to_fp32", False),
|
||||
V.set_ops_handler(low_prec_analysis),
|
||||
):
|
||||
prologue._body(*prologue.get_ranges())
|
||||
|
||||
return not low_prec_analysis.low_precision_numeric_op
|
||||
@@ -0,0 +1,340 @@
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
from collections.abc import Callable
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
from unittest import mock
|
||||
|
||||
import torch
|
||||
import torch._export
|
||||
from torch._inductor.utils import is_cpu_device
|
||||
|
||||
from .runtime.runtime_utils import cache_dir
|
||||
|
||||
|
||||
log = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@dataclass
|
||||
class AOTICompileBackend:
|
||||
compile_fn: Callable[..., str]
|
||||
load_fn: Callable[[str, str, str], list[dict[str, Any] | None]]
|
||||
|
||||
|
||||
_aoti_compile_backends: dict[str, AOTICompileBackend] = {}
|
||||
|
||||
|
||||
def register_aoti_compile_backend(
|
||||
device_type: str,
|
||||
compile_fn: Callable[..., str],
|
||||
load_fn: Callable[[str, str, str], list[dict[str, Any] | None]],
|
||||
) -> None:
|
||||
_aoti_compile_backends[device_type] = AOTICompileBackend(
|
||||
compile_fn=compile_fn,
|
||||
load_fn=load_fn,
|
||||
)
|
||||
|
||||
|
||||
def aoti_eager_cache_dir(namespace: str, device: str) -> Path:
|
||||
return Path(cache_dir()) / "aoti_eager" / namespace / device
|
||||
|
||||
|
||||
def aoti_eager_op_conf_lock(op_func_name_with_overload: str) -> Any:
|
||||
# Avoid circular import
|
||||
from torch._inductor.codecache import get_lock_dir, LOCK_TIMEOUT
|
||||
from torch.utils._filelock import FileLock
|
||||
|
||||
op_conf_lock_file = f"{op_func_name_with_overload}.lock"
|
||||
lock_dir = get_lock_dir()
|
||||
return FileLock(os.path.join(lock_dir, op_conf_lock_file), timeout=LOCK_TIMEOUT)
|
||||
|
||||
|
||||
def load_aoti_eager_cache(
|
||||
ns: str, op_func_name_with_overload: str, device_type: str
|
||||
) -> list[dict[str, Any] | None]:
|
||||
backend = _aoti_compile_backends.get(device_type)
|
||||
if backend:
|
||||
return backend.load_fn(ns, op_func_name_with_overload, device_type)
|
||||
|
||||
device_kernel_cache = aoti_eager_cache_dir(ns, device_type)
|
||||
op_conf = device_kernel_cache / f"{op_func_name_with_overload}.json"
|
||||
if not op_conf.exists():
|
||||
return []
|
||||
|
||||
try:
|
||||
with aoti_eager_op_conf_lock(op_func_name_with_overload):
|
||||
with open(op_conf) as f:
|
||||
json_data = json.load(f)
|
||||
for item in json_data:
|
||||
# Get absolution path for kernel library
|
||||
kernel_lib_abs_path = device_kernel_cache / item["kernel_path"]
|
||||
item["kernel_path"] = kernel_lib_abs_path.as_posix()
|
||||
|
||||
# Check if the kernel library exists
|
||||
if not kernel_lib_abs_path.exists():
|
||||
return []
|
||||
|
||||
for metadata in item["meta_info"]:
|
||||
if metadata.get("is_dynamic"):
|
||||
raise NotImplementedError(
|
||||
"Only support static shape for now"
|
||||
)
|
||||
if (
|
||||
"device_type" in metadata
|
||||
and metadata["device_type"] == "cpu"
|
||||
):
|
||||
metadata["device_index"] = -1
|
||||
for dtype_key in ["dtype", "dtype_value"]:
|
||||
if dtype_key in metadata:
|
||||
metadata[dtype_key] = getattr(
|
||||
torch, metadata[dtype_key].split(".")[-1]
|
||||
)
|
||||
if "layout_value" in metadata:
|
||||
metadata["layout_value"] = getattr(
|
||||
torch, metadata["layout_value"].split(".")[-1]
|
||||
)
|
||||
if "memory_format_value" in metadata:
|
||||
metadata["memory_format_value"] = getattr(
|
||||
torch, metadata["memory_format_value"].split(".")[-1]
|
||||
)
|
||||
|
||||
return json_data
|
||||
except Exception as e:
|
||||
err_msg = f"Failed to load aoti eager cache: {e}"
|
||||
log.exception(err_msg)
|
||||
return []
|
||||
|
||||
|
||||
def supported_builtin_dtype_torch_dtype() -> dict[type, torch.dtype]:
|
||||
return {int: torch.int32, float: torch.float, bool: torch.bool}
|
||||
|
||||
|
||||
def supported_scalar_types() -> tuple[type, ...]:
|
||||
type_to_torch_dtype = supported_builtin_dtype_torch_dtype()
|
||||
return tuple(type_to_torch_dtype.keys())
|
||||
|
||||
|
||||
def extract_tensor_metadata(dynamic: bool, input: torch.Tensor) -> dict[str, Any]:
|
||||
metadata: dict[str, Any] = {}
|
||||
metadata["is_dynamic"] = dynamic
|
||||
|
||||
assert isinstance(input, torch.Tensor)
|
||||
metadata["device_type"] = f"{input.device.type}"
|
||||
if is_cpu_device([input]):
|
||||
metadata["device_index"] = -1
|
||||
else:
|
||||
metadata["device_index"] = input.device.index
|
||||
metadata["dtype"] = f"{input.dtype}"
|
||||
metadata["sizes"] = list(input.size())
|
||||
metadata["strides"] = list(input.stride())
|
||||
metadata["requires_grad"] = input.requires_grad
|
||||
metadata["dispatch_key_set"] = torch._C._dispatch_keys(input).raw_repr()
|
||||
return metadata
|
||||
|
||||
|
||||
def extract_tensor_list_metadata(
|
||||
dynamic: bool,
|
||||
input: list[torch.Tensor],
|
||||
) -> dict[str, Any]:
|
||||
metadata_list = []
|
||||
for item in input:
|
||||
assert isinstance(item, torch.Tensor)
|
||||
metadata_list.append(extract_tensor_metadata(dynamic, item))
|
||||
|
||||
metadata: dict[str, Any] = {}
|
||||
metadata["tensor_list"] = metadata_list
|
||||
return metadata
|
||||
|
||||
|
||||
def extract_scalar_metadata(device_type: str, input: Any) -> dict[str, Any]:
|
||||
assert isinstance(input, supported_scalar_types())
|
||||
metadata: dict[str, Any] = {}
|
||||
metadata["is_dynamic"] = False
|
||||
# Scalar tensor
|
||||
metadata["device_type"] = device_type
|
||||
metadata["device_index"] = -1 if device_type == "cpu" else 0
|
||||
type_to_torch_dtype = supported_builtin_dtype_torch_dtype()
|
||||
metadata["dtype"] = f"{type_to_torch_dtype[type(input)]}"
|
||||
metadata["scalar_value"] = input
|
||||
return metadata
|
||||
|
||||
|
||||
def extract_string_metadata(input: str) -> dict[str, Any]:
|
||||
assert isinstance(input, str)
|
||||
metadata: dict[str, Any] = {}
|
||||
metadata["string_value"] = input
|
||||
return metadata
|
||||
|
||||
|
||||
def extract_dtype_metadata(input: torch.dtype) -> dict[str, Any]:
|
||||
assert isinstance(input, torch.dtype)
|
||||
metadata: dict[str, Any] = {}
|
||||
metadata["dtype_value"] = f"{input}"
|
||||
return metadata
|
||||
|
||||
|
||||
def extract_device_metadata(input: torch.device) -> dict[str, Any]:
|
||||
assert isinstance(input, torch.device)
|
||||
metadata: dict[str, Any] = {}
|
||||
metadata["device_type_value"] = f"{input.type}"
|
||||
metadata["device_index_value"] = input.index
|
||||
return metadata
|
||||
|
||||
|
||||
def extract_layout_metadata(input: torch.layout) -> dict[str, Any]:
|
||||
assert isinstance(input, torch.layout)
|
||||
metadata: dict[str, Any] = {}
|
||||
metadata["layout_value"] = f"{input}"
|
||||
return metadata
|
||||
|
||||
|
||||
def aoti_compile_with_persistent_cache(
|
||||
ns: str,
|
||||
op_func_name_with_overload: str,
|
||||
device_type: str,
|
||||
dynamic: bool,
|
||||
f: Callable[..., Any],
|
||||
args: tuple[Any],
|
||||
kwargs: dict[str, Any],
|
||||
*,
|
||||
dynamic_shapes: dict[str, Any] | None = None,
|
||||
options: dict[str, Any] | None = None,
|
||||
remove_runtime_assertions: bool = False,
|
||||
disable_constraint_solver: bool = False,
|
||||
) -> str:
|
||||
"""
|
||||
Compile the given function with persistent cache for AOTI eager mode.
|
||||
"""
|
||||
backend = _aoti_compile_backends.get(device_type)
|
||||
if backend:
|
||||
return backend.compile_fn(
|
||||
ns,
|
||||
op_func_name_with_overload,
|
||||
device_type,
|
||||
dynamic,
|
||||
f,
|
||||
args,
|
||||
kwargs,
|
||||
dynamic_shapes=dynamic_shapes,
|
||||
options=options,
|
||||
remove_runtime_assertions=remove_runtime_assertions,
|
||||
disable_constraint_solver=disable_constraint_solver,
|
||||
)
|
||||
|
||||
assert not dynamic, "Only support static shape for now"
|
||||
flattened_inputs = list(args) + list(kwargs.values())
|
||||
if not all(
|
||||
isinstance(
|
||||
input,
|
||||
(
|
||||
supported_scalar_types(),
|
||||
torch.Tensor,
|
||||
list,
|
||||
str,
|
||||
torch.dtype,
|
||||
torch.device,
|
||||
torch.layout,
|
||||
),
|
||||
)
|
||||
for input in flattened_inputs
|
||||
):
|
||||
err_msg = f"Unsupported input types: {flattened_inputs}"
|
||||
log.exception(err_msg)
|
||||
raise NotImplementedError(err_msg)
|
||||
|
||||
for input in flattened_inputs:
|
||||
if isinstance(input, list) and not all(
|
||||
isinstance(item, torch.Tensor) for item in input
|
||||
):
|
||||
err_msg = f"_impl_with_aoti_compile encounters unsupported input types: {flattened_inputs}"
|
||||
log.exception(err_msg)
|
||||
raise NotImplementedError(err_msg)
|
||||
|
||||
persistent_cache = aoti_eager_cache_dir(ns, device_type)
|
||||
if not persistent_cache.exists():
|
||||
persistent_cache.mkdir(parents=True)
|
||||
|
||||
persistent_cache_lib = persistent_cache / "lib"
|
||||
if not persistent_cache_lib.exists():
|
||||
persistent_cache_lib.mkdir()
|
||||
|
||||
with mock.patch.dict(
|
||||
os.environ,
|
||||
{"TORCHINDUCTOR_CACHE_DIR": persistent_cache_lib.absolute().as_posix()},
|
||||
):
|
||||
try:
|
||||
kernel_lib_path = torch._export.aot_compile(
|
||||
f,
|
||||
args,
|
||||
kwargs,
|
||||
dynamic_shapes=dynamic_shapes,
|
||||
remove_runtime_assertions=remove_runtime_assertions,
|
||||
disable_constraint_solver=disable_constraint_solver,
|
||||
# Some operations may have non-Tensor parameters like int, float, bool. These
|
||||
# non-Tensor parameters will not be the input of the graph. Therefore, we do
|
||||
# need to keep the same signature.
|
||||
same_signature=False,
|
||||
)
|
||||
assert isinstance(kernel_lib_path, str)
|
||||
|
||||
kernel_metadata_items = []
|
||||
|
||||
for idx, input in enumerate(flattened_inputs):
|
||||
if isinstance(input, torch.Tensor):
|
||||
metadata = extract_tensor_metadata(dynamic, input)
|
||||
elif isinstance(input, list):
|
||||
assert all(isinstance(item, torch.Tensor) for item in input)
|
||||
metadata = extract_tensor_list_metadata(dynamic, input)
|
||||
elif isinstance(input, supported_scalar_types()):
|
||||
metadata = extract_scalar_metadata(device_type, input)
|
||||
elif isinstance(input, str):
|
||||
metadata = extract_string_metadata(input)
|
||||
elif isinstance(input, torch.dtype):
|
||||
metadata = extract_dtype_metadata(input)
|
||||
elif isinstance(input, torch.device):
|
||||
metadata = extract_device_metadata(input)
|
||||
elif isinstance(input, torch.layout):
|
||||
metadata = extract_layout_metadata(input)
|
||||
else:
|
||||
raise NotImplementedError(f"Unsupported input type: {type(input)}")
|
||||
|
||||
metadata["arg_order"] = idx
|
||||
kernel_metadata_items.append(metadata)
|
||||
|
||||
kernel_meta_info: dict[str, Any] = {}
|
||||
kernel_meta_info["meta_info"] = kernel_metadata_items
|
||||
kernel_meta_info["kernel_path"] = (
|
||||
Path(kernel_lib_path).relative_to(persistent_cache).as_posix()
|
||||
)
|
||||
|
||||
json_data = []
|
||||
update_json = True
|
||||
op_conf = persistent_cache / f"{op_func_name_with_overload}.json"
|
||||
mode = "r" if op_conf.exists() else "w"
|
||||
with aoti_eager_op_conf_lock(op_func_name_with_overload):
|
||||
with open(op_conf, mode) as op_conf_file:
|
||||
try:
|
||||
json_data = json.load(op_conf_file)
|
||||
except Exception:
|
||||
json_data = []
|
||||
|
||||
assert isinstance(json_data, list)
|
||||
for item in json_data:
|
||||
assert isinstance(item, dict)
|
||||
# Same kernel meta info already exists in the json file
|
||||
if item["meta_info"] == kernel_metadata_items:
|
||||
update_json = False
|
||||
break
|
||||
|
||||
if update_json:
|
||||
json_data.append(kernel_meta_info)
|
||||
with open(op_conf, "w") as op_conf_file:
|
||||
json.dump(json_data, op_conf_file, indent=4)
|
||||
|
||||
return kernel_lib_path
|
||||
except Exception as e:
|
||||
err_msg = f"Failed to compile {op_func_name_with_overload}: {e}"
|
||||
log.exception(err_msg)
|
||||
return ""
|
||||
@@ -0,0 +1,783 @@
|
||||
# mypy: allow-untyped-defs
|
||||
from __future__ import annotations
|
||||
|
||||
import atexit
|
||||
import functools
|
||||
import json
|
||||
import logging
|
||||
import multiprocessing
|
||||
import os
|
||||
import re
|
||||
import sys
|
||||
from concurrent.futures import Future, ThreadPoolExecutor
|
||||
from concurrent.futures.process import BrokenProcessPool
|
||||
from functools import partial
|
||||
from time import time, time_ns
|
||||
from typing import Any, TYPE_CHECKING
|
||||
|
||||
import torch
|
||||
from torch._dynamo.device_interface import get_registered_device_interfaces
|
||||
from torch._dynamo.utils import (
|
||||
counters,
|
||||
dynamo_timed,
|
||||
get_metrics_context,
|
||||
set_feature_use,
|
||||
)
|
||||
from torch._inductor import config
|
||||
from torch._inductor.codecache import (
|
||||
_load_triton_kernel_from_source,
|
||||
code_hash,
|
||||
CodeCacheFuture,
|
||||
CppCodeCache,
|
||||
CppPythonBindingsCodeCache,
|
||||
CUDACodeCache,
|
||||
HalideCodeCache,
|
||||
LambdaFuture,
|
||||
ROCmCodeCache,
|
||||
StaticAutotunerFuture,
|
||||
torch_key,
|
||||
XPUCodeCache,
|
||||
)
|
||||
from torch._inductor.compile_worker.subproc_pool import (
|
||||
AnyPool,
|
||||
SubprocException,
|
||||
SubprocPool,
|
||||
)
|
||||
from torch._inductor.compile_worker.tracked_process_pool import (
|
||||
TrackedProcessPoolExecutor,
|
||||
)
|
||||
from torch._inductor.compile_worker.utils import _async_compile_initializer
|
||||
from torch._inductor.runtime.compile_tasks import (
|
||||
_set_triton_libdevice_path,
|
||||
_set_triton_ptxas_path,
|
||||
_worker_compile_triton,
|
||||
)
|
||||
from torch._inductor.utils import clear_on_fresh_cache
|
||||
from torch._inductor.virtualized import V
|
||||
from torch._utils_internal import log_triton_builds
|
||||
from torch.hub import _Faketqdm, tqdm
|
||||
from torch.utils._ordered_set import OrderedSet
|
||||
from torch.utils._triton import has_triton_package
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Callable
|
||||
|
||||
from torch._inductor.runtime.hints import HalideMeta
|
||||
from torch._inductor.runtime.triton_heuristics import CachingAutotuner
|
||||
|
||||
# timing metrics for time spent in the compilation
|
||||
_cumulative_compile_time = 0.0
|
||||
_t0: float | None = None
|
||||
|
||||
kernel_code_log = torch._logging.getArtifactLogger(__name__, "kernel_code")
|
||||
|
||||
log = logging.getLogger(__name__)
|
||||
|
||||
_triton_kernel_metrics: dict[str, dict[str, Any]] | None = None
|
||||
|
||||
size_hints_regex = re.compile(
|
||||
r"size_hints=(\{.*?\})",
|
||||
)
|
||||
|
||||
|
||||
def pre_fork_setup():
|
||||
"""
|
||||
Setup that must be done prior to forking with a process pool.
|
||||
"""
|
||||
# ensure properties have been calculated before processes
|
||||
# are forked
|
||||
caching_device_properties()
|
||||
|
||||
# Computing the triton key can be slow. If we call it before fork,
|
||||
# it will be cached for the forked subprocesses.
|
||||
from torch._inductor.runtime.triton_compat import HAS_TRITON, triton_key
|
||||
|
||||
if HAS_TRITON:
|
||||
triton_key()
|
||||
|
||||
|
||||
def caching_device_properties():
|
||||
for _, device_interface in get_registered_device_interfaces():
|
||||
if device_interface.is_available():
|
||||
device_interface.Worker.get_device_properties()
|
||||
|
||||
|
||||
def _compile_start() -> None:
|
||||
global _t0, _triton_kernel_metrics
|
||||
if _t0 is None:
|
||||
_t0 = time()
|
||||
if _triton_kernel_metrics is None:
|
||||
_triton_kernel_metrics = {}
|
||||
|
||||
|
||||
def _compile_end() -> None:
|
||||
global _cumulative_compile_time, _t0, _triton_kernel_metrics
|
||||
if _t0 is not None:
|
||||
t1 = time()
|
||||
_cumulative_compile_time += t1 - _t0
|
||||
_t0 = None
|
||||
# print("CUMULATIVE COMPILE TIME", _cumulative_compile_time)
|
||||
if _triton_kernel_metrics:
|
||||
# Log triton kernel info
|
||||
sorted_info = dict(sorted(_triton_kernel_metrics.items()))
|
||||
torch._logging.trace_structured(
|
||||
"artifact",
|
||||
metadata_fn=lambda: {
|
||||
"name": "triton_kernel_info",
|
||||
"encoding": "json",
|
||||
},
|
||||
payload_fn=lambda: json.dumps(sorted_info),
|
||||
)
|
||||
_triton_kernel_metrics = None
|
||||
|
||||
|
||||
def _add_triton_kernel_info(kernel_name: str, info: dict[str, Any]):
|
||||
global _triton_kernel_metrics
|
||||
# Must be called between _compile_start and _compile_end
|
||||
if _triton_kernel_metrics is not None:
|
||||
_triton_kernel_metrics[kernel_name] = info
|
||||
|
||||
|
||||
_IS_WINDOWS = sys.platform == "win32"
|
||||
|
||||
log = logging.getLogger(__name__)
|
||||
|
||||
# Used to keep track of all process pools invoked so far.
|
||||
_pool_set = OrderedSet[AnyPool]()
|
||||
|
||||
|
||||
def shutdown_compile_workers() -> None:
|
||||
"""Shut down all outstanding compile-worker pools."""
|
||||
for pool in _pool_set:
|
||||
pool.shutdown()
|
||||
AsyncCompile._ready_future = None
|
||||
after_fork()
|
||||
|
||||
|
||||
def after_fork():
|
||||
"""Reset pools to initial state without shutting them down"""
|
||||
_pool_set.clear()
|
||||
AsyncCompile.process_pool.cache_clear()
|
||||
|
||||
|
||||
try:
|
||||
os.register_at_fork(after_in_child=after_fork)
|
||||
except AttributeError:
|
||||
pass # register_at_fork does not exists on windows
|
||||
|
||||
|
||||
def get_compile_threads() -> int:
|
||||
"""
|
||||
Temporary for internal rollout. Assign config.compile_threads lazily and return it.
|
||||
TODO: remove after rollout.
|
||||
"""
|
||||
if config.compile_threads is None:
|
||||
config.compile_threads = config.decide_compile_threads()
|
||||
return config.compile_threads
|
||||
|
||||
|
||||
@clear_on_fresh_cache
|
||||
class CompiledTritonKernels:
|
||||
"""
|
||||
In memory cache for storing compiled triton kernels.
|
||||
|
||||
Each triton kernel is keyed by the hash of its source code. Each value stored
|
||||
in the cache is a return value of AsyncCompile.triton().
|
||||
|
||||
Currently, the cache stores Future objects, but it should be generalizable for any kernels.
|
||||
"""
|
||||
|
||||
_cache: dict[str, CodeCacheFuture] = {}
|
||||
|
||||
@staticmethod
|
||||
def key(kernel_src: str):
|
||||
"""
|
||||
Generates a cache key given a triton kernel's full source code.
|
||||
This source includes the inductor meta, compilation metadata, the kernel itself, etc.
|
||||
`kernel_src` should be the exact string passed to async_compile.triton()'s first argument.
|
||||
"""
|
||||
# Hashes the kernel source with torch_key into a single hash key
|
||||
return code_hash(kernel_src, extra=torch_key())
|
||||
|
||||
@staticmethod
|
||||
def save(kernel_src: str, future: CodeCacheFuture):
|
||||
"""
|
||||
Saves a compiled triton kernel to the cache.
|
||||
TODO: We store a LambdaFuture as that's the callable returned by async_compile.triton,
|
||||
but the real type we want to return here is actually an abstract triton kernel.
|
||||
|
||||
TODO: Source code here is not just the kernel's source code, but also includes the inductor preamble, etc.
|
||||
so it could be less strict.
|
||||
"""
|
||||
key = CompiledTritonKernels.key(kernel_src)
|
||||
CompiledTritonKernels._cache[key] = future
|
||||
|
||||
@staticmethod
|
||||
def get(kernel_src: str) -> CodeCacheFuture | None:
|
||||
key = CompiledTritonKernels.key(kernel_src)
|
||||
return CompiledTritonKernels._cache.get(key, None)
|
||||
|
||||
@staticmethod
|
||||
def cache_clear():
|
||||
CompiledTritonKernels._cache = {}
|
||||
|
||||
@staticmethod
|
||||
def remove_future(kernel_src: str) -> None:
|
||||
key = CompiledTritonKernels.key(kernel_src)
|
||||
|
||||
# Delete the LambdaFuture if there is one
|
||||
if key in CompiledTritonKernels._cache:
|
||||
del CompiledTritonKernels._cache[key]
|
||||
|
||||
|
||||
class AsyncCompile:
|
||||
"""
|
||||
Utilities to compile in thread pools or subprocess pools (in the case of Triton).
|
||||
"""
|
||||
|
||||
_ready_future: Future[Any] | None = None
|
||||
_metal_sources: list[tuple[str, str, list[str]]] | None = None
|
||||
|
||||
def __init__(self) -> None:
|
||||
pass
|
||||
|
||||
@staticmethod
|
||||
@functools.lru_cache(1)
|
||||
def pool() -> ThreadPoolExecutor:
|
||||
assert get_compile_threads() > 1
|
||||
return ThreadPoolExecutor(get_compile_threads())
|
||||
|
||||
@staticmethod
|
||||
def _get_ready():
|
||||
"""No-op function to help mark when the subprocess pool is ready."""
|
||||
return "ready"
|
||||
|
||||
@staticmethod
|
||||
@functools.lru_cache(1)
|
||||
def process_pool() -> AnyPool:
|
||||
assert get_compile_threads() > 1
|
||||
AsyncCompile._ready_future = None
|
||||
log.info(
|
||||
"Creating '%s' pool with %d workers",
|
||||
config.worker_start_method,
|
||||
get_compile_threads(),
|
||||
)
|
||||
|
||||
pool: AnyPool
|
||||
if config.worker_start_method == "subprocess":
|
||||
# Wrapper around ProcessPoolExecutor forks in a new process we control
|
||||
pool = SubprocPool(
|
||||
get_compile_threads(), quiesce=config.quiesce_async_compile_pool
|
||||
)
|
||||
else:
|
||||
if config.worker_start_method == "spawn":
|
||||
# Avoid creating pools in the spawned subprocs themselves:
|
||||
os.environ["TORCH_WARM_POOL"] = "0"
|
||||
pre_fork_setup()
|
||||
ctx = multiprocessing.get_context(config.worker_start_method)
|
||||
pool = TrackedProcessPoolExecutor(
|
||||
get_compile_threads(),
|
||||
mp_context=ctx,
|
||||
initializer=partial(_async_compile_initializer, os.getpid()),
|
||||
)
|
||||
# when this pool is created in a subprocess object, the normal exit handler
|
||||
# doesn't run, and we need to register our own handler.
|
||||
# exitpriority has to be high, because another one of the finalizers will
|
||||
# kill the worker thread that sends the shutdown message to the workers...
|
||||
multiprocessing.util.Finalize(None, pool.shutdown, exitpriority=sys.maxsize)
|
||||
|
||||
_pool_set.add(pool)
|
||||
return pool
|
||||
|
||||
@classmethod
|
||||
def warm_pool(cls) -> None:
|
||||
if get_compile_threads() <= 1:
|
||||
return
|
||||
_compile_start()
|
||||
# Pool is created on first access. Note for a SubprocPool, the sidecar process starts,
|
||||
# but its ProcessPoolExecutor does not initialize until a wakeup() call or the first
|
||||
# job is submitted.
|
||||
cls.process_pool()
|
||||
_compile_end()
|
||||
|
||||
@classmethod
|
||||
def wait_pool_ready(cls, timeout=120) -> None:
|
||||
cls.use_process_pool()
|
||||
if cls._ready_future is not None:
|
||||
cls._ready_future.result(timeout=timeout)
|
||||
|
||||
@classmethod
|
||||
def submit(cls, task: Callable[..., Any]) -> Any:
|
||||
if get_compile_threads() <= 1:
|
||||
return task()
|
||||
return cls.pool().submit(task)
|
||||
|
||||
@classmethod
|
||||
def use_process_pool(cls):
|
||||
if get_compile_threads() <= 1:
|
||||
return False
|
||||
|
||||
# Proton instrumentation backend requires compilation to happen in the main
|
||||
# process so it can instrument the Triton IR during JIT compilation.
|
||||
# Force synchronous compilation when proton profiling is enabled.
|
||||
if config.triton.proton_profiling:
|
||||
return False
|
||||
|
||||
# Create a dummy job to check if the pool is ready. Submit it here instead of at
|
||||
# pool creation so we don't launch the full pool of worker subprocesses until
|
||||
# we're sure they're needed.
|
||||
if not cls._ready_future:
|
||||
cls._ready_future = cls.process_pool().submit(cls._get_ready)
|
||||
return cls._ready_future.done()
|
||||
|
||||
@classmethod
|
||||
def wakeup(cls) -> None:
|
||||
"""
|
||||
If using a SubprocPool, signal the sidecar process to start up its
|
||||
ProcessPoolExecutor.
|
||||
"""
|
||||
if not cls.use_process_pool():
|
||||
return
|
||||
pool = cls.process_pool()
|
||||
if isinstance(pool, SubprocPool):
|
||||
pool.wakeup()
|
||||
|
||||
def triton(self, kernel_name: str, source_code: str, device_str: str = "cuda"):
|
||||
"""
|
||||
Async_compile.triton is more complicated than the other backends because
|
||||
we're trying to optimize compile time as much as possible for this hot callsite.
|
||||
|
||||
First of all, the function is cached by CompiledTritonKernels; if there's a kernel
|
||||
already compiled, we grab it directly from the cache and return.
|
||||
|
||||
Otherwise, if we have multiple compile threads, we kick off triton compilations on each
|
||||
worker process by giving it a kernel and source code to compile. The worker initializes
|
||||
a CachingAutotuner, runs triton compilation, and pickles the kernel back to us.
|
||||
We use TritonCompileResult to represent the objects being pickled back to us by each
|
||||
worker.
|
||||
|
||||
Some maybe not obvious things that are pickled back to us:
|
||||
- Most of the time, we can avoid sending back CachingAutotuner.fn and other metadata
|
||||
and do not have to pay the cost of loading the triton kernel on the parent. But certain
|
||||
cases, like coordesc tuning and dynamic_scale_rblock, require us to reload the function
|
||||
in the parent lazily when we require it.
|
||||
- The AutotuneCache, if enabled, is constructed on each worker per triton config
|
||||
and pickled by to us via `CachingAutotuner.save_cache_hook`.
|
||||
"""
|
||||
load_kernel = functools.partial(
|
||||
_load_triton_kernel_from_source, kernel_name, source_code
|
||||
)
|
||||
|
||||
def reload_kernel_in_parent():
|
||||
# Benchmark how often this happens
|
||||
with dynamo_timed("reload_kernel_in_parent"):
|
||||
return load_kernel()
|
||||
|
||||
counters["inductor"]["async_compile_cache_miss"] += 1
|
||||
|
||||
kernel_code_log.info("Triton Kernel:\n%s", source_code)
|
||||
_compile_start()
|
||||
|
||||
if os.environ.get("TRITON_INTERPRET", "0") == "1":
|
||||
return getattr(
|
||||
torch._inductor.codecache.PyCodeCache.load(source_code), kernel_name
|
||||
)
|
||||
|
||||
is_parallel = self.use_process_pool()
|
||||
set_feature_use("parallel_compile_post_warmup", is_parallel)
|
||||
|
||||
compile_id = torch._guards.CompileContext.current_compile_id()
|
||||
is_backward = getattr(V.graph, "is_backward", False)
|
||||
|
||||
if (future := CompiledTritonKernels.get(source_code)) is not None:
|
||||
counters["inductor"]["async_compile_cache_hit"] += 1
|
||||
# Set reload_kernel_from_src properly based on source_code
|
||||
if isinstance(future, StaticAutotunerFuture):
|
||||
# Remove the future now that we've cache hit
|
||||
CompiledTritonKernels.remove_future(source_code)
|
||||
future.reload_kernel_from_src = reload_kernel_in_parent
|
||||
if is_parallel:
|
||||
return future
|
||||
else:
|
||||
return future.result()
|
||||
|
||||
# Cache miss
|
||||
if is_parallel:
|
||||
# Ensure libdevice path is set in os.environ before passing to workers
|
||||
_set_triton_libdevice_path()
|
||||
# We want to support changing these env vars after (and while) the
|
||||
# process pool is running, so pass them to the subprocess to reset.
|
||||
env_vars = [
|
||||
"TORCHINDUCTOR_CACHE_DIR",
|
||||
"TRITON_CACHE_DIR",
|
||||
"TRITON_LIBDEVICE_PATH",
|
||||
]
|
||||
extra_env = {v: os.environ[v] for v in env_vars if v in os.environ}
|
||||
extra_config = {
|
||||
"use_static_triton_launcher": torch._inductor.config.use_static_triton_launcher
|
||||
}
|
||||
|
||||
if len(torch._inductor.config.autotune_lookup_table) > 0:
|
||||
m = size_hints_regex.search(source_code)
|
||||
if m:
|
||||
size_hints_str = m.group(1)
|
||||
else:
|
||||
size_hints_str = str(None)
|
||||
|
||||
triton_src = source_code.split("@triton.jit\n")[1]
|
||||
from torch._inductor.runtime.triton_heuristics import (
|
||||
generate_lookup_hash_from_source_code,
|
||||
)
|
||||
|
||||
fn_hash = generate_lookup_hash_from_source_code(
|
||||
size_hints_str, triton_src
|
||||
)
|
||||
|
||||
if fn_hash in torch._inductor.config.autotune_lookup_table:
|
||||
extra_config["autotune_lookup_table"] = { # type: ignore[assignment]
|
||||
fn_hash: torch._inductor.config.autotune_lookup_table[fn_hash]
|
||||
}
|
||||
|
||||
task = self.process_pool().submit(
|
||||
_worker_compile_triton,
|
||||
load_kernel,
|
||||
extra_env,
|
||||
extra_config,
|
||||
)
|
||||
|
||||
def get_result() -> CachingAutotuner:
|
||||
try:
|
||||
kernel, elapsed_us = task.result()
|
||||
except SubprocException as e:
|
||||
raise e.with_name(kernel_name) from e
|
||||
|
||||
# Now that we've compiled, we should clear the future
|
||||
# so it can't be used again
|
||||
kernel.set_compile_info(compile_id, is_backward)
|
||||
CompiledTritonKernels.remove_future(source_code)
|
||||
|
||||
kernel.restore_after_unpickle(old_values=None)
|
||||
|
||||
kernel.precompile(
|
||||
warm_cache_only=False,
|
||||
reload_kernel=reload_kernel_in_parent,
|
||||
static_triton_bundle_key=CompiledTritonKernels.key(source_code),
|
||||
)
|
||||
info = kernel.autotune_cache_info or {}
|
||||
info["compile_time_us"] = elapsed_us
|
||||
_add_triton_kernel_info(kernel_name, info)
|
||||
get_metrics_context().add_top_n(
|
||||
"triton_kernel_compile_times_us", kernel_name, elapsed_us
|
||||
)
|
||||
return kernel
|
||||
|
||||
future = LambdaFuture(get_result, future=task)
|
||||
CompiledTritonKernels.save(source_code, future)
|
||||
return future
|
||||
else:
|
||||
with dynamo_timed(
|
||||
"async_compile.precompile",
|
||||
log_pt2_compile_event=True,
|
||||
dynamo_compile_column_us="triton_compile_time_us",
|
||||
log_waitcounter=True,
|
||||
waitcounter_name_override="compile_triton",
|
||||
):
|
||||
fail = None
|
||||
try:
|
||||
start_ns = time_ns()
|
||||
_set_triton_ptxas_path()
|
||||
_set_triton_libdevice_path()
|
||||
kernel = load_kernel()
|
||||
kernel.set_compile_info(compile_id, is_backward)
|
||||
kernel.precompile(
|
||||
warm_cache_only=False,
|
||||
static_triton_bundle_key=CompiledTritonKernels.key(source_code),
|
||||
)
|
||||
elapsed_us = (time_ns() - start_ns) // 1000
|
||||
get_metrics_context().add_top_n(
|
||||
"triton_kernel_compile_times_us", kernel_name, elapsed_us
|
||||
)
|
||||
info = kernel.autotune_cache_info or {}
|
||||
info["compile_time_us"] = elapsed_us
|
||||
_add_triton_kernel_info(kernel_name, info)
|
||||
return kernel
|
||||
except Exception as e:
|
||||
fail = str(e)
|
||||
raise
|
||||
finally:
|
||||
log_triton_builds(fail=fail)
|
||||
|
||||
def multi_kernel(self, *args, **kwargs) -> Any:
|
||||
from torch._inductor.codegen.multi_kernel import MultiKernelCall
|
||||
|
||||
# no need to call this in parallel since the sub-kernels are already parallel tasks
|
||||
return MultiKernelCall(*args, **kwargs)
|
||||
|
||||
def size_hint_multi_kernel(self, *args, **kwargs) -> Any:
|
||||
from torch._inductor.codegen.multi_kernel import SizeHintMultiKernelCall
|
||||
|
||||
return SizeHintMultiKernelCall(*args, **kwargs)
|
||||
|
||||
def cpp(self, source_code: str):
|
||||
kernel_code_log.info("CPP Kernel:\n%s", source_code)
|
||||
if get_compile_threads() <= 1:
|
||||
return CppCodeCache.load(source_code).kernel
|
||||
else:
|
||||
get_result = CppCodeCache.load_async(source_code, submit_fn=self.submit)
|
||||
return LambdaFuture(lambda: get_result().kernel)
|
||||
|
||||
def cpp_pybinding(self, argtypes: list[str], source_code: str):
|
||||
kernel_code_log.info("CPP+Bindings Kernel:\n%s", source_code)
|
||||
if get_compile_threads() <= 1:
|
||||
return CppPythonBindingsCodeCache.load_pybinding(argtypes, source_code)
|
||||
else:
|
||||
get_result = CppPythonBindingsCodeCache.load_pybinding_async(
|
||||
argtypes, source_code, submit_fn=self.submit
|
||||
)
|
||||
return LambdaFuture(get_result)
|
||||
|
||||
def cutlass(self, cache_cls, source_code, dst_file_ext, aot_compile=False):
|
||||
def task():
|
||||
if aot_compile:
|
||||
# We rely on JITInductor to compile the CUDA code,
|
||||
# so that we can load it into AOTInductor.
|
||||
output_path, *_ = cache_cls.compile(source_code, "o")
|
||||
cache_cls.aot_kernels_o.append(output_path)
|
||||
return cache_cls.load(source_code, dst_file_ext)[0]
|
||||
|
||||
return self.submit(task)
|
||||
|
||||
def cuda(self, source_code, dst_file_ext, aot_compile=False):
|
||||
kernel_code_log.info("CUDA Kernel:\n%s", source_code)
|
||||
return self.cutlass(CUDACodeCache, source_code, dst_file_ext, aot_compile)
|
||||
|
||||
def xpu(self, source_code, dst_file_ext, aot_compile=False):
|
||||
kernel_code_log.info("XPU Kernel:\n%s", source_code)
|
||||
return self.cutlass(XPUCodeCache, source_code, dst_file_ext, aot_compile)
|
||||
|
||||
def rocm(
|
||||
self,
|
||||
source_code,
|
||||
dst_file_ext,
|
||||
aot_compile=False,
|
||||
):
|
||||
kernel_code_log.info("ROCm Kernel:\n%s", source_code)
|
||||
|
||||
def task():
|
||||
if aot_compile:
|
||||
output_path, *_ = ROCmCodeCache.compile(source_code, dst_file_ext="o")
|
||||
ROCmCodeCache.aot_kernels_o.append(output_path)
|
||||
if config.rocm.generate_test_runner:
|
||||
_ = ROCmCodeCache.compile(source_code, dst_file_ext="exe")
|
||||
return ROCmCodeCache.load(source_code, dst_file_ext)[0]
|
||||
|
||||
return self.submit(task)
|
||||
|
||||
def halide(self, meta: HalideMeta, source_code: str):
|
||||
kernel_code_log.info("Halide Kernel:\n%r\n%s", meta, source_code)
|
||||
if get_compile_threads() <= 1:
|
||||
return HalideCodeCache.generate_halide(meta, source_code)
|
||||
else:
|
||||
get_result = HalideCodeCache.generate_halide_async(
|
||||
meta, source_code, submit_fn=self.submit
|
||||
)
|
||||
return LambdaFuture(get_result)
|
||||
|
||||
def cutedsl(self, kernel_name: str, source_code: str):
|
||||
"""
|
||||
Compile CuteDSL (CUTLASS Python DSL) kernels.
|
||||
|
||||
Args:
|
||||
kernel_name: Name of the kernel to be defined
|
||||
source_code: Source code of the CuteDSL kernel, as a string
|
||||
|
||||
Note:
|
||||
CuteDSL currently requires source files to do its compilation, there we
|
||||
use the PyCodeCache to write the source code to a file and load it.
|
||||
"""
|
||||
from torch._inductor.codegen.cutedsl.cutedsl_kernel import (
|
||||
CuteDSLKernelWrapper,
|
||||
MAIN_SUFFIX,
|
||||
)
|
||||
|
||||
kernel_code_log.info("CuteDSL Kernel:\n%s", source_code)
|
||||
|
||||
def task():
|
||||
key, path = torch._inductor.codecache.PyCodeCache.write(source_code)
|
||||
mod = torch._inductor.codecache.PyCodeCache.load_by_key_path(key, path)
|
||||
|
||||
# Find our special entry point named function
|
||||
main_func_name = f"{kernel_name}_{MAIN_SUFFIX}"
|
||||
if not hasattr(mod, main_func_name):
|
||||
available = [name for name in dir(mod) if callable(getattr(mod, name))]
|
||||
raise RuntimeError(
|
||||
f"Could not find CuteDSL main kernel function '{main_func_name}'. Available callables: {available}"
|
||||
)
|
||||
|
||||
return CuteDSLKernelWrapper(getattr(mod, main_func_name), kernel_path=path)
|
||||
|
||||
if get_compile_threads() <= 1:
|
||||
return task()
|
||||
else:
|
||||
future = self.submit(task)
|
||||
return LambdaFuture(lambda: future.result())
|
||||
|
||||
def pallas(self, kernel_name: str, source_code: str):
|
||||
"""
|
||||
Compile Pallas (JAX experimental) kernels.
|
||||
|
||||
Args:
|
||||
kernel_name: Name of the kernel to be defined
|
||||
source_code: Source code of the Pallas kernel, as a string
|
||||
|
||||
Note:
|
||||
Pallas kernels are Python code that uses JAX and Pallas APIs.
|
||||
We use the PyCodeCache to write the source code to a file and load it.
|
||||
"""
|
||||
from torch._inductor.codegen.pallas import MAIN_SUFFIX, PallasKernelWrapper
|
||||
|
||||
kernel_code_log.info("Pallas Kernel:\n%s", source_code)
|
||||
|
||||
def task():
|
||||
key, path = torch._inductor.codecache.PyCodeCache.write(source_code)
|
||||
mod = torch._inductor.codecache.PyCodeCache.load_by_key_path(key, path)
|
||||
|
||||
# Find our special entry point named function
|
||||
main_func_name = f"{kernel_name}_{MAIN_SUFFIX}"
|
||||
if not hasattr(mod, main_func_name):
|
||||
available = [name for name in dir(mod) if callable(getattr(mod, name))]
|
||||
raise RuntimeError(
|
||||
f"Could not find Pallas main kernel function '{main_func_name}'. Available callables: {available}"
|
||||
)
|
||||
|
||||
return PallasKernelWrapper(getattr(mod, main_func_name), kernel_path=path)
|
||||
|
||||
if get_compile_threads() <= 1:
|
||||
return task()
|
||||
else:
|
||||
future = self.submit(task)
|
||||
return LambdaFuture(lambda: future.result())
|
||||
|
||||
def nv_universal_gemm(self, kernel_name: str, source_code: str):
|
||||
"""
|
||||
Compile NVIDIA Universal GEMM kernels.
|
||||
|
||||
Args:
|
||||
kernel_name: Name of the kernel to be defined
|
||||
source_code: Source code of the kernel, as a string
|
||||
|
||||
Note:
|
||||
NVIDIA Universal GEMM kernels are Python code that calls the cutlass_api library.
|
||||
We use the PyCodeCache to write the source code to a file and load it.
|
||||
"""
|
||||
from torch._inductor.codegen.nv_universal_gemm.nv_universal_gemm_kernel import (
|
||||
NVUniversalGemmKernelWrapper,
|
||||
)
|
||||
from torch._inductor.codegen.nv_universal_gemm.nv_universal_gemm_scheduling import (
|
||||
MAIN_SUFFIX,
|
||||
)
|
||||
|
||||
kernel_code_log.info("NVIDIA Universal GEMM Kernel:\n%s", source_code)
|
||||
|
||||
def task():
|
||||
key, path = torch._inductor.codecache.PyCodeCache.write(source_code)
|
||||
mod = torch._inductor.codecache.PyCodeCache.load_by_key_path(key, path)
|
||||
|
||||
# Find our special entry point named function
|
||||
main_func_name = f"{kernel_name}_{MAIN_SUFFIX}"
|
||||
if not hasattr(mod, main_func_name):
|
||||
available = [name for name in dir(mod) if callable(getattr(mod, name))]
|
||||
raise RuntimeError(
|
||||
f"Could not find NVIDIA Universal GEMM main kernel function "
|
||||
f"'{main_func_name}'. Available callables: {available}"
|
||||
)
|
||||
|
||||
return NVUniversalGemmKernelWrapper(
|
||||
getattr(mod, main_func_name), kernel_path=path
|
||||
)
|
||||
|
||||
if get_compile_threads() <= 1:
|
||||
return task()
|
||||
else:
|
||||
future = self.submit(task)
|
||||
return LambdaFuture(lambda: future.result())
|
||||
|
||||
def metal(self, kernel_name: str, source: str, headers: list[str]) -> None:
|
||||
"""Register a Metal kernel body; wait() compiles all registered kernels into one library."""
|
||||
if self._metal_sources is None:
|
||||
self._metal_sources = []
|
||||
self._metal_sources.append((kernel_name, source, headers))
|
||||
|
||||
def wait(self, scope: dict[str, Any]) -> None:
|
||||
if get_compile_threads() > 1:
|
||||
with dynamo_timed(
|
||||
"async_compile.wait",
|
||||
log_pt2_compile_event=True,
|
||||
dynamo_compile_column_us="triton_compile_time_us",
|
||||
log_waitcounter=True,
|
||||
waitcounter_name_override="compile_triton",
|
||||
):
|
||||
self._wait_futures(scope)
|
||||
|
||||
if self._metal_sources:
|
||||
from torch._inductor.runtime.runtime_utils import compile_mps_shaders
|
||||
|
||||
scope.update(compile_mps_shaders(self._metal_sources))
|
||||
self._metal_sources.clear()
|
||||
|
||||
_compile_end()
|
||||
|
||||
def _wait_futures(self, scope: dict[str, Any]) -> None:
|
||||
kernels = {
|
||||
key: value
|
||||
for key, value in scope.items()
|
||||
if isinstance(value, (Future, CodeCacheFuture))
|
||||
}
|
||||
pbar = tqdm(
|
||||
total=len(kernels),
|
||||
desc="Inductor Compilation",
|
||||
disable=config.disable_progress,
|
||||
delay=0,
|
||||
)
|
||||
for key, result in kernels.items():
|
||||
if config.verbose_progress and not isinstance(pbar, _Faketqdm):
|
||||
pbar.set_postfix_str(key)
|
||||
try:
|
||||
kernel = result.result()
|
||||
scope[key] = kernel
|
||||
except BrokenProcessPool as e:
|
||||
raise RuntimeError(
|
||||
"A compilation subprocess exited unexpectedly. This "
|
||||
"is likely due to a crash. To facilitate debugging, "
|
||||
"you can re-run with TORCHINDUCTOR_COMPILE_THREADS=1 "
|
||||
"to cause compilation to occur in the main process."
|
||||
) from e
|
||||
pbar.update(1)
|
||||
|
||||
|
||||
def maybe_warm_pool() -> None:
|
||||
if (
|
||||
os.environ.get("TORCH_TNT_IN_USE", "0") == "1"
|
||||
or os.environ.get("TORCH_WARM_POOL", "1") != "1"
|
||||
# The subprocess pool is only used for the Triton backend
|
||||
or not has_triton_package()
|
||||
# Skip for fbcode. We have internal reports of usages inside multiprocessing
|
||||
# pools that lead a multiplicative number of compile subprocesses.
|
||||
or config.is_fbcode()
|
||||
):
|
||||
return
|
||||
|
||||
AsyncCompile.warm_pool()
|
||||
# TODO: This starts the SubprocPool's internal process pool as early as possible at
|
||||
# the expense of creating a bunch of worker processes that might not be needed. We
|
||||
# could start them lazily if we're willing to lose a small amount of compile time.
|
||||
AsyncCompile.wakeup()
|
||||
|
||||
|
||||
# On exit give the workers a chance to clean themselves up. Without this the
|
||||
# resource_tracker can complain about leaked semaphores coming from the
|
||||
# ProcessPoolExecutor:
|
||||
# UserWarning: resource_tracker: There appear to be 5 leaked semaphore objects
|
||||
# to clean up at shutdown
|
||||
atexit.register(shutdown_compile_workers)
|
||||
@@ -0,0 +1,276 @@
|
||||
import logging
|
||||
from collections import defaultdict
|
||||
|
||||
import torch
|
||||
import torch.fx as fx
|
||||
from torch._logging import trace_structured
|
||||
from torch.utils._ordered_set import OrderedSet
|
||||
|
||||
|
||||
log = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class AugmentedGraphHelper:
|
||||
"""
|
||||
Graph helper that augments the original graph with additional
|
||||
dependencies and uses, plus tracks node equivalences for coalescing.
|
||||
|
||||
TODO: if this becomes too large of compile time, consider binding
|
||||
graphcycles.cc
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
graph: fx.Graph,
|
||||
node_ancestors: dict[fx.Node, OrderedSet[fx.Node]] | None = None,
|
||||
):
|
||||
# Each node starts in its own singleton set
|
||||
self.graph = graph
|
||||
self.merge_sets = {node: OrderedSet([node]) for node in graph.nodes}
|
||||
|
||||
# Extra dependencies: node depends on dep (dep must come before node)
|
||||
self.extra_deps: dict[fx.Node, OrderedSet[fx.Node]] = defaultdict(OrderedSet)
|
||||
# Extra uses: reverse of extra_deps (node is used by user)
|
||||
self.extra_uses: dict[fx.Node, OrderedSet[fx.Node]] = defaultdict(OrderedSet)
|
||||
# Note: only reflect original ancestors, not maintained through additional deps
|
||||
# or merge sets
|
||||
self.node_ancestors = node_ancestors
|
||||
|
||||
def add_extra_dep(self, *, n: fx.Node, dep: fx.Node) -> None:
|
||||
"""Add extra dependency: node depends on dep."""
|
||||
self.extra_deps[n].add(dep)
|
||||
self.extra_uses[dep].add(n)
|
||||
|
||||
def remove_extra_dep(self, *, n: fx.Node, dep: fx.Node) -> None:
|
||||
if dep in self.extra_deps[n]:
|
||||
self.extra_deps[n].discard(dep)
|
||||
self.extra_uses[dep].discard(n)
|
||||
|
||||
def merge_to_set(self, existing_node: fx.Node, new_node: fx.Node) -> None:
|
||||
"""
|
||||
Merge new_node into existing_node's set. The new node must be a singleton set.
|
||||
"""
|
||||
existing_set = self.merge_sets[existing_node]
|
||||
new_set = self.merge_sets[new_node]
|
||||
assert len(new_set) == 1
|
||||
|
||||
# Add all nodes from new_set to existing_set
|
||||
existing_set.update(new_set)
|
||||
|
||||
# Update all nodes from new_set to point to existing_set
|
||||
for node in new_set:
|
||||
self.merge_sets[node] = existing_set
|
||||
|
||||
def unmerge_node(self, node: fx.Node) -> None:
|
||||
"""Remove a node from its merge set, making it singleton."""
|
||||
old_set = self.merge_sets[node]
|
||||
|
||||
# If already singleton, nothing to do
|
||||
if len(old_set) == 1:
|
||||
return
|
||||
|
||||
# Remove from old set
|
||||
old_set.remove(node)
|
||||
|
||||
# Make node singleton
|
||||
self.merge_sets[node] = OrderedSet([node])
|
||||
|
||||
def get_merged_deps(self, node: fx.Node) -> OrderedSet[fx.Node]:
|
||||
"""
|
||||
Get all dependencies of a node considering merges and extra deps.
|
||||
Combines:
|
||||
1. Direct deps (all_input_nodes) of node and its merge equivalents
|
||||
2. Extra deps of node and its merge equivalents
|
||||
"""
|
||||
deps: OrderedSet[fx.Node] = OrderedSet()
|
||||
|
||||
# For each node in the merge set
|
||||
for merged_node in self.merge_sets[node]:
|
||||
# Add direct dependencies from all_input_nodes
|
||||
deps.update(merged_node.all_input_nodes)
|
||||
# Add extra dependencies
|
||||
deps.update(self.extra_deps[merged_node])
|
||||
|
||||
return deps
|
||||
|
||||
def has_cycle(self) -> bool:
|
||||
return torch._dynamo.graph_deduplication._has_cycle(
|
||||
self.graph, self.get_all_extra_deps()
|
||||
)
|
||||
|
||||
def _get_all_ancestors(self, node: fx.Node) -> OrderedSet[fx.Node]:
|
||||
"""Transitive ancestors through both data deps and extra deps."""
|
||||
ancestors: OrderedSet[fx.Node] = OrderedSet()
|
||||
stack: list[fx.Node] = list(node.all_input_nodes)
|
||||
stack.extend(self.extra_deps.get(node, ()))
|
||||
while stack:
|
||||
n = stack.pop()
|
||||
if n not in ancestors:
|
||||
ancestors.add(n)
|
||||
stack.extend(n.all_input_nodes)
|
||||
stack.extend(self.extra_deps.get(n, ()))
|
||||
return ancestors
|
||||
|
||||
def has_path(self, source: fx.Node, target: fx.Node) -> bool:
|
||||
"""Check if there's a path from source to target."""
|
||||
# we should not be checking path from node to itself
|
||||
assert self.merge_sets[source] is not self.merge_sets[target]
|
||||
|
||||
# search backwards from target to source
|
||||
visited: OrderedSet[fx.Node] = OrderedSet()
|
||||
queue = [target]
|
||||
visited.add(target)
|
||||
|
||||
while queue:
|
||||
current = queue.pop()
|
||||
|
||||
for dep in self.get_merged_deps(current):
|
||||
# Check if we reached source or its equivalent
|
||||
if dep in self.merge_sets[source]:
|
||||
return True
|
||||
|
||||
if dep in visited:
|
||||
continue
|
||||
|
||||
# We are searching from target, so this node is necessarily an ancestor
|
||||
# of target.
|
||||
# If dep is an ancestor of source, any path through dep to source would imply a cycle
|
||||
if self.node_ancestors:
|
||||
source_set = self.merge_sets[source]
|
||||
is_ancestor_of_source = any(
|
||||
dep in self.node_ancestors[s] for s in source_set
|
||||
)
|
||||
# Add to visited to avoid recomputing this check if we see dep again
|
||||
if is_ancestor_of_source:
|
||||
visited.add(dep)
|
||||
continue
|
||||
|
||||
visited.add(dep)
|
||||
queue.append(dep)
|
||||
|
||||
return False
|
||||
|
||||
def transfer_erased_node_deps(
|
||||
self, erased_to_new: dict[fx.Node, fx.Node | None]
|
||||
) -> None:
|
||||
"""
|
||||
Transfer all extra dependencies from erased nodes to their replacements, handling
|
||||
cross-dependencies between erased nodes correctly.
|
||||
|
||||
Skips deps where both endpoints resolve to replacement nodes from the
|
||||
same erasure batch — these are intra-bucket deps that would create
|
||||
cycles (e.g. new_start <-> new_wait within the same bucket).
|
||||
"""
|
||||
erased_merge_sets: dict[fx.Node, fx.Node | None] = {}
|
||||
|
||||
for replaced, new in erased_to_new.items():
|
||||
for equiv in self.merge_sets[replaced]:
|
||||
erased_merge_sets[equiv] = new
|
||||
|
||||
# Transfer dependencies
|
||||
for old_node, new_node in erased_merge_sets.items():
|
||||
if new_node is None:
|
||||
# Clean up references to removed node
|
||||
for extra_use in list(self.extra_uses[old_node]):
|
||||
updated_use = erased_merge_sets.get(extra_use, extra_use)
|
||||
if updated_use is not None:
|
||||
self.extra_deps[updated_use].discard(old_node)
|
||||
for extra_dep in list(self.extra_deps[old_node]):
|
||||
updated_dep = erased_merge_sets.get(extra_dep, extra_dep)
|
||||
if updated_dep is not None:
|
||||
self.extra_uses[updated_dep].discard(old_node)
|
||||
else:
|
||||
# Transfer dependencies FROM old_node (what old_node depended on)
|
||||
for extra_dep in self.extra_deps[old_node]:
|
||||
updated_dep = erased_merge_sets.get(extra_dep, extra_dep)
|
||||
if updated_dep is not None and updated_dep != new_node:
|
||||
# Skip if reverse dep already exists (extra or data)
|
||||
if new_node in self.extra_deps.get(
|
||||
updated_dep, ()
|
||||
) or new_node in OrderedSet(updated_dep.all_input_nodes):
|
||||
continue
|
||||
self.extra_deps[new_node].add(updated_dep)
|
||||
self.extra_uses[updated_dep].discard(old_node)
|
||||
self.extra_uses[updated_dep].add(new_node)
|
||||
|
||||
# Transfer dependencies TO old_node (what depended on old_node)
|
||||
for extra_use in self.extra_uses[old_node]:
|
||||
updated_use = erased_merge_sets.get(extra_use, extra_use)
|
||||
if updated_use is not None and updated_use != new_node:
|
||||
# Skip if reverse dep already exists (extra or data)
|
||||
if updated_use in self.extra_deps.get(
|
||||
new_node, ()
|
||||
) or updated_use in OrderedSet(new_node.all_input_nodes):
|
||||
continue
|
||||
self.extra_deps[updated_use].discard(old_node)
|
||||
self.extra_deps[updated_use].add(new_node)
|
||||
self.extra_uses[new_node].add(updated_use)
|
||||
|
||||
# Clean up erased nodes
|
||||
for old_node in erased_merge_sets:
|
||||
self.extra_deps[old_node].clear()
|
||||
self.extra_uses[old_node].clear()
|
||||
del self.merge_sets[old_node]
|
||||
|
||||
def remove_erased_extra_deps(self) -> None:
|
||||
"""Remove extra deps referencing erased nodes."""
|
||||
for node in list(self.extra_deps):
|
||||
if node._erased:
|
||||
for dep in list(self.extra_deps[node]):
|
||||
self.remove_extra_dep(n=node, dep=dep)
|
||||
continue
|
||||
for dep in list(self.extra_deps[node]):
|
||||
if dep._erased:
|
||||
self.remove_extra_dep(n=node, dep=dep)
|
||||
|
||||
def check_and_maybe_autofix_cyclic_extra_deps(
|
||||
self, *, autofix: bool = False
|
||||
) -> None:
|
||||
"""Check for and optionally remove extra deps that create cycles.
|
||||
|
||||
Args:
|
||||
autofix: If True, silently remove cyclic deps. If False (default),
|
||||
raise an error so the root cause gets investigated.
|
||||
"""
|
||||
if not self.has_cycle():
|
||||
return
|
||||
removed = []
|
||||
for node in list(self.extra_deps):
|
||||
for dep in list(self.extra_deps[node]):
|
||||
ancestors = self._get_all_ancestors(dep)
|
||||
if node in ancestors:
|
||||
removed.append((node.name, dep.name))
|
||||
self.remove_extra_dep(n=node, dep=dep)
|
||||
if not removed:
|
||||
return
|
||||
msg = (
|
||||
f"Overlap scheduling: detected {len(removed)} cyclic extra "
|
||||
f"dep(s): {removed}. Please report this to the overlap "
|
||||
f"scheduling developers."
|
||||
)
|
||||
log.warning(msg)
|
||||
trace_structured(
|
||||
"artifact",
|
||||
metadata_fn=lambda: {
|
||||
"name": "inductor_overlap_cyclic_extra_deps",
|
||||
"encoding": "string",
|
||||
},
|
||||
payload_fn=lambda: msg,
|
||||
)
|
||||
if not autofix:
|
||||
raise RuntimeError(
|
||||
f"{msg}\nTo unblock, set "
|
||||
f"torch._inductor.config.aten_distributed_optimizations"
|
||||
f".overlap_scheduling_autofix_cycles = True"
|
||||
)
|
||||
|
||||
def get_all_extra_deps(self) -> dict[fx.Node, OrderedSet[fx.Node]]:
|
||||
"""
|
||||
Get all extra dependencies in a format suitable for topological sort.
|
||||
Returns a copy to avoid external modifications.
|
||||
"""
|
||||
return {
|
||||
node: OrderedSet(deps)
|
||||
for node, deps in self.extra_deps.items()
|
||||
if deps # Only include nodes with non-empty deps
|
||||
}
|
||||
+296
@@ -0,0 +1,296 @@
|
||||
# flake8: noqa: B950
|
||||
# fmt: off
|
||||
# This file was generated by AutoHeuristic. Do not modify it manually!
|
||||
# To regenerate this file, take a look at the steps in the README.md file inside torchgen/_autoheuristic/mm/
|
||||
from typing import List, Optional, Tuple
|
||||
|
||||
from torch._inductor.autoheuristic.autoheuristic_utils import (
|
||||
AHContext,
|
||||
AHMetadata,
|
||||
Choice,
|
||||
)
|
||||
from torch._inductor.autoheuristic.learnedheuristic_interface import (
|
||||
LearnedHeuristicDecision,
|
||||
)
|
||||
|
||||
|
||||
class MMRankingA100(LearnedHeuristicDecision):
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.choices: list[Choice] = []
|
||||
self.fill_choices()
|
||||
|
||||
def check_precondition(self, metadata: AHMetadata, context: AHContext,) -> bool:
|
||||
return (
|
||||
metadata.name == self.get_name()
|
||||
and metadata.shared_memory == 166912
|
||||
and str(metadata.device_capa) == "(8, 0)"
|
||||
)
|
||||
|
||||
def get_confidence_threshold(self) -> float:
|
||||
return 0.0
|
||||
|
||||
def get_choice(self, idx: int) -> str | None:
|
||||
if idx < len(self.choices):
|
||||
return self.choices[idx]
|
||||
return None
|
||||
|
||||
def fill_choices(self) -> None:
|
||||
self.choices.append('extern_mm')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=128_BLOCK-N=16_numstages=4_numwarps=8')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=128_BLOCK-N=32_numstages=4_numwarps=8')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=128_BLOCK-N=64_numstages=4_numwarps=8')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=16_BLOCK-N=128_numstages=2_numwarps=8')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=16_BLOCK-N=128_numstages=3_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=16_BLOCK-N=128_numstages=3_numwarps=8')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=16_BLOCK-N=128_numstages=4_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=16_BLOCK-N=128_numstages=5_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=16_BLOCK-N=128_numstages=5_numwarps=8')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=16_BLOCK-N=16_numstages=2_numwarps=2')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=16_BLOCK-N=16_numstages=2_numwarps=8')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=16_BLOCK-N=16_numstages=3_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=16_BLOCK-N=16_numstages=3_numwarps=8')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=16_BLOCK-N=16_numstages=4_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=16_BLOCK-N=16_numstages=4_numwarps=8')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=16_BLOCK-N=16_numstages=5_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=16_BLOCK-N=16_numstages=5_numwarps=8')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=16_BLOCK-N=32_numstages=2_numwarps=2')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=16_BLOCK-N=32_numstages=2_numwarps=8')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=16_BLOCK-N=32_numstages=3_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=16_BLOCK-N=32_numstages=3_numwarps=8')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=16_BLOCK-N=32_numstages=4_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=16_BLOCK-N=32_numstages=4_numwarps=8')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=16_BLOCK-N=32_numstages=5_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=16_BLOCK-N=32_numstages=5_numwarps=8')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=16_BLOCK-N=64_numstages=2_numwarps=2')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=16_BLOCK-N=64_numstages=2_numwarps=8')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=16_BLOCK-N=64_numstages=3_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=16_BLOCK-N=64_numstages=3_numwarps=8')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=16_BLOCK-N=64_numstages=4_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=16_BLOCK-N=64_numstages=4_numwarps=8')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=16_BLOCK-N=64_numstages=5_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=16_BLOCK-N=64_numstages=5_numwarps=8')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=32_BLOCK-N=128_numstages=2_numwarps=8')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=32_BLOCK-N=128_numstages=3_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=32_BLOCK-N=128_numstages=3_numwarps=8')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=32_BLOCK-N=128_numstages=4_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=32_BLOCK-N=128_numstages=5_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=32_BLOCK-N=128_numstages=5_numwarps=8')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=32_BLOCK-N=16_numstages=2_numwarps=8')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=32_BLOCK-N=16_numstages=3_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=32_BLOCK-N=16_numstages=3_numwarps=8')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=32_BLOCK-N=16_numstages=4_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=32_BLOCK-N=16_numstages=4_numwarps=8')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=32_BLOCK-N=16_numstages=5_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=32_BLOCK-N=16_numstages=5_numwarps=8')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=32_BLOCK-N=32_numstages=2_numwarps=2')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=32_BLOCK-N=32_numstages=2_numwarps=8')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=32_BLOCK-N=32_numstages=3_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=32_BLOCK-N=32_numstages=3_numwarps=8')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=32_BLOCK-N=32_numstages=4_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=32_BLOCK-N=32_numstages=4_numwarps=8')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=32_BLOCK-N=32_numstages=5_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=32_BLOCK-N=32_numstages=5_numwarps=8')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=32_BLOCK-N=64_numstages=2_numwarps=2')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=32_BLOCK-N=64_numstages=2_numwarps=8')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=32_BLOCK-N=64_numstages=3_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=32_BLOCK-N=64_numstages=3_numwarps=8')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=32_BLOCK-N=64_numstages=4_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=32_BLOCK-N=64_numstages=4_numwarps=8')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=32_BLOCK-N=64_numstages=5_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=32_BLOCK-N=64_numstages=5_numwarps=8')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=64_BLOCK-N=128_numstages=3_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=64_BLOCK-N=128_numstages=3_numwarps=8')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=64_BLOCK-N=128_numstages=5_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=64_BLOCK-N=128_numstages=5_numwarps=8')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=64_BLOCK-N=16_numstages=3_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=64_BLOCK-N=16_numstages=3_numwarps=8')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=64_BLOCK-N=16_numstages=4_numwarps=8')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=64_BLOCK-N=16_numstages=5_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=64_BLOCK-N=16_numstages=5_numwarps=8')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=64_BLOCK-N=32_numstages=3_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=64_BLOCK-N=32_numstages=3_numwarps=8')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=64_BLOCK-N=32_numstages=4_numwarps=8')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=64_BLOCK-N=32_numstages=5_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=64_BLOCK-N=32_numstages=5_numwarps=8')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=64_BLOCK-N=64_numstages=3_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=64_BLOCK-N=64_numstages=3_numwarps=8')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=64_BLOCK-N=64_numstages=4_numwarps=8')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=64_BLOCK-N=64_numstages=5_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=64_BLOCK-N=64_numstages=5_numwarps=8')
|
||||
self.choices.append('type=triton_BLOCK-M=16_BLOCK-K=128_BLOCK-N=128_numstages=4_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=16_BLOCK-K=128_BLOCK-N=32_numstages=5_numwarps=2')
|
||||
self.choices.append('type=triton_BLOCK-M=16_BLOCK-K=128_BLOCK-N=64_numstages=3_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=16_BLOCK-K=128_BLOCK-N=64_numstages=4_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=16_BLOCK-K=128_BLOCK-N=64_numstages=5_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=16_BLOCK-K=16_BLOCK-N=128_numstages=3_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=16_BLOCK-K=16_BLOCK-N=128_numstages=3_numwarps=8')
|
||||
self.choices.append('type=triton_BLOCK-M=16_BLOCK-K=16_BLOCK-N=128_numstages=4_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=16_BLOCK-K=16_BLOCK-N=128_numstages=4_numwarps=8')
|
||||
self.choices.append('type=triton_BLOCK-M=16_BLOCK-K=16_BLOCK-N=128_numstages=5_numwarps=8')
|
||||
self.choices.append('type=triton_BLOCK-M=16_BLOCK-K=16_BLOCK-N=16_numstages=5_numwarps=1')
|
||||
self.choices.append('type=triton_BLOCK-M=16_BLOCK-K=16_BLOCK-N=32_numstages=1_numwarps=2')
|
||||
self.choices.append('type=triton_BLOCK-M=16_BLOCK-K=16_BLOCK-N=32_numstages=2_numwarps=2')
|
||||
self.choices.append('type=triton_BLOCK-M=16_BLOCK-K=16_BLOCK-N=32_numstages=3_numwarps=2')
|
||||
self.choices.append('type=triton_BLOCK-M=16_BLOCK-K=16_BLOCK-N=32_numstages=4_numwarps=2')
|
||||
self.choices.append('type=triton_BLOCK-M=16_BLOCK-K=16_BLOCK-N=32_numstages=5_numwarps=2')
|
||||
self.choices.append('type=triton_BLOCK-M=16_BLOCK-K=16_BLOCK-N=64_numstages=2_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=16_BLOCK-K=16_BLOCK-N=64_numstages=4_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=16_BLOCK-K=16_BLOCK-N=64_numstages=5_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=16_BLOCK-K=32_BLOCK-N=128_numstages=3_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=16_BLOCK-K=32_BLOCK-N=128_numstages=4_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=16_BLOCK-K=32_BLOCK-N=16_numstages=5_numwarps=1')
|
||||
self.choices.append('type=triton_BLOCK-M=16_BLOCK-K=32_BLOCK-N=32_numstages=5_numwarps=2')
|
||||
self.choices.append('type=triton_BLOCK-M=16_BLOCK-K=32_BLOCK-N=64_numstages=3_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=16_BLOCK-K=32_BLOCK-N=64_numstages=4_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=16_BLOCK-K=32_BLOCK-N=64_numstages=5_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=16_BLOCK-K=64_BLOCK-N=128_numstages=5_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=16_BLOCK-K=64_BLOCK-N=128_numstages=5_numwarps=8')
|
||||
self.choices.append('type=triton_BLOCK-M=16_BLOCK-K=64_BLOCK-N=32_numstages=2_numwarps=2')
|
||||
self.choices.append('type=triton_BLOCK-M=16_BLOCK-K=64_BLOCK-N=64_numstages=3_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=16_BLOCK-K=64_BLOCK-N=64_numstages=4_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=16_BLOCK-K=64_BLOCK-N=64_numstages=5_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=32_BLOCK-K=128_BLOCK-N=128_numstages=4_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=32_BLOCK-K=128_BLOCK-N=16_numstages=2_numwarps=2')
|
||||
self.choices.append('type=triton_BLOCK-M=32_BLOCK-K=128_BLOCK-N=32_numstages=2_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=32_BLOCK-K=128_BLOCK-N=32_numstages=5_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=32_BLOCK-K=128_BLOCK-N=64_numstages=3_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=32_BLOCK-K=128_BLOCK-N=64_numstages=4_numwarps=8')
|
||||
self.choices.append('type=triton_BLOCK-M=32_BLOCK-K=128_BLOCK-N=64_numstages=5_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=32_BLOCK-K=16_BLOCK-N=16_numstages=1_numwarps=2')
|
||||
self.choices.append('type=triton_BLOCK-M=32_BLOCK-K=16_BLOCK-N=16_numstages=2_numwarps=2')
|
||||
self.choices.append('type=triton_BLOCK-M=32_BLOCK-K=16_BLOCK-N=16_numstages=5_numwarps=2')
|
||||
self.choices.append('type=triton_BLOCK-M=32_BLOCK-K=16_BLOCK-N=32_numstages=1_numwarps=2')
|
||||
self.choices.append('type=triton_BLOCK-M=32_BLOCK-K=16_BLOCK-N=32_numstages=2_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=32_BLOCK-K=16_BLOCK-N=32_numstages=5_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=32_BLOCK-K=16_BLOCK-N=64_numstages=5_numwarps=8')
|
||||
self.choices.append('type=triton_BLOCK-M=32_BLOCK-K=32_BLOCK-N=16_numstages=2_numwarps=2')
|
||||
self.choices.append('type=triton_BLOCK-M=32_BLOCK-K=32_BLOCK-N=16_numstages=5_numwarps=2')
|
||||
self.choices.append('type=triton_BLOCK-M=32_BLOCK-K=32_BLOCK-N=64_numstages=5_numwarps=8')
|
||||
self.choices.append('type=triton_BLOCK-M=32_BLOCK-K=64_BLOCK-N=128_numstages=3_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=32_BLOCK-K=64_BLOCK-N=128_numstages=5_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=32_BLOCK-K=64_BLOCK-N=32_numstages=2_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=64_BLOCK-K=128_BLOCK-N=128_numstages=4_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=64_BLOCK-K=128_BLOCK-N=16_numstages=3_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=64_BLOCK-K=128_BLOCK-N=16_numstages=4_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=64_BLOCK-K=128_BLOCK-N=16_numstages=5_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=64_BLOCK-K=128_BLOCK-N=32_numstages=3_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=64_BLOCK-K=128_BLOCK-N=32_numstages=4_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=64_BLOCK-K=128_BLOCK-N=32_numstages=5_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=64_BLOCK-K=128_BLOCK-N=64_numstages=3_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=64_BLOCK-K=128_BLOCK-N=64_numstages=4_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=64_BLOCK-K=128_BLOCK-N=64_numstages=4_numwarps=8')
|
||||
self.choices.append('type=triton_BLOCK-M=64_BLOCK-K=128_BLOCK-N=64_numstages=5_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=64_BLOCK-K=16_BLOCK-N=128_numstages=3_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=64_BLOCK-K=16_BLOCK-N=128_numstages=4_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=64_BLOCK-K=16_BLOCK-N=128_numstages=4_numwarps=8')
|
||||
self.choices.append('type=triton_BLOCK-M=64_BLOCK-K=16_BLOCK-N=16_numstages=2_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=64_BLOCK-K=16_BLOCK-N=16_numstages=3_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=64_BLOCK-K=16_BLOCK-N=16_numstages=4_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=64_BLOCK-K=16_BLOCK-N=16_numstages=5_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=64_BLOCK-K=16_BLOCK-N=32_numstages=2_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=64_BLOCK-K=16_BLOCK-N=32_numstages=3_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=64_BLOCK-K=16_BLOCK-N=32_numstages=3_numwarps=8')
|
||||
self.choices.append('type=triton_BLOCK-M=64_BLOCK-K=16_BLOCK-N=32_numstages=4_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=64_BLOCK-K=16_BLOCK-N=32_numstages=4_numwarps=8')
|
||||
self.choices.append('type=triton_BLOCK-M=64_BLOCK-K=16_BLOCK-N=32_numstages=5_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=64_BLOCK-K=16_BLOCK-N=32_numstages=5_numwarps=8')
|
||||
self.choices.append('type=triton_BLOCK-M=64_BLOCK-K=16_BLOCK-N=64_numstages=2_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=64_BLOCK-K=16_BLOCK-N=64_numstages=3_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=64_BLOCK-K=16_BLOCK-N=64_numstages=3_numwarps=8')
|
||||
self.choices.append('type=triton_BLOCK-M=64_BLOCK-K=16_BLOCK-N=64_numstages=4_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=64_BLOCK-K=16_BLOCK-N=64_numstages=4_numwarps=8')
|
||||
self.choices.append('type=triton_BLOCK-M=64_BLOCK-K=16_BLOCK-N=64_numstages=5_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=64_BLOCK-K=32_BLOCK-N=128_numstages=3_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=64_BLOCK-K=32_BLOCK-N=128_numstages=4_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=64_BLOCK-K=32_BLOCK-N=128_numstages=4_numwarps=8')
|
||||
self.choices.append('type=triton_BLOCK-M=64_BLOCK-K=32_BLOCK-N=16_numstages=2_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=64_BLOCK-K=32_BLOCK-N=16_numstages=3_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=64_BLOCK-K=32_BLOCK-N=16_numstages=4_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=64_BLOCK-K=32_BLOCK-N=16_numstages=5_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=64_BLOCK-K=32_BLOCK-N=32_numstages=2_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=64_BLOCK-K=32_BLOCK-N=32_numstages=3_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=64_BLOCK-K=32_BLOCK-N=32_numstages=3_numwarps=8')
|
||||
self.choices.append('type=triton_BLOCK-M=64_BLOCK-K=32_BLOCK-N=32_numstages=4_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=64_BLOCK-K=32_BLOCK-N=32_numstages=4_numwarps=8')
|
||||
self.choices.append('type=triton_BLOCK-M=64_BLOCK-K=32_BLOCK-N=32_numstages=5_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=64_BLOCK-K=32_BLOCK-N=32_numstages=5_numwarps=8')
|
||||
self.choices.append('type=triton_BLOCK-M=64_BLOCK-K=32_BLOCK-N=64_numstages=2_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=64_BLOCK-K=32_BLOCK-N=64_numstages=3_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=64_BLOCK-K=32_BLOCK-N=64_numstages=3_numwarps=8')
|
||||
self.choices.append('type=triton_BLOCK-M=64_BLOCK-K=32_BLOCK-N=64_numstages=4_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=64_BLOCK-K=32_BLOCK-N=64_numstages=4_numwarps=8')
|
||||
self.choices.append('type=triton_BLOCK-M=64_BLOCK-K=32_BLOCK-N=64_numstages=5_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=64_BLOCK-K=64_BLOCK-N=128_numstages=3_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=64_BLOCK-K=64_BLOCK-N=128_numstages=4_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=64_BLOCK-K=64_BLOCK-N=16_numstages=3_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=64_BLOCK-K=64_BLOCK-N=16_numstages=4_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=64_BLOCK-K=64_BLOCK-N=16_numstages=5_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=64_BLOCK-K=64_BLOCK-N=32_numstages=3_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=64_BLOCK-K=64_BLOCK-N=32_numstages=3_numwarps=8')
|
||||
self.choices.append('type=triton_BLOCK-M=64_BLOCK-K=64_BLOCK-N=32_numstages=4_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=64_BLOCK-K=64_BLOCK-N=32_numstages=5_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=64_BLOCK-K=64_BLOCK-N=64_numstages=3_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=64_BLOCK-K=64_BLOCK-N=64_numstages=3_numwarps=8')
|
||||
self.choices.append('type=triton_BLOCK-M=64_BLOCK-K=64_BLOCK-N=64_numstages=4_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=64_BLOCK-K=64_BLOCK-N=64_numstages=5_numwarps=4')
|
||||
|
||||
def get_name(self) -> str:
|
||||
return 'mm'
|
||||
|
||||
def get_best_choices(self, context: AHContext) -> list[tuple[float, int]] | None:
|
||||
if context.get_value('arith_intensity') <= 52.6245059967041:
|
||||
if context.get_value('n') <= 34.0:
|
||||
if context.get_value('n') <= 18.0:
|
||||
if context.get_value('k*n') <= 312.0:
|
||||
return [(0.093, 12), (0.081, 16), (0.081, 148), (0.070, 10), (0.070, 17), (0.070, 149), (0.070, 151), (0.070, 150), (0.070, 14), (0.058, 11), (0.058, 15), (0.058, 13), (0.058, 122), (0.047, 121), (0.035, 123), (0.012, 92)]
|
||||
else:
|
||||
if context.get_value('k') <= 40.0:
|
||||
return [(0.083, 42), (0.083, 46), (0.083, 44), (0.083, 40), (0.083, 128), (0.067, 45), (0.067, 43), (0.067, 41), (0.067, 169), (0.067, 171), (0.067, 168), (0.067, 129), (0.067, 170), (0.033, 103), (0.017, 121)]
|
||||
else:
|
||||
return [(0.112, 137), (0.104, 136), (0.101, 0), (0.081, 1), (0.073, 135), (0.069, 67), (0.066, 187), (0.058, 41), (0.050, 71), (0.046, 68), (0.046, 70), (0.031, 44), (0.027, 43), (0.027, 170), (0.019, 189), (0.019, 188), (0.015, 169), (0.015, 171), (0.012, 115), (0.012, 168), (0.012, 69), (0.004, 103)]
|
||||
else:
|
||||
if context.get_value('mat1_stride_0') <= 20.0:
|
||||
return [(0.069, 0), (0.059, 157), (0.059, 22), (0.059, 153), (0.059, 155), (0.059, 25), (0.059, 23), (0.059, 19), (0.044, 21), (0.044, 18), (0.044, 152), (0.044, 158), (0.044, 154), (0.044, 156), (0.044, 20), (0.044, 124), (0.044, 24), (0.030, 125), (0.029, 126), (0.015, 97), (0.015, 95), (0.015, 96), (0.010, 2), (0.010, 75)]
|
||||
else:
|
||||
if context.get_value('k') <= 68.0:
|
||||
return [(0.087, 72), (0.087, 74), (0.087, 73), (0.086, 76), (0.077, 75), (0.067, 192), (0.058, 190), (0.048, 47), (0.048, 193), (0.048, 49), (0.048, 51), (0.048, 191), (0.038, 53), (0.019, 133), (0.019, 50), (0.019, 175), (0.019, 172), (0.019, 48), (0.019, 174), (0.010, 173), (0.010, 177), (0.010, 52), (0.010, 54), (0.010, 178), (0.010, 176)]
|
||||
else:
|
||||
return [(0.154, 52), (0.154, 72), (0.102, 75), (0.087, 49), (0.087, 73), (0.086, 51), (0.057, 176), (0.045, 2), (0.038, 191), (0.038, 178), (0.038, 190), (0.029, 173), (0.029, 76), (0.026, 138), (0.013, 139), (0.013, 140), (0.003, 0)]
|
||||
else:
|
||||
if context.get_value('k') <= 35.0:
|
||||
if context.get_value('k') <= 18.0:
|
||||
if context.get_value('m*n') <= 19505152.0:
|
||||
return [(0.151, 159), (0.140, 160), (0.129, 164), (0.055, 127), (0.051, 29), (0.044, 161), (0.044, 147), (0.040, 146), (0.040, 31), (0.037, 145), (0.026, 28), (0.022, 90), (0.022, 93), (0.022, 94), (0.022, 100), (0.022, 125), (0.022, 158), (0.022, 157), (0.011, 87), (0.011, 88), (0.011, 89), (0.011, 91), (0.011, 95), (0.011, 96), (0.011, 98), (0.011, 99)]
|
||||
else:
|
||||
return [(0.069, 7), (0.069, 5), (0.067, 147), (0.066, 8), (0.061, 145), (0.058, 146), (0.052, 124), (0.049, 29), (0.049, 159), (0.046, 31), (0.043, 157), (0.041, 9), (0.041, 4), (0.040, 6), (0.035, 164), (0.035, 160), (0.026, 158), (0.017, 125), (0.017, 28), (0.017, 32), (0.017, 162), (0.017, 27), (0.017, 30), (0.017, 161), (0.009, 33), (0.009, 26), (0.009, 163), (0.006, 0)]
|
||||
else:
|
||||
if context.get_value('n') <= 68.0:
|
||||
return [(0.101, 182), (0.101, 59), (0.088, 57), (0.076, 184), (0.076, 61), (0.076, 179), (0.076, 62), (0.076, 58), (0.063, 180), (0.063, 60), (0.051, 56), (0.050, 181), (0.025, 130), (0.025, 177), (0.025, 183), (0.013, 178), (0.013, 55)]
|
||||
else:
|
||||
return [(0.089, 180), (0.079, 60), (0.066, 35), (0.066, 181), (0.066, 38), (0.066, 58), (0.066, 179), (0.066, 57), (0.062, 184), (0.053, 37), (0.044, 166), (0.040, 55), (0.040, 39), (0.040, 36), (0.040, 165), (0.040, 167), (0.027, 177), (0.027, 34), (0.022, 159)]
|
||||
else:
|
||||
if context.get_value('m*n') <= 309760.0:
|
||||
return [(0.298, 0), (0.097, 140), (0.080, 83), (0.072, 86), (0.044, 84), (0.036, 178), (0.036, 117), (0.036, 82), (0.032, 120), (0.032, 85), (0.028, 119), (0.024, 130), (0.024, 109), (0.020, 108), (0.020, 118), (0.012, 104), (0.012, 116), (0.012, 141), (0.012, 144), (0.008, 105), (0.008, 106), (0.008, 111), (0.008, 114), (0.008, 107), (0.008, 132), (0.004, 101), (0.004, 102), (0.004, 110), (0.004, 112), (0.004, 113), (0.004, 131)]
|
||||
else:
|
||||
if context.get_value('n') <= 72.0:
|
||||
return [(0.227, 77), (0.118, 78), (0.102, 194), (0.086, 80), (0.059, 57), (0.054, 81), (0.049, 196), (0.048, 197), (0.048, 59), (0.043, 79), (0.032, 195), (0.027, 180), (0.022, 3), (0.021, 141), (0.016, 60), (0.016, 142), (0.011, 183), (0.011, 0), (0.011, 144)]
|
||||
else:
|
||||
return [(0.140, 186), (0.132, 185), (0.109, 63), (0.085, 65), (0.078, 37), (0.077, 35), (0.062, 197), (0.047, 194), (0.046, 165), (0.046, 57), (0.039, 78), (0.039, 79), (0.039, 66), (0.039, 64), (0.016, 195), (0.008, 159)]
|
||||
else:
|
||||
if str(context.get_value('using_tf32')) != 'False':
|
||||
if context.get_value('m*n') <= 815360.0:
|
||||
if context.get_value('k') <= 1184.0:
|
||||
return [(0.218, 140), (0.205, 0), (0.154, 144), (0.115, 141), (0.051, 185), (0.051, 104), (0.039, 78), (0.038, 116), (0.026, 165), (0.026, 130), (0.026, 178), (0.013, 57), (0.013, 195), (0.013, 167), (0.013, 186)]
|
||||
else:
|
||||
return [(0.901, 0), (0.030, 144), (0.030, 134), (0.016, 3), (0.006, 78), (0.006, 77), (0.002, 57), (0.002, 194), (0.002, 59), (0.002, 60), (0.002, 143)]
|
||||
else:
|
||||
if context.get_value('arith_intensity') <= 187.23922729492188:
|
||||
if context.get_value('mat1_stride_0') <= 198.0:
|
||||
return [(0.273, 63), (0.158, 37), (0.152, 35), (0.127, 57), (0.097, 165), (0.053, 185), (0.031, 0), (0.028, 64), (0.014, 60), (0.014, 78), (0.009, 55), (0.008, 134), (0.005, 34), (0.005, 167), (0.005, 179), (0.005, 65), (0.005, 66), (0.005, 186), (0.005, 194), (0.002, 166)]
|
||||
else:
|
||||
return [(0.296, 63), (0.235, 0), (0.132, 64), (0.074, 37), (0.069, 78), (0.051, 185), (0.051, 35), (0.030, 57), (0.020, 77), (0.016, 194), (0.008, 66), (0.007, 65), (0.003, 3), (0.003, 165), (0.003, 141), (0.001, 134), (0.001, 166)]
|
||||
else:
|
||||
return [(0.405, 0), (0.246, 37), (0.177, 63), (0.145, 35), (0.005, 185), (0.005, 65), (0.005, 64), (0.004, 57), (0.003, 66), (0.002, 165), (0.001, 78), (0.001, 55)]
|
||||
else:
|
||||
return [(0.357, 0), (0.112, 165), (0.101, 57), (0.094, 179), (0.086, 64), (0.074, 167), (0.067, 60), (0.064, 159), (0.033, 35), (0.007, 195), (0.002, 180), (0.001, 34), (0.001, 166), (0.001, 78)]
|
||||
+321
@@ -0,0 +1,321 @@
|
||||
# flake8: noqa: B950
|
||||
# fmt: off
|
||||
# This file was generated by AutoHeuristic. Do not modify it manually!
|
||||
# To regenerate this file, take a look at the steps in the README.md file inside torchgen/_autoheuristic/mm/
|
||||
from typing import List, Optional, Tuple
|
||||
|
||||
from torch._inductor.autoheuristic.autoheuristic_utils import (
|
||||
AHContext,
|
||||
AHMetadata,
|
||||
Choice,
|
||||
)
|
||||
from torch._inductor.autoheuristic.learnedheuristic_interface import (
|
||||
LearnedHeuristicDecision,
|
||||
)
|
||||
|
||||
|
||||
class MMRankingH100(LearnedHeuristicDecision):
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.choices: list[Choice] = []
|
||||
self.fill_choices()
|
||||
|
||||
def check_precondition(self, metadata: AHMetadata, context: AHContext,) -> bool:
|
||||
return (
|
||||
metadata.name == self.get_name()
|
||||
and metadata.shared_memory == 232448
|
||||
and str(metadata.device_capa) == "(9, 0)"
|
||||
)
|
||||
|
||||
def get_confidence_threshold(self) -> float:
|
||||
return 0.0
|
||||
|
||||
def get_choice(self, idx: int) -> str | None:
|
||||
if idx < len(self.choices):
|
||||
return self.choices[idx]
|
||||
return None
|
||||
|
||||
def fill_choices(self) -> None:
|
||||
self.choices.append('extern_mm')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=128_BLOCK-N=16_numstages=4_numwarps=8')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=128_BLOCK-N=32_numstages=4_numwarps=8')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=128_BLOCK-N=64_numstages=4_numwarps=8')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=16_BLOCK-N=128_numstages=2_numwarps=8')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=16_BLOCK-N=128_numstages=3_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=16_BLOCK-N=128_numstages=3_numwarps=8')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=16_BLOCK-N=128_numstages=4_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=16_BLOCK-N=128_numstages=5_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=16_BLOCK-N=128_numstages=5_numwarps=8')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=16_BLOCK-N=16_numstages=2_numwarps=2')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=16_BLOCK-N=16_numstages=2_numwarps=8')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=16_BLOCK-N=16_numstages=3_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=16_BLOCK-N=16_numstages=3_numwarps=8')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=16_BLOCK-N=16_numstages=4_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=16_BLOCK-N=16_numstages=4_numwarps=8')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=16_BLOCK-N=16_numstages=5_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=16_BLOCK-N=16_numstages=5_numwarps=8')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=16_BLOCK-N=32_numstages=2_numwarps=2')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=16_BLOCK-N=32_numstages=2_numwarps=8')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=16_BLOCK-N=32_numstages=3_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=16_BLOCK-N=32_numstages=4_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=16_BLOCK-N=32_numstages=4_numwarps=8')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=16_BLOCK-N=32_numstages=5_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=16_BLOCK-N=64_numstages=2_numwarps=2')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=16_BLOCK-N=64_numstages=2_numwarps=8')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=16_BLOCK-N=64_numstages=3_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=16_BLOCK-N=64_numstages=3_numwarps=8')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=16_BLOCK-N=64_numstages=4_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=16_BLOCK-N=64_numstages=4_numwarps=8')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=16_BLOCK-N=64_numstages=5_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=16_BLOCK-N=64_numstages=5_numwarps=8')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=32_BLOCK-N=128_numstages=2_numwarps=8')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=32_BLOCK-N=128_numstages=3_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=32_BLOCK-N=128_numstages=3_numwarps=8')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=32_BLOCK-N=128_numstages=4_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=32_BLOCK-N=128_numstages=5_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=32_BLOCK-N=128_numstages=5_numwarps=8')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=32_BLOCK-N=16_numstages=2_numwarps=2')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=32_BLOCK-N=16_numstages=2_numwarps=8')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=32_BLOCK-N=16_numstages=3_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=32_BLOCK-N=16_numstages=3_numwarps=8')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=32_BLOCK-N=16_numstages=4_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=32_BLOCK-N=16_numstages=4_numwarps=8')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=32_BLOCK-N=16_numstages=5_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=32_BLOCK-N=16_numstages=5_numwarps=8')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=32_BLOCK-N=32_numstages=2_numwarps=2')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=32_BLOCK-N=32_numstages=2_numwarps=8')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=32_BLOCK-N=32_numstages=3_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=32_BLOCK-N=32_numstages=3_numwarps=8')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=32_BLOCK-N=32_numstages=4_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=32_BLOCK-N=32_numstages=4_numwarps=8')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=32_BLOCK-N=32_numstages=5_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=32_BLOCK-N=32_numstages=5_numwarps=8')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=32_BLOCK-N=64_numstages=2_numwarps=2')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=32_BLOCK-N=64_numstages=2_numwarps=8')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=32_BLOCK-N=64_numstages=3_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=32_BLOCK-N=64_numstages=3_numwarps=8')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=32_BLOCK-N=64_numstages=4_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=32_BLOCK-N=64_numstages=4_numwarps=8')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=32_BLOCK-N=64_numstages=5_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=32_BLOCK-N=64_numstages=5_numwarps=8')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=64_BLOCK-N=128_numstages=3_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=64_BLOCK-N=128_numstages=3_numwarps=8')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=64_BLOCK-N=128_numstages=5_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=64_BLOCK-N=128_numstages=5_numwarps=8')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=64_BLOCK-N=16_numstages=3_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=64_BLOCK-N=16_numstages=3_numwarps=8')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=64_BLOCK-N=16_numstages=4_numwarps=8')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=64_BLOCK-N=16_numstages=5_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=64_BLOCK-N=16_numstages=5_numwarps=8')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=64_BLOCK-N=32_numstages=3_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=64_BLOCK-N=32_numstages=3_numwarps=8')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=64_BLOCK-N=32_numstages=4_numwarps=8')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=64_BLOCK-N=32_numstages=5_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=64_BLOCK-N=32_numstages=5_numwarps=8')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=64_BLOCK-N=64_numstages=3_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=64_BLOCK-N=64_numstages=3_numwarps=8')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=64_BLOCK-N=64_numstages=4_numwarps=8')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=64_BLOCK-N=64_numstages=5_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=64_BLOCK-N=64_numstages=5_numwarps=8')
|
||||
self.choices.append('type=triton_BLOCK-M=16_BLOCK-K=128_BLOCK-N=128_numstages=4_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=16_BLOCK-K=128_BLOCK-N=32_numstages=2_numwarps=2')
|
||||
self.choices.append('type=triton_BLOCK-M=16_BLOCK-K=128_BLOCK-N=32_numstages=5_numwarps=2')
|
||||
self.choices.append('type=triton_BLOCK-M=16_BLOCK-K=128_BLOCK-N=64_numstages=3_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=16_BLOCK-K=128_BLOCK-N=64_numstages=4_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=16_BLOCK-K=128_BLOCK-N=64_numstages=5_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=16_BLOCK-K=16_BLOCK-N=128_numstages=2_numwarps=8')
|
||||
self.choices.append('type=triton_BLOCK-M=16_BLOCK-K=16_BLOCK-N=128_numstages=3_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=16_BLOCK-K=16_BLOCK-N=128_numstages=3_numwarps=8')
|
||||
self.choices.append('type=triton_BLOCK-M=16_BLOCK-K=16_BLOCK-N=128_numstages=4_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=16_BLOCK-K=16_BLOCK-N=128_numstages=4_numwarps=8')
|
||||
self.choices.append('type=triton_BLOCK-M=16_BLOCK-K=16_BLOCK-N=128_numstages=5_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=16_BLOCK-K=16_BLOCK-N=128_numstages=5_numwarps=8')
|
||||
self.choices.append('type=triton_BLOCK-M=16_BLOCK-K=16_BLOCK-N=16_numstages=3_numwarps=1')
|
||||
self.choices.append('type=triton_BLOCK-M=16_BLOCK-K=16_BLOCK-N=16_numstages=4_numwarps=1')
|
||||
self.choices.append('type=triton_BLOCK-M=16_BLOCK-K=16_BLOCK-N=16_numstages=5_numwarps=1')
|
||||
self.choices.append('type=triton_BLOCK-M=16_BLOCK-K=16_BLOCK-N=32_numstages=1_numwarps=2')
|
||||
self.choices.append('type=triton_BLOCK-M=16_BLOCK-K=16_BLOCK-N=32_numstages=2_numwarps=2')
|
||||
self.choices.append('type=triton_BLOCK-M=16_BLOCK-K=16_BLOCK-N=32_numstages=3_numwarps=2')
|
||||
self.choices.append('type=triton_BLOCK-M=16_BLOCK-K=16_BLOCK-N=32_numstages=4_numwarps=2')
|
||||
self.choices.append('type=triton_BLOCK-M=16_BLOCK-K=16_BLOCK-N=32_numstages=5_numwarps=2')
|
||||
self.choices.append('type=triton_BLOCK-M=16_BLOCK-K=16_BLOCK-N=64_numstages=2_numwarps=2')
|
||||
self.choices.append('type=triton_BLOCK-M=16_BLOCK-K=16_BLOCK-N=64_numstages=2_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=16_BLOCK-K=16_BLOCK-N=64_numstages=3_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=16_BLOCK-K=16_BLOCK-N=64_numstages=4_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=16_BLOCK-K=16_BLOCK-N=64_numstages=5_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=16_BLOCK-K=32_BLOCK-N=128_numstages=2_numwarps=8')
|
||||
self.choices.append('type=triton_BLOCK-M=16_BLOCK-K=32_BLOCK-N=128_numstages=3_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=16_BLOCK-K=32_BLOCK-N=128_numstages=4_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=16_BLOCK-K=32_BLOCK-N=128_numstages=4_numwarps=8')
|
||||
self.choices.append('type=triton_BLOCK-M=16_BLOCK-K=32_BLOCK-N=16_numstages=4_numwarps=1')
|
||||
self.choices.append('type=triton_BLOCK-M=16_BLOCK-K=32_BLOCK-N=16_numstages=5_numwarps=1')
|
||||
self.choices.append('type=triton_BLOCK-M=16_BLOCK-K=32_BLOCK-N=32_numstages=4_numwarps=2')
|
||||
self.choices.append('type=triton_BLOCK-M=16_BLOCK-K=32_BLOCK-N=32_numstages=5_numwarps=2')
|
||||
self.choices.append('type=triton_BLOCK-M=16_BLOCK-K=32_BLOCK-N=64_numstages=2_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=16_BLOCK-K=32_BLOCK-N=64_numstages=4_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=16_BLOCK-K=32_BLOCK-N=64_numstages=5_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=16_BLOCK-K=64_BLOCK-N=128_numstages=3_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=16_BLOCK-K=64_BLOCK-N=128_numstages=3_numwarps=8')
|
||||
self.choices.append('type=triton_BLOCK-M=16_BLOCK-K=64_BLOCK-N=128_numstages=5_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=16_BLOCK-K=64_BLOCK-N=128_numstages=5_numwarps=8')
|
||||
self.choices.append('type=triton_BLOCK-M=16_BLOCK-K=64_BLOCK-N=64_numstages=3_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=32_BLOCK-K=128_BLOCK-N=16_numstages=2_numwarps=2')
|
||||
self.choices.append('type=triton_BLOCK-M=32_BLOCK-K=128_BLOCK-N=32_numstages=2_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=32_BLOCK-K=128_BLOCK-N=32_numstages=5_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=32_BLOCK-K=128_BLOCK-N=64_numstages=5_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=32_BLOCK-K=16_BLOCK-N=16_numstages=1_numwarps=2')
|
||||
self.choices.append('type=triton_BLOCK-M=32_BLOCK-K=16_BLOCK-N=16_numstages=2_numwarps=2')
|
||||
self.choices.append('type=triton_BLOCK-M=32_BLOCK-K=16_BLOCK-N=16_numstages=5_numwarps=2')
|
||||
self.choices.append('type=triton_BLOCK-M=32_BLOCK-K=16_BLOCK-N=32_numstages=1_numwarps=2')
|
||||
self.choices.append('type=triton_BLOCK-M=32_BLOCK-K=16_BLOCK-N=32_numstages=2_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=32_BLOCK-K=16_BLOCK-N=32_numstages=5_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=32_BLOCK-K=16_BLOCK-N=64_numstages=5_numwarps=8')
|
||||
self.choices.append('type=triton_BLOCK-M=32_BLOCK-K=32_BLOCK-N=16_numstages=2_numwarps=2')
|
||||
self.choices.append('type=triton_BLOCK-M=32_BLOCK-K=32_BLOCK-N=16_numstages=5_numwarps=2')
|
||||
self.choices.append('type=triton_BLOCK-M=32_BLOCK-K=32_BLOCK-N=32_numstages=2_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=32_BLOCK-K=32_BLOCK-N=32_numstages=5_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=32_BLOCK-K=32_BLOCK-N=64_numstages=5_numwarps=8')
|
||||
self.choices.append('type=triton_BLOCK-M=32_BLOCK-K=64_BLOCK-N=16_numstages=2_numwarps=2')
|
||||
self.choices.append('type=triton_BLOCK-M=32_BLOCK-K=64_BLOCK-N=32_numstages=2_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=64_BLOCK-K=128_BLOCK-N=128_numstages=4_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=64_BLOCK-K=128_BLOCK-N=16_numstages=3_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=64_BLOCK-K=128_BLOCK-N=16_numstages=4_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=64_BLOCK-K=128_BLOCK-N=16_numstages=5_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=64_BLOCK-K=128_BLOCK-N=32_numstages=3_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=64_BLOCK-K=128_BLOCK-N=32_numstages=4_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=64_BLOCK-K=128_BLOCK-N=32_numstages=5_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=64_BLOCK-K=128_BLOCK-N=64_numstages=3_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=64_BLOCK-K=128_BLOCK-N=64_numstages=4_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=64_BLOCK-K=128_BLOCK-N=64_numstages=5_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=64_BLOCK-K=16_BLOCK-N=128_numstages=3_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=64_BLOCK-K=16_BLOCK-N=128_numstages=4_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=64_BLOCK-K=16_BLOCK-N=128_numstages=4_numwarps=8')
|
||||
self.choices.append('type=triton_BLOCK-M=64_BLOCK-K=16_BLOCK-N=16_numstages=2_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=64_BLOCK-K=16_BLOCK-N=16_numstages=3_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=64_BLOCK-K=16_BLOCK-N=16_numstages=4_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=64_BLOCK-K=16_BLOCK-N=16_numstages=5_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=64_BLOCK-K=16_BLOCK-N=32_numstages=2_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=64_BLOCK-K=16_BLOCK-N=32_numstages=3_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=64_BLOCK-K=16_BLOCK-N=32_numstages=4_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=64_BLOCK-K=16_BLOCK-N=32_numstages=5_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=64_BLOCK-K=16_BLOCK-N=32_numstages=5_numwarps=8')
|
||||
self.choices.append('type=triton_BLOCK-M=64_BLOCK-K=16_BLOCK-N=64_numstages=2_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=64_BLOCK-K=16_BLOCK-N=64_numstages=3_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=64_BLOCK-K=16_BLOCK-N=64_numstages=3_numwarps=8')
|
||||
self.choices.append('type=triton_BLOCK-M=64_BLOCK-K=16_BLOCK-N=64_numstages=4_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=64_BLOCK-K=16_BLOCK-N=64_numstages=4_numwarps=8')
|
||||
self.choices.append('type=triton_BLOCK-M=64_BLOCK-K=16_BLOCK-N=64_numstages=5_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=64_BLOCK-K=32_BLOCK-N=128_numstages=3_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=64_BLOCK-K=32_BLOCK-N=128_numstages=4_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=64_BLOCK-K=32_BLOCK-N=128_numstages=4_numwarps=8')
|
||||
self.choices.append('type=triton_BLOCK-M=64_BLOCK-K=32_BLOCK-N=16_numstages=2_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=64_BLOCK-K=32_BLOCK-N=16_numstages=3_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=64_BLOCK-K=32_BLOCK-N=16_numstages=4_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=64_BLOCK-K=32_BLOCK-N=16_numstages=5_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=64_BLOCK-K=32_BLOCK-N=32_numstages=2_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=64_BLOCK-K=32_BLOCK-N=32_numstages=3_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=64_BLOCK-K=32_BLOCK-N=32_numstages=3_numwarps=8')
|
||||
self.choices.append('type=triton_BLOCK-M=64_BLOCK-K=32_BLOCK-N=32_numstages=4_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=64_BLOCK-K=32_BLOCK-N=32_numstages=4_numwarps=8')
|
||||
self.choices.append('type=triton_BLOCK-M=64_BLOCK-K=32_BLOCK-N=32_numstages=5_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=64_BLOCK-K=32_BLOCK-N=32_numstages=5_numwarps=8')
|
||||
self.choices.append('type=triton_BLOCK-M=64_BLOCK-K=32_BLOCK-N=64_numstages=2_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=64_BLOCK-K=32_BLOCK-N=64_numstages=3_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=64_BLOCK-K=32_BLOCK-N=64_numstages=3_numwarps=8')
|
||||
self.choices.append('type=triton_BLOCK-M=64_BLOCK-K=32_BLOCK-N=64_numstages=4_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=64_BLOCK-K=32_BLOCK-N=64_numstages=4_numwarps=8')
|
||||
self.choices.append('type=triton_BLOCK-M=64_BLOCK-K=32_BLOCK-N=64_numstages=5_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=64_BLOCK-K=64_BLOCK-N=128_numstages=3_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=64_BLOCK-K=64_BLOCK-N=128_numstages=4_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=64_BLOCK-K=64_BLOCK-N=16_numstages=3_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=64_BLOCK-K=64_BLOCK-N=16_numstages=4_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=64_BLOCK-K=64_BLOCK-N=16_numstages=5_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=64_BLOCK-K=64_BLOCK-N=32_numstages=3_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=64_BLOCK-K=64_BLOCK-N=32_numstages=3_numwarps=8')
|
||||
self.choices.append('type=triton_BLOCK-M=64_BLOCK-K=64_BLOCK-N=32_numstages=4_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=64_BLOCK-K=64_BLOCK-N=32_numstages=5_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=64_BLOCK-K=64_BLOCK-N=64_numstages=3_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=64_BLOCK-K=64_BLOCK-N=64_numstages=3_numwarps=8')
|
||||
self.choices.append('type=triton_BLOCK-M=64_BLOCK-K=64_BLOCK-N=64_numstages=4_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=64_BLOCK-K=64_BLOCK-N=64_numstages=5_numwarps=4')
|
||||
|
||||
def get_name(self) -> str:
|
||||
return 'mm'
|
||||
|
||||
def get_best_choices(self, context: AHContext) -> list[tuple[float, int]] | None:
|
||||
if context.get_value('arith_intensity') <= 29.89772129058838:
|
||||
if context.get_value('n') <= 34.0:
|
||||
if context.get_value('n') <= 18.0:
|
||||
if context.get_value('k*n') <= 432.0:
|
||||
if context.get_value('arith_intensity') <= 7.8700292110443115:
|
||||
return [(0.098, 128), (0.098, 129), (0.098, 127), (0.073, 14), (0.073, 16), (0.073, 12), (0.073, 154), (0.073, 156), (0.073, 157), (0.073, 155), (0.049, 10), (0.049, 94), (0.049, 95), (0.048, 96)]
|
||||
else:
|
||||
return [(0.091, 154), (0.073, 10), (0.073, 15), (0.073, 13), (0.073, 11), (0.073, 17), (0.073, 16), (0.073, 14), (0.073, 12), (0.055, 127), (0.054, 157), (0.054, 156), (0.054, 155), (0.036, 129), (0.036, 128), (0.018, 41), (0.018, 43)]
|
||||
else:
|
||||
if context.get_value('k') <= 40.0:
|
||||
return [(0.070, 39), (0.069, 45), (0.069, 41), (0.069, 43), (0.069, 111), (0.069, 112), (0.056, 38), (0.056, 40), (0.056, 42), (0.056, 44), (0.056, 174), (0.056, 173), (0.056, 175), (0.056, 134), (0.056, 172), (0.056, 135), (0.014, 154), (0.014, 127)]
|
||||
else:
|
||||
return [(0.147, 144), (0.119, 143), (0.087, 142), (0.083, 0), (0.073, 191), (0.059, 69), (0.050, 67), (0.046, 70), (0.041, 1), (0.036, 174), (0.032, 43), (0.032, 123), (0.028, 40), (0.027, 42), (0.027, 173), (0.023, 175), (0.018, 66), (0.014, 192), (0.014, 193), (0.014, 139), (0.014, 68), (0.014, 127)]
|
||||
else:
|
||||
if context.get_value('mat1_stride_0') <= 40.0:
|
||||
if context.get_value('mat1_stride_0') <= 20.0:
|
||||
return [(0.109, 23), (0.109, 21), (0.109, 20), (0.088, 0), (0.087, 131), (0.066, 18), (0.065, 130), (0.065, 132), (0.065, 159), (0.065, 160), (0.065, 161), (0.065, 158), (0.022, 22), (0.022, 19)]
|
||||
else:
|
||||
return [(0.065, 46), (0.064, 52), (0.064, 50), (0.064, 48), (0.064, 51), (0.064, 49), (0.064, 47), (0.064, 53), (0.064, 181), (0.064, 177), (0.064, 179), (0.064, 176), (0.038, 130), (0.038, 136), (0.026, 182), (0.026, 178), (0.026, 180), (0.026, 137), (0.025, 158), (0.013, 114), (0.013, 113)]
|
||||
else:
|
||||
if context.get_value('mat1_stride_0') <= 68.0:
|
||||
return [(0.138, 140), (0.125, 195), (0.100, 71), (0.100, 74), (0.100, 196), (0.100, 194), (0.100, 197), (0.075, 75), (0.062, 72), (0.062, 73), (0.012, 180), (0.012, 51), (0.012, 182)]
|
||||
else:
|
||||
return [(0.124, 180), (0.124, 182), (0.114, 75), (0.103, 74), (0.093, 51), (0.093, 71), (0.072, 72), (0.062, 194), (0.052, 145), (0.052, 195), (0.021, 48), (0.021, 50), (0.021, 47), (0.020, 124), (0.010, 147), (0.010, 146), (0.010, 46)]
|
||||
else:
|
||||
if context.get_value('k') <= 18.0:
|
||||
if context.get_value('m*k') <= 528.0:
|
||||
return [(0.097, 88), (0.087, 92), (0.077, 90), (0.058, 105), (0.058, 103), (0.058, 104), (0.058, 99), (0.058, 100), (0.058, 106), (0.058, 93), (0.057, 91), (0.057, 97), (0.057, 98), (0.057, 101), (0.048, 102), (0.029, 87), (0.029, 89)]
|
||||
else:
|
||||
if context.get_value('n') <= 80.0:
|
||||
return [(0.057, 161), (0.057, 130), (0.057, 24), (0.056, 164), (0.056, 163), (0.056, 166), (0.056, 168), (0.056, 30), (0.056, 28), (0.056, 26), (0.056, 25), (0.056, 27), (0.056, 29), (0.056, 31), (0.042, 131), (0.028, 99), (0.028, 101), (0.028, 100), (0.028, 167), (0.028, 165), (0.028, 133)]
|
||||
else:
|
||||
return [(0.110, 164), (0.108, 163), (0.106, 168), (0.069, 161), (0.066, 151), (0.060, 152), (0.055, 165), (0.050, 27), (0.050, 29), (0.048, 131), (0.043, 153), (0.037, 133), (0.037, 130), (0.028, 8), (0.028, 5), (0.027, 7), (0.026, 26), (0.016, 162), (0.012, 9), (0.007, 4), (0.005, 100), (0.005, 6), (0.005, 24)]
|
||||
else:
|
||||
if context.get_value('k') <= 36.0:
|
||||
if context.get_value('n') <= 68.0:
|
||||
return [(0.097, 184), (0.097, 56), (0.086, 186), (0.086, 183), (0.086, 188), (0.086, 58), (0.086, 60), (0.065, 54), (0.043, 187), (0.043, 185), (0.043, 57), (0.043, 61), (0.032, 55), (0.032, 130), (0.032, 59), (0.011, 181), (0.011, 163), (0.011, 136), (0.011, 138)]
|
||||
else:
|
||||
return [(0.117, 184), (0.117, 170), (0.117, 169), (0.107, 183), (0.106, 188), (0.075, 181), (0.064, 130), (0.064, 56), (0.053, 171), (0.032, 57), (0.032, 59), (0.032, 185), (0.011, 163), (0.011, 32), (0.011, 37), (0.011, 34), (0.011, 33), (0.011, 35), (0.011, 36), (0.011, 54)]
|
||||
else:
|
||||
if context.get_value('mat2_stride_0') <= 384.0:
|
||||
return [(0.244, 0), (0.061, 76), (0.061, 79), (0.030, 3), (0.030, 183), (0.030, 189), (0.030, 187), (0.030, 64), (0.030, 190), (0.030, 62), (0.030, 198), (0.030, 201), (0.030, 77), (0.030, 200), (0.030, 80), (0.030, 199), (0.030, 78), (0.030, 184), (0.020, 86), (0.020, 84), (0.020, 120), (0.020, 81), (0.020, 121), (0.020, 85), (0.020, 122), (0.010, 83), (0.010, 118), (0.010, 119), (0.010, 82)]
|
||||
else:
|
||||
return [(0.274, 83), (0.171, 86), (0.152, 0), (0.071, 85), (0.061, 125), (0.050, 84), (0.020, 109), (0.020, 117), (0.020, 81), (0.020, 118), (0.020, 121), (0.020, 108), (0.020, 115), (0.020, 116), (0.010, 110), (0.010, 120), (0.010, 103), (0.010, 107), (0.010, 119), (0.010, 122)]
|
||||
else:
|
||||
if context.get_value('arith_intensity') <= 56.995582580566406:
|
||||
if context.get_value('n') <= 68.0:
|
||||
if context.get_value('k*n') <= 4448.0:
|
||||
if context.get_value('m*n') <= 29626368.0:
|
||||
return [(0.107, 198), (0.107, 200), (0.107, 201), (0.107, 199), (0.106, 76), (0.106, 79), (0.064, 197), (0.063, 56), (0.043, 184), (0.043, 187), (0.042, 80), (0.042, 77), (0.042, 183), (0.021, 78)]
|
||||
else:
|
||||
return [(0.073, 201), (0.073, 198), (0.073, 200), (0.073, 199), (0.073, 197), (0.073, 56), (0.073, 58), (0.073, 79), (0.073, 76), (0.072, 59), (0.072, 78), (0.072, 77), (0.072, 80), (0.018, 184), (0.018, 55), (0.018, 54)]
|
||||
else:
|
||||
if context.get_value('k') <= 348.0:
|
||||
return [(0.206, 76), (0.183, 77), (0.169, 198), (0.160, 199), (0.053, 59), (0.046, 56), (0.038, 3), (0.030, 148), (0.030, 58), (0.030, 187), (0.023, 184), (0.015, 0), (0.008, 55), (0.008, 54)]
|
||||
else:
|
||||
return [(0.146, 198), (0.145, 199), (0.145, 148), (0.126, 0), (0.084, 76), (0.084, 77), (0.042, 80), (0.042, 79), (0.021, 149), (0.021, 150), (0.021, 3), (0.014, 46), (0.014, 74), (0.014, 75), (0.014, 124), (0.014, 194), (0.014, 195), (0.007, 145), (0.007, 146), (0.007, 2), (0.007, 72), (0.007, 147), (0.007, 71)]
|
||||
else:
|
||||
if context.get_value('m') <= 3264.0:
|
||||
return [(0.247, 147), (0.115, 197), (0.066, 199), (0.066, 201), (0.066, 198), (0.049, 0), (0.049, 169), (0.049, 171), (0.033, 140), (0.033, 125), (0.033, 114), (0.016, 126), (0.016, 183), (0.016, 184), (0.016, 185), (0.016, 182), (0.016, 188), (0.016, 78), (0.016, 148), (0.016, 138), (0.016, 77), (0.016, 56), (0.016, 59)]
|
||||
else:
|
||||
if context.get_value('k') <= 62.5:
|
||||
return [(0.226, 190), (0.226, 189), (0.122, 62), (0.122, 64), (0.055, 77), (0.055, 78), (0.037, 198), (0.036, 201), (0.036, 33), (0.024, 163), (0.018, 56), (0.018, 35), (0.018, 169), (0.006, 171)]
|
||||
else:
|
||||
return [(0.162, 35), (0.118, 33), (0.096, 189), (0.096, 190), (0.088, 169), (0.074, 62), (0.073, 56), (0.066, 171), (0.051, 198), (0.051, 201), (0.044, 59), (0.037, 64), (0.029, 63), (0.007, 0), (0.007, 77)]
|
||||
else:
|
||||
if context.get_value('m*n') <= 1097728.0:
|
||||
return [(0.403, 0), (0.179, 141), (0.134, 150), (0.086, 147), (0.051, 148), (0.048, 3), (0.024, 189), (0.020, 199), (0.017, 64), (0.010, 65), (0.010, 77), (0.007, 114), (0.003, 138), (0.003, 59), (0.003, 182)]
|
||||
else:
|
||||
if context.get_value('m*n') <= 3244032.0:
|
||||
return [(0.295, 189), (0.176, 64), (0.157, 65), (0.090, 0), (0.069, 62), (0.059, 63), (0.046, 77), (0.039, 169), (0.023, 199), (0.020, 35), (0.013, 33), (0.010, 171), (0.003, 141)]
|
||||
else:
|
||||
if context.get_value('n') <= 136.0:
|
||||
return [(0.197, 189), (0.197, 63), (0.161, 77), (0.157, 62), (0.061, 33), (0.044, 65), (0.039, 35), (0.039, 64), (0.030, 169), (0.026, 0), (0.017, 199), (0.017, 148), (0.009, 56), (0.004, 3)]
|
||||
else:
|
||||
return [(0.460, 0), (0.145, 62), (0.138, 63), (0.081, 35), (0.047, 33), (0.043, 189), (0.023, 64), (0.018, 77), (0.013, 169), (0.009, 65), (0.009, 56), (0.005, 32), (0.005, 59), (0.002, 183), (0.002, 163)]
|
||||
+150
@@ -0,0 +1,150 @@
|
||||
# flake8: noqa: B950
|
||||
# fmt: off
|
||||
# This file was generated by AutoHeuristic. Do not modify it manually!
|
||||
# To regenerate this file, take a look at the steps in the README.md file inside torchgen/_autoheuristic/mixed_mm/
|
||||
from typing import List, Optional, Tuple
|
||||
|
||||
from torch._inductor.autoheuristic.autoheuristic_utils import (
|
||||
AHContext,
|
||||
AHMetadata,
|
||||
Choice,
|
||||
)
|
||||
from torch._inductor.autoheuristic.learnedheuristic_interface import (
|
||||
LearnedHeuristicDecision,
|
||||
)
|
||||
|
||||
|
||||
class MixedMMA100(LearnedHeuristicDecision):
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.choices: list[Choice] = []
|
||||
self.fill_choices()
|
||||
|
||||
def check_precondition(self, metadata: AHMetadata, context: AHContext,) -> bool:
|
||||
return (
|
||||
metadata.name == self.get_name()
|
||||
and metadata.shared_memory == 166912
|
||||
and str(metadata.device_capa) == "(8, 0)"
|
||||
)
|
||||
|
||||
def get_confidence_threshold(self) -> float:
|
||||
return 0.0
|
||||
|
||||
def get_choice(self, idx: int) -> str | None:
|
||||
if idx < len(self.choices):
|
||||
return self.choices[idx]
|
||||
return None
|
||||
|
||||
def fill_choices(self) -> None:
|
||||
self.choices.append('extern_fallback_mixed_mm')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=32_BLOCK-N=128_numstages=3_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=64_BLOCK-N=128_numstages=3_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=16_BLOCK-K=128_BLOCK-N=128_numstages=4_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=16_BLOCK-K=128_BLOCK-N=32_numstages=2_numwarps=2')
|
||||
self.choices.append('type=triton_BLOCK-M=16_BLOCK-K=128_BLOCK-N=32_numstages=5_numwarps=2')
|
||||
self.choices.append('type=triton_BLOCK-M=16_BLOCK-K=128_BLOCK-N=64_numstages=5_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=16_BLOCK-K=256_BLOCK-N=128_numstages=3_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=16_BLOCK-K=256_BLOCK-N=128_numstages=5_numwarps=8')
|
||||
self.choices.append('type=triton_BLOCK-M=16_BLOCK-K=64_BLOCK-N=128_numstages=5_numwarps=8')
|
||||
self.choices.append('type=triton_BLOCK-M=16_BLOCK-K=64_BLOCK-N=64_numstages=3_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=32_BLOCK-K=128_BLOCK-N=128_numstages=4_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=32_BLOCK-K=128_BLOCK-N=32_numstages=2_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=32_BLOCK-K=128_BLOCK-N=32_numstages=5_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=64_BLOCK-K=128_BLOCK-N=128_numstages=4_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=64_BLOCK-K=128_BLOCK-N=32_numstages=5_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=64_BLOCK-K=128_BLOCK-N=64_numstages=5_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=64_BLOCK-K=32_BLOCK-N=128_numstages=3_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=64_BLOCK-K=32_BLOCK-N=128_numstages=4_numwarps=8')
|
||||
self.choices.append('type=triton_BLOCK-M=64_BLOCK-K=32_BLOCK-N=64_numstages=3_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=64_BLOCK-K=64_BLOCK-N=128_numstages=3_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=64_BLOCK-K=64_BLOCK-N=128_numstages=5_numwarps=8')
|
||||
|
||||
def get_name(self) -> str:
|
||||
return 'mixed_mm'
|
||||
|
||||
def get_best_choices(self, context: AHContext) -> list[tuple[float, int]] | None:
|
||||
if str(context.get_value('1LEQmLEQ16')) != 'True':
|
||||
if context.get_value('m') <= 32.5:
|
||||
if context.get_value('n') <= 6976.0:
|
||||
if context.get_value('n') <= 3520.0:
|
||||
if context.get_value('m*n') <= 37632.0:
|
||||
return None
|
||||
else:
|
||||
return [(1.000, 13)]
|
||||
else:
|
||||
if context.get_value('m*k') <= 452352.0:
|
||||
return [(0.590, 13), (0.256, 8), (0.103, 7), (0.051, 11)]
|
||||
else:
|
||||
return [(0.778, 8), (0.222, 13)]
|
||||
else:
|
||||
if context.get_value('k*n') <= 102776832.0:
|
||||
if context.get_value('n') <= 14656.0:
|
||||
return [(1.000, 11)]
|
||||
else:
|
||||
return [(0.889, 11), (0.111, 13)]
|
||||
else:
|
||||
return [(1.000, 11)]
|
||||
else:
|
||||
if context.get_value('m*n') <= 446464.0:
|
||||
if context.get_value('m*n') <= 223424.0:
|
||||
if context.get_value('mat1_stride_0') <= 3968.0:
|
||||
return None
|
||||
else:
|
||||
return None
|
||||
else:
|
||||
if context.get_value('m*n') <= 346112.0:
|
||||
return [(0.960, 16), (0.040, 7)]
|
||||
else:
|
||||
return [(0.750, 16), (0.136, 14), (0.114, 7)]
|
||||
else:
|
||||
if str(context.get_value('33LEQmLEQ64')) != 'True':
|
||||
if context.get_value('n') <= 6976.0:
|
||||
return [(1.000, 14)]
|
||||
else:
|
||||
return [(0.753, 2), (0.222, 1), (0.015, 7), (0.007, 16), (0.004, 12)]
|
||||
else:
|
||||
if context.get_value('n') <= 13888.0:
|
||||
return [(0.710, 14), (0.275, 21), (0.014, 12)]
|
||||
else:
|
||||
return [(0.374, 19), (0.339, 20), (0.106, 21), (0.101, 16), (0.066, 17), (0.009, 14), (0.004, 18)]
|
||||
else:
|
||||
if context.get_value('n') <= 3520.0:
|
||||
if context.get_value('arith_intensity') <= 3.994754433631897:
|
||||
if str(context.get_value('mat2_dtype')) != 'torch.uint8':
|
||||
if context.get_value('m*k') <= 18944.0:
|
||||
return [(0.577, 5), (0.423, 6)]
|
||||
else:
|
||||
return [(0.988, 5), (0.012, 6)]
|
||||
else:
|
||||
if context.get_value('arith_intensity') <= 2.9899919033050537:
|
||||
return None
|
||||
else:
|
||||
return None
|
||||
else:
|
||||
if context.get_value('arith_intensity') <= 7.956453561782837:
|
||||
if context.get_value('k*n') <= 9244032.0:
|
||||
return [(0.822, 5), (0.178, 6)]
|
||||
else:
|
||||
return [(0.977, 5), (0.023, 0)]
|
||||
else:
|
||||
if context.get_value('m*k') <= 978944.0:
|
||||
return [(1.000, 5)]
|
||||
else:
|
||||
return [(0.971, 5), (0.029, 0)]
|
||||
else:
|
||||
if context.get_value('n') <= 13632.0:
|
||||
if context.get_value('n') <= 6976.0:
|
||||
return [(1.000, 6)]
|
||||
else:
|
||||
if context.get_value('k') <= 3968.0:
|
||||
return [(0.617, 3), (0.111, 5), (0.099, 7), (0.086, 9), (0.062, 6), (0.025, 8)]
|
||||
else:
|
||||
return [(0.779, 8), (0.119, 5), (0.053, 7), (0.035, 6), (0.013, 3)]
|
||||
else:
|
||||
if context.get_value('k*n') <= 39518208.0:
|
||||
return [(0.385, 4), (0.327, 3), (0.192, 6), (0.038, 7), (0.038, 10), (0.019, 5)]
|
||||
else:
|
||||
if context.get_value('n') <= 20800.0:
|
||||
return [(0.821, 6), (0.121, 7), (0.029, 4), (0.014, 5), (0.007, 3), (0.007, 8)]
|
||||
else:
|
||||
return [(0.530, 7), (0.386, 6), (0.046, 8), (0.021, 3), (0.015, 4), (0.002, 5)]
|
||||
+149
@@ -0,0 +1,149 @@
|
||||
# flake8: noqa: B950
|
||||
# fmt: off
|
||||
# This file was generated by AutoHeuristic. Do not modify it manually!
|
||||
# To regenerate this file, take a look at the steps in the README.md file inside torchgen/_autoheuristic/mixed_mm/
|
||||
from typing import Optional
|
||||
|
||||
from torch._inductor.autoheuristic.autoheuristic_utils import (
|
||||
AHContext,
|
||||
AHMetadata,
|
||||
Choice,
|
||||
)
|
||||
from torch._inductor.autoheuristic.learnedheuristic_interface import (
|
||||
LearnedHeuristicDecision,
|
||||
)
|
||||
|
||||
|
||||
class MixedMMH100(LearnedHeuristicDecision):
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.choices: list[Choice] = []
|
||||
self.fill_choices()
|
||||
|
||||
def check_precondition(self, metadata: AHMetadata, context: AHContext,) -> bool:
|
||||
return (
|
||||
metadata.name == self.get_name()
|
||||
and metadata.shared_memory == 232448
|
||||
and str(metadata.device_capa) == "(9, 0)"
|
||||
)
|
||||
|
||||
def get_confidence_threshold(self) -> float:
|
||||
return 0.0
|
||||
|
||||
def get_choice(self, idx: int) -> str | None:
|
||||
if idx < len(self.choices):
|
||||
return self.choices[idx]
|
||||
return None
|
||||
|
||||
def fill_choices(self) -> None:
|
||||
self.choices.append('extern_fallback_mixed_mm')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=32_BLOCK-N=128_numstages=3_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=32_BLOCK-N=64_numstages=3_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=128_BLOCK-K=64_BLOCK-N=128_numstages=5_numwarps=8')
|
||||
self.choices.append('type=triton_BLOCK-M=16_BLOCK-K=128_BLOCK-N=128_numstages=4_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=16_BLOCK-K=128_BLOCK-N=32_numstages=2_numwarps=2')
|
||||
self.choices.append('type=triton_BLOCK-M=16_BLOCK-K=128_BLOCK-N=32_numstages=5_numwarps=2')
|
||||
self.choices.append('type=triton_BLOCK-M=16_BLOCK-K=128_BLOCK-N=64_numstages=5_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=16_BLOCK-K=256_BLOCK-N=128_numstages=3_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=16_BLOCK-K=256_BLOCK-N=128_numstages=5_numwarps=8')
|
||||
self.choices.append('type=triton_BLOCK-M=16_BLOCK-K=64_BLOCK-N=128_numstages=5_numwarps=8')
|
||||
self.choices.append('type=triton_BLOCK-M=16_BLOCK-K=64_BLOCK-N=64_numstages=3_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=32_BLOCK-K=128_BLOCK-N=128_numstages=4_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=32_BLOCK-K=128_BLOCK-N=32_numstages=2_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=32_BLOCK-K=128_BLOCK-N=32_numstages=5_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=32_BLOCK-K=32_BLOCK-N=64_numstages=5_numwarps=8')
|
||||
self.choices.append('type=triton_BLOCK-M=64_BLOCK-K=128_BLOCK-N=128_numstages=4_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=64_BLOCK-K=128_BLOCK-N=32_numstages=5_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=64_BLOCK-K=128_BLOCK-N=64_numstages=5_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=64_BLOCK-K=64_BLOCK-N=128_numstages=3_numwarps=4')
|
||||
self.choices.append('type=triton_BLOCK-M=64_BLOCK-K=64_BLOCK-N=64_numstages=3_numwarps=8')
|
||||
|
||||
def get_name(self) -> str:
|
||||
return 'mixed_mm'
|
||||
|
||||
def get_best_choices(self, context: AHContext) -> list[tuple[float, int]] | None:
|
||||
if context.get_value('arith_intensity') <= 15.988086223602295:
|
||||
if context.get_value('n') <= 25280.0:
|
||||
if context.get_value('n') <= 1344.0:
|
||||
if context.get_value('mat1_stride_0') <= 7808.0:
|
||||
return [(0.581, 7), (0.419, 6)]
|
||||
else:
|
||||
if context.get_value('m*n') <= 7680.0:
|
||||
return [(0.875, 0), (0.125, 6)]
|
||||
else:
|
||||
return [(0.833, 0), (0.167, 7)]
|
||||
else:
|
||||
if context.get_value('n') <= 8512.0:
|
||||
if str(context.get_value('mat2_dtype')) != 'torch.int8':
|
||||
return [(0.763, 6), (0.237, 7)]
|
||||
else:
|
||||
return [(0.725, 7), (0.275, 6)]
|
||||
else:
|
||||
if str(context.get_value('mat1_dtype')) != 'torch.bfloat16':
|
||||
return [(0.736, 7), (0.197, 9), (0.048, 6), (0.014, 8), (0.005, 10)]
|
||||
else:
|
||||
return [(0.473, 7), (0.398, 6), (0.097, 9), (0.032, 10)]
|
||||
else:
|
||||
if context.get_value('n') <= 42254.0:
|
||||
if context.get_value('n') <= 33856.0:
|
||||
if context.get_value('k*n') <= 68157440.0:
|
||||
return [(0.370, 4), (0.370, 5), (0.074, 7), (0.074, 8), (0.074, 11), (0.037, 6)]
|
||||
else:
|
||||
return [(0.916, 8), (0.036, 7), (0.036, 9), (0.012, 4)]
|
||||
else:
|
||||
return [(0.659, 5), (0.341, 6)]
|
||||
else:
|
||||
if context.get_value('k*n') <= 326052992.0:
|
||||
if context.get_value('n') <= 55232.0:
|
||||
return [(0.571, 6), (0.321, 7), (0.036, 4), (0.036, 8), (0.036, 9)]
|
||||
else:
|
||||
return [(0.506, 6), (0.325, 8), (0.104, 7), (0.039, 5), (0.026, 9)]
|
||||
else:
|
||||
if context.get_value('n') <= 57024.0:
|
||||
return [(0.462, 9), (0.385, 7), (0.115, 6), (0.038, 8)]
|
||||
else:
|
||||
return [(0.598, 8), (0.223, 9), (0.107, 6), (0.071, 7)]
|
||||
else:
|
||||
if context.get_value('m*n') <= 543936.0:
|
||||
if str(context.get_value('17LEQmLEQ32')) != 'True':
|
||||
if context.get_value('m*n') <= 262272.0:
|
||||
if context.get_value('n') <= 1592.5:
|
||||
return [(0.860, 0), (0.140, 9)]
|
||||
else:
|
||||
return None
|
||||
else:
|
||||
if context.get_value('m*k') <= 1294336.0:
|
||||
return [(0.833, 17), (0.150, 18), (0.017, 15)]
|
||||
else:
|
||||
return [(0.917, 17), (0.083, 8)]
|
||||
else:
|
||||
if context.get_value('n') <= 12416.0:
|
||||
if context.get_value('m*n') <= 43008.0:
|
||||
return None
|
||||
else:
|
||||
return [(0.853, 14), (0.147, 9)]
|
||||
else:
|
||||
return [(0.625, 12), (0.375, 14)]
|
||||
else:
|
||||
if context.get_value('m') <= 32.5:
|
||||
if context.get_value('mat2_stride_1') <= 6656.0:
|
||||
if context.get_value('n') <= 69184.0:
|
||||
return [(0.611, 12), (0.361, 14), (0.028, 13)]
|
||||
else:
|
||||
return [(1.000, 12)]
|
||||
else:
|
||||
if context.get_value('mat2_stride_1') <= 20864.0:
|
||||
return [(1.000, 12)]
|
||||
else:
|
||||
return [(0.958, 12), (0.042, 9)]
|
||||
else:
|
||||
if context.get_value('m*n') <= 1085440.0:
|
||||
if context.get_value('n') <= 9152.0:
|
||||
return [(1.000, 18)]
|
||||
else:
|
||||
return [(0.780, 18), (0.160, 16), (0.060, 20)]
|
||||
else:
|
||||
if context.get_value('m') <= 67.0:
|
||||
return [(0.650, 16), (0.203, 19), (0.122, 18), (0.016, 20), (0.008, 1)]
|
||||
else:
|
||||
return [(0.561, 3), (0.185, 16), (0.096, 20), (0.083, 19), (0.076, 2)]
|
||||
+265
@@ -0,0 +1,265 @@
|
||||
# flake8: noqa: B950
|
||||
# fmt: off
|
||||
# This file was generated by AutoHeuristic. Do not modify it manually!
|
||||
# To regenerate this file, take a look at the steps in the README.md file inside torchgen/_autoheuristic/pad_mm/
|
||||
from typing import Optional
|
||||
|
||||
from torch._inductor.autoheuristic.autoheuristic_utils import (
|
||||
AHContext,
|
||||
AHMetadata,
|
||||
Choice,
|
||||
)
|
||||
from torch._inductor.autoheuristic.learnedheuristic_interface import (
|
||||
LearnedHeuristicDecision,
|
||||
)
|
||||
|
||||
|
||||
class PadMMA100(LearnedHeuristicDecision):
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.choices: list[Choice] = []
|
||||
self.fill_choices()
|
||||
|
||||
def check_precondition(self, metadata: AHMetadata, context: AHContext,) -> bool:
|
||||
return (
|
||||
metadata.name == self.get_name()
|
||||
and metadata.shared_memory == 166912
|
||||
and str(metadata.device_capa) == "(8, 0)"
|
||||
)
|
||||
|
||||
def get_confidence_threshold(self) -> float:
|
||||
return 0.9294871794871795
|
||||
|
||||
def get_choice(self, idx: int) -> Optional[str]:
|
||||
if idx < len(self.choices):
|
||||
return self.choices[idx]
|
||||
return None
|
||||
|
||||
def fill_choices(self) -> None:
|
||||
self.choices.append('orig')
|
||||
self.choices.append('pad')
|
||||
|
||||
def get_name(self) -> str:
|
||||
return 'pad_mm'
|
||||
|
||||
def get_best_choices(self, context: AHContext) -> Optional[list[tuple[float, int]]]:
|
||||
if str(context.get_value('mat1_innermost_needs_padding')) != 'False':
|
||||
if context.get_value('arith_intensity') <= 880.0238037109375:
|
||||
if str(context.get_value('m_multiple_2')) != 'True':
|
||||
if context.get_value('n') <= 652.0:
|
||||
if context.get_value('m') <= 2022.0:
|
||||
if str(context.get_value('using_tf32')) != 'False':
|
||||
return [(0.579, 1), (0.421, 0)]
|
||||
else:
|
||||
return [(1.000, 0)]
|
||||
else:
|
||||
return [(1.000, 0)]
|
||||
else:
|
||||
if context.get_value('m*k') <= 107278336.0:
|
||||
if str(context.get_value('using_tf32')) != 'False':
|
||||
if context.get_value('m') <= 23691.0:
|
||||
return [(0.993, 1), (0.007, 0)]
|
||||
else:
|
||||
return [(0.840, 1), (0.160, 0)]
|
||||
else:
|
||||
if context.get_value('arith_intensity') <= 793.3185424804688:
|
||||
return [(0.958, 0), (0.042, 1)]
|
||||
else:
|
||||
return [(0.792, 1), (0.208, 0)]
|
||||
else:
|
||||
if context.get_value('arith_intensity') <= 795.6242370605469:
|
||||
if context.get_value('mat2_stride_1') <= 2048.5:
|
||||
return [(0.929, 0), (0.071, 1)]
|
||||
else:
|
||||
return [(1.000, 1)]
|
||||
else:
|
||||
if context.get_value('arith_intensity') <= 796.3460388183594:
|
||||
return [(0.957, 1), (0.043, 0)]
|
||||
else:
|
||||
return [(0.778, 0), (0.222, 1)]
|
||||
else:
|
||||
if context.get_value('mat2_stride_0') <= 2432.0:
|
||||
if str(context.get_value('k_multiple_2')) != 'False':
|
||||
if context.get_value('n') <= 1024.5:
|
||||
if str(context.get_value('prepadded_mat1')) != 'False':
|
||||
return [(0.580, 0), (0.420, 1)]
|
||||
else:
|
||||
return [(0.986, 0), (0.014, 1)]
|
||||
else:
|
||||
if context.get_value('mat1_stride_0') <= 5125.0:
|
||||
return [(0.551, 0), (0.449, 1)]
|
||||
else:
|
||||
return [(0.916, 0), (0.084, 1)]
|
||||
else:
|
||||
if context.get_value('mat2_align_size') <= 6.0:
|
||||
if str(context.get_value('using_tf32')) != 'True':
|
||||
return [(1.000, 0)]
|
||||
else:
|
||||
return [(0.800, 0), (0.200, 1)]
|
||||
else:
|
||||
if context.get_value('mat2_stride_1') <= 3820.0:
|
||||
return [(0.986, 1), (0.014, 0)]
|
||||
else:
|
||||
return [(0.532, 0), (0.468, 1)]
|
||||
else:
|
||||
if str(context.get_value('using_tf32')) != 'False':
|
||||
if context.get_value('m*n') <= 5244928.0:
|
||||
if str(context.get_value('k_multiple_2')) != 'True':
|
||||
return [(0.971, 1), (0.029, 0)]
|
||||
else:
|
||||
return [(0.646, 1), (0.354, 0)]
|
||||
else:
|
||||
if context.get_value('k/(m*n)') <= 9.468618827668251e-06:
|
||||
return [(0.800, 1), (0.200, 0)]
|
||||
else:
|
||||
return [(1.000, 1)]
|
||||
else:
|
||||
if context.get_value('mat1_stride_1') <= 1288.0:
|
||||
if context.get_value('k') <= 5717.0:
|
||||
return [(0.983, 0), (0.017, 1)]
|
||||
else:
|
||||
return [(0.800, 0), (0.200, 1)]
|
||||
else:
|
||||
return [(0.588, 0), (0.412, 1)]
|
||||
else:
|
||||
if str(context.get_value('using_tf32')) != 'False':
|
||||
if context.get_value('n') <= 2640.0:
|
||||
if context.get_value('m_padded_length') <= 1.5:
|
||||
if context.get_value('mat2_stride_1') <= 6021.0:
|
||||
return [(1.000, 1)]
|
||||
else:
|
||||
if context.get_value('mat2_stride_1') <= 6069.0:
|
||||
return [(0.900, 1), (0.100, 0)]
|
||||
else:
|
||||
return [(1.000, 1)]
|
||||
else:
|
||||
if str(context.get_value('m_multiple_2')) != 'False':
|
||||
if context.get_value('m*k') <= 24444928.0:
|
||||
return [(0.593, 0), (0.407, 1)]
|
||||
else:
|
||||
return [(0.923, 0), (0.077, 1)]
|
||||
else:
|
||||
return [(1.000, 1)]
|
||||
else:
|
||||
if context.get_value('m*k') <= 404182016.0:
|
||||
if str(context.get_value('mat2_innermost_needs_padding')) != 'False':
|
||||
if context.get_value('m*k') <= 12328960.0:
|
||||
return [(0.732, 1), (0.268, 0)]
|
||||
else:
|
||||
return [(0.989, 1), (0.011, 0)]
|
||||
else:
|
||||
if context.get_value('m*k') <= 389028864.0:
|
||||
return [(0.998, 1), (0.002, 0)]
|
||||
else:
|
||||
return [(0.922, 1), (0.078, 0)]
|
||||
else:
|
||||
if context.get_value('m*n') <= 137631744.0:
|
||||
if context.get_value('m*k') <= 405715968.0:
|
||||
return [(0.611, 1), (0.389, 0)]
|
||||
else:
|
||||
return [(0.946, 1), (0.054, 0)]
|
||||
else:
|
||||
return [(0.714, 0), (0.286, 1)]
|
||||
else:
|
||||
if context.get_value('mat2_stride_0') <= 3902.5:
|
||||
return [(0.941, 0), (0.059, 1)]
|
||||
else:
|
||||
return [(0.583, 1), (0.417, 0)]
|
||||
else:
|
||||
if context.get_value('n_padded_length') <= 0.5:
|
||||
if str(context.get_value('mat2_innermost_needs_padding')) != 'False':
|
||||
if str(context.get_value('k_multiple_2')) != 'True':
|
||||
if context.get_value('arith_intensity') <= 884.5185852050781:
|
||||
if context.get_value('arith_intensity') <= 743.931884765625:
|
||||
return [(1.000, 0)]
|
||||
else:
|
||||
return [(0.583, 0), (0.417, 1)]
|
||||
else:
|
||||
return [(1.000, 1)]
|
||||
else:
|
||||
if context.get_value('k/(m*n)') <= 0.00023481580865336582:
|
||||
return [(0.900, 0), (0.100, 1)]
|
||||
else:
|
||||
return [(1.000, 0)]
|
||||
else:
|
||||
if str(context.get_value('using_tf32')) != 'False':
|
||||
if context.get_value('m*k') <= 93734912.0:
|
||||
if context.get_value('mat1_stride_0') <= 1344.0:
|
||||
if context.get_value('n') <= 7168.0:
|
||||
return [(1.000, 0)]
|
||||
else:
|
||||
return [(0.970, 0), (0.030, 1)]
|
||||
else:
|
||||
if context.get_value('m') <= 22883.5:
|
||||
return [(0.977, 0), (0.023, 1)]
|
||||
else:
|
||||
return [(0.800, 0), (0.200, 1)]
|
||||
else:
|
||||
if context.get_value('arith_intensity') <= 1914.3681030273438:
|
||||
if str(context.get_value('prepadded_mat1')) != 'False':
|
||||
return [(0.981, 0), (0.019, 1)]
|
||||
else:
|
||||
return [(0.995, 0), (0.005, 1)]
|
||||
else:
|
||||
return [(1.000, 0)]
|
||||
else:
|
||||
if str(context.get_value('prepadded_mat1')) != 'False':
|
||||
if context.get_value('mat2_stride_1') <= 256.5:
|
||||
if context.get_value('m') <= 5880.5:
|
||||
return [(1.000, 0)]
|
||||
else:
|
||||
return [(0.800, 0), (0.200, 1)]
|
||||
else:
|
||||
if context.get_value('m*k') <= 6318080.0:
|
||||
return [(0.618, 1), (0.382, 0)]
|
||||
else:
|
||||
return [(0.880, 0), (0.120, 1)]
|
||||
else:
|
||||
if context.get_value('k/(m*n)') <= 0.0009747986623551697:
|
||||
if context.get_value('k/(m*n)') <= 0.0006397514371201396:
|
||||
return [(0.951, 0), (0.049, 1)]
|
||||
else:
|
||||
return [(0.857, 0), (0.143, 1)]
|
||||
else:
|
||||
return [(1.000, 0)]
|
||||
else:
|
||||
if str(context.get_value('using_tf32')) != 'False':
|
||||
if str(context.get_value('n_multiple_2')) != 'False':
|
||||
if context.get_value('m') <= 2024.0:
|
||||
if context.get_value('mat2_stride_0') <= 1629.0:
|
||||
if context.get_value('k*n') <= 1288704.0:
|
||||
return [(0.600, 0), (0.400, 1)]
|
||||
else:
|
||||
return [(0.982, 0), (0.018, 1)]
|
||||
else:
|
||||
if context.get_value('mat1_stride_0') <= 768.0:
|
||||
return [(0.619, 0), (0.381, 1)]
|
||||
else:
|
||||
return [(0.812, 1), (0.188, 0)]
|
||||
else:
|
||||
if context.get_value('m*n') <= 5803008.0:
|
||||
return [(0.500, 0), (0.500, 1)]
|
||||
else:
|
||||
if context.get_value('mat2_stride_1') <= 896.0:
|
||||
return [(1.000, 1)]
|
||||
else:
|
||||
return [(0.818, 1), (0.182, 0)]
|
||||
else:
|
||||
if context.get_value('mat2_stride_1') <= 2560.0:
|
||||
if context.get_value('num_dims_needs_padding') <= 1.5:
|
||||
if context.get_value('mat2_stride_1') <= 896.0:
|
||||
return [(1.000, 1)]
|
||||
else:
|
||||
return [(0.857, 1), (0.143, 0)]
|
||||
else:
|
||||
return [(0.727, 1), (0.273, 0)]
|
||||
else:
|
||||
return [(0.667, 1), (0.333, 0)]
|
||||
else:
|
||||
if context.get_value('k/(m*n)') <= 0.00015462777810171247:
|
||||
return [(0.857, 0), (0.143, 1)]
|
||||
else:
|
||||
if context.get_value('k/(m*n)') <= 0.0019917909521609545:
|
||||
return [(1.000, 0)]
|
||||
else:
|
||||
return [(0.900, 0), (0.100, 1)]
|
||||
+175
@@ -0,0 +1,175 @@
|
||||
# flake8: noqa: B950
|
||||
# fmt: off
|
||||
# This file was generated by AutoHeuristic. Do not modify it manually!
|
||||
# To regenerate this file, take a look at the steps in the README.md file inside torchgen/_autoheuristic/pad_mm/
|
||||
from typing import Optional
|
||||
|
||||
from torch._inductor.autoheuristic.autoheuristic_utils import (
|
||||
AHContext,
|
||||
AHMetadata,
|
||||
Choice,
|
||||
)
|
||||
from torch._inductor.autoheuristic.learnedheuristic_interface import (
|
||||
LearnedHeuristicDecision,
|
||||
)
|
||||
|
||||
|
||||
class PadMMH200(LearnedHeuristicDecision):
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.choices: list[Choice] = []
|
||||
self.fill_choices()
|
||||
|
||||
def check_precondition(self, metadata: AHMetadata, context: AHContext,) -> bool:
|
||||
return (
|
||||
metadata.name == self.get_name()
|
||||
and metadata.shared_memory == 232448
|
||||
and str(metadata.device_capa) == "(9, 0)"
|
||||
)
|
||||
|
||||
def get_confidence_threshold(self) -> float:
|
||||
return 0.7710651828298887
|
||||
|
||||
def get_choice(self, idx: int) -> Optional[str]:
|
||||
if idx < len(self.choices):
|
||||
return self.choices[idx]
|
||||
return None
|
||||
|
||||
def fill_choices(self) -> None:
|
||||
self.choices.append('orig')
|
||||
self.choices.append('pad')
|
||||
|
||||
def get_name(self) -> str:
|
||||
return 'pad_mm'
|
||||
|
||||
def get_best_choices(self, context: AHContext) -> Optional[list[tuple[float, int]]]:
|
||||
if str(context.get_value('mat1_innermost_needs_padding')) != 'True':
|
||||
if str(context.get_value('mat2_innermost_needs_padding')) != 'True':
|
||||
if context.get_value('n_padded_length') <= 0.5:
|
||||
if str(context.get_value('prepadded_mat1')) != 'True':
|
||||
if str(context.get_value('using_tf32')) != 'False':
|
||||
return [(1.000, 0)]
|
||||
else:
|
||||
if context.get_value('mat1_stride_0') <= 3584.0:
|
||||
return [(1.000, 0)]
|
||||
else:
|
||||
if context.get_value('mat2_stride_0') <= 3584.0:
|
||||
return [(1.000, 0)]
|
||||
else:
|
||||
return [(0.528, 0), (0.472, 1)]
|
||||
else:
|
||||
if context.get_value('n') <= 2304.0:
|
||||
if context.get_value('m*k') <= 25198592.0:
|
||||
if context.get_value('arith_intensity') <= 1103.9319458007812:
|
||||
return [(1.000, 0)]
|
||||
else:
|
||||
return [(0.885, 0), (0.115, 1)]
|
||||
else:
|
||||
if context.get_value('m*k') <= 25688064.0:
|
||||
return [(1.000, 1)]
|
||||
else:
|
||||
return [(0.771, 0), (0.229, 1)]
|
||||
else:
|
||||
if str(context.get_value('using_tf32')) != 'False':
|
||||
if context.get_value('m') <= 27825.0:
|
||||
return [(0.948, 0), (0.052, 1)]
|
||||
else:
|
||||
return [(0.855, 0), (0.145, 1)]
|
||||
else:
|
||||
if context.get_value('mat2_stride_0') <= 3584.0:
|
||||
return [(1.000, 0)]
|
||||
else:
|
||||
return [(0.917, 1), (0.083, 0)]
|
||||
else:
|
||||
if context.get_value('m') <= 1823.5:
|
||||
if str(context.get_value('n_multiple_2')) != 'False':
|
||||
if context.get_value('k*n') <= 7859200.0:
|
||||
return [(0.600, 0), (0.400, 1)]
|
||||
else:
|
||||
return [(1.000, 0)]
|
||||
else:
|
||||
if context.get_value('k/(m*n)') <= 0.00040277576772496104:
|
||||
return [(1.000, 1)]
|
||||
else:
|
||||
return [(0.800, 1), (0.200, 0)]
|
||||
else:
|
||||
if context.get_value('n') <= 3602.0:
|
||||
return [(0.800, 1), (0.200, 0)]
|
||||
else:
|
||||
return [(1.000, 1)]
|
||||
else:
|
||||
if str(context.get_value('using_tf32')) != 'False':
|
||||
if str(context.get_value('n_multiple_16')) != 'False':
|
||||
if str(context.get_value('k_multiple_2')) != 'True':
|
||||
if context.get_value('arith_intensity') <= 744.8332214355469:
|
||||
return [(0.600, 0), (0.400, 1)]
|
||||
else:
|
||||
return [(1.000, 1)]
|
||||
else:
|
||||
if context.get_value('m*n') <= 8912896.0:
|
||||
if context.get_value('m*k') <= 5934080.0:
|
||||
return [(0.800, 0), (0.200, 1)]
|
||||
else:
|
||||
return [(1.000, 0)]
|
||||
else:
|
||||
return [(1.000, 1)]
|
||||
else:
|
||||
return [(1.000, 1)]
|
||||
else:
|
||||
return [(1.000, 0)]
|
||||
else:
|
||||
if context.get_value('arith_intensity') <= 895.8767395019531:
|
||||
if str(context.get_value('m_multiple_2')) != 'False':
|
||||
if context.get_value('mat1_stride_1') <= 3421.0:
|
||||
if str(context.get_value('using_tf32')) != 'False':
|
||||
if context.get_value('mat2_stride_1') <= 10706.5:
|
||||
if context.get_value('mat2_stride_0') <= 1024.5:
|
||||
return [(0.816, 1), (0.184, 0)]
|
||||
else:
|
||||
return [(1.000, 1)]
|
||||
else:
|
||||
if str(context.get_value('k_multiple_2')) != 'True':
|
||||
return [(0.905, 1), (0.095, 0)]
|
||||
else:
|
||||
return [(1.000, 0)]
|
||||
else:
|
||||
if str(context.get_value('prepadded_mat2')) != 'True':
|
||||
if str(context.get_value('mat2_innermost_needs_padding')) != 'False':
|
||||
return [(1.000, 0)]
|
||||
else:
|
||||
return [(0.932, 0), (0.068, 1)]
|
||||
else:
|
||||
if context.get_value('arith_intensity') <= 742.1241760253906:
|
||||
return [(0.889, 0), (0.111, 1)]
|
||||
else:
|
||||
return [(0.765, 1), (0.235, 0)]
|
||||
else:
|
||||
if context.get_value('n') <= 1216.0:
|
||||
if str(context.get_value('using_tf32')) != 'True':
|
||||
if context.get_value('mat1_stride_1') <= 5567.0:
|
||||
return [(0.896, 0), (0.104, 1)]
|
||||
else:
|
||||
return [(0.999, 0), (0.001, 1)]
|
||||
else:
|
||||
return [(1.000, 1)]
|
||||
else:
|
||||
if str(context.get_value('using_tf32')) != 'False':
|
||||
return [(1.000, 1)]
|
||||
else:
|
||||
return [(1.000, 0)]
|
||||
else:
|
||||
if str(context.get_value('using_tf32')) != 'False':
|
||||
return [(1.000, 1)]
|
||||
else:
|
||||
if context.get_value('mat2_stride_1') <= 2688.0:
|
||||
return [(1.000, 0)]
|
||||
else:
|
||||
return [(0.500, 0), (0.500, 1)]
|
||||
else:
|
||||
if str(context.get_value('using_tf32')) != 'False':
|
||||
return [(1.000, 1)]
|
||||
else:
|
||||
if str(context.get_value('mat2_innermost_needs_padding')) != 'True':
|
||||
return [(1.000, 0)]
|
||||
else:
|
||||
return [(0.800, 0), (0.200, 1)]
|
||||
+318
@@ -0,0 +1,318 @@
|
||||
import json
|
||||
import os
|
||||
from collections.abc import Callable
|
||||
from functools import partial
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
from torch._inductor.autoheuristic.autoheuristic_utils import (
|
||||
AHContext,
|
||||
AHMetadata,
|
||||
AHOperation,
|
||||
Choice,
|
||||
CHOICE_COL,
|
||||
Feedback,
|
||||
FEEDBACK_COL,
|
||||
get_metadata_str_from_log,
|
||||
)
|
||||
from torch._inductor.autoheuristic.learned_heuristic_controller import (
|
||||
LearnedHeuristicController,
|
||||
)
|
||||
from torch._inductor.ir import ChoiceCaller
|
||||
from torch._inductor.runtime.runtime_utils import cache_dir
|
||||
from torch._inductor.utils import get_gpu_shared_memory
|
||||
|
||||
|
||||
class LocalFeedback:
|
||||
"""
|
||||
To be able to collect data for a choice, a function providing feedback given a choice has to be provided.
|
||||
LocalFeedback can be used when AutoHeuristic should immediately run the function to collect feedback for each choice
|
||||
(see pad_mm.py, where the autotuning happens locally, for an example).
|
||||
"""
|
||||
|
||||
def __init__(self, feedback_fn: Callable[[Choice], Feedback]) -> None:
|
||||
self.feedback_fn = feedback_fn
|
||||
|
||||
def __call__(self, choice: Choice) -> Feedback:
|
||||
return self.feedback_fn(choice)
|
||||
|
||||
|
||||
class InconsistentMetadata(Exception):
|
||||
"""
|
||||
Exception that is thrown when AutoHeuristic tries to log data to a file where the metadata stored in the file does
|
||||
not match the metadata it would store if the file didn't exist.
|
||||
"""
|
||||
|
||||
|
||||
class AutoHeuristic:
|
||||
"""
|
||||
AutoHeuristic is a framework that allows one to collect data, learn a heuristic (i.e. a regression tree) and
|
||||
generate the heuristic to code. This class allows one to collect data. The collected data can then be used to train
|
||||
a heuristic (see torchgen/autoheuristic/).
|
||||
"""
|
||||
|
||||
collected_feedback: dict[Choice, Feedback]
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
fallback: Callable[[], Choice],
|
||||
choices: list[Choice],
|
||||
feedback: LocalFeedback | None,
|
||||
context: AHContext,
|
||||
name: str,
|
||||
augment_context: list[AHOperation] | None = None,
|
||||
precondition: Callable[[AHMetadata, AHContext], bool] | None = None,
|
||||
) -> None:
|
||||
"""
|
||||
Initializes an instance of the AutoHeuristic class.
|
||||
|
||||
Args:
|
||||
fallback: A callable that returns a Choice when the heuristic is unsure which choice to make, or
|
||||
AutoHeuristic is in data collection mode.
|
||||
choices: A list of possible choices the heuristic can make.
|
||||
feedback: An instance of LocalFeedback that provides feedback for a given choice.
|
||||
context: Context to store with each choice and feedback.
|
||||
name: A string that identifies the heuristic.
|
||||
augment_context: An optional list of AHOperation instances that augment the context.
|
||||
precondition: A callable that returns a boolean indicating whether AutoHeuristic should run.
|
||||
"""
|
||||
self.fallback = fallback
|
||||
self.choices = choices
|
||||
self.feedback = feedback
|
||||
self.context = context
|
||||
self.name = name
|
||||
self.collected_feedback = {}
|
||||
self.augment_context = augment_context
|
||||
self.metadata = AHMetadata(
|
||||
get_gpu_shared_memory(),
|
||||
torch.cuda.get_device_capability(),
|
||||
self.choices,
|
||||
self.name,
|
||||
)
|
||||
self.precondition = precondition
|
||||
|
||||
if not self.satisfies_precondition():
|
||||
return
|
||||
|
||||
if torch._inductor.config.autoheuristic_log_path == "DEFAULT":
|
||||
self.log_path = self.get_default_log_path()
|
||||
else:
|
||||
self.log_path = torch._inductor.config.autoheuristic_log_path
|
||||
|
||||
if torch._inductor.config.collect_autoheuristic(self.name):
|
||||
if self.feedback is not None:
|
||||
for choice in self.choices:
|
||||
feedback_val = self.feedback(choice)
|
||||
self.save_data(choice, feedback_val)
|
||||
|
||||
def satisfies_precondition(self) -> bool:
|
||||
return self.precondition is None or self.precondition(
|
||||
self.metadata, self.context
|
||||
)
|
||||
|
||||
def get_choice(self) -> Choice:
|
||||
"""
|
||||
Returns the chosen option based on the value of autoheuristic_use.
|
||||
If self.name is one of the comma separated strings in autoheuristic_use,
|
||||
it queries a learned heuristic to make a decision. Otherwise, it returns the fallback option.
|
||||
"""
|
||||
|
||||
if not self.satisfies_precondition():
|
||||
return self.fallback()
|
||||
|
||||
if torch._inductor.config.use_autoheuristic(self.name):
|
||||
if self.augment_context is not None:
|
||||
self.context.apply_operations(self.augment_context)
|
||||
controller = LearnedHeuristicController(
|
||||
self.metadata,
|
||||
self.context,
|
||||
)
|
||||
decision = controller.get_decision()
|
||||
if decision not in self.choices:
|
||||
# TODO(AlnisM): We might want to allow this in the future
|
||||
return self.fallback()
|
||||
if decision is not None:
|
||||
return decision
|
||||
return self.fallback()
|
||||
|
||||
def get_top_k_choices(
|
||||
self, top_k: int, always_included: list[str] | None = None
|
||||
) -> list[Choice] | None:
|
||||
if not self.satisfies_precondition():
|
||||
return None
|
||||
if torch._inductor.config.use_autoheuristic(self.name):
|
||||
if self.augment_context is not None:
|
||||
self.context.apply_operations(self.augment_context)
|
||||
controller = LearnedHeuristicController(
|
||||
self.metadata,
|
||||
self.context,
|
||||
)
|
||||
choices = controller.get_decisions_ranked(top_k)
|
||||
if choices is None:
|
||||
return None
|
||||
if always_included is not None:
|
||||
for choice in always_included:
|
||||
if choice not in choices:
|
||||
choices.append(choice)
|
||||
return choices
|
||||
return None
|
||||
|
||||
def get_collected_feedback(self, choice: Choice) -> Any:
|
||||
return self.collected_feedback.get(choice, None)
|
||||
|
||||
@staticmethod
|
||||
def get_device_identifier() -> str:
|
||||
# a heuristic might work well for one GPU, but not for another
|
||||
# we store the collected data per GPU model and learn a heuristic per GPU model
|
||||
|
||||
# TODO(AlnisM): just using the device name for now, but the same GPU model can have different names
|
||||
device_name = torch.cuda.get_device_name().replace(" ", "_")
|
||||
return device_name
|
||||
|
||||
def get_default_log_path(self) -> str:
|
||||
device_name = self.get_device_identifier()
|
||||
path = f"{cache_dir()}/autoheuristic/{device_name}/"
|
||||
os.makedirs(path, exist_ok=True)
|
||||
path += f"{self.name}.txt"
|
||||
return path
|
||||
|
||||
def serialize_metadata(self) -> str:
|
||||
metadata_dict = self.metadata.to_dict()
|
||||
(
|
||||
num_features,
|
||||
cat_features,
|
||||
) = self.context.get_numerical_and_categorical_features()
|
||||
metadata_dict["numerical_features"] = num_features
|
||||
metadata_dict["categorical_features"] = cat_features
|
||||
return json.dumps(metadata_dict)
|
||||
|
||||
def save_data(self, choice: Choice, feedback_val: Feedback) -> None:
|
||||
self.collected_feedback[choice] = feedback_val
|
||||
log_path = self.log_path
|
||||
|
||||
lines = []
|
||||
log_exists = os.path.exists(log_path)
|
||||
if log_exists:
|
||||
# if log already exists, make sure it is consistent
|
||||
metadata = self.serialize_metadata()
|
||||
existing_metadata = get_metadata_str_from_log(self.log_path)
|
||||
if existing_metadata != metadata:
|
||||
raise InconsistentMetadata(
|
||||
"Given metadata does not match existing metadata"
|
||||
)
|
||||
else:
|
||||
lines.append(self.serialize_metadata())
|
||||
feature_header = self.context.get_feature_names_csv()
|
||||
header = feature_header + "," + CHOICE_COL + "," + FEEDBACK_COL
|
||||
lines.append(header)
|
||||
|
||||
line = ""
|
||||
feature_values = self.context.get_feature_values_csv()
|
||||
line += feature_values + "," + choice + "," + str(feedback_val)
|
||||
lines.append(line)
|
||||
|
||||
with open(log_path, "a") as f:
|
||||
f.write("\n".join(lines) + "\n")
|
||||
|
||||
|
||||
class AutoHeuristicSelectAlgorithm(AutoHeuristic):
|
||||
"""
|
||||
AutoHeuristicSelectAlgorithm is a subclass of AutoHeuristic that allows one to collect data and learn a heuristic
|
||||
when one wants to use AutoHeuristic for kernel choice selection.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
fallback: Callable[[], ChoiceCaller | None],
|
||||
choices: list[ChoiceCaller],
|
||||
input_nodes: list[Any],
|
||||
context: AHContext,
|
||||
name: str,
|
||||
augment_context: list[AHOperation] | None = None,
|
||||
precondition: Callable[[AHMetadata, AHContext], bool] | None = None,
|
||||
) -> None:
|
||||
"""
|
||||
The arguments choices, input_nodes and name have to match the ones used in the call to
|
||||
autotune_select_algorithm(), e.g. if the following call is made
|
||||
autotune_select_algorithm(name, choices, input_nodes, layout), the same name, choices and input_nodes
|
||||
have to be used here.
|
||||
"""
|
||||
self.input_nodes = input_nodes
|
||||
self.choicestr2choice: dict[str, ChoiceCaller] = {}
|
||||
for choice in choices:
|
||||
self.choicestr2choice[choice.autoheuristic_id()] = choice
|
||||
choices_str = list(self.choicestr2choice.keys())
|
||||
|
||||
def fallback_str() -> str:
|
||||
fallback_choice = fallback()
|
||||
if fallback_choice is None:
|
||||
# TODO: Find a nicer way to handle this
|
||||
return "unsure"
|
||||
return fallback_choice.autoheuristic_id()
|
||||
|
||||
super().__init__(
|
||||
fallback_str,
|
||||
choices_str,
|
||||
None,
|
||||
context,
|
||||
name,
|
||||
augment_context,
|
||||
precondition,
|
||||
)
|
||||
|
||||
if (
|
||||
torch._inductor.config.collect_autoheuristic(self.name)
|
||||
and self.satisfies_precondition()
|
||||
):
|
||||
self.register_global_feedback(input_nodes, choices)
|
||||
|
||||
def register_global_feedback(
|
||||
self, input_nodes: list[Any], choices: list[ChoiceCaller]
|
||||
) -> None:
|
||||
"""
|
||||
Registers a callback in select_algorithm, which is called with the timing of each choice.
|
||||
"""
|
||||
|
||||
from torch._inductor.select_algorithm import (
|
||||
add_feedback_saver,
|
||||
create_inputs_key,
|
||||
create_precompile_key,
|
||||
)
|
||||
|
||||
def store_global_feedback(
|
||||
ah_inputs_key: str,
|
||||
ah_precompile_key: str,
|
||||
timings: dict[ChoiceCaller, float],
|
||||
name: str,
|
||||
input_nodes: list[Any],
|
||||
choices: list[ChoiceCaller],
|
||||
profiled_time: Callable[[], dict[ChoiceCaller, float]],
|
||||
precompile_times: dict[ChoiceCaller, float],
|
||||
) -> None:
|
||||
current_inputs_key = create_inputs_key(input_nodes)
|
||||
if current_inputs_key != ah_inputs_key:
|
||||
return
|
||||
current_precompile_key = create_precompile_key(
|
||||
name, current_inputs_key, choices
|
||||
)
|
||||
if current_precompile_key != ah_precompile_key:
|
||||
return
|
||||
for choice, time in timings.items():
|
||||
self.save_data(choice.autoheuristic_id(), time)
|
||||
|
||||
inputs_key = create_inputs_key(input_nodes)
|
||||
precompile_key = create_precompile_key(self.name, inputs_key, choices)
|
||||
feedback_saver = partial(store_global_feedback, inputs_key, precompile_key)
|
||||
add_feedback_saver(feedback_saver)
|
||||
|
||||
def get_choice_caller(self) -> ChoiceCaller | None:
|
||||
choice = self.get_choice()
|
||||
return self.choicestr2choice.get(choice, None)
|
||||
|
||||
def get_top_k_choices_caller(
|
||||
self, top_k: int, always_included: list[str] | None = None
|
||||
) -> list[ChoiceCaller] | None:
|
||||
choices = self.get_top_k_choices(top_k, always_included)
|
||||
if choices is None:
|
||||
return None
|
||||
return [self.choicestr2choice[choice] for choice in choices]
|
||||
+340
@@ -0,0 +1,340 @@
|
||||
import functools
|
||||
from collections.abc import Callable
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
|
||||
|
||||
Feedback = float
|
||||
Choice = str
|
||||
Value = Any
|
||||
|
||||
CHOICE_COL = "choice"
|
||||
FEEDBACK_COL = "feedback"
|
||||
|
||||
|
||||
class AHFeature:
|
||||
"""
|
||||
The context, that AutoHeuristic stores, is a list of features. AutoHeuristic needs to know whether a feature is
|
||||
categorical (i.e., not a continuous variable) to learn a machine learning model.
|
||||
"""
|
||||
|
||||
def __init__(self, name: str, value: Value, is_categorical: bool = False) -> None:
|
||||
self.name = name
|
||||
self.value = value
|
||||
self.is_categorical = is_categorical
|
||||
|
||||
|
||||
class AHOperation:
|
||||
"""
|
||||
AHOperation can be used to augment the data collected by AutoHeuristic.
|
||||
One might for example store features like m, k, n, but also want to use
|
||||
features like m*n, or k*n, to learn a heuristic. Instead of storing features
|
||||
that can be created from the collected data, one can use AHOperation to
|
||||
create new features from the collected data.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self, name: str, func: Callable[[Any], Value], is_categorical: bool = False
|
||||
) -> None:
|
||||
self.name = name
|
||||
self.func = func
|
||||
self.is_categorical = is_categorical
|
||||
|
||||
def apply_operation(self, data: Any) -> None:
|
||||
data[self.name] = self.func(data)
|
||||
|
||||
|
||||
class AHContext:
|
||||
"""
|
||||
This class is used to specify which information AutoHeuristic should store. For each choice, AutoHeursitic will
|
||||
store the context and the collected feedback. The context could be something like the shape of a tensor, i.e.,
|
||||
information that will help to learn a heuristic.
|
||||
"""
|
||||
|
||||
features: list[AHFeature]
|
||||
context_dict: dict[str, Value]
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.features = []
|
||||
self.context_dict = {}
|
||||
|
||||
def add_feature(
|
||||
self, name: str, value: Value, is_categorical: bool = False
|
||||
) -> None:
|
||||
self.features.append(AHFeature(name, value, is_categorical=is_categorical))
|
||||
self.context_dict[name] = value
|
||||
|
||||
def get_numerical_and_categorical_features(self) -> tuple[list[str], list[str]]:
|
||||
numerical_features = []
|
||||
categorical_features = []
|
||||
for feature in self.features:
|
||||
if feature.is_categorical:
|
||||
categorical_features.append(feature.name)
|
||||
else:
|
||||
numerical_features.append(feature.name)
|
||||
|
||||
return numerical_features, categorical_features
|
||||
|
||||
def get_feature_names_csv(self) -> str:
|
||||
return ",".join(feature.name for feature in self.features)
|
||||
|
||||
def get_feature_values_csv(self) -> str:
|
||||
return ",".join(str(feature.value) for feature in self.features)
|
||||
|
||||
def get_value(self, name: str) -> Value:
|
||||
return self.context_dict[name]
|
||||
|
||||
def apply_operations(self, operations: list[AHOperation]) -> None:
|
||||
for op in operations:
|
||||
op.apply_operation(self.context_dict)
|
||||
|
||||
|
||||
class AHMetadata:
|
||||
def __init__(
|
||||
self,
|
||||
shared_memory: Any,
|
||||
device_capa: tuple[int, int],
|
||||
choices: list[Choice],
|
||||
name: str,
|
||||
) -> None:
|
||||
# use amount of shared_memory and device_capability to identify GPU
|
||||
# TODO(AlnisM): there might be a better way to do this
|
||||
self.shared_memory = shared_memory
|
||||
self.device_capa = device_capa
|
||||
self.choices = choices
|
||||
self.name = name
|
||||
|
||||
def to_dict(self) -> dict[str, Value]:
|
||||
return {
|
||||
"shared_memory": self.shared_memory,
|
||||
"device_capa": self.device_capa,
|
||||
"name": self.name,
|
||||
}
|
||||
|
||||
|
||||
def get_metadata_str_from_log(log_path: str) -> str:
|
||||
with open(log_path, newline="") as file:
|
||||
json_string = file.readline().strip()
|
||||
return json_string
|
||||
|
||||
|
||||
def check_minsize(context: AHContext, minsize: int) -> bool:
|
||||
return (
|
||||
context.get_value("m") >= minsize
|
||||
and context.get_value("k") >= minsize
|
||||
and context.get_value("n") >= minsize
|
||||
)
|
||||
|
||||
|
||||
def pad_mm_precondition(metadata: AHMetadata, context: AHContext) -> bool:
|
||||
if metadata.shared_memory == 166912 and metadata.device_capa == (8, 0):
|
||||
# A100 precondition
|
||||
return check_minsize(context, 512)
|
||||
elif metadata.shared_memory == 232448 and metadata.device_capa == (9, 0):
|
||||
# H100 precondition
|
||||
return check_minsize(context, 768)
|
||||
return True
|
||||
|
||||
|
||||
def get_mixedmm_precondition(metadata: AHMetadata, context: AHContext) -> bool:
|
||||
m = context.get_value("m")
|
||||
k = context.get_value("k")
|
||||
n = context.get_value("n")
|
||||
if m > 128 or k < 1024 or n < 1024:
|
||||
return False
|
||||
mat1_iscontig = context.get_value("mat1_iscontig")
|
||||
mat2_iscontig = context.get_value("mat2_iscontig")
|
||||
return mat1_iscontig and not mat2_iscontig
|
||||
|
||||
|
||||
def get_mult_dims_ops() -> list[AHOperation]:
|
||||
m_times_k_op = AHOperation("m*k", lambda data: data["m"] * data["k"])
|
||||
m_times_n_op = AHOperation("m*n", lambda data: data["m"] * data["n"])
|
||||
k_times_n_op = AHOperation("k*n", lambda data: data["k"] * data["n"])
|
||||
return [m_times_k_op, m_times_n_op, k_times_n_op]
|
||||
|
||||
|
||||
def get_arith_intensity(data: Any) -> float:
|
||||
m = data["m"]
|
||||
k = data["k"]
|
||||
n = data["n"]
|
||||
if m == 0 or k == 0 or n == 0:
|
||||
return 0.0
|
||||
return m * k * n / (m * k + k * n + m * n)
|
||||
|
||||
|
||||
def pad_mm_operations() -> list[AHOperation]:
|
||||
mult_dims_ops = get_mult_dims_ops()
|
||||
k_div_m_times_n_op = AHOperation(
|
||||
"k/(m*n)", lambda data: data["k"] / (data["m"] * data["n"])
|
||||
)
|
||||
|
||||
def bfloat_perf_hit(data: Any) -> bool:
|
||||
m = data["m"]
|
||||
k = data["k"]
|
||||
n = data["n"]
|
||||
is_bfloat = str(data["mat1_dtype"]) == "torch.bfloat16"
|
||||
return k > (m * 1024) and k > (n * 1024) and is_bfloat
|
||||
|
||||
bfloat_perf_hit_op = AHOperation(
|
||||
"bfloat_perf_hit", bfloat_perf_hit, is_categorical=True
|
||||
)
|
||||
|
||||
arith_intensity_op = AHOperation("arith_intensity", get_arith_intensity)
|
||||
dims_need_padding_ops = get_dims_need_padding_ops()
|
||||
dims_multiple_ops = get_dims_multiple_ops()
|
||||
is_contig_ops = get_is_contig_ops()
|
||||
|
||||
ah_operations = mult_dims_ops + [
|
||||
k_div_m_times_n_op,
|
||||
bfloat_perf_hit_op,
|
||||
arith_intensity_op,
|
||||
]
|
||||
ah_operations.extend(dims_need_padding_ops)
|
||||
ah_operations.extend(dims_multiple_ops)
|
||||
ah_operations.extend(is_contig_ops)
|
||||
return ah_operations
|
||||
|
||||
|
||||
def between_op(data: Any, dim: str, lower: int, upper: int) -> bool:
|
||||
return data[dim] >= lower and data[dim] <= upper
|
||||
|
||||
|
||||
def between_ops() -> list[AHOperation]:
|
||||
dims = ["m", "k", "n"]
|
||||
limits = [(1, 16), (17, 32), (33, 64), (65, 128), (129, 256)]
|
||||
ah_operations = []
|
||||
for dim in dims:
|
||||
for lower, upper in limits:
|
||||
between_op_fn = functools.partial(
|
||||
between_op, dim=dim, lower=lower, upper=upper
|
||||
)
|
||||
# using 'LEQ' instead of '<=' because '<=' cannot be exported to dot
|
||||
between_op_name = f"{lower}LEQ{dim}LEQ{upper}"
|
||||
ah_operations.append(
|
||||
AHOperation(between_op_name, between_op_fn, is_categorical=True)
|
||||
)
|
||||
return ah_operations
|
||||
|
||||
|
||||
def pow2_op(data: Any, dim: str, exponent: int) -> bool:
|
||||
return data[dim] == 2**exponent
|
||||
|
||||
|
||||
def mm_operations() -> list[AHOperation]:
|
||||
mult_dims_ops = get_mult_dims_ops()
|
||||
arith_intensity_op = AHOperation("arith_intensity", get_arith_intensity)
|
||||
return mult_dims_ops + [arith_intensity_op]
|
||||
|
||||
|
||||
def mixed_mm_operations() -> list[AHOperation]:
|
||||
return mm_operations() + between_ops()
|
||||
|
||||
|
||||
def is_multiple(data: Any, dim: str, mult: int) -> bool:
|
||||
return data[dim] % mult == 0
|
||||
|
||||
|
||||
def get_dims_multiple_ops() -> list[AHOperation]:
|
||||
multiples = [2, 4, 8, 16, 32]
|
||||
dims = ["m", "k", "n"]
|
||||
dims_multiple_ops = []
|
||||
for dim in dims:
|
||||
for mult in multiples:
|
||||
is_multiple_fn = functools.partial(is_multiple, dim=dim, mult=mult)
|
||||
dims_multiple_op = AHOperation(
|
||||
f"{dim}_multiple_{mult}", is_multiple_fn, is_categorical=True
|
||||
)
|
||||
dims_multiple_ops.append(dims_multiple_op)
|
||||
return dims_multiple_ops
|
||||
|
||||
|
||||
def get_dims_need_padding_ops() -> list[AHOperation]:
|
||||
def mat1_innermost_needs_padding_fn(data: Any) -> bool:
|
||||
mat1_stride_0 = data["mat1_stride_0"]
|
||||
mat1_stride_1 = data["mat1_stride_1"]
|
||||
m_padded_length = data["m_padded_length"]
|
||||
k_padded_length = data["k_padded_length"]
|
||||
mat1_innermost_needs_padding = False
|
||||
if mat1_stride_0 == 1 and m_padded_length != 0:
|
||||
mat1_innermost_needs_padding = True
|
||||
if mat1_stride_1 == 1 and k_padded_length != 0:
|
||||
mat1_innermost_needs_padding = True
|
||||
return mat1_innermost_needs_padding
|
||||
|
||||
mat1_innermost_op = AHOperation(
|
||||
"mat1_innermost_needs_padding",
|
||||
mat1_innermost_needs_padding_fn,
|
||||
is_categorical=True,
|
||||
)
|
||||
|
||||
def mat2_innermost_needs_padding_fn(data: Any) -> bool:
|
||||
mat2_stride_0 = data["mat2_stride_0"]
|
||||
mat2_stride_1 = data["mat2_stride_1"]
|
||||
k_padded_length = data["k_padded_length"]
|
||||
n_padded_length = data["n_padded_length"]
|
||||
mat2_innermost_needs_padding = False
|
||||
if mat2_stride_0 == 1 and k_padded_length != 0:
|
||||
mat2_innermost_needs_padding = True
|
||||
if mat2_stride_1 == 1 and n_padded_length != 0:
|
||||
mat2_innermost_needs_padding = True
|
||||
return mat2_innermost_needs_padding
|
||||
|
||||
mat2_innermost_op = AHOperation(
|
||||
"mat2_innermost_needs_padding",
|
||||
mat2_innermost_needs_padding_fn,
|
||||
is_categorical=True,
|
||||
)
|
||||
|
||||
def num_dims_needs_padding_fn(data: Any) -> int:
|
||||
m_padded_length = data["m_padded_length"]
|
||||
k_padded_length = data["k_padded_length"]
|
||||
n_padded_length = data["n_padded_length"]
|
||||
num_dims_needs_padding = 0
|
||||
if m_padded_length != 0:
|
||||
num_dims_needs_padding += 1
|
||||
if k_padded_length != 0:
|
||||
num_dims_needs_padding += 1
|
||||
if n_padded_length != 0:
|
||||
num_dims_needs_padding += 1
|
||||
return num_dims_needs_padding
|
||||
|
||||
num_dims_op = AHOperation("num_dims_needs_padding", num_dims_needs_padding_fn)
|
||||
return [mat1_innermost_op, mat2_innermost_op, num_dims_op]
|
||||
|
||||
|
||||
def get_is_contig_ops() -> list[AHOperation]:
|
||||
def mat1_is_contig_fn(data: Any) -> bool:
|
||||
stride_0 = data["mat1_stride_0"]
|
||||
stride_1 = data["mat1_stride_1"]
|
||||
k = data["k"]
|
||||
return stride_0 == k and stride_1 == 1
|
||||
|
||||
mat1_is_contig_op = AHOperation(
|
||||
"mat1_iscontig", mat1_is_contig_fn, is_categorical=True
|
||||
)
|
||||
|
||||
def mat2_is_contig_fn(data: Any) -> bool:
|
||||
stride_0 = data["mat2_stride_0"]
|
||||
stride_1 = data["mat2_stride_1"]
|
||||
n = data["n"]
|
||||
return stride_0 == n and stride_1 == 1
|
||||
|
||||
mat2_is_contig_op = AHOperation(
|
||||
"mat2_iscontig", mat2_is_contig_fn, is_categorical=True
|
||||
)
|
||||
|
||||
return [mat1_is_contig_op, mat2_is_contig_op]
|
||||
|
||||
|
||||
def context_add_strides(context: AHContext, name: str, stride: tuple[int, ...]) -> None:
|
||||
for i, s in enumerate(stride):
|
||||
context.add_feature(f"{name}_stride_{i}", s)
|
||||
|
||||
|
||||
def context_add_using_tf32(context: AHContext, dtype: torch.dtype) -> None:
|
||||
using_tf32 = "not_float_32"
|
||||
if dtype == torch.float32:
|
||||
using_tf32 = torch.backends.cuda.matmul.fp32_precision == "tf32"
|
||||
context.add_feature("using_tf32", using_tf32, is_categorical=True)
|
||||
+119
@@ -0,0 +1,119 @@
|
||||
import importlib
|
||||
import inspect
|
||||
import pkgutil
|
||||
from collections import defaultdict
|
||||
from typing import Any
|
||||
|
||||
from torch._inductor.autoheuristic.autoheuristic_utils import (
|
||||
AHContext,
|
||||
AHMetadata,
|
||||
Choice,
|
||||
)
|
||||
from torch._inductor.autoheuristic.learnedheuristic_interface import LearnedHeuristic
|
||||
|
||||
|
||||
def find_and_instantiate_subclasses(
|
||||
package_name: str, base_class: Any
|
||||
) -> list[LearnedHeuristic]:
|
||||
instances = []
|
||||
|
||||
package = importlib.import_module(package_name)
|
||||
for _, module_name, _ in pkgutil.walk_packages(
|
||||
package.__path__, package.__name__ + "."
|
||||
):
|
||||
try:
|
||||
module_basename = module_name.split(".")[-1]
|
||||
if not module_basename.startswith("_"):
|
||||
# learned heuristics start with an underscore
|
||||
continue
|
||||
module = importlib.import_module(module_name)
|
||||
|
||||
# look for classes that are subclasses of base_class
|
||||
for _name, obj in inspect.getmembers(module):
|
||||
if (
|
||||
inspect.isclass(obj)
|
||||
and issubclass(obj, base_class)
|
||||
and obj != base_class
|
||||
):
|
||||
instance = obj()
|
||||
instances.append(instance)
|
||||
except Exception as e:
|
||||
print(f"Error processing module {module_name}: {e}")
|
||||
|
||||
return instances
|
||||
|
||||
|
||||
class LearnedHeuristicController:
|
||||
"""
|
||||
Class that finds and instantiates all learned heuristics. It also provides
|
||||
a way to get the decision of a learned heuristic.
|
||||
"""
|
||||
|
||||
existing_heuristics: dict[str, list[LearnedHeuristic]] = defaultdict(list)
|
||||
"""
|
||||
A dictionary that stores all the learned heuristics for each optimization.
|
||||
The key is the optimization name, and the value is a list of LearnedHeuristic objects.
|
||||
"""
|
||||
|
||||
heuristics_initialized: bool = False
|
||||
"""
|
||||
A flag that indicates whether the learned heuristics have been initialized.
|
||||
Set to true when the get_decision() function is called for the first time.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
metadata: AHMetadata,
|
||||
context: AHContext,
|
||||
) -> None:
|
||||
self.metadata = metadata
|
||||
self.context = context
|
||||
|
||||
def get_heuristics(self, name: str) -> list[LearnedHeuristic]:
|
||||
"""
|
||||
Returns a list of learned heuristics for the given optimization name.
|
||||
"""
|
||||
|
||||
if not LearnedHeuristicController.heuristics_initialized:
|
||||
# learned heuristics are generated into the following package
|
||||
learned_heuristics_package = "torch._inductor.autoheuristic.artifacts"
|
||||
|
||||
# learned heuristics have to be of type LearnedHeuristic
|
||||
base_class = LearnedHeuristic
|
||||
found_heuristics = find_and_instantiate_subclasses(
|
||||
learned_heuristics_package, base_class
|
||||
)
|
||||
|
||||
for learned_heuristic in found_heuristics:
|
||||
opt_name = learned_heuristic.get_name()
|
||||
LearnedHeuristicController.existing_heuristics[opt_name].append(
|
||||
learned_heuristic
|
||||
)
|
||||
LearnedHeuristicController.heuristics_initialized = True
|
||||
|
||||
return LearnedHeuristicController.existing_heuristics[name]
|
||||
|
||||
def get_decision(self) -> Choice | None:
|
||||
"""
|
||||
Returns the decision made by the learned heuristic or None if no heuristic was found or the heuristic is unsure
|
||||
which choice to make.
|
||||
"""
|
||||
|
||||
heuristics = self.get_heuristics(self.metadata.name)
|
||||
for heuristic in heuristics:
|
||||
if heuristic.check_precondition(self.metadata, self.context):
|
||||
return heuristic.get_decision(self.context, self.metadata.choices)
|
||||
return None
|
||||
|
||||
def get_decisions_ranked(self, top_k: int) -> list[Choice] | None:
|
||||
heuristics = self.get_heuristics(self.metadata.name)
|
||||
for heuristic in heuristics:
|
||||
if heuristic.check_precondition(self.metadata, self.context):
|
||||
choices = heuristic.get_decisions_ranked(self.context)
|
||||
if choices is None:
|
||||
return None
|
||||
avail_choices = [
|
||||
choice for choice in choices if choice in self.metadata.choices
|
||||
]
|
||||
return avail_choices[:top_k]
|
||||
return None
|
||||
+82
@@ -0,0 +1,82 @@
|
||||
import operator
|
||||
|
||||
from torch._inductor.autoheuristic.autoheuristic_utils import (
|
||||
AHContext,
|
||||
AHMetadata,
|
||||
Choice,
|
||||
)
|
||||
|
||||
|
||||
class LearnedHeuristic:
|
||||
"""
|
||||
LearnedHeuristic is a base class for all learned heuristics.
|
||||
"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
pass
|
||||
|
||||
def check_precondition(
|
||||
self,
|
||||
metadata: AHMetadata,
|
||||
context: AHContext,
|
||||
) -> bool:
|
||||
return True
|
||||
|
||||
def get_decision(self, context: AHContext, choices: list[Choice]) -> Choice | None:
|
||||
return None
|
||||
|
||||
def get_confidence_threshold(self) -> float:
|
||||
return 1.0
|
||||
|
||||
def get_name(self) -> str:
|
||||
return ""
|
||||
|
||||
def get_decisions_ranked(self, context: AHContext) -> list[str] | None:
|
||||
return None
|
||||
|
||||
|
||||
class LearnedHeuristicRegression(LearnedHeuristic):
|
||||
def get_feedback(self, context: AHContext, choice: Choice) -> float:
|
||||
return 1.0
|
||||
|
||||
def get_decision(self, context: AHContext, choices: list[Choice]) -> Choice | None:
|
||||
choice2feedback = {}
|
||||
for choice in choices:
|
||||
predicted_feedback = self.get_feedback(context, choice)
|
||||
choice2feedback[choice] = predicted_feedback
|
||||
sorted_choices_feedback = sorted(
|
||||
choice2feedback.items(), key=operator.itemgetter(1)
|
||||
)
|
||||
highest_feedback = sorted_choices_feedback[-1][1]
|
||||
second_highest_feedback = sorted_choices_feedback[-2][1]
|
||||
if highest_feedback / second_highest_feedback > self.get_confidence_threshold():
|
||||
return sorted_choices_feedback[-1][0]
|
||||
# We are not sure which choice is the best one
|
||||
return None
|
||||
|
||||
|
||||
class LearnedHeuristicDecision(LearnedHeuristic):
|
||||
def get_choice(self, idx: int) -> str | None:
|
||||
return None
|
||||
|
||||
def get_decision(self, context: AHContext, choices: list[Choice]) -> Choice | None:
|
||||
best_choices = self.get_best_choices(context)
|
||||
if not best_choices:
|
||||
return None
|
||||
(best_choice_proba, best_choice_idx) = best_choices[0]
|
||||
if best_choice_proba <= self.get_confidence_threshold():
|
||||
return None
|
||||
return self.get_choice(best_choice_idx)
|
||||
|
||||
def get_decisions_ranked(self, context: AHContext) -> list[str] | None:
|
||||
feedback_idx_list = self.get_best_choices(context)
|
||||
if feedback_idx_list is None:
|
||||
return None
|
||||
choices = [
|
||||
self.get_choice(feedback_idx[1]) for feedback_idx in feedback_idx_list
|
||||
]
|
||||
choices = [choice for choice in choices if choice is not None]
|
||||
return choices
|
||||
|
||||
def get_best_choices(self, context: AHContext) -> list[tuple[float, int]] | None:
|
||||
return []
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,178 @@
|
||||
import asyncio
|
||||
import sys
|
||||
import weakref
|
||||
from asyncio import AbstractEventLoop, Future
|
||||
from collections.abc import Awaitable, Callable, Coroutine, Generator, Iterator
|
||||
from contextlib import contextmanager, ExitStack
|
||||
from contextvars import Context
|
||||
from typing import Any, Protocol, TypeVar
|
||||
|
||||
from torch.utils._ordered_set import OrderedSet
|
||||
|
||||
|
||||
T = TypeVar("T")
|
||||
TCoro = Generator[Any, None, T]
|
||||
|
||||
if sys.version_info >= (3, 11):
|
||||
|
||||
class TaskFactory(Protocol):
|
||||
def __call__(
|
||||
self,
|
||||
__loop: AbstractEventLoop,
|
||||
__factory: Coroutine[None, None, object] | Generator[None, None, object],
|
||||
__context: Context | None = None,
|
||||
/,
|
||||
) -> asyncio.futures.Future[object]: ...
|
||||
|
||||
TaskFactoryType = TaskFactory
|
||||
else:
|
||||
TaskFactoryType = Callable[[AbstractEventLoop, Generator[TCoro, None, T]], Future] # type: ignore[valid-type]
|
||||
|
||||
|
||||
def await_sync(awaitable: Awaitable[T]) -> T:
|
||||
with get_loop() as loop:
|
||||
return loop.run_until_complete(awaitable)
|
||||
|
||||
|
||||
@contextmanager
|
||||
def get_loop(
|
||||
always_create_new_loop: bool = False,
|
||||
) -> Iterator[AbstractEventLoop]:
|
||||
try:
|
||||
loop = asyncio.get_event_loop()
|
||||
except RuntimeError as re:
|
||||
if "There is no current event loop in thread" in str(re):
|
||||
with _new_loop() as loop:
|
||||
yield loop
|
||||
return
|
||||
else:
|
||||
raise
|
||||
|
||||
@contextmanager
|
||||
def _restore_loop(
|
||||
loop: asyncio.AbstractEventLoop,
|
||||
) -> Iterator[None]:
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
asyncio.set_event_loop(loop)
|
||||
|
||||
@contextmanager
|
||||
def _restore_running_loop() -> Iterator[None]:
|
||||
loop_from_events = asyncio.events._get_running_loop()
|
||||
asyncio.events._set_running_loop(None)
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
asyncio.events._set_running_loop(loop_from_events)
|
||||
|
||||
with ExitStack() as stack:
|
||||
if loop.is_running():
|
||||
stack.enter_context(_restore_running_loop())
|
||||
stack.enter_context(_restore_loop(loop=loop))
|
||||
loop = stack.enter_context(_new_loop(loop.get_task_factory())) # type: ignore[arg-type]
|
||||
elif loop.is_closed():
|
||||
loop = stack.enter_context(_new_loop()) # type: ignore[arg-type]
|
||||
elif always_create_new_loop:
|
||||
stack.enter_context(_restore_loop(loop=loop))
|
||||
loop = stack.enter_context(_new_loop()) # type: ignore[arg-type]
|
||||
yield loop
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _new_loop(
|
||||
task_factory: TaskFactoryType | None = None,
|
||||
) -> Iterator[asyncio.AbstractEventLoop]:
|
||||
loop = asyncio.new_event_loop()
|
||||
tasks = _patch_loop(loop)
|
||||
|
||||
if task_factory:
|
||||
# pyre-ignore[6]
|
||||
loop.set_task_factory(task_factory) # type: ignore[arg-type]
|
||||
|
||||
asyncio.set_event_loop(loop)
|
||||
try:
|
||||
yield loop
|
||||
finally:
|
||||
try:
|
||||
_cancel_all_tasks(loop, tasks)
|
||||
finally:
|
||||
asyncio.set_event_loop(None)
|
||||
loop.close()
|
||||
|
||||
|
||||
def _cancel_all_tasks(
|
||||
loop: AbstractEventLoop,
|
||||
tasks: OrderedSet[Future], # type: ignore[type-arg]
|
||||
) -> None:
|
||||
to_cancel = [task for task in tasks if not task.done()]
|
||||
|
||||
if not to_cancel:
|
||||
return
|
||||
|
||||
# pyre-fixme[1001]: Awaitable assigned to `task` is never awaited.
|
||||
for task in to_cancel:
|
||||
task.cancel()
|
||||
|
||||
# pyrefly: ignore [bad-argument-type]
|
||||
loop.run_until_complete(asyncio.gather(*to_cancel, return_exceptions=True))
|
||||
|
||||
for task in to_cancel:
|
||||
if task.cancelled():
|
||||
continue
|
||||
if task.exception() is not None:
|
||||
loop.call_exception_handler(
|
||||
{
|
||||
"message": "unhandled exception during asyncio.run() shutdown",
|
||||
"exception": task.exception(),
|
||||
"task": task,
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def _patch_loop(loop: AbstractEventLoop) -> OrderedSet[Future]: # type: ignore[type-arg]
|
||||
tasks: weakref.WeakSet[Future] = weakref.WeakSet() # type: ignore[type-arg]
|
||||
|
||||
task_factories: list[TaskFactoryType | None] = [None]
|
||||
|
||||
def _set_task_factory(factory: TaskFactoryType | None) -> None:
|
||||
task_factories[0] = factory
|
||||
|
||||
def _get_task_factory() -> TaskFactoryType | None:
|
||||
return task_factories[0]
|
||||
|
||||
def _safe_task_factory(
|
||||
loop: AbstractEventLoop,
|
||||
coro: TCoro, # type: ignore[type-arg]
|
||||
*,
|
||||
context: Context | None = None,
|
||||
) -> asyncio.Future: # type: ignore[valid-type, type-arg]
|
||||
task_factory = task_factories[0]
|
||||
if task_factory is None:
|
||||
if sys.version_info >= (3, 11):
|
||||
# pyrefly: ignore [bad-argument-type]
|
||||
task = asyncio.Task(coro, loop=loop, context=context)
|
||||
else:
|
||||
task = asyncio.Task(coro, loop=loop)
|
||||
# pyre-ignore[16]: `Task` has no attribute `_source_traceback`.
|
||||
if task._source_traceback: # type: ignore[attr-defined]
|
||||
del task._source_traceback[ # type: ignore[attr-defined]
|
||||
-1
|
||||
] # pragma: no cover # type: ignore[attr-defined]
|
||||
else:
|
||||
if sys.version_info >= (3, 11):
|
||||
task = task_factory(loop, coro, context=context) # type: ignore[arg-type, call-arg, assignment]
|
||||
else:
|
||||
task = task_factory(loop, coro) # type: ignore[arg-type]
|
||||
# `Union[Task[Any], Future[Any]]`.
|
||||
tasks.add(task)
|
||||
return task
|
||||
|
||||
# pyre-ignore[6]
|
||||
loop.set_task_factory(_safe_task_factory) # type: ignore[method-assign, arg-type]
|
||||
# pyre-ignore[8]
|
||||
loop.set_task_factory = _set_task_factory # type: ignore[method-assign, assignment]
|
||||
# pyre-ignore[8]
|
||||
loop.get_task_factory = _get_task_factory # type: ignore[method-assign, assignment]
|
||||
|
||||
return tasks # type: ignore[return-value]
|
||||
@@ -0,0 +1,263 @@
|
||||
import logging
|
||||
import operator
|
||||
from collections.abc import Callable
|
||||
from functools import partial
|
||||
from typing import Any
|
||||
|
||||
import sympy
|
||||
from sympy import Expr
|
||||
|
||||
import torch
|
||||
from torch.utils._sympy.value_ranges import (
|
||||
bound_sympy,
|
||||
SymPyValueRangeAnalysis,
|
||||
ValueRanges,
|
||||
)
|
||||
|
||||
from ..utils._sympy.functions import PowByNatural
|
||||
from ..utils._sympy.numbers import int_oo
|
||||
from .loop_body import InterpreterShim, LoopBody, LoopBodyBlock
|
||||
from .ops_handler import DefaultHandler, ReductionType, StoreMode
|
||||
from .utils import cache_on_self, dominated_nodes
|
||||
from .virtualized import V
|
||||
|
||||
|
||||
log = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class BoundVars:
|
||||
"""
|
||||
Performs Value Range Analysis on LoopBody's fx graph by calling BoundVars.run()
|
||||
It exposes the ranges of the nodes in the `bounds` variable
|
||||
|
||||
Note. A current limitation of this analysis is that it just works on a per-loop basis.
|
||||
We should be able to propagate the bounds between across the whole graph. This may benefit
|
||||
the case a bounded variable is returned by a kernel and fed into another.
|
||||
"""
|
||||
|
||||
def __init__(self, loop_body: LoopBody) -> None:
|
||||
def upper_bound(v: Expr | int) -> int:
|
||||
return bound_sympy(v).upper if isinstance(v, Expr) else v
|
||||
|
||||
self.loop_body = loop_body
|
||||
self.replacement_vals = {
|
||||
k: ValueRanges[Expr](0, upper_bound(v) - 1)
|
||||
for k, v in loop_body.var_ranges.items()
|
||||
}
|
||||
# avoid computing these values, pessimistically assume that they are unbounded
|
||||
self.unbounded_vars = dominated_nodes(
|
||||
node
|
||||
for node in self.loop_body.get_nodes()
|
||||
if node.target in ["load", "reduction", operator.getitem]
|
||||
or "masked_subblock" in node.target
|
||||
)
|
||||
# To access this variable call `get_bounds()`
|
||||
self._bounds: dict[torch.fx.Node, ValueRanges[Expr]] = {}
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return (
|
||||
f"{self.__class__.__name__}("
|
||||
f"loop_body={self.loop_body},\n "
|
||||
f"replacement_vals={self.replacement_vals}, \n"
|
||||
f"unbounded_vars={self.unbounded_vars}, \n"
|
||||
f"_bounds={self._bounds})"
|
||||
)
|
||||
|
||||
@cache_on_self
|
||||
def get_bounds(self) -> dict[torch.fx.Node, ValueRanges[Expr]]:
|
||||
submodules = self.swap_submodules(self.loop_body.submodules)
|
||||
|
||||
# Initialize the environment with the unbounded variables
|
||||
for node in self.unbounded_vars:
|
||||
# we need to evaluate masked_subblock to recurse, and we need to set indirect values
|
||||
if not isinstance(node.target, str) or (
|
||||
"masked_subblock" not in node.target
|
||||
and "set_indirect" not in node.target
|
||||
):
|
||||
self._bounds[node] = ValueRanges[Expr].unknown()
|
||||
|
||||
with V.set_ops_handler(ValueRangeAnalysis()):
|
||||
interpreter = InterpreterShim(self.loop_body.root_block.graph, submodules)
|
||||
log.debug("get_bounds:\n%s", self.loop_body.root_block.graph)
|
||||
interpreter.run(V.get_ops_handler(), initial_env=self._bounds)
|
||||
return self._bounds
|
||||
|
||||
def swap_submodules(
|
||||
self, submodules: dict[str, Callable[..., Any]]
|
||||
) -> dict[str, Callable[..., ValueRanges[Expr]]]:
|
||||
result: dict[str, Callable[..., ValueRanges[Expr]]] = {}
|
||||
for key in submodules:
|
||||
if key == "get_index":
|
||||
result[key] = self.get_index
|
||||
elif "masked_subblock" in key:
|
||||
subblock = self.loop_body.subblocks[key]
|
||||
# The result within the lambda will reference to the final
|
||||
# set of modules at the end of the for-loop as it stores a reference to it
|
||||
|
||||
# bind subblock in a function because python lambdas close over by reference
|
||||
# moving the lambda out of make_fn would close over the reference to subblock,
|
||||
# so all lambdas would have the same subblock reference that is the final
|
||||
# subblock in the loop
|
||||
def make_fn(
|
||||
subblock: LoopBodyBlock,
|
||||
) -> Callable[[Any, Any], ValueRanges[Expr]]:
|
||||
return lambda mask, value: self.masked_subblock(
|
||||
subblock, self._bounds, mask, value, result
|
||||
)
|
||||
|
||||
result[key] = make_fn(subblock)
|
||||
elif "set_indirect" in key:
|
||||
idx = int(key[len("set_indirect") :])
|
||||
var = self.loop_body.indirect_vars[idx]
|
||||
indirect = partial(self.set_indirect, var)
|
||||
result[key] = indirect
|
||||
else:
|
||||
assert "scan" in key
|
||||
result[key] = submodules[key]
|
||||
|
||||
return result
|
||||
|
||||
def masked_subblock(
|
||||
self,
|
||||
subblock: LoopBodyBlock,
|
||||
env: dict[torch.fx.Node, ValueRanges[Expr]],
|
||||
mask: Any,
|
||||
value: Any,
|
||||
submodules: dict[str, Callable[..., Any]],
|
||||
) -> ValueRanges[Expr]:
|
||||
interp = InterpreterShim(subblock.graph, submodules)
|
||||
interp.run(V.get_ops_handler(), initial_env=env)
|
||||
output = [node for node in subblock.graph.nodes if node.target == "output"]
|
||||
assert len(output) == 1
|
||||
# dont bother unioning with value since the load from buffer will be
|
||||
# pessimistically assumed to be inf anyway
|
||||
return interp.env[output[0]]
|
||||
|
||||
def set_indirect(self, old: Expr, new: ValueRanges[Expr]) -> ValueRanges[Expr]:
|
||||
assert isinstance(new, ValueRanges)
|
||||
self.replacement_vals[old] = new
|
||||
return new
|
||||
|
||||
def get_index(self, name: str) -> ValueRanges[Expr]:
|
||||
expr = self.loop_body.indexing_exprs[name]
|
||||
bound = self.replacement_vals.get(expr)
|
||||
if bound is None:
|
||||
bound = bound_sympy(expr, self.replacement_vals)
|
||||
# The following assertion is true at the time of this writing
|
||||
# We don't assert is as to not execute bound_sympy when bound is not None
|
||||
# assert bound is None or bound == bound_sympy(expr, self.replacement_vals)
|
||||
self.replacement_vals[name] = bound
|
||||
return bound
|
||||
|
||||
|
||||
class ValueRangeAnalysis(SymPyValueRangeAnalysis, DefaultHandler):
|
||||
def __init__(self) -> None:
|
||||
self.name = "ValueRangeAnalysis"
|
||||
boolean_operators = (
|
||||
"xor",
|
||||
"logical_and",
|
||||
"logical_or",
|
||||
"logical_not",
|
||||
)
|
||||
for op in boolean_operators:
|
||||
setattr(self, op, self.bool_handler)
|
||||
|
||||
@staticmethod
|
||||
def bool_handler(*args: Any, **kwargs: Any) -> ValueRanges[Any]:
|
||||
# just assuming bools can have both values
|
||||
return ValueRanges(sympy.false, sympy.true) # type: ignore[arg-type]
|
||||
|
||||
def _default(self, name: str, args: tuple[Any, ...], kwargs: dict[str, Any]) -> Any:
|
||||
# many ops are unlikely to show up in optimizable indexing compute,
|
||||
# so we dont have full coverage
|
||||
return ValueRanges.unknown()
|
||||
|
||||
def load(self, name: str, index: sympy.Expr) -> ValueRanges[Any]:
|
||||
return ValueRanges.unknown()
|
||||
|
||||
def store(
|
||||
self, name: str, index: sympy.Expr, value: Any, mode: StoreMode = None
|
||||
) -> None:
|
||||
return
|
||||
|
||||
def reduction(
|
||||
self,
|
||||
dtype: torch.dtype,
|
||||
src_dtype: torch.dtype,
|
||||
reduction_type: ReductionType,
|
||||
value: Any,
|
||||
) -> ValueRanges[Any]:
|
||||
return ValueRanges.unknown()
|
||||
|
||||
@classmethod
|
||||
def index_expr(cls, index: Any, dtype: torch.dtype) -> ValueRanges[Any]:
|
||||
assert isinstance(index, ValueRanges)
|
||||
return cls.to_dtype(index, dtype)
|
||||
|
||||
@staticmethod
|
||||
# pyrefly: ignore [bad-override]
|
||||
def to_dtype(
|
||||
x: Any,
|
||||
dtype: torch.dtype,
|
||||
src_dtype: torch.dtype | None = None,
|
||||
use_compute_types: bool = True,
|
||||
) -> ValueRanges[Any]:
|
||||
x = ValueRanges.wrap(x)
|
||||
|
||||
if dtype == torch.bool:
|
||||
if x.is_singleton():
|
||||
return ValueRanges.wrap(x.lower != 0)
|
||||
elif x.is_bool:
|
||||
return x
|
||||
elif 0 not in x:
|
||||
return ValueRanges.wrap(sympy.true)
|
||||
else:
|
||||
return ValueRanges(sympy.false, sympy.true)
|
||||
|
||||
def cast(x: Any, dtype: torch.dtype) -> sympy.Expr:
|
||||
# dtype is int or float
|
||||
if dtype.is_floating_point:
|
||||
return sympy.Float(x)
|
||||
else:
|
||||
if x in (int_oo, -int_oo):
|
||||
return x
|
||||
try:
|
||||
return sympy.Integer(x)
|
||||
except TypeError:
|
||||
# inf cannot be cast to Integer
|
||||
return x
|
||||
|
||||
if x.is_bool:
|
||||
if x.is_singleton():
|
||||
val = 1 if x.lower else 0
|
||||
return ValueRanges.wrap(cast(val, dtype))
|
||||
else:
|
||||
return ValueRanges(cast(0, dtype), cast(1, dtype))
|
||||
else:
|
||||
# int to float or float to int
|
||||
return ValueRanges(cast(x.lower, dtype), cast(x.upper, dtype))
|
||||
|
||||
@staticmethod
|
||||
# pyrefly: ignore [bad-override]
|
||||
def square(x: Any) -> ValueRanges[Any]:
|
||||
return ValueRanges.convex_min_zero_map(x, lambda y: PowByNatural(y, 2))
|
||||
|
||||
@staticmethod
|
||||
# pyrefly: ignore [bad-override]
|
||||
def neg(x: Any) -> ValueRanges[Any]:
|
||||
return ValueRanges.decreasing_map(x, operator.neg)
|
||||
|
||||
# TODO: this is slightly inaccurate because truncdiv operates at integer
|
||||
# precision, but we're going through float truediv which means we can
|
||||
# potentially lose precision on the bounds
|
||||
@classmethod
|
||||
def truncdiv(cls, a: Any, b: Any) -> ValueRanges[Any]:
|
||||
x = cls.truediv(a, b)
|
||||
if x == ValueRanges.unknown():
|
||||
return x
|
||||
|
||||
return cls.trunc(x)
|
||||
|
||||
@classmethod
|
||||
def sub(cls, a: Any, b: Any) -> ValueRanges[Any]:
|
||||
return cls.add(a, cls.neg(b))
|
||||
@@ -0,0 +1,419 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import pickle
|
||||
from abc import ABC, abstractmethod
|
||||
from ast import literal_eval
|
||||
from functools import cached_property
|
||||
from hashlib import sha256
|
||||
from os import getenv
|
||||
from pathlib import Path
|
||||
from tempfile import gettempdir
|
||||
from threading import Lock
|
||||
from typing import Any, Generic, TYPE_CHECKING, TypeVar
|
||||
from typing_extensions import assert_never, override, Self
|
||||
|
||||
from torch.utils._filelock import FileLock
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from concurrent.futures import Future, ThreadPoolExecutor
|
||||
|
||||
|
||||
# TypeVars can't be recursive, so generic types that fall within
|
||||
# Key or Value can't be bound properly; for example, Key should
|
||||
# only take tuples of other Key types: tuple[Key, ...]. this is
|
||||
# a known shortcoming of torch's typing
|
||||
Key = TypeVar("Key", str, int, tuple[Any, ...])
|
||||
Value = TypeVar("Value", str, int, tuple[Any, ...], bytes, dict[Any, Any], list[Any])
|
||||
|
||||
|
||||
class CacheError(ValueError):
|
||||
"""
|
||||
Exception raised for errors encountered during cache operations.
|
||||
"""
|
||||
|
||||
|
||||
class Cache(ABC, Generic[Key, Value]):
|
||||
"""
|
||||
Abstract base class for cache implementations.
|
||||
Provides the interface for cache operations.
|
||||
"""
|
||||
|
||||
@abstractmethod
|
||||
def get(self: Self, key: Key) -> Value | None:
|
||||
"""
|
||||
Retrieve a value from the cache.
|
||||
Args:
|
||||
key (Key): The key to look up.
|
||||
Returns:
|
||||
Value | None: The cached value if present, else None.
|
||||
"""
|
||||
|
||||
@abstractmethod
|
||||
def insert(self: Self, key: Key, value: Value) -> bool:
|
||||
"""
|
||||
Insert a value into the cache.
|
||||
Args:
|
||||
key (Key): The key to insert.
|
||||
value (Value): The value to associate with the key.
|
||||
Returns:
|
||||
bool: True if the value was inserted, False if the key already exists.
|
||||
"""
|
||||
|
||||
|
||||
class InMemoryCache(Cache[Key, Value]):
|
||||
"""
|
||||
In-memory cache implementation using a dictionary and thread lock.
|
||||
"""
|
||||
|
||||
def __init__(self: Self) -> None:
|
||||
"""
|
||||
Initialize an empty in-memory cache.
|
||||
"""
|
||||
self._cache: dict[Key, Value] = {}
|
||||
self._lock: Lock = Lock()
|
||||
|
||||
def get(self: Self, key: Key) -> Value | None:
|
||||
"""
|
||||
Retrieve a value from the cache.
|
||||
Args:
|
||||
key (Key): The key to look up.
|
||||
Returns:
|
||||
Value | None: The cached value if present, else None.
|
||||
"""
|
||||
with self._lock:
|
||||
if (value := self._cache.get(key)) is not None:
|
||||
return value
|
||||
return None
|
||||
|
||||
def insert(self: Self, key: Key, value: Value) -> bool:
|
||||
"""
|
||||
Insert a value into the cache.
|
||||
Args:
|
||||
key (Key): The key to insert.
|
||||
value (Value): The value to associate with the key.
|
||||
Returns:
|
||||
bool: True if the value was inserted, False if the key already exists.
|
||||
"""
|
||||
with self._lock:
|
||||
if key in self._cache:
|
||||
# no overwrites for insert!
|
||||
return False
|
||||
self._cache[key] = value
|
||||
return True
|
||||
|
||||
@classmethod
|
||||
def from_env_var(cls, env_var: str) -> Self:
|
||||
"""
|
||||
Create an in-memory cache from an environment variable.
|
||||
Args:
|
||||
env_var (str): Name of the environment variable containing cache data.
|
||||
Returns:
|
||||
InMemoryCache: An instance populated from the environment variable.
|
||||
Raises:
|
||||
CacheError: If the environment variable is malformed or contains invalid data.
|
||||
"""
|
||||
cache = cls()
|
||||
|
||||
if (env_val := getenv(env_var)) is None:
|
||||
# env_var doesn't exist = empty cache
|
||||
return cache
|
||||
|
||||
for kv_pair in env_val.split(";"):
|
||||
# ignore whitespace prefix/suffix
|
||||
kv_pair = kv_pair.strip()
|
||||
|
||||
if not kv_pair:
|
||||
# kv_pair could be '' if env_val is '' or has ; suffix
|
||||
continue
|
||||
|
||||
try:
|
||||
# keys and values should be comma separated
|
||||
key_bytes_repr, value_bytes_repr = kv_pair.split(",", 1)
|
||||
except ValueError as err:
|
||||
raise CacheError(
|
||||
f"Malformed kv_pair {kv_pair!r} from env_var {env_var!r}, likely missing comma separator."
|
||||
) from err
|
||||
|
||||
# ignore whitespace prefix/suffix, again
|
||||
key_bytes_repr, value_bytes_repr = (
|
||||
key_bytes_repr.strip(),
|
||||
value_bytes_repr.strip(),
|
||||
)
|
||||
|
||||
try:
|
||||
# check that key_bytes_str is an actual, legitimate encoding
|
||||
key_bytes = literal_eval(key_bytes_repr)
|
||||
except (ValueError, SyntaxError) as err:
|
||||
raise CacheError(
|
||||
f"Malformed key_bytes_repr {key_bytes_repr!r} in kv_pair {kv_pair!r}, encoding is invalid."
|
||||
) from err
|
||||
try:
|
||||
# check that value_bytes_str is an actual, legitimate encoding
|
||||
value_bytes = literal_eval(value_bytes_repr)
|
||||
except (ValueError, SyntaxError) as err:
|
||||
raise CacheError(
|
||||
f"Malformed value_bytes_repr {value_bytes_repr!r} in kv_pair {kv_pair!r}, encoding is invalid."
|
||||
) from err
|
||||
|
||||
try:
|
||||
key = pickle.loads(key_bytes)
|
||||
except pickle.UnpicklingError as err:
|
||||
raise CacheError(
|
||||
f"Malformed key_bytes_repr {key_bytes_repr!r} in kv_pair {kv_pair!r}, not un-pickle-able."
|
||||
) from err
|
||||
try:
|
||||
value = pickle.loads(value_bytes)
|
||||
except pickle.UnpicklingError as err:
|
||||
raise CacheError(
|
||||
f"Malformed value_bytes_repr {value_bytes_repr!r} in kv_pair {kv_pair!r}, not un-pickle-able."
|
||||
) from err
|
||||
|
||||
# true duplicates, i.e. multiple occurrences of the same key => value
|
||||
# mapping are ok and treated as a no-op; key duplicates with differing
|
||||
# values, i.e. key => value_1 and key => value_2 where value_1 != value_2,
|
||||
# are not okay since we don't allow overwriting cached values (it's bad regardless)
|
||||
if (not cache.insert(key, value)) and (cache.get(key) != value):
|
||||
raise CacheError(
|
||||
f"Multiple values for key {key!r} found, got {cache.get(key)!r} and {value!r}."
|
||||
)
|
||||
|
||||
return cache
|
||||
|
||||
@classmethod
|
||||
def from_file_path(cls, fpath: Path) -> Self:
|
||||
"""
|
||||
Create an in-memory cache from a file path.
|
||||
Args:
|
||||
fpath (Path): Path to the file containing pickled cache data.
|
||||
Returns:
|
||||
InMemoryCache: An instance populated from the file.
|
||||
Raises:
|
||||
CacheError: If the file is not a valid pickled dictionary.
|
||||
"""
|
||||
cache = cls()
|
||||
|
||||
if not fpath.is_file():
|
||||
# fpath doesn't exit = empty cache
|
||||
return cache
|
||||
|
||||
try:
|
||||
with open(fpath, "rb") as fp:
|
||||
cache._cache = pickle.load(fp)
|
||||
except pickle.UnpicklingError as err:
|
||||
raise CacheError(
|
||||
f"Failed to create cache from file path {fpath}, file contents are un-pickle-able."
|
||||
) from err
|
||||
|
||||
if not isinstance(cache._cache, dict):
|
||||
raise CacheError(
|
||||
f"Failed to create cache from file path {fpath}, file contents not pickled dict[Key, Value]."
|
||||
)
|
||||
|
||||
return cache
|
||||
|
||||
|
||||
class AsyncCache(Cache[Key, Value]):
|
||||
"""
|
||||
Asynchronous cache implementation using ThreadPoolExecutor.
|
||||
"""
|
||||
|
||||
def get_async(
|
||||
self: Self, key: Key, executor: ThreadPoolExecutor
|
||||
) -> Future[Value | None]:
|
||||
"""
|
||||
Retrieve a value from the cache asynchronously.
|
||||
Args:
|
||||
key (Key): The key to look up.
|
||||
executor (ThreadPoolExecutor): Executor for async execution.
|
||||
Returns:
|
||||
Future[Value | None]: Future for the cached value or None.
|
||||
"""
|
||||
return executor.submit(self.get, key)
|
||||
|
||||
def insert_async(
|
||||
self: Self, key: Key, value: Value, executor: ThreadPoolExecutor
|
||||
) -> Future[bool]:
|
||||
"""
|
||||
Insert a value into the cache asynchronously.
|
||||
Args:
|
||||
key (Key): The key to insert.
|
||||
value (Value): The value to associate with the key.
|
||||
executor (ThreadPoolExecutor): Executor for async execution.
|
||||
Returns:
|
||||
Future[bool]: Future for the result of insertion.
|
||||
"""
|
||||
return executor.submit(self.insert, key, value)
|
||||
|
||||
|
||||
class OnDiskCache(AsyncCache[Key, Value]):
|
||||
"""
|
||||
On-disk cache implementation using files and file locks.
|
||||
Stores cache data in files on disk, with atomic operations and versioning.
|
||||
Supports custom cache directory names.
|
||||
Attributes:
|
||||
version (int): The version used for cache versioning.
|
||||
name (str): The name of the cache directory.
|
||||
"""
|
||||
|
||||
version: int = 0
|
||||
|
||||
def __init__(self: Self, name: str | None = None) -> None:
|
||||
"""
|
||||
Initialize an on-disk cache instance.
|
||||
Args:
|
||||
name (str | None, optional): The name of the cache directory. If None,
|
||||
defaults to "on_disk_cache".
|
||||
"""
|
||||
self.name = name or "on_disk_cache"
|
||||
|
||||
@cached_property
|
||||
def base_dir(self: Self) -> Path:
|
||||
"""
|
||||
Get the base directory for the cache.
|
||||
Returns:
|
||||
Path: The base directory path for storing cache files.
|
||||
"""
|
||||
return Path(gettempdir()) / "cache" / self.name
|
||||
|
||||
def _fpath_from_key(self: Self, key: Key) -> Path:
|
||||
"""
|
||||
Get the file path for a given key.
|
||||
Args:
|
||||
key (Key): The key to convert to a file path.
|
||||
Returns:
|
||||
Path: The file path for the key.
|
||||
Raises:
|
||||
CacheError: If the key is not pickle-able.
|
||||
"""
|
||||
try:
|
||||
return self.base_dir / sha256(pickle.dumps(key)).hexdigest()[:32]
|
||||
except (AttributeError, pickle.PicklingError) as err:
|
||||
raise CacheError(
|
||||
f"Failed to get fpath for key {key!r}, key is not pickle-able."
|
||||
) from err
|
||||
# pyrefly: ignore [bad-argument-type]
|
||||
assert_never(key)
|
||||
|
||||
def _flock_from_fpath(self: Self, fpath: Path) -> FileLock:
|
||||
"""
|
||||
Get a file lock for a given file path.
|
||||
Args:
|
||||
fpath (Path): The file path.
|
||||
Returns:
|
||||
FileLock: The file lock for the path.
|
||||
"""
|
||||
# fpath.name is a hex digest, meaning there are 16^4 potential values
|
||||
# for fpath.name[:4]; this is more than enough unique locks to not
|
||||
# cause additional overhead from shared locks and it also saves our
|
||||
# cache dir from becoming 50 percent locks
|
||||
# pyrefly: ignore [bad-return]
|
||||
return FileLock(str(fpath.parent / "locks" / fpath.name[:4]) + ".lock")
|
||||
|
||||
@property
|
||||
def version_prefix(self: Self) -> bytes:
|
||||
"""
|
||||
Get the version prefix for the cache.
|
||||
Returns:
|
||||
bytes: The version prefix as bytes, derived from the cache version string.
|
||||
"""
|
||||
return sha256(str(OnDiskCache.version).encode()).digest()[:4]
|
||||
|
||||
@override
|
||||
def get(self: Self, key: Key) -> Value | None:
|
||||
"""
|
||||
Retrieve a value from the cache.
|
||||
Args:
|
||||
key (Key): The key to look up.
|
||||
Returns:
|
||||
Value | None: The cached value if present and version matches, else None.
|
||||
Raises:
|
||||
CacheError: If the value is corrupted or cannot be unpickled.
|
||||
Side Effects:
|
||||
Removes stale cache files if the version prefix does not match.
|
||||
"""
|
||||
fpath = self._fpath_from_key(key)
|
||||
flock = self._flock_from_fpath(fpath)
|
||||
|
||||
with flock:
|
||||
if not fpath.is_file():
|
||||
return None
|
||||
|
||||
value_bytes = None
|
||||
prefix_length = len(self.version_prefix)
|
||||
with open(fpath, "rb") as fp:
|
||||
if fp.read(prefix_length) == self.version_prefix:
|
||||
value_bytes = fp.read()
|
||||
|
||||
if value_bytes is None:
|
||||
# version_prefix did not match, so we can't read the stale
|
||||
# cached value; we should also remove the stale cached value,
|
||||
# so that key can be re-cached by the newer version
|
||||
fpath.unlink()
|
||||
return None
|
||||
|
||||
try:
|
||||
value = pickle.loads(value_bytes)
|
||||
except pickle.UnpicklingError as err:
|
||||
raise CacheError(
|
||||
f"Failed to get key {key!r}, value is potentially corrupted (value is not un-pickle-able)."
|
||||
) from err
|
||||
|
||||
return value
|
||||
|
||||
@override
|
||||
def insert(self: Self, key: Key, value: Value) -> bool:
|
||||
"""
|
||||
Insert a value into the cache.
|
||||
Args:
|
||||
key (Key): The key to insert.
|
||||
value (Value): The value to associate with the key.
|
||||
Returns:
|
||||
bool: True if the value was inserted, False if the key already exists.
|
||||
Raises:
|
||||
CacheError: If the value is not pickle-able.
|
||||
Side Effects:
|
||||
Creates the cache directory if it does not exist.
|
||||
"""
|
||||
fpath = self._fpath_from_key(key)
|
||||
flock = self._flock_from_fpath(fpath)
|
||||
fpath.parent.mkdir(parents=True, exist_ok=True)
|
||||
try:
|
||||
# "x" mode is exclusive creation, meaning the file will be created
|
||||
# iff the file does not already exist (atomic w/o overwrite); use
|
||||
# flock for added atomicity guarantee and to prevent partial writes
|
||||
with flock as _, open(fpath, "xb") as fp:
|
||||
fp.write(self.version_prefix)
|
||||
pickle.dump(value, fp)
|
||||
except pickle.PicklingError as err:
|
||||
raise CacheError(
|
||||
f"Failed to insert key {key!r} with value {value!r}, value is not pickle-able."
|
||||
) from err
|
||||
except FileExistsError:
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
class InductorOnDiskCache(OnDiskCache[Key, Value]):
|
||||
"""
|
||||
Inductor-specific on-disk cache implementation.
|
||||
Uses a custom base directory for Inductor cache files.
|
||||
"""
|
||||
|
||||
def __init__(self: Self) -> None:
|
||||
"""
|
||||
Initialize an inductor on-disk cache instance.
|
||||
Sets the cache directory name to "inductor_on_disk_cache".
|
||||
"""
|
||||
super().__init__("inductor_on_disk_cache")
|
||||
|
||||
@cached_property
|
||||
def base_dir(self: Self) -> Path:
|
||||
"""
|
||||
Get the base directory for the Inductor cache.
|
||||
Returns:
|
||||
Path: The base directory path for Inductor cache files.
|
||||
"""
|
||||
from torch._inductor.runtime.runtime_utils import default_cache_dir
|
||||
|
||||
return Path(default_cache_dir(), "cache", self.name)
|
||||
@@ -0,0 +1,689 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import dataclasses
|
||||
import typing
|
||||
from typing import Any, TYPE_CHECKING
|
||||
|
||||
import sympy
|
||||
|
||||
import torch
|
||||
from torch._inductor.runtime.runtime_utils import next_power_of_2
|
||||
from torch._inductor.scheduler import MixOrderReduction
|
||||
from torch.utils._sympy.value_ranges import bound_sympy
|
||||
|
||||
from . import config
|
||||
from .codecache import write_text
|
||||
from .kernel_inputs import KernelInputs # noqa: TC001
|
||||
from .kernel_template_choice import make_ktc_generator
|
||||
from .metrics import get_metric_table, is_metric_table_enabled
|
||||
from .runtime.hints import DeviceProperties, ReductionHint
|
||||
from .scheduler import BaseSchedulerNode, Scheduler, WhyNoFuse
|
||||
from .select_algorithm import ExternKernelChoice
|
||||
from .template_heuristics import get_template_heuristic
|
||||
from .template_heuristics.triton import (
|
||||
BaseConfigHeuristic,
|
||||
CPUConfigHeuristic,
|
||||
CUDAConfigHeuristic,
|
||||
MTIAConfigHeuristic,
|
||||
ROCmConfigHeuristic,
|
||||
XPUConfigHeuristic,
|
||||
)
|
||||
from .utils import _use_autotune_backend
|
||||
from .virtualized import V
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Generator
|
||||
from functools import partial
|
||||
|
||||
from triton import Config as TritonConfig
|
||||
|
||||
from .codegen.common import KernelTemplate
|
||||
from .codegen.simd_kernel_features import SIMDKernelFeatures
|
||||
from .codegen.triton import TritonKernel
|
||||
from .ir import ChoiceCaller
|
||||
from .kernel_template_choice import KernelTemplateChoice
|
||||
|
||||
from torch.utils._ordered_set import OrderedSet # isort: skip
|
||||
|
||||
|
||||
class Sortable(typing.Protocol):
|
||||
"""Anything that can be used as a list.sort() key (int/tuple/etc)"""
|
||||
|
||||
def __lt__(self, other: typing.Self) -> bool: ...
|
||||
|
||||
|
||||
@dataclasses.dataclass
|
||||
class FusionScore:
|
||||
template_score: int
|
||||
node_type_score: bool
|
||||
memory_score: int
|
||||
buffer_overlap_score: int
|
||||
proximity_score: int
|
||||
|
||||
def __lt__(self, other):
|
||||
"""
|
||||
node_type_score has higher priority than memory_score unless
|
||||
the memory_score differs too much.
|
||||
|
||||
buffer_overlap_score is prioritized below memory_score so that
|
||||
strict global memory savings (exact dep matches) are preferred
|
||||
over buffer overlap scoring (same buffer, different indexing).
|
||||
"""
|
||||
threshold = 16
|
||||
if self.template_score != other.template_score:
|
||||
return self.template_score < other.template_score
|
||||
|
||||
if (
|
||||
max(self.memory_score, other.memory_score)
|
||||
> min(self.memory_score, other.memory_score) * threshold
|
||||
):
|
||||
return self.memory_score < other.memory_score
|
||||
|
||||
return (
|
||||
self.node_type_score,
|
||||
self.memory_score,
|
||||
self.buffer_overlap_score,
|
||||
self.proximity_score,
|
||||
) < (
|
||||
other.node_type_score,
|
||||
other.memory_score,
|
||||
other.buffer_overlap_score,
|
||||
other.proximity_score,
|
||||
)
|
||||
|
||||
|
||||
class InductorChoices:
|
||||
"""
|
||||
This class contains a collection of default heuristics that affect performance of our generated
|
||||
code. We try to not put correctness requirements in this file.
|
||||
|
||||
You can override the choices made here by doing:
|
||||
|
||||
class MyHeuristics(InductorChoices):
|
||||
...
|
||||
|
||||
torch._inductor.virtualized.V.set_choices_handler(MyHeuristics())
|
||||
|
||||
Subclasses used with inductor_choices_class must implement uuid() for
|
||||
cache key computation.
|
||||
"""
|
||||
|
||||
def get_config_heuristics(
|
||||
self, device_type: str | None = "cuda"
|
||||
) -> BaseConfigHeuristic:
|
||||
if device_type == "cuda":
|
||||
if torch.version.hip is None:
|
||||
return CUDAConfigHeuristic()
|
||||
else:
|
||||
return ROCmConfigHeuristic()
|
||||
elif device_type == "xpu":
|
||||
return XPUConfigHeuristic()
|
||||
elif device_type == "cpu":
|
||||
return CPUConfigHeuristic()
|
||||
elif device_type == "mtia":
|
||||
return MTIAConfigHeuristic()
|
||||
else:
|
||||
return BaseConfigHeuristic()
|
||||
|
||||
# Conv configs
|
||||
def get_conv_configs(
|
||||
self, device_type: str | None = "cuda"
|
||||
) -> partial[Generator[TritonConfig, None, None]]:
|
||||
conv_heuristics = self.get_config_heuristics(device_type)
|
||||
return conv_heuristics.get_conv_configs()
|
||||
|
||||
def get_depthwise_conv_configs(self, device_type: str | None = "cuda") -> list[Any]:
|
||||
heuristics = self.get_config_heuristics(device_type)
|
||||
return heuristics.get_depthwise_conv_configs()
|
||||
|
||||
# Flex attention configs
|
||||
# TODO(coconutruben): break out flexattention/decode configs into the new retrieval mechanism
|
||||
def get_flex_attention_fwd_configs(
|
||||
self, head_dim: int, dtype: torch.dtype, device_type: str | None = "cuda"
|
||||
) -> list[Any]:
|
||||
flex_heuristics = self.get_config_heuristics(device_type)
|
||||
return flex_heuristics.get_flex_attn_fwd_configs(head_dim, dtype)
|
||||
|
||||
def get_flex_attention_bwd_configs(
|
||||
self, head_dim: int, dtype: torch.dtype, device_type: str | None = "cuda"
|
||||
) -> list[Any]:
|
||||
flex_heuristics = self.get_config_heuristics(device_type)
|
||||
return flex_heuristics.get_flex_attn_bwd_configs(head_dim, dtype)
|
||||
|
||||
def get_flex_decode_configs(
|
||||
self, head_dim: int, dtype: torch.dtype, device_type: str | None = "cuda"
|
||||
) -> list[Any]:
|
||||
flex_heuristics = self.get_config_heuristics(device_type)
|
||||
return flex_heuristics.get_flex_decode_configs(head_dim, dtype)
|
||||
|
||||
def _finalize_template_configs(
|
||||
self,
|
||||
template_choices: dict[str, Generator[KernelTemplateChoice, None, None]],
|
||||
kernel_inputs: KernelInputs,
|
||||
templates: list[KernelTemplate | ExternKernelChoice],
|
||||
op_name: str,
|
||||
kwarg_overrides: dict[str, dict[str, Any]] | None = None,
|
||||
) -> list[KernelTemplateChoice]:
|
||||
"""
|
||||
This method can be subclassed to perform any override/modification of the choices.
|
||||
The incoming parameters are cheap (generators), so you can do any overrides without
|
||||
incurring too much cost. Override this method to customize the kernel template choices
|
||||
before they are converted to ChoiceCaller objects, which is expensive on template codegen.
|
||||
|
||||
The full list of arguments are here to facilitate any overrides you may want to do,
|
||||
as they can be used to start from scratch for each template if so desired.
|
||||
|
||||
Args:
|
||||
template_choices: Dictionary mapping template UIDs to generators of KernelTemplateChoice objects
|
||||
kernel_inputs: MMKernelInputs containing input tensor nodes and matrix indices
|
||||
templates: List of template objects (KernelTemplate or ExternKernelChoice) in use
|
||||
op_name: Operation name (e.g., "bmm", "baddbmm", "addmm")
|
||||
kwarg_overrides: Optional dict of kwargs to override for each template heuristic
|
||||
|
||||
Returns:
|
||||
Flattened list of KernelTemplateChoice objects across all templates
|
||||
"""
|
||||
choices: list[KernelTemplateChoice] = []
|
||||
for choice_gen in template_choices.values():
|
||||
choices.extend(choice_gen)
|
||||
return choices
|
||||
|
||||
def get_ktc(
|
||||
self,
|
||||
kernel_inputs: KernelInputs,
|
||||
template: KernelTemplate | ExternKernelChoice,
|
||||
op_name: str,
|
||||
kwarg_overrides: dict[str, Any] | None = None,
|
||||
) -> Generator[KernelTemplateChoice, None, None]:
|
||||
"""
|
||||
Utility to get the KernelTemplateChoice generator for a specific input.
|
||||
|
||||
This is a per template/op call, whereas get_template_configs is an op wide call (all templates).
|
||||
Consider when overriding/using at which level you need to make decisions
|
||||
"""
|
||||
# Extract device_type from kernel_inputs
|
||||
device_type = kernel_inputs.device_type
|
||||
assert device_type is not None, "get_ktc requires a valid device type"
|
||||
# Extract template_name from the template object
|
||||
template_name = template.uid
|
||||
|
||||
# Get the appropriate template-specific heuristic
|
||||
heuristic = get_template_heuristic(template_name, device_type, op_name)
|
||||
cs = heuristic.get_template_configs(
|
||||
kernel_inputs,
|
||||
op_name,
|
||||
)
|
||||
# adjust the kernel inputs to the template-specific heuristic, if needed
|
||||
# default here is to just return the kernel_inputs as is
|
||||
inputs_val = heuristic.adjust_kernel_inputs(kernel_inputs, op_name)
|
||||
extra_kwargs = heuristic.get_extra_kwargs(kernel_inputs, op_name)
|
||||
# Create KernelTemplateChoice generator using the moved function
|
||||
overrides = kwarg_overrides or {}
|
||||
return make_ktc_generator(
|
||||
template=template,
|
||||
cs=cs,
|
||||
extra_kwargs=extra_kwargs,
|
||||
overrides=overrides,
|
||||
layout=kernel_inputs.output_layout(),
|
||||
inputs=inputs_val,
|
||||
)
|
||||
|
||||
def _need_to_fix_layout(
|
||||
self,
|
||||
adjusted_choices: list[KernelTemplateChoice],
|
||||
op_name: str,
|
||||
) -> bool:
|
||||
"""
|
||||
Check if we need to fix the layout instead of keeping it flexible
|
||||
|
||||
Args:
|
||||
ktc: KernelTemplateChoice object
|
||||
|
||||
Returns:
|
||||
True if we need to fix the layout, False otherwise
|
||||
"""
|
||||
# TODO: debug and fix
|
||||
# NOTE: on mps, we see issues with flexible layouts on baddmm. This check just makes sure
|
||||
# that for mps, everything stays as it was before this optimization
|
||||
if len(adjusted_choices) > 0:
|
||||
if adjusted_choices[0].inputs.device_type == "mps" and op_name not in [
|
||||
"mm",
|
||||
"addmm",
|
||||
]:
|
||||
return True
|
||||
|
||||
# Since the following backends are not using get_mm_configs yet through the singular call,
|
||||
if not (config.max_autotune or config.max_autotune_gemm):
|
||||
# no danger of using other backends than ATEN
|
||||
if not config.max_autotune_allow_flexible_layouts and op_name not in [
|
||||
# The historical implementation for mm and addmm allowed had flexible layouts in the
|
||||
# not max-autotune world
|
||||
"mm",
|
||||
"addmm",
|
||||
]:
|
||||
# TODO: deprecate this by migrating users to the new behavior
|
||||
return True
|
||||
return False
|
||||
|
||||
if not config.max_autotune_allow_flexible_layouts:
|
||||
# we always need to fix the layout
|
||||
return True
|
||||
|
||||
# Since the following backends are not using get_template_configs yet through the singular call,
|
||||
# we don't know if they are a valid choice or not. Instead, just skip the optimization
|
||||
# defensively.
|
||||
# TODO(coconutruben): remove this once CPP,CK,CUTLASS are supported
|
||||
if _use_autotune_backend("CUTLASS"):
|
||||
return True
|
||||
if _use_autotune_backend("CK") or _use_autotune_backend("CKTILE"):
|
||||
return True
|
||||
if _use_autotune_backend("CPP"):
|
||||
return True
|
||||
return any(
|
||||
not isinstance(ktc.template, ExternKernelChoice) for ktc in adjusted_choices
|
||||
)
|
||||
|
||||
def get_template_configs(
|
||||
self,
|
||||
kernel_inputs: KernelInputs,
|
||||
templates: list[KernelTemplate | ExternKernelChoice],
|
||||
op_name: str,
|
||||
kwarg_overrides: dict[str, dict[str, Any]] | None = None,
|
||||
) -> list[ChoiceCaller]:
|
||||
"""
|
||||
Get list of ChoiceCallers for MM templates using template-specific heuristics.
|
||||
|
||||
Args:
|
||||
kernel_inputs: MMKernelInputs containing input tensor nodes and matrix indices
|
||||
layout: Output layout
|
||||
templates: List of template objects (KernelTemplate or ExternKernelChoice)
|
||||
op_name: Operation name (e.g., "bmm", "baddbmm", "addmm", "mm_plus_mm")
|
||||
kwarg_overrides: Optional dict of kwargs to override for each template heuristic,
|
||||
indexed by template.uid. These only override the per config kwargs, not the extra kwargs
|
||||
Returns:
|
||||
List of ChoiceCaller objects from the templates
|
||||
"""
|
||||
if kwarg_overrides is None:
|
||||
kwarg_overrides = {}
|
||||
input_tensors = kernel_inputs.nodes()
|
||||
if len(input_tensors) < 2:
|
||||
raise ValueError(f"Need at least 2 input tensors, got {len(input_tensors)}")
|
||||
layout = kernel_inputs.output_layout()
|
||||
# First pass: Create dict of template.uid to generator of KernelTemplateChoice objects
|
||||
template_choices = {}
|
||||
for template in templates:
|
||||
template_choices[template.uid] = self.get_ktc(
|
||||
kernel_inputs,
|
||||
template,
|
||||
op_name,
|
||||
kwarg_overrides.get(template.uid, {}),
|
||||
)
|
||||
|
||||
# Second pass: Adjust the template choices
|
||||
adjusted_choices = self._finalize_template_configs(
|
||||
template_choices,
|
||||
kernel_inputs,
|
||||
templates,
|
||||
op_name,
|
||||
kwarg_overrides,
|
||||
)
|
||||
# Layout optimization: if all choices are ExternKernelChoice and layout is FixedLayout, convert to FlexibleLayout
|
||||
if self._need_to_fix_layout(adjusted_choices, op_name):
|
||||
layout = kernel_inputs.output_layout(flexible=False)
|
||||
for ktc in adjusted_choices:
|
||||
ktc.layout = layout
|
||||
# for good measure, delete the cached ChoiceCaller from the ktc if it existed.
|
||||
# ExternKernelChoice are cheap to generate
|
||||
if hasattr(ktc, "_choice"):
|
||||
del ktc._choice
|
||||
# Third pass: Convert to ChoiceCaller objects
|
||||
return [ktc.choice for ktc in adjusted_choices if ktc.choice is not None]
|
||||
|
||||
def triton_kernel_kwargs(
|
||||
self,
|
||||
kernel_cls: type[TritonKernel],
|
||||
features: SIMDKernelFeatures,
|
||||
groups: list[sympy.Expr],
|
||||
kernel_kwargs: dict[str, Any],
|
||||
) -> dict[str, Any]:
|
||||
"""Hook to change the kwargs passed to TritonKernel, used to apply fixed configurations"""
|
||||
return kernel_kwargs
|
||||
|
||||
def override_best_choice(
|
||||
self,
|
||||
best_choice: ChoiceCaller,
|
||||
timings: dict[ChoiceCaller, float],
|
||||
) -> ChoiceCaller:
|
||||
"""Hook to override the autotuning best choice after benchmarking."""
|
||||
return best_choice
|
||||
|
||||
def customize_fused_kernel_name(self, fused_name: str, src_code: str) -> str:
|
||||
"""Hook to transform fused kernel names during codegen"""
|
||||
return fused_name
|
||||
|
||||
@staticmethod
|
||||
def should_use_cooperative_reduction(features: SIMDKernelFeatures) -> bool:
|
||||
"""Heuristic to decide if a cooperative reduction should be used."""
|
||||
if config.triton.force_cooperative_reductions:
|
||||
return True
|
||||
if (
|
||||
not config.triton.cooperative_reductions
|
||||
or V.graph.get_current_device_or_throw().type == "cpu"
|
||||
):
|
||||
return False
|
||||
|
||||
xhint = V.graph.sizevars.optimization_hint(features.numel, fallback=2)
|
||||
if xhint <= 8:
|
||||
threshold = 32768 * xhint
|
||||
elif xhint <= 16:
|
||||
threshold = 2097152
|
||||
else:
|
||||
return False
|
||||
# TODO(jansel): should this default on for dynamic shapes?
|
||||
# TODO(laith) What if hint(features.reduction_numel) >= threshold ?
|
||||
# shall we compare hints instead
|
||||
return V.graph.sizevars.statically_known_geq(
|
||||
features.reduction_numel, threshold
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def should_use_persistent_reduction(
|
||||
features: SIMDKernelFeatures, cooperative_reduction: bool
|
||||
) -> bool:
|
||||
"""
|
||||
Heuristic to decide if a persistent reduction should be used.
|
||||
"""
|
||||
if not config.triton.persistent_reductions:
|
||||
return False
|
||||
threshold = {
|
||||
ReductionHint.INNER: 1024,
|
||||
}.get(features.get_reduction_hint(), 64)
|
||||
|
||||
if features.get_reduction_hint() not in (
|
||||
ReductionHint.INNER,
|
||||
ReductionHint.OUTER_TINY,
|
||||
):
|
||||
bounds = bound_sympy(features.reduction_numel)
|
||||
lower = bounds.lower
|
||||
upper = bounds.upper
|
||||
|
||||
if not all(
|
||||
(
|
||||
(isinstance(bound, int) or bound.is_constant())
|
||||
and not torch.utils._sympy.numbers.is_infinite(bound)
|
||||
)
|
||||
for bound in (lower, upper)
|
||||
):
|
||||
return False
|
||||
|
||||
lower = next_power_of_2(int(lower))
|
||||
upper = next_power_of_2(int(upper))
|
||||
|
||||
# If we are are coalescing on xblock (not ReductionHint.INNER) and this is not a tiny kernel
|
||||
# (not ReductionHint.OUTER_TINY), do not use persistent reduction if it induces tile
|
||||
# quantization. Persistent reduction forces rblock == rnumel, if the bounds between lower
|
||||
# and upper are large, for the lower values we will be masking off large % of read/writes,
|
||||
# when we could expand the coalescing xblock instead.
|
||||
if lower != upper:
|
||||
return False
|
||||
|
||||
if cooperative_reduction:
|
||||
# The RSPLIT of cooperative reductions means each thread block is operating on fewer elements
|
||||
# The default fallback will be used if optimizations hint is not provided. The default fallback
|
||||
# is >> 32.
|
||||
threshold *= 32 // min(
|
||||
V.graph.sizevars.optimization_hint(features.numel), 32
|
||||
)
|
||||
|
||||
# If multi_kernel is enabled, we do more aggressive persistent reduction.
|
||||
# This may result in some persistent reductions slower than the
|
||||
# corresponding non-persistent reductions. MultiKernel will do benchmarking
|
||||
# to pick the faster one.
|
||||
if config.triton.multi_kernel:
|
||||
threshold *= 16
|
||||
|
||||
return V.graph.sizevars.statically_known_leq(
|
||||
features.reduction_numel, threshold
|
||||
) # type: ignore[arg-types]
|
||||
|
||||
@staticmethod
|
||||
def reduction_split_factor(
|
||||
device: torch.device,
|
||||
reduction_numel_hint: int,
|
||||
numel_hint: int,
|
||||
inner_reduction: bool,
|
||||
) -> int:
|
||||
"""Heuristic to decide the RSPLIT used for split reductions.
|
||||
When a reduction has a small number of outputs there is not enough parallelism,
|
||||
so we will do the reduction in two phases."""
|
||||
props = DeviceProperties.create(device)
|
||||
num_sm = props.multi_processor_count
|
||||
warp_size = props.warp_size if props.warp_size is not None else 32
|
||||
max_threads_per_sm = (
|
||||
props.max_threads_per_multi_processor
|
||||
if props.max_threads_per_multi_processor is not None
|
||||
else 2048
|
||||
)
|
||||
min_elements_per_thread = warp_size
|
||||
max_elements_per_thread = 512
|
||||
threads_per_sm = max_threads_per_sm
|
||||
min_elements_per_device = min_elements_per_thread * num_sm * threads_per_sm
|
||||
max_elements_per_device = max_elements_per_thread * num_sm * threads_per_sm
|
||||
num_warps = 8
|
||||
num_threads = warp_size * num_warps
|
||||
|
||||
if inner_reduction:
|
||||
# do heuristics that's close to eager mode for split inner reduction
|
||||
# we leak reduction autotune configs here, and will need to refactor to avoid this later
|
||||
if numel_hint >= 2 * num_sm: # don't split if there are enough outputs
|
||||
return 1
|
||||
if reduction_numel_hint <= 8192:
|
||||
return 1
|
||||
if reduction_numel_hint * numel_hint <= min_elements_per_device:
|
||||
split_size = min_elements_per_thread
|
||||
elif reduction_numel_hint * numel_hint < max_elements_per_device:
|
||||
target_blocks = num_sm * threads_per_sm // (2 * num_threads)
|
||||
blocks_per_output = (target_blocks + numel_hint - 1) // numel_hint
|
||||
tmp_split_size = (
|
||||
reduction_numel_hint + num_threads * blocks_per_output - 1
|
||||
) // (num_threads * blocks_per_output)
|
||||
divisors = sympy.divisors(reduction_numel_hint)
|
||||
closest = min(divisors, key=lambda x: abs(x - tmp_split_size))
|
||||
if abs(closest - tmp_split_size) < 30:
|
||||
# prefer even splits, but never smalle than min_elements_per_thread
|
||||
split_size = max(closest, min_elements_per_thread)
|
||||
else:
|
||||
split_size = tmp_split_size
|
||||
else:
|
||||
divisors = sympy.divisors(reduction_numel_hint)
|
||||
closest = min(divisors, key=lambda x: abs(x - max_elements_per_thread))
|
||||
if abs(closest - max_elements_per_thread) < 50:
|
||||
# prefer even splits
|
||||
split_size = closest
|
||||
else:
|
||||
split_size = max_elements_per_thread
|
||||
return (reduction_numel_hint + split_size * num_threads - 1) // (
|
||||
split_size * num_threads
|
||||
)
|
||||
else:
|
||||
# TODO the best heuristic currently has XBLOCK (corresponding to numel_hint) 128
|
||||
# extend to even smaller number of outputs
|
||||
rvals_per_thread = 4 # comes from heuristics, refactor to not leak here
|
||||
xvals_per_block = 128
|
||||
xblocks = (numel_hint + xvals_per_block - 1) // xvals_per_block
|
||||
if reduction_numel_hint * numel_hint < min_elements_per_device:
|
||||
split_size = min_elements_per_thread
|
||||
elif reduction_numel_hint * numel_hint < max_elements_per_device:
|
||||
target_blocks = num_sm * threads_per_sm // (num_threads)
|
||||
target_blocks = (target_blocks + xblocks - 1) // xblocks
|
||||
tmp_split_size = (
|
||||
reduction_numel_hint + rvals_per_thread * target_blocks - 1
|
||||
) // (rvals_per_thread * target_blocks)
|
||||
divisors = sympy.divisors(reduction_numel_hint)
|
||||
closest = min(divisors, key=lambda x: abs(x - tmp_split_size))
|
||||
if abs(tmp_split_size - closest) < 20:
|
||||
split_size = max(closest, min_elements_per_thread)
|
||||
else:
|
||||
split_size = tmp_split_size
|
||||
else:
|
||||
divisors = sympy.divisors(reduction_numel_hint)
|
||||
closest = min(divisors, key=lambda x: abs(x - max_elements_per_thread))
|
||||
if abs(closest - max_elements_per_thread) < 50:
|
||||
# prefer even splits
|
||||
split_size = closest
|
||||
else:
|
||||
split_size = max_elements_per_thread
|
||||
|
||||
return (reduction_numel_hint + rvals_per_thread * split_size - 1) // (
|
||||
rvals_per_thread * split_size
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def can_fuse(
|
||||
scheduler: Scheduler,
|
||||
node1: BaseSchedulerNode,
|
||||
node2: BaseSchedulerNode,
|
||||
shared_data_score: int,
|
||||
) -> bool:
|
||||
"""
|
||||
Heuristics to prevent fusion applied to both horizontal and vertical fusions. Heuristics here should not
|
||||
be needed for correctness and tweaking them may yield additional performance.
|
||||
|
||||
See also some related heuristics that can be changed via config:
|
||||
- config.triton.tiling_prevents_pointwise_fusion
|
||||
- config.triton.tiling_prevents_reduction_fusion
|
||||
- config.aggressive_fusion (will cause this function to be called more times)
|
||||
"""
|
||||
if shared_data_score == 0 and (
|
||||
not config.aggressive_fusion or node1.is_reduction() or node2.is_reduction()
|
||||
):
|
||||
if is_metric_table_enabled("fusion_failure_due_to_indexing_mismatch"):
|
||||
common_buf_names: OrderedSet[str] = (
|
||||
node1.read_writes.buffer_names() & node2.read_writes.buffer_names()
|
||||
)
|
||||
if len(common_buf_names) > 0:
|
||||
get_metric_table("fusion_failure_due_to_indexing_mismatch").add_row(
|
||||
# pyrefly: ignore [bad-argument-type]
|
||||
lambda: {
|
||||
"pre_grad_graph_id": V.graph.graph_id,
|
||||
"post_grad_graph_id": V.graph.post_grad_graph_id,
|
||||
"node1_name": node1.get_name(),
|
||||
"node2_name": node2.get_name(),
|
||||
"node1_debug_str": write_text(node1.debug_str()),
|
||||
"node2_debug_str": write_text(node2.debug_str()),
|
||||
"common_buffer_names": list(common_buf_names), # type: ignore[dict-item]
|
||||
"failure_reason": scheduler.decide_fusion_fail_reason(
|
||||
node1, node2, common_buf_names
|
||||
),
|
||||
}
|
||||
)
|
||||
|
||||
WhyNoFuse(node1, node2)("no shared data due to indexing mismatch")
|
||||
return False
|
||||
WhyNoFuse(node1, node2)("no shared data")
|
||||
return False # heuristic not needed for correctness
|
||||
|
||||
if (
|
||||
not node1.is_foreach()
|
||||
and not node2.is_foreach()
|
||||
and len(node1.get_nodes()) + len(node2.get_nodes()) > config.max_fusion_size
|
||||
):
|
||||
WhyNoFuse(node1, node2)("exceeds max fusion")
|
||||
return False # heuristic not needed for correctness
|
||||
|
||||
if scheduler.can_fusion_increase_peak_memory(node1, node2):
|
||||
WhyNoFuse(node1, node2)("Fusion will increase peak memory")
|
||||
return False
|
||||
|
||||
if (
|
||||
config.max_fusion_unique_io_buffers is not None
|
||||
and scheduler.fusion_prevent_too_many_reads_and_writes(
|
||||
node1,
|
||||
node2,
|
||||
config.max_fusion_unique_io_buffers,
|
||||
)
|
||||
):
|
||||
WhyNoFuse(node1, node2)("fusion_prevent_too_many_reads_and_writes")
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
@staticmethod
|
||||
def can_fuse_vertical(
|
||||
scheduler: Scheduler,
|
||||
node1: BaseSchedulerNode,
|
||||
node2: BaseSchedulerNode,
|
||||
shared_data_score: int,
|
||||
) -> bool:
|
||||
"""Hook for heuristics to prevent vertical (producer/consumer) fusions"""
|
||||
return True
|
||||
|
||||
@staticmethod
|
||||
def can_fuse_horizontal(
|
||||
scheduler: Scheduler,
|
||||
node1: BaseSchedulerNode,
|
||||
node2: BaseSchedulerNode,
|
||||
shared_data_score: int,
|
||||
) -> bool:
|
||||
"""Hook for heuristics to prevent horizontal (consumer/consumer) fusions"""
|
||||
if MixOrderReduction.can_fuse(node1, node2):
|
||||
# For mix order reduction, we disregard shared data or
|
||||
# distance.
|
||||
return True
|
||||
if shared_data_score < config.score_fusion_memory_threshold:
|
||||
WhyNoFuse(node1, node2)("score_fusion_memory_threshold")
|
||||
return False
|
||||
if scheduler.are_long_distant_nodes(node1, node2):
|
||||
WhyNoFuse(node1, node2)(
|
||||
"Nodes are too far away. Fusing them may increase peak memory."
|
||||
)
|
||||
return False
|
||||
return True
|
||||
|
||||
@staticmethod
|
||||
def score_fusion(
|
||||
scheduler: Scheduler,
|
||||
node1: BaseSchedulerNode,
|
||||
node2: BaseSchedulerNode,
|
||||
) -> Sortable:
|
||||
"""
|
||||
Assign a score (higher comes first) to the fusion of node1 and node2.
|
||||
When different fusions conflict with each other, this is the way we
|
||||
decide what order to run them in.
|
||||
|
||||
Our current score is based on:
|
||||
- The type of fusion (template/reduction/etc)
|
||||
- Estimate of the saved memory operations
|
||||
- Fusions closer together in original graph order
|
||||
"""
|
||||
|
||||
memory_score, buffer_overlap_score, is_mix_order_reduction = typing.cast(
|
||||
tuple[int, int, bool],
|
||||
scheduler.score_fusion_memory(
|
||||
node1, node2, return_is_mix_order_reduction=True
|
||||
),
|
||||
)
|
||||
proximity_score = -max(
|
||||
abs(node1.min_order - node2.max_order),
|
||||
abs(node2.min_order - node1.max_order),
|
||||
)
|
||||
|
||||
# prologue fusion always last
|
||||
if node2.is_template():
|
||||
template_score = 0
|
||||
else:
|
||||
template_score = 1 + (
|
||||
(node1.is_template() == config.epilogue_fusion_first)
|
||||
and memory_score > 0
|
||||
)
|
||||
|
||||
type_score = node1.is_reduction() == node2.is_reduction() and memory_score > 0
|
||||
|
||||
return FusionScore(
|
||||
template_score,
|
||||
type_score,
|
||||
memory_score,
|
||||
buffer_overlap_score,
|
||||
proximity_score,
|
||||
)
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,36 @@
|
||||
import re
|
||||
|
||||
import torch
|
||||
|
||||
|
||||
# It is not a good idea to directly apply hipify_torch to codegen, which will be vulnerable to cases like:
|
||||
# "...
|
||||
# from ..codecache import CudaKernelParamCache
|
||||
# ..."
|
||||
# In such cases, we do not need to hipify_torch the original class/file name in codegen/codecache
|
||||
|
||||
|
||||
def maybe_hipify_code_wrapper(source_codes: str, force_hipify: bool = False) -> str:
|
||||
if torch.version.hip is None and not force_hipify:
|
||||
return source_codes
|
||||
|
||||
try:
|
||||
from torch.utils.hipify.hipify_python import PYTORCH_MAP, PYTORCH_TRIE
|
||||
except ImportError:
|
||||
# hipify not available for non-AMD builds
|
||||
return source_codes
|
||||
|
||||
def c2_repl(m: re.Match[str]) -> object:
|
||||
return PYTORCH_MAP[m.group(0)]
|
||||
|
||||
# We need to redefine RE_PYTORCH_PREPROCESSOR here since in hipify_torch,
|
||||
# it will apply positive lookbehind (?<=\W) to the pattern to avoid matching
|
||||
# keyword at the beginning of code line. However, this can happen in codegen,
|
||||
# which will cause the pattern to not match.
|
||||
|
||||
# Note that lookahead (?=\W) is still needed to keep hipification idomponent, for example
|
||||
# we need to skip replacing "getStreamFromExternal" in "getStreamFromExternalMasqueradingAsCUDA"
|
||||
RE_PYTORCH_PREPROCESSOR = re.compile(rf"({PYTORCH_TRIE.export_to_regex()})(?=\W)")
|
||||
|
||||
source_codes = RE_PYTORCH_PREPROCESSOR.sub(c2_repl, source_codes) # type: ignore[arg-type]
|
||||
return source_codes
|
||||
+488
@@ -0,0 +1,488 @@
|
||||
// Definition of AOTI runtime interface functions
|
||||
|
||||
#include <torch/csrc/inductor/aoti_runtime/interface.h>
|
||||
#include <torch/csrc/inductor/aoti_runtime/model_container.h>
|
||||
|
||||
#include <iostream>
|
||||
#include <vector>
|
||||
|
||||
#define CONVERT_EXCEPTION_TO_ERROR_CODE(...) \
|
||||
try { \
|
||||
__VA_ARGS__ \
|
||||
} catch (const std::exception& e) { \
|
||||
std::cerr << "Error: " << e.what() << '\n'; \
|
||||
return AOTI_RUNTIME_FAILURE; \
|
||||
} catch (...) { \
|
||||
std::cerr << "Unknown exception occurred.\n"; \
|
||||
return AOTI_RUNTIME_FAILURE; \
|
||||
} \
|
||||
return AOTI_RUNTIME_SUCCESS;
|
||||
|
||||
#define AOTI_VECTOR_SIZE_CHECK(actual_size, expected_size, name) \
|
||||
do { \
|
||||
AOTI_RUNTIME_CHECK( \
|
||||
actual_size == expected_size, \
|
||||
"expected " + std::string(name) + " vector size to be " + \
|
||||
std::to_string(expected_size) + ", but got " + \
|
||||
std::to_string(actual_size)); \
|
||||
} while (0)
|
||||
|
||||
// AOTInductor uses at::addmm_out, which doesn't supports
|
||||
// arguments that requires gradient. For this reason, we
|
||||
// enforce no_grad context for run APIs.
|
||||
//
|
||||
// A RAII, thread local (!) guard that enables or disables grad mode upon
|
||||
// construction, and sets it back to the original value upon destruction.
|
||||
struct AOTINoGradGuard {
|
||||
AOTINoGradGuard() {
|
||||
aoti_torch_grad_mode_set_enabled(false);
|
||||
}
|
||||
AOTINoGradGuard(const AOTINoGradGuard&) = delete;
|
||||
AOTINoGradGuard(AOTINoGradGuard&&) noexcept = delete;
|
||||
~AOTINoGradGuard() {
|
||||
aoti_torch_grad_mode_set_enabled(prev_mode);
|
||||
}
|
||||
AOTINoGradGuard& operator=(const AOTINoGradGuard&) = delete;
|
||||
AOTINoGradGuard& operator=(AOTINoGradGuard&&) noexcept = delete;
|
||||
bool prev_mode{aoti_torch_grad_mode_is_enabled()};
|
||||
};
|
||||
|
||||
extern "C" {
|
||||
|
||||
AOTIRuntimeError AOTInductorModelContainerCreate(
|
||||
AOTInductorModelContainerHandle* container_handle,
|
||||
size_t num_models,
|
||||
bool is_cpu,
|
||||
const char* cubin_dir) {
|
||||
return AOTInductorModelContainerCreateWithDevice(
|
||||
container_handle,
|
||||
num_models,
|
||||
is_cpu ? "cpu" : "cuda",
|
||||
cubin_dir);
|
||||
}
|
||||
|
||||
AOTIRuntimeError AOTInductorModelContainerCreateWithDevice(
|
||||
AOTInductorModelContainerHandle* container_handle,
|
||||
size_t num_models,
|
||||
const char* device_str,
|
||||
const char* cubin_dir) {
|
||||
|
||||
if (num_models == 0) {
|
||||
std::cerr << "Error: num_models must be positive, but got 0\n";
|
||||
return AOTI_RUNTIME_FAILURE;
|
||||
}
|
||||
CONVERT_EXCEPTION_TO_ERROR_CODE({
|
||||
std::optional<std::string> cubin_dir_opt;
|
||||
if (cubin_dir != nullptr) {
|
||||
cubin_dir_opt.emplace(cubin_dir);
|
||||
}
|
||||
auto* container = new torch::aot_inductor::AOTInductorModelContainer(
|
||||
num_models, std::string(device_str), cubin_dir_opt);
|
||||
*container_handle =
|
||||
reinterpret_cast<AOTInductorModelContainerHandle>(container);
|
||||
})
|
||||
}
|
||||
|
||||
|
||||
AOTIRuntimeError AOTInductorModelContainerDelete(
|
||||
AOTInductorModelContainerHandle container_handle) {
|
||||
CONVERT_EXCEPTION_TO_ERROR_CODE({
|
||||
auto* container =
|
||||
reinterpret_cast<torch::aot_inductor::AOTInductorModelContainer*>(
|
||||
container_handle);
|
||||
delete container;
|
||||
});
|
||||
}
|
||||
|
||||
AOTIRuntimeError AOTInductorModelContainerRun(
|
||||
AOTInductorModelContainerHandle container_handle,
|
||||
AtenTensorHandle* input_handles, // array of input AtenTensorHandle; handles
|
||||
// are stolen; the array itself is borrowed
|
||||
size_t num_inputs,
|
||||
AtenTensorHandle*
|
||||
output_handles, // array for writing output AtenTensorHandle; handles
|
||||
// will be stolen by the caller; the array itself is
|
||||
// borrowed
|
||||
size_t num_outputs,
|
||||
AOTInductorStreamHandle stream_handle,
|
||||
AOTIProxyExecutorHandle proxy_executor_handle) {
|
||||
auto* container =
|
||||
reinterpret_cast<torch::aot_inductor::AOTInductorModelContainer*>(
|
||||
container_handle);
|
||||
AOTI_VECTOR_SIZE_CHECK(num_inputs, container->num_inputs(), "inputs");
|
||||
AOTI_VECTOR_SIZE_CHECK(num_outputs, container->num_outputs(), "outputs");
|
||||
|
||||
auto stream =
|
||||
reinterpret_cast<torch::aot_inductor::DeviceStreamType>(stream_handle);
|
||||
CONVERT_EXCEPTION_TO_ERROR_CODE({
|
||||
AOTINoGradGuard guard;
|
||||
container->run(
|
||||
input_handles, output_handles, stream, proxy_executor_handle);
|
||||
})
|
||||
}
|
||||
|
||||
AOTIRuntimeError AOTInductorModelContainerRunSingleThreaded(
|
||||
AOTInductorModelContainerHandle container_handle,
|
||||
AtenTensorHandle* input_handles, // array of input AtenTensorHandle; handles
|
||||
// are stolen; the array itself is borrowed
|
||||
size_t num_inputs,
|
||||
AtenTensorHandle*
|
||||
output_handles, // array for writing output AtenTensorHandle; handles
|
||||
// will be stolen by the caller; the array itself is
|
||||
// borrowed
|
||||
size_t num_outputs,
|
||||
AOTInductorStreamHandle stream_handle,
|
||||
AOTIProxyExecutorHandle proxy_executor_handle) {
|
||||
auto* container =
|
||||
reinterpret_cast<torch::aot_inductor::AOTInductorModelContainer*>(
|
||||
container_handle);
|
||||
AOTI_VECTOR_SIZE_CHECK(num_inputs, container->num_inputs(), "inputs");
|
||||
AOTI_VECTOR_SIZE_CHECK(num_outputs, container->num_outputs(), "outputs");
|
||||
|
||||
auto stream =
|
||||
reinterpret_cast<torch::aot_inductor::DeviceStreamType>(stream_handle);
|
||||
CONVERT_EXCEPTION_TO_ERROR_CODE({
|
||||
AOTINoGradGuard guard;
|
||||
container->run_single_threaded(
|
||||
input_handles, output_handles, stream, proxy_executor_handle);
|
||||
})
|
||||
}
|
||||
|
||||
AOTIRuntimeError AOTInductorModelContainerGetNumConstants(
|
||||
AOTInductorModelContainerHandle container_handle,
|
||||
size_t* num_constants) {
|
||||
auto* container =
|
||||
reinterpret_cast<torch::aot_inductor::AOTInductorModelContainer*>(
|
||||
container_handle);
|
||||
CONVERT_EXCEPTION_TO_ERROR_CODE(
|
||||
{ *num_constants = container->num_constants(); })
|
||||
}
|
||||
|
||||
AOTIRuntimeError AOTInductorModelContainerGetConstantName(
|
||||
AOTInductorModelContainerHandle container_handle,
|
||||
size_t idx,
|
||||
const char** name) {
|
||||
auto* container =
|
||||
reinterpret_cast<torch::aot_inductor::AOTInductorModelContainer*>(
|
||||
container_handle);
|
||||
CONVERT_EXCEPTION_TO_ERROR_CODE(
|
||||
{ *name = container->constant_name(idx); })
|
||||
}
|
||||
|
||||
AOTIRuntimeError AOTInductorModelContainerGetConstantOriginalFQN(
|
||||
AOTInductorModelContainerHandle container_handle,
|
||||
size_t idx,
|
||||
const char** original_fqn) {
|
||||
auto* container =
|
||||
reinterpret_cast<torch::aot_inductor::AOTInductorModelContainer*>(
|
||||
container_handle);
|
||||
CONVERT_EXCEPTION_TO_ERROR_CODE(
|
||||
{ *original_fqn = container->constant_original_fqn(idx); })
|
||||
}
|
||||
|
||||
AOTIRuntimeError AOTInductorModelContainerGetConstantFromFolded(
|
||||
AOTInductorModelContainerHandle container_handle,
|
||||
size_t idx,
|
||||
bool* from_folded) {
|
||||
auto* container =
|
||||
reinterpret_cast<torch::aot_inductor::AOTInductorModelContainer*>(container_handle);
|
||||
CONVERT_EXCEPTION_TO_ERROR_CODE({ *from_folded = container->constant_from_folded(idx); })
|
||||
}
|
||||
|
||||
AOTIRuntimeError AOTInductorModelContainerGetConstantType(
|
||||
AOTInductorModelContainerHandle container_handle,
|
||||
size_t idx,
|
||||
int32_t* type) {
|
||||
auto* container =
|
||||
reinterpret_cast<torch::aot_inductor::AOTInductorModelContainer*>(container_handle);
|
||||
CONVERT_EXCEPTION_TO_ERROR_CODE({ *type = container->constant_type(idx); })
|
||||
}
|
||||
|
||||
AOTIRuntimeError AOTInductorModelContainerGetConstantDtype(
|
||||
AOTInductorModelContainerHandle container_handle,
|
||||
size_t idx,
|
||||
int32_t* dtype) {
|
||||
auto* container =
|
||||
reinterpret_cast<torch::aot_inductor::AOTInductorModelContainer*>(
|
||||
container_handle);
|
||||
CONVERT_EXCEPTION_TO_ERROR_CODE(
|
||||
{ *dtype = container->constant_dtype(idx); })
|
||||
}
|
||||
|
||||
AOTIRuntimeError AOTInductorModelContainerGetConstantDataSize(
|
||||
AOTInductorModelContainerHandle container_handle,
|
||||
size_t idx,
|
||||
size_t* data_size) {
|
||||
auto* container =
|
||||
reinterpret_cast<torch::aot_inductor::AOTInductorModelContainer*>(
|
||||
container_handle);
|
||||
CONVERT_EXCEPTION_TO_ERROR_CODE(
|
||||
{ *data_size = container->constant_data_size(idx); })
|
||||
}
|
||||
|
||||
AOTIRuntimeError AOTInductorModelContainerExtractConstantsMap(
|
||||
AOTInductorModelContainerHandle container_handle,
|
||||
AOTInductorConstantMapHandle constant_map_handle,
|
||||
bool use_inactive) {
|
||||
auto* container =
|
||||
reinterpret_cast<torch::aot_inductor::AOTInductorModelContainer*>(
|
||||
container_handle);
|
||||
auto constants_map = reinterpret_cast<std::unordered_map<std::string, AtenTensorHandle>*>(constant_map_handle);
|
||||
CONVERT_EXCEPTION_TO_ERROR_CODE(
|
||||
{ const auto ret = container->extract_constants_map(use_inactive);
|
||||
for (const auto& pair: ret) {
|
||||
constants_map->emplace(pair.first, pair.second);
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
AOTIRuntimeError AOTInductorModelContainerUpdateUserManagedConstantBuffer(
|
||||
AOTInductorModelContainerHandle container_handle,
|
||||
AOTInductorConstantMapHandle constant_map_handle,
|
||||
bool use_inactive,
|
||||
bool validate_full_update) {
|
||||
auto* container =
|
||||
reinterpret_cast<torch::aot_inductor::AOTInductorModelContainer*>(
|
||||
container_handle);
|
||||
auto input_map = reinterpret_cast<std::unordered_map<std::string, AtenTensorHandle>*>(constant_map_handle);
|
||||
CONVERT_EXCEPTION_TO_ERROR_CODE({
|
||||
container->update_constant_buffer(
|
||||
*input_map, use_inactive, validate_full_update, /* user_managed = */ true);
|
||||
})
|
||||
}
|
||||
|
||||
AOTIRuntimeError AOTInductorModelContainerUpdateUserManagedConstantBufferPairs(
|
||||
AOTInductorModelContainerHandle container_handle,
|
||||
const AOTInductorConstantMapEntry* pairs,
|
||||
size_t num_pairs,
|
||||
bool use_inactive,
|
||||
bool validate_full_update) {
|
||||
auto* container =
|
||||
reinterpret_cast<torch::aot_inductor::AOTInductorModelContainer*>(container_handle);
|
||||
// Build a local unordered_map inside
|
||||
std::unordered_map<std::string, AtenTensorHandle> input_map;
|
||||
input_map.reserve(num_pairs);
|
||||
for (size_t i = 0; i < num_pairs; ++i) {
|
||||
input_map.emplace(pairs[i].name, pairs[i].handle);
|
||||
}
|
||||
CONVERT_EXCEPTION_TO_ERROR_CODE({
|
||||
container->update_constant_buffer(
|
||||
input_map, use_inactive, validate_full_update, /*user_managed=*/true);
|
||||
})
|
||||
}
|
||||
|
||||
AOTIRuntimeError AOTInductorModelContainerUpdateConstantBuffer(
|
||||
AOTInductorModelContainerHandle container_handle,
|
||||
AOTInductorConstantMapHandle constant_map_handle,
|
||||
bool use_inactive,
|
||||
bool validate_full_update) {
|
||||
auto* container =
|
||||
reinterpret_cast<torch::aot_inductor::AOTInductorModelContainer*>(
|
||||
container_handle);
|
||||
auto input_map = reinterpret_cast<std::unordered_map<std::string, AtenTensorHandle>*>(constant_map_handle);
|
||||
CONVERT_EXCEPTION_TO_ERROR_CODE({
|
||||
container->update_constant_buffer(
|
||||
*input_map, use_inactive, validate_full_update);
|
||||
})
|
||||
}
|
||||
|
||||
AOTIRuntimeError AOTInductorModelContainerUpdateInactiveConstantBuffer(
|
||||
AOTInductorModelContainerHandle container_handle,
|
||||
AOTInductorConstantMapHandle constant_map_handle) {
|
||||
return AOTInductorModelContainerUpdateConstantBuffer(container_handle,
|
||||
constant_map_handle,
|
||||
/*use_inactive*/ true,
|
||||
/*validate_full_update*/ true);
|
||||
}
|
||||
|
||||
AOTIRuntimeError AOTInductorModelContainerFreeInactiveConstantBuffer(
|
||||
AOTInductorModelContainerHandle container_handle) {
|
||||
auto* container =
|
||||
reinterpret_cast<torch::aot_inductor::AOTInductorModelContainer*>(
|
||||
container_handle);
|
||||
CONVERT_EXCEPTION_TO_ERROR_CODE({
|
||||
container->free_inactive_constant_buffer();
|
||||
})
|
||||
}
|
||||
|
||||
AOTIRuntimeError AOTInductorModelContainerRunConstantFolding(
|
||||
AOTInductorModelContainerHandle container_handle,
|
||||
bool use_inactive,
|
||||
AOTInductorStreamHandle stream_handle,
|
||||
AOTIProxyExecutorHandle proxy_executor_handle) {
|
||||
auto* container =
|
||||
reinterpret_cast<torch::aot_inductor::AOTInductorModelContainer*>(
|
||||
container_handle);
|
||||
auto stream =
|
||||
reinterpret_cast<torch::aot_inductor::DeviceStreamType>(stream_handle);
|
||||
CONVERT_EXCEPTION_TO_ERROR_CODE({
|
||||
AOTINoGradGuard guard;
|
||||
container->run_const_fold(use_inactive, stream, proxy_executor_handle);
|
||||
})
|
||||
}
|
||||
|
||||
AOTIRuntimeError AOTInductorModelContainerSwapConstantBuffer(
|
||||
AOTInductorModelContainerHandle container_handle) {
|
||||
auto* container =
|
||||
reinterpret_cast<torch::aot_inductor::AOTInductorModelContainer*>(
|
||||
container_handle);
|
||||
CONVERT_EXCEPTION_TO_ERROR_CODE({
|
||||
container->swap_constant_buffer();
|
||||
})
|
||||
}
|
||||
|
||||
AOTIRuntimeError AOTInductorModelContainerGetNumInputs(
|
||||
AOTInductorModelContainerHandle container_handle,
|
||||
size_t* ret_num_inputs) {
|
||||
auto* container =
|
||||
reinterpret_cast<torch::aot_inductor::AOTInductorModelContainer*>(
|
||||
container_handle);
|
||||
CONVERT_EXCEPTION_TO_ERROR_CODE(
|
||||
{ *ret_num_inputs = container->num_inputs(); })
|
||||
}
|
||||
|
||||
AOTIRuntimeError AOTInductorModelContainerGetInputName(
|
||||
AOTInductorModelContainerHandle container_handle,
|
||||
size_t input_idx,
|
||||
const char** ret_input_names) {
|
||||
auto* container =
|
||||
reinterpret_cast<torch::aot_inductor::AOTInductorModelContainer*>(
|
||||
container_handle);
|
||||
CONVERT_EXCEPTION_TO_ERROR_CODE(
|
||||
{ *ret_input_names = container->input_name(input_idx); })
|
||||
}
|
||||
|
||||
AOTIRuntimeError AOTInductorModelContainerGetNumOutputs(
|
||||
AOTInductorModelContainerHandle container_handle,
|
||||
size_t* ret_num_outputs) {
|
||||
auto* container =
|
||||
reinterpret_cast<torch::aot_inductor::AOTInductorModelContainer*>(
|
||||
container_handle);
|
||||
CONVERT_EXCEPTION_TO_ERROR_CODE(
|
||||
{ *ret_num_outputs = container->num_outputs(); })
|
||||
}
|
||||
|
||||
AOTIRuntimeError AOTInductorModelContainerGetOutputName(
|
||||
AOTInductorModelContainerHandle container_handle,
|
||||
size_t output_idx,
|
||||
const char** ret_output_names) {
|
||||
auto* container =
|
||||
reinterpret_cast<torch::aot_inductor::AOTInductorModelContainer*>(
|
||||
container_handle);
|
||||
CONVERT_EXCEPTION_TO_ERROR_CODE(
|
||||
{ *ret_output_names = container->output_name(output_idx); })
|
||||
}
|
||||
|
||||
AOTIRuntimeError AOTInductorModelContainerGetCallSpec(
|
||||
AOTInductorModelContainerHandle container_handle,
|
||||
const char** in_spec,
|
||||
const char** out_spec) {
|
||||
auto* container =
|
||||
reinterpret_cast<torch::aot_inductor::AOTInductorModelContainer*>(
|
||||
container_handle);
|
||||
CONVERT_EXCEPTION_TO_ERROR_CODE({
|
||||
*in_spec = container->get_in_spec();
|
||||
*out_spec = container->get_out_spec();
|
||||
})
|
||||
}
|
||||
|
||||
AOTIRuntimeError AOTInductorModelCreate(
|
||||
AOTInductorModelHandle* model_handle,
|
||||
AOTInductorConstantMapHandle constant_map_handle){
|
||||
CONVERT_EXCEPTION_TO_ERROR_CODE({
|
||||
auto constant_map = std::make_shared<torch::aot_inductor::ConstantMap>();
|
||||
auto constant_array = std::make_shared<std::vector<torch::aot_inductor::ConstantHandle>>();
|
||||
auto input_map = reinterpret_cast<std::unordered_map<std::string, AtenTensorHandle>*>(constant_map_handle);
|
||||
|
||||
auto model = new torch::aot_inductor::AOTInductorModel(
|
||||
constant_map,
|
||||
constant_array,
|
||||
"cpu", // device_str is hardcoded, as AOTInductorModelCreate is only use for CPU models
|
||||
""
|
||||
);
|
||||
|
||||
if (input_map) {
|
||||
for (auto const& kv : *input_map) {
|
||||
constant_map->emplace(kv.first, kv.second);
|
||||
}
|
||||
} else {
|
||||
model->load_constants();
|
||||
}
|
||||
|
||||
*model_handle = reinterpret_cast<AOTInductorModelHandle>(model);
|
||||
})}
|
||||
|
||||
AOTIRuntimeError AOTInductorModelRun(
|
||||
AOTInductorModelHandle model_handle,
|
||||
AtenTensorHandle* input_handles,
|
||||
AtenTensorHandle* output_handles) {
|
||||
auto model =
|
||||
reinterpret_cast<torch::aot_inductor::AOTInductorModel*>(model_handle);
|
||||
CONVERT_EXCEPTION_TO_ERROR_CODE({
|
||||
AOTINoGradGuard guard;
|
||||
model->run_impl(
|
||||
input_handles,
|
||||
output_handles,
|
||||
(torch::aot_inductor::DeviceStreamType) nullptr,
|
||||
nullptr);
|
||||
})
|
||||
}
|
||||
|
||||
AOTIRuntimeError AOTInductorModelDelete(AOTInductorModelHandle model_handle){
|
||||
CONVERT_EXCEPTION_TO_ERROR_CODE({
|
||||
auto model = reinterpret_cast<torch::aot_inductor::AOTInductorModel*>(
|
||||
model_handle);
|
||||
delete model;
|
||||
})}
|
||||
|
||||
AOTIRuntimeError AOTInductorModelGetNumOutputs(
|
||||
AOTInductorModelHandle model_handle,
|
||||
size_t* ret_num_outputs) {
|
||||
CONVERT_EXCEPTION_TO_ERROR_CODE({
|
||||
auto model = reinterpret_cast<torch::aot_inductor::AOTInductorModel*>(model_handle);
|
||||
*ret_num_outputs = model->num_outputs();
|
||||
})
|
||||
}
|
||||
|
||||
AOTIRuntimeError AOTInductorModelUpdateConstantsMap(
|
||||
AOTInductorModelHandle model_handle,
|
||||
AOTInductorConstantMapHandle constant_map_handle) {
|
||||
auto model =
|
||||
reinterpret_cast<torch::aot_inductor::AOTInductorModel*>(model_handle);
|
||||
CONVERT_EXCEPTION_TO_ERROR_CODE({
|
||||
auto constant_map = std::make_shared<torch::aot_inductor::ConstantMap>();
|
||||
auto input_map =
|
||||
reinterpret_cast<std::unordered_map<std::string, AtenTensorHandle>*>(
|
||||
constant_map_handle);
|
||||
|
||||
for (auto const& kv : *input_map) {
|
||||
constant_map->emplace(kv.first, kv.second);
|
||||
}
|
||||
model->update_constants_map(std::move(constant_map));
|
||||
})
|
||||
}
|
||||
|
||||
AOTIRuntimeError AOTInductorModelContainerGetConstantsBlobSize(
|
||||
AOTInductorModelContainerHandle container_handle,
|
||||
uint64_t* ret_size) {
|
||||
auto* container =
|
||||
reinterpret_cast<torch::aot_inductor::AOTInductorModelContainer*>(
|
||||
container_handle);
|
||||
CONVERT_EXCEPTION_TO_ERROR_CODE(
|
||||
{ *ret_size = container->constant_blob_size(); })
|
||||
}
|
||||
|
||||
|
||||
// Load weights from a single blob in weight_blob_ptr
|
||||
AOTIRuntimeError AOTInductorModelUpdateConstantsFromBlob(
|
||||
AOTInductorModelContainerHandle container_handle,
|
||||
const uint8_t* weight_blob_ptr){
|
||||
auto* container =
|
||||
reinterpret_cast<torch::aot_inductor::AOTInductorModelContainer*>(
|
||||
container_handle);
|
||||
CONVERT_EXCEPTION_TO_ERROR_CODE(
|
||||
{container->update_constants_from_blob(weight_blob_ptr); })
|
||||
}
|
||||
|
||||
|
||||
} // extern "C"
|
||||
@@ -0,0 +1,218 @@
|
||||
import collections
|
||||
import functools
|
||||
import textwrap
|
||||
|
||||
import sympy
|
||||
from sympy import Expr, Symbol
|
||||
|
||||
from torch.utils._ordered_set import OrderedSet
|
||||
from torch.utils._sympy.functions import FloorDiv, ModularIndexing
|
||||
|
||||
from ..utils import sympy_dot, sympy_subs
|
||||
from ..virtualized import V
|
||||
|
||||
|
||||
class BlockPatternMatcher:
|
||||
"""
|
||||
Matches block indexing expressions.
|
||||
"""
|
||||
|
||||
_indexing_wild_signed_int = functools.partial(
|
||||
sympy.Wild, properties=[lambda x: x.is_integer]
|
||||
)
|
||||
_indexing_wild_unsigned_int = functools.partial(
|
||||
sympy.Wild, properties=[lambda x: x.is_integer and x.is_nonnegative]
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def get_subexpr_involving_symbol(cls, expr: Expr, symbol: Symbol) -> Expr:
|
||||
"""
|
||||
Given a sympy expression, return the subexpression comprised only of terms
|
||||
involving the specified symbol.
|
||||
|
||||
For example, if `expr` is `x * 5 + x ** 2 + y * 2 + 5`, and `symbol` is `x`,
|
||||
this returns `x * 5 + x ** 2`.
|
||||
"""
|
||||
expr = cls._preprocess(expr)
|
||||
return sympy.S.Zero + sum(
|
||||
term for term in sympy.Add.make_args(expr) if symbol in term.free_symbols
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def factor_index_expr(expr: sympy.Expr, index_var: Symbol) -> sympy.Expr:
|
||||
"""
|
||||
Given an index expression, factor the expression around
|
||||
- FloorDiv(index_var, ...)
|
||||
- ModularIndexing(index_var, ...)
|
||||
- xindex
|
||||
|
||||
e.g. FloorDiv(index_var, d0)*s0 + FloorDiv(index_var, d0)*s1 ->
|
||||
FloorDiv(index_var, d0) * (s0 + s1)
|
||||
"""
|
||||
centres = OrderedSet()
|
||||
for sub in sympy.preorder_traversal(expr):
|
||||
if isinstance(sub, FloorDiv) and sub.args[0] == index_var:
|
||||
centres.add(sub)
|
||||
elif isinstance(sub, ModularIndexing) and sub.args[0] == index_var:
|
||||
centres.add(sub)
|
||||
centres.add(index_var)
|
||||
|
||||
expr_out = expr
|
||||
for c in centres:
|
||||
expr_out = sympy.collect(expr_out, c)
|
||||
return expr_out
|
||||
|
||||
@staticmethod
|
||||
def get_slice_numels(dims: list[Expr]) -> list[Expr]:
|
||||
"""
|
||||
Compute the cumulative size of each dimension's slice.
|
||||
This proceeds from the last dim up to the second.
|
||||
"""
|
||||
numels = collections.deque([sympy.S.One])
|
||||
for dim in dims[:0:-1]:
|
||||
numel = dim * numels[0]
|
||||
numels.appendleft(numel)
|
||||
return [*numels]
|
||||
|
||||
@staticmethod
|
||||
def _preprocess(expr: Expr) -> Expr:
|
||||
# Remove any Identity nodes, e.g. expand x + (5 * y) to x + 5 * y.
|
||||
# Disable mul and multinomial as those expansions affect op trees:
|
||||
# e.g. sympy expects to match Mul(a, b), but expansion has simplified
|
||||
# to Add(...). Even though they may be algebraically equivalent,
|
||||
# sympy `match` performs structural pattern matching
|
||||
return expr.expand(mul=False, multinomial=False, identity=True)
|
||||
|
||||
@classmethod
|
||||
def match_mod_div_block_expr(
|
||||
cls,
|
||||
index: Expr,
|
||||
index_var: Symbol,
|
||||
numel: Expr,
|
||||
num_dims: int,
|
||||
) -> tuple[list[Expr], list[Expr], list[Expr]] | None:
|
||||
"""
|
||||
Matches modular indexing expressions, converting them to implied block dimensions and strides.
|
||||
See triton.py for more information.
|
||||
|
||||
Warning: this function requires that `index`, `numel` and any other sympy
|
||||
expression does not have precomputed replacements since otherwise block
|
||||
pattern matching may fail.
|
||||
See [Note: Precomputed replacements with BlockPatternMatch]
|
||||
"""
|
||||
index = cls._preprocess(index)
|
||||
|
||||
# Pattern match to find the strides and offset.
|
||||
wild_unsigned_int = functools.partial(
|
||||
cls._indexing_wild_unsigned_int, exclude=[index_var]
|
||||
)
|
||||
wild_signed_int = functools.partial(
|
||||
cls._indexing_wild_signed_int, exclude=[index_var]
|
||||
)
|
||||
dims: list[Expr] = [
|
||||
wild_unsigned_int(f"dim_mod{idx}") for idx in range(num_dims)
|
||||
]
|
||||
strides: list[Expr] = [
|
||||
wild_signed_int(f"stride_mod{idx}") for idx in range(num_dims)
|
||||
]
|
||||
|
||||
# The first dimension's index is computed by division.
|
||||
# The remaining are computed by modulo.
|
||||
slice_numels = cls.get_slice_numels(dims[:num_dims])
|
||||
block_index_exprs = [FloorDiv(index_var, slice_numels[0])] + [
|
||||
ModularIndexing(index_var, numel, dim)
|
||||
for dim, numel in zip(dims[1:], slice_numels[1:])
|
||||
]
|
||||
|
||||
# Calculate a linear index from block indices.
|
||||
match_expr = sympy_dot(strides, block_index_exprs)
|
||||
|
||||
# Heuristic: if the number of dimensions is high, check that the minimum requirements
|
||||
# are met before attempting an expensive full match. see triton.py:match_mod_div_block
|
||||
# for more details. In short, here we check that each subexpression in sympy.Add contains
|
||||
# only FloorDiv or ModularIndexing expressions.
|
||||
if num_dims >= 5:
|
||||
stride = sympy.symbols("stride", cls=wild_signed_int)
|
||||
denom, other = sympy.symbols("denominator other", cls=wild_unsigned_int)
|
||||
mod_div_pattern = stride * ModularIndexing(index_var, denom, other)
|
||||
floor_div_pattern = stride * FloorDiv(index_var, denom)
|
||||
first_dim_floor_div_matched = False
|
||||
match_failed = False
|
||||
for arg in sympy.Add.make_args(index):
|
||||
if arg.match(floor_div_pattern):
|
||||
# There should only be a single FloorDiv(index, denom) expression
|
||||
# corresponding to the first dimension
|
||||
if first_dim_floor_div_matched:
|
||||
match_failed = True
|
||||
break
|
||||
first_dim_floor_div_matched = True
|
||||
elif arg.match(mod_div_pattern):
|
||||
continue
|
||||
else:
|
||||
match_failed = True
|
||||
break
|
||||
|
||||
if match_failed:
|
||||
return None
|
||||
|
||||
# Pattern match.
|
||||
match = index.match(match_expr)
|
||||
if match is None:
|
||||
return None
|
||||
|
||||
# Provide default values for unmatched dims and strides.
|
||||
for dim in dims[1:]:
|
||||
if dim not in match:
|
||||
match[dim] = sympy.S.One
|
||||
for stride in strides[1:]:
|
||||
if stride not in match:
|
||||
match[stride] = sympy.S.Zero
|
||||
|
||||
# Replace wildcards with matched expressions.
|
||||
dims = [dims[0]] + [match[dim] for dim in dims[1:]]
|
||||
strides = [match[stride] for stride in strides]
|
||||
slice_numels = cls.get_slice_numels(dims)
|
||||
block_index_exprs = [sympy_subs(expr, match) for expr in block_index_exprs]
|
||||
|
||||
sizevars = V.graph.sizevars
|
||||
|
||||
# The leading dimension is not directly matched in our expression.
|
||||
# We solve for it by dividing the range tree numel by the product of
|
||||
# all other dimensions. We quit if they are not known to be divisible.
|
||||
assert dims[0] not in match, "Expected not to match the leading dimension!"
|
||||
if not sizevars.statically_known_multiple_of(numel, slice_numels[0]):
|
||||
return None
|
||||
dims[0] = numel / slice_numels[0]
|
||||
|
||||
# Sanity check that we can recover the index from the matched subexpressions.
|
||||
matched_index = sympy_dot(strides, block_index_exprs)
|
||||
assert sizevars.statically_known_equals(
|
||||
matched_index,
|
||||
index,
|
||||
), textwrap.dedent(
|
||||
f"""
|
||||
Invalid match!
|
||||
Index: {index}
|
||||
Matched expression: {matched_index}
|
||||
"""
|
||||
)
|
||||
|
||||
return dims, strides, block_index_exprs
|
||||
|
||||
@classmethod
|
||||
def match_affine_block_expr(
|
||||
cls,
|
||||
index: Expr,
|
||||
index_var: Symbol,
|
||||
) -> Expr | None:
|
||||
"""
|
||||
Matches simple expressions of the form stride * index, returning the
|
||||
stride.
|
||||
"""
|
||||
index = cls._preprocess(index)
|
||||
stride = cls._indexing_wild_signed_int(name="stride", exclude=[index_var])
|
||||
m = index.match(index_var * stride)
|
||||
if m is None:
|
||||
return None
|
||||
|
||||
return m[stride]
|
||||
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,268 @@
|
||||
# mypy: allow-untyped-defs
|
||||
import contextlib
|
||||
import itertools
|
||||
from collections.abc import Callable
|
||||
from typing import Any
|
||||
from unittest.mock import patch
|
||||
|
||||
import sympy
|
||||
|
||||
from .. import ir
|
||||
from ..select_algorithm import PartialRender
|
||||
from ..virtualized import V
|
||||
from .common import ArgName
|
||||
from .cpp_gemm_template import CppGemmTemplate, GEMM_TEMPLATE
|
||||
from .cpp_micro_gemm import LayoutType
|
||||
from .cpp_template_kernel import CppTemplateKernel
|
||||
from .cpp_utils import DTYPE_TO_CPP, GemmBlocking
|
||||
|
||||
|
||||
# We pass all sizevars present in BY to the GEMM templates so variables are not renamed in the BMM definition
|
||||
GEMM_SINGLE_THREAD_MM_STUB = r"""
|
||||
{{kernel.def_kernel(
|
||||
inputs={"X": X, "W": W},
|
||||
outputs={"Y": Y_2d},
|
||||
aliases=aliases,
|
||||
function_name=kernel_name+"_single_thread_mm",
|
||||
extra_sizevars=BY_sizevars + [b_index],
|
||||
placeholder="<SINGLE_THREAD_MM_DEF_FOR_BMM>")}}"""
|
||||
|
||||
GEMM_THREADED_MM_STUB = r"""
|
||||
{{kernel.def_kernel(
|
||||
inputs={"X": X, "W": W},
|
||||
outputs={"Y": Y_2d},
|
||||
aliases=aliases,
|
||||
function_name=kernel_name+"_threaded_mm",
|
||||
extra_sizevars=BY_sizevars + [b_index],
|
||||
placeholder="<THREADED_MM_DEF_FOR_BMM>")}}"""
|
||||
|
||||
BMM_TEMPLATE = r"""
|
||||
{{ template.codegen_microkernel_def() }}
|
||||
{{ template.codegen_single_thread_gemm() }}
|
||||
{{ template.codegen_multi_thread_gemm() }}
|
||||
|
||||
extern "C"
|
||||
{{kernel.def_kernel(inputs={"X": BX, "W": BW}, outputs={"Y": BY}, aliases=aliases)}}
|
||||
{
|
||||
const int64_t B = {{kernel.size(BY_2d, 0)}};
|
||||
{%- if num_threads > 1 %}
|
||||
constexpr int64_t num_threads = {{num_threads}};
|
||||
int64_t B_single_thread_block = (B / num_threads) * num_threads;
|
||||
|
||||
{%- set use_dynamic_threads = ((config.cpp.threads < 1) and (num_threads == cpu_count)) or config.cpp.dynamic_threads %}
|
||||
{%- if use_dynamic_threads %}
|
||||
#pragma omp parallel for
|
||||
{%- else %}
|
||||
#pragma omp parallel for num_threads({{num_threads}})
|
||||
{%- endif %}
|
||||
{%- else %}
|
||||
int64_t B_single_thread_block = B;
|
||||
{%- endif %}
|
||||
for (int64_t b_start = 0; b_start < B_single_thread_block; ++b_start) {
|
||||
{{template.get_gemm_function_call(
|
||||
kernel,
|
||||
kernel_name+"_single_thread_mm",
|
||||
"<SINGLE_THREAD_CALL_FOR_BMM>",
|
||||
b_index="b_start",
|
||||
)}}
|
||||
}
|
||||
for (int64_t b_start = B_single_thread_block; b_start < B; ++b_start) {
|
||||
{{template.get_gemm_function_call(
|
||||
kernel,
|
||||
kernel_name+"_threaded_mm",
|
||||
"<THREADED_MM_CALL_FOR_BMM>",
|
||||
b_index="b_start",
|
||||
)}}
|
||||
}
|
||||
}
|
||||
"""
|
||||
|
||||
|
||||
class CppBmmTemplate(CppGemmTemplate):
|
||||
def __init__(
|
||||
self,
|
||||
input_nodes,
|
||||
layout: ir.Layout,
|
||||
num_threads: int,
|
||||
register_blocking: GemmBlocking,
|
||||
beta=1,
|
||||
alpha=1,
|
||||
has_bias=False,
|
||||
epilogue_creator: Callable[[ir.Buffer], ir.Pointwise] | None = None,
|
||||
should_block_weights: bool = False,
|
||||
name="bmm",
|
||||
):
|
||||
"""
|
||||
In order to simplify the implementation and increase code reuse, the BMM template implements
|
||||
two versions of the GEMM kernel: a single-threaded version and a multi-threaded version.
|
||||
GEMM kernels are called in a loop over the batch dimension, with single-threaded GEMM calls
|
||||
for all but the last (B % num_threads), which are handled by the multi-threaded GEMM kernel.
|
||||
|
||||
We use an extra sizevar `b_index` to index the batch dimension, which we pass into the GEMM
|
||||
template as a sympy.Symbol. This allows us to slice the 3D batch tensors in the GEMM template
|
||||
without any changes to the GEMM template itself.
|
||||
"""
|
||||
super().__init__(
|
||||
input_nodes,
|
||||
layout,
|
||||
num_threads,
|
||||
register_blocking,
|
||||
beta=beta,
|
||||
alpha=alpha,
|
||||
has_bias=has_bias,
|
||||
epilogue_creator=epilogue_creator,
|
||||
should_block_weights=should_block_weights,
|
||||
name=name,
|
||||
)
|
||||
self.b_index = sympy.Symbol("s_b_index", integer=True, nonnegative=True)
|
||||
|
||||
@staticmethod
|
||||
def get_padded_size(n, block_n, k, should_block_weight):
|
||||
if should_block_weight:
|
||||
# Tensor is constant or not contiguous, so we will pad and block
|
||||
new_size, padded_n = CppGemmTemplate.get_padded_size(
|
||||
n, block_n, k, should_block_weight
|
||||
)
|
||||
# Add the new batch dimension
|
||||
new_size.insert(0, -1)
|
||||
return new_size, padded_n
|
||||
else:
|
||||
new_size = [-1, k, n]
|
||||
return new_size, n
|
||||
|
||||
@staticmethod
|
||||
def check_if_block_weight(W, micro_gemm):
|
||||
assert isinstance(W, ir.IRNode)
|
||||
_, n = W.get_size()[-2:]
|
||||
result = (
|
||||
not W.get_layout().is_contiguous()
|
||||
or W.get_name() in V.graph.constants
|
||||
or (
|
||||
n % micro_gemm.register_blocking.block_n != 0
|
||||
and micro_gemm.get_b_layout != LayoutType.NORMAL
|
||||
)
|
||||
)
|
||||
return result
|
||||
|
||||
def get_gemm_function_call(
|
||||
self,
|
||||
kernel: CppTemplateKernel,
|
||||
function_name: str,
|
||||
placeholder: str,
|
||||
b_index: str,
|
||||
) -> str:
|
||||
"""
|
||||
Similar to 'def_kernel' in cpp_template_kernel, but instead of generating a function definition,
|
||||
generate a function call for the GEMM kernel.
|
||||
Args:
|
||||
placeholder: The string to replace the function call with
|
||||
b_index: The index for slicing the 3D batch tensors
|
||||
"""
|
||||
|
||||
def hook():
|
||||
arg_defs, call_args, _, _ = kernel.args.python_argdefs()
|
||||
for i, buf in enumerate(call_args):
|
||||
if buf == self.b_index:
|
||||
arg_defs[i] = ArgName(b_index)
|
||||
call = f"{function_name}({', '.join(x.full_name() for x in arg_defs)});"
|
||||
return call
|
||||
|
||||
assert placeholder not in kernel.render_hooks
|
||||
kernel.render_hooks[placeholder] = hook
|
||||
return placeholder
|
||||
|
||||
def get_default_reindexers(self, epilogue_nodes):
|
||||
def reindexer(args):
|
||||
# if epilogue nodes exist, they have 3D ranges but args are 2D, so add 0 index
|
||||
return [self.b_index] + args
|
||||
|
||||
return [reindexer] * len(epilogue_nodes)
|
||||
|
||||
def get_options(
|
||||
self,
|
||||
kernel: CppTemplateKernel,
|
||||
template_buffer_node: ir.CppTemplateBuffer | None = None,
|
||||
flag_template_buffer_has_other_users: bool | None = None,
|
||||
epilogue_nodes: list[ir.IRNode] | None = None,
|
||||
**kwargs,
|
||||
) -> dict[str, Any]:
|
||||
options = super().get_options(
|
||||
kernel=kernel,
|
||||
template_buffer_node=template_buffer_node,
|
||||
flag_template_buffer_has_other_users=flag_template_buffer_has_other_users,
|
||||
epilogue_nodes=epilogue_nodes,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
BX, BW, BY = options["X"], options["W"], options["Y"]
|
||||
options["BX"], options["BW"], options["BY"] = BX, BW, BY
|
||||
options["BY_2d"] = options["Y_2d"]
|
||||
for kword in ["X", "W", "GemmOut", "Y_2d"]:
|
||||
options[kword] = kernel.select(options[kword], 0, self.b_index)
|
||||
for kword in ["X", "W", "Y_2d"]:
|
||||
options[kword + "_dtype"] = DTYPE_TO_CPP[options[kword].dtype]
|
||||
options["b_index"] = self.b_index
|
||||
options["BY_sizevars"] = [
|
||||
s
|
||||
for sym in itertools.chain(BY.get_size(), BY.get_stride())
|
||||
if isinstance(sym, sympy.Expr)
|
||||
for s in sym.free_symbols
|
||||
]
|
||||
options["kernel_name"] = kernel.kernel_name
|
||||
|
||||
return options
|
||||
|
||||
def render( # type: ignore[override, return]
|
||||
self,
|
||||
kernel: CppTemplateKernel,
|
||||
template_buffer_node: ir.CppTemplateBuffer | None = None,
|
||||
flag_template_buffer_has_other_users: bool | None = None,
|
||||
epilogue_nodes: list[ir.IRNode] | None = None,
|
||||
**kwargs,
|
||||
) -> str:
|
||||
options = self.get_options(
|
||||
kernel=kernel,
|
||||
template_buffer_node=template_buffer_node,
|
||||
flag_template_buffer_has_other_users=flag_template_buffer_has_other_users,
|
||||
epilogue_nodes=epilogue_nodes,
|
||||
**kwargs,
|
||||
)
|
||||
self.render_options = options
|
||||
|
||||
with contextlib.ExitStack() as stack:
|
||||
for buf in options["fake_buffers"]:
|
||||
stack.enter_context(
|
||||
patch.object(V.graph, "get_dtype", self._fake_get_dtype(buf))
|
||||
)
|
||||
result = self._template_from_string(BMM_TEMPLATE).render(**options)
|
||||
|
||||
# Finalize the function definitions for the gemm routines
|
||||
sub_mm_hooks = {
|
||||
name: hook
|
||||
for name, hook in kernel.render_hooks.items()
|
||||
if "FOR_BMM" in name
|
||||
}
|
||||
result = PartialRender(result, sub_mm_hooks).finalize_all()
|
||||
for name in sub_mm_hooks:
|
||||
del kernel.render_hooks[name]
|
||||
del kernel.args.sizevars[options["b_index"]]
|
||||
return result
|
||||
|
||||
def codegen_single_thread_gemm(self):
|
||||
stub = self._template_from_string(GEMM_SINGLE_THREAD_MM_STUB).render(
|
||||
self.render_options
|
||||
)
|
||||
return stub + self._template_from_string(GEMM_TEMPLATE).render(
|
||||
{**self.render_options, "num_threads": 1}
|
||||
)
|
||||
|
||||
def codegen_multi_thread_gemm(self):
|
||||
stub = self._template_from_string(GEMM_THREADED_MM_STUB).render(
|
||||
self.render_options
|
||||
)
|
||||
return stub + self._template_from_string(GEMM_TEMPLATE).render(
|
||||
self.render_options
|
||||
)
|
||||
|
||||
def codegen_gemm_stub_def(self):
|
||||
return ""
|
||||
+1089
File diff suppressed because it is too large
Load Diff
+1837
File diff suppressed because it is too large
Load Diff
+521
@@ -0,0 +1,521 @@
|
||||
import contextlib
|
||||
import logging
|
||||
import os
|
||||
from collections.abc import Callable
|
||||
from typing import Any, cast, TypeVar
|
||||
from unittest.mock import patch
|
||||
|
||||
import torch
|
||||
import torch.utils
|
||||
from torch.utils._ordered_set import OrderedSet
|
||||
|
||||
from ..._dynamo.utils import counters
|
||||
from .. import config, ir
|
||||
from ..kernel.mm_common import mm_args
|
||||
from ..select_algorithm import ChoiceCaller, DataProcessorTemplateWrapper
|
||||
from ..utils import parallel_num_threads
|
||||
from ..virtualized import V
|
||||
from .cpp import get_export_declaration
|
||||
from .cpp_gemm_template import (
|
||||
CppGemmTemplate,
|
||||
expand_bias,
|
||||
gen_2d_view_of_epilogue_buf,
|
||||
prune_tensors,
|
||||
transpose_w,
|
||||
)
|
||||
from .cpp_micro_gemm import CppMicroGemmAMX, create_micro_gemm
|
||||
from .cpp_template_kernel import CppTemplateKernel
|
||||
from .cpp_utils import (
|
||||
create_epilogue_with_attr,
|
||||
DTYPE_TO_CPP,
|
||||
GemmBlocking,
|
||||
get_gemm_template_output_and_compute_dtype,
|
||||
)
|
||||
|
||||
|
||||
log = logging.getLogger(__name__)
|
||||
|
||||
GEMM_TEMPLATE = r"""
|
||||
{{template.header().getvalue()}}
|
||||
{{micro_gemm.codegen_define(kernel)}}
|
||||
|
||||
extern "C" {{export_declaration}}
|
||||
{{kernel.def_kernel(inputs=kernel_args, outputs=Y_list, aliases=aliases)}}
|
||||
{
|
||||
{{kernel.maybe_codegen_profile()}}
|
||||
{{ template.codegen_blocks(
|
||||
num_threads, N, K, micro_gemm, is_dynamic_M, kernel, GemmOuts[0], config, L1_cache_size, L2_cache_size, X_list[0], W_list[0]
|
||||
) }}
|
||||
{%- if num_threads > 1 %}
|
||||
{%- set use_dynamic_threads = ((config.cpp.threads < 1) and (num_threads == cpu_count)) or config.cpp.dynamic_threads %}
|
||||
{%- if use_dynamic_threads %}
|
||||
#pragma omp parallel
|
||||
{%- else %}
|
||||
#pragma omp parallel num_threads({{num_threads}})
|
||||
{%- endif %}
|
||||
{
|
||||
{{ template.codegen_multi_threads_params()|indent(8, false) }}
|
||||
{%- else %}
|
||||
{
|
||||
{{ template.codegen_single_thread_params(is_dynamic_M)|indent(8, false) }}
|
||||
{%- endif %}
|
||||
{{ micro_gemm.codegen_init(kernel) }}
|
||||
{%- set acc_buf_name_list=[] %}
|
||||
{%- set acc_buf_name_prefix = "local_acc_buf_" %}
|
||||
{%- for gemm_idx in range(0, gemm_grouped_num, 1) %}
|
||||
{%- set acc_buf_name = acc_buf_name_prefix + gemm_idx|string %}
|
||||
{{ kernel.define_buffer(acc_buf_name, ["Mc_blocks*Mr", "Nc_blocks*Nr"], acc_buf_dtype) }}
|
||||
{%- set acc_buf_name_list=acc_buf_name_list.append(acc_buf_name) %}
|
||||
{%- endfor %}
|
||||
for (int64_t mc_block_id = 0; mc_block_id < num_Mc_blocks_per_thread; mc_block_id++) {
|
||||
{{ template.codegen_m_loop_params()|indent(12, false) }}
|
||||
for (int64_t nc = n_block_start; nc < n_block_end; nc += Nc_blocks) {
|
||||
{{ template.codegen_n_loop_params()|indent(16, false) }}
|
||||
{%- set acc_list=[] %}
|
||||
{%- for gemm_idx in range(0, gemm_grouped_num, 1) %}
|
||||
{%- set acc_list = acc_list.append( kernel.local_buffers[acc_buf_name_list[gemm_idx]] ) %}
|
||||
{{ kernel.reinit_buffer_if_null(acc_buf_name_list[gemm_idx]) }}
|
||||
{%- endfor %}
|
||||
for (int64_t kc = k_block_start; kc < k_block_end; kc += Kc_blocks) {
|
||||
int64_t k_start = kc * Kr;
|
||||
int64_t k_end = std::min(std::min(kc + Kc_blocks, k_block_end) * Kr, K);
|
||||
{%- set tile_X_list=[] %}
|
||||
{%- for gemm_idx in range(0, gemm_grouped_num, 1) %}
|
||||
{%- set tile_X_list = tile_X_list.append( kernel.slice_nd(X_list[gemm_idx], [("m_start", "m_end"), ("k_start", "k_end")]) ) %}
|
||||
{%- endfor %}
|
||||
for (int64_t nci = nc; nci < nc_block_end; nci++) {
|
||||
{%- set tile_W_3d_list=[] %}
|
||||
{%- set tile_W_list=[] %}
|
||||
{%- set acc_slice_list=[] %}
|
||||
{%- for gemm_idx in range(0, gemm_grouped_num, 1) %}
|
||||
{%- set acc_slice_list = acc_slice_list.append(
|
||||
kernel.slice_nd(acc_list[gemm_idx], [("0", "m_end - m_start"), ("(nci - nc)*Nr", "(nci - nc + 1)*Nr")])
|
||||
) %}
|
||||
{%- set tile_W_3d_list = tile_W_3d_list.append(
|
||||
kernel.slice_nd(W_list[gemm_idx], [("nci", "nci + 1"), ("k_start", "k_end"), ()])
|
||||
) %}
|
||||
{%- endfor %}
|
||||
{%- for gemm_idx in range(0, gemm_grouped_num, 1) %}
|
||||
{%- set tile_W_list = tile_W_list.append(
|
||||
kernel.view(tile_W_3d_list[gemm_idx], ["k_end - k_start", micro_gemm.register_blocking.block_n])
|
||||
) %}
|
||||
{%- endfor %}
|
||||
if (kc == k_block_start) {
|
||||
{%- for gemm_idx in range(0, gemm_grouped_num, 1) %}
|
||||
{{ micro_gemm.codegen_call(
|
||||
kernel, tile_X_list[gemm_idx], tile_W_list[gemm_idx], acc_slice_list[gemm_idx], accum=False
|
||||
)|indent(28, false) }}
|
||||
{%- endfor %}
|
||||
} else {
|
||||
{%- for gemm_idx in range(0, gemm_grouped_num, 1) %}
|
||||
{{ micro_gemm.codegen_call(
|
||||
kernel, tile_X_list[gemm_idx], tile_W_list[gemm_idx], acc_slice_list[gemm_idx], accum=True
|
||||
)|indent(28, false) }}
|
||||
{%- endfor %}
|
||||
}
|
||||
}
|
||||
}
|
||||
{
|
||||
{%- set tile_acc_list = [] %}
|
||||
{%- set tile_Y_list = [] %}
|
||||
{%- for gemm_idx in range(0, gemm_grouped_num, 1) %}
|
||||
{%- set tile_acc_list = tile_acc_list.append(
|
||||
kernel.slice_nd(acc_list[gemm_idx], [("0", "m_end - m_start"), ("0", "n_end - n_start")])
|
||||
) %}
|
||||
{%- set tile_Y_list = tile_Y_list.append(
|
||||
kernel.slice_nd(Y_2d_list[gemm_idx], [("m_start", "m_end"), ("n_start", "n_end")])
|
||||
) %}
|
||||
{%- endfor %}
|
||||
{{ kernel.store_outputs(
|
||||
tile_Y_list,
|
||||
tile_acc_list,
|
||||
GemmOuts,
|
||||
epilogue_nodes,
|
||||
offsets=("m_start", "n_start"),
|
||||
reindexers=reindexers,
|
||||
multi_output_buffers=multi_output_buffers
|
||||
)|indent(20, false)
|
||||
}}
|
||||
}
|
||||
}
|
||||
}
|
||||
{{ micro_gemm.codegen_finalize(kernel) }}
|
||||
}
|
||||
}
|
||||
"""
|
||||
|
||||
|
||||
def get_deduplicated_act(act_mapping: dict[int, ir.IRNode]) -> list[ir.IRNode]:
|
||||
act_deduplicated = []
|
||||
act_deduplicated_name: OrderedSet[str] = OrderedSet()
|
||||
for act_idx in range(len(act_mapping.values())):
|
||||
act = act_mapping[act_idx]
|
||||
if act.get_name() not in act_deduplicated_name:
|
||||
act_deduplicated.append(act)
|
||||
act_deduplicated_name.add(act.get_name())
|
||||
return act_deduplicated
|
||||
|
||||
|
||||
class CppGroupedGemmTemplate(CppGemmTemplate):
|
||||
def __init__(
|
||||
self,
|
||||
input_nodes: list[ir.IRNode],
|
||||
layout: ir.Layout,
|
||||
num_threads: int,
|
||||
register_blocking: GemmBlocking,
|
||||
beta: int = 1,
|
||||
alpha: int = 1,
|
||||
has_bias: bool = False,
|
||||
epilogue_creator: Callable[[ir.Buffer], ir.Pointwise] | None = None,
|
||||
act_mapping: dict[int, ir.IRNode] | None = None,
|
||||
gemm_grouped_num: int = 1,
|
||||
) -> None:
|
||||
"""
|
||||
Template for Group of GEMMs:
|
||||
* Each GEMM has the same dimensions (m, n, k) and the same leading dimensions (lda, ldb, ldc)
|
||||
for their A, B, and C matrices.
|
||||
* Each GEMM has distinct or shared activations, has distinct weight, has unique bias or no bias, has distinct epilogues.
|
||||
* In the current implementation, the outputs of all GEMMs are accumulated using pointwise epilogues.
|
||||
This behavior can be extended in the future if needed.
|
||||
"""
|
||||
super().__init__(
|
||||
input_nodes,
|
||||
layout,
|
||||
num_threads,
|
||||
register_blocking,
|
||||
beta,
|
||||
alpha,
|
||||
has_bias,
|
||||
epilogue_creator,
|
||||
)
|
||||
self.act_mapping = act_mapping
|
||||
self.gemm_grouped_num = gemm_grouped_num
|
||||
# pyrefly: ignore [bad-override]
|
||||
self.output_node: list[ir.Buffer] = [
|
||||
ir.Buffer(name="buf_out" + str(idx), layout=layout)
|
||||
for idx in range(gemm_grouped_num)
|
||||
]
|
||||
|
||||
@classmethod
|
||||
# pyrefly: ignore [bad-override]
|
||||
def add_choices(
|
||||
cls,
|
||||
choices: list[ChoiceCaller],
|
||||
layout: ir.Layout,
|
||||
input_nodes: list[ir.IRNode],
|
||||
beta: int = 1,
|
||||
alpha: int = 1,
|
||||
has_bias: tuple[bool, ...] = (False, False),
|
||||
trans_w: bool = False,
|
||||
input_indices: list[int] | None = None,
|
||||
epilogue_creator: Callable[[ir.Buffer], ir.Pointwise] | None = None,
|
||||
act_mapping: dict[int, ir.IRNode] | None = None, # gemm idx to its act buf
|
||||
) -> DataProcessorTemplateWrapper:
|
||||
# Input nodes order: x, optional[x1], ... w0, w1, ... optional[b0], optional[b1], ...
|
||||
gemm_grouped_num = len(has_bias)
|
||||
assert act_mapping
|
||||
act_deduplicated = get_deduplicated_act(act_mapping)
|
||||
wgt_start_idx = len(act_deduplicated)
|
||||
bias_start_idx = wgt_start_idx + gemm_grouped_num
|
||||
input_indices = list(range(len(input_nodes)))
|
||||
|
||||
_T = TypeVar("_T", ir.IRNode, torch.Tensor)
|
||||
_U = TypeVar("_U", ir.Layout, torch.Tensor)
|
||||
|
||||
def reorder_and_filter(
|
||||
inputs: list[_T],
|
||||
layout_or_out: _U,
|
||||
) -> tuple[list[_T], _U]:
|
||||
assert input_indices is not None, "input_indices must be set"
|
||||
return [inputs[idx] for idx in input_indices], layout_or_out
|
||||
|
||||
new_inputs, new_layout = reorder_and_filter(input_nodes, layout)
|
||||
|
||||
def maybe_to_dense(
|
||||
inputs: list[_T],
|
||||
layout_or_out: _U,
|
||||
) -> tuple[list[_T], _U]:
|
||||
new_inputs = list(inputs)
|
||||
for idx in range(wgt_start_idx, wgt_start_idx + gemm_grouped_num):
|
||||
if isinstance(inputs[idx], torch.Tensor):
|
||||
W = inputs[idx]
|
||||
assert isinstance(W, torch.Tensor), "W must be a torch.Tensor"
|
||||
# pyrefly: ignore [unsupported-operation]
|
||||
new_inputs[idx] = W.to_dense() if W.is_mkldnn else W
|
||||
return new_inputs, layout_or_out
|
||||
|
||||
def normalize_shapes(
|
||||
inputs: list[_T],
|
||||
layout_or_out: _U,
|
||||
) -> tuple[list[_T], _U]:
|
||||
new_inputs: list[_T] = list(inputs)
|
||||
if not trans_w:
|
||||
return new_inputs, layout_or_out
|
||||
X = new_inputs[0]
|
||||
for wgt_idx in range(wgt_start_idx, wgt_start_idx + gemm_grouped_num):
|
||||
new_input = new_inputs[wgt_idx]
|
||||
new_inputs[wgt_idx] = transpose_w(new_input, trans_w)
|
||||
for bias_idx in range(bias_start_idx, len(new_inputs)):
|
||||
# pyrefly: ignore [bad-argument-type]
|
||||
new_bias = expand_bias(new_inputs[bias_idx], X)
|
||||
assert new_bias is not None
|
||||
# pyrefly: ignore [unsupported-operation]
|
||||
new_inputs[bias_idx] = new_bias
|
||||
return new_inputs, layout_or_out
|
||||
|
||||
num_threads = parallel_num_threads()
|
||||
new_inputs, _ = normalize_shapes(*maybe_to_dense(new_inputs, new_layout))
|
||||
m, n, k, *_ = mm_args(new_inputs[0], new_inputs[wgt_start_idx])
|
||||
output_dtype, compute_dtype = get_gemm_template_output_and_compute_dtype(
|
||||
new_inputs[0].get_dtype()
|
||||
)
|
||||
micro_gemm = create_micro_gemm(
|
||||
"micro_gemm",
|
||||
m,
|
||||
n,
|
||||
k,
|
||||
input_dtype=new_inputs[0].get_dtype(),
|
||||
input2_dtype=new_inputs[wgt_start_idx].get_dtype(),
|
||||
output_dtype=output_dtype,
|
||||
compute_dtype=compute_dtype,
|
||||
alpha=alpha,
|
||||
num_threads=num_threads,
|
||||
)
|
||||
assert micro_gemm is not None
|
||||
_, block_n, _ = micro_gemm.register_blocking
|
||||
new_size, padded_n = cls.get_padded_size(
|
||||
n, block_n, k, should_block_weight=True
|
||||
)
|
||||
padding = padded_n - n
|
||||
|
||||
def pack_weight(
|
||||
inputs: list[_T],
|
||||
layout_or_out: _U,
|
||||
) -> tuple[list[_T], _U]:
|
||||
new_W_list = []
|
||||
new_inputs = list(inputs)
|
||||
W_list = new_inputs[wgt_start_idx : wgt_start_idx + gemm_grouped_num]
|
||||
for W in W_list:
|
||||
blocked_w = cls.block_weight(W, new_size, padding)
|
||||
new_W_list.append(cls.pack_vnni_weight(blocked_w, micro_gemm, new_size))
|
||||
new_inputs[wgt_start_idx : wgt_start_idx + gemm_grouped_num] = new_W_list
|
||||
return new_inputs, layout_or_out
|
||||
|
||||
def preprocessor(
|
||||
inputs: list[_T],
|
||||
layout: _U,
|
||||
) -> tuple[list[_T], _U]:
|
||||
return pack_weight(
|
||||
*normalize_shapes(*maybe_to_dense(*reorder_and_filter(inputs, layout)))
|
||||
)
|
||||
|
||||
def postprocessor(output: _T) -> _T:
|
||||
if isinstance(output, ir.TensorBox):
|
||||
template_buffer = ir.InputsKernel.unwrap_storage_for_input(output)
|
||||
assert isinstance(template_buffer, ir.CppTemplateBuffer)
|
||||
new_input_nodes, _ = reorder_and_filter(input_nodes, layout)
|
||||
W_nodes = new_input_nodes[
|
||||
wgt_start_idx : wgt_start_idx + gemm_grouped_num
|
||||
]
|
||||
W_tensor = []
|
||||
for W_node in W_nodes:
|
||||
assert W_node.get_name() in V.graph.constants
|
||||
# pyrefly: ignore [bad-argument-type]
|
||||
W_tensor.append(V.graph.constants[W_node.get_name()])
|
||||
# pyrefly: ignore [unsupported-operation]
|
||||
new_input_nodes[wgt_start_idx : wgt_start_idx + gemm_grouped_num] = (
|
||||
W_tensor # type: ignore[assignment]
|
||||
)
|
||||
new_input_nodes, _ = pack_weight(
|
||||
*normalize_shapes(*maybe_to_dense(new_input_nodes, layout))
|
||||
)
|
||||
# Prune unused tensors
|
||||
prune_tensors(input_nodes, new_input_nodes)
|
||||
for idx in range(wgt_start_idx, wgt_start_idx + gemm_grouped_num):
|
||||
W_packed = new_input_nodes[idx]
|
||||
assert isinstance(W_packed, torch.Tensor)
|
||||
W_packed_constant = V.graph.add_tensor_constant(W_packed)
|
||||
template_buffer.inputs[idx] = (
|
||||
ir.InputsKernel.unwrap_storage_for_input(W_packed_constant)
|
||||
)
|
||||
|
||||
return output
|
||||
|
||||
template = DataProcessorTemplateWrapper(
|
||||
CppGroupedGemmTemplate,
|
||||
preprocessor,
|
||||
postprocessor,
|
||||
input_nodes=input_nodes,
|
||||
layout=layout,
|
||||
num_threads=num_threads,
|
||||
register_blocking=micro_gemm.register_blocking,
|
||||
beta=beta,
|
||||
alpha=alpha,
|
||||
has_bias=has_bias,
|
||||
epilogue_creator=epilogue_creator,
|
||||
act_mapping=act_mapping,
|
||||
gemm_grouped_num=gemm_grouped_num,
|
||||
)
|
||||
template.maybe_append_choice(choices)
|
||||
return template
|
||||
|
||||
def render( # type: ignore[override,return,no-untyped-def]
|
||||
self,
|
||||
kernel: CppTemplateKernel,
|
||||
template_buffer_node: ir.CppTemplateBuffer | None = None,
|
||||
flag_template_buffer_has_other_users: bool | None = None,
|
||||
epilogue_nodes: list[ir.IRNode] | None = None,
|
||||
**kwargs,
|
||||
) -> str:
|
||||
assert self.act_mapping
|
||||
act_deduplicated = get_deduplicated_act(self.act_mapping)
|
||||
wgt_start_idx = len(act_deduplicated)
|
||||
bias_start_idx = wgt_start_idx + self.gemm_grouped_num
|
||||
X_list = list(self.act_mapping.values())
|
||||
W_list = self.input_nodes[wgt_start_idx : wgt_start_idx + self.gemm_grouped_num]
|
||||
inp_list = []
|
||||
cur_idx = bias_start_idx
|
||||
for inp_idx in range(self.gemm_grouped_num):
|
||||
inp = None
|
||||
# pyrefly: ignore [bad-index, index-error]
|
||||
if self.has_bias[inp_idx]:
|
||||
inp = self.input_nodes[cur_idx]
|
||||
cur_idx += 1
|
||||
inp_list.append(inp)
|
||||
|
||||
Y_list = self.output_node
|
||||
multi_output_buffers = None
|
||||
if template_buffer_node is not None:
|
||||
W_list = template_buffer_node.inputs[
|
||||
wgt_start_idx : wgt_start_idx + self.gemm_grouped_num
|
||||
]
|
||||
assert isinstance(template_buffer_node.outputs, list)
|
||||
Y_list = template_buffer_node.outputs
|
||||
counters["inductor"]["cpp_grouped_gemm_template"] += 1
|
||||
multi_output_buffers = template_buffer_node.outputs
|
||||
|
||||
template_buffer = Y_list[0]
|
||||
fake_buffers: list[ir.Buffer] = []
|
||||
Y_2d_list = Y_list
|
||||
output_dtype, compute_dtype = get_gemm_template_output_and_compute_dtype(
|
||||
X_list[0].get_dtype()
|
||||
)
|
||||
micro_gemm = create_micro_gemm(
|
||||
f"{kernel.kernel_name}_micro_gemm",
|
||||
self.m,
|
||||
self.n,
|
||||
self.k,
|
||||
input_dtype=X_list[0].get_dtype(),
|
||||
input2_dtype=W_list[0].get_dtype(),
|
||||
output_dtype=output_dtype,
|
||||
compute_dtype=compute_dtype,
|
||||
alpha=self.alpha,
|
||||
num_threads=self.num_threads,
|
||||
)
|
||||
assert micro_gemm is not None
|
||||
assert self.register_blocking == micro_gemm.register_blocking
|
||||
self.log_blockings()
|
||||
if isinstance(micro_gemm, CppMicroGemmAMX):
|
||||
counters["inductor"]["cpp_micro_gemm_amx_counter"] += 1
|
||||
|
||||
L1_cache_size = torch.cpu.get_capabilities().get(
|
||||
"l1d_cache_size", 0
|
||||
) # per core cache size in Bytes
|
||||
assert L1_cache_size > 0, f"Expect L1_cache_size > 0 but got {L1_cache_size}"
|
||||
|
||||
L2_cache_size = torch.cpu.get_capabilities().get(
|
||||
"l2_cache_size", 0
|
||||
) # per core cache size in Bytes
|
||||
assert L2_cache_size > 0, f"Expect L2_cache_size > 0 but got {L2_cache_size}"
|
||||
|
||||
epilogues: list[ir.IRNode] = []
|
||||
reindexers: list[Callable[[list[Any]], list[Any]] | None] = []
|
||||
gemm_output_buffers: list[ir.Buffer] = []
|
||||
for out_buf_idx in range(self.gemm_grouped_num):
|
||||
gemm_output_name = f"{template_buffer.get_name()}_GemmOut" + str(
|
||||
out_buf_idx
|
||||
)
|
||||
gemm_output_buffers.append(
|
||||
ir.Buffer(name=gemm_output_name, layout=template_buffer.layout)
|
||||
)
|
||||
|
||||
assert not self.epilogue_creator, (
|
||||
"epilogue_creator is not supported yet in Grouped GEMM Template"
|
||||
)
|
||||
|
||||
kernel_args: dict[str, ir.IRNode | None] = {}
|
||||
for x_idx in range(wgt_start_idx):
|
||||
kernel_args["X" + str(x_idx)] = act_deduplicated[x_idx]
|
||||
for w_idx in range(self.gemm_grouped_num):
|
||||
kernel_args["W" + str(w_idx)] = W_list[w_idx]
|
||||
for inp_idx in range(self.gemm_grouped_num):
|
||||
kernel_args["inp" + str(inp_idx)] = inp_list[inp_idx]
|
||||
|
||||
def _bias_add_epilogue(buf: ir.IRNode, inp: ir.IRNode) -> ir.Pointwise:
|
||||
return create_epilogue_with_attr(
|
||||
buf, "bias_add", other=inp, beta=self.beta, dtype=self.layout.dtype
|
||||
)
|
||||
|
||||
for gemm_idx, inp in enumerate(inp_list):
|
||||
if inp:
|
||||
buffer_name = Y_list[gemm_idx].get_name()
|
||||
epilogues.append(
|
||||
ir.ComputedBuffer(
|
||||
name=buffer_name,
|
||||
layout=template_buffer.layout,
|
||||
data=_bias_add_epilogue(gemm_output_buffers[gemm_idx], inp),
|
||||
)
|
||||
)
|
||||
reindexers.append(None)
|
||||
|
||||
if epilogue_nodes:
|
||||
epilogues.extend(epilogue_nodes)
|
||||
for epilogue_node in epilogue_nodes:
|
||||
Y = cast(ir.Buffer, epilogue_node)
|
||||
_, reindexers = gen_2d_view_of_epilogue_buf(
|
||||
Y,
|
||||
template_buffer,
|
||||
[
|
||||
epilogue_node,
|
||||
],
|
||||
reindexers,
|
||||
default_reindexers=[
|
||||
None,
|
||||
],
|
||||
)
|
||||
|
||||
options = dict(
|
||||
N=self.n,
|
||||
K=self.k,
|
||||
PADDED_N=self.padded_n,
|
||||
aliases={},
|
||||
beta=self.beta,
|
||||
alpha=self.alpha,
|
||||
num_threads=self.num_threads,
|
||||
micro_gemm=micro_gemm,
|
||||
is_dynamic_M=self.is_dynamic_M,
|
||||
template=self,
|
||||
kernel=kernel,
|
||||
export_declaration=get_export_declaration(),
|
||||
acc_buf_dtype=torch.float,
|
||||
DTYPE_TO_CPP=DTYPE_TO_CPP,
|
||||
L1_cache_size=L1_cache_size,
|
||||
L2_cache_size=L2_cache_size,
|
||||
config=config,
|
||||
epilogue_nodes=epilogues,
|
||||
GemmOuts=gemm_output_buffers,
|
||||
reindexers=reindexers,
|
||||
kernel_args=kernel_args,
|
||||
X_list=X_list,
|
||||
W_list=W_list,
|
||||
gemm_grouped_num=self.gemm_grouped_num,
|
||||
Y_list={"Y" + str(idx): Y for idx, Y in enumerate(Y_list)},
|
||||
Y_2d_list=Y_2d_list,
|
||||
multi_output_buffers=multi_output_buffers,
|
||||
cpu_count=os.cpu_count(),
|
||||
)
|
||||
with contextlib.ExitStack() as stack:
|
||||
stack.enter_context(
|
||||
patch.object(V.graph, "get_dtype", self._fake_get_dtype(fake_buffers))
|
||||
)
|
||||
return self._template_from_string(GEMM_TEMPLATE).render(**options)
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,139 @@
|
||||
# mypy: allow-untyped-defs
|
||||
import ctypes
|
||||
import functools
|
||||
import itertools
|
||||
import logging
|
||||
import sys
|
||||
from collections.abc import Callable, Iterable
|
||||
from unittest.mock import patch
|
||||
|
||||
import sympy
|
||||
|
||||
from .. import config, ir
|
||||
from ..autotune_process import CppBenchmarkRequest, TensorMeta
|
||||
from ..utils import IndentedBuffer, Placeholder, unique
|
||||
from ..virtualized import V
|
||||
from .common import KernelTemplate
|
||||
from .cpp_template_kernel import CppTemplateCaller, CppTemplateKernel
|
||||
|
||||
|
||||
log = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class CppTemplate(KernelTemplate):
|
||||
index_counter = itertools.count()
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
name: str,
|
||||
input_nodes,
|
||||
layout: ir.Layout,
|
||||
num_threads: int,
|
||||
epilogue_creator: Callable[[ir.Buffer], ir.Pointwise] | None = None,
|
||||
) -> None:
|
||||
super().__init__(name)
|
||||
self.input_nodes = input_nodes
|
||||
self.index = next(self.index_counter)
|
||||
self.output_node: ir.Buffer | list[ir.Buffer] = ir.Buffer(
|
||||
name=f"buf_out{self.index}", layout=layout
|
||||
)
|
||||
self.layout = layout
|
||||
self.num_threads = num_threads
|
||||
self.epilogue_creator = epilogue_creator
|
||||
|
||||
def generate(self, **kwargs):
|
||||
kernel_name = f"cpp_{self.name}"
|
||||
with (
|
||||
patch.object(V.graph, "get_dtype", self._fake_get_dtype(self.output_node)),
|
||||
patch.object(ir.FlexibleLayout, "allow_indexing", True),
|
||||
V.graph.set_current_device(self.layout.device),
|
||||
CppTemplateKernel(
|
||||
kernel_name=kernel_name, num_threads=self.num_threads
|
||||
) as kernel,
|
||||
):
|
||||
code = kernel.render(self, **kwargs)
|
||||
_, call_args, _, _ = kernel.args.python_argdefs()
|
||||
log.debug("Generated Code:\n%s", code)
|
||||
log.debug(
|
||||
"Args: cpp_argdefs: %s, python_argdefs: %s",
|
||||
kernel.args.cpp_argdefs(),
|
||||
kernel.args.python_argdefs(),
|
||||
)
|
||||
|
||||
expected_args = list(
|
||||
unique(input_node.get_name() for input_node in self.input_nodes)
|
||||
)
|
||||
if isinstance(self.output_node, Iterable):
|
||||
expected_args.extend([node.get_name() for node in self.output_node])
|
||||
else:
|
||||
expected_args.extend([self.output_node.get_name()])
|
||||
assert list(call_args)[: len(expected_args)] == expected_args, (
|
||||
call_args,
|
||||
expected_args,
|
||||
)
|
||||
# extra_args are only used for benchmarking, not compiled kernel correctness
|
||||
extra_args = V.graph.sizevars.optimization_hints(
|
||||
map(sympy.expand, call_args[len(expected_args) :])
|
||||
)
|
||||
# Cast the size hint from int to ctypes.c_ulonglong explicitly
|
||||
# since in cpp kernel, we bind it to C long
|
||||
extra_args = tuple(ctypes.c_ulonglong(x) for x in extra_args)
|
||||
|
||||
kernel_hash_name = f"cpp_{self.name}_{self.index}"
|
||||
|
||||
# Create the BenchmarkRequest for CPP
|
||||
bmreq = CppBenchmarkRequest(
|
||||
kernel_name=kernel_name,
|
||||
input_tensor_meta=TensorMeta.from_irnodes(self.input_nodes),
|
||||
output_tensor_meta=TensorMeta.from_irnodes(self.output_node),
|
||||
extra_args=extra_args,
|
||||
source_code=code,
|
||||
)
|
||||
|
||||
def make_kernel_render(
|
||||
template_node: ir.CppTemplateBuffer,
|
||||
flag_template_buffer_has_other_users: bool,
|
||||
epilogue_nodes: list[ir.IRNode] | None = None,
|
||||
):
|
||||
kernel = CppTemplateKernel(
|
||||
kernel_name=str(Placeholder.KERNEL_NAME), num_threads=self.num_threads
|
||||
)
|
||||
render = functools.partial(
|
||||
kernel.render,
|
||||
self,
|
||||
template_buffer_node=template_node,
|
||||
flag_template_buffer_has_other_users=flag_template_buffer_has_other_users,
|
||||
epilogue_nodes=epilogue_nodes,
|
||||
**kwargs,
|
||||
)
|
||||
return kernel, render
|
||||
|
||||
return CppTemplateCaller(
|
||||
kernel_hash_name,
|
||||
self.name,
|
||||
self.input_nodes,
|
||||
# pyrefly: ignore [bad-index, index-error]
|
||||
self.output_node[0].get_layout()
|
||||
if isinstance(self.output_node, Iterable)
|
||||
else self.output_node.get_layout(),
|
||||
make_kernel_render,
|
||||
bmreq,
|
||||
self,
|
||||
)
|
||||
|
||||
def header(self) -> IndentedBuffer:
|
||||
res = IndentedBuffer()
|
||||
res.writeline("#include <torch/csrc/inductor/cpp_prefix.h>")
|
||||
# TODO: add c10::ForcedUnroll test to test_aoti_abi_check
|
||||
res.splice("""#include <c10/util/Unroll.h>""")
|
||||
res.splice("""#include <torch/csrc/inductor/aoti_torch/c/shim.h>""")
|
||||
enable_kernel_profile = config.cpp.enable_kernel_profile and sys.platform in [
|
||||
"linux",
|
||||
"win32",
|
||||
]
|
||||
if enable_kernel_profile:
|
||||
res.writelines(["#include <torch/csrc/inductor/aoti_runtime/utils.h>"])
|
||||
return res
|
||||
|
||||
def render(self, **kwargs) -> str:
|
||||
raise NotImplementedError
|
||||
+621
@@ -0,0 +1,621 @@
|
||||
# mypy: allow-untyped-defs
|
||||
import itertools
|
||||
from collections.abc import Callable, Iterable
|
||||
from typing import Any
|
||||
from unittest.mock import patch
|
||||
|
||||
import sympy
|
||||
from sympy.parsing.sympy_parser import parse_expr
|
||||
|
||||
import torch
|
||||
from torch._inductor.utils import do_bench_using_profiling
|
||||
from torch.utils._ordered_set import OrderedSet
|
||||
from torch.utils._sympy.symbol import SymT
|
||||
|
||||
from .. import config, cpp_builder, ir, lowering as L
|
||||
from ..autotune_process import CppBenchmarkRequest
|
||||
from ..loop_body import LoopBody
|
||||
from ..select_algorithm import PartialRender
|
||||
from ..utils import sympy_index_symbol, sympy_index_symbol_with_prefix
|
||||
from ..virtualized import V
|
||||
from .common import REMOVED
|
||||
from .cpp import CppKernel, CppKernelProxy, KernelGroup, ParallelDepth
|
||||
from .cpp_utils import cexpr_index, DTYPE_TO_CPP, LocalBufferContext
|
||||
|
||||
|
||||
def parse_expr_with_index_symbols(expr):
|
||||
if isinstance(expr, sympy.Expr):
|
||||
return expr
|
||||
elif isinstance(expr, (list, tuple)):
|
||||
return [parse_expr_with_index_symbols(e) for e in expr]
|
||||
else:
|
||||
expr = parse_expr(str(expr))
|
||||
int_symbols = {sym: sympy_index_symbol(sym.name) for sym in expr.free_symbols}
|
||||
return expr.subs(int_symbols)
|
||||
|
||||
|
||||
def wrap_with_tensorbox(node) -> ir.TensorBox:
|
||||
return (
|
||||
ir.TensorBox.create(node) if isinstance(node, ir.Buffer) else ir.TensorBox(node)
|
||||
)
|
||||
|
||||
|
||||
class CppTemplateKernel(CppKernel):
|
||||
def __init__(self, kernel_name, num_threads):
|
||||
super().__init__(None, num_threads)
|
||||
self.kernel_name = kernel_name
|
||||
self.render_hooks = {}
|
||||
self.local_buffers = {}
|
||||
|
||||
def render(self, template, **kwargs):
|
||||
return PartialRender(
|
||||
template.render(kernel=self, **kwargs), self.render_hooks
|
||||
).finalize_all()
|
||||
|
||||
def def_kernel(
|
||||
self,
|
||||
inputs: dict[str, ir.Buffer],
|
||||
outputs: dict[str, ir.Buffer],
|
||||
aliases: dict[str, str] | None = None,
|
||||
function_name: str = "",
|
||||
extra_sizevars: list[sympy.Expr] | None = None,
|
||||
placeholder: str = "<DEF_KERNEL>",
|
||||
) -> str:
|
||||
if len(function_name) == 0:
|
||||
function_name = str(self.kernel_name)
|
||||
for name, inp in inputs.items():
|
||||
if inp is not None:
|
||||
self.args.input_buffers[inp.get_name()] = name
|
||||
for name, out in outputs.items():
|
||||
self.args.output_buffers[out.get_name()] = name
|
||||
if aliases is not None:
|
||||
for alias, orig in aliases.items():
|
||||
if orig in self.args.input_buffers:
|
||||
self.args.input_buffers[alias] = self.args.input_buffers[orig]
|
||||
if orig in self.args.output_buffers:
|
||||
self.args.output_buffers[alias] = self.args.output_buffers[orig]
|
||||
|
||||
unique_sizevars = OrderedSet(
|
||||
s
|
||||
for input in inputs.values()
|
||||
if input is not None
|
||||
for sym in itertools.chain(input.get_size(), input.get_stride())
|
||||
if isinstance(sym, sympy.Expr)
|
||||
for s in sym.free_symbols
|
||||
)
|
||||
unique_sizevars.update(
|
||||
s
|
||||
for sym in extra_sizevars or []
|
||||
if isinstance(sym, sympy.Expr)
|
||||
for s in sym.free_symbols
|
||||
)
|
||||
unique_sizevars.update(
|
||||
s
|
||||
for output in outputs.values()
|
||||
for sym in itertools.chain(output.get_size(), output.get_stride())
|
||||
if isinstance(sym, sympy.Expr)
|
||||
for s in sym.free_symbols
|
||||
)
|
||||
sizevars = sorted(unique_sizevars, key=str)
|
||||
for sizevar in sizevars:
|
||||
self.args.sizevars[sizevar] = f"k{sizevar}"
|
||||
|
||||
def hook():
|
||||
# remove all aliases before generate function definition
|
||||
if aliases is not None:
|
||||
for alias in aliases:
|
||||
if alias in self.args.input_buffers:
|
||||
raise AssertionError(
|
||||
f"input_buffers cannot be removed: {alias}"
|
||||
)
|
||||
if alias in self.args.output_buffers:
|
||||
self.args.output_buffers[alias] = REMOVED
|
||||
cpp_argdefs, _, _ = self.args.cpp_argdefs()
|
||||
return f"void {function_name}({', '.join(cpp_argdefs)})"
|
||||
|
||||
assert placeholder not in self.render_hooks
|
||||
self.render_hooks[placeholder] = hook
|
||||
return placeholder
|
||||
|
||||
def call_kernel(self, name: str, node: ir.CppTemplateBuffer):
|
||||
wrapper = V.graph.wrapper_code
|
||||
_, call_args, arg_types = self.args.cpp_argdefs()
|
||||
wrapper.generate_kernel_call(name, call_args, triton=False, arg_types=arg_types)
|
||||
|
||||
def dtype(self, node: ir.Buffer) -> str:
|
||||
return DTYPE_TO_CPP[node.get_dtype()]
|
||||
|
||||
def acc_dtype(self, node: ir.Buffer) -> str:
|
||||
if node.get_dtype() in [torch.float32, torch.bfloat16, torch.half]:
|
||||
return "float"
|
||||
else:
|
||||
raise NotImplementedError(f"Unsupported dtype: {node.get_dtype()}")
|
||||
|
||||
def size(self, node: ir.Buffer, dim: int) -> str:
|
||||
return cexpr_index(self.rename_indexing(node.get_size()[dim]))
|
||||
|
||||
def stride(self, node: ir.Buffer, dim: int) -> str:
|
||||
return cexpr_index(self.rename_indexing(node.get_stride()[dim]))
|
||||
|
||||
def index(self, node: ir.Buffer, indices: list[Any]) -> str:
|
||||
indexer = node.get_layout().as_fixed().make_indexer()
|
||||
index = indexer(parse_expr_with_index_symbols(indices))
|
||||
index = self.rename_indexing(index)
|
||||
outer_name = node.get_name()
|
||||
inner_name = (
|
||||
outer_name
|
||||
if outer_name in self.local_buffers
|
||||
else self.args.input(node.get_name())
|
||||
)
|
||||
return f"{inner_name}[{cexpr_index(index)}]"
|
||||
|
||||
def slice_nd(self, node, ranges: list[tuple[Any, Any]]) -> ir.ReinterpretView:
|
||||
"""
|
||||
Slice the given node with a list of ranges (start and end) corresponding to its dims.
|
||||
The dim is not sliced if the corresponding range is empty.
|
||||
"""
|
||||
assert len(ranges) == len(node.get_size()), f"{ranges=}, {node=}"
|
||||
sliced = wrap_with_tensorbox(node)
|
||||
for dim, _range in enumerate(ranges):
|
||||
if len(_range) == 0:
|
||||
continue
|
||||
assert len(_range) == 2
|
||||
start, end = parse_expr_with_index_symbols(_range)
|
||||
sliced = L.slice_(sliced, dim, start, end, clamp=False)
|
||||
assert isinstance(sliced, ir.TensorBox)
|
||||
assert isinstance(sliced.data, ir.ReinterpretView), sliced.data
|
||||
return sliced.data
|
||||
|
||||
def select(self, node, dim: int, idx: int) -> ir.ReinterpretView:
|
||||
# We avoid using L.select here because we need clamp=False so the dim after slicing
|
||||
# is 1 instead of a sympy expression of symbol - dim_size.
|
||||
node = wrap_with_tensorbox(node)
|
||||
idx = ir.View.handle_negative_index(idx, node.get_size()[dim])
|
||||
sliced = L.squeeze(L.slice_(node, dim, idx, idx + 1, clamp=False), dim)
|
||||
assert isinstance(sliced.data, ir.ReinterpretView), sliced.data
|
||||
return sliced.data
|
||||
|
||||
def view(self, node, sizes: list[Any]) -> ir.IRNode:
|
||||
node = wrap_with_tensorbox(node)
|
||||
sizes = parse_expr_with_index_symbols(sizes)
|
||||
return L.view(node, sizes).data # type: ignore[arg-type]
|
||||
|
||||
def permute(self, node, dims):
|
||||
node = wrap_with_tensorbox(node)
|
||||
permuted = L.permute(node, dims).data
|
||||
assert isinstance(permuted, ir.ReinterpretView)
|
||||
return permuted
|
||||
|
||||
def maybe_codegen_profile(self) -> str:
|
||||
if config.cpp.enable_kernel_profile:
|
||||
graph_id = V.graph.graph_id
|
||||
prefix = "graph_" + str(graph_id) + "_" if graph_id is not None else ""
|
||||
handle_str = (
|
||||
"torch::aot_inductor::RAIIAtenRecordFunctionHandle "
|
||||
f'record_{prefix}{self.kernel_name}_("{prefix}{self.kernel_name}", nullptr);'
|
||||
)
|
||||
return handle_str
|
||||
else:
|
||||
return ""
|
||||
|
||||
def unroll_pragma(self, unroll):
|
||||
if cpp_builder.is_gcc():
|
||||
return f"#pragma GCC unroll {unroll}"
|
||||
else:
|
||||
return f"#pragma unroll {unroll}"
|
||||
|
||||
def define_buffer(self, name, sizes: list[Any], dtype=torch.float) -> str:
|
||||
"""Define kernel local buffer"""
|
||||
sizes = parse_expr_with_index_symbols(sizes)
|
||||
buf = ir.Buffer(
|
||||
name=name, layout=ir.FixedLayout(torch.device("cpu"), dtype, sizes)
|
||||
)
|
||||
self.local_buffers[name] = buf
|
||||
ctype = f"{DTYPE_TO_CPP[dtype]}"
|
||||
numel = f"{cexpr_index(buf.get_numel())}"
|
||||
return f"auto _{name} = std::make_unique<{ctype}[]>({numel}); auto {name} = _{name}.get();"
|
||||
|
||||
def define_stack_allocated_buffer(
|
||||
self, name, sizes: list[Any], dtype=torch.float
|
||||
) -> str:
|
||||
"""Define stack-allocated buffer"""
|
||||
sizes = parse_expr_with_index_symbols(sizes)
|
||||
buf = ir.Buffer(
|
||||
name=name, layout=ir.FixedLayout(torch.device("cpu"), dtype, sizes)
|
||||
)
|
||||
self.local_buffers[name] = buf
|
||||
ctype = f"{DTYPE_TO_CPP[dtype]}"
|
||||
numel = f"{cexpr_index(buf.get_numel())}"
|
||||
return f"alignas(64) {ctype} _{name}[{numel}]; {ctype}* {name} = _{name};"
|
||||
|
||||
def reinit_buffer_if_null(self, name):
|
||||
"""Reinit the previously defined local buffer if it is null"""
|
||||
assert name in self.local_buffers
|
||||
buf = self.local_buffers[name]
|
||||
ctype = f"{DTYPE_TO_CPP[buf.layout.dtype]}"
|
||||
numel = f"{cexpr_index(buf.get_numel())}"
|
||||
return f"if (_{name} == nullptr) {{ _{name} = std::make_unique<{ctype}[]>({numel}); {name} = _{name}.get(); }}"
|
||||
|
||||
def release_buffer(self, name):
|
||||
"""Codegen the code to release the ownership of a local buffer to others"""
|
||||
assert name in self.local_buffers
|
||||
return f"_{name}.release()"
|
||||
|
||||
def store_pointwise_nodes(
|
||||
self,
|
||||
dst: ir.Buffer,
|
||||
nodes: list[ir.IRNode],
|
||||
offsets: list[sympy.Expr] | None = None,
|
||||
reindexers: list[Callable[[list[Any]], list[Any]] | None] | None = None,
|
||||
) -> str:
|
||||
var_sizes = (tuple(dst.get_size()), ())
|
||||
var_ranges = {
|
||||
sympy_index_symbol_with_prefix(SymT.INDEX, i): sz
|
||||
for i, sz in enumerate(var_sizes[0])
|
||||
}
|
||||
if not offsets:
|
||||
offsets = [sympy.S.Zero] * len(var_sizes[0])
|
||||
if not reindexers:
|
||||
reindexers = [None] * len(nodes)
|
||||
assert len(offsets) == len(var_sizes[0])
|
||||
output_index = dst.get_layout().make_indexer()([*var_ranges.keys()])
|
||||
kernel_group = KernelGroup()
|
||||
kernel_group.args = self.args
|
||||
cpp_kernel_proxy = CppKernelProxy(kernel_group)
|
||||
bodies = []
|
||||
var_sizes_list = []
|
||||
for i, node in enumerate(nodes):
|
||||
output_name = node.get_name() if i < len(nodes) - 1 else dst.get_name()
|
||||
node = node.data if isinstance(node, ir.ComputedBuffer) else node
|
||||
assert isinstance(node, ir.Pointwise), node
|
||||
|
||||
def fn(*args):
|
||||
assert len(args) == 2
|
||||
assert len(args[0]) == len(var_sizes[0])
|
||||
assert len(args[1]) == 0
|
||||
new_args = [arg + offset for arg, offset in zip(args[0], offsets)] # type: ignore[arg-type]
|
||||
if reindexers[i] is not None:
|
||||
new_args = reindexers[i](new_args) # type: ignore[misc]
|
||||
V.ops.store(
|
||||
output_name,
|
||||
output_index,
|
||||
node.make_loader()(new_args).value,
|
||||
)
|
||||
|
||||
body = LoopBody(
|
||||
fn,
|
||||
(list(var_ranges.keys()), ()),
|
||||
var_ranges,
|
||||
list(var_ranges.keys()),
|
||||
tuple(),
|
||||
)
|
||||
bodies.append(body)
|
||||
var_sizes_list.append(var_sizes)
|
||||
|
||||
cpp_kernel_proxy.codegen_loop_bodies(bodies, var_sizes_list)
|
||||
|
||||
def max_parallel_depth():
|
||||
return ParallelDepth(parallel_depth=0, start_depth=0)
|
||||
|
||||
# This loop is not parallelized since it is not the outermost loop.
|
||||
with patch.object(
|
||||
cpp_kernel_proxy.loop_nest, "max_parallel_depth", max_parallel_depth
|
||||
):
|
||||
kernel_group.finalize_kernel(cpp_kernel_proxy, [])
|
||||
return kernel_group.loops_code.getvalue()
|
||||
|
||||
def store_grouped_gemm_pointwise_nodes(
|
||||
self,
|
||||
dst: tuple[ir.Buffer],
|
||||
nodes: list[ir.IRNode],
|
||||
offsets: list[sympy.Expr],
|
||||
reindexers: list[Callable[[list[Any]], list[Any]] | None],
|
||||
output_names: list[str],
|
||||
) -> str:
|
||||
ref_dst = dst[0]
|
||||
var_sizes = (tuple(ref_dst.get_size()), ())
|
||||
var_ranges = {
|
||||
sympy_index_symbol_with_prefix(SymT.INDEX, i): sz
|
||||
for i, sz in enumerate(var_sizes[0])
|
||||
}
|
||||
assert offsets, "offsets should be set outside"
|
||||
assert all(len(offset) == len(var_sizes[0]) for offset in offsets)
|
||||
output_index = ref_dst.get_layout().make_indexer()([*var_ranges.keys()])
|
||||
kernel_group = KernelGroup()
|
||||
kernel_group.args = self.args
|
||||
cpp_kernel_proxy = CppKernelProxy(kernel_group)
|
||||
bodies = []
|
||||
var_sizes_list = []
|
||||
for i, node in enumerate(nodes):
|
||||
output_name = output_names[i]
|
||||
node = node.data if isinstance(node, ir.ComputedBuffer) else node
|
||||
assert isinstance(node, ir.Pointwise), node
|
||||
|
||||
def fn(*args):
|
||||
assert len(args) == 2
|
||||
assert len(args[0]) == len(var_sizes[0])
|
||||
assert len(args[1]) == 0
|
||||
new_args = [arg + offset for arg, offset in zip(args[0], offsets[i])] # type: ignore[arg-type]
|
||||
if reindexers[i] is not None:
|
||||
new_args = reindexers[i](new_args) # type: ignore[misc]
|
||||
V.ops.store(
|
||||
output_name,
|
||||
output_index,
|
||||
node.make_loader()(new_args).value,
|
||||
)
|
||||
|
||||
body = LoopBody(
|
||||
fn,
|
||||
(list(var_ranges.keys()), ()),
|
||||
var_ranges,
|
||||
list(var_ranges.keys()),
|
||||
tuple(),
|
||||
)
|
||||
bodies.append(body)
|
||||
var_sizes_list.append(var_sizes)
|
||||
|
||||
cpp_kernel_proxy.codegen_loop_bodies(bodies, var_sizes_list)
|
||||
|
||||
def max_parallel_depth():
|
||||
return ParallelDepth(parallel_depth=0, start_depth=0)
|
||||
|
||||
# This loop is not parallelized since it is not the outermost loop.
|
||||
with patch.object(
|
||||
cpp_kernel_proxy.loop_nest, "max_parallel_depth", max_parallel_depth
|
||||
):
|
||||
kernel_group.finalize_kernel(cpp_kernel_proxy, [])
|
||||
return kernel_group.loops_code.getvalue()
|
||||
|
||||
def store_output(
|
||||
self,
|
||||
dst: ir.Buffer,
|
||||
src: ir.Buffer,
|
||||
orig_src: ir.Buffer | None = None,
|
||||
epilogue_nodes: list[ir.IRNode] | None = None,
|
||||
offsets: list[Any] | None = None,
|
||||
reindexers: list[Callable[[list[Any]], list[Any]] | None] | None = None,
|
||||
):
|
||||
"""
|
||||
Store the `src` buffer to the `dst` buffer. The size of `src` and `dst` should match.
|
||||
If `epilogue_nodes` is provided, the `src` buffer is firstly computed with the epilogues
|
||||
before stored to `dst`. The `epilogues_nodes` are all pointwise.
|
||||
|
||||
Notes:
|
||||
1. `src` and `dst` buffer could be the same buffer in which case we are doing in-place compute
|
||||
and stores. In case `epilogue_nodes` are not provided, we do nothing.
|
||||
2. The `epilogue_nodes`, if exist, have computations on `src` before storing to `dst` but since
|
||||
they come form the original Inductor IR, they might need to be adjusted before working with
|
||||
`src` and `dst` as outlined below:
|
||||
a) `src` or `dst` buffer could be a sub-slice of the ranges the `epilogue_nodes`work on.
|
||||
In this case, the `offsets` could be provided to adjust the indices passed to
|
||||
`epilogue_nodes` during codegen and the data ranges are also configured according to
|
||||
the sizes of `src` and `dst`.
|
||||
b) `dst` might be indexed in a different way as the `epilogue_nodes`, hence a `reindexer` is
|
||||
needed on the indices to `epilogue_nodes` to match the indexing of `dst`.
|
||||
c) If `src` is local, we need to add a local buffer for it and localize the `orig_src` buffer
|
||||
in `epilogue_nodes` with `src`.
|
||||
"""
|
||||
assert isinstance(dst, (ir.Buffer, ir.ReinterpretView))
|
||||
assert dst.get_size() == src.get_size(), f"{dst=}, {src=}"
|
||||
if offsets:
|
||||
offsets = parse_expr_with_index_symbols(offsets)
|
||||
if epilogue_nodes:
|
||||
with LocalBufferContext(self.args) as scope:
|
||||
assert orig_src is not None
|
||||
if orig_src.get_name() != src.get_name():
|
||||
scope.add_local_buffer(
|
||||
src,
|
||||
[
|
||||
orig_src,
|
||||
],
|
||||
)
|
||||
epilogue_nodes = scope.localize_nodes(epilogue_nodes)
|
||||
return self.store_pointwise_nodes(
|
||||
dst,
|
||||
epilogue_nodes, # type: ignore[arg-type]
|
||||
offsets,
|
||||
reindexers,
|
||||
)
|
||||
else:
|
||||
if dst.get_name() != src.get_name():
|
||||
# src is local
|
||||
copy = L.copy(dst, src).data.data
|
||||
with LocalBufferContext(self.args) as scope:
|
||||
scope.add_local_buffer(src)
|
||||
|
||||
return self.store_pointwise_nodes(dst, [copy])
|
||||
else:
|
||||
assert dst.layout == src.layout, f"{dst=}, {src=}"
|
||||
return ""
|
||||
|
||||
def store_outputs(
|
||||
self,
|
||||
dst: tuple[ir.Buffer],
|
||||
src: tuple[ir.IRNode],
|
||||
orig_src: tuple[ir.IRNode] | None = None,
|
||||
epilogue_nodes: list[ir.IRNode] | None = None,
|
||||
offsets: list[Any] | None = None,
|
||||
reindexers: list[Callable[[list[Any]], list[Any]] | None] | None = None,
|
||||
multi_output_buffers: tuple[ir.MultiOutput, ...] | None = None,
|
||||
):
|
||||
assert isinstance(dst, Iterable)
|
||||
assert all(_dst.get_size() == _src.get_size() for _src, _dst in zip(src, dst))
|
||||
if offsets:
|
||||
offsets = parse_expr_with_index_symbols(offsets)
|
||||
gemm_num = len(src)
|
||||
final_offsets = []
|
||||
output_names = []
|
||||
if epilogue_nodes:
|
||||
if not reindexers:
|
||||
reindexers = [None] * len(epilogue_nodes)
|
||||
with LocalBufferContext(self.args) as scope:
|
||||
assert orig_src is not None
|
||||
localize_epilogue_nodes = []
|
||||
all_read_names = []
|
||||
for epilogue in epilogue_nodes:
|
||||
all_read_names.extend(list(epilogue.get_read_names()))
|
||||
localize_epilogue_nodes.extend(scope.localize_nodes(epilogue_nodes))
|
||||
final_offsets.extend([offsets] * len(localize_epilogue_nodes))
|
||||
output_names.extend(
|
||||
[node.get_name() for node in localize_epilogue_nodes]
|
||||
)
|
||||
for gemm_idx in range(gemm_num):
|
||||
if orig_src[gemm_idx].get_name() != src[gemm_idx].get_name():
|
||||
if orig_src[gemm_idx].get_name() in all_read_names or (
|
||||
multi_output_buffers
|
||||
and multi_output_buffers[gemm_idx].get_name()
|
||||
in all_read_names
|
||||
):
|
||||
# If any of the Epilogue nodes use this GEMM output, let's localize the GEMM output
|
||||
global_buffers = [orig_src[gemm_idx]]
|
||||
if (
|
||||
multi_output_buffers
|
||||
and multi_output_buffers[gemm_idx].get_name()
|
||||
in all_read_names
|
||||
and orig_src[gemm_idx].get_name() not in all_read_names
|
||||
):
|
||||
# Epilogue might directly read the MultiOutput, Locallize MultiOutput to the local Buffer
|
||||
# if this MultiOutput has not been stored by in-template epilogue
|
||||
# otherwise, use the cse store cache if it will be stored before used
|
||||
global_buffers.append(multi_output_buffers[gemm_idx])
|
||||
scope.add_local_buffer(
|
||||
src[gemm_idx],
|
||||
global_buffers,
|
||||
)
|
||||
else:
|
||||
scope.add_local_buffer(src[gemm_idx])
|
||||
localize_epilogue_nodes.extend(
|
||||
[L.copy(dst[gemm_idx], src[gemm_idx]).data.data]
|
||||
)
|
||||
reindexers.append(None)
|
||||
output_names.append(dst[gemm_idx].get_name())
|
||||
final_offsets.append(
|
||||
[sympy.S.Zero] * len(dst[gemm_idx].get_size())
|
||||
)
|
||||
res = self.store_grouped_gemm_pointwise_nodes(
|
||||
dst,
|
||||
localize_epilogue_nodes,
|
||||
final_offsets,
|
||||
reindexers,
|
||||
output_names=output_names,
|
||||
)
|
||||
for gemm_idx in range(gemm_num):
|
||||
if (
|
||||
multi_output_buffers
|
||||
and multi_output_buffers[gemm_idx].get_name() in all_read_names
|
||||
):
|
||||
# If the MultiOutput is used in the Epilogue, let's remove it from args
|
||||
multi_output_name = multi_output_buffers[gemm_idx].get_name()
|
||||
if (
|
||||
multi_output_name in self.args.output_buffers
|
||||
and self.args.output_buffers[multi_output_name]
|
||||
is not REMOVED
|
||||
):
|
||||
self.remove_buffer(multi_output_name)
|
||||
return res
|
||||
else:
|
||||
if dst[0].get_name() != src[0].get_name():
|
||||
copy_list = []
|
||||
with LocalBufferContext(self.args) as scope:
|
||||
for _src, _dst in zip(src, dst):
|
||||
copy_list.extend([L.copy(_dst, _src).data.data])
|
||||
scope.add_local_buffer(_src)
|
||||
output_names.append(_dst.get_name())
|
||||
final_offsets.append([sympy.S.Zero] * len(_dst.get_size()))
|
||||
reindexers = [None] * len(copy_list)
|
||||
return self.store_grouped_gemm_pointwise_nodes(
|
||||
dst,
|
||||
nodes=copy_list,
|
||||
offsets=final_offsets,
|
||||
reindexers=reindexers,
|
||||
output_names=output_names,
|
||||
)
|
||||
else:
|
||||
assert all(
|
||||
_src.get_name() == _dst.get_name() for _src, _dst in zip(src, dst)
|
||||
)
|
||||
assert all(
|
||||
_src.get_layout() == _dst.get_layout()
|
||||
for _src, _dst in zip(src, dst)
|
||||
)
|
||||
return ""
|
||||
|
||||
def check_bounds(self, expr, size, lower, upper):
|
||||
# CppTemplateKernel does not need codegen related operations
|
||||
return
|
||||
|
||||
|
||||
class CppTemplateCaller(ir.ChoiceCaller):
|
||||
"""
|
||||
CppTemplateCaller
|
||||
|
||||
This class represents a caller for CPP template kernels. It is a subclass of ir.ChoiceCaller.
|
||||
Attributes:
|
||||
name (str): The name of the caller.
|
||||
category (str): The category of the caller.
|
||||
bmreq (CppBenchmarkRequest): The benchmark request for the caller.
|
||||
template_buffer (ir.CppTemplateBuffer): The template buffer for the caller.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
name: str,
|
||||
category: str,
|
||||
input_nodes: list[ir.Buffer],
|
||||
layout: ir.Layout,
|
||||
make_kernel_render: Callable[
|
||||
[
|
||||
ir.CppTemplateBuffer,
|
||||
bool,
|
||||
list[ir.IRNode] | None,
|
||||
],
|
||||
str,
|
||||
],
|
||||
bmreq: CppBenchmarkRequest,
|
||||
template: "CppTemplate", # type: ignore[name-defined] # noqa: F821
|
||||
info_kwargs: dict[str, ir.PrimitiveInfoType | list[ir.PrimitiveInfoType]]
|
||||
| None = None,
|
||||
):
|
||||
super().__init__(name, input_nodes, layout, description="")
|
||||
self.category = category
|
||||
self.make_kernel_render = make_kernel_render
|
||||
self.bmreq = bmreq
|
||||
self.template = template
|
||||
self.info_kwargs = info_kwargs
|
||||
|
||||
def precompile(self) -> None:
|
||||
assert self.bmreq is not None
|
||||
self.bmreq.precompile()
|
||||
|
||||
def benchmark(self, *args, out) -> float:
|
||||
assert self.bmreq is not None
|
||||
if config.profile_bandwidth_with_do_bench_using_profiling:
|
||||
algo = self.bmreq.make_run_fn(*args, out=out)
|
||||
return do_bench_using_profiling(algo)
|
||||
return self.bmreq.benchmark(*args, out=out)
|
||||
|
||||
def hash_key(self) -> str:
|
||||
return "-".join(
|
||||
[
|
||||
self.category,
|
||||
self.bmreq.hash_key,
|
||||
]
|
||||
)
|
||||
|
||||
def info_dict(
|
||||
self,
|
||||
) -> dict[str, ir.PrimitiveInfoType | list[ir.PrimitiveInfoType]]:
|
||||
return {"backend": "CPP", "op_type": "unknown"}
|
||||
|
||||
def output_node(self) -> ir.TensorBox:
|
||||
buffer = ir.CppTemplateBuffer(
|
||||
layout=self.layout,
|
||||
inputs=self.input_nodes,
|
||||
make_kernel_render=self.make_kernel_render,
|
||||
template=self.template,
|
||||
choice=self,
|
||||
)
|
||||
# Pass KTC annotation to the buffer for encoding
|
||||
if "ktc" in self.annotations:
|
||||
buffer.annotations["ktc"] = self.annotations["ktc"]
|
||||
return ir.TensorBox.create(buffer)
|
||||
@@ -0,0 +1,790 @@
|
||||
# mypy: allow-untyped-defs
|
||||
import contextlib
|
||||
import dataclasses
|
||||
import functools
|
||||
import math
|
||||
import sys
|
||||
from collections import namedtuple
|
||||
from collections.abc import Callable, Sequence
|
||||
from typing import Any
|
||||
from unittest.mock import patch
|
||||
|
||||
import sympy
|
||||
|
||||
import torch
|
||||
from torch._prims_common import is_integer_dtype
|
||||
from torch.utils._ordered_set import OrderedSet
|
||||
from torch.utils._sympy.printers import CppPrinter as _CppPrinter
|
||||
from torch.utils._sympy.symbol import symbol_is_type, SymT
|
||||
from torch.utils._sympy.value_ranges import ValueRanges
|
||||
|
||||
from .. import ir
|
||||
from ..dependencies import Dep
|
||||
from ..loop_body import LoopBody
|
||||
from ..scheduler import BaseSchedulerNode, SchedulerBuffer
|
||||
from ..shape_propagation import BlockShapeType
|
||||
from ..utils import IndentedBuffer, sympy_index_symbol_with_prefix, sympy_subs
|
||||
from ..virtualized import ops, OpsValue, V
|
||||
from .common import CSEVariable, Kernel, KernelArgs, OptimizationContext
|
||||
|
||||
|
||||
DTYPE_TO_CPP = {
|
||||
torch.float32: "float",
|
||||
torch.float64: "double",
|
||||
torch.float16: "at::Half",
|
||||
torch.int64: "int64_t",
|
||||
torch.int32: "int32_t",
|
||||
torch.int16: "int16_t",
|
||||
torch.int8: "int8_t",
|
||||
torch.uint64: "uint64_t",
|
||||
torch.uint32: "uint32_t",
|
||||
torch.uint16: "uint16_t",
|
||||
torch.uint8: "uint8_t",
|
||||
torch.bool: "bool",
|
||||
torch.bfloat16: "at::BFloat16",
|
||||
torch.complex32: "at::complex<at::Half>",
|
||||
torch.complex64: "at::complex<float>",
|
||||
torch.complex128: "at::complex<double>",
|
||||
torch.float8_e4m3fn: "at::Float8_e4m3fn",
|
||||
torch.float8_e5m2: "at::Float8_e5m2",
|
||||
torch.float8_e4m3fnuz: "at::Float8_e4m3fnuz",
|
||||
torch.float8_e5m2fnuz: "at::Float8_e5m2fnuz",
|
||||
}
|
||||
|
||||
DTYPE_TO_ATEN = {
|
||||
torch.float32: "at::kFloat",
|
||||
torch.float64: "at::kDouble",
|
||||
torch.float16: "at::kHalf",
|
||||
torch.int64: "at::kLong",
|
||||
torch.int32: "at::kInt",
|
||||
torch.int16: "at::kShort",
|
||||
torch.int8: "at::kChar",
|
||||
torch.uint64: "at::kUInt64",
|
||||
torch.uint32: "at::kUInt32",
|
||||
torch.uint16: "at::kUInt16",
|
||||
torch.uint8: "at::kByte",
|
||||
torch.uint32: "at::kUInt32",
|
||||
torch.uint64: "at::kUInt64",
|
||||
torch.bool: "at::kBool",
|
||||
torch.bfloat16: "at::kBFloat16",
|
||||
torch.complex32: "at::kComplexHalf",
|
||||
torch.complex64: "at::kComplexFloat",
|
||||
torch.complex128: "at::kComplexDouble",
|
||||
torch.float8_e4m3fn: "at::kFloat8_e4m3fn",
|
||||
torch.float8_e5m2: "at::kFloat8_e5m2",
|
||||
torch.float8_e4m3fnuz: "at::kFloat8_e4m3fnuz",
|
||||
torch.float8_e5m2fnuz: "at::kFloat8_e5m2fnuz",
|
||||
}
|
||||
|
||||
DEVICE_TO_ATEN = {
|
||||
"meta": "at::kMeta",
|
||||
"cpu": "at::kCPU",
|
||||
"cuda": "at::kCUDA",
|
||||
"xpu": "at::kXPU",
|
||||
"mps": "at::kMPS",
|
||||
}
|
||||
|
||||
LAYOUT_TO_ATEN = {
|
||||
torch.strided: "at::kStrided",
|
||||
torch._mkldnn: "at::kMkldnn", # type: ignore[attr-defined]
|
||||
}
|
||||
|
||||
# matches c10/core/DeviceType.h
|
||||
DEVICE_TO_INT = {"cpu": 0, "cuda": 1}
|
||||
|
||||
_IS_WINDOWS = sys.platform == "win32"
|
||||
|
||||
INDEX_TYPE = "int64_t"
|
||||
|
||||
GemmBlocking = namedtuple("GemmBlocking", ["block_m", "block_n", "block_k"])
|
||||
|
||||
|
||||
def get_promote_dtype(args):
|
||||
return (
|
||||
# pyrefly: ignore [no-matching-overload]
|
||||
functools.reduce(
|
||||
torch.promote_types, # type: ignore[arg-type]
|
||||
# pyrefly: ignore [bad-argument-type]
|
||||
[n.dtype for n in args if isinstance(n, CppCSEVariable)],
|
||||
)
|
||||
if all(n.dtype is not None for n in args if isinstance(n, CppCSEVariable))
|
||||
else None # not enough info to calculate the promote dtype
|
||||
)
|
||||
|
||||
|
||||
def promote_args(new_args):
|
||||
def promote_arg(arg, promote_type):
|
||||
if (
|
||||
isinstance(arg, CppCSEVariable)
|
||||
and arg.dtype
|
||||
and promote_type
|
||||
and arg.dtype != promote_type
|
||||
):
|
||||
arg = ops.to_dtype(arg, promote_type)
|
||||
arg = arg.value if isinstance(arg, OpsValue) else arg
|
||||
arg.dtype = promote_type
|
||||
return arg
|
||||
|
||||
promote_type = get_promote_dtype(new_args)
|
||||
promote_fn = functools.partial(
|
||||
promote_arg,
|
||||
promote_type=promote_type,
|
||||
)
|
||||
if (
|
||||
all(
|
||||
new_arg.dtype is not None
|
||||
for new_arg in new_args
|
||||
if isinstance(new_arg, CppCSEVariable)
|
||||
)
|
||||
and promote_type
|
||||
):
|
||||
new_args = list(map(promote_fn, new_args))
|
||||
return new_args
|
||||
|
||||
|
||||
class CppCSEVariable(CSEVariable):
|
||||
def __init__(
|
||||
self,
|
||||
name,
|
||||
bounds: ValueRanges[Any],
|
||||
dtype: torch.dtype | None = None,
|
||||
shape: BlockShapeType = None,
|
||||
) -> None:
|
||||
super().__init__(name, bounds, dtype, shape=shape)
|
||||
self.is_vec = False
|
||||
self.dependent_itervars = OrderedSet[sympy.Symbol]()
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return (
|
||||
f"CppCSEVariable(name: {self.name}, bounds: {self.bounds}, is_vec: {self.is_vec}, dtype: {self.dtype}, "
|
||||
f"dependent_itervars: {self.dependent_itervars})"
|
||||
)
|
||||
|
||||
def update_on_args(self, name, args, kwargs):
|
||||
if name == "load":
|
||||
# args[2] is index
|
||||
self._set_dependent_itervars(args[2])
|
||||
else:
|
||||
# propagate relevant itervars and is_vec from args
|
||||
self.dependent_itervars.update(
|
||||
*[
|
||||
arg.dependent_itervars
|
||||
for arg in args
|
||||
if isinstance(arg, CppCSEVariable)
|
||||
]
|
||||
)
|
||||
if name == "index_expr":
|
||||
self._set_dependent_itervars(args[0])
|
||||
if any(arg.is_vec for arg in args if isinstance(arg, CppCSEVariable)):
|
||||
self.is_vec = True
|
||||
|
||||
def _set_dependent_itervars(self, index: sympy.Expr):
|
||||
"""
|
||||
Set the relevant itervars for this variable based on the `index` expression.
|
||||
This includes the itervars directly used in the `index` as well as relevant itervars
|
||||
of other cse variables used in the `index`.
|
||||
"""
|
||||
for s in index.free_symbols:
|
||||
if s in V.kernel.itervars:
|
||||
self.dependent_itervars.add(s) # type: ignore[arg-type]
|
||||
elif s.name in V.kernel.cse.varname_map: # type: ignore[attr-defined]
|
||||
self.dependent_itervars.update(
|
||||
V.kernel.cse.varname_map[s.name].dependent_itervars # type: ignore[attr-defined]
|
||||
)
|
||||
|
||||
def depends_on(self, itervar: sympy.Symbol):
|
||||
return itervar in self.dependent_itervars
|
||||
|
||||
|
||||
class CppPrinter(_CppPrinter):
|
||||
def doprint(self, expr, *, simplify: bool = True, p=True):
|
||||
# TODO: why are people passing strings to the printer here :think:
|
||||
if simplify and isinstance(expr, sympy.Expr) and hasattr(V.graph, "sizevars"):
|
||||
expr = V.graph.sizevars.simplify(expr)
|
||||
return super().doprint(expr)
|
||||
|
||||
def parenthesize(self, item: sympy.Expr, level: int, strict: bool = False) -> str:
|
||||
if isinstance(item, sympy.Mod):
|
||||
# use parenthesis to enforce precedence.
|
||||
# in sympy 1.13.3, -2*Mod(x,y) becomes -2*x%y, which is wrong.
|
||||
# pyrefly: ignore [missing-attribute]
|
||||
return f"({self._print(item)})"
|
||||
else:
|
||||
return super().parenthesize(item, level, strict)
|
||||
|
||||
|
||||
# A function to print, useful for printing sympy symbols.
|
||||
cexpr = CppPrinter().doprint
|
||||
|
||||
|
||||
def cexpr_index(index):
|
||||
return f"static_cast<{INDEX_TYPE}>({cexpr(index)})"
|
||||
|
||||
|
||||
def value_to_cpp(value, cpp_type):
|
||||
if value == float("-inf"):
|
||||
return f"-std::numeric_limits<{cpp_type}>::infinity()"
|
||||
elif value == float("inf"):
|
||||
return f"std::numeric_limits<{cpp_type}>::infinity()"
|
||||
elif isinstance(value, bool):
|
||||
return f"static_cast<{cpp_type}>({str(value).lower()})"
|
||||
elif math.isnan(value):
|
||||
return f"std::numeric_limits<{cpp_type}>::quiet_NaN()"
|
||||
else:
|
||||
return f"static_cast<{cpp_type}>({repr(value)})"
|
||||
|
||||
|
||||
def rewrite_index_for_function(
|
||||
localize_buffer_handler: "LocalizeBufferHandler",
|
||||
index: sympy.Expr,
|
||||
global_buf_name: str,
|
||||
):
|
||||
# Local buffer at the inner dimensions
|
||||
snode = V.graph.scheduler.name_to_buf[global_buf_name].defining_op
|
||||
assert snode is not None
|
||||
local_buf = localize_buffer_handler.global_to_local[global_buf_name]
|
||||
scheduler_nodes = snode.get_nodes()
|
||||
_, (group, reduction_group) = max(
|
||||
scheduler_nodes, key=lambda x: int(x.is_reduction())
|
||||
).group
|
||||
call_ranges = tuple(group) + tuple(reduction_group)
|
||||
indices_to_keep = [
|
||||
f"x{len(call_ranges) - (idx + 1)}"
|
||||
for idx in range(len(local_buf.get_layout().size))
|
||||
]
|
||||
sorted_symbols = sorted(index.free_symbols, key=lambda s: s.name) # type: ignore[attr-defined]
|
||||
replacements = {}
|
||||
for x in sorted_symbols:
|
||||
if x.name.startswith("x") and x.name not in indices_to_keep: # type: ignore[attr-defined]
|
||||
# Only keep index used by local buffer
|
||||
replacements[x] = sympy.core.numbers.Zero()
|
||||
index = sympy_subs(index, replacements) # type: ignore[arg-type]
|
||||
return index
|
||||
|
||||
|
||||
def rewrite_index_for_nodes(
|
||||
localize_buffer_handler: "LocalizeBufferHandler",
|
||||
index: sympy.Expr,
|
||||
global_buf_name: str,
|
||||
):
|
||||
used_vars = OrderedSet(
|
||||
s for s in index.free_symbols if symbol_is_type(s, SymT.INDEX)
|
||||
)
|
||||
index_vars = []
|
||||
local_buf = localize_buffer_handler.global_to_local[global_buf_name]
|
||||
for i in range(len(local_buf.get_size())):
|
||||
var = sympy_index_symbol_with_prefix(SymT.INDEX, i)
|
||||
index_vars.append(var if var in used_vars else 0)
|
||||
index = local_buf.get_layout().make_indexer()(index_vars)
|
||||
return index
|
||||
|
||||
|
||||
class LocalizeBufferHandler(V.WrapperHandler): # type: ignore[name-defined]
|
||||
def __init__(
|
||||
self,
|
||||
inner,
|
||||
global_to_local: dict[str, ir.Buffer],
|
||||
rewrite_index: Callable[["LocalizeBufferHandler", sympy.Expr, str], sympy.Expr],
|
||||
) -> None:
|
||||
super().__init__(inner)
|
||||
self.global_to_local = global_to_local
|
||||
self.rewrite_index = rewrite_index
|
||||
|
||||
def localize(self, name: str, index: sympy.Expr):
|
||||
if self.global_to_local and name in self.global_to_local:
|
||||
assert self.rewrite_index is not None
|
||||
index = self.rewrite_index(self, index, name)
|
||||
name = self.global_to_local[name].get_name()
|
||||
return name, index
|
||||
|
||||
def load(self, name: str, index: sympy.Expr):
|
||||
return self._inner.load(*self.localize(name, index))
|
||||
|
||||
def store(self, name, index, value, mode=None):
|
||||
local_buffer_name, local_buffer_index = self.localize(name, index)
|
||||
res = self._inner.store(local_buffer_name, local_buffer_index, value, mode)
|
||||
if (
|
||||
self.global_to_local
|
||||
and name in self.global_to_local
|
||||
and isinstance(V.kernel, Kernel)
|
||||
):
|
||||
# Remove name of local buffer from Kernel.store_buffer_names
|
||||
# local_buffer_name is added to Kernel.store_buffer_names in Kernel.CSEProxy.store.
|
||||
V.kernel.store_buffer_names.discard(local_buffer_name)
|
||||
return res
|
||||
|
||||
def store_reduction(self, name, index, value):
|
||||
# pyrefly: ignore [bad-argument-count]
|
||||
return self._inner.store_reduction(*self.localize(name, index), value)
|
||||
|
||||
|
||||
class LocalBufferContext:
|
||||
"""
|
||||
This class creates a context that helps to generate code involving Inductor IR with
|
||||
function local buffers. These buffers are constructed during the codegen process and
|
||||
are used to store intermediate results such as local accumulators. We do not want to
|
||||
add them to `V.graph` since they are not global and we do not want to add them as
|
||||
function arguments either. So we patch the codegen processes under this scope to support
|
||||
these buffers without exposure to the outside world.
|
||||
"""
|
||||
|
||||
def __init__(self, kernel_args: KernelArgs) -> None:
|
||||
self.kernel_args = kernel_args
|
||||
self.exit_stack = contextlib.ExitStack()
|
||||
# map local buffer name to local buffer
|
||||
self.local_buffers: dict[str, ir.Buffer] = {}
|
||||
# map global buffer name to global buffer
|
||||
self.global_buffers: dict[str, ir.Buffer] = {}
|
||||
# map global buffer name to local buffer
|
||||
self.global_to_local: dict[str, ir.Buffer] = {}
|
||||
# record the global buffers that are removed by this LocalBufferContext
|
||||
self.removed_buffers: OrderedSet[str] = OrderedSet()
|
||||
|
||||
def __enter__(self):
|
||||
self.exit_stack.__enter__()
|
||||
original_get_dtype = V.graph.get_dtype
|
||||
|
||||
def get_dtype(name):
|
||||
if name in self.local_buffers:
|
||||
return self.local_buffers[name].get_dtype()
|
||||
return original_get_dtype(name)
|
||||
|
||||
self.exit_stack.enter_context(patch.object(V.graph, "get_dtype", get_dtype))
|
||||
|
||||
original_input = self.kernel_args.input
|
||||
|
||||
def input(name):
|
||||
if name in self.local_buffers:
|
||||
return name
|
||||
return original_input(name)
|
||||
|
||||
self.exit_stack.enter_context(patch.object(self.kernel_args, "input", input))
|
||||
|
||||
original_output = self.kernel_args.output
|
||||
|
||||
def output(name):
|
||||
if name in self.local_buffers:
|
||||
return name
|
||||
return original_output(name)
|
||||
|
||||
self.exit_stack.enter_context(patch.object(self.kernel_args, "output", output))
|
||||
|
||||
# Set current LocalBufferContext into V
|
||||
self.exit_stack.enter_context(V.set_local_buffer_context(self))
|
||||
|
||||
return self
|
||||
|
||||
def __exit__(self, exc_type, exc_val, exc_tb):
|
||||
self.local_buffers.clear()
|
||||
self.exit_stack.__exit__(exc_type, exc_val, exc_tb)
|
||||
|
||||
def add_local_buffer(
|
||||
self, local_buffer: ir.Buffer, global_buffers: list[ir.Buffer] | None = None
|
||||
):
|
||||
assert local_buffer.get_name() not in self.local_buffers
|
||||
self.local_buffers[local_buffer.get_name()] = local_buffer
|
||||
if global_buffers:
|
||||
for global_buffer in global_buffers:
|
||||
global_buffer_name = global_buffer.get_name()
|
||||
assert (
|
||||
global_buffer_name not in self.global_buffers
|
||||
and global_buffer_name not in self.global_to_local
|
||||
)
|
||||
self.global_buffers[global_buffer_name] = global_buffer
|
||||
self.global_to_local[global_buffer_name] = local_buffer
|
||||
if global_buffer_name not in V.graph.removed_buffers:
|
||||
# Record the global buffers that are removed by this LocalBufferContext
|
||||
# since which may need to restore. Refer to issue:
|
||||
# https://github.com/pytorch/pytorch/issues/144186
|
||||
self.removed_buffers.add(global_buffer_name)
|
||||
V.graph.removed_buffers.add(global_buffer_name)
|
||||
|
||||
def localize_function(
|
||||
self,
|
||||
fn: Callable[..., Any],
|
||||
rewrite_index: Callable[
|
||||
["LocalizeBufferHandler", sympy.Expr, str], sympy.Expr
|
||||
] = rewrite_index_for_function,
|
||||
):
|
||||
def inner(*args, **kwargs):
|
||||
with V.set_ops_handler(
|
||||
LocalizeBufferHandler(
|
||||
V.get_ops_handler(),
|
||||
global_to_local=self.global_to_local,
|
||||
rewrite_index=rewrite_index,
|
||||
)
|
||||
):
|
||||
return fn(*args, **kwargs)
|
||||
|
||||
return inner
|
||||
|
||||
def localize_nodes(
|
||||
self,
|
||||
nodes: list[ir.IRNode],
|
||||
rewrite_index: Callable[
|
||||
["LocalizeBufferHandler", sympy.Expr, str], sympy.Expr
|
||||
] = rewrite_index_for_nodes,
|
||||
) -> list[ir.IRNode]:
|
||||
"""
|
||||
Given `local_buf` and `global_buf` registered in current `LocalBufferContext`
|
||||
though the method of `add_local_buffer`, localizes the `global_buf` to `local_buf`
|
||||
for the given `nodes` and returns a new list of IR nodes that work on `local_buf`
|
||||
instead of `global_buf`, i.e., all the loads and stores are redirected to
|
||||
`local_buf`. This helps the fused loops to work on smaller-sized local buffers
|
||||
for better data locality.
|
||||
|
||||
The data access of `local_buf` is assumed to be contiguous with the
|
||||
same order as the `global_buf`.
|
||||
"""
|
||||
assert len(nodes) > 0
|
||||
|
||||
def wrap_inner_fn_for_node(node: ir.IRNode):
|
||||
loops = node.data if isinstance(node, ir.ComputedBuffer) else node
|
||||
assert isinstance(loops, ir.Loops)
|
||||
new_inner_fn = self.localize_function(
|
||||
loops.inner_fn,
|
||||
rewrite_index,
|
||||
)
|
||||
|
||||
new_loops = dataclasses.replace(loops, inner_fn=new_inner_fn)
|
||||
if isinstance(node, ir.ComputedBuffer):
|
||||
new_node = ir.ComputedBuffer(
|
||||
name=node.get_name(), layout=node.get_layout(), data=new_loops
|
||||
)
|
||||
else:
|
||||
new_node = new_loops # type: ignore[assignment]
|
||||
|
||||
return new_node
|
||||
|
||||
return [wrap_inner_fn_for_node(node) for node in nodes]
|
||||
|
||||
|
||||
def unify_mask_base_type(
|
||||
buffer: IndentedBuffer,
|
||||
vars: tuple[CSEVariable, ...],
|
||||
dtype=torch.float,
|
||||
):
|
||||
"""
|
||||
Given list of cse variables,
|
||||
Cast each to new mask base dtype and return casted cse variable.
|
||||
"""
|
||||
new_vars = (
|
||||
V.kernel.cse.generate(
|
||||
buffer,
|
||||
f"{V.kernel._get_mask_cast(var, dtype)}",
|
||||
)
|
||||
for var in vars
|
||||
)
|
||||
return new_vars
|
||||
|
||||
|
||||
def may_unify_binary_op_mask_type(a, b):
|
||||
"""
|
||||
Given two cse variables, when dtype is bool, unify them to the same mask dtype and return casted cse variable.
|
||||
"""
|
||||
if a.dtype == torch.bool:
|
||||
assert b.dtype == torch.bool
|
||||
mask_dtype = torch.int32
|
||||
return unify_mask_base_type(V.kernel.compute, (a, b), mask_dtype)
|
||||
return a, b
|
||||
|
||||
|
||||
def codegen_rand(offset, code, rand_function, dst_dtype=torch.float32):
|
||||
assert is_integer_dtype(offset.dtype)
|
||||
code.writeline("[&]()")
|
||||
with code.indent():
|
||||
code.writeline(
|
||||
f"{DTYPE_TO_CPP[offset.dtype]} offset[{V.kernel.tiling_factor}];"
|
||||
)
|
||||
code.writeline(f"{DTYPE_TO_CPP[dst_dtype]} result[{V.kernel.tiling_factor}];")
|
||||
code.writeline(f"{offset}.store(offset);")
|
||||
code.writeline(
|
||||
f"for( {DTYPE_TO_CPP[offset.dtype]} offset_idx = 0; offset_idx < {V.kernel.tiling_factor}; offset_idx++ )"
|
||||
)
|
||||
with code.indent():
|
||||
code.writeline(rand_function)
|
||||
num_vectors = V.kernel._get_num_vectors(dtype=dst_dtype)
|
||||
if num_vectors == 1:
|
||||
code.writeline(
|
||||
f"return at::vec::Vectorized<{DTYPE_TO_CPP[dst_dtype]}>::loadu(result);"
|
||||
)
|
||||
else:
|
||||
code.writeline(
|
||||
f"return at::vec::VectorizedN<{DTYPE_TO_CPP[dst_dtype]}, {num_vectors}>::loadu(result);"
|
||||
)
|
||||
code.writeline("()")
|
||||
return code
|
||||
|
||||
|
||||
def get_gemm_template_output_and_compute_dtype(input_dtype):
|
||||
if input_dtype in [torch.uint8, torch.int8]:
|
||||
return (torch.int32, torch.int32)
|
||||
else:
|
||||
return (torch.float32, torch.float32)
|
||||
|
||||
|
||||
def create_epilogue_with_attr(input_buffer, attr, **kwargs):
|
||||
input_loader = input_buffer.make_loader()
|
||||
dtype = input_buffer.get_dtype()
|
||||
if attr == "relu":
|
||||
|
||||
def inner_fn(index):
|
||||
input = input_loader(index)
|
||||
zero = ops.constant(0, dtype)
|
||||
return ops.maximum(input, zero)
|
||||
|
||||
elif attr == "gelu":
|
||||
assert "algorithm" in kwargs
|
||||
if kwargs["algorithm"] == "none":
|
||||
|
||||
def inner_fn(index):
|
||||
input = input_loader(index)
|
||||
if dtype != torch.float:
|
||||
input = ops.to_dtype(input, torch.float)
|
||||
half = ops.constant(0.5, torch.float)
|
||||
one = ops.constant(1.0, torch.float)
|
||||
const = ops.constant(0.7071067811865476, torch.float)
|
||||
result = input * half * (ops.erf(input * const) + one)
|
||||
if dtype != torch.float:
|
||||
result = ops.to_dtype(result, dtype)
|
||||
return result
|
||||
|
||||
else:
|
||||
assert kwargs["algorithm"] == "tanh"
|
||||
|
||||
def inner_fn(index):
|
||||
input = input_loader(index)
|
||||
if dtype != torch.float:
|
||||
input = ops.to_dtype(input, torch.float)
|
||||
half = ops.constant(0.5, torch.float)
|
||||
one = ops.constant(1.0, torch.float)
|
||||
const1 = ops.constant(0.7978845608028654, torch.float)
|
||||
const2 = ops.constant(0.044715, torch.float)
|
||||
result = (
|
||||
half
|
||||
* input
|
||||
* (
|
||||
one
|
||||
+ ops.tanh(const1 * (input + const2 * input * input * input))
|
||||
)
|
||||
)
|
||||
if dtype != torch.float:
|
||||
result = ops.to_dtype(result, dtype)
|
||||
return result
|
||||
|
||||
elif attr == "swish":
|
||||
|
||||
def inner_fn(index):
|
||||
input = input_loader(index)
|
||||
result = input * ops.sigmoid(input)
|
||||
return result
|
||||
|
||||
elif attr == "sigmoid":
|
||||
|
||||
def inner_fn(index):
|
||||
return ops.sigmoid(input_loader(index))
|
||||
|
||||
elif attr == "tanh":
|
||||
|
||||
def inner_fn(index):
|
||||
return ops.tanh(input_loader(index))
|
||||
|
||||
elif attr == "hardswish" or attr == "hardsigmoid":
|
||||
|
||||
def hardsigmoid_float(input):
|
||||
zero = ops.constant(0, torch.float)
|
||||
six = ops.constant(6, torch.float)
|
||||
three = ops.constant(3, torch.float)
|
||||
one_over_six = ops.constant(0.16666666666666666, torch.float)
|
||||
max = ops.maximum(input + three, zero)
|
||||
min = ops.minimum(max, six)
|
||||
return min * one_over_six
|
||||
|
||||
def inner_fn(index):
|
||||
input = input_loader(index)
|
||||
if dtype != torch.float:
|
||||
input = ops.to_dtype(input, torch.float)
|
||||
result = hardsigmoid_float(input)
|
||||
if attr == "hardswish":
|
||||
result = input * result
|
||||
if dtype != torch.float:
|
||||
result = ops.to_dtype(result, dtype)
|
||||
return result
|
||||
|
||||
elif attr == "leaky_relu":
|
||||
assert "scalars" in kwargs
|
||||
assert len(kwargs["scalars"]) == 1
|
||||
negative_slope = kwargs["scalars"][0]
|
||||
|
||||
def inner_fn(index):
|
||||
input = input_loader(index)
|
||||
if dtype != torch.float:
|
||||
input = ops.to_dtype(input, torch.float)
|
||||
zero = ops.constant(0, torch.float)
|
||||
result = ops.where(
|
||||
input > zero, input, input * ops.constant(negative_slope, torch.float)
|
||||
)
|
||||
if dtype != torch.float:
|
||||
result = ops.to_dtype(result, dtype)
|
||||
return result
|
||||
|
||||
elif attr == "hardtanh":
|
||||
assert "scalars" in kwargs
|
||||
assert len(kwargs["scalars"]) == 2
|
||||
min_value = kwargs["scalars"][0]
|
||||
max_value = kwargs["scalars"][1]
|
||||
|
||||
def inner_fn(index):
|
||||
input = input_loader(index)
|
||||
if dtype != torch.float:
|
||||
input = ops.to_dtype(input, torch.float)
|
||||
result = ops.minimum(
|
||||
ops.maximum(input, ops.constant(min_value, torch.float)),
|
||||
ops.constant(max_value, torch.float),
|
||||
)
|
||||
if dtype != torch.float:
|
||||
result = ops.to_dtype(result, dtype)
|
||||
return result
|
||||
|
||||
elif attr in ["add", "sub", "mul"]:
|
||||
assert "other" in kwargs
|
||||
other = kwargs["other"]
|
||||
num_input_dims = len(input_buffer.get_size())
|
||||
num_other_dims = len(other.get_size())
|
||||
dims_diff = num_input_dims - num_other_dims
|
||||
other_loader = other.make_loader()
|
||||
|
||||
def inner_fn(index):
|
||||
op = getattr(ops, attr)
|
||||
if dims_diff != 0:
|
||||
return op(input_loader(index), other_loader(index[dims_diff:]))
|
||||
else:
|
||||
return op(input_loader(index), other_loader(index))
|
||||
|
||||
elif attr == "bias_add":
|
||||
assert "other" in kwargs
|
||||
assert "beta" in kwargs
|
||||
assert "dtype" in kwargs
|
||||
beta = kwargs["beta"]
|
||||
other = kwargs["other"]
|
||||
dtype = kwargs["dtype"]
|
||||
bias_loader = other.make_loader()
|
||||
|
||||
def inner_fn(index):
|
||||
bias = bias_loader(index)
|
||||
input = input_loader(index)
|
||||
if beta != 1:
|
||||
result = ops.constant(beta, torch.float) * bias + input
|
||||
else:
|
||||
result = bias + input
|
||||
return result
|
||||
|
||||
else:
|
||||
raise ValueError(f"Unsupported epilogue attribute: {attr}")
|
||||
return ir.Pointwise(
|
||||
device=input_buffer.get_device(),
|
||||
dtype=dtype,
|
||||
inner_fn=inner_fn,
|
||||
ranges=input_buffer.get_size(),
|
||||
)
|
||||
|
||||
|
||||
def _get_loop_body(fn_list):
|
||||
if all(isinstance(fn, LoopBody) for fn in fn_list):
|
||||
loop_bodies = fn_list
|
||||
else:
|
||||
if hasattr(fn_list[0], "original_fn"):
|
||||
# For the case of local buffer, we wrap the fn with localize_function
|
||||
assert all(hasattr(fn, "original_fn") for fn in fn_list)
|
||||
assert all(
|
||||
isinstance(fn.original_fn.args[0]._body, LoopBody) for fn in fn_list
|
||||
)
|
||||
loop_bodies = [fn.original_fn.args[0]._body for fn in fn_list]
|
||||
else:
|
||||
assert all(isinstance(fn, functools.partial) for fn in fn_list)
|
||||
assert all(isinstance(fn.args[0]._body, LoopBody) for fn in fn_list)
|
||||
loop_bodies = [fn.args[0]._body for fn in fn_list]
|
||||
assert loop_bodies is not None
|
||||
return loop_bodies
|
||||
|
||||
|
||||
def _get_dtype_from_loopbodies(loop_bodies):
|
||||
dtypes = OrderedSet[torch.dtype]()
|
||||
for loop_body in loop_bodies:
|
||||
graphs = [loop_body.root_block.graph] + [
|
||||
body.graph for body in list(loop_body.subblocks.values())
|
||||
]
|
||||
for graph in graphs:
|
||||
for node in graph.nodes:
|
||||
if node.op != "call_method":
|
||||
continue
|
||||
dtypes.add(node.meta[OptimizationContext.key].dtype)
|
||||
return dtypes
|
||||
|
||||
|
||||
def template_fusion_with_epilogues_supported(
|
||||
template: BaseSchedulerNode, epilogues: list[BaseSchedulerNode]
|
||||
) -> tuple[bool, bool]:
|
||||
def _get_indexes_of_template_buf_read(
|
||||
epilogue_node: ir.Operation, template_buf_names: list[str]
|
||||
) -> list[sympy.Expr]:
|
||||
return [
|
||||
read.index
|
||||
for read in epilogue_node.get_reads()
|
||||
if read.name in template_buf_names
|
||||
]
|
||||
|
||||
def _check_supported_and_same_indexes(
|
||||
index_of_template_buf_read: Sequence[sympy.Expr],
|
||||
epilogue_writes: OrderedSet[Dep],
|
||||
) -> tuple[bool, bool]:
|
||||
num_indexes = len(OrderedSet(index_of_template_buf_read))
|
||||
|
||||
if num_indexes > 1:
|
||||
same_index = False
|
||||
supported = False # Different read indexes not supported
|
||||
elif num_indexes == 0:
|
||||
same_index = True
|
||||
supported = True # No reads, automatically supported
|
||||
elif num_indexes == 1:
|
||||
iotbr = index_of_template_buf_read[0]
|
||||
same_index = all(write.index == iotbr for write in epilogue_writes)
|
||||
# TODO: Add support of fusion when the read of template buffer and the write of epilogue output
|
||||
# in the epilogue node don't have the same index and change supported to True
|
||||
supported = same_index
|
||||
else:
|
||||
raise AssertionError("Should not reach here")
|
||||
|
||||
return supported, same_index
|
||||
|
||||
def _template_fusion_supported(
|
||||
template_outputs: Sequence[SchedulerBuffer], epilogue_nodes: list[ir.Operation]
|
||||
) -> tuple[bool, bool]:
|
||||
template_buf_names = [x.get_name() for x in template_outputs]
|
||||
indexes_of_template_buf_reads = [
|
||||
_get_indexes_of_template_buf_read(epilogue_node, template_buf_names)
|
||||
for epilogue_node in epilogue_nodes
|
||||
]
|
||||
epilogue_nodes_writes = [
|
||||
epilogue_node.get_read_writes().writes for epilogue_node in epilogue_nodes
|
||||
]
|
||||
|
||||
results = [
|
||||
_check_supported_and_same_indexes(reads, writes)
|
||||
for reads, writes in zip(
|
||||
indexes_of_template_buf_reads, epilogue_nodes_writes
|
||||
)
|
||||
]
|
||||
supported, same_indexes = zip(*results)
|
||||
return all(supported), all(same_indexes)
|
||||
|
||||
assert template.is_template()
|
||||
template_outputs = template.get_outputs()
|
||||
|
||||
epilogue_nodes = [
|
||||
n.node
|
||||
for epilogue in epilogues
|
||||
for n in epilogue.get_nodes()
|
||||
if n.node is not None
|
||||
]
|
||||
return _template_fusion_supported(template_outputs, epilogue_nodes)
|
||||
File diff suppressed because it is too large
Load Diff
+899
@@ -0,0 +1,899 @@
|
||||
# mypy: allow-untyped-defs
|
||||
from collections.abc import Callable, Sequence
|
||||
from typing import Any
|
||||
|
||||
import sympy
|
||||
|
||||
import torch
|
||||
import torch._inductor.async_compile # noqa: F401 required to warm up AsyncCompile pools
|
||||
import torch._ops
|
||||
|
||||
from .. import config, ir
|
||||
from ..utils import sympy_product
|
||||
from ..virtualized import V
|
||||
from .cpp_utils import DTYPE_TO_CPP
|
||||
from .cpp_wrapper_cpu import CppWrapperCpu
|
||||
from .wrapper import (
|
||||
BufferLike,
|
||||
EnterSubgraphLine,
|
||||
ExitSubgraphLine,
|
||||
MemoryPlanningLine,
|
||||
MemoryPlanningState,
|
||||
PythonWrapperCodegen,
|
||||
)
|
||||
|
||||
|
||||
BufferName = str
|
||||
|
||||
# Default thread stack sizes vary by platform:
|
||||
# - Linux: 8 MB
|
||||
# - macOS: 512 KB
|
||||
# - Windows: 1 MB
|
||||
# Just pick something comfortably smaller than the smallest for now.
|
||||
MAX_STACK_ALLOCATION_SIZE = 1024 * 100
|
||||
|
||||
|
||||
class CppWrapperCpuArrayRef(CppWrapperCpu):
|
||||
"""
|
||||
Generates cpp wrapper for running on CPU and calls cpp kernels
|
||||
|
||||
This class is forked from CppWrapperCpu, with a difference that tensors may be
|
||||
represented as ArrayRef, see torch/csrc/inductor/aoti_runtime/arrayref_tensor.h
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
assert self.device == "cpu", "ArrayRefTensor only supported on CPU!"
|
||||
self.allow_stack_allocation = config.aot_inductor.allow_stack_allocation
|
||||
self.stack_allocated_buffers: dict[BufferName, BufferLike] = {}
|
||||
|
||||
@staticmethod
|
||||
def create(
|
||||
is_subgraph: bool,
|
||||
subgraph_name: str | None,
|
||||
parent_wrapper: PythonWrapperCodegen | None,
|
||||
partition_signatures: ir.GraphPartitionSignature | None = None,
|
||||
):
|
||||
# TODO - support subgraph codegen by lifting functions. Check the
|
||||
# comment at CppWrapperCpu `codegen_subgraph` function.
|
||||
return CppWrapperCpuArrayRef()
|
||||
|
||||
@staticmethod
|
||||
def get_input_cpp_type(input):
|
||||
assert config.aot_inductor.use_minimal_arrayref_interface
|
||||
|
||||
if isinstance(input, sympy.Expr):
|
||||
from ..graph import may_get_constant_buffer_dtype
|
||||
|
||||
dtype = may_get_constant_buffer_dtype(input)
|
||||
assert dtype is not None, f"Failed to get the dtype of sympy.Expr: {input}"
|
||||
return DTYPE_TO_CPP[dtype]
|
||||
return f"ArrayRefTensor<{DTYPE_TO_CPP[input.get_dtype()]}>"
|
||||
|
||||
@staticmethod
|
||||
def get_device_include_path(device: str) -> str:
|
||||
assert device == "cpu", "ArrayRef only supported on CPU!"
|
||||
if V.graph.aot_mode:
|
||||
return "#include <torch/csrc/inductor/aoti_include/array_ref.h>"
|
||||
return "#include <torch/csrc/inductor/cpp_wrapper/array_ref.h>"
|
||||
|
||||
def codegen_input_numel_asserts(self):
|
||||
for name, buf in V.graph.graph_inputs.items():
|
||||
if isinstance(buf, sympy.Expr):
|
||||
continue
|
||||
|
||||
# comparing strides for 0 size tensor is tricky. Ignore them for now.
|
||||
if sympy_product(buf.get_size()) == 0:
|
||||
continue
|
||||
numel = buf.get_numel()
|
||||
self.prefix.writeline(f"assert_numel({name}, {numel});")
|
||||
|
||||
def generate_extern_kernel_alloc(self, *args, **kwargs):
|
||||
# Disable stack allocation for extern kernels.
|
||||
self.allow_stack_allocation = False
|
||||
super().generate_extern_kernel_alloc(*args, **kwargs)
|
||||
|
||||
def generate_extern_kernel_out(self, *args, **kwargs):
|
||||
# Disable stack allocation for extern kernels.
|
||||
self.allow_stack_allocation = False
|
||||
super().generate_extern_kernel_out(*args, **kwargs)
|
||||
|
||||
def generate_fallback_kernel(self, node: ir.FallbackKernel) -> None:
|
||||
# Disable stack allocation for extern kernels.
|
||||
self.allow_stack_allocation = False
|
||||
super().generate_fallback_kernel(node)
|
||||
|
||||
def _generate_kernel_call_helper(
|
||||
self,
|
||||
kernel_name: str,
|
||||
call_args,
|
||||
*,
|
||||
device=None,
|
||||
triton=True,
|
||||
arg_types=None,
|
||||
raw_keys=None,
|
||||
raw_args=None,
|
||||
triton_meta=None,
|
||||
inductor_meta=None,
|
||||
graph_name="",
|
||||
original_fxnode_name=None,
|
||||
current_stream_idx=None,
|
||||
):
|
||||
"""
|
||||
Generates kernel call code.
|
||||
|
||||
triton: Defines whether the GPU backend uses Triton for codegen.
|
||||
Otherwise it uses the CUDA language for codegen.
|
||||
Only valid when cuda == True.
|
||||
"""
|
||||
assert not triton, (
|
||||
"CppWrapperCpuArrayRef.generate_kernel_call does not support GPU"
|
||||
)
|
||||
assert arg_types is not None and len(call_args) == len(arg_types), (
|
||||
"Mismatch call_args and arg_types in generate_kernel_call"
|
||||
)
|
||||
new_args = []
|
||||
for idx, arg in enumerate(call_args):
|
||||
if "*" in arg_types[idx]:
|
||||
var_name = f"var_{next(self.arg_var_id)}"
|
||||
self.writeline(f"auto* {var_name} = get_data_ptr_wrapper({arg});")
|
||||
new_args.append(f"({arg_types[idx]})({var_name})")
|
||||
else:
|
||||
# arg is a scalar
|
||||
new_args.append(arg)
|
||||
# debug printer related logic for cpp kernel type.
|
||||
debug_printer_manager = V.graph.wrapper_code.debug_printer
|
||||
debug_printer_manager.set_printer_args(
|
||||
call_args,
|
||||
kernel_name,
|
||||
None,
|
||||
None,
|
||||
"cpp",
|
||||
)
|
||||
with debug_printer_manager:
|
||||
self.writeline(self.wrap_kernel_call(kernel_name, new_args))
|
||||
|
||||
def write_wrapper_decl(self):
|
||||
inputs_len = len(V.graph.graph_inputs.keys())
|
||||
if V.graph.aot_mode:
|
||||
if (
|
||||
config.aot_inductor.use_minimal_arrayref_interface
|
||||
and not V.graph.is_const_graph
|
||||
):
|
||||
input_cpp_types = ", ".join(
|
||||
f"{CppWrapperCpuArrayRef.get_input_cpp_type(x)}"
|
||||
for x in V.graph.graph_inputs.values()
|
||||
)
|
||||
output_arrayref_types = ", ".join(
|
||||
f"ArrayRefTensor<{DTYPE_TO_CPP[x.get_dtype()]}>"
|
||||
for x in V.graph.graph_outputs
|
||||
)
|
||||
|
||||
self.prefix.splice(
|
||||
f"""
|
||||
using AOTInductorModelInputs = std::tuple<{input_cpp_types}>;
|
||||
using AOTInductorModelOutputs = std::tuple<{output_arrayref_types}>;
|
||||
"""
|
||||
)
|
||||
|
||||
if V.graph.const_module:
|
||||
self.header.splice(V.graph.const_module.wrapper_code.header)
|
||||
|
||||
assert V.graph.const_wrapper_code is not None
|
||||
self.prefix.splice(V.graph.const_wrapper_code)
|
||||
|
||||
assert V.graph.const_kernel_code is not None
|
||||
self.kernel_declarations.splice(V.graph.const_kernel_code)
|
||||
|
||||
if V.graph.is_const_graph:
|
||||
self.prefix.splice(
|
||||
"""
|
||||
void AOTInductorModel::_const_run_impl(
|
||||
std::vector<AtenTensorHandle>& output_handles,
|
||||
DeviceStreamType stream,
|
||||
AOTIProxyExecutorHandle proxy_executor
|
||||
) {
|
||||
"""
|
||||
)
|
||||
else:
|
||||
if not config.aot_inductor.use_runtime_constant_folding:
|
||||
# If we do not split the constant graph, we'll just create
|
||||
# an empty implementation when wrapping the main module.
|
||||
self.prefix.splice(
|
||||
"""
|
||||
void AOTInductorModel::_const_run_impl(
|
||||
std::vector<AtenTensorHandle>& output_handles,
|
||||
DeviceStreamType stream,
|
||||
AOTIProxyExecutorHandle proxy_executor
|
||||
) {}
|
||||
|
||||
"""
|
||||
)
|
||||
|
||||
run_impl_proto = """
|
||||
void AOTInductorModel::run_impl(
|
||||
AtenTensorHandle*
|
||||
input_handles, // array of input AtenTensorHandle; handles
|
||||
// are stolen; the array itself is borrowed
|
||||
AtenTensorHandle*
|
||||
output_handles, // array for writing output AtenTensorHandle; handles
|
||||
// will be stolen by the caller; the array itself is
|
||||
// borrowed
|
||||
DeviceStreamType stream,
|
||||
AOTIProxyExecutorHandle proxy_executor
|
||||
) {
|
||||
"""
|
||||
|
||||
self.generate_input_output_runtime_checks()
|
||||
run_impl_proto += """
|
||||
__check_inputs_outputs(input_handles, output_handles);
|
||||
"""
|
||||
|
||||
if config.aot_inductor.use_minimal_arrayref_interface:
|
||||
self.prefix.splice(
|
||||
"""
|
||||
template <>
|
||||
AOTInductorModelOutputs AOTInductorModel::run_impl_minimal_arrayref_interface<
|
||||
AOTInductorModelInputs, AOTInductorModelOutputs>(
|
||||
const AOTInductorModelInputs& inputs,
|
||||
DeviceStreamType stream,
|
||||
AOTIProxyExecutorHandle proxy_executor
|
||||
) {
|
||||
"""
|
||||
)
|
||||
self.suffix.splice(run_impl_proto)
|
||||
self.suffix.splice(
|
||||
"""
|
||||
AOTInductorModelInputs inputs;
|
||||
convert_handles_to_inputs(input_handles, inputs);
|
||||
auto outputs = run_impl_minimal_arrayref_interface<AOTInductorModelInputs, AOTInductorModelOutputs>(
|
||||
inputs, stream, proxy_executor);
|
||||
// NOTE: outputs is full of ArrayRef to thread_local storage. If in the future we need this
|
||||
// interface to perform well for a DSO using the minimal arrayref interface, all we need
|
||||
// to do is provide ThreadLocalCachedTensor for each one!
|
||||
convert_outputs_to_handles(outputs, output_handles);
|
||||
}
|
||||
"""
|
||||
)
|
||||
|
||||
self.suffix.splice(
|
||||
"""
|
||||
extern "C" AOTIRuntimeError AOTInductorModelRunMinimalArrayrefInterface(
|
||||
AOTInductorModelHandle model_handle,
|
||||
const AOTInductorModelInputs& inputs,
|
||||
AOTInductorModelOutputs& outputs) {
|
||||
auto model = reinterpret_cast<torch::aot_inductor::AOTInductorModel*>(model_handle);
|
||||
CONVERT_EXCEPTION_TO_ERROR_CODE({
|
||||
outputs = model->run_impl_minimal_arrayref_interface<AOTInductorModelInputs, AOTInductorModelOutputs>(
|
||||
inputs,
|
||||
(torch::aot_inductor::DeviceStreamType)nullptr,
|
||||
nullptr);
|
||||
})
|
||||
}
|
||||
"""
|
||||
)
|
||||
else:
|
||||
self.prefix.splice(run_impl_proto)
|
||||
else:
|
||||
# cpp entry function for JIT with cpp wrapper
|
||||
self.prefix.splice(
|
||||
"""
|
||||
void inductor_entry_impl(
|
||||
AtenTensorHandle*
|
||||
input_handles, // array of input AtenTensorHandle; handles
|
||||
// are stolen; the array itself is borrowed
|
||||
AtenTensorHandle*
|
||||
output_handles // array for writing output AtenTensorHandle; handles
|
||||
// will be stolen by the caller; the array itself is
|
||||
// borrowed)
|
||||
) {
|
||||
"""
|
||||
)
|
||||
with self.prefix.indent():
|
||||
# assign inputs and outputs in both cases so the later codegen can be simplified
|
||||
if not config.aot_inductor.use_minimal_arrayref_interface:
|
||||
if not V.graph.is_const_graph:
|
||||
if V.graph.aot_mode:
|
||||
num_args = len(V.graph.graph_inputs)
|
||||
else:
|
||||
# Weights are promoted in the JIT mode
|
||||
num_args = len(V.graph.graph_inputs) + len(V.graph.constants)
|
||||
# release GIL to support multiple instances inference (in different threads of the same process)
|
||||
self.prefix.splice("py::gil_scoped_release_simple release;")
|
||||
|
||||
self.prefix.splice(
|
||||
f"""
|
||||
auto inputs = steal_from_raw_handles_to_raii_handles(input_handles, {num_args});
|
||||
"""
|
||||
)
|
||||
|
||||
if inputs_len != 0:
|
||||
for idx, input_key in enumerate(V.graph.graph_inputs.keys()):
|
||||
if config.aot_inductor.use_minimal_arrayref_interface:
|
||||
self.prefix.writeline(
|
||||
f"auto {input_key} = std::get<{idx}>(inputs);"
|
||||
)
|
||||
continue
|
||||
# unwrap input tensor back to scalar
|
||||
if isinstance(V.graph.graph_inputs[input_key], sympy.Expr):
|
||||
from ..graph import may_get_constant_buffer_dtype
|
||||
|
||||
dtype = may_get_constant_buffer_dtype(
|
||||
V.graph.graph_inputs[input_key] # type: ignore[arg-type]
|
||||
)
|
||||
assert dtype is not None, (
|
||||
"Fails to get the dtype of the sympy.Expr"
|
||||
)
|
||||
self.codegen_tensor_item(
|
||||
dtype, f"inputs[{idx}]", input_key, self.prefix
|
||||
)
|
||||
else:
|
||||
self.prefix.writeline(
|
||||
f"auto {input_key} = std::move(inputs[{idx}]);"
|
||||
)
|
||||
|
||||
assert all(
|
||||
isinstance(v, torch.Tensor) for v in list(V.graph.constants.values())
|
||||
), "Expect all constants to be Tensor"
|
||||
for idx, constants_key in enumerate(V.graph.constants.keys()):
|
||||
if V.graph.aot_mode:
|
||||
# Weights are stored in constants_ and owned by RAIIAtenTensorHandle there.
|
||||
# Don't call std::move here because it will cause constants_ to lose the ownership.
|
||||
self.prefix.writeline(
|
||||
f"""auto {constants_key} = constants_->at({idx});"""
|
||||
)
|
||||
else:
|
||||
# Append constants as inputs to the graph
|
||||
constants_idx = inputs_len + idx
|
||||
self.prefix.writeline(
|
||||
f"auto {constants_key} = std::move(inputs[{constants_idx}]);"
|
||||
)
|
||||
|
||||
self.codegen_inputs()
|
||||
|
||||
if V.graph.aot_mode:
|
||||
if not V.graph.is_const_graph:
|
||||
if config.aot_inductor.use_minimal_arrayref_interface:
|
||||
# TODO: input shape checking for regular tensor interface as well?
|
||||
self.codegen_input_numel_asserts()
|
||||
else:
|
||||
self.prefix.writeline("inputs.clear();")
|
||||
self.prefix.writeline(
|
||||
"[[maybe_unused]] auto& kernels = static_cast<AOTInductorModelKernels&>(*this->kernels_.get());"
|
||||
)
|
||||
|
||||
def generate_return(self, output_refs: list[str]):
|
||||
cst_names = V.graph.constants.keys()
|
||||
arr_iface = (
|
||||
not V.graph.is_const_graph
|
||||
and config.aot_inductor.use_minimal_arrayref_interface
|
||||
) # For brevity.
|
||||
|
||||
def use_thread_local_cached_output_tensor(idx, output):
|
||||
cached_output_name = f"cached_output_{next(self.cached_output_id)}"
|
||||
cache_type = "Array" if arr_iface else "Tensor"
|
||||
self.wrapper_call.writeline(
|
||||
f"thread_local ThreadLocalCachedOutput{cache_type}<std::decay_t<decltype({output})>> "
|
||||
f"{cached_output_name}({output});"
|
||||
)
|
||||
if arr_iface:
|
||||
self.wrapper_call.writeline(
|
||||
f"{cached_output_name}.copy_data_from({output});"
|
||||
)
|
||||
output_entry = f"std::get<{idx}>(output_arrayref_tensors)"
|
||||
element_type = f"std::decay_t<decltype({output_entry}.data()[0])>"
|
||||
self.wrapper_call.writeline(
|
||||
f"{output_entry} = {cached_output_name}.arrayref_tensor<{element_type}>();"
|
||||
)
|
||||
else:
|
||||
self.wrapper_call.writeline(
|
||||
f"{cached_output_name}.copy_data_from({output});"
|
||||
)
|
||||
self.wrapper_call.writeline(
|
||||
f"AOTI_TORCH_ERROR_CODE_CHECK(aoti_torch_new_uninitialized_tensor(&output_handles[{idx}]));"
|
||||
)
|
||||
self.wrapper_call.writeline(
|
||||
f"AOTI_TORCH_ERROR_CODE_CHECK(aoti_torch_assign_tensors({cached_output_name}.tensor(), "
|
||||
f"output_handles[{idx}]));"
|
||||
)
|
||||
|
||||
if arr_iface:
|
||||
self.wrapper_call.writeline(
|
||||
"AOTInductorModelOutputs output_arrayref_tensors;"
|
||||
)
|
||||
|
||||
output2idx: dict[str, int] = {}
|
||||
for idx, output in enumerate(output_refs):
|
||||
if output == "nullptr":
|
||||
continue
|
||||
|
||||
is_constant_buffer = output in cst_names
|
||||
output_buffer = V.graph.graph_outputs[idx]
|
||||
if isinstance(output_buffer, ir.BaseView):
|
||||
output_storage = output_buffer.unwrap_view()
|
||||
assert isinstance(output_storage, (ir.BaseView, ir.MutableBox))
|
||||
if isinstance(output_storage.data, ir.ConstantBuffer):
|
||||
is_constant_buffer = True
|
||||
|
||||
if isinstance(output_buffer, ir.ShapeAsConstantBuffer):
|
||||
# Need to wrap scalar into tensor as the main function returns a vector of tensors
|
||||
output_tensor = self.codegen_scalar_to_tensor(output)
|
||||
self.wrapper_call.writeline(
|
||||
f"output_handles[{idx}] = {output_tensor}.release();"
|
||||
)
|
||||
continue
|
||||
|
||||
output_is_tensor_handle_expr = (
|
||||
f"std::is_same_v<std::decay_t<decltype({output})>,"
|
||||
"RAIIAtenTensorHandle> || "
|
||||
f"std::is_same_v<std::decay_t<decltype({output})>,"
|
||||
"AtenTensorHandle> || "
|
||||
f"std::is_same_v<std::decay_t<decltype({output})>,"
|
||||
"ConstantHandle>"
|
||||
)
|
||||
self.wrapper_call.writeline(
|
||||
f"if constexpr ({output_is_tensor_handle_expr}) {{"
|
||||
)
|
||||
with self.wrapper_call.indent():
|
||||
if arr_iface:
|
||||
cached_output_name = f"cached_output_{next(self.cached_output_id)}"
|
||||
self.wrapper_call.writeline(
|
||||
f"thread_local RAIIAtenTensorHandle {cached_output_name};"
|
||||
)
|
||||
if is_constant_buffer:
|
||||
# NOTE(return_constant): In some rare cases where we return
|
||||
# a constant, we have to return a copy of this constant,
|
||||
# because (1) constants are not owned by the Model instance
|
||||
# (2) constants remain the same cross inference runs,
|
||||
# assuming they are not updated at runtime Basically, we
|
||||
# cannot release or transfer the ownership of any original
|
||||
# constant to the user.
|
||||
self.wrapper_call.writeline(
|
||||
f"AtenTensorHandle {cached_output_name}_tmp;"
|
||||
)
|
||||
self.wrapper_call.writeline(
|
||||
f"aoti_torch_clone({output}, &{cached_output_name}_tmp);"
|
||||
)
|
||||
self.wrapper_call.writeline(
|
||||
f"{cached_output_name} = {cached_output_name}_tmp;"
|
||||
)
|
||||
else:
|
||||
self.wrapper_call.writeline(
|
||||
f"{cached_output_name} = {output}.release();"
|
||||
)
|
||||
self.wrapper_call.writeline(
|
||||
f"convert_handle_to_arrayref_tensor({cached_output_name}, "
|
||||
f"std::get<{idx}>(output_arrayref_tensors));"
|
||||
)
|
||||
else:
|
||||
if is_constant_buffer:
|
||||
# See NOTE(return_constant) above.
|
||||
self.wrapper_call.writeline(
|
||||
f"aoti_torch_clone({output}, &output_handles[{idx}]);"
|
||||
)
|
||||
else:
|
||||
if output in output2idx:
|
||||
src_idx = output2idx[output]
|
||||
self.wrapper_call.writeline(
|
||||
f"output_handles[{idx}] = output_handles[{src_idx}];"
|
||||
)
|
||||
else:
|
||||
self.wrapper_call.writeline(
|
||||
f"output_handles[{idx}] = {output}.release();"
|
||||
)
|
||||
self.wrapper_call.writeline("} else {")
|
||||
with self.wrapper_call.indent():
|
||||
use_thread_local_cached_output_tensor(idx, output)
|
||||
self.wrapper_call.writeline("}")
|
||||
|
||||
if output not in output2idx:
|
||||
output2idx[output] = idx
|
||||
if arr_iface:
|
||||
self.wrapper_call.writeline("return output_arrayref_tensors;")
|
||||
|
||||
def memory_plan(self):
|
||||
from .memory_planning import MemoryPlanner
|
||||
|
||||
self.lines = MemoryPlanner(self).plan(self.lines)
|
||||
# TODO: integrate memory planning & stack allocation?
|
||||
self.allow_stack_allocation = False
|
||||
|
||||
def memory_plan_reuse(self):
|
||||
out_names = V.graph.get_output_names()
|
||||
|
||||
while (
|
||||
self.lines
|
||||
and isinstance(self.lines[-1], MemoryPlanningLine)
|
||||
# TODO: this seems legit, NullLine has no node
|
||||
and self.lines[-1].node.name not in out_names # type: ignore[attr-defined]
|
||||
):
|
||||
# these lines will be pointless
|
||||
self.lines.pop()
|
||||
|
||||
# codegen allocations in two passes
|
||||
planning_states = [MemoryPlanningState()]
|
||||
past_planning_states = []
|
||||
for i in range(len(self.lines)):
|
||||
line = self.lines[i]
|
||||
if isinstance(line, MemoryPlanningLine):
|
||||
self.lines[i] = line.plan(planning_states[-1])
|
||||
elif isinstance(line, EnterSubgraphLine):
|
||||
planning_states.append(MemoryPlanningState())
|
||||
elif isinstance(line, ExitSubgraphLine):
|
||||
past_planning_states.append(planning_states.pop())
|
||||
past_planning_states.append(planning_states.pop())
|
||||
assert len(planning_states) == 0
|
||||
|
||||
# conservatively use the sum of all allocated buffer sizes
|
||||
# in potentially nested scopes as the total allocated size
|
||||
total_allocated_buffer_size = sum(
|
||||
s.total_allocated_buffer_size for s in past_planning_states
|
||||
)
|
||||
|
||||
self.allow_stack_allocation = (
|
||||
self.allow_stack_allocation is not False
|
||||
and config.aot_inductor.allow_stack_allocation
|
||||
and total_allocated_buffer_size <= MAX_STACK_ALLOCATION_SIZE
|
||||
)
|
||||
|
||||
def can_stack_allocate_buffer(self, buffer):
|
||||
return (
|
||||
self.allow_stack_allocation
|
||||
and buffer.get_device().type == "cpu"
|
||||
and self.can_prove_buffer_has_static_shape(buffer)
|
||||
and ir.is_contiguous_strides_for_shape(
|
||||
buffer.get_stride(), buffer.get_size()
|
||||
)
|
||||
)
|
||||
|
||||
def make_buffer_free(self, buffer):
|
||||
return (
|
||||
""
|
||||
if isinstance(buffer.get_output_spec(), ir.MultiOutputLayout)
|
||||
or (V.graph.aot_mode and buffer.get_name() in self.stack_allocated_buffers)
|
||||
or (
|
||||
config.aot_inductor.use_minimal_arrayref_interface
|
||||
and V.graph.aot_mode
|
||||
and buffer.get_name() in V.graph.graph_inputs
|
||||
)
|
||||
else f"{buffer.get_name()}.reset();"
|
||||
)
|
||||
|
||||
def make_buffer_allocation(self, buffer):
|
||||
return self.make_allocation(
|
||||
buffer.get_name(),
|
||||
buffer.get_device(),
|
||||
buffer.get_dtype(),
|
||||
buffer.get_size(),
|
||||
buffer.get_stride(),
|
||||
buffer if self.can_stack_allocate_buffer(buffer) else None,
|
||||
buffer.get_is_pinned(),
|
||||
)
|
||||
|
||||
def make_allocation(
|
||||
self,
|
||||
name,
|
||||
device,
|
||||
dtype,
|
||||
shape,
|
||||
stride,
|
||||
buffer_if_can_stack_allocate=None,
|
||||
is_pinned=False,
|
||||
):
|
||||
orig_stride = stride
|
||||
device_str = self.codegen_device(device)
|
||||
dtype_code = self.codegen_dtype(dtype)
|
||||
size = self.codegen_shape_tuple(shape)
|
||||
stride = self.codegen_shape_tuple(orig_stride)
|
||||
size_array_var = self.codegen_int_array_var(
|
||||
size,
|
||||
self.wrapper_call.writeline,
|
||||
known_statically=self.is_statically_known_list_of_ints(shape),
|
||||
graph=self.get_codegened_graph(),
|
||||
)
|
||||
stride_array_var = self.codegen_int_array_var(
|
||||
stride,
|
||||
self.wrapper_call.writeline,
|
||||
known_statically=self.is_statically_known_list_of_ints(orig_stride),
|
||||
graph=self.get_codegened_graph(),
|
||||
)
|
||||
device_type, device_id = device_str.split(",")
|
||||
device_idx = "this->device_idx_" if V.graph.aot_mode else device_id
|
||||
if buffer_if_can_stack_allocate is not None:
|
||||
self.stack_allocated_buffers[name] = buffer_if_can_stack_allocate
|
||||
cpp_type = DTYPE_TO_CPP[dtype]
|
||||
numel = buffer_if_can_stack_allocate.get_numel()
|
||||
# Note: we don't zero storage because empty_strided doesn't zero either.
|
||||
self.wrapper_call.writeline(f"{cpp_type} {name}_storage[{numel}];")
|
||||
args = [
|
||||
f"{name}_storage",
|
||||
size_array_var,
|
||||
stride_array_var,
|
||||
device_type,
|
||||
device_idx,
|
||||
]
|
||||
return f"ArrayRefTensor<{cpp_type}> {name}({', '.join(args)});"
|
||||
|
||||
args = [
|
||||
str(len(shape)),
|
||||
size_array_var,
|
||||
stride_array_var,
|
||||
dtype_code,
|
||||
device_type,
|
||||
device_idx,
|
||||
f"&{name}_handle",
|
||||
]
|
||||
|
||||
self.wrapper_call.writeline(f"AtenTensorHandle {name}_handle;")
|
||||
pinned_str = "_pinned" if is_pinned else ""
|
||||
self.wrapper_call.writeline(
|
||||
f"AOTI_TORCH_ERROR_CODE_CHECK(aoti_torch_empty_strided{pinned_str}({', '.join(args)}));"
|
||||
)
|
||||
|
||||
return f"RAIIAtenTensorHandle {name}({name}_handle);"
|
||||
|
||||
def make_buffer_reuse(self, old: BufferLike, new: BufferLike, delete_old: bool):
|
||||
assert old.get_dtype() == new.get_dtype()
|
||||
old_name = old.get_name()
|
||||
new_name = new.get_name()
|
||||
del_line = ";"
|
||||
if old_name not in V.graph.get_output_names() and delete_old:
|
||||
del_line = f"; {self.make_buffer_free(old)}"
|
||||
|
||||
if old.get_size() == new.get_size() and old.get_stride() == new.get_stride():
|
||||
if old_name in self.stack_allocated_buffers:
|
||||
self.stack_allocated_buffers[new_name] = new
|
||||
return self.codegen_exact_buffer_reuse(old_name, new_name, del_line)
|
||||
|
||||
reinterpret_view = self.codegen_reinterpret_view(
|
||||
old, new.get_size(), new.get_stride(), 0, self.wrapper_call.writeline
|
||||
)
|
||||
if reinterpret_view in self.stack_allocated_buffers:
|
||||
self.stack_allocated_buffers[new_name] = new
|
||||
# The only way to get into this case is via an exact buffer reuse, since all
|
||||
# other options result in a new tensor handle.
|
||||
return self.codegen_exact_buffer_reuse(old_name, new_name, del_line)
|
||||
return f"{self.declare}{new_name} = {reinterpret_view}{del_line} // reuse"
|
||||
|
||||
def _assert_safe_to_use_borrow_arrayref_tensor_as_tensor(self):
|
||||
# Borrowing arguments to shim functions is only safe because we know
|
||||
# that the arguments can't be stack-allocated. Otherwise, to be sure
|
||||
# we can't return a dangling pointer, we need to either 1) be
|
||||
# certain that the shim function cannot return an alias of a
|
||||
# borrowed argument, or 2) be certain that the returned Tensor from
|
||||
# the shim function cannot escape.
|
||||
assert self.is_safe_to_use_borrow_arrayref_tensor_as_tensor(), (
|
||||
"borrowing arguments to shim functions is unsafe with "
|
||||
"stack allocation on! (see comment above this assertion)"
|
||||
)
|
||||
|
||||
def is_safe_to_use_borrow_arrayref_tensor_as_tensor(self):
|
||||
return not self.allow_stack_allocation and not self.stack_allocated_buffers
|
||||
|
||||
def generate_c_shim_extern_kernel_call(
|
||||
self, kernel: str, args: list[str], device: str, **_
|
||||
) -> None:
|
||||
# In the abi_compatible mode, we call fallback aten ops through a C shim layer
|
||||
# Setting self.allow_stack_allocation to False because the exchange between
|
||||
# ArrayRefTensor and at::Tensor is still fragile.
|
||||
self.allow_stack_allocation = False
|
||||
|
||||
wrapped_args = []
|
||||
for arg in args:
|
||||
# We only really *need* borrow_arrayref_tensor_as_tensor for
|
||||
# ArrayRefTensors. The code flowing into here uses `0` for nullptr, which
|
||||
# borrow_arrayref_tensor_as_tensor would blindly coerce to int, so just
|
||||
# avoid wrapping integers. Name matching is to find tensor is hacky, but
|
||||
# fixing all the ArrayRefTensor issues is not a priority for now.
|
||||
if isinstance(arg, str) and arg.startswith(
|
||||
("buf", "arg", "wrap_with_raii_handle_if_needed")
|
||||
):
|
||||
self._assert_safe_to_use_borrow_arrayref_tensor_as_tensor()
|
||||
arg = f"borrow_arrayref_tensor_as_tensor({arg})"
|
||||
wrapped_args.append(arg)
|
||||
|
||||
super().generate_c_shim_extern_kernel_call(
|
||||
kernel, wrapped_args, device, debug_args=args
|
||||
)
|
||||
|
||||
def generate_scatter_fallback(self, node: ir.ScatterFallback):
|
||||
# No stack allocation when there is a fallback op
|
||||
self.allow_stack_allocation = False
|
||||
super().generate_scatter_fallback(node)
|
||||
|
||||
def _generate_scatter_fallback(
|
||||
self,
|
||||
output,
|
||||
inputs,
|
||||
cpp_kernel_name,
|
||||
python_kernel_name,
|
||||
src_is_tensor,
|
||||
reduce,
|
||||
kwargs,
|
||||
device,
|
||||
):
|
||||
reduce = self._get_scatter_reduce_enum(reduce)
|
||||
|
||||
# call the ABI shim function instead of the ATen one
|
||||
self.add_device_include(device)
|
||||
cpp_kernel_name = self.get_c_shim_func_name(cpp_kernel_name, device)
|
||||
|
||||
# TODO: consider remove "_out" and add missing inplace variants to fallback_ops.py
|
||||
cpp_kernel_name = cpp_kernel_name.replace("__", "_") + "_out"
|
||||
self._assert_safe_to_use_borrow_arrayref_tensor_as_tensor()
|
||||
inputs_wrapped = [
|
||||
(f"borrow_arrayref_tensor_as_tensor({x})" if isinstance(x, str) else str(x))
|
||||
for x in inputs
|
||||
]
|
||||
line = f"{cpp_kernel_name}(borrow_arrayref_tensor_as_tensor({output}), {','.join(inputs_wrapped)}"
|
||||
|
||||
if python_kernel_name.startswith("aten.scatter_reduce"):
|
||||
line += f", {','.join(kwargs)}"
|
||||
else:
|
||||
if src_is_tensor:
|
||||
if reduce:
|
||||
line += f", {V.graph.wrapper_code.val_to_arg_str(reduce)}"
|
||||
else:
|
||||
assert reduce is None, (
|
||||
"Expect reduce to be None for aten.scatter_ with scalar src"
|
||||
)
|
||||
line += ");"
|
||||
self.writeline(line)
|
||||
|
||||
def generate_index_put_fallback(self, node: ir.IndexPutFallback) -> None:
|
||||
# No stack allocation when there is a fallback op
|
||||
self.allow_stack_allocation = False
|
||||
super().generate_index_put_fallback(node)
|
||||
|
||||
def _generate_index_put_fallback(self, kernel, x, indices, values, accumulate):
|
||||
self._assert_safe_to_use_borrow_arrayref_tensor_as_tensor()
|
||||
# TODO: update aoti_torch_index_put_out in ir.py to use autogen out version
|
||||
# See the comment in codegen_reinterpret_view about why having something like
|
||||
# RAIIAtenTensorHandle(tmp_tensor_handle_2) in a tmp array can cause the corresponding
|
||||
# tensor prematurely deallocated, thus the temporary array trick here.
|
||||
indices_str = self._generate_temporary_array_pointer(
|
||||
"AtenTensorHandle",
|
||||
[f"borrow_arrayref_tensor_as_tensor({i})" for i in indices],
|
||||
)
|
||||
args = [
|
||||
f"borrow_arrayref_tensor_as_tensor({x})",
|
||||
indices_str,
|
||||
str(len(indices)),
|
||||
f"borrow_arrayref_tensor_as_tensor({values})",
|
||||
accumulate,
|
||||
]
|
||||
args.insert(
|
||||
0, f"borrow_arrayref_tensor_as_tensor({x})"
|
||||
) # set x as the output tensor, this fallback mutates x.
|
||||
self.writeline(self.wrap_kernel_call(kernel, args))
|
||||
|
||||
def generate_fallback_kernel_with_runtime_lookup(
|
||||
self,
|
||||
buf_name: str,
|
||||
python_kernel_name: str,
|
||||
get_args: Callable[[], Sequence[str]],
|
||||
op_overload: torch._ops.OpOverload | torch._ops.HigherOrderOperator,
|
||||
raw_args: Sequence[Any],
|
||||
outputs: Sequence[ir.Buffer],
|
||||
) -> None:
|
||||
# No stack allocation when there is a fallback op
|
||||
self.allow_stack_allocation = False
|
||||
super().generate_fallback_kernel_with_runtime_lookup(
|
||||
buf_name, python_kernel_name, get_args, op_overload, raw_args, outputs
|
||||
)
|
||||
|
||||
def codegen_device_copy(self, src, dst, non_blocking: bool | str):
|
||||
# aoti_torch_tensor_copy_ takes AtenTensorHandle as input,
|
||||
# while stack-allocation results in ArrayRefTensor
|
||||
# so disable stack allocation here
|
||||
self.allow_stack_allocation = False
|
||||
self.writeline(
|
||||
f"AOTI_TORCH_ERROR_CODE_CHECK(aoti_torch_copy_(expensive_copy_to_tensor_if_needed({dst}), {src}, {non_blocking}));"
|
||||
)
|
||||
|
||||
def codegen_reinterpret_view(
|
||||
self,
|
||||
data,
|
||||
size,
|
||||
stride,
|
||||
offset,
|
||||
writeline: Callable[..., None],
|
||||
dtype=None,
|
||||
) -> str:
|
||||
"""Returns a newly-created, temporary RAII tensor handle containing the
|
||||
reinterpreted tensor data. Callers of this function are responsible for saving
|
||||
the handle if persistent access is needed."""
|
||||
dim = str(len(size))
|
||||
|
||||
def create_reinterpret_call() -> str:
|
||||
args = [
|
||||
f"{data.get_name()}",
|
||||
dim,
|
||||
self.codegen_int_array_var(
|
||||
self.codegen_shape_tuple(size),
|
||||
writeline,
|
||||
known_statically=self.is_statically_known_list_of_ints(size),
|
||||
graph=self.get_codegened_graph(),
|
||||
),
|
||||
self.codegen_int_array_var(
|
||||
self.codegen_shape_tuple(stride),
|
||||
writeline,
|
||||
known_statically=self.is_statically_known_list_of_ints(stride),
|
||||
graph=self.get_codegened_graph(),
|
||||
),
|
||||
offset,
|
||||
]
|
||||
return f"wrap_with_raii_handle_if_needed(reinterpret_tensor_wrapper({', '.join(args)}))"
|
||||
|
||||
def create_new_tensor_handle() -> tuple[str, list[str]]:
|
||||
# Calling reset() on ArrayRefTensor does nothing, since the array is
|
||||
# const-allocated on the stack. Thus, it's safe to return a reference to
|
||||
# the original array.
|
||||
if (name := data.get_name()) in self.stack_allocated_buffers:
|
||||
return name, []
|
||||
|
||||
tmp_AtenTensorHandle = f"tmp_{name}_{next(self.tmp_tensor_id)}"
|
||||
tmp_call_strs = [
|
||||
f"AtenTensorHandle {tmp_AtenTensorHandle};",
|
||||
f"AOTI_TORCH_ERROR_CODE_CHECK(aoti_torch_new_tensor_handle({data.get_name()}, &{tmp_AtenTensorHandle}));",
|
||||
]
|
||||
return f"RAIIAtenTensorHandle({tmp_AtenTensorHandle})", tmp_call_strs
|
||||
|
||||
if (
|
||||
size == data.layout.size
|
||||
and stride == data.layout.stride
|
||||
and offset == data.layout.offset
|
||||
and (dtype is None or dtype == data.dtype)
|
||||
):
|
||||
final_tensor_str, call_strs = create_new_tensor_handle()
|
||||
for line in call_strs:
|
||||
writeline(line)
|
||||
return final_tensor_str
|
||||
|
||||
return super().codegen_reinterpret_view(
|
||||
data, size, stride, offset, writeline, dtype
|
||||
)
|
||||
|
||||
def val_to_arg_str(self, val, type_=None) -> str:
|
||||
if (
|
||||
val is not None
|
||||
and isinstance(type_, torch.OptionalType)
|
||||
and isinstance(type_.getElementType(), torch.TensorType)
|
||||
):
|
||||
# Handle optional tensors as a special case, as in the parent class.
|
||||
base_handle = self.val_to_arg_str(val, torch.TensorType)
|
||||
if config.aot_inductor.use_minimal_arrayref_interface:
|
||||
if self.is_safe_to_use_borrow_arrayref_tensor_as_tensor():
|
||||
base_handle = f"borrow_arrayref_tensor_as_tensor({base_handle})"
|
||||
else:
|
||||
base_handle = f"copy_arrayref_tensor_to_tensor({base_handle})"
|
||||
return f"&temporary_reference({base_handle}.get())"
|
||||
|
||||
return super().val_to_arg_str(val, type_)
|
||||
|
||||
def codegen_tensor_item(
|
||||
self, dtype: torch.dtype, tensor: str, scalar: str, indented_buffer=None
|
||||
):
|
||||
dtype_str = str(dtype).split(".")[-1]
|
||||
writer = indented_buffer or self
|
||||
|
||||
if dtype == torch.float16 or dtype == torch.bfloat16:
|
||||
scalar_tmp = f"{scalar}_tmp"
|
||||
writer.writeline(f"{DTYPE_TO_CPP[dtype]} {scalar_tmp};")
|
||||
|
||||
# We know that item_ doesn't alias the input, so borrowing should be safe.
|
||||
tensor = f"borrow_arrayref_tensor_as_tensor({tensor})"
|
||||
|
||||
writer.writeline(
|
||||
f"AOTI_TORCH_ERROR_CODE_CHECK(aoti_torch_item_{dtype_str}({tensor}, &{scalar_tmp}));"
|
||||
)
|
||||
writer.writeline(f"float {scalar} = float({scalar_tmp});")
|
||||
else:
|
||||
writer.writeline(f"{DTYPE_TO_CPP[dtype]} {scalar};")
|
||||
|
||||
# We know that item_ doesn't alias the input, so borrowing should be safe.
|
||||
tensor = f"borrow_arrayref_tensor_as_tensor({tensor})"
|
||||
|
||||
writer.writeline(
|
||||
f"AOTI_TORCH_ERROR_CODE_CHECK(aoti_torch_item_{dtype_str}({tensor}, &{scalar}));"
|
||||
)
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,304 @@
|
||||
from typing import Any
|
||||
|
||||
import sympy
|
||||
|
||||
import torch
|
||||
from torch.utils._ordered_set import OrderedSet
|
||||
|
||||
from ..ir import GraphPartitionSignature
|
||||
from ..virtualized import V
|
||||
from .cpp_wrapper_cpu import CppWrapperCpu
|
||||
from .cpp_wrapper_gpu import CppWrapperGpu
|
||||
from .wrapper import KernelCallLine, PythonWrapperCodegen
|
||||
|
||||
|
||||
class CppWrapperMps(CppWrapperGpu):
|
||||
"""
|
||||
Generates cpp wrapper for running on MPS and calls metal kernels
|
||||
"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
self._used_kernel_names: OrderedSet[str] = OrderedSet()
|
||||
self._lambda_counter: int = 0
|
||||
|
||||
@staticmethod
|
||||
def create(
|
||||
is_subgraph: bool,
|
||||
subgraph_name: str | None,
|
||||
parent_wrapper: PythonWrapperCodegen | None,
|
||||
partition_signatures: GraphPartitionSignature | None = None,
|
||||
) -> "CppWrapperMps":
|
||||
return CppWrapperMps()
|
||||
|
||||
def _generate_kernel_call_helper(
|
||||
self,
|
||||
kernel_name: str,
|
||||
call_args: list[str],
|
||||
*,
|
||||
device: torch.device | None = None,
|
||||
triton: bool = True,
|
||||
arg_types: tuple[Any, ...] | None = None,
|
||||
raw_keys: tuple[Any, ...] | None = None,
|
||||
raw_args: tuple[Any, ...] | None = None,
|
||||
triton_meta: dict[str, Any] | None = None,
|
||||
inductor_meta: dict[str, Any] | None = None,
|
||||
graph_name: str = "",
|
||||
original_fxnode_name: str | None = None,
|
||||
current_stream_idx: int | None = None,
|
||||
) -> None:
|
||||
"""
|
||||
Generates MPS kernel call code. It should look something like:
|
||||
```
|
||||
auto mps_lib_0_lambda = [&](AOTIMetalKernelFunctionHandle handle) {
|
||||
aoti_torch_mps_start_encoding(handle);
|
||||
aoti_torch_mps_set_arg_tensor(handle, 0, buf0);
|
||||
aoti_torch_mps_set_arg_tensor(handle, 1, arg0_1);
|
||||
aoti_torch_mps_set_arg_tensor(handle, 2, arg1_1);
|
||||
aoti_torch_mps_dispatch_single(handle, static_cast<uint64_t>(10LL));
|
||||
};
|
||||
|
||||
std::function<void(AOTIMetalKernelFunctionHandle)> mps_lib_0_func_wrapper = mps_lib_0_lambda;
|
||||
aoti_torch_mps_run_command_block(get_mps_lib_0_handle(), aoti_torch_mps_shared_callback, &mps_lib_0_func_wrapper);
|
||||
```
|
||||
"""
|
||||
device = device or V.graph.get_current_device_or_throw()
|
||||
if device.type == "cpu":
|
||||
# Even in CppWrapperGpu, we may see cpp kernels
|
||||
return CppWrapperCpu._generate_kernel_call_helper(
|
||||
self,
|
||||
kernel_name,
|
||||
call_args,
|
||||
device=device,
|
||||
triton=triton,
|
||||
arg_types=arg_types,
|
||||
raw_keys=raw_keys,
|
||||
raw_args=raw_args,
|
||||
triton_meta=triton_meta,
|
||||
inductor_meta=inductor_meta,
|
||||
)
|
||||
|
||||
assert device.type == "mps"
|
||||
|
||||
assert arg_types is not None
|
||||
|
||||
new_args = []
|
||||
for idx, (arg, arg_type) in enumerate(zip(call_args[:-2], arg_types[:-2])):
|
||||
if isinstance(arg_type, torch.dtype):
|
||||
new_args.append(f"aoti_torch_mps_set_arg_tensor(handle, {idx}, {arg});")
|
||||
elif arg_type in (int, sympy.core.symbol.Symbol):
|
||||
new_args.append(f"aoti_torch_mps_set_arg_int(handle, {idx}, {arg});")
|
||||
else:
|
||||
raise NotImplementedError(
|
||||
f"Unsupported arg type {arg_type} for arg {arg} for kernel {kernel_name}"
|
||||
)
|
||||
|
||||
threads, group_size = call_args[-2], call_args[-1]
|
||||
if threads is None:
|
||||
raise NotImplementedError("No threads or group_size provided")
|
||||
|
||||
# Check if threads is a single value or an array-like structure
|
||||
threads_str = str(threads)
|
||||
is_single_value = (
|
||||
threads_str.startswith("{")
|
||||
and threads_str.endswith("}")
|
||||
and threads_str.count(",") == 0
|
||||
) or not threads_str.startswith(("{", "["))
|
||||
|
||||
if is_single_value:
|
||||
# Extract single value from braces if present
|
||||
if threads_str.startswith("{") and threads_str.endswith("}"):
|
||||
single_value = threads_str[1:-1].strip() # Remove braces
|
||||
else:
|
||||
single_value = threads_str
|
||||
|
||||
if group_size is None:
|
||||
new_args.append(
|
||||
f"aoti_torch_mps_dispatch_single(handle, {single_value});"
|
||||
)
|
||||
else:
|
||||
# Extract group size value if it's also in braces
|
||||
group_size_str = str(group_size)
|
||||
if group_size_str.startswith("{") and group_size_str.endswith("}"):
|
||||
group_size_value = group_size_str[1:-1].strip()
|
||||
else:
|
||||
group_size_value = group_size_str
|
||||
new_args.append(
|
||||
f"aoti_torch_mps_dispatch_single_with_group_size(handle, {single_value}, {group_size_value});"
|
||||
)
|
||||
else:
|
||||
# Handle array case - need to convert initializer list to array
|
||||
# Use kernel name to make variable names unique
|
||||
threads_var = f"{kernel_name}_threads_array"
|
||||
group_size_var = f"{kernel_name}_group_size_array"
|
||||
|
||||
# Extract array size from the initializer list string
|
||||
def get_array_size(array_str: str) -> int:
|
||||
# Remove braces and whitespace
|
||||
content = array_str.strip()
|
||||
if content.startswith("{") and content.endswith("}"):
|
||||
content = content[1:-1].strip()
|
||||
|
||||
if not content: # Empty array
|
||||
return 0
|
||||
|
||||
# Count elements by counting commas, accounting for nested structures
|
||||
depth = 0
|
||||
comma_count = 0
|
||||
for char in content:
|
||||
if char in "({[<":
|
||||
depth += 1
|
||||
elif char in ")}]>":
|
||||
depth -= 1
|
||||
elif char == "," and depth == 0:
|
||||
comma_count += 1
|
||||
|
||||
return comma_count + 1 # Number of elements = commas + 1
|
||||
|
||||
threads_size = get_array_size(threads_str)
|
||||
|
||||
if group_size is None:
|
||||
new_args.append("{")
|
||||
new_args.append(f" uint64_t {threads_var}[] = {threads};")
|
||||
new_args.append(
|
||||
f" aoti_torch_mps_dispatch_array(handle, {threads_var}, {threads_size});"
|
||||
)
|
||||
new_args.append("}")
|
||||
else:
|
||||
group_size_str = str(group_size)
|
||||
group_size_size = get_array_size(group_size_str)
|
||||
new_args.append("{")
|
||||
new_args.append(f" uint64_t {threads_var}[] = {threads};")
|
||||
new_args.append(f" uint64_t {group_size_var}[] = {group_size};")
|
||||
dispatch_args = f"handle, {threads_var}, {threads_size}, {group_size_var}, {group_size_size}"
|
||||
new_args.append(
|
||||
f" aoti_torch_mps_dispatch_array_with_group_size({dispatch_args});"
|
||||
)
|
||||
new_args.append("}")
|
||||
|
||||
# debug printer related logic for cpp kernel type.
|
||||
debug_printer_manager = V.graph.wrapper_code.debug_printer
|
||||
debug_printer_manager.set_printer_args(
|
||||
call_args[:-2],
|
||||
kernel_name,
|
||||
None,
|
||||
None,
|
||||
"cpp",
|
||||
)
|
||||
with debug_printer_manager:
|
||||
self.write_mps_kernel_call(kernel_name, new_args)
|
||||
|
||||
def write_mps_kernel_call(self, name: str, call_args: list[str]) -> None:
|
||||
# Generate unique variable names to avoid duplicate declarations
|
||||
# when the same MPS lib is used multiple times
|
||||
unique_suffix = self._lambda_counter
|
||||
self._lambda_counter += 1
|
||||
|
||||
lambda_name = f"{name}_lambda_{unique_suffix}"
|
||||
wrapper_name = f"{name}_func_wrapper_{unique_suffix}"
|
||||
|
||||
# Generate the function call code (in current location)
|
||||
# Create lambda that captures by reference and pass its pointer through void*
|
||||
self.writeline(
|
||||
f"auto {lambda_name} = [&](AOTIMetalKernelFunctionHandle handle) {{"
|
||||
)
|
||||
self.writeline(" aoti_torch_mps_start_encoding(handle);")
|
||||
|
||||
# Output call args directly since we're capturing by reference
|
||||
for call_arg in call_args:
|
||||
self.writeline(f" {call_arg}")
|
||||
self.writeline("};")
|
||||
self.writeline("")
|
||||
|
||||
# Pass lambda pointer through void*
|
||||
self.writeline(
|
||||
f"std::function<void(AOTIMetalKernelFunctionHandle)> {wrapper_name} = {lambda_name};"
|
||||
)
|
||||
self.writeline(
|
||||
f"aoti_torch_mps_run_command_block(get_{name}_handle(), aoti_torch_mps_shared_callback, &{wrapper_name});"
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def get_device_include_path(device: str) -> str:
|
||||
assert V.graph.aot_mode
|
||||
return (
|
||||
"#include <torch/csrc/inductor/aoti_include/mps.h>\n"
|
||||
"#include <torch/csrc/inductor/aoti_torch/c/shim_mps.h>"
|
||||
)
|
||||
|
||||
def codegen_additional_funcs(self) -> None:
|
||||
"""
|
||||
Generate thread-safe lazy singleton pattern for MPS shader libraries with RAII cleanup.
|
||||
|
||||
The generated code will look like:
|
||||
```
|
||||
AOTIMetalKernelFunctionHandle get_mps_lib_0_handle() {
|
||||
static auto kernel_handle = []() {
|
||||
AOTIMetalShaderLibraryHandle lib_handle = nullptr;
|
||||
AOTIMetalKernelFunctionHandle kern_handle = nullptr;
|
||||
|
||||
aoti_torch_mps_create_shader_library(mps_lib_0_source, &lib_handle);
|
||||
aoti_torch_mps_get_kernel_function(lib_handle, "generated_kernel", &kern_handle);
|
||||
|
||||
// RAII wrapper with custom deleter
|
||||
auto lib_deleter = [](AOTIMetalShaderLibraryHandle h) {
|
||||
if (h) aoti_torch_mps_delete_shader_library(h);
|
||||
};
|
||||
|
||||
using LibDeleter = decltype(lib_deleter);
|
||||
using LibPtr = std::unique_ptr<AOTIMetalShaderLibraryOpaque, LibDeleter>;
|
||||
|
||||
// Return pair of kernel handle and library smart pointer for cleanup
|
||||
return std::make_pair(kern_handle, LibPtr(lib_handle, lib_deleter));
|
||||
}();
|
||||
return kernel_handle.first;
|
||||
}
|
||||
```
|
||||
"""
|
||||
|
||||
# Add shimified handles and functions
|
||||
shader_libraries: OrderedSet[str] = OrderedSet()
|
||||
for line in self.lines:
|
||||
if not isinstance(line, KernelCallLine):
|
||||
continue
|
||||
if line.device.type != "mps":
|
||||
continue
|
||||
|
||||
# Extract library name from kernel name (e.g., "mps_lib_0" from kernel calls)
|
||||
if line.kernel_name not in self._used_kernel_names:
|
||||
self._used_kernel_names.add(line.kernel_name)
|
||||
shader_libraries.add(line.kernel_name)
|
||||
|
||||
# NOTE: For shimified version, we expect the shader source constant to be generated
|
||||
# by the existing MPS shader generation process, but instead of instantiating the
|
||||
# DynamicMetalShaderLibrary directly, we'll use our shim functions.
|
||||
# The existing codegen should produce something like:
|
||||
# const char* mps_lib_0_source = R"MTL(...shader_source...)MTL";
|
||||
# instead of:
|
||||
# at::native::mps::DynamicMetalShaderLibrary mps_lib_0(R"MTL(...shader_source...)MTL");
|
||||
|
||||
# Generate thread-safe lazy singleton with RAII for each library
|
||||
for lib_name in shader_libraries:
|
||||
self.prefix.splice(f"""
|
||||
AOTIMetalKernelFunctionHandle get_{lib_name}_handle() {{
|
||||
static auto kernel_handle = []() {{
|
||||
AOTIMetalShaderLibraryHandle lib_handle = nullptr;
|
||||
AOTIMetalKernelFunctionHandle kern_handle = nullptr;
|
||||
|
||||
aoti_torch_mps_create_shader_library({lib_name}_source, &lib_handle);
|
||||
aoti_torch_mps_get_kernel_function(lib_handle, "generated_kernel", &kern_handle);
|
||||
|
||||
// RAII wrapper with custom deleter
|
||||
auto lib_deleter = [](AOTIMetalShaderLibraryHandle h) {{
|
||||
if (h) aoti_torch_mps_delete_shader_library(h);
|
||||
}};
|
||||
|
||||
using LibDeleter = decltype(lib_deleter);
|
||||
using LibPtr = std::unique_ptr<AOTIMetalShaderLibraryOpaque, LibDeleter>;
|
||||
|
||||
// Return pair of kernel handle and library smart pointer for cleanup
|
||||
return std::make_pair(kern_handle, LibPtr(lib_handle, lib_deleter));
|
||||
}}();
|
||||
return kernel_handle.first;
|
||||
}}
|
||||
""")
|
||||
+30
@@ -0,0 +1,30 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from textwrap import dedent
|
||||
|
||||
from .common import DeviceOpOverrides, register_device_op_overrides
|
||||
|
||||
|
||||
class CpuDeviceOpOverrides(DeviceOpOverrides):
|
||||
def import_get_raw_stream_as(self, name: str) -> str:
|
||||
return dedent(
|
||||
"""
|
||||
def get_raw_stream(_):
|
||||
return 0
|
||||
"""
|
||||
)
|
||||
|
||||
def cpp_kernel_type(self) -> str:
|
||||
return "void*"
|
||||
|
||||
def set_device(self, device_idx: int) -> str:
|
||||
return "pass"
|
||||
|
||||
def synchronize(self) -> str:
|
||||
return "pass"
|
||||
|
||||
def device_guard(self, device_idx: int) -> str:
|
||||
return "pass"
|
||||
|
||||
|
||||
register_device_op_overrides("cpu", CpuDeviceOpOverrides())
|
||||
+271
@@ -0,0 +1,271 @@
|
||||
# mypy: allow-untyped-defs
|
||||
import logging
|
||||
import os
|
||||
import shutil
|
||||
from pathlib import Path
|
||||
|
||||
import torch
|
||||
from torch._inductor import config
|
||||
from torch._inductor.codegen.cuda import cuda_env
|
||||
from torch._inductor.cpp_builder import _set_gpu_runtime_env, _transform_cuda_paths
|
||||
from torch._inductor.utils import is_linux
|
||||
|
||||
|
||||
if config.is_fbcode():
|
||||
from triton.fb.build import build_paths
|
||||
|
||||
|
||||
log = logging.getLogger(__name__)
|
||||
autotuning_log = torch._logging.getArtifactLogger(__name__, "autotuning")
|
||||
|
||||
|
||||
def use_re_build() -> bool:
|
||||
"""
|
||||
Use for CUTLASS compilation only right now.
|
||||
"""
|
||||
if config.is_fbcode() and not cuda_env.nvcc_exist(_cuda_compiler()):
|
||||
from triton.fb.re_build_helper import should_build_locally
|
||||
|
||||
return not should_build_locally()
|
||||
return False
|
||||
|
||||
|
||||
def _cutlass_path() -> str:
|
||||
if config.is_fbcode():
|
||||
from libfb.py import parutil
|
||||
|
||||
return parutil.get_dir_path("cutlass-4-headers")
|
||||
else:
|
||||
return config.cutlass.cutlass_dir
|
||||
|
||||
|
||||
def _cutlass_paths() -> list[str]:
|
||||
return [
|
||||
"include",
|
||||
"tools/library/include",
|
||||
"tools/library/src",
|
||||
"tools/util/include",
|
||||
]
|
||||
|
||||
|
||||
def _clone_cutlass_paths(build_root: str) -> list[str]:
|
||||
paths = _cutlass_paths()
|
||||
cutlass_root = _cutlass_path()
|
||||
for path in _cutlass_paths():
|
||||
old_path = os.path.join(cutlass_root, path)
|
||||
new_path = os.path.join(build_root, path)
|
||||
shutil.copytree(old_path, new_path, dirs_exist_ok=True)
|
||||
return paths
|
||||
|
||||
|
||||
def _cutlass_include_paths() -> list[str]:
|
||||
cutlass_path = _cutlass_path()
|
||||
return [
|
||||
# Use realpath to get canonical absolute paths, in order not to mess up cache keys
|
||||
os.path.realpath(os.path.join(cutlass_path, path))
|
||||
for path in _cutlass_paths()
|
||||
]
|
||||
|
||||
|
||||
def _cuda_compiler() -> str | None:
|
||||
if cuda_env.nvcc_exist(config.cuda.cuda_cxx):
|
||||
return config.cuda.cuda_cxx
|
||||
if config.is_fbcode():
|
||||
return os.path.join(build_paths.sdk_home, "bin", "nvcc")
|
||||
if cuda_env.nvcc_exist(os.getenv("CUDACXX")):
|
||||
return os.getenv("CUDACXX", "")
|
||||
if cuda_env.nvcc_exist(os.getenv("CUDA_HOME")):
|
||||
return os.path.realpath(os.path.join(os.getenv("CUDA_HOME", ""), "bin/nvcc"))
|
||||
return "nvcc"
|
||||
|
||||
|
||||
def _cuda_lib_options() -> list[str]:
|
||||
"""
|
||||
Util function for CUTLASS backend to find the correct CUDA libraries.
|
||||
"""
|
||||
_set_gpu_runtime_env() # cpp_extension consults the env
|
||||
from torch.utils import cpp_extension
|
||||
|
||||
lpaths = cpp_extension.library_paths(device_type="cuda")
|
||||
if use_re_build():
|
||||
lpaths += [
|
||||
build_paths.sdk_lib,
|
||||
os.path.join(build_paths.sdk_lib, "stubs"),
|
||||
]
|
||||
extra_ldflags: list[str] = []
|
||||
if is_linux():
|
||||
_transform_cuda_paths(lpaths)
|
||||
for path in lpaths:
|
||||
if "torch/lib" in path:
|
||||
# don't want to depend on pytorch
|
||||
continue
|
||||
extra_ldflags.append(f"-L{path}")
|
||||
# -rpath ensures the DLL can find its dependencies when loaded, even
|
||||
# if the library path is non-standard.
|
||||
# But do not add the stubs folder to rpath as the driver is expected to be found at runtime
|
||||
if os.path.basename(path) != "stubs":
|
||||
extra_ldflags.extend(["-Xlinker", f"-rpath={path}"])
|
||||
extra_ldflags.append("-lcuda")
|
||||
extra_ldflags.append("-lcudart")
|
||||
else:
|
||||
raise NotImplementedError(
|
||||
"Unsupported env, failed to find cuda libs! Currently only Linux is supported."
|
||||
)
|
||||
return extra_ldflags
|
||||
|
||||
|
||||
def _nvcc_host_compiler_options() -> list[str]:
|
||||
return [
|
||||
"-fPIC",
|
||||
"-fno-strict-aliasing",
|
||||
"-fvisibility=hidden",
|
||||
"-Wconversion",
|
||||
]
|
||||
|
||||
|
||||
def _nvcc_arch_as_compile_option() -> str:
|
||||
arch = cuda_env.get_cuda_arch()
|
||||
if arch == "90":
|
||||
# Required by cutlass compilation.
|
||||
return "90a"
|
||||
if arch == "100":
|
||||
return "100a"
|
||||
if arch == "101":
|
||||
return "101a"
|
||||
if arch == "103":
|
||||
return "103a"
|
||||
if arch == "110":
|
||||
return "110a"
|
||||
if arch == "120":
|
||||
return "120a"
|
||||
if arch == "121":
|
||||
return "121a"
|
||||
return arch
|
||||
|
||||
|
||||
def _nvcc_compiler_options() -> list[str]:
|
||||
arch = _nvcc_arch_as_compile_option()
|
||||
code = [f"sm_{arch}", f"compute_{arch}"]
|
||||
if config.cuda.enable_cuda_lto:
|
||||
code += [f"lto_{arch}"]
|
||||
options = [
|
||||
"-t=0",
|
||||
"-DCUTLASS_ENABLE_TENSOR_CORE_MMA=1",
|
||||
"-DCUTLASS_ENABLE_SM90_EXTENDED_MMA_SHAPES=1",
|
||||
"-DCUTE_SM90_EXTENDED_MMA_SHAPES_ENABLED",
|
||||
"-w",
|
||||
f"-gencode=arch=compute_{arch},code=[{','.join(code)}]",
|
||||
config.cutlass.compile_opt_level,
|
||||
"-std=c++20",
|
||||
"--expt-relaxed-constexpr",
|
||||
"-DNDEBUG",
|
||||
]
|
||||
if config.is_fbcode():
|
||||
options.extend(["-ccbin", os.path.dirname(build_paths.gcc)])
|
||||
if config.cutlass.enable_debug_info:
|
||||
options.extend(["-lineinfo", "-g", "-DCUTLASS_DEBUG_TRACE_LEVEL=1"])
|
||||
if config.cuda.enable_ptxas_info:
|
||||
options.extend(
|
||||
[
|
||||
"--keep", # Keep the intermediate files for debugging (including ptx, sass, cubin etc.)
|
||||
"--ptxas-options=--warn-on-local-memory-usage", # warn us if local memory is used in CUDA Kernels
|
||||
"--ptxas-options=--warn-on-spills", # warn us if register spilling happens in CUDA Kernels
|
||||
"--resource-usage", # Report on CUDA resource usage (shared mem, registers etc.)
|
||||
"--source-in-ptx",
|
||||
]
|
||||
) # Annotate the ptx file with source information
|
||||
if config.cutlass.use_fast_math:
|
||||
options.extend(
|
||||
[
|
||||
"--use_fast_math",
|
||||
"-DCUTLASS_USE_TANH_FOR_SIGMOID=1",
|
||||
]
|
||||
)
|
||||
return options
|
||||
|
||||
|
||||
def cuda_compile_command(
|
||||
src_files: list[str],
|
||||
dst_file: str,
|
||||
dst_file_ext: str,
|
||||
extra_args: list[str] | None = None,
|
||||
) -> str:
|
||||
if extra_args is None:
|
||||
extra_args = []
|
||||
if use_re_build():
|
||||
build_path = os.path.dirname(dst_file)
|
||||
include_paths = _clone_cutlass_paths(build_path)
|
||||
src_files = [os.path.basename(src_file) for src_file in src_files]
|
||||
dst_file = os.path.basename(dst_file)
|
||||
else:
|
||||
include_paths = _cutlass_include_paths()
|
||||
cuda_lib_options = _cuda_lib_options()
|
||||
nvcc_host_compiler_options = _nvcc_host_compiler_options()
|
||||
nvcc_compiler_options = _nvcc_compiler_options()
|
||||
options = (
|
||||
nvcc_compiler_options
|
||||
+ extra_args
|
||||
+ [
|
||||
f"-Xcompiler {opt}" if "=" in opt else f"-Xcompiler={opt}"
|
||||
for opt in nvcc_host_compiler_options
|
||||
]
|
||||
+ ["-I" + path for path in include_paths]
|
||||
+ cuda_lib_options
|
||||
)
|
||||
src_file = " ".join(src_files)
|
||||
res = ""
|
||||
if dst_file_ext == "o":
|
||||
res = f"{_cuda_compiler()} {' '.join(options)} -c -o {dst_file} {src_file}"
|
||||
elif dst_file_ext == "so":
|
||||
options.append("-shared")
|
||||
res = f"{_cuda_compiler()} {' '.join(options)} -o {dst_file} {src_file}"
|
||||
elif dst_file_ext == "exe":
|
||||
res = f"{_cuda_compiler()} {' '.join(options)} -o {dst_file} {src_file}"
|
||||
else:
|
||||
raise NotImplementedError(f"Unsupported output file suffix {dst_file_ext}!")
|
||||
if log.isEnabledFor(logging.DEBUG):
|
||||
log.debug("CUDA command: %s", res)
|
||||
else:
|
||||
autotuning_log.debug("CUDA command: %s", res)
|
||||
return res
|
||||
|
||||
|
||||
class CUDACompileSourceCapturingContext:
|
||||
# Helper class for Benchmarking and Testing CUTLASS Kernels in isolation.
|
||||
# Can be used to capture the sourcecode passed to CUDACodeCache.compile
|
||||
|
||||
def __init__(self):
|
||||
self.sources = []
|
||||
self._compile_patch = None
|
||||
|
||||
def __enter__(self, *args, **kwargs):
|
||||
import unittest.mock as mock
|
||||
|
||||
import torch._inductor.codecache
|
||||
|
||||
_compile_method_orig = torch._inductor.codecache.CUDACodeCache.compile
|
||||
|
||||
def my_compile(source_code, dst_file_ext, extra_args: list[str] | None = None):
|
||||
self.sources.append(source_code)
|
||||
return _compile_method_orig(source_code, dst_file_ext)
|
||||
|
||||
# pyrefly: ignore [bad-assignment]
|
||||
self._compile_patch = mock.patch(
|
||||
"torch._inductor.codecache.CUDACodeCache.compile", my_compile
|
||||
)
|
||||
self._compile_patch.__enter__(*args, **kwargs) # type: ignore[union-attr]
|
||||
return self
|
||||
|
||||
def __exit__(self, *args, **kwargs):
|
||||
self._compile_patch.__exit__(*args, **kwargs) # type: ignore[union-attr]
|
||||
|
||||
|
||||
def cuda_standalone_runner_compile_command(srcpath: Path, exepath: Path):
|
||||
# returns command string to compile a (captured) CUDA GEMM Kernel source to a standalone executable that's ready to run
|
||||
# Passes the correct preprocessor define to nvcc to ensure the standalone runner is enabled.
|
||||
|
||||
extra_args = ["-DGENERATE_STANDALONE_RUNNER=1", "-DCUTLASS_DEBUG_TRACE_LEVEL=1"]
|
||||
compile_command = cuda_compile_command(
|
||||
[str(srcpath)], str(exepath), "exe", extra_args=extra_args
|
||||
)
|
||||
return compile_command
|
||||
@@ -0,0 +1,54 @@
|
||||
import functools
|
||||
import logging
|
||||
import shutil
|
||||
|
||||
import torch
|
||||
from torch._inductor.utils import clear_on_fresh_cache
|
||||
|
||||
from ... import config
|
||||
|
||||
|
||||
log = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@clear_on_fresh_cache
|
||||
@functools.lru_cache(1)
|
||||
def get_cuda_arch() -> str | None:
|
||||
try:
|
||||
cuda_arch = config.cuda.arch
|
||||
if cuda_arch is None:
|
||||
# Get Compute Capability of the first Visible device
|
||||
major, minor = torch.cuda.get_device_capability(0)
|
||||
return str(major * 10 + minor)
|
||||
return str(cuda_arch)
|
||||
except Exception:
|
||||
log.exception("Error getting cuda arch")
|
||||
return None
|
||||
|
||||
|
||||
@clear_on_fresh_cache
|
||||
@functools.lru_cache(1)
|
||||
def is_datacenter_blackwell_arch() -> bool:
|
||||
arch = get_cuda_arch()
|
||||
if arch is None:
|
||||
return False
|
||||
arch_number = int(arch)
|
||||
return arch_number >= 100 and arch_number < 110
|
||||
|
||||
|
||||
@clear_on_fresh_cache
|
||||
@functools.lru_cache(1)
|
||||
def get_cuda_version() -> str | None:
|
||||
try:
|
||||
cuda_version = config.cuda.version
|
||||
if cuda_version is None:
|
||||
cuda_version = torch.version.cuda
|
||||
return cuda_version
|
||||
except Exception:
|
||||
log.exception("Error getting cuda version")
|
||||
return None
|
||||
|
||||
|
||||
@functools.cache
|
||||
def nvcc_exist(nvcc_path: str | None = "nvcc") -> bool:
|
||||
return nvcc_path is not None and shutil.which(nvcc_path) is not None
|
||||
+365
@@ -0,0 +1,365 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import torch
|
||||
|
||||
from ..common import (
|
||||
DeviceOpOverrides,
|
||||
register_device_op_overrides,
|
||||
TritonScratchWorkspace,
|
||||
)
|
||||
|
||||
|
||||
class CUDADeviceOpOverrides(DeviceOpOverrides):
|
||||
"""
|
||||
CUDA-specific codegen functions, see DeviceOpOverrides for details
|
||||
"""
|
||||
|
||||
def import_get_raw_stream_as(self, name: str) -> str:
|
||||
return f"from torch._C import _cuda_getCurrentRawStream as {name}"
|
||||
|
||||
def set_device(self, device_idx: int) -> str:
|
||||
return f"torch.cuda.set_device({device_idx})"
|
||||
|
||||
def synchronize(self) -> str:
|
||||
return "torch.cuda.synchronize()"
|
||||
|
||||
def device_guard(self, device_idx: int) -> str:
|
||||
return f"torch.cuda._DeviceGuard({device_idx})"
|
||||
|
||||
def cpp_device_guard(self) -> str:
|
||||
return "at::cuda::CUDAGuard"
|
||||
|
||||
def cpp_aoti_device_guard(self) -> str:
|
||||
return "AOTICudaGuard"
|
||||
|
||||
def cpp_stream_guard(self) -> str:
|
||||
return "at::cuda::CUDAStreamGuard"
|
||||
|
||||
def cpp_aoti_stream_guard(self) -> str:
|
||||
return "AOTICudaStreamGuard"
|
||||
|
||||
def cpp_getStreamFromExternal(self) -> str:
|
||||
return "at::cuda::getStreamFromExternal"
|
||||
|
||||
def kernel_header(self) -> str:
|
||||
source_codes = """
|
||||
#include <c10/cuda/CUDAGuard.h>
|
||||
#include <c10/cuda/CUDAStream.h>
|
||||
#include <ATen/cuda/EmptyTensor.h>
|
||||
"""
|
||||
return source_codes
|
||||
|
||||
def kernel_driver(self) -> str:
|
||||
source_codes = """
|
||||
#define CUDA_DRIVER_CHECK(EXPR) \\
|
||||
do { \\
|
||||
CUresult code = EXPR; \\
|
||||
const char *msg; \\
|
||||
CUresult code_get_error = cuGetErrorString(code, &msg); \\
|
||||
if (code_get_error != CUDA_SUCCESS) { \\
|
||||
throw std::runtime_error( \\
|
||||
std::string("CUDA driver error: ") + \\
|
||||
std::string("invalid error code!")); \\
|
||||
} \\
|
||||
if (code != CUDA_SUCCESS) { \\
|
||||
throw std::runtime_error( \\
|
||||
std::string("CUDA driver error: ") + \\
|
||||
std::string(msg)); \\
|
||||
} \\
|
||||
} while (0);
|
||||
|
||||
static inline CUfunction loadKernel(
|
||||
std::string filePath,
|
||||
const std::string &funcName,
|
||||
uint32_t sharedMemBytes,
|
||||
const std::optional<std::string> &cubinDir = std::nullopt) {
|
||||
if (cubinDir) {
|
||||
std::filesystem::path p1{*cubinDir};
|
||||
std::filesystem::path p2{filePath};
|
||||
filePath = (p1 / p2.filename()).string();
|
||||
}
|
||||
|
||||
CUmodule mod;
|
||||
CUfunction func;
|
||||
CUDA_DRIVER_CHECK(cuModuleLoad(&mod, filePath.c_str()));
|
||||
CUDA_DRIVER_CHECK(cuModuleGetFunction(&func, mod, funcName.c_str()));
|
||||
if (sharedMemBytes > 0) {
|
||||
CUDA_DRIVER_CHECK(cuFuncSetAttribute(
|
||||
func,
|
||||
CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
|
||||
sharedMemBytes
|
||||
))
|
||||
}
|
||||
return func;
|
||||
}
|
||||
|
||||
static inline CUfunction loadKernel(const void* start, const std::string &funcName, uint32_t sharedMemBytes) {
|
||||
CUmodule mod;
|
||||
CUfunction func;
|
||||
CUDA_DRIVER_CHECK(cuModuleLoadData(&mod, start));
|
||||
CUDA_DRIVER_CHECK(cuModuleGetFunction(&func, mod, funcName.c_str()));
|
||||
if (sharedMemBytes > 0) {
|
||||
CUDA_DRIVER_CHECK(cuFuncSetAttribute(
|
||||
func,
|
||||
CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
|
||||
sharedMemBytes
|
||||
))
|
||||
}
|
||||
return func;
|
||||
}
|
||||
|
||||
static inline void launchKernel(
|
||||
CUfunction func,
|
||||
uint32_t gridX,
|
||||
uint32_t gridY,
|
||||
uint32_t gridZ,
|
||||
uint32_t numWarps,
|
||||
uint32_t sharedMemBytes,
|
||||
void* args[],
|
||||
cudaStream_t stream) {
|
||||
CUDA_DRIVER_CHECK(cuLaunchKernel(
|
||||
func, gridX, gridY, gridZ, 32*numWarps, 1, 1, sharedMemBytes, stream, args, nullptr
|
||||
));
|
||||
}
|
||||
"""
|
||||
if torch.version.hip is not None:
|
||||
# Adjusting the warp size to GPU supported wavefront size on AMD GPU
|
||||
prop = torch.cuda.get_device_properties(torch.cuda.current_device())
|
||||
source_codes = source_codes.replace(
|
||||
"32*numWarps", str(prop.warp_size) + "*numWarps"
|
||||
)
|
||||
return source_codes
|
||||
|
||||
def tma_descriptor_helpers(self) -> str:
|
||||
"""
|
||||
CUDA helper functions for initializing TMA Descriptors on host side
|
||||
"""
|
||||
if torch.version.hip is not None:
|
||||
raise RuntimeError("Host-side TMA descriptors not supported on HIP.")
|
||||
|
||||
# helper functions for initializing 1D and 2D TMA descriptors in C++. borrowed from the Triton code here:
|
||||
# Old APIs (fill(1|2)DTMADescriptor):
|
||||
# https://github.com/triton-lang/triton/blob/6af4f88591c85de079d8a36a4d7dba67918e2b39/third_party/nvidia/backend/driver.c#L283
|
||||
# New APIs (fillTMADescriptor):
|
||||
# https://github.com/triton-lang/triton/blob/main/third_party/nvidia/backend/driver.c#L283
|
||||
return """
|
||||
#if !defined(USE_ROCM) && defined(CUDA_VERSION) && CUDA_VERSION >= 12000
|
||||
[[maybe_unused]] static void init1DTMADescriptor(
|
||||
CUtensorMap* m,
|
||||
void* globalAddress,
|
||||
uint64_t dim,
|
||||
uint32_t blockDim,
|
||||
uint32_t elementSize) {
|
||||
uint64_t dims[1] = {dim};
|
||||
uint64_t globalStrides[1] = {dim * elementSize};
|
||||
uint32_t tensorDims[1] = {blockDim};
|
||||
uint32_t elementStrides[1] = {1};
|
||||
|
||||
CUtensorMapDataType type;
|
||||
switch (elementSize) {
|
||||
case 1:
|
||||
type = CU_TENSOR_MAP_DATA_TYPE_UINT8;
|
||||
break;
|
||||
case 2:
|
||||
type = CU_TENSOR_MAP_DATA_TYPE_UINT16;
|
||||
break;
|
||||
case 4:
|
||||
type = CU_TENSOR_MAP_DATA_TYPE_UINT32;
|
||||
break;
|
||||
default:
|
||||
throw std::runtime_error("elementSize must be 1, 2, or 4");
|
||||
}
|
||||
|
||||
if (elementSize * blockDim < 32) {
|
||||
throw std::runtime_error("block size too small");
|
||||
}
|
||||
|
||||
int rank = 1;
|
||||
|
||||
CUDA_DRIVER_CHECK(cuTensorMapEncodeTiled(
|
||||
m, type, rank, globalAddress, dims,
|
||||
globalStrides, tensorDims, elementStrides, CU_TENSOR_MAP_INTERLEAVE_NONE,
|
||||
CU_TENSOR_MAP_SWIZZLE_NONE, CU_TENSOR_MAP_L2_PROMOTION_NONE,
|
||||
CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE));
|
||||
}
|
||||
|
||||
[[maybe_unused]] static void init2DTMADescriptor(
|
||||
CUtensorMap* m,
|
||||
void* globalAddress,
|
||||
uint64_t dim1,
|
||||
uint64_t dim0,
|
||||
uint32_t blockDim1,
|
||||
uint32_t blockDim0,
|
||||
uint32_t elementSize) {
|
||||
uint64_t dims[2] = {dim0, dim1};
|
||||
uint32_t tensorDims[2] = {blockDim0, blockDim1};
|
||||
uint64_t globalStrides[2] = {dims[0] * elementSize,
|
||||
dims[0] * dims[1] * elementSize};
|
||||
uint32_t elementStrides[2] = {1, 1};
|
||||
|
||||
CUtensorMapDataType type;
|
||||
switch (elementSize) {
|
||||
case 1:
|
||||
type = CU_TENSOR_MAP_DATA_TYPE_UINT8;
|
||||
break;
|
||||
case 2:
|
||||
type = CU_TENSOR_MAP_DATA_TYPE_UINT16;
|
||||
break;
|
||||
case 4:
|
||||
type = CU_TENSOR_MAP_DATA_TYPE_UINT32;
|
||||
break;
|
||||
default:
|
||||
throw std::runtime_error("elementSize must be 1, 2, or 4");
|
||||
}
|
||||
|
||||
int rank = 2;
|
||||
|
||||
CUtensorMapSwizzle swizzle = CU_TENSOR_MAP_SWIZZLE_128B;
|
||||
uint32_t contigDimSizeInByte = elementSize * tensorDims[0];
|
||||
if (contigDimSizeInByte >= 128) {
|
||||
swizzle = CU_TENSOR_MAP_SWIZZLE_128B;
|
||||
} else if (contigDimSizeInByte >= 64) {
|
||||
swizzle = CU_TENSOR_MAP_SWIZZLE_64B;
|
||||
} else if (contigDimSizeInByte >= 32) {
|
||||
swizzle = CU_TENSOR_MAP_SWIZZLE_32B;
|
||||
} else {
|
||||
throw std::runtime_error("block size too small");
|
||||
}
|
||||
|
||||
if (contigDimSizeInByte > 128) {
|
||||
tensorDims[0] = 128 / elementSize;
|
||||
}
|
||||
|
||||
CUDA_DRIVER_CHECK(cuTensorMapEncodeTiled(
|
||||
m, type, rank, globalAddress, dims,
|
||||
globalStrides, tensorDims, elementStrides, CU_TENSOR_MAP_INTERLEAVE_NONE,
|
||||
swizzle, CU_TENSOR_MAP_L2_PROMOTION_L2_128B,
|
||||
CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE));
|
||||
}
|
||||
|
||||
[[maybe_unused]] static void initTMADescriptor(
|
||||
CUtensorMap* m,
|
||||
void* globalAddress,
|
||||
int elemSize,
|
||||
int rank,
|
||||
uint32_t* blockSize,
|
||||
uint64_t* shape,
|
||||
uint64_t* stride
|
||||
) {
|
||||
uint32_t elementStrides[5] = {1, 1, 1, 1, 1};
|
||||
uint32_t blockSizeInt[5];
|
||||
uint64_t shapeInt[5];
|
||||
uint64_t stridesLL[5];
|
||||
|
||||
// Reorder blockSize (reverse the order)
|
||||
for (int i = 0; i < rank; ++i) {
|
||||
blockSizeInt[rank - i - 1] = blockSize[i];
|
||||
}
|
||||
|
||||
// Reorder shape (reverse the order)
|
||||
for (int i = 0; i < rank; ++i) {
|
||||
shapeInt[rank - i - 1] = shape[i];
|
||||
}
|
||||
|
||||
// Reorder and calculate strides
|
||||
for (int i = 0; i + 1 < rank; ++i) {
|
||||
stridesLL[rank - i - 2] = elemSize * stride[i];
|
||||
}
|
||||
stridesLL[rank - 1] =
|
||||
shapeInt[rank - 1] * (rank == 1 ? elemSize : stridesLL[rank - 2]);
|
||||
|
||||
CUtensorMapDataType type;
|
||||
// In Triton this is computed ahead of time; but for simplicity
|
||||
// in the PyTorch version we copied this code from the old
|
||||
// TMA API handling (i.e. init2DTMADescriptor)
|
||||
switch (elemSize) {
|
||||
case 1:
|
||||
type = CU_TENSOR_MAP_DATA_TYPE_UINT8;
|
||||
break;
|
||||
case 2:
|
||||
type = CU_TENSOR_MAP_DATA_TYPE_UINT16;
|
||||
break;
|
||||
case 4:
|
||||
type = CU_TENSOR_MAP_DATA_TYPE_UINT32;
|
||||
break;
|
||||
default:
|
||||
throw std::runtime_error("elemSize must be 1, 2, or 4");
|
||||
}
|
||||
|
||||
// Calculate the size of the most contiguous dimension in bytes
|
||||
CUtensorMapSwizzle swizzle = CU_TENSOR_MAP_SWIZZLE_128B;
|
||||
uint32_t contigDimSizeInByte = elemSize * blockSizeInt[0];
|
||||
if (rank == 1) {
|
||||
// rank 1 should not be swizzled
|
||||
swizzle = CU_TENSOR_MAP_SWIZZLE_NONE;
|
||||
} else if (contigDimSizeInByte >= 128) {
|
||||
swizzle = CU_TENSOR_MAP_SWIZZLE_128B;
|
||||
} else if (contigDimSizeInByte >= 64) {
|
||||
swizzle = CU_TENSOR_MAP_SWIZZLE_64B;
|
||||
} else if (contigDimSizeInByte >= 32) {
|
||||
swizzle = CU_TENSOR_MAP_SWIZZLE_32B;
|
||||
} else {
|
||||
throw std::runtime_error("block size too small");
|
||||
}
|
||||
|
||||
CUDA_DRIVER_CHECK(cuTensorMapEncodeTiled(
|
||||
m, type, rank, globalAddress,
|
||||
shapeInt, stridesLL, blockSizeInt, elementStrides,
|
||||
CU_TENSOR_MAP_INTERLEAVE_NONE, (CUtensorMapSwizzle)swizzle,
|
||||
CU_TENSOR_MAP_L2_PROMOTION_L2_128B, CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE));
|
||||
}
|
||||
|
||||
struct StableTMADescriptor {
|
||||
CUtensorMap m;
|
||||
uint32_t block_shape[5];
|
||||
uint64_t global_shape[5];
|
||||
uint64_t strides[5];
|
||||
};
|
||||
#endif
|
||||
"""
|
||||
|
||||
def cpp_stream_type(self) -> str:
|
||||
return "cudaStream_t"
|
||||
|
||||
def aoti_get_stream(self) -> str:
|
||||
return "aoti_torch_get_current_cuda_stream"
|
||||
|
||||
def cpp_kernel_type(self) -> str:
|
||||
return "CUfunction"
|
||||
|
||||
def cpp_device_ptr(self) -> str:
|
||||
return "CUdeviceptr"
|
||||
|
||||
def cpp_scratch(
|
||||
self, idx: int, workspace: TritonScratchWorkspace, prefix: str | None = None
|
||||
) -> tuple[list[str], str] | None:
|
||||
prefix = f"{prefix}_" if prefix else ""
|
||||
var_name = f"{prefix}scratch_{idx}"
|
||||
if workspace.size > 0:
|
||||
size_expr = (
|
||||
f"static_cast<int64_t>({workspace.size}) * grid_0 * grid_1 * grid_2"
|
||||
)
|
||||
size_array = f"int64_t {var_name}_size[] = {{{size_expr}}};"
|
||||
stride_array = f"int64_t {var_name}_stride[] = {{1}};"
|
||||
device_type = "cached_torch_device_type_cuda"
|
||||
device_idx = "device_idx_"
|
||||
|
||||
return (
|
||||
[
|
||||
f"{size_array}",
|
||||
f"{stride_array}",
|
||||
f"AtenTensorHandle {var_name}_handle;",
|
||||
(
|
||||
f"AOTI_TORCH_ERROR_CODE_CHECK(aoti_torch_empty_strided(1, {var_name}_size, {var_name}_stride, "
|
||||
f"{workspace.generate_dtype_str()}, {device_type}, {device_idx}, &{var_name}_handle));"
|
||||
),
|
||||
f"RAIIAtenTensorHandle {var_name}_tensor({var_name}_handle);",
|
||||
f"CUdeviceptr {var_name} = reinterpret_cast<CUdeviceptr>({var_name}_tensor.data_ptr());",
|
||||
],
|
||||
var_name,
|
||||
)
|
||||
else:
|
||||
return [f"CUdeviceptr {var_name} = 0;"], var_name
|
||||
|
||||
|
||||
register_device_op_overrides("cuda", CUDADeviceOpOverrides())
|
||||
+186
@@ -0,0 +1,186 @@
|
||||
# mypy: allow-untyped-defs
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any, TYPE_CHECKING
|
||||
|
||||
from ..scheduler import (
|
||||
BaseSchedulerNode,
|
||||
BaseScheduling,
|
||||
FusedSchedulerNode,
|
||||
Scheduler,
|
||||
SchedulerNode,
|
||||
)
|
||||
from .cutedsl.cutedsl_scheduling import CuteDSLScheduling
|
||||
from .cutlass.scheduling import CUTLASSScheduling
|
||||
from .nv_universal_gemm.nv_universal_gemm_scheduling import NVUniversalGemmScheduling
|
||||
from .rocm.rocm_cpp_scheduling import ROCmCPPScheduling
|
||||
from .triton import TritonScheduling
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Sequence
|
||||
from typing import TypeAlias
|
||||
|
||||
from sympy import Expr
|
||||
|
||||
import torch
|
||||
from torch.utils._ordered_set import OrderedSet
|
||||
|
||||
from .common import BackendFeature
|
||||
|
||||
_IntLike: TypeAlias = int | Expr
|
||||
|
||||
|
||||
class CUDACombinedScheduling(BaseScheduling):
|
||||
"""
|
||||
Scheduler for CUDA Kernels, which delegates calls as appropriate
|
||||
to the CUDA-C++ and Triton Schedulers, which both work for CUDA devices
|
||||
and use a unified-wrapper for codegen.
|
||||
|
||||
If Scheduling code needs to be specialized for the case of mixed Triton / CUDA C++ code,
|
||||
this would also be the place to do it.
|
||||
"""
|
||||
|
||||
def __init__(self, scheduler: Scheduler | None) -> None:
|
||||
super().__init__(scheduler)
|
||||
self._triton_scheduling = TritonScheduling(scheduler)
|
||||
self._cutlass_scheduling = CUTLASSScheduling(scheduler)
|
||||
self._rocm_cpp_scheduling = ROCmCPPScheduling(scheduler)
|
||||
self._cutedsl_scheduling = CuteDSLScheduling(scheduler)
|
||||
self._nv_universal_gemm_scheduling = NVUniversalGemmScheduling(scheduler)
|
||||
|
||||
def get_backend_features(self, device: torch.device) -> OrderedSet[BackendFeature]:
|
||||
return self._triton_scheduling.get_backend_features(device)
|
||||
|
||||
def choose_node_backend(self, node: BaseSchedulerNode) -> BaseScheduling:
|
||||
if self._cutlass_scheduling.is_cutlass_template(node):
|
||||
return self._cutlass_scheduling
|
||||
if self._rocm_cpp_scheduling.is_rocm_cpp_template(node):
|
||||
return self._rocm_cpp_scheduling
|
||||
if self._cutedsl_scheduling.is_cutedsl_template(node):
|
||||
return self._cutedsl_scheduling
|
||||
if self._nv_universal_gemm_scheduling.is_nv_universal_gemm_template(node):
|
||||
return self._nv_universal_gemm_scheduling
|
||||
return self._triton_scheduling
|
||||
|
||||
def can_fuse_vertical(
|
||||
self, node1: BaseSchedulerNode, node2: BaseSchedulerNode
|
||||
) -> bool:
|
||||
if self._cutlass_scheduling.can_fuse_vertical(node1, node2):
|
||||
return True
|
||||
elif self._cutlass_scheduling.is_cutlass_template(
|
||||
node1
|
||||
) or self._cutlass_scheduling.is_cutlass_template(node2):
|
||||
return False
|
||||
# CuteDSL doesn't support vertical fusion currently
|
||||
elif self._cutedsl_scheduling.is_cutedsl_template(
|
||||
node1
|
||||
) or self._cutedsl_scheduling.is_cutedsl_template(node2):
|
||||
return False
|
||||
# NVIDIA Universal GEMM doesn't support vertical fusion currently
|
||||
elif self._nv_universal_gemm_scheduling.is_nv_universal_gemm_template(
|
||||
node1
|
||||
) or self._nv_universal_gemm_scheduling.is_nv_universal_gemm_template(node2):
|
||||
return False
|
||||
return self._triton_scheduling.can_fuse_vertical(node1, node2)
|
||||
|
||||
def can_fuse_horizontal(
|
||||
self, node1: BaseSchedulerNode, node2: BaseSchedulerNode
|
||||
) -> bool:
|
||||
for node in (node1, node2):
|
||||
if self._cutlass_scheduling.is_cutlass_template(node):
|
||||
return self._cutlass_scheduling.can_fuse_horizontal(
|
||||
node1, node2
|
||||
) # always False at the moment
|
||||
if self._cutedsl_scheduling.is_cutedsl_template(node):
|
||||
return self._cutedsl_scheduling.can_fuse_horizontal(
|
||||
node1, node2
|
||||
) # always False at the moment
|
||||
if self._nv_universal_gemm_scheduling.is_nv_universal_gemm_template(node):
|
||||
return self._nv_universal_gemm_scheduling.can_fuse_horizontal(
|
||||
node1, node2
|
||||
) # always False at the moment
|
||||
return self._triton_scheduling.can_fuse_horizontal(node1, node2)
|
||||
|
||||
def group_fn(
|
||||
self, sizes: Sequence[Sequence[_IntLike]]
|
||||
) -> tuple[tuple[_IntLike, ...], ...]:
|
||||
return self._triton_scheduling.group_fn(sizes)
|
||||
|
||||
def codegen_template(
|
||||
self,
|
||||
template_node: BaseSchedulerNode,
|
||||
epilogue_nodes: Sequence[BaseSchedulerNode],
|
||||
prologue_nodes: Sequence[BaseSchedulerNode],
|
||||
) -> str | None:
|
||||
if self._cutlass_scheduling.is_cutlass_template(template_node):
|
||||
assert not prologue_nodes
|
||||
return self._cutlass_scheduling.codegen_template(
|
||||
template_node, epilogue_nodes, prologue_nodes
|
||||
)
|
||||
elif self._rocm_cpp_scheduling.is_rocm_cpp_template(template_node):
|
||||
assert not epilogue_nodes
|
||||
assert not prologue_nodes
|
||||
return self._rocm_cpp_scheduling.codegen_template(
|
||||
template_node, epilogue_nodes, prologue_nodes
|
||||
)
|
||||
elif self._cutedsl_scheduling.is_cutedsl_template(template_node):
|
||||
# TODO remove this when we add epilogue support
|
||||
assert not epilogue_nodes
|
||||
assert not prologue_nodes
|
||||
return self._cutedsl_scheduling.codegen_template(
|
||||
template_node, epilogue_nodes, prologue_nodes
|
||||
)
|
||||
elif self._nv_universal_gemm_scheduling.is_nv_universal_gemm_template(
|
||||
template_node
|
||||
):
|
||||
# NVIDIA Universal GEMM doesn't support epilogue/prologue fusion yet
|
||||
assert not epilogue_nodes
|
||||
assert not prologue_nodes
|
||||
return self._nv_universal_gemm_scheduling.codegen_template(
|
||||
template_node, epilogue_nodes, prologue_nodes
|
||||
)
|
||||
else:
|
||||
return self._triton_scheduling.codegen_template(
|
||||
template_node, epilogue_nodes, prologue_nodes
|
||||
)
|
||||
|
||||
def codegen_mix_order_reduction(self, node):
|
||||
return self._triton_scheduling.codegen_mix_order_reduction(node)
|
||||
|
||||
def codegen_node(self, node: FusedSchedulerNode | SchedulerNode) -> None:
|
||||
return self._triton_scheduling.codegen_node(node)
|
||||
|
||||
def codegen_sync(self) -> None:
|
||||
return self._triton_scheduling.codegen_sync()
|
||||
|
||||
def flush(self) -> None:
|
||||
return self._triton_scheduling.flush()
|
||||
|
||||
def codegen_combo_kernel(self, *args: Any, **kwargs: Any) -> None:
|
||||
return self._triton_scheduling.codegen_combo_kernel(*args, **kwargs)
|
||||
|
||||
def benchmark_fused_nodes(
|
||||
self, nodes: Sequence[BaseSchedulerNode]
|
||||
) -> tuple[float, str]:
|
||||
return self._triton_scheduling.benchmark_fused_nodes(nodes)
|
||||
|
||||
def benchmark_codegened_module(self, module):
|
||||
return self._triton_scheduling.benchmark_codegened_module(module)
|
||||
|
||||
def generate_kernel_code_from_nodes(
|
||||
self,
|
||||
nodes: Sequence[Any],
|
||||
benchmark_kernel: bool = False,
|
||||
hint_override: int | None = None,
|
||||
) -> str:
|
||||
return self._triton_scheduling.generate_kernel_code_from_nodes(
|
||||
nodes, benchmark_kernel, hint_override=hint_override
|
||||
)
|
||||
|
||||
def benchmark_combo_kernel(
|
||||
self, node_list: Sequence[BaseSchedulerNode], node_benchmark_results
|
||||
) -> tuple[float, float, list[str | None]]:
|
||||
return self._triton_scheduling.benchmark_combo_kernel(
|
||||
node_list, node_benchmark_results
|
||||
)
|
||||
+101
@@ -0,0 +1,101 @@
|
||||
"""
|
||||
Custom extern kernel codegen registry.
|
||||
|
||||
This module provides a registry that maps operators to their custom codegen
|
||||
implementations. This allows us to keep all custom implementations in one place
|
||||
and easily extend support for both Python and C++ wrappers.
|
||||
|
||||
To add a new custom implementation:
|
||||
1. Create a codegen function with signature:
|
||||
def my_codegen(node: ir.FallbackKernel, writeline: Callable[[str], None]) -> None:
|
||||
2. Register it in CUSTOM_EXTERN_KERNEL_CODEGEN with the operator name as key
|
||||
|
||||
Example:
|
||||
CUSTOM_EXTERN_KERNEL_CODEGEN = {
|
||||
"torch.ops.higher_order.print": CustomCodegen(
|
||||
python=generate_print_python,
|
||||
cpp=generate_print_cpp, # Optional
|
||||
),
|
||||
}
|
||||
|
||||
Usage:
|
||||
codegen = CUSTOM_EXTERN_KERNEL_CODEGEN[op]
|
||||
codegen.python(node, writeline) # Direct attribute access
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, TYPE_CHECKING
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Callable
|
||||
|
||||
from .. import ir
|
||||
|
||||
# Type alias for codegen function signature (only used for type checking)
|
||||
CodegenFunc = Callable[[ir.FallbackKernel, Callable[[str], None]], None]
|
||||
|
||||
|
||||
@dataclass
|
||||
class CustomCodegen:
|
||||
"""
|
||||
Container for custom codegen implementations.
|
||||
|
||||
Attributes:
|
||||
python: Codegen function for Python wrapper (optional)
|
||||
cpp: Codegen function for C++ wrapper (optional)
|
||||
"""
|
||||
|
||||
python: Any = None
|
||||
cpp: Any = None
|
||||
|
||||
|
||||
def generate_print_python(
|
||||
node: ir.FallbackKernel,
|
||||
writeline: Callable[[str], None],
|
||||
) -> None:
|
||||
"""
|
||||
Generate a builtin print call for the print HOP fallback (Python wrapper).
|
||||
|
||||
This function generates Python code that calls the builtin print function
|
||||
with format string interpolation.
|
||||
|
||||
Args:
|
||||
node: The FallbackKernel IR node representing the print HOP call.
|
||||
writeline: A function that writes a line of code to the output buffer.
|
||||
|
||||
Example generated code:
|
||||
print('x = {}, y = {}'.format(buf0, buf1))
|
||||
print('x = {x}, y = {y}'.format(x=buf0, y=buf1))
|
||||
print('x = {}, y = {y}'.format(buf0, y=buf1))
|
||||
"""
|
||||
codegen_args: list[str] = node.codegen_args()
|
||||
codegen_kwargs: list[str] = node.codegen_kwargs()
|
||||
|
||||
# First arg is the format string
|
||||
if not codegen_args:
|
||||
raise ValueError(
|
||||
"generate_print_python requires a format string as the first positional argument"
|
||||
)
|
||||
format_str: str = codegen_args[0]
|
||||
|
||||
# Remaining args are positional arguments for .format()
|
||||
positional_args = codegen_args[1:]
|
||||
|
||||
args_str = ", ".join(positional_args + codegen_kwargs)
|
||||
writeline(
|
||||
f"print({format_str}.format({args_str}))"
|
||||
if args_str
|
||||
else f"print({format_str})"
|
||||
)
|
||||
|
||||
|
||||
# Registry mapping operator names to their custom codegen implementations
|
||||
# Usage: CUSTOM_EXTERN_KERNEL_CODEGEN[op_name].python(node, writeline)
|
||||
CUSTOM_EXTERN_KERNEL_CODEGEN: dict[str, CustomCodegen] = {
|
||||
"torch.ops.higher_order.print": CustomCodegen(
|
||||
python=generate_print_python,
|
||||
),
|
||||
}
|
||||
@@ -0,0 +1,8 @@
|
||||
# mypy: allow-untyped-defs
|
||||
from .cutedsl_template import CuteDSLTemplate, CuteDSLTemplateCaller
|
||||
|
||||
|
||||
__all__ = [
|
||||
"CuteDSLTemplate",
|
||||
"CuteDSLTemplateCaller",
|
||||
]
|
||||
+29
@@ -0,0 +1,29 @@
|
||||
# mypy: disable-error-code=import-not-found
|
||||
# pyrefly: ignore [import-error, missing-import]
|
||||
import cutlass.cute as cute
|
||||
|
||||
|
||||
@cute.jit # type: ignore[misc]
|
||||
def ssa_to_indexable(ssa_value: cute.TensorSSA, dtype: str) -> cute.Numeric:
|
||||
"""
|
||||
Convert SSA form to indexable non-SSA form.
|
||||
|
||||
Workaround for lack of gather support: SSA values cannot be used directly
|
||||
as indices in tensor loads. This converts SSA → fragment → scalar for indexing.
|
||||
"""
|
||||
frag = cute.make_rmem_tensor(1, dtype)
|
||||
frag.store(ssa_value)
|
||||
return frag[0]
|
||||
|
||||
|
||||
@cute.jit # type: ignore[misc]
|
||||
def result_to_ssa(value: cute.Numeric, dtype: str) -> cute.TensorSSA:
|
||||
"""
|
||||
Convert non-SSA result back to SSA form.
|
||||
|
||||
After performing operations with non-SSA values (like indexed loads),
|
||||
convert the result back to SSA form for further computation.
|
||||
"""
|
||||
frag = cute.make_rmem_tensor(1, dtype)
|
||||
frag[0] = value
|
||||
return frag.load()
|
||||
+613
@@ -0,0 +1,613 @@
|
||||
# mypy: allow-untyped-defs
|
||||
import contextlib
|
||||
import dataclasses
|
||||
import logging
|
||||
import textwrap
|
||||
from collections.abc import Callable
|
||||
from typing import Any
|
||||
|
||||
import sympy
|
||||
|
||||
import torch
|
||||
from torch._inductor import config
|
||||
from torch._inductor.codegen.common import (
|
||||
CSE,
|
||||
CSEVariable,
|
||||
IndentedBuffer,
|
||||
Kernel,
|
||||
ValueRanges,
|
||||
)
|
||||
from torch._inductor.ir import (
|
||||
BaseView,
|
||||
Buffer,
|
||||
ComputedBuffer,
|
||||
ExternKernel,
|
||||
InputBuffer,
|
||||
MutableBox,
|
||||
ReinterpretView,
|
||||
)
|
||||
from torch._inductor.ops_handler import StoreMode
|
||||
from torch._inductor.utils import OrderedSet
|
||||
from torch._inductor.virtualized import V
|
||||
|
||||
from ...utils import sympy_index_symbol
|
||||
from .cutedsl_op_overrides import CuteDSLOpOverrides
|
||||
|
||||
|
||||
# TODO setting the 'main' kernel w/ this suffix. We have 3 should probably just auto generate this
|
||||
MAIN_SUFFIX = "main"
|
||||
|
||||
|
||||
log = logging.getLogger(__name__)
|
||||
kernel_code_log = torch._logging.getArtifactLogger(__name__, "kernel_code")
|
||||
|
||||
|
||||
class CuteDSLKernelWrapper:
|
||||
"""Wrapper to provide .run() interface for CuteDSL kernels"""
|
||||
|
||||
def __init__(self, kernel_fn: Callable[..., Any], kernel_path: str | None = None):
|
||||
self.kernel_fn = kernel_fn
|
||||
self.kernel_path = kernel_path
|
||||
kernel_code_log.info("CuteDSL kernel path: %s", kernel_path)
|
||||
|
||||
def run(self, *args, stream=None, **kwargs):
|
||||
"""
|
||||
Execute the CuteDSL kernel.
|
||||
|
||||
Args:
|
||||
*args: Arguments to pass to the kernel function
|
||||
stream: CUDA stream to pass to the kernel function
|
||||
**kwargs: Additional keyword arguments for the kernel
|
||||
|
||||
Returns:
|
||||
Result of the kernel execution
|
||||
"""
|
||||
return self.kernel_fn(*args, stream=stream, **kwargs)
|
||||
|
||||
|
||||
@dataclasses.dataclass
|
||||
class CuteDSLSubgraphInfo:
|
||||
"""Minimal subgraph info for CuteDSL kernels."""
|
||||
|
||||
body: IndentedBuffer
|
||||
template_mask: str | None = None
|
||||
template_out: str | None = None
|
||||
cse: CSE[Any] | None = None
|
||||
|
||||
def __post_init__(self):
|
||||
self.only_copy_if_non_none_fields = ("cse",)
|
||||
|
||||
def to_dict(self):
|
||||
return {
|
||||
field.name: getattr(self, field.name) for field in dataclasses.fields(self)
|
||||
}
|
||||
|
||||
|
||||
class CuteDSLTemplateKernel(Kernel):
|
||||
"""
|
||||
Template kernel implementation for CuteDSL (CUTLASS Python DSL).
|
||||
Handles code generation and argument management for CuteDSL CUDA kernels.
|
||||
Provides CuteDSL-specific functionality for tensor conversion and kernel configuration.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
kernel_name: str,
|
||||
input_nodes: list[Buffer],
|
||||
output_node: Buffer,
|
||||
subgraphs: list[Buffer] | None = None,
|
||||
) -> None:
|
||||
# Call parent Kernel constructor
|
||||
super().__init__()
|
||||
self.kernel_name = kernel_name
|
||||
self.input_nodes = input_nodes
|
||||
self.output_node = output_node
|
||||
self.subgraphs = subgraphs
|
||||
self.subgraph_bodies: dict[str, CuteDSLSubgraphInfo] = {}
|
||||
|
||||
# Template attributes
|
||||
self.body: IndentedBuffer = IndentedBuffer()
|
||||
self.template_mask: str | None = None
|
||||
self.template_out: str | None = None
|
||||
self.template_indices: list[Any] | None = None
|
||||
self.render_hooks: dict[str, Any] = {}
|
||||
|
||||
# TODO Additional attributes needed by template system
|
||||
self.prologue_fused_inputs: OrderedSet[str] = OrderedSet()
|
||||
self.prologue_fused_inputs_preserve_zero: OrderedSet[str] = OrderedSet()
|
||||
self.named_input_nodes: dict[str, Buffer] = {}
|
||||
|
||||
# Create named input nodes mapping
|
||||
for i, input_node in enumerate(input_nodes):
|
||||
node_name = getattr(input_node, "name", f"input_{i}")
|
||||
self.named_input_nodes[node_name] = input_node
|
||||
|
||||
self.cse = CSE(name_prefix="tmp")
|
||||
|
||||
# Track all tensor buffers added during modification processing
|
||||
self.collected_tensor_buffers: list[str] = []
|
||||
|
||||
def kexpr(self, expr: sympy.Expr) -> str:
|
||||
"""Convert sympy expression to CuteDSL string representation."""
|
||||
return str(expr)
|
||||
|
||||
def gen_imports(self) -> str:
|
||||
"""Generate common imports for CuteDSL templates."""
|
||||
imports = IndentedBuffer()
|
||||
imports.splice(
|
||||
"""
|
||||
import torch
|
||||
import cutlass
|
||||
import cutlass.cute as cute
|
||||
from cutlass.cute.runtime import from_dlpack
|
||||
import cuda.bindings.driver as cuda
|
||||
from cutlass._mlir.dialects import math as mlir_math
|
||||
import operator
|
||||
from torch._inductor.codegen.cutedsl._cutedsl_utils import ssa_to_indexable, result_to_ssa
|
||||
"""
|
||||
)
|
||||
return imports.getvalue()
|
||||
|
||||
def gen_defines(self, **kwargs) -> str:
|
||||
"""Generate CuteDSL parameter definitions from kwargs, similar to Triton's gen_defines."""
|
||||
params = IndentedBuffer()
|
||||
for name, val in kwargs.items():
|
||||
params.writeline(f"{name}: cutlass.Constexpr = {val}")
|
||||
return params.getvalue()
|
||||
|
||||
def render(self, template, **kwargs):
|
||||
from torch._inductor.select_algorithm import PartialRender
|
||||
|
||||
"""Render the kernel using the template, returning PartialRender object with hooks."""
|
||||
# Available {{}} hooks for jinja rendering
|
||||
template_env = {
|
||||
"def_kernel": self.def_kernel,
|
||||
"gen_defines": lambda: self.gen_defines(**kwargs),
|
||||
"get_output": self.get_output,
|
||||
"get_tensor_buffers": self.get_tensor_buffers,
|
||||
"unpack_buffers": self.unpack_buffers,
|
||||
"modification": self.modification,
|
||||
"set_cute_hash": self.set_cute_hash,
|
||||
}
|
||||
|
||||
# Render the template with the environment and provided kwargs
|
||||
rendered_code = template.render(
|
||||
kernel_name=self.kernel_name,
|
||||
input_nodes=self.input_nodes,
|
||||
output_node=self.output_node,
|
||||
**template_env,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
# Always prepend the common imports
|
||||
imports = self.gen_imports()
|
||||
full_code = imports + rendered_code
|
||||
|
||||
return PartialRender(full_code, self.render_hooks)
|
||||
|
||||
@contextlib.contextmanager
|
||||
def set_subgraph_body(self, body_name: str):
|
||||
"""Set the active subgraph body for template processing."""
|
||||
assert all(
|
||||
hasattr(self, field.name)
|
||||
for field in dataclasses.fields(CuteDSLSubgraphInfo)
|
||||
)
|
||||
old_state = {
|
||||
key.name: getattr(self, key.name)
|
||||
for key in dataclasses.fields(CuteDSLSubgraphInfo)
|
||||
}
|
||||
|
||||
if body_name not in self.subgraph_bodies:
|
||||
self.subgraph_bodies[body_name] = CuteDSLSubgraphInfo(
|
||||
body=IndentedBuffer(),
|
||||
template_mask=None,
|
||||
template_out=None,
|
||||
cse=None,
|
||||
)
|
||||
|
||||
subgraph = self.subgraph_bodies[body_name]
|
||||
for key, value in subgraph.to_dict().items():
|
||||
if value is None and key in getattr(
|
||||
subgraph, "only_copy_if_non_none_fields", ()
|
||||
):
|
||||
continue
|
||||
setattr(self, key, value)
|
||||
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
# Save current state back to subgraph
|
||||
self.subgraph_bodies[body_name] = CuteDSLSubgraphInfo(
|
||||
**{
|
||||
key.name: getattr(self, key.name)
|
||||
for key in dataclasses.fields(CuteDSLSubgraphInfo)
|
||||
}
|
||||
)
|
||||
# Restore old state
|
||||
for key, value in old_state.items():
|
||||
setattr(self, key, value)
|
||||
|
||||
@contextlib.contextmanager
|
||||
def create_subgraph_body(self, body_name: str, *, clear_cse: bool = False):
|
||||
"""Create a new subgraph body for template processing."""
|
||||
assert body_name not in self.subgraph_bodies, (
|
||||
f"Subgraph body '{body_name}' already exists"
|
||||
)
|
||||
new_cse = self.cse.clone() if clear_cse else None
|
||||
self.subgraph_bodies[body_name] = CuteDSLSubgraphInfo(
|
||||
body=IndentedBuffer(),
|
||||
template_mask=None,
|
||||
template_out=None,
|
||||
cse=new_cse,
|
||||
)
|
||||
with self.set_subgraph_body(body_name):
|
||||
yield
|
||||
|
||||
def _get_reinterpret_view(self, node) -> ReinterpretView | None:
|
||||
"""Extract or convert to ReinterpretView from a node, handling all views."""
|
||||
while isinstance(node, MutableBox):
|
||||
node = node.data
|
||||
if isinstance(node, BaseView):
|
||||
return ExternKernel.convert_to_reinterpret_view(node)
|
||||
return None
|
||||
|
||||
def def_kernel(self, *argnames):
|
||||
"""Define kernel function signature for CuteDSL templates.
|
||||
|
||||
When inputs are ReinterpretViews of the same underlying buffer (e.g., Q/K/V
|
||||
from fused QKV projection), we generate separate arguments for each input
|
||||
even though they share the same underlying buffer.
|
||||
"""
|
||||
renames = IndentedBuffer(initial_indent=1)
|
||||
|
||||
# Track template input args - each input gets its own arg even if buffers are shared
|
||||
self._template_input_args: list[tuple[str, Buffer]] = []
|
||||
self._seen_input_args: OrderedSet[str] = OrderedSet()
|
||||
|
||||
for i, input_node in enumerate(self.input_nodes):
|
||||
buf_name = input_node.get_name()
|
||||
# Register with args system (may deduplicate, but we track separately)
|
||||
self.args.input(buf_name)
|
||||
|
||||
if i < len(argnames):
|
||||
template_name = argnames[i]
|
||||
arg_name = f"arg_{template_name}"
|
||||
self.args.input_buffers[buf_name] = arg_name
|
||||
renames.writeline(f"{template_name} = {arg_name}")
|
||||
self._template_input_args.append((arg_name, input_node))
|
||||
self._seen_input_args.add(arg_name)
|
||||
|
||||
if self.output_node:
|
||||
self.args.output(self.output_node.get_name())
|
||||
|
||||
def hook():
|
||||
# Generate signature with template input args plus additional args (output, sizevars)
|
||||
code = IndentedBuffer()
|
||||
code.writeline(f"# Kernel function signature: {self.kernel_name}")
|
||||
|
||||
# Start with template input args
|
||||
params = [arg_name for arg_name, _ in self._template_input_args]
|
||||
|
||||
# Get additional args from python_argdefs (output, sizevars, etc.)
|
||||
arg_defs, _, _, _ = self.args.python_argdefs()
|
||||
for arg_def in arg_defs:
|
||||
if arg_def.full_name() not in self._seen_input_args:
|
||||
params.append(arg_def.full_name())
|
||||
|
||||
params.append("stream")
|
||||
code.writeline(
|
||||
f"def {self.kernel_name}_{MAIN_SUFFIX}({', '.join(params)}):"
|
||||
)
|
||||
with code.indent():
|
||||
code.splice(renames.getvalue())
|
||||
return code.getvalue()
|
||||
|
||||
assert "<DEF_KERNEL>" not in self.render_hooks
|
||||
# Placeholder-based rendering: hook will be called when template encounters "<DEF_KERNEL>"
|
||||
self.render_hooks["<DEF_KERNEL>"] = hook
|
||||
return "<DEF_KERNEL>"
|
||||
|
||||
def get_output(self):
|
||||
"""Get the actual argument name for the output buffer."""
|
||||
assert self.output_node, "Output node must exist to get output buffer name"
|
||||
buf_name = self.output_node.get_name()
|
||||
output = self.args.output_buffers.get(buf_name, None)
|
||||
if output is None:
|
||||
raise ValueError(f"Output buffer '{buf_name}' not found in args")
|
||||
return output
|
||||
|
||||
def set_cute_hash(self, func_name: str, suffix: str = ""):
|
||||
"""Generate code to set __cute_hash__ on a codegen function.
|
||||
|
||||
This allows hash_callable in flash_attn to skip expensive runtime hashing
|
||||
for Inductor-generated functions. The hash is based on the kernel name
|
||||
which already contains a unique hash suffix.
|
||||
"""
|
||||
hash_value = f"{self.kernel_name}_{suffix}" if suffix else self.kernel_name
|
||||
return f'{func_name}.__cute_hash__ = "{hash_value}"'
|
||||
|
||||
def get_tensor_buffers(self):
|
||||
"""Get list of tensor buffer names that were collected during modifications."""
|
||||
return self.collected_tensor_buffers
|
||||
|
||||
def unpack_buffers(self, buffer_list_name: str, *, indent_width: int = 4):
|
||||
"""Generate buffer unpacking code via render hook."""
|
||||
|
||||
def hook():
|
||||
tensor_buffers = self.get_tensor_buffers()
|
||||
if not tensor_buffers:
|
||||
return ""
|
||||
|
||||
# Generate unpacking assignments: in_ptr4 = buffers[0], etc.
|
||||
unpacking_lines = []
|
||||
for i, buffer_name in enumerate(tensor_buffers):
|
||||
# pyrefly: ignore [bad-argument-type]
|
||||
unpacking_lines.append(f"{buffer_name} = {buffer_list_name}[{i}]")
|
||||
|
||||
indent = " " * indent_width
|
||||
return "\n" + indent + ("\n" + indent).join(unpacking_lines)
|
||||
|
||||
# Register the hook and return placeholder
|
||||
placeholder = "<UNPACK_BUFFERS>"
|
||||
# TODO: I think double invoking is fine for this specific hook
|
||||
# assert placeholder not in self.render_hooks
|
||||
self.render_hooks[placeholder] = hook
|
||||
return placeholder
|
||||
|
||||
def call_kernel(self, name: str, node=None):
|
||||
"""Call the kernel function. Simplified version of TritonTemplateKernel.call_kernel.
|
||||
|
||||
For inputs that are ReinterpretViews (e.g., Q/K/V slices from fused QKV),
|
||||
we generate reinterpret_tensor() calls to properly handle the views.
|
||||
"""
|
||||
wrapper = V.graph.wrapper_code
|
||||
|
||||
# Build call args matching the signature generated in `def_kernel`
|
||||
call_args = []
|
||||
arg_types = []
|
||||
|
||||
for _, input_node in self._template_input_args:
|
||||
reinterpret_view = self._get_reinterpret_view(input_node)
|
||||
if reinterpret_view is not None:
|
||||
call_args.append(reinterpret_view.codegen_reference())
|
||||
else:
|
||||
call_args.append(input_node.get_name())
|
||||
arg_types.append(V.graph.get_dtype(input_node.get_name()))
|
||||
|
||||
# Add additional args from python_argdefs (output, sizevars, ..)
|
||||
orig_arg_defs, orig_call_args, _, orig_arg_types = self.args.python_argdefs()
|
||||
for arg_def, call_arg, arg_type in zip(
|
||||
orig_arg_defs, orig_call_args, orig_arg_types
|
||||
):
|
||||
# dedupe
|
||||
if arg_def.full_name() not in self._seen_input_args:
|
||||
call_args.append(call_arg)
|
||||
arg_types.append(arg_type)
|
||||
|
||||
# TODO this karg really should not be called `triton`
|
||||
wrapper.generate_kernel_call(name, call_args, triton=True, arg_types=arg_types)
|
||||
|
||||
def _get_subgraph(self, subgraph_number: int):
|
||||
"""Get subgraph by number for modification processing."""
|
||||
assert isinstance(subgraph_number, int)
|
||||
assert isinstance(self.subgraphs, list)
|
||||
assert subgraph_number < len(self.subgraphs), (
|
||||
f"Invalid subgraph number provided to create_modification, {subgraph_number} must be < {len(self.subgraphs)}"
|
||||
)
|
||||
assert self.body.getvalue() == "", (
|
||||
"Body should be clear before adding a modification"
|
||||
)
|
||||
return self.subgraphs[subgraph_number]
|
||||
|
||||
def modification(
|
||||
self,
|
||||
subgraph_number: int,
|
||||
output_name: str | None,
|
||||
mask: str | None = None,
|
||||
**fixed_inputs,
|
||||
) -> str:
|
||||
"""Generate CuteDSL code for a subgraph modification."""
|
||||
# Find unique name to avoid collisions between multiple modifications of same subgraph
|
||||
num = 0
|
||||
while f"mod_{subgraph_number}_{num}" in self.subgraph_bodies:
|
||||
num += 1
|
||||
|
||||
with self.create_subgraph_body(f"mod_{subgraph_number}_{num}", clear_cse=True):
|
||||
subgraph = self._get_subgraph(subgraph_number)
|
||||
modification_handler = ModificationWrapperCuteDSL(
|
||||
self, subgraph_number, fixed_inputs, mask
|
||||
)
|
||||
with V.set_kernel_handler(self), V.set_ops_handler(modification_handler):
|
||||
assert isinstance(subgraph, (ComputedBuffer, list)), (
|
||||
f"Expected ComputedBuffer or List[ComputedBuffer], got {type(subgraph)}"
|
||||
)
|
||||
|
||||
if isinstance(subgraph, list):
|
||||
raise NotImplementedError(
|
||||
"Scatter graphs are not supported for CuteDSL"
|
||||
)
|
||||
|
||||
if isinstance(subgraph.data, InputBuffer):
|
||||
# grad_score_mod can be InputBuffers
|
||||
out = subgraph.data.make_loader()(())
|
||||
else:
|
||||
# Inline a pointwise lowering into the template
|
||||
out = subgraph.data.inner_fn(())
|
||||
|
||||
if output_name is not None:
|
||||
assert out is not None, (
|
||||
f"Expected computation result for named output {output_name}"
|
||||
)
|
||||
self.body.writeline(f"{output_name} = {out.value}")
|
||||
else:
|
||||
# Side-effect only: no output assignment (currently only for scatter operations)
|
||||
raise NotImplementedError(
|
||||
"Side-effect only modifications not yet supported for CuteDSL"
|
||||
)
|
||||
|
||||
# Add Buffers that were added during modification
|
||||
self.collected_tensor_buffers.extend(modification_handler.tensor_buffers)
|
||||
|
||||
return self.body.getvalue()
|
||||
|
||||
|
||||
class ModificationWrapperCuteDSL(V.WrapperHandler): # type: ignore[name-defined]
|
||||
"""
|
||||
Wrapper handler that enables CuteDSL code generation during subgraph modifications.
|
||||
|
||||
This class sits between the PyTorch IR and CuteDSL code generation, providing:
|
||||
1. Operation substitution: converts PyTorch ops to CuteDSL equivalents via CuteDSLOpOverrides
|
||||
2. Placeholder handling: resolves fixed_inputs during template processing
|
||||
3. Limited operation support: currently restricted to pointwise operations
|
||||
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
kernel,
|
||||
subgraph_number: int,
|
||||
fixed_inputs: dict[str, Any],
|
||||
mask: str | None,
|
||||
):
|
||||
cutedsl_ops = CuteDSLOpOverrides()
|
||||
super().__init__(cutedsl_ops)
|
||||
self.name = f"CuteDSLPlaceholderSubstitution_{subgraph_number}"
|
||||
self.kernel = kernel
|
||||
self.fixed_inputs = fixed_inputs
|
||||
self.mask = mask
|
||||
# Track tensor buffers that get added during modification processing
|
||||
self.tensor_buffers: list[str] = []
|
||||
|
||||
def _get_input_dtype(self, name: str) -> torch.dtype:
|
||||
"""Get the dtype for an input from the kernel's named_input_nodes."""
|
||||
if name in self.kernel.named_input_nodes:
|
||||
return self.kernel.named_input_nodes[name].dtype
|
||||
# TODO: Fallback for common dimension names - should be replaced with proper dtype tracking
|
||||
return torch.float32 if name not in ("b", "h", "m", "n") else torch.int32
|
||||
|
||||
def load(self, name: str, index: sympy.Expr):
|
||||
"""Handle loading from tensor or fixed(template args) input for CuteDSL."""
|
||||
from torch._inductor.kernel.flex.flex_flash_attention import HierarchicalIndex
|
||||
|
||||
if name not in self.fixed_inputs:
|
||||
var = self._add_kernel_input(name)
|
||||
buffer = V.graph.get_buffer(name)
|
||||
var_dtype = buffer.dtype
|
||||
|
||||
cute_dtype = CuteDSLOpOverrides.TORCH_TO_CUTE_DTYPE.get(
|
||||
var_dtype, "cutlass.Float32"
|
||||
)
|
||||
idx_vars = [
|
||||
self._emit_scalar_fragment(
|
||||
self.kernel.kexpr(self.kernel.rename_indexing(dim_index)),
|
||||
"cutlass.Int32",
|
||||
torch.int32,
|
||||
)
|
||||
for dim_index in (
|
||||
index.args if isinstance(index, HierarchicalIndex) else (index,)
|
||||
)
|
||||
]
|
||||
|
||||
val_frag = self.kernel.cse.newvar(dtype=var_dtype)
|
||||
self.kernel.body.writeline(
|
||||
f"{val_frag} = cute.make_rmem_tensor(1, {cute_dtype})"
|
||||
)
|
||||
self.kernel.body.writeline(
|
||||
f"{val_frag}[0] = ({var}[{', '.join(idx_vars)}])"
|
||||
)
|
||||
|
||||
final_expr = f"{val_frag}.load()"
|
||||
|
||||
if (
|
||||
var_dtype in (torch.float16, torch.bfloat16)
|
||||
and config.triton.codegen_upcast_to_fp32
|
||||
):
|
||||
final_expr = f"({final_expr}).to(cutlass.Float32)"
|
||||
var_dtype = torch.float32
|
||||
|
||||
out = self.kernel.cse.generate(
|
||||
self.kernel.body,
|
||||
final_expr,
|
||||
dtype=var_dtype,
|
||||
bounds=ValueRanges.unknown(),
|
||||
)
|
||||
return out
|
||||
|
||||
value = self.fixed_inputs[name]
|
||||
dtype = self._get_input_dtype(name)
|
||||
|
||||
return self.kernel.cse.generate(
|
||||
self.kernel.body, value, bounds=ValueRanges.unknown(), dtype=dtype
|
||||
)
|
||||
|
||||
def _emit_scalar_fragment(
|
||||
self, expr_str: str, cute_dtype: str, torch_dtype: torch.dtype
|
||||
) -> str:
|
||||
"""
|
||||
Convert expression to indexable scalar for tensor loads.
|
||||
|
||||
Workaround for lack of gather support: SSA values cannot be used directly
|
||||
as indices in tensor loads. This generates code to convert SSA → indexable
|
||||
scalar. Compile-time integer constants are already indexable and are
|
||||
returned directly without the SSA round-trip.
|
||||
"""
|
||||
# Constant integer expressions (e.g. sympy-folded offsets like "0")
|
||||
# are already valid indices — skip the ssa_to_indexable round-trip
|
||||
# which only accepts TensorSSA, not bare Python ints.
|
||||
if expr_str.lstrip("-").isdigit():
|
||||
return expr_str
|
||||
|
||||
result = self.kernel.cse.newvar(dtype=torch_dtype)
|
||||
self.kernel.body.writeline(
|
||||
f"{result} = ssa_to_indexable({expr_str}, {cute_dtype})"
|
||||
)
|
||||
return str(result)
|
||||
|
||||
def indirect_indexing(self, index_var: str, size, check, wrap_neg=True):
|
||||
"""Convert index variable to symbolic form."""
|
||||
return sympy_index_symbol(str(index_var))
|
||||
|
||||
# pyrefly: ignore [bad-override]
|
||||
def store(
|
||||
self, name: str, index: sympy.Expr, value: CSEVariable, mode: StoreMode = None
|
||||
) -> str:
|
||||
raise NotImplementedError(
|
||||
"Store operations not supported - CuteDSL limited to read-only operations"
|
||||
)
|
||||
|
||||
def _add_kernel_input(self, name: str):
|
||||
"""Add name as input to kernel and return input ref."""
|
||||
# Get the remapped name that will be used in the kernel
|
||||
remapped_name = self.kernel.args.input(name)
|
||||
# Track the remapped name for later collection
|
||||
if remapped_name not in self.tensor_buffers:
|
||||
self.tensor_buffers.append(remapped_name)
|
||||
return remapped_name
|
||||
|
||||
def _process_indexing(self, index):
|
||||
"""Process and rename indexing, adding symbols as kernel inputs."""
|
||||
renamed = self.kernel.rename_indexing(index)
|
||||
return self.kernel.kexpr(renamed)
|
||||
|
||||
def _default(self, name: str, args: tuple[Any, ...], kwargs: dict[str, Any]) -> Any:
|
||||
try:
|
||||
return getattr(self._inner, name)(*args, **kwargs)
|
||||
except NotImplementedError as e:
|
||||
bar = "=" * 80
|
||||
msg = textwrap.dedent(f"""
|
||||
{bar}
|
||||
UNSUPPORTED CUTEDSL OPERATION: '{name}'
|
||||
{bar}
|
||||
This operation is not yet implemented in Inductor.
|
||||
|
||||
Please open an issue at: https://github.com/pytorch/pytorch/issues
|
||||
with the following information:
|
||||
|
||||
Operation: {name}
|
||||
Args: {args!r}
|
||||
Kwargs: {kwargs!r}
|
||||
|
||||
Title your issue: [CuteDSL] Missing operation: {name}
|
||||
{bar}
|
||||
""").strip()
|
||||
raise NotImplementedError(msg) from e
|
||||
+528
@@ -0,0 +1,528 @@
|
||||
# mypy: allow-untyped-defs
|
||||
"""
|
||||
CuteDSL-specific operation overrides for pointwise operations.
|
||||
|
||||
This module provides CuteDSL implementations of common operations used in
|
||||
template kernels, particularly for flex attention modifications.
|
||||
"""
|
||||
|
||||
import math
|
||||
|
||||
import sympy
|
||||
|
||||
import torch
|
||||
from torch._inductor.codegen.common import CSEVariable, OpOverrides
|
||||
from torch._inductor.virtualized import OpsValue, V
|
||||
from torch.utils._sympy.value_ranges import ValueRanges
|
||||
|
||||
|
||||
CuteDSLArg = CSEVariable | str | bool | float | int
|
||||
|
||||
|
||||
def upcast_compute_type(dtype: torch.dtype) -> torch.dtype:
|
||||
"""Maybe upcast [b]float16 to float32"""
|
||||
if dtype in (torch.float16, torch.bfloat16):
|
||||
return torch.float32
|
||||
return dtype
|
||||
|
||||
|
||||
class CuteDSLOpOverrides(OpOverrides):
|
||||
"""
|
||||
CuteDSL-specific operation overrides that generate code using CuteDSL syntax.
|
||||
|
||||
CuteDSL TensorSSA objects have built-in operator overloads (__add__, __mul__, etc.)
|
||||
and math functions (cute.math.exp, cute.math.sqrt, etc.)
|
||||
"""
|
||||
|
||||
TORCH_TO_CUTE_DTYPE = {
|
||||
torch.float16: "cutlass.Float16",
|
||||
torch.bfloat16: "cutlass.BFloat16",
|
||||
torch.float32: "cutlass.Float32",
|
||||
torch.float64: "cutlass.Float64",
|
||||
torch.int8: "cutlass.Int8",
|
||||
torch.int16: "cutlass.Int16",
|
||||
torch.int32: "cutlass.Int32",
|
||||
torch.int64: "cutlass.Int64",
|
||||
torch.uint8: "cutlass.Uint8",
|
||||
torch.bool: "cutlass.Boolean",
|
||||
torch.float8_e4m3fn: "cutlass.Float8E4M3FN",
|
||||
torch.float8_e5m2: "cutlass.Float8E5M2",
|
||||
}
|
||||
|
||||
# Math constants
|
||||
LOG2_E = 1.4426950408889634 # 1/ln(2) for converting natural exp to base-2 exp
|
||||
|
||||
@staticmethod
|
||||
def _get_cse_var(arg: CuteDSLArg) -> CSEVariable | None:
|
||||
"""Extract CSEVariable from arg if it's a tensor (either direct or wrapped in OpsValue)."""
|
||||
if isinstance(arg, CSEVariable):
|
||||
return arg
|
||||
if isinstance(arg, OpsValue) and isinstance(arg.value, CSEVariable):
|
||||
return arg.value
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
def _as_expr(arg: CuteDSLArg) -> str:
|
||||
cse_var = CuteDSLOpOverrides._get_cse_var(arg)
|
||||
if cse_var is not None:
|
||||
return str(cse_var)
|
||||
return str(arg)
|
||||
|
||||
@staticmethod
|
||||
def _node_tensor_flags() -> tuple[bool, bool] | None:
|
||||
node = V.current_node
|
||||
if not isinstance(node, torch.fx.Node) or len(node.args) < 2:
|
||||
return None
|
||||
|
||||
def _is_tensor(raw_arg: object) -> bool:
|
||||
if hasattr(raw_arg, "meta"):
|
||||
val = raw_arg.meta.get("val")
|
||||
return isinstance(val, torch.Tensor) and val.ndim > 0
|
||||
return False
|
||||
|
||||
return _is_tensor(node.args[0]), _is_tensor(node.args[1])
|
||||
|
||||
@staticmethod
|
||||
def _ensure_tensor_ssa(
|
||||
arg: CuteDSLArg, template_tensor: CuteDSLArg, *, is_tensor: bool
|
||||
) -> str:
|
||||
"""
|
||||
Convert scalar arguments to TensorSSA using cute.full_like if needed.
|
||||
|
||||
Args:
|
||||
arg: The argument to check (CSEVariable for tensors, str for scalars, or OpsValue wrapper)
|
||||
template_tensor: A tensor argument to use as template for full_like
|
||||
|
||||
Returns:
|
||||
String representation suitable for CuteDSL operations
|
||||
"""
|
||||
arg_expr = CuteDSLOpOverrides._as_expr(arg)
|
||||
if is_tensor:
|
||||
return arg_expr
|
||||
template_expr = CuteDSLOpOverrides._as_expr(template_tensor)
|
||||
return f"cute.full_like({template_expr}, {arg_expr})"
|
||||
|
||||
@staticmethod
|
||||
def _extract_dtype_and_bounds(
|
||||
*args: CuteDSLArg,
|
||||
) -> tuple[torch.dtype | None, ValueRanges[sympy.Expr]]:
|
||||
"""Extract dtype and bounds from CSEVariable arguments (including OpsValue wrappers)."""
|
||||
for arg in args:
|
||||
cse_var = CuteDSLOpOverrides._get_cse_var(arg)
|
||||
if cse_var is not None:
|
||||
return cse_var.dtype, cse_var.bounds
|
||||
return None, ValueRanges.unknown()
|
||||
|
||||
@staticmethod
|
||||
def _apply_binary_op(a: CuteDSLArg, b: CuteDSLArg, op_format: str) -> CuteDSLArg:
|
||||
"""
|
||||
Apply a binary operation with automatic scalar-to-tensor conversion.
|
||||
|
||||
CuteDSL requires both operands to be TensorSSA objects for tensor operations.
|
||||
This helper automatically converts scalar arguments to TensorSSA using
|
||||
cute.full_like when at least one argument is a tensor (CSEVariable or OpsValue).
|
||||
|
||||
Args:
|
||||
a: First operand (CSEVariable for tensors, str for scalars, or OpsValue wrapper)
|
||||
b: Second operand (CSEVariable for tensors, str for scalars, or OpsValue wrapper)
|
||||
op_format: Format string with {a} and {b} placeholders for the operation
|
||||
|
||||
Returns:
|
||||
CSEVariable if at least one operand is a tensor, otherwise string
|
||||
"""
|
||||
a_cse = CuteDSLOpOverrides._get_cse_var(a)
|
||||
b_cse = CuteDSLOpOverrides._get_cse_var(b)
|
||||
|
||||
node_flags = CuteDSLOpOverrides._node_tensor_flags()
|
||||
if node_flags is not None:
|
||||
a_is_tensor, b_is_tensor = node_flags
|
||||
else:
|
||||
a_is_tensor = a_cse is not None
|
||||
b_is_tensor = b_cse is not None
|
||||
|
||||
tensor_arg = a if a_is_tensor else (b if b_is_tensor else None)
|
||||
if tensor_arg is None:
|
||||
tensor_arg = a_cse or b_cse
|
||||
|
||||
if tensor_arg is not None:
|
||||
if a_cse is None and b_cse is None:
|
||||
return op_format.format(
|
||||
a=CuteDSLOpOverrides._as_expr(a),
|
||||
b=CuteDSLOpOverrides._as_expr(b),
|
||||
)
|
||||
|
||||
a_ssa = CuteDSLOpOverrides._ensure_tensor_ssa(
|
||||
a, tensor_arg, is_tensor=a_is_tensor
|
||||
)
|
||||
b_ssa = CuteDSLOpOverrides._ensure_tensor_ssa(
|
||||
b, tensor_arg, is_tensor=b_is_tensor
|
||||
)
|
||||
result_expr = op_format.format(a=a_ssa, b=b_ssa)
|
||||
|
||||
dtype, bounds = CuteDSLOpOverrides._extract_dtype_and_bounds(a, b)
|
||||
expected = CuteDSLOpOverrides._expected_tensor_val()
|
||||
if dtype is None:
|
||||
dtype = expected.dtype if expected is not None else torch.int32
|
||||
if a_cse is not None:
|
||||
shape = a_cse.shape
|
||||
elif b_cse is not None:
|
||||
shape = b_cse.shape
|
||||
else:
|
||||
shape = tuple(expected.size()) if expected is not None else None
|
||||
|
||||
# Create and return CSEVariable using CSE generation for caching
|
||||
return V.kernel.cse.generate(
|
||||
V.kernel.body, result_expr, bounds=bounds, dtype=dtype, shape=shape
|
||||
)
|
||||
|
||||
return op_format.format(a=a, b=b)
|
||||
|
||||
@staticmethod
|
||||
def _expected_tensor_val() -> torch.Tensor | None:
|
||||
"""Return the fake-tensor value from the current FX node's metadata, if any."""
|
||||
node = V.current_node
|
||||
if not isinstance(node, torch.fx.Node):
|
||||
return None
|
||||
val = node.meta.get("val")
|
||||
return val if isinstance(val, torch.Tensor) else None
|
||||
|
||||
@staticmethod
|
||||
def _cast_expr(expr: str, dtype: torch.dtype) -> str:
|
||||
cute_type = CuteDSLOpOverrides.TORCH_TO_CUTE_DTYPE.get(dtype)
|
||||
if cute_type is None:
|
||||
return expr
|
||||
return f"{cute_type}({expr})"
|
||||
|
||||
@staticmethod
|
||||
def _apply_unary_op(x: CuteDSLArg, op_format: str) -> CuteDSLArg:
|
||||
"""
|
||||
Apply a unary operation, returning CSEVariable if input is a tensor.
|
||||
|
||||
Args:
|
||||
x: Input operand (CSEVariable for tensors, str for scalars, or OpsValue wrapper)
|
||||
op_format: Format string with {x} placeholder for the operation
|
||||
|
||||
Returns:
|
||||
CSEVariable if input is a tensor, otherwise string
|
||||
"""
|
||||
cse_var = CuteDSLOpOverrides._get_cse_var(x)
|
||||
if cse_var is not None:
|
||||
result_expr = op_format.format(x=str(cse_var))
|
||||
return V.kernel.cse.generate(
|
||||
V.kernel.body, result_expr, bounds=cse_var.bounds, dtype=cse_var.dtype
|
||||
)
|
||||
|
||||
return op_format.format(x=x)
|
||||
|
||||
@staticmethod
|
||||
def constant(value: bool | float | int, dtype: torch.dtype) -> str:
|
||||
"""Generate CuteDSL constant representation."""
|
||||
if value == float("-inf"):
|
||||
return "float('-inf')"
|
||||
elif value == float("inf"):
|
||||
return "float('inf')"
|
||||
elif math.isnan(value):
|
||||
return "float('nan')"
|
||||
return repr(value)
|
||||
|
||||
@staticmethod
|
||||
def add(a: CuteDSLArg, b: CuteDSLArg) -> CuteDSLArg:
|
||||
return CuteDSLOpOverrides._apply_binary_op(a, b, "({a} + {b})")
|
||||
|
||||
@staticmethod
|
||||
def mul(a: CuteDSLArg, b: CuteDSLArg) -> CuteDSLArg:
|
||||
return CuteDSLOpOverrides._apply_binary_op(a, b, "({a} * {b})")
|
||||
|
||||
@staticmethod
|
||||
def sub(a: CuteDSLArg, b: CuteDSLArg) -> CuteDSLArg:
|
||||
return CuteDSLOpOverrides._apply_binary_op(a, b, "({a} - {b})")
|
||||
|
||||
@staticmethod
|
||||
def truediv(a: CuteDSLArg, b: CuteDSLArg) -> CuteDSLArg:
|
||||
return CuteDSLOpOverrides._apply_binary_op(a, b, "({a} / {b})")
|
||||
|
||||
@staticmethod
|
||||
def floordiv(a: CuteDSLArg, b: CuteDSLArg) -> CuteDSLArg:
|
||||
return CuteDSLOpOverrides._apply_binary_op(a, b, "({a} // {b})")
|
||||
|
||||
@staticmethod
|
||||
def mod(a: CuteDSLArg, b: CuteDSLArg) -> CuteDSLArg:
|
||||
return CuteDSLOpOverrides._apply_binary_op(a, b, "({a} % {b})")
|
||||
|
||||
@staticmethod
|
||||
def remainder(a, b):
|
||||
return CuteDSLOpOverrides._apply_binary_op(a, b, "({a} % {b})")
|
||||
|
||||
@staticmethod
|
||||
# pyrefly: ignore [bad-override]
|
||||
def exp(x: CuteDSLArg) -> CuteDSLArg:
|
||||
"""Exponential using CuteDSL cute.math.exp2 with log2(e) scaling."""
|
||||
if CuteDSLOpOverrides._get_cse_var(x) is None:
|
||||
x = CuteDSLOpOverrides._cast_expr(str(x), torch.float32)
|
||||
return CuteDSLOpOverrides._apply_unary_op(
|
||||
x, f"cute.math.exp2({{x}} * {CuteDSLOpOverrides.LOG2_E})"
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
# pyrefly: ignore [bad-override]
|
||||
def sqrt(x: CuteDSLArg) -> CuteDSLArg:
|
||||
"""Square root using CuteDSL cute.math.sqrt function."""
|
||||
return CuteDSLOpOverrides._apply_unary_op(x, "cute.math.sqrt({x})")
|
||||
|
||||
@staticmethod
|
||||
# pyrefly: ignore [bad-override]
|
||||
def log(x: CuteDSLArg) -> CuteDSLArg:
|
||||
"""Natural logarithm using CuteDSL cute.math.log function."""
|
||||
return CuteDSLOpOverrides._apply_unary_op(x, "cute.math.log({x})")
|
||||
|
||||
@staticmethod
|
||||
# pyrefly: ignore [bad-override]
|
||||
def cos(x: CuteDSLArg) -> CuteDSLArg:
|
||||
"""Cosine using CuteDSL cute.math.cos function."""
|
||||
return CuteDSLOpOverrides._apply_unary_op(x, "cute.math.cos({x})")
|
||||
|
||||
@staticmethod
|
||||
# pyrefly: ignore [bad-override]
|
||||
def sin(x: CuteDSLArg) -> CuteDSLArg:
|
||||
"""Sine using CuteDSL cute.math.sin function."""
|
||||
return CuteDSLOpOverrides._apply_unary_op(x, "cute.math.sin({x})")
|
||||
|
||||
@staticmethod
|
||||
# pyrefly: ignore [bad-override]
|
||||
def erf(x: CuteDSLArg) -> CuteDSLArg:
|
||||
"""Error function using CuteDSL cute.math.erf function."""
|
||||
return CuteDSLOpOverrides._apply_unary_op(x, "cute.math.erf({x})")
|
||||
|
||||
@staticmethod
|
||||
# pyrefly: ignore [bad-override]
|
||||
def sigmoid(x: CuteDSLArg) -> CuteDSLArg:
|
||||
"""Sigmoid with fp32 compute and cast-back to expected output dtype."""
|
||||
x_cse = CuteDSLOpOverrides._get_cse_var(x)
|
||||
if x_cse is not None:
|
||||
x_fp32: CuteDSLArg = CuteDSLOpOverrides.to_dtype(
|
||||
x_cse,
|
||||
torch.float32,
|
||||
use_compute_types=False,
|
||||
)
|
||||
else:
|
||||
x_fp32 = CuteDSLOpOverrides._cast_expr(str(x), torch.float32)
|
||||
|
||||
result = CuteDSLOpOverrides._apply_unary_op(
|
||||
x_fp32,
|
||||
f"(1.0 / (1.0 + cute.math.exp2(-{{x}} * {CuteDSLOpOverrides.LOG2_E})))",
|
||||
)
|
||||
|
||||
expected = CuteDSLOpOverrides._expected_tensor_val()
|
||||
expected_dtype = expected.dtype if expected is not None else None
|
||||
if expected_dtype is not None and expected_dtype != torch.float32:
|
||||
result_cse = CuteDSLOpOverrides._get_cse_var(result)
|
||||
if result_cse is not None:
|
||||
return CuteDSLOpOverrides.to_dtype(
|
||||
result_cse,
|
||||
expected_dtype,
|
||||
use_compute_types=False,
|
||||
)
|
||||
return CuteDSLOpOverrides._cast_expr(str(result), expected_dtype)
|
||||
return result
|
||||
|
||||
@staticmethod
|
||||
def _minmax(a: CuteDSLArg, b: CuteDSLArg, *, op: str) -> CuteDSLArg:
|
||||
tensor_arg = CuteDSLOpOverrides._get_cse_var(
|
||||
a
|
||||
) or CuteDSLOpOverrides._get_cse_var(b)
|
||||
if tensor_arg is not None:
|
||||
return CuteDSLOpOverrides._apply_binary_op(
|
||||
a, b, f"cute.where(({{a}}) {op} ({{b}}), {{a}}, {{b}})"
|
||||
)
|
||||
|
||||
lhs = str(a)
|
||||
rhs = str(b)
|
||||
expected = CuteDSLOpOverrides._expected_tensor_val()
|
||||
if expected is not None:
|
||||
lhs = CuteDSLOpOverrides._cast_expr(lhs, expected.dtype)
|
||||
rhs = CuteDSLOpOverrides._cast_expr(rhs, expected.dtype)
|
||||
return f"({lhs} if {lhs} {op} {rhs} else {rhs})"
|
||||
|
||||
@staticmethod
|
||||
# pyrefly: ignore [bad-override]
|
||||
def maximum(a: CuteDSLArg, b: CuteDSLArg) -> CuteDSLArg:
|
||||
return CuteDSLOpOverrides._minmax(a, b, op=">")
|
||||
|
||||
@staticmethod
|
||||
# pyrefly: ignore [bad-override]
|
||||
def minimum(a: CuteDSLArg, b: CuteDSLArg) -> CuteDSLArg:
|
||||
return CuteDSLOpOverrides._minmax(a, b, op="<")
|
||||
|
||||
@staticmethod
|
||||
# pyrefly: ignore [bad-override]
|
||||
def where(
|
||||
condition: CuteDSLArg,
|
||||
a: CuteDSLArg,
|
||||
b: CuteDSLArg,
|
||||
) -> CuteDSLArg:
|
||||
"""Conditional selection - handles CSEVariable, OpsValue, and string inputs."""
|
||||
a_cse = CuteDSLOpOverrides._get_cse_var(a)
|
||||
b_cse = CuteDSLOpOverrides._get_cse_var(b)
|
||||
cond_cse = CuteDSLOpOverrides._get_cse_var(condition)
|
||||
tensor_arg = a_cse or b_cse or cond_cse
|
||||
|
||||
if tensor_arg is not None:
|
||||
a_ssa = CuteDSLOpOverrides._ensure_tensor_ssa(
|
||||
a, tensor_arg, is_tensor=a_cse is not None
|
||||
)
|
||||
b_ssa = CuteDSLOpOverrides._ensure_tensor_ssa(
|
||||
b, tensor_arg, is_tensor=b_cse is not None
|
||||
)
|
||||
cond_ssa = CuteDSLOpOverrides._ensure_tensor_ssa(
|
||||
condition, tensor_arg, is_tensor=cond_cse is not None
|
||||
)
|
||||
result_expr = f"cute.where({cond_ssa}, {a_ssa}, {b_ssa})"
|
||||
|
||||
dtype, bounds = CuteDSLOpOverrides._extract_dtype_and_bounds(
|
||||
a, b, condition
|
||||
)
|
||||
|
||||
return V.kernel.cse.generate(
|
||||
V.kernel.body, result_expr, bounds=bounds, dtype=dtype
|
||||
)
|
||||
|
||||
return f"cute.where({condition}, {a}, {b})"
|
||||
|
||||
@staticmethod
|
||||
def pow(a: CuteDSLArg, b: CuteDSLArg):
|
||||
return CuteDSLOpOverrides._apply_binary_op(a, b, "({a} ** {b})")
|
||||
|
||||
@staticmethod
|
||||
# pyrefly: ignore [bad-override]
|
||||
def abs(x: CuteDSLArg) -> CuteDSLArg:
|
||||
"""Absolute value using CuteDSL cute.math.abs function."""
|
||||
if isinstance(x, CSEVariable):
|
||||
x_dtype = x.dtype
|
||||
elif isinstance(x, OpsValue) and isinstance(x.value, CSEVariable):
|
||||
x_dtype = x.value.dtype
|
||||
else:
|
||||
x_dtype = torch.float32
|
||||
|
||||
abs_op = (
|
||||
"mlir_math.absf"
|
||||
if x_dtype in (torch.float16, torch.bfloat16, torch.float32)
|
||||
else "mlir_math.absi"
|
||||
)
|
||||
return CuteDSLOpOverrides._apply_unary_op(
|
||||
x,
|
||||
f"cute.TensorSSA({abs_op}({{x}}), {{x}}.shape, {{x}}.dtype)",
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def neg(x: CuteDSLArg) -> CuteDSLArg:
|
||||
"""Negation for both TensorSSA and scalar-like expressions."""
|
||||
# TensorSSA path: avoid relying on __neg__ directly due upstream issue.
|
||||
if CuteDSLOpOverrides._get_cse_var(x) is not None:
|
||||
return CuteDSLOpOverrides._apply_unary_op(
|
||||
x, "cute.TensorSSA(-{x}, {x}.shape, {x}.dtype)"
|
||||
)
|
||||
# Scalar path: shape/dtype attributes are unavailable.
|
||||
return CuteDSLOpOverrides._apply_unary_op(x, "(-{x})")
|
||||
|
||||
@staticmethod
|
||||
def to_dtype(
|
||||
x: CuteDSLArg, dtype: torch.dtype, src_dtype=None, use_compute_types=True
|
||||
) -> CuteDSLArg:
|
||||
"""Type conversion using CuteDSL TensorSSA.to(Type[Numeric]).
|
||||
|
||||
Maps torch dtypes to cutlass.cute.typing numeric types and emits
|
||||
`{x}.to(cute.typing.<Type>)`.
|
||||
|
||||
Raises NotImplementedError for unsigned integer and unsupported dtypes.
|
||||
"""
|
||||
if use_compute_types:
|
||||
dtype = upcast_compute_type(dtype)
|
||||
|
||||
cute_type = CuteDSLOpOverrides.TORCH_TO_CUTE_DTYPE.get(dtype)
|
||||
if cute_type is None:
|
||||
raise NotImplementedError(
|
||||
f"CuteDSL dtype cast not implemented for torch dtype: {dtype}"
|
||||
)
|
||||
|
||||
if isinstance(x, CSEVariable):
|
||||
result_expr = f"{str(x)}.to({cute_type})"
|
||||
return V.kernel.cse.generate(
|
||||
V.kernel.body, result_expr, bounds=x.bounds, dtype=dtype
|
||||
)
|
||||
|
||||
return f"{x}.to({cute_type})"
|
||||
|
||||
@staticmethod
|
||||
def tanh(x0: CuteDSLArg) -> CuteDSLArg:
|
||||
"""Hyperbolic tangent using CuteDSL cute.math.tanh function."""
|
||||
return CuteDSLOpOverrides._apply_unary_op(x0, "cute.math.tanh({x})")
|
||||
|
||||
# Logical operations
|
||||
@staticmethod
|
||||
def logical_and(x0: CuteDSLArg, x1: CuteDSLArg) -> CuteDSLArg:
|
||||
return CuteDSLOpOverrides._apply_binary_op(x0, x1, "({a} & {b})")
|
||||
|
||||
@staticmethod
|
||||
def logical_or(x0: CuteDSLArg, x1: CuteDSLArg) -> CuteDSLArg:
|
||||
return CuteDSLOpOverrides._apply_binary_op(x0, x1, "({a} | {b})")
|
||||
|
||||
# Bitwise operations (override parent class to properly CSE)
|
||||
@staticmethod
|
||||
# pyrefly: ignore [bad-override]
|
||||
def bitwise_and(x: CuteDSLArg, y: CuteDSLArg) -> CuteDSLArg:
|
||||
return CuteDSLOpOverrides._apply_binary_op(x, y, "({a} & {b})")
|
||||
|
||||
@staticmethod
|
||||
# pyrefly: ignore [bad-override]
|
||||
def bitwise_or(x: CuteDSLArg, y: CuteDSLArg) -> CuteDSLArg:
|
||||
return CuteDSLOpOverrides._apply_binary_op(x, y, "({a} | {b})")
|
||||
|
||||
@staticmethod
|
||||
# pyrefly: ignore [bad-override]
|
||||
def bitwise_xor(x: CuteDSLArg, y: CuteDSLArg) -> CuteDSLArg:
|
||||
return CuteDSLOpOverrides._apply_binary_op(x, y, "({a} ^ {b})")
|
||||
|
||||
@staticmethod
|
||||
# pyrefly: ignore [bad-override]
|
||||
def bitwise_not(x: CuteDSLArg) -> CuteDSLArg:
|
||||
return CuteDSLOpOverrides._apply_unary_op(x, "(~{x})")
|
||||
|
||||
@staticmethod
|
||||
# pyrefly: ignore [bad-override]
|
||||
def bitwise_left_shift(x: CuteDSLArg, y: CuteDSLArg) -> CuteDSLArg:
|
||||
return CuteDSLOpOverrides._apply_binary_op(x, y, "({a} << {b})")
|
||||
|
||||
@staticmethod
|
||||
# pyrefly: ignore [bad-override]
|
||||
def bitwise_right_shift(x: CuteDSLArg, y: CuteDSLArg) -> CuteDSLArg:
|
||||
return CuteDSLOpOverrides._apply_binary_op(x, y, "({a} >> {b})")
|
||||
|
||||
@staticmethod
|
||||
def logical_not(a):
|
||||
"""Logical NOT."""
|
||||
return CuteDSLOpOverrides._apply_unary_op(a, "({x} == 0)")
|
||||
|
||||
# Comparison operations
|
||||
@staticmethod
|
||||
def eq(a: CuteDSLArg, b: CuteDSLArg) -> CuteDSLArg:
|
||||
return CuteDSLOpOverrides._apply_binary_op(a, b, "operator.eq({a}, {b})")
|
||||
|
||||
@staticmethod
|
||||
def ne(a: CuteDSLArg, b: CuteDSLArg) -> CuteDSLArg:
|
||||
return CuteDSLOpOverrides._apply_binary_op(a, b, "operator.ne({a}, {b})")
|
||||
|
||||
@staticmethod
|
||||
def lt(a: CuteDSLArg, b: CuteDSLArg) -> CuteDSLArg:
|
||||
return CuteDSLOpOverrides._apply_binary_op(a, b, "operator.lt({a}, {b})")
|
||||
|
||||
@staticmethod
|
||||
def le(a: CuteDSLArg, b: CuteDSLArg) -> CuteDSLArg:
|
||||
return CuteDSLOpOverrides._apply_binary_op(a, b, "operator.le({a}, {b})")
|
||||
|
||||
@staticmethod
|
||||
def gt(a: CuteDSLArg, b: CuteDSLArg) -> CuteDSLArg:
|
||||
return CuteDSLOpOverrides._apply_binary_op(a, b, "operator.gt({a}, {b})")
|
||||
|
||||
@staticmethod
|
||||
def ge(a: CuteDSLArg, b: CuteDSLArg) -> CuteDSLArg:
|
||||
return CuteDSLOpOverrides._apply_binary_op(a, b, "operator.ge({a}, {b})")
|
||||
+141
@@ -0,0 +1,141 @@
|
||||
# mypy: allow-untyped-defs
|
||||
import hashlib
|
||||
import logging
|
||||
from collections.abc import Sequence
|
||||
from typing import cast
|
||||
|
||||
from torch._inductor.utils import Placeholder
|
||||
from torch.utils._ordered_set import OrderedSet
|
||||
|
||||
from ... import config
|
||||
from ...codecache import code_hash, get_path
|
||||
from ...ir import CuteDSLTemplateBuffer
|
||||
from ...scheduler import (
|
||||
BaseSchedulerNode,
|
||||
BaseScheduling,
|
||||
FusedSchedulerNode,
|
||||
SchedulerNode,
|
||||
)
|
||||
from ...select_algorithm import PartialRender
|
||||
from ...utils import get_fused_kernel_name, get_kernel_metadata
|
||||
from ...virtualized import V
|
||||
from ..common import BackendFeature, IndentedBuffer
|
||||
|
||||
|
||||
log = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class CuteDSLScheduling(BaseScheduling):
|
||||
"""
|
||||
Scheduling implementation for CuteDSL (CUTLASS Python DSL) kernels.
|
||||
This class is intended to be used in combination with other schedulers,
|
||||
and delegated to by CUDACombinedScheduling.
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def get_backend_features(cls, device) -> OrderedSet[BackendFeature]:
|
||||
return OrderedSet()
|
||||
|
||||
@staticmethod
|
||||
def is_cutedsl_template(node: BaseSchedulerNode) -> bool:
|
||||
"""Check if a node is a CuteDSL template."""
|
||||
return isinstance(node, SchedulerNode) and isinstance(
|
||||
node.node, CuteDSLTemplateBuffer
|
||||
)
|
||||
|
||||
def is_cutedsl_fused_template(self, node: BaseSchedulerNode) -> bool:
|
||||
"""Check if a node is a fused CuteDSL template."""
|
||||
return isinstance(node, FusedSchedulerNode) and self.is_cutedsl_template(node)
|
||||
|
||||
def can_fuse_vertical(
|
||||
self, node1: BaseSchedulerNode, node2: BaseSchedulerNode
|
||||
) -> bool:
|
||||
"""
|
||||
TODO CuteDSL doesn't support vertical fusion yet.
|
||||
This could be extended in the future for epilogue fusion.
|
||||
"""
|
||||
return False
|
||||
|
||||
def define_kernel(self, src_code_str: str, node_schedule) -> str:
|
||||
"""Produce the kernel string
|
||||
Args:
|
||||
src_code_str: The finalized kernel code string
|
||||
node_schedule: List of nodes in the schedule
|
||||
|
||||
Note:
|
||||
This is a little weird since async_compile.cutedsl() has to write the string to
|
||||
a file in order to cute compile it. Feels bad to have two...
|
||||
"""
|
||||
wrapper = V.graph.wrapper_code
|
||||
|
||||
# Use the string as the key for caching
|
||||
if src_code_str in wrapper.src_to_kernel:
|
||||
kernel_name = wrapper.src_to_kernel[src_code_str]
|
||||
else:
|
||||
fused_name = (
|
||||
get_fused_kernel_name(node_schedule, config.triton.descriptive_names)
|
||||
if config.triton.descriptive_names
|
||||
else ""
|
||||
)
|
||||
|
||||
kernel_hash = hashlib.sha256(src_code_str.encode("utf-8")).hexdigest()[:8]
|
||||
if fused_name == "fused":
|
||||
kernel_name = f"cutedsl_{kernel_hash}"
|
||||
else:
|
||||
kernel_name = f"cutedsl_{fused_name}_{kernel_hash}"
|
||||
wrapper.src_to_kernel[src_code_str] = kernel_name
|
||||
src_code_str = src_code_str.replace(
|
||||
str(Placeholder.KERNEL_NAME), kernel_name
|
||||
)
|
||||
|
||||
_, _, kernel_path = get_path(code_hash(src_code_str), "py")
|
||||
|
||||
compile_wrapper = IndentedBuffer()
|
||||
compile_wrapper.writeline(f"async_compile.cutedsl({kernel_name!r}, r'''")
|
||||
compile_wrapper.splice(src_code_str, strip=True)
|
||||
compile_wrapper.writeline("''')")
|
||||
|
||||
metadata_comment = f"# kernel path: {kernel_path}"
|
||||
origins, detailed_origins = get_kernel_metadata(node_schedule, wrapper)
|
||||
metadata_comment += "\n" + origins + "\n" + detailed_origins
|
||||
wrapper.define_kernel(
|
||||
kernel_name, compile_wrapper.getvalue(), metadata_comment
|
||||
)
|
||||
return kernel_name
|
||||
|
||||
def codegen_template(
|
||||
self,
|
||||
template_node: BaseSchedulerNode,
|
||||
epilogue_nodes: Sequence[BaseSchedulerNode],
|
||||
prologue_nodes: Sequence[BaseSchedulerNode],
|
||||
):
|
||||
"""
|
||||
Codegen a CuteDSL template. Currently doesn't support fusion.
|
||||
"""
|
||||
assert self.is_cutedsl_template(template_node), (
|
||||
"Template node passed to CuteDSLScheduling.codegen_template must be a "
|
||||
"SchedulerNode that wraps a CuteDSLTemplateBuffer"
|
||||
)
|
||||
# TODO remove when supported
|
||||
assert not epilogue_nodes, "CuteDSL doesn't support epilogue fusion yet"
|
||||
assert not prologue_nodes, "CuteDSL doesn't support prologue fusion yet"
|
||||
|
||||
template_node = cast(SchedulerNode, template_node)
|
||||
ctb: CuteDSLTemplateBuffer = cast(CuteDSLTemplateBuffer, template_node.node)
|
||||
|
||||
kernel, render = ctb.make_kernel_render(ctb) # type: ignore[misc]
|
||||
template_node.mark_run()
|
||||
src_code = render()
|
||||
# Finalize PartialRender if needed
|
||||
if isinstance(src_code, PartialRender):
|
||||
src_code_str = src_code.finalize_all()
|
||||
else:
|
||||
src_code_str = src_code
|
||||
|
||||
with V.set_kernel_handler(kernel):
|
||||
node_schedule = [template_node]
|
||||
kernel_name = self.define_kernel(src_code_str, node_schedule)
|
||||
self.codegen_comment(node_schedule, kernel_name)
|
||||
kernel.call_kernel(kernel_name, ctb)
|
||||
V.graph.removed_buffers |= kernel.removed_buffers
|
||||
self.free_buffers_in_scheduler()
|
||||
+212
@@ -0,0 +1,212 @@
|
||||
# mypy: allow-untyped-defs
|
||||
import functools
|
||||
import itertools
|
||||
from collections.abc import Iterable
|
||||
from typing import Any
|
||||
from unittest.mock import patch
|
||||
|
||||
from torch._inductor.utils import Placeholder
|
||||
from torch._inductor.virtualized import V
|
||||
from torch._logging import getArtifactLogger
|
||||
|
||||
from ...autotune_process import CuteDSLBenchmarkRequest, TensorMeta
|
||||
from ...ir import Buffer, ChoiceCaller, CuteDSLTemplateBuffer, IRNode, Layout, TensorBox
|
||||
from ..common import KernelTemplate
|
||||
from .cutedsl_kernel import CuteDSLTemplateKernel
|
||||
|
||||
|
||||
log = getArtifactLogger(__name__, "output_code")
|
||||
|
||||
|
||||
class CuteDSLTemplate(KernelTemplate):
|
||||
"""Template for generating CuteDSL (CUTLASS Python DSL) kernels."""
|
||||
|
||||
kernel_type: type[Any] = CuteDSLTemplateKernel
|
||||
index_counter = itertools.count()
|
||||
all_templates: dict[str, "CuteDSLTemplate"] = {}
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
name: str,
|
||||
source: str,
|
||||
subgraph_fn: Any | None = None,
|
||||
mask_fn: Any | None = None,
|
||||
) -> None:
|
||||
super().__init__(name)
|
||||
self.source = source
|
||||
self.subgraph_fn = subgraph_fn
|
||||
self.mask_fn = mask_fn
|
||||
self.template = CuteDSLTemplate._template_from_string(source)
|
||||
assert name not in self.all_templates, f"duplicate template name, {name}"
|
||||
CuteDSLTemplate.all_templates[name] = self
|
||||
|
||||
@staticmethod
|
||||
@functools.lru_cache(None)
|
||||
# pyrefly: ignore [bad-override]
|
||||
def _template_from_string(source: str) -> Any:
|
||||
return KernelTemplate._template_from_string(source)
|
||||
|
||||
def maybe_append_choice(
|
||||
self, choices: list[Any], **kwargs: Any
|
||||
) -> NotImplementedError | None:
|
||||
"""
|
||||
Maybe generates a new ChoiceCaller and appends it into existing choices.
|
||||
Returns None if success, otherwise returns the error.
|
||||
"""
|
||||
try:
|
||||
choices.append(self.generate(**kwargs))
|
||||
return None
|
||||
except NotImplementedError as e:
|
||||
log.debug("CuteDSL template choice generation failed: %s", e) # noqa: G200
|
||||
return e
|
||||
except Exception as e:
|
||||
log.debug("CuteDSL template choice generation error: %s", e) # noqa: G200
|
||||
return NotImplementedError(f"CuteDSL template failed: {e}")
|
||||
|
||||
def generate(self, **kwargs: Any) -> ChoiceCaller:
|
||||
"""Generate the CuteDSL kernel caller."""
|
||||
input_nodes = kwargs.pop("input_nodes")
|
||||
layout = kwargs.pop("layout")
|
||||
mutated_inputs = kwargs.pop("mutated_inputs", None)
|
||||
subgraphs = kwargs.pop("subgraphs", None)
|
||||
template_kwargs = dict(kwargs)
|
||||
|
||||
kernel_name = f"cutedsl_{self.name}_{next(self.index_counter)}"
|
||||
|
||||
if self.template is None:
|
||||
raise RuntimeError("Template compilation failed (Jinja2 required)")
|
||||
|
||||
self.output_node: Buffer = Buffer(name="buf_out", layout=layout)
|
||||
# Patch V.graph.get_dtype to handle the fake buf_out buffer
|
||||
with patch.object(
|
||||
V.graph, "get_dtype", KernelTemplate._fake_get_dtype(self.output_node)
|
||||
):
|
||||
kernel = self.kernel_type(
|
||||
kernel_name=kernel_name,
|
||||
input_nodes=input_nodes,
|
||||
output_node=self.output_node,
|
||||
subgraphs=subgraphs,
|
||||
)
|
||||
code = kernel.render(self.template, **kwargs)
|
||||
|
||||
log.debug("Generated CuteDSL Code:\n%s", code)
|
||||
|
||||
bmreq = CuteDSLBenchmarkRequest(
|
||||
kernel_name=kernel_name,
|
||||
input_tensor_meta=TensorMeta.from_irnodes(input_nodes),
|
||||
output_tensor_meta=TensorMeta.from_irnodes(self.output_node),
|
||||
extra_args=tuple(),
|
||||
source_code=code,
|
||||
)
|
||||
|
||||
def make_kernel_render(out_node, hint_override: int | None = None):
|
||||
"""
|
||||
Factory function that creates a kernel renderer for the final output.
|
||||
|
||||
This closure captures the current template and parameters, but allows
|
||||
the output node to be specified later. This is used during the final
|
||||
kernel selection phase when the actual output buffer is available.
|
||||
"""
|
||||
render_kernel = self.kernel_type(
|
||||
kernel_name=str(Placeholder.KERNEL_NAME),
|
||||
input_nodes=input_nodes,
|
||||
output_node=out_node,
|
||||
subgraphs=subgraphs,
|
||||
)
|
||||
|
||||
def render():
|
||||
return render_kernel.render(self.template, **kwargs)
|
||||
|
||||
return render_kernel, render
|
||||
|
||||
return CuteDSLTemplateCaller(
|
||||
name=kernel_name,
|
||||
input_nodes=input_nodes,
|
||||
layout=layout,
|
||||
make_kernel_render=make_kernel_render,
|
||||
bmreq=bmreq,
|
||||
template=self,
|
||||
mutated_inputs=mutated_inputs,
|
||||
template_kwargs=template_kwargs,
|
||||
)
|
||||
|
||||
|
||||
class CuteDSLTemplateCaller(ChoiceCaller):
|
||||
"""Caller for CuteDSL templates that integrates with the autotuning system."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
name: str,
|
||||
input_nodes: list[Buffer],
|
||||
layout: Layout,
|
||||
make_kernel_render: Any,
|
||||
bmreq: CuteDSLBenchmarkRequest,
|
||||
template: "CuteDSLTemplate",
|
||||
mutated_inputs: Iterable[IRNode] | None = None,
|
||||
template_kwargs: dict[str, Any] | None = None,
|
||||
):
|
||||
description = self._build_description(name, template_kwargs)
|
||||
super().__init__(
|
||||
name=name,
|
||||
input_nodes=input_nodes,
|
||||
layout=layout,
|
||||
description=description,
|
||||
)
|
||||
self.make_kernel_render = make_kernel_render
|
||||
self.bmreq = bmreq
|
||||
self.template = template
|
||||
self.mutated_inputs = mutated_inputs
|
||||
|
||||
def _build_description(
|
||||
self, name: str, template_kwargs: dict[str, Any] | None
|
||||
) -> str:
|
||||
if not template_kwargs:
|
||||
return f"CuteDSL template {name}"
|
||||
kwargs_desc = ", ".join(f"{k}={v}" for k, v in template_kwargs.items())
|
||||
return f"CuteDSL template {name} ({kwargs_desc})"
|
||||
|
||||
def __str__(self) -> str:
|
||||
return f"CuteDSLTemplateCaller({self.name})"
|
||||
|
||||
def benchmark(self, *args, out) -> float:
|
||||
"""Benchmark the kernel execution."""
|
||||
return self.bmreq.benchmark(*args, out=out)
|
||||
|
||||
def output_node(self) -> TensorBox:
|
||||
"""Create the output node for this template choice."""
|
||||
buffer = CuteDSLTemplateBuffer(
|
||||
layout=self.layout,
|
||||
inputs=self.input_nodes,
|
||||
make_kernel_render=self.make_kernel_render,
|
||||
template=self.template,
|
||||
mutated_inputs=self.mutated_inputs,
|
||||
)
|
||||
# Pass KTC annotation to the buffer for encoding
|
||||
if "ktc" in self.annotations:
|
||||
buffer.annotations["ktc"] = self.annotations["ktc"]
|
||||
return TensorBox.create(buffer)
|
||||
|
||||
def call_name(self) -> str:
|
||||
"""Return the kernel call name."""
|
||||
return self.name
|
||||
|
||||
def to_callable(self) -> Any:
|
||||
"""Return callable that can execute this kernel."""
|
||||
return self.make_kernel_render
|
||||
|
||||
def hash_key(self) -> str:
|
||||
"""Return unique hash key for this choice."""
|
||||
return "-".join(
|
||||
[
|
||||
self.name.rsplit("_", 1)[0],
|
||||
self.bmreq.module_cache_key,
|
||||
]
|
||||
)
|
||||
|
||||
def info_dict(self) -> dict[str, Any]:
|
||||
"""Return information about this kernel."""
|
||||
return {
|
||||
"name": self.name,
|
||||
"backend": "CuteDSL",
|
||||
"template": self.template.name,
|
||||
}
|
||||
@@ -0,0 +1,122 @@
|
||||
# mypy: allow-untyped-defs
|
||||
import functools
|
||||
import hashlib
|
||||
import inspect
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import time
|
||||
from typing import Any
|
||||
|
||||
import torch._inductor.config as config
|
||||
from torch._inductor.codecache import cutlass_key
|
||||
from torch._inductor.codegen.cutlass import serialization, utils
|
||||
from torch._inductor.codegen.cutlass.serialization import (
|
||||
get_cutlass_operation_serializer,
|
||||
)
|
||||
from torch._inductor.runtime.cache_dir_utils import cache_dir
|
||||
from torch._inductor.utils import clear_on_fresh_cache
|
||||
|
||||
|
||||
log = logging.getLogger(__name__)
|
||||
|
||||
|
||||
CONFIG_PREFIX: str = "configs"
|
||||
|
||||
|
||||
def get_config_request_key(
|
||||
arch: str,
|
||||
toolkit_version: str,
|
||||
instantiation_level: str,
|
||||
) -> str:
|
||||
"""
|
||||
Return a key for the full ops, based on cutlass key, arch, toolkit version, instantiation level, and serialization.py file hash.
|
||||
"""
|
||||
|
||||
# Get hash of serialization.py and cutlass_utils.py files using their module file paths
|
||||
def get_file_hash(file_module):
|
||||
file_path = inspect.getfile(file_module)
|
||||
with open(file_path, "rb") as f:
|
||||
return hashlib.sha256(f.read()).hexdigest()
|
||||
|
||||
serialization_hash = get_file_hash(serialization)
|
||||
cutlass_utils_hash = get_file_hash(utils)
|
||||
|
||||
hash_target = "-".join(
|
||||
[
|
||||
cutlass_key().hex(),
|
||||
arch,
|
||||
toolkit_version,
|
||||
instantiation_level,
|
||||
serialization_hash,
|
||||
cutlass_utils_hash,
|
||||
]
|
||||
)
|
||||
return hashlib.sha256(hash_target.encode("utf-8")).hexdigest()[0:8]
|
||||
|
||||
|
||||
def _generate_config_filename(request_key: str) -> str:
|
||||
"""
|
||||
Generate a filename for the full ops.
|
||||
"""
|
||||
return f"{CONFIG_PREFIX}_{request_key}.json"
|
||||
|
||||
|
||||
@clear_on_fresh_cache
|
||||
@functools.cache
|
||||
def maybe_fetch_ops(device_type: str) -> list[Any] | None:
|
||||
"""
|
||||
Fetch ops from databases.
|
||||
"""
|
||||
if config.force_disable_caches:
|
||||
return None
|
||||
|
||||
# setup
|
||||
arch: str = utils.cutlass_arch(device_type)
|
||||
version: str = utils.toolkit_version(device_type)
|
||||
if device_type == "cuda":
|
||||
# get_cuda_version might return "12.4.0" or "12.4"
|
||||
# but we want to use "12.4"
|
||||
version = ".".join(version.split(".")[:2])
|
||||
instantiation_level: str = config.cutlass.cutlass_instantiation_level
|
||||
|
||||
# filename and filepath
|
||||
request_key: str = get_config_request_key(arch, version, instantiation_level)
|
||||
filename: str = _generate_config_filename(request_key)
|
||||
filepath: str = os.path.join(cache_dir(), filename)
|
||||
|
||||
# try fetch
|
||||
serialized_ops: list[str] | None = None
|
||||
start_time = time.time()
|
||||
if os.path.isfile(filepath):
|
||||
# locally
|
||||
try:
|
||||
with open(filepath) as f:
|
||||
serialized_ops = json.load(f)
|
||||
|
||||
assert isinstance(serialized_ops, list), (
|
||||
f"Expected serialized ops is a list, got {type(serialized_ops)}"
|
||||
)
|
||||
except Exception:
|
||||
log.warning(
|
||||
"Failed to load CUTLASS config %s from local cache",
|
||||
filename,
|
||||
exc_info=True,
|
||||
)
|
||||
serialized_ops = None
|
||||
elif config.is_fbcode():
|
||||
from torch._inductor.fb.cutlass_remote_cache import (
|
||||
maybe_fetch_cutlass_configs_from_remote,
|
||||
)
|
||||
|
||||
# from remote
|
||||
serialized_ops = maybe_fetch_cutlass_configs_from_remote(filepath)
|
||||
|
||||
if serialized_ops is None:
|
||||
return None
|
||||
|
||||
# deserialize
|
||||
serializer = get_cutlass_operation_serializer()
|
||||
full_ops = [serializer.deserialize(x) for x in serialized_ops] # type: ignore[union-attr]
|
||||
log.info("Loaded ops from %s cache in %.3fs", filename, time.time() - start_time)
|
||||
return full_ops
|
||||
+2040
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,687 @@
|
||||
# mypy: allow-untyped-defs
|
||||
import functools
|
||||
import itertools
|
||||
import logging
|
||||
from collections import defaultdict
|
||||
from collections.abc import Callable
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Literal, TYPE_CHECKING
|
||||
|
||||
from sympy import Expr, symbols
|
||||
|
||||
import torch._inductor.config as config
|
||||
from torch import dtype as torch_dtype
|
||||
from torch._inductor.codegen.common import get_device_op_overrides
|
||||
from torch._inductor.codegen.cpp_wrapper_cpu import CppWrapperCpu
|
||||
from torch._inductor.scheduler import BaseSchedulerNode
|
||||
from torch._inductor.utils import do_bench_using_profiling, OrderedSet, Placeholder
|
||||
from torch.utils._sympy.value_ranges import ValueRanges
|
||||
|
||||
from .utils import DTYPE_TO_CUTLASS_TYPE
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from .template import ArgInfo
|
||||
|
||||
from ...autotune_process import CUTLASSBenchmarkRequest
|
||||
from ...ir import (
|
||||
Buffer,
|
||||
ChoiceCaller,
|
||||
CUTLASSTemplateBuffer,
|
||||
IRNode,
|
||||
Layout,
|
||||
PrimitiveInfoType,
|
||||
TensorBox,
|
||||
)
|
||||
from ...utils import sympy_product
|
||||
from ...virtualized import V
|
||||
from ..common import (
|
||||
CSEVariable,
|
||||
IndentedBuffer,
|
||||
Kernel,
|
||||
OpOverrides,
|
||||
WorkspaceArg,
|
||||
WorkspaceZeroMode,
|
||||
)
|
||||
from ..cpp_utils import CppPrinter, DTYPE_TO_CPP
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from torch._inductor.codegen.cutlass.template import CUTLASSTemplate
|
||||
|
||||
log = logging.getLogger(__name__)
|
||||
|
||||
cexpr = CppPrinter().doprint
|
||||
|
||||
|
||||
def _normalize_idx(index: int, total_length: int) -> int:
|
||||
return index if index >= 0 else index + total_length
|
||||
|
||||
|
||||
ValidLayoutSymbols = Literal["M", "N", "K", "B", "lda", "ldb", "ldc", "ldd"]
|
||||
ValidLayoutAttrs = Literal["size", "stride"]
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class LayoutArg:
|
||||
node: IRNode
|
||||
symbol: ValidLayoutSymbols
|
||||
attr: ValidLayoutAttrs
|
||||
dim: int
|
||||
|
||||
def matches(self, node, attr, dim) -> bool:
|
||||
return self.node == node and self.attr == attr and self.dim == dim
|
||||
|
||||
|
||||
class CUTLASSKernel(Kernel):
|
||||
"""
|
||||
Baseclass for Cutlass based Kernels
|
||||
"""
|
||||
|
||||
overrides = OpOverrides # type: ignore[assignment]
|
||||
|
||||
def __init__(self, *args, **kwargs) -> None:
|
||||
super().__init__(*args, **kwargs)
|
||||
self.layout_args: dict[str, list[LayoutArg]] = defaultdict(list)
|
||||
self.size_args: list[Expr | int] = []
|
||||
# Mapping from arg name to IRNode.
|
||||
self.named_nodes: dict[str, IRNode] = {}
|
||||
|
||||
def find_symbol(self, node: IRNode, attr: ValidLayoutAttrs, dim: int) -> str | None:
|
||||
arg = self.find_layout_arg(node, attr, dim)
|
||||
return arg.symbol if arg else None
|
||||
|
||||
def find_layout_arg(
|
||||
self, node: IRNode, attr: ValidLayoutAttrs, dim: int
|
||||
) -> LayoutArg | None:
|
||||
matches = [
|
||||
arg
|
||||
for arg in itertools.chain.from_iterable(self.layout_args.values())
|
||||
if arg.matches(node, attr, dim)
|
||||
]
|
||||
if len(matches) >= 1:
|
||||
# Verify all matches have the same node, attribute, and dimension
|
||||
# And if they come from the same node, whichever symbol we use is fine.
|
||||
# if in runtime the logic changes, this would trigger guard
|
||||
first_match = matches[0]
|
||||
if not all(
|
||||
match.node == first_match.node
|
||||
and match.attr == first_match.attr
|
||||
and match.dim == first_match.dim
|
||||
for match in matches
|
||||
):
|
||||
raise AssertionError("All matching layout args should be identical")
|
||||
return first_match
|
||||
return None
|
||||
|
||||
def add_layout_arg(
|
||||
self, symbol: ValidLayoutSymbols, node: IRNode, attr: ValidLayoutAttrs, dim: int
|
||||
):
|
||||
arg = LayoutArg(node, symbol, attr, dim)
|
||||
self.layout_args[symbol].append(arg)
|
||||
|
||||
def init_layout_args(self) -> None:
|
||||
X = self.named_nodes["X"]
|
||||
W = self.named_nodes["W"]
|
||||
Y = self.named_nodes["Y"]
|
||||
Bias = self.named_nodes.get("Bias", None)
|
||||
x_mdim = _normalize_idx(-2, len(X.get_size()))
|
||||
x_kdim = _normalize_idx(-1, len(X.get_size()))
|
||||
w_kdim = _normalize_idx(-2, len(W.get_size()))
|
||||
w_ndim = _normalize_idx(-1, len(W.get_size()))
|
||||
y_mdim = _normalize_idx(-2, len(Y.get_size()))
|
||||
y_ndim = _normalize_idx(-1, len(Y.get_size()))
|
||||
self.add_layout_arg("M", X, "size", x_mdim)
|
||||
self.add_layout_arg("K", X, "size", x_kdim)
|
||||
self.add_layout_arg("K", W, "size", w_kdim)
|
||||
self.add_layout_arg("N", W, "size", w_ndim)
|
||||
self.add_layout_arg("M", Y, "size", y_mdim)
|
||||
self.add_layout_arg("N", Y, "size", y_ndim)
|
||||
if len(X.get_size()) > 2:
|
||||
self.add_layout_arg("B", X, "size", 0)
|
||||
|
||||
lda_dim = self.find_ld_idx(X)
|
||||
ldb_dim = self.find_ld_idx(W)
|
||||
ldc_dim = self.find_ld_idx(Bias) if Bias else None
|
||||
ldd_dim = self.find_ld_idx(Y)
|
||||
self.add_layout_arg("lda", X, "stride", lda_dim)
|
||||
self.add_layout_arg("ldb", W, "stride", ldb_dim)
|
||||
if Bias is not None and ldc_dim is not None:
|
||||
self.add_layout_arg("ldc", Bias, "stride", ldc_dim)
|
||||
self.add_layout_arg("ldd", Y, "stride", ldd_dim)
|
||||
|
||||
def get_layout_args(self) -> tuple[Expr | int, ...]:
|
||||
X = self.named_nodes["X"]
|
||||
W = self.named_nodes["W"]
|
||||
Y = self.named_nodes["Y"]
|
||||
Bias = self.named_nodes.get("Bias", None)
|
||||
mdim = _normalize_idx(-2, len(X.get_size()))
|
||||
ndim = _normalize_idx(-1, len(W.get_size()))
|
||||
kdim = _normalize_idx(-1, len(X.get_size()))
|
||||
|
||||
def get_ld(node) -> Expr | int:
|
||||
dim = self.find_ld_idx(node)
|
||||
return node.get_stride()[dim]
|
||||
|
||||
M = X.get_size()[mdim]
|
||||
N = W.get_size()[ndim]
|
||||
K = X.get_size()[kdim]
|
||||
B = X.get_size()[0] if len(X.get_size()) > 2 else 1
|
||||
LDA = get_ld(X)
|
||||
LDB = get_ld(W)
|
||||
LDC = get_ld(Bias) if Bias else 0
|
||||
LDD = get_ld(Y)
|
||||
return (M, N, K, B, LDA, LDB, LDC, LDD)
|
||||
|
||||
def get_dynamic_shape_args(self) -> list[Expr | int]:
|
||||
return [*self.get_layout_args(), *self.size_args]
|
||||
|
||||
def get_offset_args(self) -> list[Expr]:
|
||||
return [node.get_layout().offset for node in self.named_nodes.values()]
|
||||
|
||||
@staticmethod
|
||||
def find_ld_idx(node: IRNode) -> int:
|
||||
strides = node.get_stride()
|
||||
# Handle 1D tensor case
|
||||
if V.graph.sizevars.statically_known_equals(strides[-1], 1):
|
||||
return _normalize_idx(-2, len(strides))
|
||||
|
||||
assert V.graph.sizevars.statically_known_equals(strides[-2], 1), strides[-2]
|
||||
return _normalize_idx(-1, len(strides))
|
||||
|
||||
|
||||
class CUTLASSTemplateKernel(CUTLASSKernel):
|
||||
"""
|
||||
Template kernels defined by Cutlass in C++.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
kernel_name: str,
|
||||
runtime_arg_info: list["ArgInfo"],
|
||||
runtime_arg_values: list[Any],
|
||||
device_type: str = "cuda", # type: ignore[assignment]
|
||||
) -> None:
|
||||
"""
|
||||
Initializes a new instance of the CUTLASSTemplateKernel class.
|
||||
|
||||
Args:
|
||||
kernel_name (str): The name of the kernel.
|
||||
"""
|
||||
super().__init__()
|
||||
self.kernel_name = kernel_name
|
||||
self.runtime_arg_info = runtime_arg_info
|
||||
self.runtime_arg_values = runtime_arg_values
|
||||
self.device_type = device_type
|
||||
self.device_codegen = get_device_op_overrides(self.device_type)
|
||||
self._EXTRA_CPP_ARGS = f"size_t* workspace_size, uint8_t* workspace, {self.device_codegen.cpp_stream_type()} stream"
|
||||
|
||||
def check_not_null(self, node: IRNode) -> str:
|
||||
"""
|
||||
Generates code to check that a node is not null.
|
||||
"""
|
||||
if node is None:
|
||||
return ""
|
||||
|
||||
size_str = self.size(node, 0, -1)
|
||||
name_str = self.arg_name(node)
|
||||
if name_str is None:
|
||||
return ""
|
||||
|
||||
res = IndentedBuffer(initial_indent=2)
|
||||
res.tabwidth = 1
|
||||
res.splice(
|
||||
f"""
|
||||
{{
|
||||
if (!{name_str}) {{
|
||||
int64_t {name_str}_size = {size_str};
|
||||
if ({name_str}_size > 0) {{
|
||||
throw std::runtime_error("input {name_str} is null but size is not 0!");
|
||||
}}
|
||||
}}
|
||||
}}
|
||||
"""
|
||||
)
|
||||
return res.getvalue()
|
||||
|
||||
def get_signature(self) -> str:
|
||||
return self.signature
|
||||
|
||||
def def_kernel(
|
||||
self,
|
||||
inputs: list[IRNode],
|
||||
outputs: list[IRNode],
|
||||
names_str: str = "",
|
||||
input_reorder: list[int] | None = None,
|
||||
) -> str:
|
||||
"""
|
||||
Hook called from template code to generate function definition and
|
||||
needed args.
|
||||
|
||||
Args:
|
||||
inputs: List of input IRNodes
|
||||
outputs: List of output IRNodes
|
||||
names_str: Comma separated list of input + output argument names.
|
||||
input_reorder: The actual order of input nodes.
|
||||
e.g. The template might have input argument defined as [X, W, Bias],
|
||||
and the actual input passed into this template could be [Bias, X, W].
|
||||
In this case, the `input_reorder` would be [2, 0, 1].
|
||||
additional_size_args: Additional size arguments for epilogue inputs
|
||||
"""
|
||||
# NB: name order matters here, it's used to match up offsets
|
||||
names = [x.strip() for x in names_str.strip().split(",")]
|
||||
if len(inputs) + len(outputs) != len(names):
|
||||
raise RuntimeError(
|
||||
f"{len(inputs) + len(outputs)=} != {len(names)=}, {inputs=}, {outputs=}, {names=}"
|
||||
)
|
||||
|
||||
if input_reorder is not None:
|
||||
assert len(inputs) == len(input_reorder)
|
||||
else:
|
||||
input_reorder = list(range(len(inputs)))
|
||||
|
||||
for idx in input_reorder:
|
||||
name = names[idx]
|
||||
node = inputs[idx]
|
||||
if node is not None:
|
||||
self.named_nodes[name] = node
|
||||
self.args.input_buffers[node.get_name()] = name
|
||||
|
||||
free_symbols: OrderedSet[Expr] = OrderedSet()
|
||||
for name, node in zip(names[len(inputs) : len(inputs) + len(outputs)], outputs):
|
||||
if node is not None:
|
||||
# NB: named nodes must be populated in the order of names
|
||||
self.named_nodes[name] = node
|
||||
self.args.output_buffers[node.get_name()] = name
|
||||
|
||||
if name not in (
|
||||
"X",
|
||||
"W",
|
||||
"Bias",
|
||||
"Y",
|
||||
): # we handle these symbolic shapes explicitly
|
||||
for expr in itertools.chain(node.get_size(), node.get_stride()):
|
||||
if isinstance(expr, Expr):
|
||||
for s in expr.free_symbols:
|
||||
free_symbols.add(s) # type: ignore[arg-type]
|
||||
|
||||
arg_defs, *_ = self.args.cpp_argdefs(DTYPE_TO_CUTLASS_TYPE)
|
||||
|
||||
self.init_layout_args()
|
||||
size_vars = ["M", "N", "K", "B", "lda", "ldb", "ldc", "ldd"]
|
||||
size_vars.extend(str(s) for s in free_symbols)
|
||||
self.size_args.extend(free_symbols)
|
||||
size_args = [f"const int {s}" for s in size_vars]
|
||||
offset_args = [f"const int {name}_offset" for name in self.named_nodes]
|
||||
runtime_arg_decls = ",".join(
|
||||
[f"{arg.ty} {arg.name}" for arg in self.runtime_arg_info]
|
||||
)
|
||||
if runtime_arg_decls:
|
||||
runtime_arg_decls += ", "
|
||||
|
||||
signature = (
|
||||
f"int {self.kernel_name}({', '.join(arg_defs + size_args + offset_args)},\
|
||||
{runtime_arg_decls}{self._EXTRA_CPP_ARGS})"
|
||||
)
|
||||
self.signature = signature
|
||||
return signature
|
||||
|
||||
def call_kernel(
|
||||
self,
|
||||
name: str,
|
||||
node: "CUTLASSTemplateBuffer", # type: ignore[name-defined]
|
||||
) -> None:
|
||||
"""
|
||||
Generates code to call the kernel through V.graph.wrapper_code.
|
||||
used from within torch._inductor.wrapper.PythonWrapperCodegen
|
||||
|
||||
name: Name of kernel function.
|
||||
node: The CUTLASSTemplateBuffer node which contains information about the kernel, it's fused epilogue nodes
|
||||
as well as all required inputs and outputs.
|
||||
"""
|
||||
wrapper = V.graph.wrapper_code
|
||||
|
||||
arg_types: list[Any]
|
||||
if V.graph.cpp_wrapper:
|
||||
# Make sure we initialize these kernels since they're exported as
|
||||
# C-style symbol names.
|
||||
assert isinstance(wrapper, CppWrapperCpu)
|
||||
wrapper.initialized_kernels[name] = self
|
||||
# We always originally initialize name with "KERNEL_NAME". So, we
|
||||
# we replace with the real kernel name passed as an arg to this function.
|
||||
self.signature = self.signature.replace(str(Placeholder.KERNEL_NAME), name)
|
||||
_, call_args, arg_types = self.args.cpp_argdefs(DTYPE_TO_CUTLASS_TYPE)
|
||||
else:
|
||||
_, call_args, _, arg_types = self.args.python_argdefs()
|
||||
|
||||
dynamic_shape_args = self.get_dynamic_shape_args()
|
||||
offset_args = self.get_offset_args()
|
||||
call_args.extend(dynamic_shape_args) # type: ignore[arg-type]
|
||||
call_args.extend(offset_args) # type: ignore[arg-type]
|
||||
for arg in self.runtime_arg_values:
|
||||
call_args.append(str(arg))
|
||||
arg_types.extend("const int" for _ in dynamic_shape_args)
|
||||
arg_types.extend("const int" for _ in offset_args)
|
||||
for arg in self.runtime_arg_info:
|
||||
arg_types.append(arg.ty)
|
||||
# dynamo wraps unspec variable as 0d CPU tensor, need convert to scalar
|
||||
for i in range(len(call_args)):
|
||||
if V.graph.is_unspec_arg(call_args[i]):
|
||||
call_args[i] = call_args[i] + ".item()"
|
||||
elif isinstance(arg_types[i], torch_dtype):
|
||||
call_args[i] = (
|
||||
call_args[i]
|
||||
if V.graph.cpp_wrapper
|
||||
else f"c_void_p({call_args[i]}.data_ptr())"
|
||||
)
|
||||
|
||||
# workspace_size ptr is NULL to mark this call is not intended for retrieving workspace_size.
|
||||
# workspace_size should have already been retrieved prior to this call.
|
||||
# workspace_size is here.
|
||||
call_args.append("nullptr" if V.graph.cpp_wrapper else "None")
|
||||
if V.graph.cpp_wrapper:
|
||||
arg_types.append("size_t*")
|
||||
|
||||
if node.get_workspace_size() > 0:
|
||||
ws = WorkspaceArg(
|
||||
count=node.get_workspace_size(),
|
||||
device=V.graph.get_current_device_or_throw(),
|
||||
zero_mode=WorkspaceZeroMode.UNINITIALIZED,
|
||||
outer_name=WorkspaceArg.unique_name(),
|
||||
)
|
||||
wrapper.generate_workspace_allocation(ws)
|
||||
workspace = str(ws.outer_name)
|
||||
call_args.append(
|
||||
workspace
|
||||
if V.graph.cpp_wrapper
|
||||
else f"c_void_p({workspace}.data_ptr())"
|
||||
)
|
||||
else:
|
||||
ws = None
|
||||
call_args.append("nullptr" if V.graph.cpp_wrapper else "None")
|
||||
if V.graph.cpp_wrapper:
|
||||
arg_types.append("uint8_t*")
|
||||
|
||||
wrapper.generate_kernel_call(
|
||||
name,
|
||||
call_args,
|
||||
triton=False,
|
||||
arg_types=arg_types,
|
||||
)
|
||||
if ws:
|
||||
wrapper.generate_workspace_deallocation(ws)
|
||||
|
||||
def dtype(self, node: IRNode) -> str | None:
|
||||
"""
|
||||
Generates code which represents dtype of a given node.
|
||||
"""
|
||||
|
||||
if node is None:
|
||||
return "void"
|
||||
return DTYPE_TO_CPP.get(node.get_layout().dtype)
|
||||
|
||||
def cutlass_dtype(self, node: IRNode, default_dtype="void") -> str | None:
|
||||
# Helper method, called into from CUTLASSGemmTemplate
|
||||
if node is None:
|
||||
return default_dtype
|
||||
from torch._inductor.codegen.cutlass.template import CUTLASSTemplate
|
||||
|
||||
return CUTLASSTemplate._DTYPE_TO_CUTLASS[node.get_layout().dtype]
|
||||
|
||||
def max_valid_index(self, node: IRNode, default=-1):
|
||||
# Helper method, called into from CUTLASSGemmTemplate
|
||||
if node is None:
|
||||
return default
|
||||
max_valid_offset = 0
|
||||
for i in range(len(node.get_size())):
|
||||
max_valid_offset += (node.get_size()[i] - 1) * node.get_stride()[i]
|
||||
return max_valid_offset
|
||||
|
||||
def ptr(self, node: IRNode) -> str:
|
||||
"""
|
||||
Generates code which represents pointer of a given node.
|
||||
"""
|
||||
|
||||
if node is None:
|
||||
return "nullptr"
|
||||
arg_name = self.arg_name(node)
|
||||
if arg_name is None:
|
||||
return "nullptr"
|
||||
return f"{arg_name} + {arg_name}_offset"
|
||||
|
||||
def size(
|
||||
self,
|
||||
node: IRNode,
|
||||
start_index: int,
|
||||
end_index: int | None = None,
|
||||
default_value: int = 0,
|
||||
) -> str:
|
||||
"""
|
||||
Hook called from template code to get the size of an arg.
|
||||
Generates code which represents size of a given node in [start_index, end_index).
|
||||
If node is None, returns default_value.
|
||||
|
||||
TODO: Will add needed args to pass it in if it is dynamic.
|
||||
"""
|
||||
|
||||
if node is None:
|
||||
return str(default_value)
|
||||
|
||||
start_index = _normalize_idx(start_index, len(node.get_size()))
|
||||
if end_index is None:
|
||||
end_index = start_index
|
||||
end_index = _normalize_idx(end_index, len(node.get_size()))
|
||||
sizes = [
|
||||
self.find_symbol(node, "size", dim=i) or node.get_size()[i]
|
||||
for i in range(start_index, end_index + 1)
|
||||
]
|
||||
if len(sizes) == 0:
|
||||
return str(default_value)
|
||||
|
||||
sizes = [symbols(v) if isinstance(v, str) else v for v in sizes]
|
||||
val = sympy_product(sizes)
|
||||
return val
|
||||
|
||||
def stride(self, node: IRNode, index: int, default_value: int = 0) -> str:
|
||||
"""
|
||||
Hook called from template code to get the stride of an arg.
|
||||
Generates code which represents stride of a given node at index.
|
||||
If node is None, returns default_value.
|
||||
|
||||
TODO: Will add needed args to pass it in if it is dynamic.
|
||||
"""
|
||||
|
||||
if node is None:
|
||||
return str(default_value)
|
||||
|
||||
index = _normalize_idx(index, len(node.get_size()))
|
||||
if index < 0:
|
||||
return str(default_value)
|
||||
|
||||
stride = node.get_stride()[index]
|
||||
if V.graph.sizevars.statically_known_leq(stride, 1):
|
||||
return str(stride)
|
||||
return self.find_symbol(node, "stride", dim=index) or str(stride)
|
||||
|
||||
def batch_stride(self, node: IRNode, default_value: int = 0) -> str:
|
||||
"""
|
||||
Hook called from template code to get the batch stride of an arg.
|
||||
Returns 0 if batch dim is not present.
|
||||
|
||||
This method assumes that batch stride is the largest stride.
|
||||
"""
|
||||
|
||||
if node is None:
|
||||
return str(default_value)
|
||||
|
||||
if len(node.get_size()) < 3:
|
||||
return str(default_value)
|
||||
|
||||
batch_stride = node.get_stride()[0]
|
||||
if V.graph.sizevars.statically_known_leq(batch_stride, 1):
|
||||
return str(batch_stride)
|
||||
|
||||
return "{}*{}".format(
|
||||
self.find_symbol(node, "size", dim=1) or node.get_size()[1],
|
||||
self.find_symbol(node, "size", dim=2) or node.get_size()[2],
|
||||
)
|
||||
|
||||
def row_or_column_stride(self, node: IRNode, default_value: int = 0) -> str:
|
||||
"""
|
||||
Hook called from template code to get the row or column stride of an arg.
|
||||
This is required by some CUTLASS 2.X APIs.
|
||||
If the node is in row_major, it returns stride[-2].
|
||||
If the node is in column_major, it returns stride[-1].
|
||||
|
||||
TODO: Will add needed args to pass it in if it is dynamic.
|
||||
"""
|
||||
|
||||
if node is None or len(node.get_stride()) < 2:
|
||||
return str(default_value)
|
||||
|
||||
stride0 = node.get_stride()[-1]
|
||||
stride1 = node.get_stride()[-2]
|
||||
if stride0 == 1:
|
||||
return cexpr(self.rename_indexing(stride1))
|
||||
elif stride1 == 1:
|
||||
return cexpr(self.rename_indexing(stride0))
|
||||
else:
|
||||
raise RuntimeError(
|
||||
f"At least 1 stride should be 1. Strides: {node.get_stride()=}"
|
||||
)
|
||||
|
||||
def load(self, name: str, index: Expr, mode: Any = None) -> CSEVariable:
|
||||
"""
|
||||
Mock load function for memory planning to optimize allocations properly.
|
||||
"""
|
||||
return self.create_cse_var(name, bounds=ValueRanges.unknown())
|
||||
|
||||
def store(self, name: str, index: Expr, value: Any, mode: Any = None) -> None:
|
||||
"""
|
||||
Mock store function for memory planning to optimize allocations properly.
|
||||
"""
|
||||
self.store_buffer_names.add(name)
|
||||
|
||||
|
||||
class CUTLASSTemplateCaller(ChoiceCaller):
|
||||
"""
|
||||
CUTLASSTemplateCaller
|
||||
|
||||
This class represents a caller for CUTLASS template kernels. It is a subclass of ChoiceCaller.
|
||||
Attributes:
|
||||
name (str): The name of the caller.
|
||||
category (str): The category of the caller.
|
||||
bmreq (CUTLASSBenchmarkRequest): The benchmark request for the caller.
|
||||
template_buffer (CUTLASSTemplateBuffer): The template buffer for the caller.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
name: str,
|
||||
category: str,
|
||||
input_nodes: list[Buffer],
|
||||
layout: Layout,
|
||||
make_kernel_render: Callable[
|
||||
[CUTLASSTemplateBuffer, list[BaseSchedulerNode] | None],
|
||||
tuple[CUTLASSTemplateKernel, functools.partial[str]],
|
||||
],
|
||||
bmreq: CUTLASSBenchmarkRequest,
|
||||
supports_epilogue_fusion: bool,
|
||||
template: "CUTLASSTemplate", # type: ignore[name-defined]
|
||||
info_kwargs: dict[str, PrimitiveInfoType | list[PrimitiveInfoType]] | None, # type: ignore[type-arg]
|
||||
description: str,
|
||||
) -> None:
|
||||
super().__init__(name, input_nodes, layout, description)
|
||||
self.category = category
|
||||
self.make_kernel_render = make_kernel_render
|
||||
self.bmreq = bmreq
|
||||
self.supports_epilogue_fusion = supports_epilogue_fusion
|
||||
self.template = template
|
||||
self.info_kwargs = info_kwargs
|
||||
|
||||
def precompile(self) -> None:
|
||||
assert self.bmreq is not None
|
||||
self.bmreq.precompile()
|
||||
|
||||
def benchmark(self, *args, out) -> float:
|
||||
assert self.bmreq is not None
|
||||
if config.profile_bandwidth_with_do_bench_using_profiling:
|
||||
algo = self.bmreq.make_run_fn(*args, out=out)
|
||||
return do_bench_using_profiling(algo)
|
||||
return self.bmreq.benchmark(*args, out=out)
|
||||
|
||||
def __str__(self) -> str:
|
||||
return f"CUTLASSTemplateCaller(source_file={self.bmreq.source_file})"
|
||||
|
||||
def call_name(self) -> str:
|
||||
return f"cutlass_template_kernels.{self.name}"
|
||||
|
||||
def kernel_hash_key(self) -> str:
|
||||
"""
|
||||
Return kernel hash key that does not depend on swizzle.
|
||||
"""
|
||||
return "-".join(
|
||||
[
|
||||
self.category,
|
||||
self.bmreq.hash_key,
|
||||
]
|
||||
)
|
||||
|
||||
def hash_key(self) -> str:
|
||||
"""
|
||||
Return kernel hash key that does not depend on swizzle.
|
||||
"""
|
||||
swizzle_str: str = (
|
||||
str(self.info_kwargs.get("swizzle"))
|
||||
if isinstance(self.info_kwargs, dict)
|
||||
else "None"
|
||||
)
|
||||
return "-".join(
|
||||
[
|
||||
self.category,
|
||||
self.bmreq.hash_key,
|
||||
swizzle_str,
|
||||
]
|
||||
)
|
||||
|
||||
def info_dict(self) -> dict[str, PrimitiveInfoType | list[PrimitiveInfoType]]:
|
||||
"""
|
||||
Information returned here is logged to the autotune log file when that is enabled.
|
||||
|
||||
In general, we should avoid calling this function as it is expensive to compute,
|
||||
and can add up very fast.
|
||||
"""
|
||||
if self.info_kwargs is not None and "op" in self.info_kwargs:
|
||||
op: Any = self.info_kwargs["op"]
|
||||
return {
|
||||
"backend": "CUDA",
|
||||
"op_type": type(op).__name__,
|
||||
"op_conf_name": str(op.configuration_name()),
|
||||
"op_arch": str(op.arch),
|
||||
"tile_shape": str(op.tile_description.tile_shape),
|
||||
"epilogue_schedule": str(op.epilogue_schedule),
|
||||
"kernel_schedule": str(op.kernel_schedule),
|
||||
"element_accumulator": str(op.accumulator_type()),
|
||||
"op_name": str(op.procedural_name()),
|
||||
"instruction_shape": str(
|
||||
op.tile_description.math_instruction.instruction_shape
|
||||
),
|
||||
"swizzle": str(self.info_kwargs["swizzle"]),
|
||||
}
|
||||
else:
|
||||
return {"backend": "CUDA", "op_type": "unknown"}
|
||||
|
||||
def output_node(self) -> TensorBox:
|
||||
self.bmreq.update_workspace_size()
|
||||
buffer = CUTLASSTemplateBuffer(
|
||||
layout=self.layout,
|
||||
inputs=self.input_nodes,
|
||||
make_kernel_render=self.make_kernel_render,
|
||||
workspace_size=self.bmreq.workspace_size,
|
||||
supports_epilogue_fusion=self.supports_epilogue_fusion,
|
||||
template=self.template,
|
||||
)
|
||||
# Pass KTC annotation to the buffer for encoding
|
||||
if "ktc" in self.annotations:
|
||||
buffer.annotations["ktc"] = self.annotations["ktc"]
|
||||
return TensorBox.create(buffer)
|
||||
+6
@@ -0,0 +1,6 @@
|
||||
import torch
|
||||
|
||||
|
||||
__version__ = torch.version.cuda
|
||||
|
||||
from .cuda import * # noqa: F403
|
||||
+24
@@ -0,0 +1,24 @@
|
||||
# mypy: disable-error-code="no-untyped-def"
|
||||
# flake8: noqa
|
||||
import torch
|
||||
|
||||
|
||||
class CUdeviceptr:
|
||||
pass
|
||||
|
||||
|
||||
class CUstream:
|
||||
def __init__(self, v):
|
||||
pass
|
||||
|
||||
|
||||
class CUresult:
|
||||
CUDA_SUCCESS = True
|
||||
|
||||
|
||||
class nvrtc:
|
||||
pass
|
||||
|
||||
|
||||
def cuDeviceGetCount():
|
||||
return (CUresult.CUDA_SUCCESS, torch.cuda.device_count())
|
||||
+17
@@ -0,0 +1,17 @@
|
||||
# mypy: disable-error-code="no-untyped-def"
|
||||
import torch.cuda
|
||||
|
||||
|
||||
class cudaError_t:
|
||||
cudaSuccess = True
|
||||
|
||||
|
||||
def cudaFree(n):
|
||||
return (cudaError_t.cudaSuccess,)
|
||||
|
||||
|
||||
def cudaGetDeviceProperties(d):
|
||||
class DummyError:
|
||||
value = False
|
||||
|
||||
return (DummyError(), torch.cuda.get_device_properties(d))
|
||||
+2
@@ -0,0 +1,2 @@
|
||||
# mypy: disable-error-code="var-annotated"
|
||||
Dot = None
|
||||
+3
@@ -0,0 +1,3 @@
|
||||
# typing: ignore
|
||||
# flake8: noqa
|
||||
from .special import *
|
||||
+2
@@ -0,0 +1,2 @@
|
||||
# mypy: disable-error-code="var-annotated"
|
||||
erf = None
|
||||
+289
@@ -0,0 +1,289 @@
|
||||
from collections.abc import Callable
|
||||
from typing import Any
|
||||
|
||||
from sympy import Expr
|
||||
|
||||
from torch._inductor.ir import (
|
||||
ComputedBuffer,
|
||||
InputBuffer,
|
||||
is_contiguous_strides_for_shape,
|
||||
)
|
||||
from torch.utils._ordered_set import OrderedSet
|
||||
|
||||
from ..utils import torch_dtype_to_cutlass_type, try_import_cutlass
|
||||
|
||||
|
||||
EpilogueFunctor = Any # EpilogueFunctor local class defined in _trace
|
||||
Buffer = ComputedBuffer | InputBuffer
|
||||
CutlassTupleType = Any # cutlass.backend.c_types.tuple_factory_.<locals>.TupleType
|
||||
CutlassVisitorType = Any # cutlass.backend.c_types.visitor_factory.<locals>.VisitorType
|
||||
CutlassArgType = (
|
||||
Any # Can be a CutlassTupleType, CutlassVisitorType, EmptyByte, or ctype.c_void_p
|
||||
)
|
||||
|
||||
|
||||
if try_import_cutlass():
|
||||
import ast
|
||||
import ctypes
|
||||
import textwrap
|
||||
|
||||
from cutlass_cppgen.backend.c_types import ( # type: ignore[import-not-found]
|
||||
EmptyByte,
|
||||
)
|
||||
from cutlass_cppgen.backend.epilogue import ( # type: ignore[import-not-found]
|
||||
dtype2ctype,
|
||||
)
|
||||
from cutlass_cppgen.backend.evt import ( # type: ignore[import-not-found]
|
||||
EpilogueFunctorVisitor,
|
||||
)
|
||||
from cutlass_cppgen.backend.evt.backend.emitter_base import ( # type: ignore[import-not-found]
|
||||
FusionCallbacks,
|
||||
)
|
||||
from cutlass_cppgen.backend.evt.backend.sm100_emitter import ( # type: ignore[import-not-found]
|
||||
Sm100CollectiveEpilogue,
|
||||
)
|
||||
from cutlass_cppgen.backend.evt.backend.sm90_emitter import ( # type: ignore[import-not-found]
|
||||
CollectiveEpilogue,
|
||||
)
|
||||
from cutlass_cppgen.backend.evt.frontend import ( # type: ignore[import-not-found]
|
||||
PythonASTFrontend,
|
||||
)
|
||||
from cutlass_cppgen.backend.evt.ir.tensor import ( # type: ignore[import-not-found]
|
||||
Tensor as CutlassTensor,
|
||||
)
|
||||
from cutlass_library import (
|
||||
DataType,
|
||||
EpilogueScheduleType,
|
||||
LayoutType,
|
||||
TileDescription,
|
||||
)
|
||||
|
||||
from torch._inductor.codegen.cuda import cuda_env
|
||||
from torch._inductor.utils import IndentedBuffer
|
||||
|
||||
_CUTLASS_C_DTYPES = OrderedSet(dtype2ctype.values()) # type: ignore[var-annotated]
|
||||
|
||||
class EVTArgRenames:
|
||||
"""Handles mapping buffer names to variable names in the cpp kernel signature and body"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.buf_renames: dict[str, str] = {}
|
||||
|
||||
def new_name(self, name: str) -> str:
|
||||
if name in self.buf_renames:
|
||||
return self.buf_renames[name]
|
||||
else:
|
||||
new_name = f"ptr_{len(self.buf_renames)}"
|
||||
self.buf_renames[name] = new_name
|
||||
return new_name
|
||||
|
||||
def get(self, name: str) -> str:
|
||||
return self.buf_renames.get(name, name)
|
||||
|
||||
def create_example_tensors(
|
||||
var_name_to_buffer_name: dict[str, str],
|
||||
name_to_buffer: dict[str, Buffer],
|
||||
size_hint_fn: Callable[[Expr | int], int],
|
||||
) -> dict[str, CutlassTensor]:
|
||||
def cutlass_tensor_from_buffer(
|
||||
buffer: Buffer,
|
||||
) -> CutlassTensor:
|
||||
shape = buffer.get_layout().size
|
||||
stride = buffer.get_layout().stride
|
||||
shape = tuple(size_hint_fn(x) for x in shape)
|
||||
stride = tuple(size_hint_fn(x) for x in stride)
|
||||
|
||||
is_row_major = is_contiguous_strides_for_shape(stride, shape)
|
||||
is_column_major = is_contiguous_strides_for_shape(stride[::-1], shape[::-1])
|
||||
|
||||
if not is_row_major and not is_column_major:
|
||||
raise RuntimeError(
|
||||
f"Cannot create example tensor for {buffer.get_name()} with \
|
||||
non-contiguous layout, received stride: {stride} and shape: {shape}"
|
||||
)
|
||||
|
||||
return CutlassTensor(
|
||||
shape=shape,
|
||||
layout_tag=(
|
||||
LayoutType.RowMajor if is_row_major else LayoutType.ColumnMajor
|
||||
),
|
||||
element=torch_dtype_to_cutlass_type(buffer.get_layout().dtype),
|
||||
)
|
||||
|
||||
return {
|
||||
key: cutlass_tensor_from_buffer(name_to_buffer[name])
|
||||
for key, name in var_name_to_buffer_name.items()
|
||||
}
|
||||
|
||||
def trace(
|
||||
fn_src: str,
|
||||
example_tensors: dict[str, CutlassTensor],
|
||||
accum_type: DataType,
|
||||
output_type: DataType,
|
||||
tile_description: TileDescription,
|
||||
epilogue_schedule: EpilogueScheduleType,
|
||||
name_to_buffer: dict[str, Buffer],
|
||||
size_hint_fn: Callable[[Expr | int], int],
|
||||
kernel_schedule: Any | None = None,
|
||||
**kwargs: dict[str, Any],
|
||||
) -> tuple[str, str, str, EVTArgRenames]:
|
||||
cuda_arch = int(cuda_env.get_cuda_arch()) # type: ignore[arg-type]
|
||||
assert cuda_arch >= 90, "Only SM90+ is supported for EVT"
|
||||
epilogue_functor = _trace(fn_src, example_tensors, cuda_arch, **kwargs)
|
||||
visitor = EpilogueFunctorVisitor(cuda_arch, epilogue_functor)
|
||||
fusion_callbacks = FusionCallbacks(visitor.graph, cuda_arch, emit_CD=False)
|
||||
if cuda_arch < 100:
|
||||
collective_epilogue = CollectiveEpilogue(
|
||||
tile_description,
|
||||
epilogue_schedule,
|
||||
accum_type,
|
||||
output_type,
|
||||
fusion_callbacks,
|
||||
)
|
||||
else:
|
||||
collective_epilogue = Sm100CollectiveEpilogue(
|
||||
tile_description=tile_description,
|
||||
kernel_schedule=kernel_schedule,
|
||||
epilogue_schedule=epilogue_schedule,
|
||||
element_accumulator=accum_type,
|
||||
element_d=output_type,
|
||||
fusion_callbacks=fusion_callbacks,
|
||||
)
|
||||
evt_name, evt_code = collective_epilogue.emit()
|
||||
evt_args, arg_renames = _render_argument_type(
|
||||
epilogue_functor, name_to_buffer, size_hint_fn
|
||||
)
|
||||
return evt_name, evt_args, evt_code, arg_renames
|
||||
|
||||
# Based off of
|
||||
# https://github.com/NVIDIA/cutlass/blob/df18f5e4f5de76bed8be1de8e4c245f2f5ec3020/python/cutlass/epilogue/epilogue.py#L117
|
||||
# This is modified to enable directly passing the source code of the epilogue vs getting it from a bona-fide python function
|
||||
# The reason for this is that inspect.getsource does not work with functions defined at runtime via exec/eval
|
||||
def _trace(
|
||||
fn_src: str,
|
||||
example_tensors: dict[str, CutlassTensor],
|
||||
cc: int,
|
||||
**kwargs: Any,
|
||||
) -> EpilogueFunctor:
|
||||
class EpilogueFunctor(PythonASTFrontend):
|
||||
def __init__(self, cc: int, **kwargs: Any):
|
||||
self.source = textwrap.dedent(fn_src)
|
||||
super().__init__(cc, **kwargs)
|
||||
|
||||
def parse(
|
||||
self,
|
||||
example_inputs: dict[str, CutlassTensor],
|
||||
) -> None:
|
||||
self.example_inputs = example_inputs
|
||||
self.ast = ast.parse(self.source)
|
||||
# pyrefly: ignore [missing-attribute]
|
||||
self.visit(self.ast)
|
||||
|
||||
cc = int(cuda_env.get_cuda_arch())
|
||||
epilogue_functor = EpilogueFunctor(cc=cc, **kwargs)
|
||||
epilogue_functor.trace(example_tensors)
|
||||
return epilogue_functor
|
||||
|
||||
def _render_argument_type(
|
||||
epilogue_functor: EpilogueFunctor,
|
||||
name_to_buffer: dict[str, Buffer],
|
||||
size_hint_fn: Callable[[Expr | int], int],
|
||||
) -> tuple[str, EVTArgRenames]:
|
||||
epilogue_thread_type = epilogue_functor.epilogue_thread_type
|
||||
arg_renames = EVTArgRenames()
|
||||
|
||||
# Fragile, but this is the only way to guarantee t is expected type because t is a local class
|
||||
def is_nested_visitor_type(t: type) -> bool:
|
||||
return (
|
||||
".".join([t.__module__, t.__qualname__])
|
||||
== "cutlass_cppgen.backend.c_types.visitor_factory.<locals>.VisitorType"
|
||||
)
|
||||
|
||||
buffer = IndentedBuffer()
|
||||
with buffer.set_tabwidth(2):
|
||||
|
||||
def render_argument_type(name: str, t: CutlassArgType) -> None:
|
||||
if issubclass(t, ctypes.c_byte):
|
||||
buffer.writeline(f"{{}}, /* {name} */")
|
||||
else:
|
||||
fields = [
|
||||
(
|
||||
fname,
|
||||
_get_arg_from_node(
|
||||
ty, name_to_buffer[name], size_hint_fn, arg_renames
|
||||
),
|
||||
)
|
||||
for fname, ty in t._fields_
|
||||
]
|
||||
field_strs = [
|
||||
f"/* {fname} */ {str(field)}" for fname, field in fields
|
||||
]
|
||||
buffer.writeline(f"{{{', '.join(field_strs)}}}, /* {name} */")
|
||||
|
||||
def render_thread_type(name: str, t: CutlassArgType) -> None:
|
||||
if is_nested_visitor_type(t):
|
||||
buffer.writeline(f"{{ /* {name} */")
|
||||
with buffer.indent():
|
||||
for name, inner_t in t._fields_:
|
||||
render_thread_type(name, inner_t)
|
||||
buffer.writeline("},")
|
||||
else:
|
||||
render_argument_type(name, t)
|
||||
|
||||
# unroll the recursion once to address special case formatting
|
||||
# namely, no ending comma and no indentation for the outermost thread type
|
||||
buffer.writeline("{ /* thread */")
|
||||
with buffer.indent(3):
|
||||
if is_nested_visitor_type(epilogue_thread_type):
|
||||
with buffer.indent():
|
||||
for name, inner_t in epilogue_thread_type._fields_:
|
||||
render_thread_type(name, inner_t)
|
||||
else:
|
||||
render_argument_type("thread", epilogue_thread_type)
|
||||
buffer.writeline("}")
|
||||
|
||||
return buffer.getvalue(), arg_renames
|
||||
|
||||
def _get_arg_from_node(
|
||||
arg_ty: type,
|
||||
node: Buffer,
|
||||
size_hint_fn: Callable[[Expr | int], int],
|
||||
arg_renames: EVTArgRenames,
|
||||
) -> str:
|
||||
from ..template import CUTLASSTemplate
|
||||
|
||||
# Today, arguments are either a pointer to the
|
||||
# node's memory, a stride tuple, the datatype
|
||||
# Once again, need to check for local class type for stride tuple
|
||||
if (
|
||||
str(arg_ty)
|
||||
== "<class 'cutlass_cppgen.backend.c_types.tuple_factory_.<locals>.TupleType'>"
|
||||
):
|
||||
DEFAULT_STRIDE_LEN = 3
|
||||
assert len(node.get_layout().stride) <= DEFAULT_STRIDE_LEN
|
||||
stride = [size_hint_fn(x) for x in node.get_layout().stride]
|
||||
for _ in range(DEFAULT_STRIDE_LEN - len(stride)):
|
||||
stride.append(0)
|
||||
|
||||
def render_stride(x: int) -> str:
|
||||
# Handle EBO for 0 and 1
|
||||
if x == 0:
|
||||
return "_0{}"
|
||||
elif x == 1:
|
||||
return "_1{}"
|
||||
else:
|
||||
return str(x)
|
||||
|
||||
return f"{{{', '.join([render_stride(x) for x in stride])}}}"
|
||||
|
||||
elif issubclass(arg_ty, ctypes.c_void_p):
|
||||
name = arg_renames.new_name(node.get_name())
|
||||
return f"({CUTLASSTemplate._DTYPE_TO_CUTLASS[node.get_layout().dtype]}*) ({name} + {name}_offset)"
|
||||
elif (
|
||||
arg_ty in _CUTLASS_C_DTYPES
|
||||
): # Assumption: this is the element dtype, this holds for all cutlass ir nodes currently
|
||||
return f"{CUTLASSTemplate._DTYPE_TO_CUTLASS[node.get_layout().dtype]}(0)"
|
||||
elif issubclass(arg_ty, EmptyByte):
|
||||
return "{}"
|
||||
|
||||
raise NotImplementedError(f"Unsupported arg type: {arg_ty}")
|
||||
+418
@@ -0,0 +1,418 @@
|
||||
# mypy: ignore-errors
|
||||
from ..utils import try_import_cutlass
|
||||
|
||||
|
||||
# copied / modified from original at
|
||||
# https://github.com/NVIDIA/cutlass/blob/8783c41851cd3582490e04e69e0cd756a8c1db7f/tools/library/scripts/gemm_operation.py#L658
|
||||
|
||||
if try_import_cutlass():
|
||||
import enum
|
||||
|
||||
from cutlass_library.gemm_operation import * # noqa: F401, F403
|
||||
from cutlass_library.library import * # noqa: F401, F403
|
||||
|
||||
_LOGGER = logging.getLogger(__name__)
|
||||
|
||||
class EmitGemmUniversal3xInstanceWithEVT:
|
||||
"""Responsible for emitting a CUTLASS 3.x template definition"""
|
||||
|
||||
def __init__(self, operation_suffix="", evt_name=None, device_type="cuda"):
|
||||
self.operation_suffix = operation_suffix
|
||||
self.includes = [
|
||||
"cutlass/cutlass.h",
|
||||
"cutlass/gemm/gemm.h",
|
||||
"cutlass/numeric_types.h",
|
||||
"cutlass/gemm/kernel/gemm_universal.hpp",
|
||||
"cutlass/gemm/collective/collective_builder.hpp",
|
||||
"cutlass/epilogue/collective/collective_builder.hpp",
|
||||
]
|
||||
self.builtin_epilogue_functor_template = """${epilogue_functor}<
|
||||
${element_d},
|
||||
${element_epilogue},
|
||||
${element_c},
|
||||
${element_epilogue}
|
||||
>"""
|
||||
self.evt_name = evt_name
|
||||
self.device_type = device_type
|
||||
self.gemm_template = """
|
||||
using ${operation_name}_epilogue =
|
||||
typename cutlass::epilogue::collective::CollectiveBuilder<
|
||||
${arch}, ${opcode_class_epi},
|
||||
cute::Shape<cute::_${tile_shape_m}, cute::_${tile_shape_n}, cute::_${tile_shape_k}>,
|
||||
cute::Shape<${cluster_shape_m}, ${cluster_shape_n}, ${cluster_shape_k}>,
|
||||
${epi_tile_mn},
|
||||
${element_accumulator}, ${element_epilogue},
|
||||
${element_c}, ${layout_c}, ${align_c},
|
||||
${element_d}, ${layout_d}, ${align_d},
|
||||
${epilogue_schedule},
|
||||
${epilogue_functor}
|
||||
>::CollectiveOp;
|
||||
|
||||
${mixed_dtype_prepare_code}
|
||||
|
||||
using ${operation_name}_mainloop =
|
||||
typename cutlass::gemm::collective::CollectiveBuilder<
|
||||
${arch}, ${opcode_class_main},
|
||||
${element_a}, ${layout_a}, ${align_a},
|
||||
${element_b}, ${layout_b}, ${align_b},
|
||||
${element_accumulator},
|
||||
cute::Shape<cute::_${tile_shape_m}, cute::_${tile_shape_n}, cute::_${tile_shape_k}>,
|
||||
cute::Shape<${cluster_shape_m}, ${cluster_shape_n}, ${cluster_shape_k}>,
|
||||
${stages},
|
||||
${kernel_schedule}
|
||||
>::CollectiveOp;
|
||||
|
||||
// Gemm operator ${operation_name}
|
||||
using ${operation_name}_base = cutlass::gemm::kernel::GemmUniversal<
|
||||
${problem_shape},
|
||||
${operation_name}_mainloop,
|
||||
${operation_name}_epilogue,
|
||||
${tile_scheduler}>;
|
||||
|
||||
// Define named type
|
||||
struct ${operation_name} :
|
||||
public ${operation_name}_base { };
|
||||
|
||||
"""
|
||||
|
||||
#
|
||||
def instance_template(self):
|
||||
return """
|
||||
${compile_guard_start}
|
||||
{
|
||||
using GemmKernel = cutlass::gemm::device::GemmUniversalAdapter<${operation_name}>;
|
||||
manifest.append(
|
||||
new ${gemm_kind}<GemmKernel>("${operation_name}"));
|
||||
}
|
||||
${compile_guard_end}
|
||||
"""
|
||||
|
||||
def emit_block_scale_epilogue_functor(self, operation):
|
||||
block_scaled_template = """
|
||||
${epilogue_functor}<
|
||||
${epi_vs},
|
||||
${element_d},
|
||||
${element_accumulator},
|
||||
${element_sfd},
|
||||
${layout_sfd},
|
||||
${element_c},
|
||||
${element_scalar}
|
||||
>
|
||||
"""
|
||||
block_scaled_values = {
|
||||
"epi_vs": str(operation.ScaleFactorVectorSize),
|
||||
"element_d": str(DataTypeTag[operation.D.element]),
|
||||
"element_sfd": str(DataTypeTag[operation.ScaleFactorD.element]),
|
||||
"layout_sfd": LayoutTag[operation.ScaleFactorD.layout],
|
||||
"epilogue_functor": EpilogueFunctor3xTag[
|
||||
EpilogueFunctor3x.LinearCombinationBlockScaleFactor
|
||||
],
|
||||
"element_accumulator": str(DataTypeTag[operation.accumulator_type()]),
|
||||
"element_scalar": str(DataTypeTag[operation.accumulator_type()]),
|
||||
"element_c": str(DataTypeTag[operation.C.element]),
|
||||
}
|
||||
return SubstituteTemplate(block_scaled_template, block_scaled_values)
|
||||
|
||||
@staticmethod
|
||||
def pointerize_if_grouped(operation, layout):
|
||||
return layout if not is_grouped(operation.gemm_kind) else layout + "* "
|
||||
|
||||
@staticmethod
|
||||
def problem_shape(operation):
|
||||
gemm_shape_type = "cute::Shape<int,int,int,int>"
|
||||
grouped_gemm_shape_type = "cute::Shape<int,int,int>"
|
||||
grouped_gemm_shape_type = (
|
||||
"cutlass::gemm::GroupProblemShape<" + grouped_gemm_shape_type + ">"
|
||||
)
|
||||
|
||||
return (
|
||||
gemm_shape_type
|
||||
if not is_grouped(operation.gemm_kind)
|
||||
else grouped_gemm_shape_type
|
||||
)
|
||||
|
||||
def emit(self, operation):
|
||||
"""Given a gem operation, emits a template definition of the operation"""
|
||||
|
||||
opcode_class_main = operation.tile_description.math_instruction.opcode_class
|
||||
opcode_class_epi = opcode_class_main
|
||||
|
||||
tile_shape = operation.tile_description.tile_shape
|
||||
instruction_shape = (
|
||||
operation.tile_description.math_instruction.instruction_shape
|
||||
)
|
||||
cluster_m = operation.tile_description.cluster_shape[0]
|
||||
cluster_n = operation.tile_description.cluster_shape[1]
|
||||
|
||||
tile_shape_m, tile_shape_n, tile_shape_k = tile_shape
|
||||
|
||||
# account for static/dynamic cluster shapes
|
||||
cta_m = tile_shape[0] // cluster_m if cluster_m > 0 else tile_shape[0]
|
||||
cta_n = tile_shape[1] // cluster_n if cluster_n > 0 else tile_shape[1]
|
||||
|
||||
# Shape passed to epilogue builder
|
||||
is_sm100_kernel = operation.arch == 100
|
||||
if is_sm100_kernel:
|
||||
cta_m_per_mma_instruction = (
|
||||
2 if "2sm" in operation.procedural_name() else 1
|
||||
)
|
||||
if cluster_m <= 0:
|
||||
cta_m = cta_m // cta_m_per_mma_instruction
|
||||
|
||||
if opcode_class_main in [
|
||||
OpcodeClass.TensorOp,
|
||||
OpcodeClass.BlockScaledTensorOp,
|
||||
]:
|
||||
tile_shape_m = instruction_shape[0]
|
||||
tile_shape_n = instruction_shape[1]
|
||||
|
||||
# stage count set to zero indicates builder automatic stage selection
|
||||
if operation.tile_description.stages > 0:
|
||||
stage_count_string = f"cutlass::gemm::collective::StageCount<\
|
||||
{str(operation.tile_description.stages)}>"
|
||||
else:
|
||||
stage_count_string = (
|
||||
f"cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(\
|
||||
sizeof(typename {str(operation.procedural_name())}_epilogue::SharedStorage))>"
|
||||
)
|
||||
if self.device_type == "xpu":
|
||||
stage_count_string = "cutlass::gemm::collective::StageCountAuto"
|
||||
|
||||
epi_tile_mn = "cutlass::epilogue::collective::EpilogueTileAuto"
|
||||
|
||||
(
|
||||
instance_layout_A,
|
||||
instance_layout_B,
|
||||
instance_layout_C,
|
||||
instance_layout_D,
|
||||
) = (
|
||||
operation.A.layout,
|
||||
operation.B.layout,
|
||||
operation.C.layout,
|
||||
operation.D.layout,
|
||||
)
|
||||
|
||||
# 3.0 profiler integration only supports trivial epilogues for now
|
||||
epilogue_vector_length = 1
|
||||
|
||||
# Support built-in epilogue functors or user-defined functions
|
||||
if isinstance(operation.epilogue_functor, enum.Enum):
|
||||
values = {
|
||||
"element_epilogue": str(DataTypeTag[operation.element_epilogue]),
|
||||
"epilogue_functor": EpilogueFunctor3xTag[
|
||||
operation.epilogue_functor
|
||||
],
|
||||
}
|
||||
epilogue_functor = SubstituteTemplate(
|
||||
self.builtin_epilogue_functor_template, values
|
||||
)
|
||||
|
||||
if (
|
||||
is_block_scaled(operation.gemm_kind)
|
||||
and operation.ScaleFactorD.element != DataType.void
|
||||
):
|
||||
epilogue_functor = self.emit_block_scale_epilogue_functor(operation)
|
||||
else:
|
||||
epilogue_functor = self.epilogue_functor.emit_declaration()
|
||||
|
||||
if (
|
||||
is_block_scaled(operation.gemm_kind)
|
||||
and operation.ScaleFactorD.element != DataType.void
|
||||
):
|
||||
epilogue_functor = self.emit_block_scale_epilogue_functor(operation)
|
||||
|
||||
#
|
||||
# Cutlass3x complex kernels' ElementA(B) is a tuple in collective mainloop builder,
|
||||
# e.g. cute::tuple<Element, Transform>, Transform : cute::identity / cute::conjugate.
|
||||
element_a = (
|
||||
DataTypeTag[operation.A.element]
|
||||
if not operation.is_complex()
|
||||
else f"cute::tuple<{str(DataTypeTag[operation.A.element])},\
|
||||
{str(ComplexTransformTag3x[operation.A.complex_transform])}>"
|
||||
)
|
||||
element_b = (
|
||||
DataTypeTag[operation.B.element]
|
||||
if not operation.is_complex()
|
||||
else f"cute::tuple<{str(DataTypeTag[operation.B.element])},\
|
||||
{str(ComplexTransformTag3x[operation.B.complex_transform])}>"
|
||||
)
|
||||
epilogue_schedule_type = EpilogueScheduleTag[operation.epilogue_schedule]
|
||||
|
||||
if opcode_class_main == OpcodeClass.BlockScaledTensorOp:
|
||||
is_no_smem_epilogue = operation.epilogue_schedule in [
|
||||
EpilogueScheduleType.NoSmemWarpSpecialized1Sm,
|
||||
EpilogueScheduleType.NoSmemWarpSpecialized2Sm,
|
||||
]
|
||||
grouped = is_grouped(operation.gemm_kind)
|
||||
if cta_n == 256 and operation.kernel_schedule == to_grouped_schedule(
|
||||
KernelScheduleType.Nvf4TmaWarpSpecialized1SmSm100, grouped
|
||||
):
|
||||
epi_tile_mn = "cute::Shape<cute::_128,cute::_64>"
|
||||
if not is_no_smem_epilogue:
|
||||
epilogue_schedule_type = EpilogueScheduleTag[
|
||||
to_grouped_schedule(
|
||||
EpilogueScheduleType.TmaWarpSpecialized1Sm, grouped
|
||||
)
|
||||
]
|
||||
if cta_n == 256 and operation.kernel_schedule == to_grouped_schedule(
|
||||
KernelScheduleType.Nvf4TmaWarpSpecialized2SmSm100, grouped
|
||||
):
|
||||
epi_tile_mn = "cute::Shape<cute::_128,cute::_64>"
|
||||
if not is_no_smem_epilogue:
|
||||
epilogue_schedule_type = EpilogueScheduleTag[
|
||||
to_grouped_schedule(
|
||||
EpilogueScheduleType.TmaWarpSpecialized2Sm, grouped
|
||||
)
|
||||
]
|
||||
element_a = f"cute::tuple<{str(element_a)},{str(DataTypeTag[operation.ScaleFactorA])}>"
|
||||
element_b = f"cute::tuple<{str(element_b)},{str(DataTypeTag[operation.ScaleFactorB])}>"
|
||||
|
||||
operation_name_str = operation.procedural_name()
|
||||
layout_a_str = LayoutTag[instance_layout_A]
|
||||
layout_b_str = LayoutTag[instance_layout_B]
|
||||
mixed_dtype_prepare_code = ""
|
||||
if operation.mixed_input_mode is not None:
|
||||
A_dtype = operation.A.element
|
||||
B_dtype = operation.B.element
|
||||
A_dtype_bits = DataTypeSize[A_dtype]
|
||||
B_dtype_bits = DataTypeSize[B_dtype]
|
||||
is_A_dtype_narrow = A_dtype_bits < B_dtype_bits
|
||||
if is_A_dtype_narrow:
|
||||
narrow_dtype, wide_dtype = (A_dtype, B_dtype)
|
||||
narrow_dtype_bits, wide_dtype_bits = (A_dtype_bits, B_dtype_bits)
|
||||
else:
|
||||
narrow_dtype, wide_dtype = (B_dtype, A_dtype)
|
||||
narrow_dtype_bits, wide_dtype_bits = (B_dtype_bits, A_dtype_bits)
|
||||
|
||||
narrow_tag = DataTypeTag[narrow_dtype]
|
||||
wide_tag = DataTypeTag[wide_dtype]
|
||||
scale_tag = DataTypeTag[wide_dtype]
|
||||
zero_tag = DataTypeTag[wide_dtype]
|
||||
|
||||
do_shuffle = False
|
||||
value_shuffle_str = ""
|
||||
if narrow_dtype_bits == 4 and wide_dtype_bits == 16:
|
||||
value_shuffle_str = "cute::Layout<cute::Shape<cute::_2,cute::_4>, \
|
||||
cute::Stride<cute::_4,cute::_1>>"
|
||||
do_shuffle = True
|
||||
if narrow_dtype_bits == 8 and wide_dtype_bits == 16:
|
||||
value_shuffle_str = "cute::Layout<cute::Shape<cute::_2,cute::_2>, \
|
||||
cute::Stride<cute::_2,cute::_1>>"
|
||||
do_shuffle = True
|
||||
do_shuffle = operation.mixed_input_shuffle and do_shuffle
|
||||
|
||||
if do_shuffle:
|
||||
if is_A_dtype_narrow:
|
||||
stride_narrow_str = (
|
||||
f"cutlass::detail::TagToStrideA_t<{layout_a_str}>"
|
||||
)
|
||||
layout_a_str = f"{operation_name_str}_LayoutNarrowReordered"
|
||||
else:
|
||||
stride_narrow_str = (
|
||||
f"cutlass::detail::TagToStrideB_t<{layout_b_str}>"
|
||||
)
|
||||
layout_b_str = f"{operation_name_str}_LayoutNarrowReordered"
|
||||
# The {operation_name_str}_ prefixs in mixed_dtype_prepare_code and
|
||||
# layout_{a, b}_str are to prevent errors in Windows platform unity build
|
||||
mixed_dtype_prepare_code = f"""
|
||||
using {operation_name_str}_StrideNarrow = {stride_narrow_str};
|
||||
using {operation_name_str}_ValueShuffle = {value_shuffle_str};
|
||||
static constexpr int {operation_name_str}_NumShuffleAtoms = 1;
|
||||
using {operation_name_str}_MmaAtomShape = \
|
||||
cute::Layout<cute::Shape<cute::_1, cute::Int<{operation_name_str}_NumShuffleAtoms>>>;
|
||||
using {operation_name_str}_LayoutAtomQuant = \
|
||||
decltype(cutlass::compute_memory_reordering_atom<{wide_tag}, {operation_name_str}_MmaAtomShape, \
|
||||
{operation_name_str}_ValueShuffle>());
|
||||
using {operation_name_str}_LayoutNarrowReordered = \
|
||||
decltype(cute::tile_to_shape({operation_name_str}_LayoutAtomQuant{{}}, \
|
||||
cute::Layout<cute::Shape<int,int,int>, {operation_name_str}_StrideNarrow>{{}}));
|
||||
"""
|
||||
|
||||
mixed_input_modes_to_element = {
|
||||
MixedInputMode.ConvertOnly: narrow_tag,
|
||||
MixedInputMode.ScaleOnly: f"cute::tuple<{narrow_tag}, {scale_tag}>",
|
||||
MixedInputMode.ScaleWithZeroPoint: f"cute::tuple<{narrow_tag}, {scale_tag}, {zero_tag}>",
|
||||
}
|
||||
narrow_element = mixed_input_modes_to_element.get(
|
||||
operation.mixed_input_mode, narrow_tag
|
||||
)
|
||||
|
||||
if narrow_dtype == DataType.s4 and (
|
||||
wide_dtype == DataType.e4m3 or wide_dtype == DataType.e5m2
|
||||
):
|
||||
narrow_element = (
|
||||
f"cute::tuple<{narrow_tag}, cutlass::Array<{scale_tag}, 8>>"
|
||||
)
|
||||
|
||||
if is_A_dtype_narrow:
|
||||
element_a = narrow_element
|
||||
else:
|
||||
element_b = narrow_element
|
||||
|
||||
if self.evt_name:
|
||||
epilogue_functor = self.evt_name
|
||||
|
||||
if self.device_type == "xpu":
|
||||
arch = f"cutlass::arch::Xe{operation.arch}"
|
||||
else:
|
||||
arch = f"cutlass::arch::Sm{operation.arch}"
|
||||
|
||||
values = {
|
||||
"operation_name": operation_name_str,
|
||||
"operation_suffix": self.operation_suffix,
|
||||
"problem_shape": self.problem_shape(operation),
|
||||
"element_a": element_a,
|
||||
"layout_a": self.pointerize_if_grouped(operation, layout_a_str),
|
||||
"element_b": element_b,
|
||||
"layout_b": self.pointerize_if_grouped(operation, layout_b_str),
|
||||
"element_c": DataTypeTag[operation.C.element],
|
||||
"layout_c": self.pointerize_if_grouped(
|
||||
operation, LayoutTag[instance_layout_C]
|
||||
),
|
||||
"element_d": DataTypeTag[operation.D.element],
|
||||
"layout_d": self.pointerize_if_grouped(
|
||||
operation, LayoutTag[instance_layout_D]
|
||||
),
|
||||
"element_accumulator": DataTypeTag[operation.accumulator_type()],
|
||||
"opcode_class_main": OpcodeClassTag[opcode_class_main],
|
||||
"opcode_class_epi": OpcodeClassTag[opcode_class_epi],
|
||||
"arch": arch,
|
||||
"tile_shape_m": str(tile_shape_m),
|
||||
"tile_shape_n": str(tile_shape_n),
|
||||
"tile_shape_k": str(tile_shape_k),
|
||||
"cluster_shape_m": "cute::_"
|
||||
+ str(operation.tile_description.cluster_shape[0])
|
||||
if operation.tile_description.cluster_shape[0] > 0
|
||||
else "int",
|
||||
"cluster_shape_n": "cute::_"
|
||||
+ str(operation.tile_description.cluster_shape[1])
|
||||
if operation.tile_description.cluster_shape[1] > 0
|
||||
else "int",
|
||||
"cluster_shape_k": "cute::_"
|
||||
+ str(operation.tile_description.cluster_shape[2])
|
||||
if operation.tile_description.cluster_shape[2] > 0
|
||||
else "int",
|
||||
"instruction_shape_m": str(instruction_shape[0]),
|
||||
"instruction_shape_n": str(instruction_shape[1]),
|
||||
"instruction_shape_k": str(instruction_shape[2]),
|
||||
"kernel_schedule": str(KernelScheduleTag[operation.kernel_schedule]),
|
||||
"epilogue_schedule": str(epilogue_schedule_type),
|
||||
"epi_tile_mn": epi_tile_mn,
|
||||
"epilogue_functor": epilogue_functor,
|
||||
"stages": stage_count_string,
|
||||
"align_a": str(operation.A.alignment),
|
||||
"align_b": str(operation.B.alignment),
|
||||
"align_c": str(operation.C.alignment),
|
||||
"align_d": str(operation.D.alignment),
|
||||
"transform_a": ComplexTransformTag[operation.A.complex_transform],
|
||||
"transform_b": ComplexTransformTag[operation.B.complex_transform],
|
||||
"math_operation": MathOperationTag[
|
||||
operation.tile_description.math_instruction.math_operation
|
||||
],
|
||||
"epilogue_vector_length": str(epilogue_vector_length),
|
||||
"element_epilogue": str(DataTypeTag[operation.element_epilogue]),
|
||||
"tile_scheduler": str(TileSchedulerTag[operation.tile_scheduler]),
|
||||
"mixed_dtype_prepare_code": mixed_dtype_prepare_code,
|
||||
}
|
||||
|
||||
return SubstituteTemplate(self.gemm_template, values)
|
||||
+327
@@ -0,0 +1,327 @@
|
||||
import itertools
|
||||
from collections.abc import Generator, Iterable, Iterator, Sequence
|
||||
from contextlib import contextmanager
|
||||
from os import linesep
|
||||
from typing import Any
|
||||
|
||||
import sympy
|
||||
|
||||
import torch
|
||||
import torch._inductor.virtualized as virtualized
|
||||
from torch._inductor.ir import ComputedBuffer, Pointwise
|
||||
from torch._inductor.ops_handler import DefaultHandler, WrapperHandler
|
||||
from torch._inductor.scheduler import BaseSchedulerNode
|
||||
from torch._inductor.utils import DelayReplaceLine, IndentedBuffer, OrderedSet
|
||||
from torch._inductor.virtualized import OpsValue
|
||||
|
||||
from ...virtualized import V
|
||||
|
||||
|
||||
_ACCUMULATOR_ARG_NAME = "accum"
|
||||
|
||||
|
||||
def scaled_mm_evt(
|
||||
scale_A_name: str, scale_B_name: str, bias_name: str | None, output_name: str
|
||||
) -> tuple[list[str], dict[str, Any], str]:
|
||||
evt_read_names = [scale_A_name, scale_B_name]
|
||||
var_name_to_buffer_name = {n: n for n in [scale_A_name, scale_B_name]}
|
||||
var_name_to_buffer_name["D"] = output_name
|
||||
var_name_to_buffer_name[_ACCUMULATOR_ARG_NAME] = output_name
|
||||
expr = f"accum * {scale_A_name} * {scale_B_name}{linesep}"
|
||||
if bias_name:
|
||||
expr = f"({expr}) + {bias_name}"
|
||||
evt_read_names.append(bias_name)
|
||||
var_name_to_buffer_name[bias_name] = bias_name
|
||||
|
||||
evt_py_code = f"def fn(accum, {','.join(evt_read_names)}):{linesep}\
|
||||
D = {expr}{linesep}\
|
||||
return D{linesep}"
|
||||
|
||||
return evt_read_names, var_name_to_buffer_name, evt_py_code
|
||||
|
||||
|
||||
class CutlassEVTOpsMixIn:
|
||||
@staticmethod
|
||||
def _infix_bin_op(op: str, a: str, b: str) -> str:
|
||||
return f"{a} {op} {b}"
|
||||
|
||||
@staticmethod
|
||||
def _prefix_bin_op(op: str, a: str, b: str) -> str:
|
||||
return f"{op}({a}, {b})"
|
||||
|
||||
@staticmethod
|
||||
def _prefix_un_op(op: str, a: str) -> str:
|
||||
return f"{op}({a})"
|
||||
|
||||
@staticmethod
|
||||
def to_dtype(
|
||||
x: str,
|
||||
dtype: Any,
|
||||
src_dtype: torch.dtype | None = None,
|
||||
use_compute_types: bool = False,
|
||||
) -> str:
|
||||
return x
|
||||
|
||||
@staticmethod
|
||||
def constant(value: Any, dtype: Any) -> str:
|
||||
raise NotImplementedError
|
||||
|
||||
@staticmethod
|
||||
def mul(x0: str, x1: str) -> str:
|
||||
return CutlassEVTOpsMixIn._infix_bin_op("*", x0, x1)
|
||||
|
||||
@staticmethod
|
||||
def truediv(x0: str, x1: str) -> str:
|
||||
return CutlassEVTOpsMixIn._infix_bin_op("/", x0, x1)
|
||||
|
||||
@staticmethod
|
||||
def ge(x0: str, x1: str) -> str:
|
||||
raise NotImplementedError
|
||||
|
||||
@staticmethod
|
||||
def add(x0: str, x1: str) -> str:
|
||||
return CutlassEVTOpsMixIn._infix_bin_op("+", x0, x1)
|
||||
|
||||
@staticmethod
|
||||
def relu(x0: str) -> str:
|
||||
return CutlassEVTOpsMixIn._prefix_un_op("relu", x0)
|
||||
|
||||
@staticmethod
|
||||
def sigmoid(x0: str) -> str:
|
||||
return CutlassEVTOpsMixIn._prefix_un_op("sigmoid", x0)
|
||||
|
||||
@staticmethod
|
||||
def sub(x0: str, x1: str) -> str:
|
||||
return CutlassEVTOpsMixIn._infix_bin_op("-", x0, x1)
|
||||
|
||||
@staticmethod
|
||||
def tanh(x0: str) -> str:
|
||||
return CutlassEVTOpsMixIn._prefix_un_op("tanh", x0)
|
||||
|
||||
@staticmethod
|
||||
def exp(x0: str) -> str:
|
||||
return CutlassEVTOpsMixIn._prefix_un_op("exp", x0)
|
||||
|
||||
|
||||
class MockCutlassHandler(CutlassEVTOpsMixIn, WrapperHandler):
|
||||
"""Passthrough handler for cutlass ops, used for running epilogue nodes for memory planning"""
|
||||
|
||||
|
||||
class _AssignmentFormatter(DefaultHandler):
|
||||
def __init__(self, parent_handler: "CutlassEVTCodegen"):
|
||||
self.parent_handler = parent_handler
|
||||
|
||||
def _default(self, name: str, args: tuple[Any, ...], kwargs: dict[str, Any]) -> Any:
|
||||
# Handle op dispatch here
|
||||
if hasattr(self.parent_handler, name):
|
||||
fn = getattr(self.parent_handler, name)
|
||||
line = fn(*args, **kwargs)
|
||||
if name in ("load", "store"):
|
||||
return OpsValue(line)
|
||||
else:
|
||||
var = self.parent_handler._tmp_var()
|
||||
line = DelayReplaceLine(
|
||||
var,
|
||||
lambda: "D"
|
||||
if var == self.parent_handler.last_stored_var_name
|
||||
else var,
|
||||
f"{var} = {line}",
|
||||
)
|
||||
self.parent_handler.body.writeline(line)
|
||||
return OpsValue(var)
|
||||
else:
|
||||
raise NotImplementedError(name)
|
||||
|
||||
|
||||
class CutlassEVTCodegen(CutlassEVTOpsMixIn):
|
||||
"""
|
||||
Notes:
|
||||
* Used by CUTLASSGemmTemplate.
|
||||
* This class should not be instantiated by users, it is intended to be used
|
||||
by calling CutlassEVTCodegen.ir_to_evt_python_code(...)
|
||||
which instantiates this class as an ops handler for virtualized.V.ops.[op-name]
|
||||
* Extend this with more _op_<whatever> nodes to add support for new pointwise operations.
|
||||
"""
|
||||
|
||||
def __init__(self, accumulator_node_name: str, removed_buffers: OrderedSet[str]):
|
||||
"""
|
||||
|
||||
Initializes a CutlassEVTEpilogueArgumentFormatter object. Do not instantiate directly.
|
||||
Use the CutlassEVTCodegen.ir_to_evt_python_code static method.
|
||||
|
||||
Args:
|
||||
accumulator_node_name: The name of the accumulator node which should contain
|
||||
the Matmul result before fusion according to the IR graph.
|
||||
epilogue_nodes: The list of scheduler nodes to be fused into the epilogue
|
||||
"""
|
||||
self.accumulator_node_name: str = accumulator_node_name #
|
||||
self.body: IndentedBuffer = IndentedBuffer(1) # The body buffer for codegen
|
||||
self.var_counter: Iterator[int] = itertools.count()
|
||||
self.store_name_to_value: dict[str, OpsValue] = (
|
||||
dict()
|
||||
) # Aliases for subexpression functors
|
||||
self.reads: OrderedSet[str] = OrderedSet([])
|
||||
# Used for creating example tensors
|
||||
self.var_name_to_buffer_name: dict[str, str] = {
|
||||
_ACCUMULATOR_ARG_NAME: accumulator_node_name
|
||||
}
|
||||
self.removed_buffers: OrderedSet[str] = removed_buffers
|
||||
self.cur_node: ComputedBuffer | None = None
|
||||
self.name_to_buffer = V.graph.name_to_buffer | V.graph.graph_inputs
|
||||
for name in V.graph.constants:
|
||||
# pyrefly: ignore [unsupported-operation]
|
||||
self.name_to_buffer[name] = V.graph.add_tensor_constant(
|
||||
V.graph.constants[name], name
|
||||
)
|
||||
self.is_D_assigned = False
|
||||
self.D_var_name = None
|
||||
|
||||
if accumulator_node_name not in removed_buffers:
|
||||
# cannot return accumulator directly, so alias it
|
||||
var = self._tmp_var()
|
||||
self.body.writeline(f"{var} = {_ACCUMULATOR_ARG_NAME}")
|
||||
self.store(accumulator_node_name, value=OpsValue(var))
|
||||
|
||||
@staticmethod
|
||||
def ir_to_evt_python_code(
|
||||
cutlass_template_node_name: str,
|
||||
epilogue_nodes: list[BaseSchedulerNode],
|
||||
removed_buffers: OrderedSet[str],
|
||||
) -> tuple[list[str], list[str], dict[str, Any], str]:
|
||||
codegen = CutlassEVTCodegen(cutlass_template_node_name, removed_buffers)
|
||||
handler = _AssignmentFormatter(codegen)
|
||||
|
||||
with virtualized.V.set_ops_handler(handler):
|
||||
for s_node in epilogue_nodes:
|
||||
node = s_node.node
|
||||
assert isinstance(node, ComputedBuffer)
|
||||
with codegen.set_cur_node(node):
|
||||
index_vars = CutlassEVTCodegen.get_index_vars(node)
|
||||
node.get_store_function()(index_vars)
|
||||
|
||||
codegen.finalize()
|
||||
|
||||
return (
|
||||
codegen.get_reads(),
|
||||
codegen.get_writes(),
|
||||
codegen.get_renames(),
|
||||
codegen.get_value(),
|
||||
)
|
||||
|
||||
def get_value(self) -> str:
|
||||
return linesep.join(
|
||||
[
|
||||
self._render_input_signature(),
|
||||
self.body.getvalue(),
|
||||
self._render_return_statement(),
|
||||
]
|
||||
)
|
||||
|
||||
def finalize(self) -> None:
|
||||
# Rename the last store to D
|
||||
# no other code references this store
|
||||
# to workaround https://github.com/NVIDIA/cutlass/issues/2288
|
||||
# Note: the delayed line will automatically rewrite the last assignment to
|
||||
# be to D
|
||||
buffer_name = self.var_name_to_buffer_name[self.last_stored_var_name]
|
||||
self.var_name_to_buffer_name.pop(self.last_stored_var_name)
|
||||
self.var_name_to_buffer_name["D"] = buffer_name
|
||||
self.store_name_to_value[buffer_name] = OpsValue("D")
|
||||
|
||||
@contextmanager
|
||||
def set_cur_node(self, node: ComputedBuffer) -> Generator[None, Any, Any]:
|
||||
prev_node = self.cur_node
|
||||
try:
|
||||
self.cur_node = node
|
||||
yield
|
||||
finally:
|
||||
self.cur_node = prev_node
|
||||
|
||||
def get_renames(self) -> dict[str, str]:
|
||||
return dict(self.var_name_to_buffer_name)
|
||||
|
||||
def get_reads(self) -> list[str]:
|
||||
return list(self.reads.difference(self.store_name_to_value.keys()))
|
||||
|
||||
def get_writes(self) -> list[str]:
|
||||
return list(self.store_name_to_value.keys())
|
||||
|
||||
def load(self, name: str, index: Any) -> str:
|
||||
self._check_indexing(name, index)
|
||||
if name in self.store_name_to_value:
|
||||
return self.store_name_to_value[name].value
|
||||
elif name == self.accumulator_node_name:
|
||||
return _ACCUMULATOR_ARG_NAME
|
||||
else:
|
||||
self.reads.add(name)
|
||||
self.var_name_to_buffer_name[name] = name
|
||||
return name
|
||||
|
||||
def store(
|
||||
self, name: Any, index: Any = None, value: Any = None, mode: Any = None
|
||||
) -> None:
|
||||
if name not in self.removed_buffers:
|
||||
if index:
|
||||
self._check_indexing(name, index)
|
||||
assert value.value != _ACCUMULATOR_ARG_NAME, (
|
||||
"Cannot store accumulator arg name"
|
||||
)
|
||||
self.var_name_to_buffer_name[value.value] = name
|
||||
self.store_name_to_value[name] = value
|
||||
self.last_stored_var_name = value.value
|
||||
return None
|
||||
|
||||
def _get_cur_node(self) -> ComputedBuffer:
|
||||
assert self.cur_node
|
||||
return self.cur_node
|
||||
|
||||
@staticmethod
|
||||
def get_index_vars(node: ComputedBuffer) -> Sequence[sympy.Expr]:
|
||||
data = node.data
|
||||
# TODO mlazos: relax this, cutlass supports reductions and other ops
|
||||
assert isinstance(data, Pointwise)
|
||||
return data._index(data.ranges)
|
||||
|
||||
def _get_current_index_vars(self) -> Sequence[sympy.Expr]:
|
||||
return self.get_index_vars(self._get_cur_node())
|
||||
|
||||
def _check_indexing(self, name: str, index: sympy.Expr) -> None:
|
||||
# We only support indexing that matches the layout today because
|
||||
# CUTLASS doesn't support arbitrary indexing
|
||||
buffer_name = (
|
||||
self.accumulator_node_name if name == _ACCUMULATOR_ARG_NAME else name
|
||||
)
|
||||
buffer = self.name_to_buffer[buffer_name]
|
||||
index_strides = V.graph.sizevars.stride_vars(
|
||||
index, self._get_current_index_vars()
|
||||
)
|
||||
stride = buffer.get_layout().stride
|
||||
if not self._stride_compatible(stride, index_strides):
|
||||
raise NotImplementedError(
|
||||
f"Unsupported indexing for {name} with index {index}, index strides {index_strides}, and layout stride {stride}"
|
||||
)
|
||||
|
||||
def _stride_compatible(
|
||||
self, left: Iterable[sympy.Expr], right: Iterable[sympy.Expr]
|
||||
) -> bool:
|
||||
return all(
|
||||
sympy.Eq(l, r) or sympy.Eq(l, 0) or sympy.Eq(r, 0)
|
||||
for l, r in (zip(left, right))
|
||||
)
|
||||
|
||||
def _render_input_signature(self) -> str:
|
||||
arguments = ", ".join(
|
||||
[_ACCUMULATOR_ARG_NAME]
|
||||
+ [name for name in self.reads if name != self.accumulator_node_name]
|
||||
)
|
||||
return f"def fn({arguments}):"
|
||||
|
||||
def _render_return_statement(self) -> str:
|
||||
return_vars = OrderedSet(
|
||||
op_v.value for op_v in self.store_name_to_value.values()
|
||||
)
|
||||
assert "D" in return_vars
|
||||
return f"return {', '.join(return_vars)}"
|
||||
|
||||
def _tmp_var(self) -> str:
|
||||
return f"tmp_{next(self.var_counter)}"
|
||||
+295
@@ -0,0 +1,295 @@
|
||||
# mypy: allow-untyped-defs
|
||||
import hashlib
|
||||
import logging
|
||||
from collections.abc import Sequence
|
||||
from typing import cast, TypeGuard
|
||||
|
||||
from torch._inductor.codegen.cutlass.python_evt import (
|
||||
CutlassEVTCodegen,
|
||||
MockCutlassHandler,
|
||||
)
|
||||
from torch._inductor.utils import Placeholder
|
||||
from torch.utils._ordered_set import OrderedSet
|
||||
|
||||
from ...._dynamo.utils import counters
|
||||
from ... import config
|
||||
from ...codecache import code_hash, get_path
|
||||
from ...ir import Buffer, ComputedBuffer, CUTLASSTemplateBuffer, Pointwise
|
||||
from ...scheduler import (
|
||||
BaseSchedulerNode,
|
||||
BaseScheduling,
|
||||
FusedSchedulerNode,
|
||||
SchedulerNode,
|
||||
WhyNoFuse,
|
||||
)
|
||||
from ...utils import get_fused_kernel_name, get_kernel_metadata, sympy_product
|
||||
from ...virtualized import V
|
||||
from ..common import BackendFeature, IndentedBuffer
|
||||
|
||||
|
||||
log = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class WhyNoFuseNames(WhyNoFuse):
|
||||
def __init__(self, name1: str, name2: str) -> None:
|
||||
self.name1 = name1
|
||||
self.name2 = name2
|
||||
|
||||
|
||||
class CUTLASSScheduling(BaseScheduling):
|
||||
"""
|
||||
Partial Scheduling implementation for cutlass C++ Kernels.
|
||||
This class is intended to be used in combination with TritonScheduling,
|
||||
and delegated to by CUDACombinedScheduling/XPUCombinedScheduling.
|
||||
|
||||
It handles fusion decisions and cutlass C++ specific template code generation.
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def get_backend_features(cls, device) -> OrderedSet[BackendFeature]:
|
||||
return OrderedSet()
|
||||
|
||||
def group_fn(self, sizes):
|
||||
return tuple(V.graph.sizevars.simplify(sympy_product(s)) for s in sizes)
|
||||
|
||||
@staticmethod
|
||||
def is_cutlass_template(node: BaseSchedulerNode) -> TypeGuard[SchedulerNode]:
|
||||
return isinstance(node, SchedulerNode) and isinstance(
|
||||
node.node, CUTLASSTemplateBuffer
|
||||
)
|
||||
|
||||
def is_cutlass_fused_template(self, node: BaseSchedulerNode) -> bool:
|
||||
return isinstance(node, FusedSchedulerNode) and self.is_cutlass_template(node)
|
||||
|
||||
def can_fuse_vertical(
|
||||
self, node1: BaseSchedulerNode, node2: BaseSchedulerNode
|
||||
) -> bool:
|
||||
if self.is_cutlass_template(node1) and isinstance(node2, BaseSchedulerNode):
|
||||
assert node1.node, "node1.node should not be None"
|
||||
return self._can_fuse_epilogue_impl(
|
||||
cast(CUTLASSTemplateBuffer, node1.node),
|
||||
[],
|
||||
node2, # type: ignore[arg-type]
|
||||
)
|
||||
elif self.is_cutlass_fused_template(node1) and isinstance(
|
||||
node2, BaseSchedulerNode
|
||||
):
|
||||
assert node1.node, "node1.node should not be None"
|
||||
assert node2.node, "node2.node should not be None"
|
||||
fnode1 = cast(FusedSchedulerNode, node1)
|
||||
return self._can_fuse_epilogue_impl(
|
||||
fnode1.get_template_node(), # type: ignore[arg-type]
|
||||
self._unwrap_epilogue_nodes(fnode1),
|
||||
node2, # type: ignore[arg-type]
|
||||
)
|
||||
|
||||
return False
|
||||
|
||||
def define_kernel(self, src_code: str, node_schedule) -> str:
|
||||
wrapper = V.graph.wrapper_code
|
||||
if src_code in wrapper.src_to_kernel:
|
||||
kernel_name = wrapper.src_to_kernel[src_code]
|
||||
else:
|
||||
fused_name = (
|
||||
get_fused_kernel_name(node_schedule, config.triton.descriptive_names)
|
||||
if config.triton.descriptive_names
|
||||
else ""
|
||||
)
|
||||
|
||||
# use the original src_code as the key
|
||||
kernel_hash = hashlib.sha256(src_code.encode("utf-8")).hexdigest()[:8]
|
||||
if fused_name == "fused":
|
||||
# no EVT kernel, use the original kernel name
|
||||
kernel_name = f"cutlass_{kernel_hash}"
|
||||
else:
|
||||
kernel_name = f"cutlass_{fused_name}_{kernel_hash}"
|
||||
wrapper.src_to_kernel[src_code] = kernel_name
|
||||
src_code = src_code.replace(str(Placeholder.KERNEL_NAME), kernel_name)
|
||||
|
||||
_, _, kernel_path = get_path(code_hash(src_code), "py")
|
||||
|
||||
compile_wrapper = IndentedBuffer()
|
||||
compile_wrapper.writeline(f"async_compile.{V.graph.device_type}(r'''")
|
||||
compile_wrapper.splice(src_code, strip=True)
|
||||
compile_wrapper.writeline(
|
||||
f"''', 'so', aot_compile={str(V.graph.aot_mode)})"
|
||||
)
|
||||
|
||||
metadata_comment = f"# kernel path: {kernel_path}"
|
||||
origins, detailed_origins = get_kernel_metadata(node_schedule, wrapper)
|
||||
metadata_comment += "\n" + origins + "\n" + detailed_origins
|
||||
wrapper.define_kernel(
|
||||
kernel_name, compile_wrapper.getvalue(), metadata_comment
|
||||
)
|
||||
return kernel_name
|
||||
|
||||
def codegen_template(
|
||||
self,
|
||||
template_node: BaseSchedulerNode,
|
||||
epilogue_nodes: Sequence[BaseSchedulerNode],
|
||||
prologue_nodes: Sequence[BaseSchedulerNode],
|
||||
):
|
||||
"""
|
||||
Codegen a cutlass template, possibly with fused epilogues
|
||||
"""
|
||||
counters["inductor"]["cutlass_epilogue_fusion_counter"] += len(epilogue_nodes)
|
||||
assert self.is_cutlass_template(template_node), (
|
||||
"Template node passed to CUTLASSScheduling.codegen_template must be a SchedulerNode that wraps a CUTLASSTemplateBuffer"
|
||||
)
|
||||
_, (_numel, rnumel) = template_node.group
|
||||
assert rnumel == 1
|
||||
ctb: CUTLASSTemplateBuffer = cast(CUTLASSTemplateBuffer, template_node.node)
|
||||
epilogue_ir_nodes: list[Buffer] = [n.node for n in epilogue_nodes] # type: ignore[misc]
|
||||
assert all(isinstance(n, ComputedBuffer) for n in epilogue_ir_nodes), (
|
||||
"Epilogue nodes must all be instances of ir.ComputedBuffer"
|
||||
)
|
||||
kernel, render = ctb.make_kernel_render( # type: ignore[misc]
|
||||
ctb, epilogue_nodes=epilogue_nodes
|
||||
)
|
||||
with kernel:
|
||||
for node in [template_node, *epilogue_nodes]:
|
||||
node.mark_run()
|
||||
|
||||
# typically there is a codegen pass which runs after mark_run
|
||||
# for this kernel we've already generated the C++ code, but we still
|
||||
# need to let the kernel know about loads/stores that occur in the fused
|
||||
# kernel for memory planning to properly optimize allocations
|
||||
ctb.emulate_store_fn()
|
||||
for node in epilogue_ir_nodes:
|
||||
with V.set_ops_handler(MockCutlassHandler(V.get_ops_handler())):
|
||||
assert isinstance(
|
||||
node, ComputedBuffer
|
||||
) # Not sure why we need to do this again
|
||||
node.get_store_function()(CutlassEVTCodegen.get_index_vars(node))
|
||||
|
||||
with V.set_kernel_handler(kernel):
|
||||
src_code = render()
|
||||
node_schedule = [template_node, *epilogue_nodes]
|
||||
kernel_name = self.define_kernel(src_code, node_schedule)
|
||||
|
||||
# debug printing values of intermediate tensors
|
||||
_, call_args, arg_signatures, _ = kernel.args.python_argdefs()
|
||||
debug_printer_manager = V.graph.wrapper_code.debug_printer
|
||||
debug_printer_manager.set_printer_args(
|
||||
call_args, kernel_name, arg_signatures, kernel
|
||||
)
|
||||
with debug_printer_manager:
|
||||
self.codegen_comment(node_schedule, kernel_name)
|
||||
kernel.call_kernel(kernel_name, ctb)
|
||||
|
||||
V.graph.removed_buffers |= kernel.removed_buffers
|
||||
self.free_buffers_in_scheduler()
|
||||
|
||||
@staticmethod
|
||||
def _unwrap_epilogue_nodes(
|
||||
fused_node: FusedSchedulerNode,
|
||||
) -> list[BaseSchedulerNode]:
|
||||
nodes = fused_node.get_nodes()
|
||||
template_node = fused_node.get_template_node()
|
||||
assert all(n.node is not None for n in nodes), (
|
||||
"All epilogue nodes should have an IRNode"
|
||||
)
|
||||
# pyrefly: ignore [redundant-cast]
|
||||
return cast(
|
||||
list[BaseSchedulerNode], [n for n in nodes if n.node is not template_node]
|
||||
)
|
||||
|
||||
def _can_fuse_epilogue_impl(
|
||||
self,
|
||||
cutlass_template_buffer: CUTLASSTemplateBuffer,
|
||||
existing_epilogue_nodes: list[BaseSchedulerNode],
|
||||
node_to_fuse: BaseSchedulerNode,
|
||||
) -> bool:
|
||||
"""
|
||||
Check if the given node can be fused with the epilogue. At the moment, Kernels
|
||||
support fusion with Pointwise operations, wrapped in (named) ComputedBuffer nodes.
|
||||
|
||||
Args:
|
||||
cutlass_template_buffer : A CUTLASSTemplateBuffer object representing the CUTLASS template and it's result buffer
|
||||
existing_epilogue_nodes : List[SchedulerNode]: The list of already fused epilogue nodes.
|
||||
node_to_fuse: The SchedulerNode node to be checked if it can be fused with the epilogue.
|
||||
Returns:
|
||||
- bool: True if the given node can be fused with the epilogue, False otherwise.
|
||||
|
||||
"""
|
||||
why = WhyNoFuseNames(
|
||||
cutlass_template_buffer.get_name(), node_to_fuse.get_name()
|
||||
)
|
||||
|
||||
scheduler_nodes_to_fuse = node_to_fuse.get_nodes()
|
||||
|
||||
assert isinstance(cutlass_template_buffer, CUTLASSTemplateBuffer)
|
||||
|
||||
# Checks on constituent nodes
|
||||
for s_node in scheduler_nodes_to_fuse:
|
||||
node = s_node.node
|
||||
|
||||
if not isinstance(node, ComputedBuffer):
|
||||
why(f"{node} is not a ComputedBuffer")
|
||||
return False
|
||||
elif not isinstance(node.data, Pointwise):
|
||||
why(f"{node} is not a Pointwise op")
|
||||
return False
|
||||
elif not node.get_computed_buffer_name(): # type: ignore[attr-defined]
|
||||
why(f"{node} does not have a computed buffer name")
|
||||
return False
|
||||
|
||||
name = node.get_computed_buffer_name() # type: ignore[attr-defined]
|
||||
# dtype can differ, and strides can differ as long as they are broadcastable
|
||||
if node.get_size() != cutlass_template_buffer.get_size():
|
||||
why(
|
||||
f"{name}'s size: {node.get_size()} differs from {cutlass_template_buffer.get_name()}'s \
|
||||
size: {cutlass_template_buffer.get_size()}"
|
||||
)
|
||||
return False
|
||||
|
||||
assert len(
|
||||
existing_epilogue_nodes
|
||||
) or cutlass_template_buffer.get_name() in OrderedSet(
|
||||
[rd.name for rd in node_to_fuse.read_writes.reads]
|
||||
), "First epilogue node must read from cutlass template buffer"
|
||||
|
||||
if node_to_fuse.has_aliasing_or_mutation():
|
||||
why(f"{node_to_fuse.get_name()} has aliasing or mutation")
|
||||
return False
|
||||
elif node_to_fuse.is_reduction():
|
||||
why(
|
||||
f"{node_to_fuse.get_name()} is a reduction which is not yet supported by EVT"
|
||||
)
|
||||
return False
|
||||
elif (
|
||||
not config.cutlass.cutlass_epilogue_fusion_enabled
|
||||
or not config.epilogue_fusion
|
||||
):
|
||||
why("cutlass epilogue fusion is not enabled")
|
||||
return False
|
||||
elif not cutlass_template_buffer.supports_epilogue_fusion:
|
||||
why("epilogue fusion is only supported for TMA-enabled gemm ops")
|
||||
return False
|
||||
|
||||
try:
|
||||
from torch._inductor.codegen.cutlass.python_evt import CutlassEVTCodegen
|
||||
|
||||
CutlassEVTCodegen.ir_to_evt_python_code(
|
||||
cutlass_template_buffer.get_name(),
|
||||
existing_epilogue_nodes + list(node_to_fuse.get_nodes()),
|
||||
OrderedSet(),
|
||||
)
|
||||
|
||||
except NotImplementedError as e:
|
||||
not_implemented_op = str(e)
|
||||
if not_implemented_op.startswith("_op_"):
|
||||
not_implemented_op = not_implemented_op[4:]
|
||||
why(
|
||||
f"Cannot fuse epilogue node {node_to_fuse} into {cutlass_template_buffer.name}, \
|
||||
likely due to unsupported operation: {not_implemented_op}" # noqa: G004, B950
|
||||
)
|
||||
return False
|
||||
else: # Likely due to unsupported dtype.
|
||||
why(
|
||||
f"Cannot fuse epilogue node {node_to_fuse} into {cutlass_template_buffer.name}. \
|
||||
Reason: {not_implemented_op}" # noqa: G004, B950
|
||||
)
|
||||
return False
|
||||
|
||||
return True
|
||||
+507
@@ -0,0 +1,507 @@
|
||||
# mypy: allow-untyped-defs
|
||||
import functools
|
||||
import json
|
||||
from enum import Enum
|
||||
from typing import Any, Optional
|
||||
|
||||
from torch._inductor.codegen.cutlass.utils import try_import_cutlass
|
||||
|
||||
|
||||
class CUTLASSOperationSerializer:
|
||||
"""Serializes and deserializes CUTLASS GEMM operations to/from JSON.
|
||||
|
||||
Handles GemmOperation objects and their nested components (TileDescription, TensorDescription).
|
||||
"""
|
||||
|
||||
# not used, but keeping in case we want to generalize the serializer
|
||||
_SUPPORTED_CLASSES: list[str] = [
|
||||
"GemmOperation",
|
||||
"GemmKind",
|
||||
"TileDescription",
|
||||
"TensorDescription",
|
||||
"DataType",
|
||||
"EpilogueFunctor",
|
||||
"EpilogueFunctor3x",
|
||||
"SwizzlingFunctor",
|
||||
"KernelScheduleType",
|
||||
"EpilogueScheduleType",
|
||||
"TileSchedulerType",
|
||||
]
|
||||
|
||||
@classmethod
|
||||
def serialize(cls, operation: "GemmOperation") -> str: # type: ignore[name-defined] # noqa: F821
|
||||
"""Serialize a GEMM operation to JSON string.
|
||||
|
||||
Args:
|
||||
operation: GemmOperation object
|
||||
|
||||
Returns:
|
||||
str: JSON string representation of the operation
|
||||
"""
|
||||
assert operation.__class__.__qualname__ == "GemmOperation", (
|
||||
"Only GemmOperation objects are supported via the main API"
|
||||
)
|
||||
return json.dumps(cls._gemm_operation_to_json(operation))
|
||||
|
||||
@classmethod
|
||||
def deserialize(cls, json_str: str) -> "GemmOperation": # type: ignore[name-defined] # noqa: F821
|
||||
"""Deserialize JSON string to a GEMM operation.
|
||||
|
||||
Args:
|
||||
json_str: JSON string of a GEMM operation
|
||||
|
||||
Returns:
|
||||
GemmOperation: Reconstructed operation
|
||||
"""
|
||||
json_dict = json.loads(json_str)
|
||||
return cls._json_to_gemm_operation(json_dict)
|
||||
|
||||
@classmethod
|
||||
def _gemm_operation_to_json(cls, operation: "GemmOperation") -> dict[str, Any]: # type: ignore[name-defined] # noqa: F821
|
||||
"""Convert GemmOperation to JSON-serializable dict.
|
||||
|
||||
Args:
|
||||
operation: GemmOperation object
|
||||
|
||||
Returns:
|
||||
dict: Dictionary representation
|
||||
"""
|
||||
from cutlass_library.library import TensorDescription
|
||||
|
||||
# Create the main dictionary with required and optional parameters
|
||||
result = {
|
||||
# Required parameters
|
||||
"gemm_kind": cls._enum_to_json(operation.gemm_kind),
|
||||
"arch": operation.arch,
|
||||
"tile_description": cls._tile_description_to_json(
|
||||
operation.tile_description
|
||||
),
|
||||
"A": cls._tensor_description_to_json(operation.A),
|
||||
"B": cls._tensor_description_to_json(operation.B),
|
||||
"C": cls._tensor_description_to_json(operation.C),
|
||||
"element_epilogue": cls._enum_to_json(operation.element_epilogue),
|
||||
# Optional parameters
|
||||
"epilogue_functor": cls._enum_to_json(operation.epilogue_functor),
|
||||
"swizzling_functor": cls._enum_to_json(operation.swizzling_functor),
|
||||
"D": cls._tensor_description_to_json(operation.D) if operation.D else None,
|
||||
"kernel_schedule": cls._enum_to_json(operation.kernel_schedule),
|
||||
"epilogue_schedule": cls._enum_to_json(operation.epilogue_schedule),
|
||||
"tile_scheduler": cls._enum_to_json(operation.tile_scheduler),
|
||||
}
|
||||
|
||||
# Process optional attributes
|
||||
optional_attrs = [
|
||||
"mixed_input_mode",
|
||||
"mixed_input_shuffle",
|
||||
"ScaleFactorA",
|
||||
"ScaleFactorB",
|
||||
"ScaleFactorD",
|
||||
"ScaleFactorMVecSize",
|
||||
"ScaleFactorNVecSize",
|
||||
"ScaleFactorKVecSize",
|
||||
"ScaleFactorVectorSize",
|
||||
"is_3x",
|
||||
]
|
||||
|
||||
for attr in optional_attrs:
|
||||
if not hasattr(operation, attr):
|
||||
continue
|
||||
|
||||
value = getattr(operation, attr)
|
||||
|
||||
if isinstance(value, TensorDescription):
|
||||
result[attr] = cls._tensor_description_to_json(value)
|
||||
elif isinstance(value, Enum):
|
||||
result[attr] = cls._enum_to_json(value)
|
||||
else:
|
||||
result[attr] = value
|
||||
|
||||
return result
|
||||
|
||||
@classmethod
|
||||
def _json_to_gemm_operation(cls, json_dict: dict[str, Any]) -> "GemmOperation": # type: ignore[name-defined] # noqa: F821
|
||||
"""Convert JSON dict to GemmOperation object.
|
||||
|
||||
Args:
|
||||
json_dict: Dictionary representation
|
||||
|
||||
Returns:
|
||||
GemmOperation: Reconstructed object
|
||||
"""
|
||||
from cutlass_library import DataType
|
||||
from cutlass_library.gemm_operation import GemmKind, GemmOperation
|
||||
from cutlass_library.library import (
|
||||
EpilogueFunctor,
|
||||
EpilogueFunctor3x,
|
||||
EpilogueScheduleType,
|
||||
KernelScheduleType,
|
||||
MixedInputMode,
|
||||
SwizzlingFunctor,
|
||||
TileSchedulerType,
|
||||
)
|
||||
|
||||
# Extract constructor parameters from the JSON dictionary
|
||||
gemm_kind = cls._json_to_enum(json_dict["gemm_kind"], GemmKind)
|
||||
arch = json_dict["arch"]
|
||||
tile_description = cls._json_to_tile_description(json_dict["tile_description"])
|
||||
A = cls._json_to_tensor_description(json_dict.get("A"), "A")
|
||||
B = cls._json_to_tensor_description(json_dict.get("B"), "B")
|
||||
C = cls._json_to_tensor_description(json_dict.get("C"), "C")
|
||||
element_epilogue = cls._json_to_enum(json_dict["element_epilogue"], DataType)
|
||||
|
||||
# Get optional parameters with defaults
|
||||
epilogue_functor = cls._json_to_enum(
|
||||
json_dict.get("epilogue_functor"),
|
||||
EpilogueFunctor3x if json_dict.get("is_3x") else EpilogueFunctor,
|
||||
)
|
||||
swizzling_functor = cls._json_to_enum(
|
||||
json_dict.get("swizzling_functor"), SwizzlingFunctor
|
||||
)
|
||||
D = cls._json_to_tensor_description(json_dict.get("D"), "D")
|
||||
kernel_schedule = cls._json_to_enum(
|
||||
json_dict.get("kernel_schedule"), KernelScheduleType
|
||||
)
|
||||
epilogue_schedule = cls._json_to_enum(
|
||||
json_dict.get("epilogue_schedule"), EpilogueScheduleType
|
||||
)
|
||||
tile_scheduler = cls._json_to_enum(
|
||||
json_dict.get("tile_scheduler"), TileSchedulerType
|
||||
)
|
||||
|
||||
mixed_input_mode = cls._json_to_enum(
|
||||
json_dict.get("mixed_input_mode"), MixedInputMode
|
||||
)
|
||||
mixed_input_shuffle = json_dict.get("mixed_input_shuffle", False)
|
||||
|
||||
# Scale factors
|
||||
ScaleFactorA = cls._json_to_enum(json_dict.get("ScaleFactorA"), DataType)
|
||||
ScaleFactorB = cls._json_to_enum(json_dict.get("ScaleFactorB"), DataType)
|
||||
|
||||
ScaleFactorD = None
|
||||
if "ScaleFactorD" in json_dict and "ScaleFactorVectorSize" in json_dict:
|
||||
ScaleFactorD = {
|
||||
"tensor": cls._json_to_tensor_description(
|
||||
json_dict.get("ScaleFactorD"), "ScaleFactorD"
|
||||
),
|
||||
"vector_size": json_dict.get("ScaleFactorVectorSize"),
|
||||
}
|
||||
|
||||
ScaleFactorMVecSize = json_dict.get("ScaleFactorMVecSize")
|
||||
ScaleFactorNVecSize = json_dict.get("ScaleFactorNVecSize")
|
||||
ScaleFactorKVecSize = json_dict.get("ScaleFactorKVecSize")
|
||||
|
||||
# Create the GemmOperation with the extracted parameters
|
||||
operation = GemmOperation(
|
||||
gemm_kind=gemm_kind,
|
||||
arch=arch,
|
||||
tile_description=tile_description,
|
||||
A=A,
|
||||
B=B,
|
||||
C=C,
|
||||
element_epilogue=element_epilogue,
|
||||
epilogue_functor=epilogue_functor,
|
||||
swizzling_functor=swizzling_functor,
|
||||
D=D,
|
||||
kernel_schedule=kernel_schedule,
|
||||
epilogue_schedule=epilogue_schedule,
|
||||
tile_scheduler=tile_scheduler,
|
||||
mixed_input_mode=mixed_input_mode,
|
||||
mixed_input_shuffle=mixed_input_shuffle,
|
||||
ScaleFactorA=ScaleFactorA,
|
||||
ScaleFactorB=ScaleFactorB,
|
||||
ScaleFactorD=ScaleFactorD,
|
||||
ScaleFactorMVecSize=ScaleFactorMVecSize,
|
||||
ScaleFactorNVecSize=ScaleFactorNVecSize,
|
||||
ScaleFactorKVecSize=ScaleFactorKVecSize,
|
||||
)
|
||||
|
||||
return operation
|
||||
|
||||
@classmethod
|
||||
@functools.lru_cache(None)
|
||||
def _tile_description_to_json(cls, tile_desc: "TileDescription") -> str: # type: ignore[name-defined] # noqa: F821
|
||||
"""
|
||||
Convert TileDescription to JSON string.
|
||||
|
||||
Args:
|
||||
tile_desc: TileDescription object
|
||||
|
||||
Returns:
|
||||
str: JSON string representation
|
||||
"""
|
||||
|
||||
# Create the main dictionary with field names matching TileDescription constructor parameters
|
||||
result = {
|
||||
"threadblock_shape": tile_desc.threadblock_shape,
|
||||
"stages": tile_desc.stages,
|
||||
"warp_count": tile_desc.warp_count,
|
||||
"math_instruction": cls._math_instruction_to_json(
|
||||
tile_desc.math_instruction
|
||||
),
|
||||
"min_compute": tile_desc.minimum_compute_capability, # Store as min_compute for constructor
|
||||
"max_compute": tile_desc.maximum_compute_capability, # Store as max_compute for constructor
|
||||
"cluster_shape": tile_desc.cluster_shape,
|
||||
"explicit_vector_sizes": tile_desc.explicit_vector_sizes,
|
||||
}
|
||||
|
||||
# Add tile_shape if it exists and differs from threadblock_shape
|
||||
if (
|
||||
hasattr(tile_desc, "tile_shape")
|
||||
and tile_desc.tile_shape != tile_desc.threadblock_shape
|
||||
):
|
||||
result["tile_shape"] = tile_desc.tile_shape
|
||||
|
||||
return json.dumps(result)
|
||||
|
||||
@classmethod
|
||||
@functools.lru_cache(None)
|
||||
def _json_to_tile_description(
|
||||
cls, json_dict: str | None
|
||||
) -> Optional["TileDescription"]: # type: ignore[name-defined] # noqa: F821
|
||||
"""
|
||||
Convert JSON dict to TileDescription object.
|
||||
|
||||
Args:
|
||||
json_dict: Dictionary representation
|
||||
|
||||
Returns:
|
||||
TileDescription: Reconstructed object
|
||||
"""
|
||||
if json_dict is None:
|
||||
return None
|
||||
|
||||
tile_dict = json.loads(json_dict)
|
||||
|
||||
from cutlass_library.library import TileDescription
|
||||
|
||||
math_instruction = cls._json_to_math_instruction(tile_dict["math_instruction"])
|
||||
|
||||
# Get compute capability values, checking both naming conventions
|
||||
min_compute = tile_dict.get(
|
||||
"min_compute", tile_dict.get("minimum_compute_capability")
|
||||
)
|
||||
max_compute = tile_dict.get(
|
||||
"max_compute", tile_dict.get("maximum_compute_capability")
|
||||
)
|
||||
|
||||
# Get cluster shape with default value
|
||||
cluster_shape = tile_dict.get("cluster_shape", [1, 1, 1])
|
||||
|
||||
# Create the TileDescription object
|
||||
tile_desc = TileDescription(
|
||||
threadblock_shape=tile_dict["threadblock_shape"],
|
||||
stages=tile_dict["stages"],
|
||||
warp_count=tile_dict["warp_count"],
|
||||
math_instruction=math_instruction,
|
||||
min_compute=min_compute,
|
||||
max_compute=max_compute,
|
||||
cluster_shape=cluster_shape,
|
||||
explicit_vector_sizes=tile_dict.get("explicit_vector_sizes"),
|
||||
)
|
||||
|
||||
# Set tile_shape if it exists and differs from threadblock_shape
|
||||
if (
|
||||
"tile_shape" in tile_dict
|
||||
and tile_dict["tile_shape"] != tile_dict["threadblock_shape"]
|
||||
):
|
||||
tile_desc.tile_shape = tile_dict["tile_shape"]
|
||||
|
||||
return tile_desc
|
||||
|
||||
@classmethod
|
||||
@functools.lru_cache(None)
|
||||
def _math_instruction_to_json(
|
||||
cls,
|
||||
math_instruction: Optional["MathInstruction"], # type: ignore[name-defined] # noqa: F821
|
||||
) -> str | None:
|
||||
"""Convert MathInstruction to JSON string.
|
||||
|
||||
Args:
|
||||
math_instruction: MathInstruction object
|
||||
|
||||
Returns:
|
||||
Optional[str]: JSON string representation or None
|
||||
"""
|
||||
if math_instruction is None:
|
||||
return None
|
||||
|
||||
result = {
|
||||
"instruction_shape": math_instruction.instruction_shape,
|
||||
"element_a": cls._enum_to_json(math_instruction.element_a),
|
||||
"element_b": cls._enum_to_json(math_instruction.element_b),
|
||||
"element_accumulator": cls._enum_to_json(
|
||||
math_instruction.element_accumulator
|
||||
),
|
||||
"opcode_class": cls._enum_to_json(math_instruction.opcode_class),
|
||||
"math_operation": cls._enum_to_json(math_instruction.math_operation),
|
||||
"element_scale_factor": cls._enum_to_json(
|
||||
math_instruction.element_scale_factor
|
||||
),
|
||||
}
|
||||
|
||||
return json.dumps(result)
|
||||
|
||||
@classmethod
|
||||
@functools.lru_cache(None)
|
||||
def _json_to_math_instruction(
|
||||
cls, json_dict: str | None
|
||||
) -> Optional["MathInstruction"]: # type: ignore[name-defined] # noqa: F821
|
||||
"""Convert JSON string to MathInstruction object.
|
||||
|
||||
Args:
|
||||
json_dict: JSON string representation
|
||||
|
||||
Returns:
|
||||
Optional[MathInstruction]: Reconstructed object or None
|
||||
"""
|
||||
if json_dict is None:
|
||||
return None
|
||||
|
||||
from cutlass_library import DataType
|
||||
from cutlass_library.library import MathInstruction, MathOperation, OpcodeClass
|
||||
|
||||
mi_dict = json.loads(json_dict)
|
||||
|
||||
# Convert string enum names back to enum values
|
||||
element_a = cls._json_to_enum(mi_dict["element_a"], DataType)
|
||||
element_b = cls._json_to_enum(mi_dict["element_b"], DataType)
|
||||
element_acc = cls._json_to_enum(mi_dict["element_accumulator"], DataType)
|
||||
|
||||
# Get the opcode_class enum
|
||||
opcode_class = cls._json_to_enum(mi_dict["opcode_class"], OpcodeClass)
|
||||
|
||||
# Get the math_operation enum
|
||||
math_op = cls._json_to_enum(mi_dict["math_operation"], MathOperation)
|
||||
|
||||
# Create the MathInstruction object
|
||||
math_instruction_obj = MathInstruction(
|
||||
instruction_shape=mi_dict["instruction_shape"],
|
||||
element_a=element_a,
|
||||
element_b=element_b,
|
||||
element_accumulator=element_acc,
|
||||
opcode_class=opcode_class,
|
||||
math_operation=math_op,
|
||||
)
|
||||
|
||||
# Add element_scale_factor if it exists
|
||||
if (
|
||||
"element_scale_factor" in mi_dict
|
||||
and mi_dict["element_scale_factor"] is not None
|
||||
):
|
||||
math_instruction_obj.element_scale_factor = cls._json_to_enum(
|
||||
mi_dict["element_scale_factor"], DataType
|
||||
)
|
||||
|
||||
return math_instruction_obj
|
||||
|
||||
@classmethod
|
||||
@functools.lru_cache(None)
|
||||
def _tensor_description_to_json(
|
||||
cls,
|
||||
tensor_desc: Optional["TensorDescription"], # type: ignore[name-defined] # noqa: F821
|
||||
) -> str | None:
|
||||
"""Convert TensorDescription to JSON string.
|
||||
|
||||
Args:
|
||||
tensor_desc: TensorDescription object
|
||||
|
||||
Returns:
|
||||
Optional[str]: JSON string representation or None
|
||||
"""
|
||||
if tensor_desc is None:
|
||||
return None
|
||||
|
||||
result = {
|
||||
"element": cls._enum_to_json(tensor_desc.element),
|
||||
"layout": cls._enum_to_json(tensor_desc.layout),
|
||||
"alignment": tensor_desc.alignment,
|
||||
"complex_transform": cls._enum_to_json(tensor_desc.complex_transform),
|
||||
}
|
||||
|
||||
return json.dumps(result)
|
||||
|
||||
@classmethod
|
||||
@functools.lru_cache(None)
|
||||
def _json_to_tensor_description(
|
||||
cls,
|
||||
json_dict: str | None,
|
||||
tensor_name: str | None = None,
|
||||
) -> Optional["TensorDescription"]: # type: ignore[name-defined] # noqa: F821
|
||||
"""Convert JSON string to TensorDescription object.
|
||||
|
||||
Args:
|
||||
json_dict: JSON string representation
|
||||
tensor_name: Name of the tensor to avoid cache in the same op
|
||||
|
||||
Returns:
|
||||
Optional[TensorDescription]: Reconstructed object or None
|
||||
"""
|
||||
if json_dict is None:
|
||||
return None
|
||||
|
||||
tensor_dict = json.loads(json_dict)
|
||||
|
||||
from cutlass_library import DataType
|
||||
from cutlass_library.library import (
|
||||
ComplexTransform,
|
||||
LayoutType,
|
||||
TensorDescription,
|
||||
)
|
||||
|
||||
element = cls._json_to_enum(tensor_dict["element"], DataType)
|
||||
layout = cls._json_to_enum(tensor_dict["layout"], LayoutType)
|
||||
alignment = tensor_dict["alignment"]
|
||||
complex_transform = cls._json_to_enum(
|
||||
tensor_dict["complex_transform"], ComplexTransform
|
||||
)
|
||||
|
||||
return TensorDescription(element, layout, alignment, complex_transform)
|
||||
|
||||
@classmethod
|
||||
@functools.lru_cache(None)
|
||||
def _enum_to_json(cls, enum_value: Enum | None) -> str | None:
|
||||
"""Convert enum value to JSON string.
|
||||
|
||||
Args:
|
||||
enum_value: Enum value
|
||||
|
||||
Returns:
|
||||
Optional[str]: JSON string representation or None
|
||||
"""
|
||||
if enum_value is None:
|
||||
return None
|
||||
|
||||
result = {
|
||||
"type": enum_value.__class__.__name__,
|
||||
"name": enum_value.name,
|
||||
}
|
||||
|
||||
return json.dumps(result)
|
||||
|
||||
@classmethod
|
||||
@functools.lru_cache(None)
|
||||
def _json_to_enum(cls, json_dict: str | None, enum_class: Any) -> Enum | None:
|
||||
"""Convert JSON string to enum value.
|
||||
|
||||
Format: {name: "EnumName", value: 1}
|
||||
|
||||
Args:
|
||||
json_dict: JSON string representation
|
||||
enum_class: Target enum class
|
||||
|
||||
Returns:
|
||||
Optional[Enum]: Reconstructed enum value or None
|
||||
"""
|
||||
if json_dict is None:
|
||||
return None
|
||||
|
||||
enum_dict = json.loads(json_dict)
|
||||
|
||||
return enum_class[enum_dict["name"]]
|
||||
|
||||
|
||||
@functools.lru_cache(1)
|
||||
def get_cutlass_operation_serializer() -> CUTLASSOperationSerializer | None:
|
||||
if not try_import_cutlass():
|
||||
return None
|
||||
return CUTLASSOperationSerializer()
|
||||
@@ -0,0 +1,383 @@
|
||||
# mypy: allow-untyped-defs
|
||||
import functools
|
||||
import hashlib
|
||||
import itertools
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, TYPE_CHECKING
|
||||
from unittest.mock import patch
|
||||
|
||||
import sympy
|
||||
|
||||
import torch
|
||||
from torch._inductor import config
|
||||
from torch._inductor.utils import clear_on_fresh_cache, Placeholder
|
||||
from torch._logging import getArtifactLogger
|
||||
|
||||
from ...autotune_process import CUTLASSBenchmarkRequest, TensorMeta
|
||||
from ...ir import Buffer, CUTLASSTemplateBuffer, IRNode, Layout
|
||||
from ...utils import IndentedBuffer, unique
|
||||
from ...virtualized import V
|
||||
from ..common import KernelTemplate
|
||||
from .kernel import CUTLASSTemplateCaller, CUTLASSTemplateKernel
|
||||
from .utils import DTYPE_TO_CUTLASS_TYPE
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ...scheduler import BaseSchedulerNode # noqa: TC004
|
||||
else:
|
||||
BaseSchedulerNode = Any
|
||||
|
||||
GemmOperation = Any
|
||||
|
||||
autotuning_log = getArtifactLogger(__name__, "autotuning")
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ArgInfo:
|
||||
name: str
|
||||
ty: str
|
||||
|
||||
|
||||
@clear_on_fresh_cache
|
||||
class CUTLASSTemplate(KernelTemplate):
|
||||
"""
|
||||
CUTLASSTemplate is a class that provides a template for generating CUTLASS Templates. Used as a baseclass for the
|
||||
CUTLASSGemmTemplate, providing functionality that might also be relevant for non-GEMM CUTLASS Kernels.
|
||||
"""
|
||||
|
||||
index_counter = itertools.count()
|
||||
# dict of cache key to (code, size_args)
|
||||
code_cache: dict[str, tuple[str, tuple[int, ...], tuple[int, ...]]] = {}
|
||||
cache_clear = staticmethod(code_cache.clear)
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
name: str,
|
||||
input_nodes: list[Buffer],
|
||||
layout: Layout,
|
||||
input_reorder: list[int] | None = None,
|
||||
) -> None:
|
||||
"""
|
||||
Baseclass for CUTLASS C++ Templates, derived from KernelTemplate.
|
||||
Not to be instantiated directly.
|
||||
|
||||
Args:
|
||||
name (str): The name of the CUTLASSTemplate object.
|
||||
input_nodes (List[IRNode]): A list of input IRNodes.
|
||||
layout (Layout): The layout of the output buffer / tensor.
|
||||
input_reorder (Optional[List[int]]): An optional list that specifies
|
||||
the order of the input nodes.
|
||||
"""
|
||||
super().__init__(name)
|
||||
self.input_nodes = input_nodes
|
||||
self.output_node: Buffer = Buffer(name="buf_out", layout=layout)
|
||||
self.input_reorder = input_reorder
|
||||
self.layout = layout
|
||||
self.device_type = layout.device.type
|
||||
|
||||
@classmethod
|
||||
@functools.lru_cache(None)
|
||||
# pyrefly: ignore [bad-override]
|
||||
def _template_from_string(cls, source: str) -> Any:
|
||||
return KernelTemplate._template_from_string(source)
|
||||
|
||||
@staticmethod
|
||||
def supports_epilogue_fusion(op: GemmOperation) -> bool:
|
||||
return False
|
||||
|
||||
def make_key(self, name: str, input_key: str, layout_repr: str) -> str:
|
||||
"""
|
||||
Make a key for the code cache. The idea of the method is to cache
|
||||
everything that matters but doesn't include runtime param values, i.e.,
|
||||
self.get_runtime_arg_values().
|
||||
|
||||
Args:
|
||||
kwargs: Additional keyword arguments. Including op (GemmOperation).
|
||||
"""
|
||||
return hashlib.sha256(
|
||||
str(
|
||||
(
|
||||
input_key,
|
||||
self.input_reorder,
|
||||
# output layout, same as self.output_node.get_layout()
|
||||
layout_repr,
|
||||
self.get_runtime_arg_info(),
|
||||
name,
|
||||
)
|
||||
).encode("utf-8")
|
||||
).hexdigest()
|
||||
|
||||
def generate_code_and_args(
|
||||
self, name: str, input_key: str, layout_repr: str, **kwargs
|
||||
) -> tuple[str, tuple[int, ...]]:
|
||||
"""
|
||||
Generate code and args with caching. We cache the code even if runtime
|
||||
args are different.
|
||||
"""
|
||||
key: str | None = None
|
||||
if config.cutlass.enable_caching_codegen:
|
||||
key = self.make_key(name=name, input_key=input_key, layout_repr=layout_repr)
|
||||
|
||||
if key is not None and key in self.code_cache:
|
||||
code, size_args, offset_args = self.code_cache[key]
|
||||
extra_args = tuple(
|
||||
list(size_args)
|
||||
+ list(offset_args)
|
||||
+ list(self.get_runtime_arg_values(**kwargs))
|
||||
)
|
||||
return code, extra_args
|
||||
|
||||
kernel_name = str(Placeholder.KERNEL_NAME)
|
||||
kernel = CUTLASSTemplateKernel(
|
||||
kernel_name=kernel_name,
|
||||
runtime_arg_info=self.get_runtime_arg_info(),
|
||||
runtime_arg_values=self.get_runtime_arg_values(**kwargs),
|
||||
device_type=self.device_type,
|
||||
)
|
||||
with patch.object(V.graph, "get_dtype", self._fake_get_dtype(self.output_node)):
|
||||
code = self.render(kernel=kernel, **kwargs)
|
||||
_, call_args, _, _ = kernel.args.python_argdefs()
|
||||
autotuning_log.debug("Generated Code:\n%s", code)
|
||||
autotuning_log.debug(
|
||||
"Args: cpp_argdefs: %s, python_argdefs: %s",
|
||||
kernel.args.cpp_argdefs(DTYPE_TO_CUTLASS_TYPE),
|
||||
kernel.args.python_argdefs(),
|
||||
)
|
||||
|
||||
input_reorder = (
|
||||
self.input_reorder
|
||||
if self.input_reorder is not None
|
||||
else list(range(len(self.input_nodes)))
|
||||
)
|
||||
expected_args = list(
|
||||
unique(self.input_nodes[idx].get_name() for idx in input_reorder)
|
||||
)
|
||||
expected_args.extend([self.output_node.get_name()])
|
||||
assert list(call_args)[: len(expected_args)] == expected_args, (
|
||||
call_args,
|
||||
expected_args,
|
||||
)
|
||||
# Resolve symbolic sizes to concrete ints for benchmarking only.
|
||||
V.graph.sizevars.optimization_hints(
|
||||
map(sympy.expand, call_args[len(expected_args) :])
|
||||
)
|
||||
size_args = V.graph.sizevars.optimization_hints(kernel.get_dynamic_shape_args())
|
||||
offset_args = V.graph.sizevars.optimization_hints(kernel.get_offset_args())
|
||||
|
||||
if key is not None:
|
||||
self.code_cache[key] = code, size_args, offset_args
|
||||
|
||||
# extra args has runtime params, which shouldn't be cached
|
||||
extra_args = tuple(
|
||||
list(size_args) + list(offset_args) + self.get_runtime_arg_values(**kwargs)
|
||||
)
|
||||
|
||||
return code, extra_args
|
||||
|
||||
def generate( # type: ignore[override]
|
||||
self,
|
||||
name: str,
|
||||
description: str,
|
||||
input_key: str,
|
||||
layout_repr: str,
|
||||
input_tensor_meta: TensorMeta | list[TensorMeta],
|
||||
output_tensor_meta: TensorMeta | list[TensorMeta],
|
||||
**kwargs,
|
||||
) -> CUTLASSTemplateCaller:
|
||||
"""
|
||||
Generates the CUDA template caller object for the given GEMM template and operation.
|
||||
This CUTLASSTemplateCaller may be used to call and benchmark the generated CUDA kernel
|
||||
in a standalone manner to enable Autotuning.
|
||||
|
||||
Args:
|
||||
description: op name followed by swizzle.
|
||||
kwargs: Additional keyword arguments.
|
||||
|
||||
Returns:
|
||||
A CUTLASSTemplateCaller object representing the generated CUDA template caller.
|
||||
"""
|
||||
code, extra_args = self.generate_code_and_args(
|
||||
name=name,
|
||||
input_key=input_key,
|
||||
layout_repr=layout_repr,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
# not caching since kernel name is needed below
|
||||
kernel_hash = hashlib.sha256(code.encode("utf-8")).hexdigest()[:8]
|
||||
kernel_name = f"cutlass_{kernel_hash}"
|
||||
code = code.replace(self.name, kernel_name)
|
||||
|
||||
# create the BenchmarkRequest
|
||||
bmreq = CUTLASSBenchmarkRequest(
|
||||
kernel_name=kernel_name,
|
||||
input_tensor_meta=input_tensor_meta,
|
||||
output_tensor_meta=output_tensor_meta,
|
||||
extra_args=extra_args,
|
||||
source_code=code,
|
||||
device_type=self.device_type,
|
||||
)
|
||||
|
||||
# kwargs has "op" argument in case of CUTLASSGemmTemplate
|
||||
op = kwargs["op"]
|
||||
if not op:
|
||||
supports_epilogue_fusion = False
|
||||
else:
|
||||
# epilogue fusion is only supported for TMA kernels
|
||||
supports_epilogue_fusion = self.supports_epilogue_fusion(op)
|
||||
|
||||
def make_kernel_render(
|
||||
template_node: CUTLASSTemplateBuffer,
|
||||
epilogue_nodes: list[BaseSchedulerNode] | None = None,
|
||||
) -> tuple[CUTLASSTemplateKernel, functools.partial[str]]:
|
||||
assert supports_epilogue_fusion or not epilogue_nodes, (
|
||||
"epilogue fusion is not supported for this kernel"
|
||||
)
|
||||
kernel = CUTLASSTemplateKernel(
|
||||
kernel_name=str(Placeholder.KERNEL_NAME),
|
||||
runtime_arg_info=self.get_runtime_arg_info(),
|
||||
runtime_arg_values=self.get_runtime_arg_values(**kwargs),
|
||||
device_type=self.device_type,
|
||||
)
|
||||
render = functools.partial(
|
||||
self.render,
|
||||
kernel=kernel,
|
||||
template_buffer_node=template_node,
|
||||
epilogue_nodes=epilogue_nodes,
|
||||
**kwargs, # includes "op" argument in case of CUTLASSGemmTemplate
|
||||
)
|
||||
return kernel, render
|
||||
|
||||
return CUTLASSTemplateCaller(
|
||||
kernel_name,
|
||||
"cutlass_gemm",
|
||||
self.input_nodes,
|
||||
self.output_node.get_layout(),
|
||||
make_kernel_render,
|
||||
bmreq,
|
||||
supports_epilogue_fusion,
|
||||
self,
|
||||
kwargs,
|
||||
description,
|
||||
)
|
||||
|
||||
def header(self) -> IndentedBuffer:
|
||||
res = IndentedBuffer()
|
||||
res.splice(
|
||||
"""
|
||||
#include <exception>
|
||||
#include <iostream>
|
||||
#include <memory>
|
||||
#include <random>
|
||||
#include <vector>
|
||||
"""
|
||||
)
|
||||
res.splice(
|
||||
"""
|
||||
#include "cute/tensor.hpp"
|
||||
#include "cutlass/cutlass.h"
|
||||
#include "cutlass/numeric_types.h"
|
||||
#include "cutlass/tensor_ref.h"
|
||||
#include "cutlass/util/host_tensor.h"
|
||||
#include "cutlass/util/reference/host/tensor_fill.h"
|
||||
#include "cutlass/util/reference/device/tensor_fill.h"
|
||||
#include "cutlass/util/device_memory.h"
|
||||
"""
|
||||
)
|
||||
return res
|
||||
|
||||
def globals(self) -> IndentedBuffer:
|
||||
res = IndentedBuffer()
|
||||
res.splice(
|
||||
"""
|
||||
// We compile all models with -fvisibility=hidden. Any symbols that need to be
|
||||
// exposed in the final shared library must be declared with PT_EXPORT to make
|
||||
// them visible.
|
||||
#ifdef __GNUC__ // Applies to any compiler with GNU extensions (clang and g++)
|
||||
#define PT_EXPORT __attribute__((__visibility__("default")))
|
||||
#else
|
||||
#ifdef _WIN32
|
||||
#define PT_EXPORT __declspec(dllexport)
|
||||
#else
|
||||
#define PT_EXPORT
|
||||
#endif
|
||||
#endif
|
||||
"""
|
||||
)
|
||||
res.splice(
|
||||
"""
|
||||
using namespace cute;
|
||||
#define CUTLASS_CHECK(status) \\
|
||||
{ \\
|
||||
cutlass::Status error = status; \\
|
||||
if (error != cutlass::Status::kSuccess) { \\
|
||||
auto msg = std::string("[") + __FILE__ + "] Got cutlass error: " + \\
|
||||
cutlassGetStatusString(error) + " at: " + std::to_string(__LINE__); \\
|
||||
throw std::runtime_error(msg); \\
|
||||
} \\
|
||||
}
|
||||
|
||||
// Used as pass-through functor in EVT just for type casting / rounding
|
||||
template <typename T>
|
||||
struct identity_op {
|
||||
CUTLASS_HOST_DEVICE
|
||||
T operator()(T val) const { return val; }
|
||||
};
|
||||
|
||||
"""
|
||||
)
|
||||
return res
|
||||
|
||||
def cute_int(self, int_str: str, var_name: str) -> str:
|
||||
res = ""
|
||||
if int_str in ("1", "1L"):
|
||||
res = "cute::Int<1>{}"
|
||||
else:
|
||||
res = int_str
|
||||
|
||||
return f"{res} /* {var_name} */"
|
||||
|
||||
_DTYPE_TO_CUTLASS = {
|
||||
torch.float32: "float",
|
||||
torch.float64: "double",
|
||||
torch.float16: "cutlass::half_t",
|
||||
torch.int32: "int32_t",
|
||||
torch.int16: "int16_t",
|
||||
torch.int8: "int8_t",
|
||||
torch.uint8: "uint8_t",
|
||||
torch.bool: "bool",
|
||||
torch.bfloat16: "cutlass::bfloat16_t",
|
||||
torch.float8_e4m3fn: "cutlass::float_e4m3_t",
|
||||
torch.float8_e5m2: "cutlass::float_e5m2_t",
|
||||
}
|
||||
|
||||
_DTYPE_TO_CUTLASS_SPARSE_META = {
|
||||
torch.int32: "uint32_t",
|
||||
torch.int16: "uint16_t",
|
||||
}
|
||||
|
||||
def cutlass_type_cast(self, node: IRNode, ptr: str) -> str:
|
||||
if node is None:
|
||||
return ptr
|
||||
else:
|
||||
return f"({self._DTYPE_TO_CUTLASS.get(node.get_dtype())}*)({ptr})"
|
||||
|
||||
def cutlass_sparse_meta_type_cast(self, node: IRNode, ptr: str) -> str:
|
||||
if node is None:
|
||||
return ptr
|
||||
else:
|
||||
return (
|
||||
f"({self._DTYPE_TO_CUTLASS_SPARSE_META.get(node.get_dtype())}*)({ptr})"
|
||||
)
|
||||
|
||||
def render(self, **kwargs) -> str:
|
||||
raise NotImplementedError
|
||||
|
||||
def get_runtime_arg_info(self) -> list[ArgInfo]:
|
||||
return [ArgInfo("swizzle", "const uint8_t")]
|
||||
|
||||
def get_runtime_arg_values(self, **kwargs) -> list[Any]:
|
||||
"""
|
||||
Helper method to retrieve runtime args from generate kwargs
|
||||
"""
|
||||
return [kwargs[arg.name] for arg in self.get_runtime_arg_info()]
|
||||
@@ -0,0 +1,536 @@
|
||||
# mypy: allow-untyped-defs
|
||||
import atexit
|
||||
import functools
|
||||
import logging
|
||||
import os
|
||||
import shutil
|
||||
import sys
|
||||
import time
|
||||
from dataclasses import dataclass
|
||||
from typing import Any
|
||||
from typing_extensions import TypeIs
|
||||
|
||||
import sympy
|
||||
|
||||
import torch
|
||||
from torch._inductor.runtime.runtime_utils import dynamo_timed
|
||||
from torch._inductor.utils import clear_on_fresh_cache
|
||||
from torch.utils._ordered_set import OrderedSet
|
||||
|
||||
from ... import config
|
||||
from ...ir import Layout
|
||||
from ...runtime.runtime_utils import cache_dir
|
||||
from ...virtualized import V
|
||||
from ..cuda.cuda_env import get_cuda_arch, get_cuda_version
|
||||
from ..xpu.xpu_env import get_xpu_arch, get_xpu_version
|
||||
|
||||
|
||||
log = logging.getLogger(__name__)
|
||||
|
||||
CUTLASS_OPERATION_KIND: str = "gemm"
|
||||
ACCUMULATOR_DTYPES: OrderedSet[torch.dtype] = OrderedSet([torch.float, torch.int32])
|
||||
XW_DTYPES: OrderedSet[torch.dtype] = OrderedSet(
|
||||
[torch.half, torch.bfloat16, torch.float8_e4m3fn, torch.int8, torch.float8_e5m2]
|
||||
)
|
||||
|
||||
|
||||
@atexit.register
|
||||
def move_cutlass_compiled_cache() -> None:
|
||||
"""Move CUTLASS compiled cache file to the cache directory if it exists."""
|
||||
if try_import_cutlass.cache_info().currsize == 0:
|
||||
return
|
||||
|
||||
try:
|
||||
import cutlass_cppgen # type: ignore[import-not-found]
|
||||
except ImportError:
|
||||
return
|
||||
|
||||
# Check if the CACHE_FILE attribute exists in cutlass_cppgen and if the file exists
|
||||
if not hasattr(cutlass_cppgen, "CACHE_FILE") or not os.path.exists(
|
||||
cutlass_cppgen.CACHE_FILE
|
||||
):
|
||||
return
|
||||
|
||||
try:
|
||||
filename = os.path.basename(cutlass_cppgen.CACHE_FILE)
|
||||
shutil.move(cutlass_cppgen.CACHE_FILE, os.path.join(cache_dir(), filename))
|
||||
log.debug("Moved CUTLASS compiled cache file to %s", cache_dir())
|
||||
except OSError:
|
||||
log.warning("Failed to move CUTLASS compiled cache file", exc_info=True)
|
||||
|
||||
|
||||
def _rename_cutlass_import(content: str, cutlass_modules: list[str]) -> str:
|
||||
for cutlass_module in cutlass_modules:
|
||||
content = content.replace(
|
||||
f"from {cutlass_module} import ",
|
||||
f"from cutlass_library.{cutlass_module} import ",
|
||||
)
|
||||
return content
|
||||
|
||||
|
||||
@functools.cache
|
||||
def try_import_cutlass() -> bool:
|
||||
"""
|
||||
We want to support three ways of passing in CUTLASS:
|
||||
1. fbcode, handled by the internal build system.
|
||||
2. User specifies cutlass_dir. The default is ../third_party/cutlass/,
|
||||
which is the directory when developers build from source.
|
||||
"""
|
||||
if config.is_fbcode():
|
||||
try:
|
||||
import cutlass_cppgen # type: ignore[import-not-found] # noqa: F401
|
||||
import cutlass_library # type: ignore[import-not-found]
|
||||
except ImportError as e:
|
||||
log.warning( # noqa: G200
|
||||
"Failed to import CUTLASS packages in fbcode: %s, ignoring the CUTLASS backend.",
|
||||
e,
|
||||
)
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
# Copy CUTLASS python scripts to a temp dir and add the temp dir to Python search path.
|
||||
# This is a temporary hack to avoid CUTLASS module naming conflicts.
|
||||
# TODO(ipiszy): remove this hack when CUTLASS solves Python scripts packaging structure issues.
|
||||
|
||||
# TODO(mlazos): epilogue visitor tree currently lives in python/cutlass,
|
||||
# but will be moved to python/cutlass_library in the future (later 2025)
|
||||
def path_join(path0, path1):
|
||||
return os.path.abspath(os.path.join(path0, path1))
|
||||
|
||||
# contains both cutlass and cutlass_library
|
||||
# we need cutlass for eVT
|
||||
cutlass_dir = (
|
||||
config.xpu.cutlass_dir
|
||||
if torch.xpu._is_compiled()
|
||||
else config.cutlass.cutlass_dir
|
||||
)
|
||||
cutlass_python_path = path_join(cutlass_dir, "python")
|
||||
torch_root = os.path.abspath(os.path.dirname(torch.__file__))
|
||||
mock_src_path = os.path.join(
|
||||
torch_root,
|
||||
"_inductor",
|
||||
"codegen",
|
||||
"cutlass",
|
||||
"lib_extensions",
|
||||
"cutlass_mock_imports",
|
||||
)
|
||||
|
||||
cutlass_library_src_path = path_join(cutlass_python_path, "cutlass_library")
|
||||
cutlass_cppgen_src_path = path_join(cutlass_python_path, "cutlass_cppgen")
|
||||
pycute_src_path = path_join(cutlass_python_path, "pycute")
|
||||
|
||||
tmp_cutlass_full_path = os.path.abspath(os.path.join(cache_dir(), "torch_cutlass"))
|
||||
|
||||
dst_link_library = path_join(tmp_cutlass_full_path, "cutlass_library")
|
||||
dst_link_cutlass_cppgen = path_join(tmp_cutlass_full_path, "cutlass_cppgen")
|
||||
dst_link_pycute = path_join(tmp_cutlass_full_path, "pycute")
|
||||
|
||||
# mock modules to import cutlass
|
||||
mock_modules = ["cuda", "scipy", "pydot"]
|
||||
|
||||
if os.path.isdir(cutlass_python_path):
|
||||
if tmp_cutlass_full_path not in sys.path:
|
||||
|
||||
def link_and_append(dst_link, src_path, parent_dir):
|
||||
if os.path.lexists(dst_link):
|
||||
assert os.path.islink(dst_link), (
|
||||
f"{dst_link} is not a symlink. Try to remove {dst_link} manually and try again."
|
||||
)
|
||||
assert os.path.realpath(os.readlink(dst_link)) == os.path.realpath(
|
||||
src_path,
|
||||
), f"Symlink at {dst_link} does not point to {src_path}"
|
||||
else:
|
||||
os.makedirs(parent_dir, exist_ok=True)
|
||||
os.symlink(src_path, dst_link)
|
||||
|
||||
if parent_dir not in sys.path:
|
||||
sys.path.append(parent_dir)
|
||||
|
||||
link_and_append(
|
||||
dst_link_library, cutlass_library_src_path, tmp_cutlass_full_path
|
||||
)
|
||||
link_and_append(
|
||||
dst_link_cutlass_cppgen, cutlass_cppgen_src_path, tmp_cutlass_full_path
|
||||
)
|
||||
link_and_append(dst_link_pycute, pycute_src_path, tmp_cutlass_full_path)
|
||||
|
||||
for module in mock_modules:
|
||||
link_and_append(
|
||||
path_join(tmp_cutlass_full_path, module), # dst_link
|
||||
path_join(mock_src_path, module), # src_path
|
||||
tmp_cutlass_full_path, # parent
|
||||
)
|
||||
|
||||
try:
|
||||
import cutlass_cppgen # type: ignore[import-not-found] # noqa: F401, F811
|
||||
import cutlass_library.generator # noqa: F401
|
||||
import cutlass_library.library # noqa: F401
|
||||
import cutlass_library.manifest # noqa: F401
|
||||
import pycute # type: ignore[import-not-found] # noqa: F401
|
||||
|
||||
return True
|
||||
except ImportError as e:
|
||||
log.debug( # noqa: G200
|
||||
"Failed to import CUTLASS packages: %s, ignoring the CUTLASS backend.",
|
||||
e,
|
||||
)
|
||||
else:
|
||||
log.debug(
|
||||
"Failed to import CUTLASS packages: CUTLASS repo does not exist: %s",
|
||||
cutlass_python_path,
|
||||
)
|
||||
return False
|
||||
|
||||
|
||||
def _normalize_xpu_arch(arch: str) -> str:
|
||||
if arch.startswith("Xe"):
|
||||
return arch[2:]
|
||||
if 12 <= int(arch) and int(arch) <= 50:
|
||||
return arch
|
||||
else:
|
||||
raise NotImplementedError(f"Unsupported xpu arch: {arch}")
|
||||
|
||||
|
||||
def _normalize_cuda_arch(arch: str) -> str:
|
||||
arch_num = arch
|
||||
if isinstance(arch, str):
|
||||
digits = "".join(ch for ch in arch if ch.isdigit())
|
||||
if not digits:
|
||||
raise ValueError(f"Unrecognized cuda arch: {arch}")
|
||||
arch_num = int(digits)
|
||||
else:
|
||||
arch_num = int(arch)
|
||||
|
||||
if arch_num > 103:
|
||||
log.warning("Detected CUDA architecture > 103: %s. Please file an issue.", arch)
|
||||
return str(arch_num)
|
||||
if arch_num >= 103:
|
||||
return "103"
|
||||
if arch_num >= 100:
|
||||
return "100"
|
||||
if arch_num >= 90:
|
||||
return "90"
|
||||
if arch_num >= 80:
|
||||
return "80"
|
||||
if arch_num >= 75:
|
||||
return "75"
|
||||
if arch_num >= 70:
|
||||
return "70"
|
||||
raise NotImplementedError(f"Unsupported cuda arch: {arch}")
|
||||
|
||||
|
||||
@functools.lru_cache(8)
|
||||
def cutlass_arch(device_type: str) -> str:
|
||||
if device_type == "xpu":
|
||||
arch = get_xpu_arch()
|
||||
return _normalize_xpu_arch(arch)
|
||||
else:
|
||||
arch = get_cuda_arch()
|
||||
return _normalize_cuda_arch(arch)
|
||||
|
||||
|
||||
@functools.lru_cache(1)
|
||||
def toolkit_version(device_type: str) -> str:
|
||||
if device_type == "xpu":
|
||||
return get_xpu_version()
|
||||
else:
|
||||
return get_cuda_version()
|
||||
|
||||
|
||||
@dataclass
|
||||
class CUTLASSArgs:
|
||||
"""
|
||||
CUTLASS args used to initialize a CUTLASS Manifest.
|
||||
"""
|
||||
|
||||
architectures: str | None = None
|
||||
toolkit_version: str | None = None
|
||||
instantiation_level: str | None = None
|
||||
operations: str | None = None
|
||||
|
||||
build_dir = ""
|
||||
curr_build_dir = ""
|
||||
generator_target = ""
|
||||
kernels = "all"
|
||||
ignore_kernels = ""
|
||||
exclude_kernels = ""
|
||||
# TODO: these three look dead?
|
||||
kernel_filter_file: None = None
|
||||
selected_kernel_list: None = None
|
||||
interface_dir: None = None
|
||||
filter_by_cc = True
|
||||
disable_full_archs_compilation = False
|
||||
device_type: str = "cuda"
|
||||
|
||||
def __post_init__(self):
|
||||
if self.architectures is None or self.toolkit_version is None:
|
||||
raise RuntimeError(
|
||||
f"{self.architectures=} or {self.toolkit_version=} is None!"
|
||||
)
|
||||
|
||||
|
||||
@clear_on_fresh_cache
|
||||
@functools.cache
|
||||
def _gen_ops_cached(arch: str, version: str, device_type: str) -> dict[Any, Any]:
|
||||
# Note: Cache needs to be specific for cuda architecture and version
|
||||
|
||||
# Import cutlass python scripts.
|
||||
assert try_import_cutlass()
|
||||
import cutlass_library.generator as cutlass_generator
|
||||
import cutlass_library.manifest as cutlass_manifest
|
||||
|
||||
if arch is None or version is None:
|
||||
log.error(
|
||||
"Cannot detect cuda arch %s or version %s. "
|
||||
"Will discard all cutlass ops. "
|
||||
"Please consider setting _inductor.cuda.arch and _inductor.cuda.version configs.",
|
||||
arch,
|
||||
version,
|
||||
)
|
||||
return {}
|
||||
|
||||
gen_arch = (
|
||||
"100" if arch == "103" else arch
|
||||
) # CUTLASS SM103 generator only covers NVFB4; fallback to SM100 set
|
||||
instantiation_level: str = config.cutlass.cutlass_instantiation_level
|
||||
args = CUTLASSArgs(
|
||||
architectures=gen_arch,
|
||||
toolkit_version=version,
|
||||
instantiation_level=instantiation_level,
|
||||
operations=CUTLASS_OPERATION_KIND,
|
||||
device_type=device_type,
|
||||
)
|
||||
manifest = cutlass_manifest.Manifest(args)
|
||||
|
||||
start_time = time.time()
|
||||
if device_type == "xpu":
|
||||
if hasattr(cutlass_generator, "GenerateIntelXe"):
|
||||
cutlass_generator.GenerateIntelXe(
|
||||
manifest, args.toolkit_version, arch=int(arch)
|
||||
)
|
||||
else:
|
||||
raise NotImplementedError(
|
||||
"Arch " + arch + " is not supported by current cutlass lib."
|
||||
)
|
||||
|
||||
elif arch == "100":
|
||||
if hasattr(cutlass_generator, "GenerateSM100"):
|
||||
cutlass_generator.GenerateSM100(manifest, args.toolkit_version)
|
||||
cutlass_generator.GenerateSM90(manifest, args.toolkit_version)
|
||||
else:
|
||||
try:
|
||||
func = getattr(cutlass_generator, "GenerateSM" + gen_arch)
|
||||
func(manifest, args.toolkit_version)
|
||||
except AttributeError as e:
|
||||
raise NotImplementedError(
|
||||
"Arch " + gen_arch + " is not supported by current cutlass lib."
|
||||
) from e
|
||||
|
||||
log.info(
|
||||
"CUTLASS library generated a dict of %d operation kinds in %.2f seconds",
|
||||
len(manifest.operations),
|
||||
time.time() - start_time,
|
||||
)
|
||||
return manifest.operations
|
||||
|
||||
|
||||
def gen_ops(device_type: str) -> dict[Any, Any]:
|
||||
"""
|
||||
Generates all supported CUTLASS operations.
|
||||
"""
|
||||
with dynamo_timed("cutlass_utils.gen_ops"):
|
||||
arch = cutlass_arch(device_type)
|
||||
version = toolkit_version(device_type)
|
||||
return _gen_ops_cached(arch, version, device_type)
|
||||
|
||||
|
||||
from ..cpp_utils import DTYPE_TO_CPP
|
||||
|
||||
|
||||
if torch.xpu._is_compiled():
|
||||
DTYPE_TO_CUTLASS_TYPE = {
|
||||
**DTYPE_TO_CPP,
|
||||
torch.float16: "uint16_t",
|
||||
torch.bfloat16: "uint16_t",
|
||||
torch.float8_e4m3fn: "uint8_t",
|
||||
torch.float8_e5m2: "uint8_t",
|
||||
}
|
||||
else:
|
||||
DTYPE_TO_CUTLASS_TYPE = {
|
||||
**DTYPE_TO_CPP,
|
||||
torch.float16: "__half",
|
||||
torch.bfloat16: "__nv_bfloat16",
|
||||
torch.float8_e4m3fn: "__nv_fp8_e4m3",
|
||||
torch.float8_e5m2: "__nv_fp8_e5m2",
|
||||
}
|
||||
|
||||
|
||||
@functools.lru_cache(32)
|
||||
def torch_dtype_to_cutlass_type(
|
||||
torch_dtype: torch.dtype,
|
||||
) -> "cutlass_library.library.DataType": # type: ignore[name-defined] # noqa: F821
|
||||
# Import cutlass python scripts.
|
||||
assert try_import_cutlass()
|
||||
import cutlass_library # type: ignore[import]
|
||||
|
||||
if torch_dtype == torch.float:
|
||||
return cutlass_library.library.DataType.f32
|
||||
elif torch_dtype == torch.half:
|
||||
return cutlass_library.library.DataType.f16
|
||||
elif torch_dtype == torch.bfloat16:
|
||||
return cutlass_library.library.DataType.bf16
|
||||
elif torch_dtype == torch.float8_e4m3fn:
|
||||
return cutlass_library.library.DataType.e4m3
|
||||
elif torch_dtype == torch.float8_e5m2:
|
||||
return cutlass_library.library.DataType.e5m2
|
||||
else:
|
||||
raise NotImplementedError(f"Unsupported data type: {torch_dtype=}")
|
||||
|
||||
|
||||
@functools.lru_cache(32)
|
||||
def dtype_match(
|
||||
torch_dtype: torch.dtype | None,
|
||||
cutlass_dtype: "cutlass_library.library.DataType", # type: ignore[name-defined] # noqa: F821
|
||||
) -> bool:
|
||||
# Import cutlass python scripts.
|
||||
assert try_import_cutlass()
|
||||
import cutlass_library
|
||||
|
||||
if torch_dtype == torch.float:
|
||||
return (
|
||||
cutlass_dtype == cutlass_library.library.DataType.f32
|
||||
or cutlass_dtype == cutlass_library.library.DataType.tf32
|
||||
)
|
||||
elif torch_dtype == torch.half:
|
||||
return cutlass_dtype == cutlass_library.library.DataType.f16
|
||||
elif torch_dtype == torch.bfloat16:
|
||||
return cutlass_dtype == cutlass_library.library.DataType.bf16
|
||||
elif torch_dtype == torch.int8:
|
||||
return cutlass_dtype == cutlass_library.library.DataType.s8
|
||||
elif torch_dtype == torch.uint8:
|
||||
return cutlass_dtype == cutlass_library.library.DataType.u8
|
||||
elif torch_dtype == torch.int32:
|
||||
return cutlass_dtype == cutlass_library.library.DataType.s32
|
||||
elif torch_dtype == torch.float8_e4m3fn:
|
||||
return cutlass_dtype == cutlass_library.library.DataType.e4m3
|
||||
elif torch_dtype == torch.float8_e5m2:
|
||||
return cutlass_dtype == cutlass_library.library.DataType.e5m2
|
||||
else:
|
||||
return False
|
||||
|
||||
|
||||
def get_accumulator_dtype(
|
||||
input_torch_dtypes: list[torch.dtype],
|
||||
) -> torch.dtype | None:
|
||||
"""
|
||||
Given a pair of input torch dtypes, returns the inferred accumulator torch dtype.
|
||||
"""
|
||||
|
||||
assert OrderedSet(input_torch_dtypes) <= XW_DTYPES, (
|
||||
f"{input_torch_dtypes=} is not supported"
|
||||
)
|
||||
|
||||
if len(input_torch_dtypes) != 2:
|
||||
return None
|
||||
|
||||
if OrderedSet(input_torch_dtypes) == OrderedSet(
|
||||
[torch.float8_e5m2, torch.float8_e4m3fn]
|
||||
):
|
||||
return torch.float
|
||||
|
||||
torch_dtype = None
|
||||
if input_torch_dtypes[0] == input_torch_dtypes[1]:
|
||||
torch_dtype = input_torch_dtypes[0]
|
||||
else:
|
||||
size0 = torch.tensor([], dtype=input_torch_dtypes[0]).element_size()
|
||||
size1 = torch.tensor([], dtype=input_torch_dtypes[1]).element_size()
|
||||
if size0 > size1:
|
||||
dtype0, dtype1 = input_torch_dtypes
|
||||
else:
|
||||
dtype1, dtype0 = input_torch_dtypes
|
||||
if dtype0 in [torch.half, torch.bfloat16] and dtype1 in [
|
||||
torch.int8,
|
||||
torch.uint8,
|
||||
]:
|
||||
torch_dtype = dtype0
|
||||
|
||||
if torch_dtype in (
|
||||
torch.float16,
|
||||
torch.bfloat16,
|
||||
torch.float,
|
||||
torch.float8_e4m3fn,
|
||||
torch.float8_e5m2,
|
||||
):
|
||||
accumulator_dtype = torch.float
|
||||
elif torch_dtype == torch.int8:
|
||||
accumulator_dtype = torch.int32
|
||||
else:
|
||||
raise NotImplementedError(f"Unsupported data types: {input_torch_dtypes=}")
|
||||
|
||||
assert accumulator_dtype in ACCUMULATOR_DTYPES, (
|
||||
f"{accumulator_dtype=} is not supported"
|
||||
)
|
||||
return accumulator_dtype
|
||||
|
||||
|
||||
@functools.lru_cache(32)
|
||||
def get_alignments(torch_dtype: torch.dtype) -> list[int]:
|
||||
"""
|
||||
Returns all possible valid CUTLASS alignments in terms of the number of elements for a given dtype.
|
||||
CUTLASS gemm / conv SM80 APIs support 16 bytes max alignment, and 2 bytes min alignment.
|
||||
"""
|
||||
|
||||
if torch_dtype in (torch.half, torch.bfloat16):
|
||||
return [8, 4, 2, 1]
|
||||
elif torch_dtype == torch.float:
|
||||
return [4, 2, 1]
|
||||
elif torch_dtype in (
|
||||
torch.uint8,
|
||||
torch.int8,
|
||||
torch.float8_e4m3fn,
|
||||
torch.float8_e5m2,
|
||||
):
|
||||
return [16, 8, 4, 2]
|
||||
elif torch_dtype == torch.int32:
|
||||
return [4, 2, 1]
|
||||
else:
|
||||
raise NotImplementedError(f"unsupported {torch_dtype=} for alignments")
|
||||
|
||||
|
||||
def get_max_alignment(inductor_layout: Layout) -> int:
|
||||
"""
|
||||
Returns the max alignment (in terms of number of elements) for a given Inductor Layout.
|
||||
"""
|
||||
|
||||
dtype = inductor_layout.dtype
|
||||
size = inductor_layout.size
|
||||
offset = inductor_layout.offset
|
||||
|
||||
def is_static_int(number: object) -> TypeIs[int | sympy.Integer]:
|
||||
return isinstance(number, (int | sympy.Integer))
|
||||
|
||||
def a_factor_of(x, alignment):
|
||||
if is_static_int(x) and is_static_int(alignment):
|
||||
return x % alignment == 0
|
||||
rem = sympy.Mod(x, alignment)
|
||||
return V.graph.sizevars.evaluate_expr(sympy.Eq(rem, 0))
|
||||
|
||||
try:
|
||||
contiguous_dim = inductor_layout.stride.index(1)
|
||||
except ValueError:
|
||||
# No dim with stride 1 found, return 1
|
||||
return 1
|
||||
alignments = get_alignments(dtype)
|
||||
for alignment in alignments:
|
||||
if not a_factor_of(size[contiguous_dim], alignment) or not a_factor_of(
|
||||
offset, alignment
|
||||
):
|
||||
continue
|
||||
if all(
|
||||
(dim == contiguous_dim)
|
||||
or a_factor_of(inductor_layout.stride[dim], alignment)
|
||||
for dim in range(len(size))
|
||||
):
|
||||
return alignment
|
||||
return 1
|
||||
@@ -0,0 +1,290 @@
|
||||
# mypy: allow-untyped-defs
|
||||
from __future__ import annotations
|
||||
|
||||
import functools
|
||||
import logging
|
||||
import os
|
||||
from enum import Enum
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import torch
|
||||
from torch import dtype as torch_dtype
|
||||
|
||||
from .. import config
|
||||
from ..virtualized import V
|
||||
from .multi_kernel import MultiKernel
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Callable
|
||||
|
||||
|
||||
log = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def _print_debugging_tensor_value_info(msg, arg):
|
||||
# helper for printing debugging stats for intermediate tensor values
|
||||
# at jit inductor level codegen
|
||||
max_numel_to_print = 64
|
||||
print(msg)
|
||||
if not isinstance(arg, torch.Tensor):
|
||||
print("Value: ", arg)
|
||||
return
|
||||
numel = arg.float().numel()
|
||||
# print the debug printing stats
|
||||
if numel <= max_numel_to_print:
|
||||
print(arg)
|
||||
print("Number of elements: ", numel)
|
||||
print("Size: ", arg.float().size())
|
||||
print("Dtype: ", arg.float().mean().item())
|
||||
print("Mean: ", arg.float().mean().item())
|
||||
print("Min: ", arg.float().min().item())
|
||||
print("Max: ", arg.float().max().item())
|
||||
print("Std: ", arg.float().std().item())
|
||||
|
||||
|
||||
# AOTI debug printing related configs
|
||||
class IntermediateValueDebuggingLevel(Enum):
|
||||
# OFF: No intermediate tensor value debug info will be printed or saved.
|
||||
OFF = "0"
|
||||
# LEVEL 1: Save all intermediate tensor values to individual `.pt` files. No debug printing will be displayed.
|
||||
SAVE_ONLY = "1"
|
||||
# LEVEL 2: Print all intermediate tensor values by default to the console. No debug saving will be performed.
|
||||
PRINT_ONLY = "2"
|
||||
# LEVEL 3: Print all kernel names to the console only. No debug saving/printing for input tensor value info will be performed.
|
||||
# This mode can be helpful in cases when you just want to pinpointing what kernel is running into a CUDA IMA issue, etc.
|
||||
PRINT_KERNEL_NAMES_ONLY = "3"
|
||||
|
||||
|
||||
class DebugPrinterManager:
|
||||
def __init__(
|
||||
self,
|
||||
debug_printer_level,
|
||||
use_array_ref: bool,
|
||||
writeline: Callable[..., None] | None = None,
|
||||
args_to_print_or_save: list[str] | None = None,
|
||||
kernel_name: str = "",
|
||||
kernel=None,
|
||||
arg_signatures: list[type] | None = None,
|
||||
kernel_type=None,
|
||||
):
|
||||
self.debug_printer_level = IntermediateValueDebuggingLevel(debug_printer_level)
|
||||
self.use_array_ref = use_array_ref
|
||||
if args_to_print_or_save is None:
|
||||
args_to_print_or_save = []
|
||||
self.args_to_print_or_save = args_to_print_or_save
|
||||
self.kernel_name = kernel_name
|
||||
self.arg_signatures: list[type] | None = None
|
||||
self.kernel = kernel
|
||||
self.filtered_kernel_names_to_print = self._get_debug_filtered_kernel_names()
|
||||
self.kernel_type = None
|
||||
|
||||
def __enter__(self):
|
||||
self._perform_debug_print_or_save_helper(
|
||||
self.args_to_print_or_save,
|
||||
self.kernel_name,
|
||||
before_launch=True,
|
||||
arg_signatures=self.arg_signatures,
|
||||
)
|
||||
|
||||
def __exit__(self, args_to_print_or_save, kernel_name, arg_signatures):
|
||||
self._perform_debug_print_or_save_helper(
|
||||
args_to_print_or_save,
|
||||
kernel_name,
|
||||
before_launch=False,
|
||||
arg_signatures=arg_signatures,
|
||||
)
|
||||
|
||||
def _perform_debug_print_or_save_helper(
|
||||
self,
|
||||
args_to_print_or_save,
|
||||
kernel_name,
|
||||
before_launch,
|
||||
arg_signatures: list[type] | None = None,
|
||||
):
|
||||
if self.debug_printer_level == IntermediateValueDebuggingLevel.OFF:
|
||||
return
|
||||
if self.debug_printer_level == IntermediateValueDebuggingLevel.SAVE_ONLY:
|
||||
# by default save all the tensor values before launch
|
||||
self.codegen_intermediate_tensor_value_save(
|
||||
self.args_to_print_or_save,
|
||||
self.kernel_name,
|
||||
before_launch,
|
||||
arg_signatures=self.arg_signatures,
|
||||
)
|
||||
if self.debug_printer_level == IntermediateValueDebuggingLevel.PRINT_ONLY:
|
||||
# by default print all the tensor values before launch
|
||||
self.codegen_intermediate_tensor_value_print(
|
||||
self.args_to_print_or_save,
|
||||
self.kernel_name,
|
||||
before_launch,
|
||||
arg_signatures=self.arg_signatures,
|
||||
)
|
||||
if (
|
||||
self.debug_printer_level
|
||||
== IntermediateValueDebuggingLevel.PRINT_KERNEL_NAMES_ONLY
|
||||
):
|
||||
# Print all kernel names to the console only
|
||||
self.codegen_intermediate_tensor_value_print(
|
||||
[],
|
||||
self.kernel_name,
|
||||
before_launch,
|
||||
)
|
||||
|
||||
@functools.lru_cache # noqa: B019
|
||||
def _get_debug_filtered_kernel_names(self) -> list[str]:
|
||||
if config.aot_inductor.filtered_kernel_names is None:
|
||||
return []
|
||||
return [
|
||||
x.strip()
|
||||
for x in config.aot_inductor.filtered_kernel_names.lower().split(",")
|
||||
]
|
||||
|
||||
def set_printer_args(
|
||||
self,
|
||||
args_to_print_or_save: list[str],
|
||||
kernel_name: str,
|
||||
arg_signatures: list[type] | None,
|
||||
kernel,
|
||||
kernel_type=None,
|
||||
):
|
||||
# Note: MultiKernel debug printing is not supported for now
|
||||
if isinstance(kernel, MultiKernel):
|
||||
log.info(
|
||||
"MultiKernel type is not supported in AOTI debug printer tool yet."
|
||||
)
|
||||
self.debug_printer_level = IntermediateValueDebuggingLevel.OFF
|
||||
|
||||
self.kernel_type = kernel_type
|
||||
# Note: if the kernel type is an extern kernel (or cpp kernel), we do a special handling to
|
||||
# get the list of args_to_print_or_save
|
||||
# TODO: Find a more reliable way to detect kernel args types to print for extern kernel calls
|
||||
if kernel_type == "extern":
|
||||
args_to_print_or_save_extern = [
|
||||
arg
|
||||
for arg in args_to_print_or_save
|
||||
if isinstance(arg, str) and arg.startswith(("buf", "arg"))
|
||||
]
|
||||
self.args_to_print_or_save = args_to_print_or_save_extern
|
||||
elif kernel_type == "cpp":
|
||||
self.args_to_print_or_save = [
|
||||
(
|
||||
f"copy_arrayref_tensor_to_tensor({arg})"
|
||||
if self.use_array_ref
|
||||
else arg
|
||||
)
|
||||
for arg in args_to_print_or_save
|
||||
if isinstance(arg, str) and arg.startswith(("buf", "arg"))
|
||||
]
|
||||
else:
|
||||
self.args_to_print_or_save = args_to_print_or_save
|
||||
self.kernel_name = kernel_name
|
||||
self.arg_signatures = arg_signatures
|
||||
self.kernel = kernel
|
||||
|
||||
def codegen_model_inputs_value_print(self, input_args_to_print: list[str]) -> None:
|
||||
if self.debug_printer_level != IntermediateValueDebuggingLevel.PRINT_ONLY:
|
||||
return
|
||||
for arg in input_args_to_print:
|
||||
if V.graph.cpp_wrapper:
|
||||
V.graph.wrapper_code.prefix.writeline(
|
||||
f'aoti_torch_print_tensor_handle({arg}, "aoti_model_inputs - {arg}");'
|
||||
)
|
||||
|
||||
def codegen_intermediate_tensor_value_save(
|
||||
self,
|
||||
args_to_save,
|
||||
kernel_name,
|
||||
before_launch=True,
|
||||
arg_signatures: list[type] | None = None,
|
||||
) -> None:
|
||||
for i, arg in enumerate(args_to_save):
|
||||
if arg_signatures is not None and not isinstance(
|
||||
arg_signatures[i], torch_dtype
|
||||
):
|
||||
# infer from the arg data type (has torch.dtype) to see if it is a tensor type
|
||||
continue
|
||||
launch_prefix = "before_launch" if before_launch else "after_launch"
|
||||
if V.graph.cpp_wrapper:
|
||||
V.graph.wrapper_code.writeline(
|
||||
f'aoti_torch_save_tensor_handle({arg}, "{arg}", "{launch_prefix}", "{kernel_name}");'
|
||||
)
|
||||
else:
|
||||
cwd = os.getcwd()
|
||||
saved_dir = cwd + "/tmp/jit_inductor/"
|
||||
if not os.path.exists(saved_dir):
|
||||
log.info(
|
||||
"Creating directory to save inductor intermediate tensor values."
|
||||
)
|
||||
os.makedirs(saved_dir)
|
||||
# Save the model to the directory
|
||||
saved_path = saved_dir + f"{launch_prefix}_{kernel_name}_{arg}.pt"
|
||||
log.info(
|
||||
"Saved intermediate tensor %s for %s to %s",
|
||||
arg,
|
||||
kernel_name,
|
||||
saved_path,
|
||||
)
|
||||
line = f"torch.save({arg}, '{saved_path}')"
|
||||
V.graph.wrapper_code.writeline(line)
|
||||
|
||||
def codegen_intermediate_tensor_value_print(
|
||||
self,
|
||||
args_to_print,
|
||||
kernel_name,
|
||||
before_launch=True,
|
||||
arg_signatures: list[type] | None = None,
|
||||
) -> None:
|
||||
launch_prefix = "before_launch" if before_launch else "after_launch"
|
||||
|
||||
# if the debug printing level is PRINT_KERNEL_NAMES_ONLY
|
||||
# we only print the kernel name to the console
|
||||
if (
|
||||
self.debug_printer_level
|
||||
== IntermediateValueDebuggingLevel.PRINT_KERNEL_NAMES_ONLY
|
||||
):
|
||||
if V.graph.cpp_wrapper:
|
||||
V.graph.wrapper_code.writeline(
|
||||
f'printf("[ {launch_prefix}: {kernel_name} ]\\n");'
|
||||
)
|
||||
return
|
||||
|
||||
if self.debug_printer_level != IntermediateValueDebuggingLevel.PRINT_ONLY:
|
||||
return
|
||||
for i, arg in enumerate(args_to_print):
|
||||
# when debug printing is enabled i.e. IntermediateValueDebuggingLevel.PRINT_ONLY,
|
||||
# check if filtered kernel name list is provided
|
||||
if (
|
||||
len(self.filtered_kernel_names_to_print) > 0
|
||||
and kernel_name.lower() not in self.filtered_kernel_names_to_print
|
||||
):
|
||||
continue
|
||||
if V.graph.cpp_wrapper:
|
||||
if arg_signatures is not None and isinstance(
|
||||
arg_signatures[i], torch_dtype
|
||||
):
|
||||
# infer from the arg data type (has torch.dtype) to see if it is a tensor type
|
||||
V.graph.wrapper_code.writeline(
|
||||
f'aoti_torch_print_tensor_handle({arg}, "{launch_prefix} - {kernel_name} - {arg}");'
|
||||
)
|
||||
elif arg_signatures is not None and isinstance(
|
||||
arg_signatures[i],
|
||||
(
|
||||
type(torch._inductor.codegen.wrapper.SymbolicCallArg),
|
||||
type(int),
|
||||
type(float),
|
||||
type(bool),
|
||||
),
|
||||
):
|
||||
V.graph.wrapper_code.writeline(
|
||||
f'printf("[ {launch_prefix} - {kernel_name} - {arg}: %ld ]", {arg}); printf("\\\\n");'
|
||||
)
|
||||
else:
|
||||
if arg_signatures is None and self.kernel_type in ("cpp", "extern"):
|
||||
V.graph.wrapper_code.writeline(
|
||||
f'aoti_torch_print_tensor_handle({arg}, "{launch_prefix} - {kernel_name} - {arg}");'
|
||||
)
|
||||
else:
|
||||
V.graph.wrapper_code.writeline(
|
||||
f'_print_debugging_tensor_value_info("inductor: {launch_prefix} - {kernel_name} - {arg}", {arg})'
|
||||
)
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,816 @@
|
||||
# mypy: allow-untyped-defs
|
||||
from __future__ import annotations
|
||||
|
||||
import collections
|
||||
import dataclasses
|
||||
import itertools
|
||||
import pprint
|
||||
from typing import Any, Protocol, TYPE_CHECKING
|
||||
|
||||
import sympy
|
||||
|
||||
import torch
|
||||
from torch.fx.experimental.symbolic_shapes import free_unbacked_symbols
|
||||
from torch.utils._ordered_set import OrderedSet
|
||||
|
||||
from .. import config
|
||||
from ..utils import _align, align, cache_on_self, CachedMethod, IndentedBuffer
|
||||
from ..virtualized import V
|
||||
from .wrapper import (
|
||||
AllocateLine,
|
||||
BufferLike,
|
||||
FreeIfNotReusedLine,
|
||||
MemoryPlanningLine,
|
||||
NullLine,
|
||||
ReuseLine,
|
||||
)
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Iterable
|
||||
|
||||
|
||||
@dataclasses.dataclass
|
||||
class LiveRange:
|
||||
"""
|
||||
A range where a given tensor is live. Begin and end are both counters
|
||||
representing points in the program of grouped memory operations.
|
||||
Begin is inclusive, end is exclusive.
|
||||
|
||||
Invariant: begin <= end
|
||||
"""
|
||||
|
||||
begin: float # int | +/-inf
|
||||
end: float # int | +/-inf
|
||||
|
||||
def contains(self, other: LiveRange):
|
||||
"""Is other entirely within self"""
|
||||
return self.begin <= other.begin and other.end <= self.end
|
||||
|
||||
def join(self, other: LiveRange):
|
||||
"""Combine two ranges using a union operation"""
|
||||
return LiveRange(min(self.begin, other.begin), max(self.end, other.end))
|
||||
|
||||
def __len__(self):
|
||||
return self.end - self.begin
|
||||
|
||||
|
||||
class LiveRanges:
|
||||
"""
|
||||
A collection of LiveRange regions, allowing for non-contiguous
|
||||
live regions.
|
||||
|
||||
Invariant: LiveRanges.ranges is in sorted order and non-overlapping
|
||||
"""
|
||||
|
||||
def __init__(self, ranges: Iterable[LiveRange]):
|
||||
ranges = [*sorted(ranges, key=lambda x: x.begin)]
|
||||
self.ranges = ranges[:1]
|
||||
for r in ranges[1:]:
|
||||
assert self.ranges[-1].begin <= r.begin
|
||||
if self.ranges[-1].end >= r.begin:
|
||||
self.ranges[-1] = LiveRange.join(self.ranges[-1], r)
|
||||
else:
|
||||
self.ranges.append(r)
|
||||
|
||||
def overlaps(self, other: LiveRanges):
|
||||
"""Check if any pair of ranges in self and other overlap"""
|
||||
left = collections.deque(self.ranges)
|
||||
right = collections.deque(other.ranges)
|
||||
while left and right:
|
||||
if left[0].begin > right[0].begin:
|
||||
left, right = right, left
|
||||
assert left[0].begin <= right[0].begin
|
||||
if left[0].end > right[0].begin:
|
||||
return True
|
||||
left.popleft()
|
||||
return False
|
||||
|
||||
@property
|
||||
def begin(self):
|
||||
return self.ranges[0].begin
|
||||
|
||||
@property
|
||||
def end(self):
|
||||
return self.ranges[-1].end
|
||||
|
||||
def __repr__(self):
|
||||
return f"{self.__class__.__name__}([{', '.join(map(repr, self.ranges))}])"
|
||||
|
||||
|
||||
class AllocationTreeNode:
|
||||
"""
|
||||
Abstract base class for nodes in allocation pool.
|
||||
"""
|
||||
|
||||
def allocate(self, block: Allocation, is_last: bool) -> bool:
|
||||
"""
|
||||
Try to assign block to a memory location in this bool. Return True if
|
||||
an assignment was made.
|
||||
"""
|
||||
return False
|
||||
|
||||
def get_live_ranges(self) -> LiveRanges:
|
||||
"""Aggregate LiveRanges for all objects below this in tree"""
|
||||
raise NotImplementedError
|
||||
|
||||
def get_size_hint(self) -> int:
|
||||
"""Number of bytes used for example inputs"""
|
||||
raise NotImplementedError
|
||||
|
||||
def get_symbolic_size(self) -> sympy.Expr:
|
||||
"""Number of bytes needed at runtime"""
|
||||
raise NotImplementedError
|
||||
|
||||
def finalize(self, pool, offset) -> AllocationTreeNode:
|
||||
"""Called after all allocations have been made"""
|
||||
return self
|
||||
|
||||
def is_empty(self):
|
||||
return False
|
||||
|
||||
|
||||
@dataclasses.dataclass
|
||||
class Allocation(AllocationTreeNode):
|
||||
"""
|
||||
Represents memory allocated to a given node in the allocation pool.
|
||||
"""
|
||||
|
||||
node: BufferLike
|
||||
live_range: LiveRange
|
||||
size_hint: int
|
||||
symbolic_size: sympy.Expr
|
||||
allocated: bool = False
|
||||
pool: AllocationPool | None = None
|
||||
offset: sympy.Expr | None = None
|
||||
earliest_available: float | None = None
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
has_unbacked_sym = False
|
||||
for s in self.node.get_layout().size:
|
||||
if free_unbacked_symbols(s):
|
||||
has_unbacked_sym = True
|
||||
break
|
||||
|
||||
if has_unbacked_sym:
|
||||
self.earliest_available = self.get_live_ranges().begin
|
||||
|
||||
@property
|
||||
def device(self):
|
||||
return self.node.get_device()
|
||||
|
||||
def get_live_ranges(self):
|
||||
return LiveRanges([self.live_range])
|
||||
|
||||
def get_size_hint(self):
|
||||
return self.size_hint
|
||||
|
||||
def get_symbolic_size(self):
|
||||
return self.symbolic_size
|
||||
|
||||
def mark_allocated(self):
|
||||
assert not self.allocated
|
||||
self.allocated = True
|
||||
|
||||
def finalize(self, pool, offset):
|
||||
assert self.pool is None and self.offset is None
|
||||
self.pool = pool
|
||||
self.offset = offset
|
||||
return self
|
||||
|
||||
def codegen_alloc_from_pool(self, wrapper):
|
||||
assert self.pool
|
||||
node = self.node
|
||||
shape = tuple(node.get_size())
|
||||
stride = tuple(node.get_stride())
|
||||
return wrapper.codegen_alloc_from_pool(
|
||||
self.pool.name, self.offset, node.get_dtype(), shape, stride
|
||||
)
|
||||
|
||||
def __repr__(self):
|
||||
return (
|
||||
f"{self.__class__.__name__}("
|
||||
f"node={self.node.get_name()}, "
|
||||
f"live_range={self.live_range}, "
|
||||
f"size_hint={self.size_hint}, "
|
||||
f"symbolic_size={self.symbolic_size}, "
|
||||
f"pool={self.pool.name if self.pool else None}, "
|
||||
f"offset={self.offset})"
|
||||
)
|
||||
|
||||
def get_earliest_available(self):
|
||||
return self.earliest_available
|
||||
|
||||
|
||||
@dataclasses.dataclass
|
||||
class Empty(AllocationTreeNode):
|
||||
"""
|
||||
Placeholder to represent empty space in the allocation pool.
|
||||
Only exists to get the size_hint correct in parent nodes.
|
||||
"""
|
||||
|
||||
size_hint: int
|
||||
|
||||
def get_live_ranges(self):
|
||||
return LiveRanges([])
|
||||
|
||||
def get_size_hint(self):
|
||||
return self.size_hint
|
||||
|
||||
def get_symbolic_size(self):
|
||||
return 0
|
||||
|
||||
def is_empty(self):
|
||||
return True
|
||||
|
||||
|
||||
class MemorySplitProtocol(Protocol):
|
||||
get_live_ranges: CachedMethod[[], LiveRanges]
|
||||
get_size_hint: CachedMethod[[], int]
|
||||
get_symbolic_size: CachedMethod[[], sympy.Expr]
|
||||
|
||||
def _allocate(self, block: Allocation, is_last: bool) -> bool: ...
|
||||
|
||||
|
||||
class ClearCacheOnAllocateMixin(MemorySplitProtocol):
|
||||
"""
|
||||
Helper to assist in caching get_live_ranges, get_size_hint, and
|
||||
get_symbolic_size.
|
||||
"""
|
||||
|
||||
def allocate(self, block: Allocation, is_last: bool):
|
||||
is_allocated = self._allocate(block, is_last)
|
||||
if is_allocated:
|
||||
self.clear_cache()
|
||||
return is_allocated
|
||||
|
||||
def clear_cache(self):
|
||||
self.get_live_ranges.clear_cache(self)
|
||||
self.get_size_hint.clear_cache(self)
|
||||
self.get_symbolic_size.clear_cache(self)
|
||||
|
||||
|
||||
@dataclasses.dataclass
|
||||
class TemporalSplit(ClearCacheOnAllocateMixin, AllocationTreeNode):
|
||||
"""
|
||||
Contains a list of allocations not overlapping in LiveRanges.
|
||||
|
||||
Invariant: no pair (a,b) in self.allocations will have:
|
||||
a.get_live_ranges().overlaps(b.get_live_ranges())
|
||||
"""
|
||||
|
||||
allocations: list[AllocationTreeNode]
|
||||
|
||||
def _allocate(self, block: Allocation, is_last: bool):
|
||||
slot_size = self.get_size_hint()
|
||||
block_size = block.get_size_hint()
|
||||
if not is_last and block_size > slot_size:
|
||||
return False # doesn't fit
|
||||
|
||||
block_live = block.get_live_ranges()
|
||||
overlapping = [
|
||||
s for s in self.allocations if s.get_live_ranges().overlaps(block_live)
|
||||
]
|
||||
if len(overlapping) > 1:
|
||||
# TODO(jansel): we could try harder here by merging overlapping in space
|
||||
return False
|
||||
elif len(overlapping) == 1:
|
||||
return overlapping[0].allocate(block, is_last)
|
||||
else:
|
||||
block.mark_allocated()
|
||||
|
||||
if len(self.allocations) == 1 and isinstance(self.allocations[-1], Empty):
|
||||
self.allocations.pop()
|
||||
|
||||
if slot_size == block_size:
|
||||
# perfect fit
|
||||
self.allocations.append(block)
|
||||
elif slot_size > block_size:
|
||||
self.allocations.append(
|
||||
SpatialSplit.create(block, slot_size - block_size)
|
||||
)
|
||||
else: # grow this allocation
|
||||
assert is_last
|
||||
self.allocations = [
|
||||
*(
|
||||
SpatialSplit.create(a, block_size - slot_size)
|
||||
for a in self.allocations
|
||||
),
|
||||
block,
|
||||
]
|
||||
return True
|
||||
|
||||
@cache_on_self
|
||||
def get_live_ranges(self) -> LiveRanges:
|
||||
return LiveRanges(
|
||||
itertools.chain.from_iterable(
|
||||
x.get_live_ranges().ranges for x in self.allocations
|
||||
)
|
||||
)
|
||||
|
||||
@cache_on_self
|
||||
def get_size_hint(self) -> int:
|
||||
if not self.allocations:
|
||||
return 0
|
||||
return max(x.get_size_hint() for x in self.allocations)
|
||||
|
||||
@cache_on_self
|
||||
def get_symbolic_size(self) -> sympy.Expr:
|
||||
if not self.allocations:
|
||||
return 0 # type: ignore[return-value]
|
||||
return sympy.Max(*[x.get_symbolic_size() for x in self.allocations])
|
||||
|
||||
def is_empty(self):
|
||||
return len(self.allocations) == 1 and self.allocations[0].is_empty()
|
||||
|
||||
def finalize(self, pool, offset):
|
||||
self.allocations = [block.finalize(pool, offset) for block in self.allocations]
|
||||
self.clear_cache()
|
||||
if len(self.allocations) == 1:
|
||||
return self.allocations[0]
|
||||
return self
|
||||
|
||||
|
||||
@dataclasses.dataclass
|
||||
class SpatialSplit(ClearCacheOnAllocateMixin, AllocationTreeNode):
|
||||
"""
|
||||
Contains two allocations, left and right, that do not overlap in space.
|
||||
Right will be allocated immediately after left in memory.
|
||||
"""
|
||||
|
||||
left: TemporalSplit
|
||||
right: TemporalSplit
|
||||
|
||||
@staticmethod
|
||||
def create(left, extra_space):
|
||||
assert isinstance(left, AllocationTreeNode)
|
||||
assert isinstance(extra_space, int) and extra_space >= 1
|
||||
return SpatialSplit(TemporalSplit([left]), TemporalSplit([Empty(extra_space)]))
|
||||
|
||||
def _allocate(self, block: Allocation, is_last: bool):
|
||||
return self.left.allocate(block, False) or self.right.allocate(block, is_last)
|
||||
|
||||
@cache_on_self
|
||||
def get_live_ranges(self):
|
||||
return LiveRanges(
|
||||
itertools.chain(
|
||||
self.left.get_live_ranges().ranges, self.right.get_live_ranges().ranges
|
||||
)
|
||||
)
|
||||
|
||||
@cache_on_self
|
||||
def get_size_hint(self) -> int:
|
||||
return _align(self.left.get_size_hint()) + self.right.get_size_hint()
|
||||
|
||||
@cache_on_self
|
||||
def get_symbolic_size(self) -> sympy.Expr:
|
||||
return align(self.left.get_symbolic_size()) + self.right.get_symbolic_size()
|
||||
|
||||
def finalize(self, pool, offset):
|
||||
self.left = self.left.finalize(pool, offset)
|
||||
self.right = self.right.finalize(
|
||||
pool, offset + align(self.left.get_symbolic_size())
|
||||
)
|
||||
self.clear_cache()
|
||||
if self.right.is_empty():
|
||||
return self.left
|
||||
return self
|
||||
|
||||
|
||||
@dataclasses.dataclass
|
||||
class AllocationPool:
|
||||
"""
|
||||
Represents a pool of allocations that will be generated by a single
|
||||
call to torch.empty.
|
||||
"""
|
||||
|
||||
device: torch.device
|
||||
root: TemporalSplit
|
||||
can_expand: bool = True
|
||||
restrict_live_range: LiveRange | None = None
|
||||
name: str | None = None
|
||||
names_to_del: list[str] = dataclasses.field(default_factory=list)
|
||||
creation_cache: dict[str, str] = dataclasses.field(default_factory=dict)
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
for block in self.root.allocations:
|
||||
if isinstance(block, Allocation):
|
||||
self.update_restrict_live_range(block)
|
||||
|
||||
def allocate(self, block: Allocation, is_last: bool):
|
||||
if (
|
||||
self.restrict_live_range is not None
|
||||
and not self.restrict_live_range.contains(block.live_range)
|
||||
):
|
||||
return False
|
||||
|
||||
block_earliest_available = block.get_earliest_available()
|
||||
pool_begin = self.root.get_live_ranges().begin
|
||||
if block_earliest_available and block_earliest_available > pool_begin:
|
||||
return False
|
||||
|
||||
is_last = self.can_expand and is_last
|
||||
if self.root.allocate(block, is_last):
|
||||
self.update_restrict_live_range(block)
|
||||
return True
|
||||
|
||||
if is_last:
|
||||
return self.allocate_at_end(block)
|
||||
|
||||
return False
|
||||
|
||||
def update_restrict_live_range(self, block: Allocation):
|
||||
if block_earliest_available := block.get_earliest_available():
|
||||
if self.restrict_live_range is None:
|
||||
self.restrict_live_range = LiveRange(
|
||||
block_earliest_available, float("inf")
|
||||
)
|
||||
else:
|
||||
self.restrict_live_range = LiveRange(
|
||||
min(self.restrict_live_range.begin, block_earliest_available),
|
||||
self.restrict_live_range.end,
|
||||
)
|
||||
|
||||
def allocate_at_end(self, block):
|
||||
block.mark_allocated()
|
||||
self.root = TemporalSplit([SpatialSplit(self.root, TemporalSplit([block]))])
|
||||
self.update_restrict_live_range(block)
|
||||
return True
|
||||
|
||||
def finalize(self, name):
|
||||
assert not self.name
|
||||
self.name = name
|
||||
self.names_to_del.append(name)
|
||||
self.root.finalize(self, 0)
|
||||
|
||||
def codegen_create(self, wrapper, code: IndentedBuffer):
|
||||
assert self.name
|
||||
nbytes = self.root.get_symbolic_size()
|
||||
for block in self.root.allocations:
|
||||
if isinstance(block, Allocation) and nbytes == block.get_symbolic_size():
|
||||
node = block.node
|
||||
code.writeline(
|
||||
wrapper.make_allocation(
|
||||
self.name,
|
||||
device=self.device,
|
||||
dtype=node.get_dtype(),
|
||||
shape=tuple(node.get_size()),
|
||||
stride=tuple(node.get_stride()),
|
||||
)
|
||||
)
|
||||
return
|
||||
else:
|
||||
code.writeline(
|
||||
wrapper.make_allocation(
|
||||
self.name,
|
||||
device=self.device,
|
||||
dtype=torch.uint8,
|
||||
shape=(nbytes,),
|
||||
stride=(1,),
|
||||
)
|
||||
)
|
||||
|
||||
def codegen_destroy(self, wrapper, code: IndentedBuffer):
|
||||
code.writeline(wrapper.make_free_by_names(self.names_to_del))
|
||||
|
||||
def __eq__(self, other):
|
||||
return self is other
|
||||
|
||||
def __hash__(self):
|
||||
return id(self)
|
||||
|
||||
|
||||
@dataclasses.dataclass
|
||||
class AllocationPools:
|
||||
"""
|
||||
Collection of many AllocationPool objects grouped by device.
|
||||
"""
|
||||
|
||||
device_to_pools: dict[torch.device, list[AllocationPool]] = dataclasses.field(
|
||||
default_factory=dict
|
||||
)
|
||||
|
||||
def get_pools(self, block):
|
||||
if block.device not in self.device_to_pools:
|
||||
self.device_to_pools[block.device] = []
|
||||
return self.device_to_pools[block.device]
|
||||
|
||||
def allocate(self, block: Allocation):
|
||||
pools = self.get_pools(block)
|
||||
|
||||
for pool in pools:
|
||||
if pool.allocate(block, is_last=pool is pools[-1]):
|
||||
return
|
||||
|
||||
# everything is full, make a new pool
|
||||
pools.append(
|
||||
AllocationPool(
|
||||
block.device,
|
||||
TemporalSplit([block]),
|
||||
can_expand=config.memory_pool != "none",
|
||||
)
|
||||
)
|
||||
block.mark_allocated()
|
||||
|
||||
def allocate_output(self, block: Allocation):
|
||||
"""Outputs get different pools so memory gets freed properly"""
|
||||
pools = self.get_pools(block)
|
||||
if pools and config.memory_pool in ("outputs", "combined"):
|
||||
pools[-1].allocate_at_end(block)
|
||||
else:
|
||||
# create a new pool
|
||||
block.mark_allocated()
|
||||
pools.append(
|
||||
AllocationPool(
|
||||
block.device,
|
||||
TemporalSplit([block]),
|
||||
can_expand=config.memory_pool == "combined",
|
||||
)
|
||||
)
|
||||
|
||||
def finalize(self):
|
||||
"""Called at the end of allocation process"""
|
||||
for i, pool in enumerate(
|
||||
itertools.chain.from_iterable(self.device_to_pools.values())
|
||||
):
|
||||
pool.finalize(f"pool{i}")
|
||||
|
||||
def pprint(self):
|
||||
for pool in itertools.chain.from_iterable(self.device_to_pools.values()):
|
||||
print()
|
||||
print(pool.name)
|
||||
print(pool.root.get_live_ranges())
|
||||
pprint.pprint(pool.root)
|
||||
|
||||
|
||||
class BufferGroup:
|
||||
"""
|
||||
Due to inplace reuse an allocated buffer can have many names.
|
||||
This tracks these collections of buffers sharing underlying memory.
|
||||
"""
|
||||
|
||||
def __init__(self, node: BufferLike):
|
||||
self.node = node
|
||||
self.names = [node.get_name()]
|
||||
self.is_output = False
|
||||
self.allocation: Allocation | None = None
|
||||
self.live_range = LiveRange(float("inf"), -float("inf"))
|
||||
|
||||
def update_usage(self, timestep: int):
|
||||
"""Expand self.live_range to include timestep"""
|
||||
self.live_range = LiveRange(
|
||||
min(timestep, self.live_range.begin),
|
||||
max(timestep, self.live_range.end),
|
||||
)
|
||||
|
||||
def sym_nbytes(self):
|
||||
return self.node.get_layout().storage_size() * self.node.get_dtype().itemsize
|
||||
|
||||
def make_allocation(self):
|
||||
assert not self.allocation, "multiple allocations"
|
||||
assert isinstance(self.live_range.begin, int), "live ranges not computed"
|
||||
nbytes = self.sym_nbytes()
|
||||
# For now, fallback value will be used if we encounter an unbacked SymInt. The longer-term plan is to have
|
||||
# size_hint() use better heuristics for unbackeds, at which point the fallback value will be ignored.
|
||||
size_hint = V.graph.sizevars.optimization_hint(nbytes, fallback=64)
|
||||
self.allocation = Allocation(
|
||||
self.node,
|
||||
self.live_range,
|
||||
size_hint=size_hint,
|
||||
symbolic_size=nbytes,
|
||||
)
|
||||
|
||||
def __repr__(self):
|
||||
return (
|
||||
f"{self.__class__.__name__}({self.names!r}, is_output={self.is_output}, "
|
||||
f"live_range={self.live_range}"
|
||||
)
|
||||
|
||||
|
||||
@dataclasses.dataclass
|
||||
class PoolMemoryPlanningLine(MemoryPlanningLine):
|
||||
"""Abstract base class for {Alloc,Dealloc}FromPoolLine"""
|
||||
|
||||
group: BufferGroup
|
||||
timestep: int | None = None
|
||||
|
||||
@property
|
||||
def node(self):
|
||||
return self.group.node
|
||||
|
||||
|
||||
@dataclasses.dataclass
|
||||
class AllocFromPoolLine(PoolMemoryPlanningLine):
|
||||
"""Similar to AllocationLine, but takes memory from a pool"""
|
||||
|
||||
is_first_pool_usage: bool = False
|
||||
|
||||
def codegen(self, code: IndentedBuffer):
|
||||
allocation = self.group.allocation
|
||||
assert allocation and allocation.pool
|
||||
pool = allocation.pool
|
||||
name = self.node.get_name()
|
||||
|
||||
if self.is_first_pool_usage:
|
||||
pool.codegen_create(self.wrapper, code)
|
||||
|
||||
pool.names_to_del.extend(self.group.names)
|
||||
alloc_from_pool, allocation_lines_to_write = allocation.codegen_alloc_from_pool(
|
||||
self.wrapper
|
||||
)
|
||||
code.writelines(allocation_lines_to_write)
|
||||
if alloc_from_pool in pool.creation_cache:
|
||||
code.writeline(
|
||||
self.wrapper.make_tensor_alias(
|
||||
name, pool.creation_cache[alloc_from_pool], "alloc"
|
||||
)
|
||||
)
|
||||
else:
|
||||
pool.creation_cache[alloc_from_pool] = name
|
||||
code.writeline(
|
||||
f"{self.wrapper.declare}{name} = {alloc_from_pool}{self.wrapper.ending}"
|
||||
)
|
||||
|
||||
|
||||
@dataclasses.dataclass
|
||||
class DeallocFromPoolLine(PoolMemoryPlanningLine):
|
||||
"""Similar to FreeIfNotReusedLine, but takes memory from a pool"""
|
||||
|
||||
is_last_pool_usage: bool = False
|
||||
|
||||
def codegen(self, code: IndentedBuffer):
|
||||
if self.is_last_pool_usage:
|
||||
assert self.group.allocation and self.group.allocation.pool
|
||||
self.group.allocation.pool.codegen_destroy(self.wrapper, code)
|
||||
|
||||
|
||||
@dataclasses.dataclass
|
||||
class MemoryPlanner:
|
||||
"""
|
||||
Coordination object to run memory planning passes during wrapper
|
||||
codegen.
|
||||
"""
|
||||
|
||||
wrapper: Any
|
||||
pools: AllocationPools = dataclasses.field(default_factory=AllocationPools)
|
||||
buffer_groups: list[BufferGroup] | None = None
|
||||
|
||||
def plan(self, lines: list[Any]) -> list[Any]:
|
||||
"""Call all the memory planning passes in sequence"""
|
||||
lines = [*lines]
|
||||
self.drop_removed_buffers(lines)
|
||||
self.convert_to_pool_lines(lines)
|
||||
self.compute_live_ranges(lines)
|
||||
self.allocate_groups()
|
||||
self.mark_first_last_usage(lines)
|
||||
return lines
|
||||
|
||||
def drop_removed_buffers(self, lines):
|
||||
"""
|
||||
Replace any memory planning lines in V.graph.removed_buffers with NullLine
|
||||
"""
|
||||
# drop any removed buffers
|
||||
for i, line in enumerate(lines):
|
||||
if isinstance(line, (AllocateLine, FreeIfNotReusedLine, ReuseLine)):
|
||||
if line.node.get_name() in V.graph.removed_buffers:
|
||||
lines[i] = NullLine(self.wrapper)
|
||||
|
||||
def compute_buffer_groups(self, lines):
|
||||
"""
|
||||
Populates self.buffer_groups with BufferGroup objects that join
|
||||
allocations with common storage (due to inplace reuse) into a
|
||||
single object.
|
||||
"""
|
||||
name_to_group = {}
|
||||
for line in lines:
|
||||
if isinstance(line, AllocateLine):
|
||||
name = line.node.get_name()
|
||||
assert name not in name_to_group
|
||||
name_to_group[name] = BufferGroup(line.node)
|
||||
elif isinstance(line, ReuseLine):
|
||||
old_name = line.node.get_name()
|
||||
new_name = line.reused_as.get_name()
|
||||
assert new_name not in name_to_group
|
||||
# TODO(jansel): we should support reusing buffers created via ExternKernelAlloc
|
||||
if old_name in name_to_group:
|
||||
name_to_group[old_name].names.append(new_name)
|
||||
name_to_group[new_name] = name_to_group[old_name]
|
||||
|
||||
outputs = OrderedSet(V.graph.get_output_names())
|
||||
unique_groups = [*{id(g): g for g in name_to_group.values()}.values()]
|
||||
for group in unique_groups:
|
||||
group.is_output = any(x in outputs for x in group.names)
|
||||
|
||||
assert self.buffer_groups is None
|
||||
self.buffer_groups = unique_groups
|
||||
return name_to_group
|
||||
|
||||
def convert_to_pool_lines(self, lines):
|
||||
"""
|
||||
Convert AllocateLine/FreeIfNotReusedLine/ReuseLine into their
|
||||
pool-based counterparts.
|
||||
"""
|
||||
name_to_group = self.compute_buffer_groups(lines)
|
||||
for i, line in enumerate(lines):
|
||||
if isinstance(line, AllocateLine):
|
||||
if line.node.get_name() in name_to_group:
|
||||
lines[i] = AllocFromPoolLine(
|
||||
self.wrapper, name_to_group[line.node.get_name()]
|
||||
)
|
||||
elif isinstance(line, FreeIfNotReusedLine):
|
||||
assert not line.is_reused
|
||||
if line.node.get_name() in name_to_group:
|
||||
lines[i] = DeallocFromPoolLine(
|
||||
self.wrapper, name_to_group[line.node.get_name()]
|
||||
)
|
||||
elif isinstance(line, ReuseLine):
|
||||
if line.node.get_name() in name_to_group:
|
||||
line.delete_old = False
|
||||
|
||||
def compute_live_ranges(self, lines):
|
||||
"""Populate every BufferGroup.live_ranges field based on first/last usage"""
|
||||
timestep = 0
|
||||
worklist = collections.deque(lines)
|
||||
while worklist:
|
||||
if isinstance(worklist[0], MemoryPlanningLine):
|
||||
timestep += 1
|
||||
while worklist and isinstance(worklist[0], MemoryPlanningLine):
|
||||
line = worklist.popleft()
|
||||
if isinstance(line, PoolMemoryPlanningLine):
|
||||
line.group.update_usage(timestep)
|
||||
line.timestep = timestep
|
||||
else:
|
||||
worklist.popleft()
|
||||
|
||||
timestep += 1
|
||||
assert self.buffer_groups is not None
|
||||
for group in self.buffer_groups:
|
||||
if group.is_output:
|
||||
group.update_usage(timestep)
|
||||
|
||||
def allocate_groups(self):
|
||||
"""
|
||||
Assign every allocation to a specific location in a specific AllocationPool.
|
||||
"""
|
||||
assert config.memory_pool in ("none", "intermediates", "outputs", "combined")
|
||||
assert self.buffer_groups is not None
|
||||
|
||||
for group in self.buffer_groups:
|
||||
group.make_allocation()
|
||||
|
||||
outputs: list[Allocation] = []
|
||||
intermediates: list[Allocation] = []
|
||||
for group in self.buffer_groups:
|
||||
assert group.allocation
|
||||
if group.is_output and config.memory_pool != "combined":
|
||||
outputs.append(group.allocation)
|
||||
else:
|
||||
intermediates.append(group.allocation)
|
||||
|
||||
for block in sorted(
|
||||
outputs,
|
||||
key=lambda x: (
|
||||
x.size_hint,
|
||||
-len(x.live_range),
|
||||
),
|
||||
):
|
||||
self.pools.allocate_output(block)
|
||||
|
||||
for block in sorted(
|
||||
intermediates,
|
||||
key=lambda x: (
|
||||
-x.size_hint,
|
||||
-len(x.live_range),
|
||||
),
|
||||
):
|
||||
self.pools.allocate(block)
|
||||
|
||||
self.pools.finalize()
|
||||
|
||||
def mark_first_last_usage(self, lines):
|
||||
"""
|
||||
Populate the AllocFromPoolLine.is_first_pool_usage and
|
||||
DeallocFromPoolLine.is_last_pool_usage fields so that pools
|
||||
are created/destroyed.
|
||||
"""
|
||||
seen = OrderedSet[AllocationPool]()
|
||||
for line in lines:
|
||||
if isinstance(line, AllocFromPoolLine):
|
||||
assert line.group.allocation
|
||||
pool = line.group.allocation.pool
|
||||
assert pool is not None
|
||||
if pool not in seen:
|
||||
line.is_first_pool_usage = True
|
||||
seen.add(pool)
|
||||
|
||||
seen = OrderedSet[AllocationPool]()
|
||||
for line in reversed(lines):
|
||||
if isinstance(line, DeallocFromPoolLine):
|
||||
assert line.group.allocation
|
||||
pool = line.group.allocation.pool
|
||||
assert pool is not None
|
||||
if pool not in seen:
|
||||
line.is_last_pool_usage = (
|
||||
pool.root.get_live_ranges().end <= line.timestep
|
||||
)
|
||||
seen.add(pool)
|
||||
File diff suppressed because it is too large
Load Diff
+24
@@ -0,0 +1,24 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from .common import DeviceOpOverrides, register_device_op_overrides
|
||||
|
||||
|
||||
class MPSDeviceOpOverrides(DeviceOpOverrides):
|
||||
def device_guard(self, device_idx: int) -> str:
|
||||
assert device_idx == 0
|
||||
return "torch._ops.contextlib.nullcontext()"
|
||||
|
||||
def set_device(self, device_idx: int) -> str:
|
||||
assert device_idx == 0
|
||||
return "pass # MPS set device"
|
||||
|
||||
def kernel_driver(self) -> str:
|
||||
return """
|
||||
#include <ATen/native/mps/MetalShaderLibrary.h>
|
||||
"""
|
||||
|
||||
def cpp_kernel_type(self) -> str:
|
||||
return "MTLFunction_t"
|
||||
|
||||
|
||||
register_device_op_overrides("mps", MPSDeviceOpOverrides())
|
||||
+20
@@ -0,0 +1,20 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from ..common import DeviceOpOverrides, register_device_op_overrides
|
||||
|
||||
|
||||
class MTIADeviceOpOverrides(DeviceOpOverrides):
|
||||
def import_get_raw_stream_as(self, name: str) -> str:
|
||||
return f"from torch._C import _mtia_getCurrentRawStream as {name}"
|
||||
|
||||
def set_device(self, device_idx: int) -> str:
|
||||
return f"torch.mtia.set_device({device_idx})"
|
||||
|
||||
def synchronize(self) -> str:
|
||||
return "torch.mtia.synchronize()"
|
||||
|
||||
def device_guard(self, device_idx: int) -> str:
|
||||
return f"torch.mtia.device({device_idx})"
|
||||
|
||||
|
||||
register_device_op_overrides("mtia", MTIADeviceOpOverrides())
|
||||
@@ -0,0 +1,610 @@
|
||||
# mypy: allow-untyped-defs
|
||||
import functools
|
||||
import logging
|
||||
import math
|
||||
import os
|
||||
import pathlib
|
||||
from typing import Any
|
||||
|
||||
from torch._inductor.ir import MultiTemplateBuffer
|
||||
from torch._inductor.metrics import get_metric_table, is_metric_table_enabled
|
||||
from torch.utils._ordered_set import OrderedSet
|
||||
|
||||
from .. import config
|
||||
from ..codecache import code_hash, CodeCacheFuture, get_path, write_atomic
|
||||
from ..runtime.benchmarking import benchmarker
|
||||
from ..utils import cache_on_self, IndentedBuffer
|
||||
from ..virtualized import V
|
||||
from .common import TensorArg, WorkspaceArg
|
||||
|
||||
|
||||
log = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class MultiKernelState:
|
||||
"""
|
||||
Maintain state of multi-kernel compilation so we don't define duplicated
|
||||
multi-kernel for the same set of sub-kernels.
|
||||
|
||||
V.graph.wrapper_code has a reference to MultiKernelState instance.
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
self.subkernel_to_kernel_name = {}
|
||||
self.kernel_defs = IndentedBuffer()
|
||||
|
||||
def define_kernel(
|
||||
self,
|
||||
kernels: list[Any],
|
||||
kernel_shape_keys: list[None | tuple[tuple[int, ...], ...]] | None = None,
|
||||
) -> str:
|
||||
"""
|
||||
Previously we name the multi kernel as "multi_kernel_{kernel_names[0]}".
|
||||
This has some minor issue.
|
||||
|
||||
E.g. for persistent reduction https://gist.github.com/shunting314/39e7c00ff8bb2055942ed5a3255d61ca ,
|
||||
there are 2 flavors of non-persistent reduction:
|
||||
https://gist.github.com/shunting314/056d43d35907e87efb883970b35c17d4
|
||||
and
|
||||
https://gist.github.com/shunting314/02ee753b65c513c54e695626afe682bd
|
||||
|
||||
The only different is cache eviction policy.
|
||||
|
||||
We should name the multi-kernel differently in these 2 cases.
|
||||
|
||||
kernels:
|
||||
A list of kernels
|
||||
kernel_shape_keys:
|
||||
Specified for size-hint multi-kernels.
|
||||
Each list element is a shape key, corresponding to the concrete input & output size hints each kernel was tuned for.
|
||||
"""
|
||||
# Prevent circular import
|
||||
from ..select_algorithm import TritonTemplateKernel
|
||||
|
||||
kernel_names = tuple(k.kernel_name for k in kernels)
|
||||
if kernel_names in self.subkernel_to_kernel_name:
|
||||
return self.subkernel_to_kernel_name[kernel_names]
|
||||
|
||||
# name the multi kernel based on the first kernel
|
||||
multi_kernel_name = f"multi_kernel_{len(self.subkernel_to_kernel_name)}"
|
||||
self.subkernel_to_kernel_name[kernel_names] = multi_kernel_name
|
||||
|
||||
if V.graph.cpp_wrapper and not config.triton.autotune_at_compile_time:
|
||||
# we should not generate any python code for multi-kernel during
|
||||
# the second pass of cpp-wrapper.
|
||||
return multi_kernel_name
|
||||
|
||||
arg_index: dict[int, list[slice]] = {}
|
||||
_, call_args, _, arg_types = kernels[0].args.python_argdefs()
|
||||
if isinstance(kernels[0], TritonTemplateKernel) and isinstance(
|
||||
kernels[0].output_node, MultiTemplateBuffer
|
||||
):
|
||||
for i, kernel in enumerate(kernels):
|
||||
additional_call_args, _ = kernel.additional_call_args_and_types()
|
||||
if i not in arg_index:
|
||||
arg_index[i] = []
|
||||
arg_index[i].append(slice(0, len(call_args)))
|
||||
arg_index[i].append(
|
||||
slice(
|
||||
len(call_args) + i * len(additional_call_args),
|
||||
len(call_args) + (i + 1) * len(additional_call_args),
|
||||
)
|
||||
)
|
||||
else:
|
||||
kernels[0].add_numel_to_call_args(multi_kernel_name, call_args, arg_types)
|
||||
for i in range(len(kernels)):
|
||||
arg_index[i] = [slice(0, len(call_args))]
|
||||
|
||||
keyed_by_sizes = kernel_shape_keys is not None
|
||||
buf = self.kernel_defs
|
||||
buf.writeline("")
|
||||
buf.writeline("arg_index = {")
|
||||
for key, slice_list in arg_index.items():
|
||||
slice_reprs = ", ".join(repr(s) for s in slice_list)
|
||||
buf.writeline(f" {key}: [{slice_reprs}],")
|
||||
buf.writeline("}")
|
||||
|
||||
if not keyed_by_sizes: # no size hint keys, just call with list of kernels
|
||||
buf.writeline(
|
||||
f"{multi_kernel_name} = async_compile.multi_kernel({multi_kernel_name!r}, ["
|
||||
)
|
||||
with buf.indent():
|
||||
for name in kernel_names:
|
||||
buf.writeline(f"{name},")
|
||||
buf.writeline("], arg_index=arg_index)")
|
||||
else: # call with dict[size hint key, kernel]
|
||||
assert isinstance(kernels[0], TritonTemplateKernel)
|
||||
assert isinstance(kernel_shape_keys, list)
|
||||
assert len(kernels) == len(kernel_shape_keys)
|
||||
buf.writeline(
|
||||
f"{multi_kernel_name} = async_compile.size_hint_multi_kernel({multi_kernel_name!r}, {{"
|
||||
)
|
||||
with buf.indent():
|
||||
for shape_key, name in zip(kernel_shape_keys, kernel_names):
|
||||
buf.writeline(f"{shape_key}: {name},")
|
||||
buf.writeline("}, arg_index=arg_index)")
|
||||
|
||||
if config.triton.autotune_at_compile_time:
|
||||
V.graph.wrapper_code.src_to_kernel["\n".join(kernel_names)] = (
|
||||
multi_kernel_name
|
||||
)
|
||||
|
||||
return multi_kernel_name
|
||||
|
||||
|
||||
class MultiKernel:
|
||||
"""
|
||||
This class maintains the compile time state for multi kernels.
|
||||
|
||||
Assume we do codegen for a MultiKernel encapsulating kernel1 and kernel2.
|
||||
The generated definition for the multi-kernel will looks like:
|
||||
```
|
||||
multi_kernel_kernel1 = MultiKernelCall(
|
||||
[kernel1, kernel2], multi_kernel_definition_code
|
||||
)
|
||||
```
|
||||
|
||||
Here is a concrete example: https://gist.github.com/shunting314/d9f3fb6bc6cee3dbae005825ca196d39
|
||||
"""
|
||||
|
||||
def __init__(self, kernels):
|
||||
assert len(kernels) >= 2
|
||||
|
||||
self.kernels = kernels
|
||||
self.kernel_name = V.graph.wrapper_code.multi_kernel_state.define_kernel(
|
||||
kernels
|
||||
)
|
||||
|
||||
# need this since some code in inductor check if the kernel object has an args
|
||||
# attribute to decide if it's a non-null kernel.
|
||||
self.args = object()
|
||||
|
||||
@staticmethod
|
||||
def _merge_workspace_args(left: list[WorkspaceArg], right: list[WorkspaceArg]):
|
||||
if left == right:
|
||||
return left
|
||||
result = {x.inner_name: x for x in left}
|
||||
for arg in right:
|
||||
if arg.inner_name in result:
|
||||
result[arg.inner_name] = WorkspaceArg.maximum(
|
||||
result[arg.inner_name], arg
|
||||
)
|
||||
else:
|
||||
result[arg.inner_name] = arg
|
||||
return [*result.values()]
|
||||
|
||||
@staticmethod
|
||||
def merge_workspaces_inplace(kernels):
|
||||
if len(kernels) < 2:
|
||||
return
|
||||
# All kernels must share the same workspace
|
||||
workspace_args = functools.reduce(
|
||||
MultiKernel._merge_workspace_args,
|
||||
[kernel.args.workspace_args for kernel in kernels],
|
||||
)
|
||||
for kernel in kernels:
|
||||
kernel.args.workspace_args = workspace_args
|
||||
return workspace_args
|
||||
|
||||
def call_kernel(self, kernel_name):
|
||||
"""
|
||||
Collect the union of arguments from all subkernels as the arguments
|
||||
for the multi-kernel.
|
||||
"""
|
||||
# Prevent circular import
|
||||
from ..select_algorithm import TritonTemplateKernel
|
||||
|
||||
assert kernel_name == self.kernel_name
|
||||
V.graph.wrapper_code.write_triton_header_once()
|
||||
_, call_args, _, arg_types = self.kernels[0].args.python_argdefs()
|
||||
for kernel in self.kernels[1:]:
|
||||
_, other_call_args, _, other_arg_types = kernel.args.python_argdefs()
|
||||
assert call_args == other_call_args, (call_args, other_call_args)
|
||||
assert arg_types == other_arg_types
|
||||
|
||||
if V.graph.cpp_wrapper and not config.triton.autotune_at_compile_time:
|
||||
# for the second pass of cpp-wrapper codegen, we should call
|
||||
# the fast kernel directly
|
||||
kernel_name = MultiKernelCall.lookup_choice(self.kernel_name)
|
||||
|
||||
if isinstance(self.kernels[0], TritonTemplateKernel) and isinstance(
|
||||
self.kernels[0].output_node, MultiTemplateBuffer
|
||||
):
|
||||
# For matmuls the grid arguments are passed in as additional arguments
|
||||
# to the kernel run method. These grids change based on the various
|
||||
# parameters of the matmul. So we need to pass each kernel's grid into
|
||||
# the multi call kernel.
|
||||
multi_call_args = call_args
|
||||
multi_call_arg_types = arg_types
|
||||
for kernel in self.kernels:
|
||||
additional_call_args, additional_arg_types = (
|
||||
kernel.additional_call_args_and_types()
|
||||
)
|
||||
multi_call_args.extend(list(additional_call_args))
|
||||
multi_call_arg_types.extend(list(additional_arg_types))
|
||||
else:
|
||||
# numels for all subkernels should be the same. Use kernels[0] here
|
||||
self.kernels[0].add_numel_to_call_args(kernel_name, call_args, arg_types)
|
||||
multi_call_args = call_args
|
||||
multi_call_arg_types = arg_types
|
||||
|
||||
for ws in self.kernels[0].args.workspace_args:
|
||||
V.graph.wrapper_code.generate_workspace_allocation(ws)
|
||||
|
||||
if V.graph.cpp_wrapper:
|
||||
# We have already selected the best kernel at compile time
|
||||
# so we only have one set of call args. NB: this currently
|
||||
# doesn't work with MultiTemplateBuffer kernels. @bobrenjc93
|
||||
# will add it in a subsequent PR.
|
||||
V.graph.wrapper_code.generate_kernel_call(
|
||||
kernel_name, call_args, arg_types=arg_types
|
||||
)
|
||||
else:
|
||||
V.graph.wrapper_code.generate_kernel_call(
|
||||
kernel_name, multi_call_args, arg_types=multi_call_arg_types
|
||||
)
|
||||
|
||||
for ws in reversed(self.kernels[0].args.workspace_args):
|
||||
V.graph.wrapper_code.generate_workspace_deallocation(ws)
|
||||
|
||||
def codegen_nan_check(self):
|
||||
wrapper = V.graph.wrapper_code
|
||||
seen: OrderedSet[str] = OrderedSet()
|
||||
for k in self.kernels:
|
||||
_, call_args, precompile_args, _ = k.args.python_argdefs()
|
||||
for arg, precompile_arg in zip(call_args, precompile_args):
|
||||
if arg in seen:
|
||||
continue
|
||||
seen.add(arg)
|
||||
if isinstance(precompile_arg, TensorArg):
|
||||
line = f"assert not {arg}.isnan().any().item()"
|
||||
wrapper.writeline(line)
|
||||
line = f"assert not {arg}.isinf().any().item()"
|
||||
wrapper.writeline(line)
|
||||
|
||||
@property
|
||||
def removed_buffers(self):
|
||||
return OrderedSet.intersection(*[k.removed_buffers for k in self.kernels])
|
||||
|
||||
@property
|
||||
def inplaced_to_remove(self):
|
||||
return OrderedSet.intersection(*[k.inplaced_to_remove for k in self.kernels])
|
||||
|
||||
@property
|
||||
@cache_on_self
|
||||
def inplace_update_buffers(self):
|
||||
"""
|
||||
Make sure all kernels have the same inplace update mappings.
|
||||
"""
|
||||
for k in self.kernels[1:]:
|
||||
assert k.inplace_update_buffers == self.kernels[0].inplace_update_buffers
|
||||
return self.kernels[0].inplace_update_buffers
|
||||
|
||||
def warn_mix_layout(self, kernel_name: str):
|
||||
pass
|
||||
|
||||
|
||||
class MultiKernelCall:
|
||||
"""
|
||||
This class is called at run time to actually run the kernel
|
||||
"""
|
||||
|
||||
def __init__(self, multi_kernel_name, kernels, arg_index):
|
||||
assert len(kernels) >= 1
|
||||
self._kernels = kernels
|
||||
self.multi_kernel_name = multi_kernel_name
|
||||
|
||||
self.disable_cache = os.environ.get(
|
||||
"TORCHINDUCTOR_DISABLE_MULTI_KERNEL_CACHE"
|
||||
) == "1" or is_metric_table_enabled("persistent_red_perf")
|
||||
|
||||
self.picked_kernel = None
|
||||
self.arg_index = arg_index
|
||||
if config.triton.multi_kernel > 1:
|
||||
# manually force a subkernel to ease perf testing
|
||||
picked_by_config = config.triton.multi_kernel - 2
|
||||
assert picked_by_config < len(self._kernels)
|
||||
# pyrefly: ignore [bad-assignment]
|
||||
self.picked_kernel = picked_by_config
|
||||
elif not self.disable_cache:
|
||||
self.load_cache()
|
||||
|
||||
self._recorded = False
|
||||
|
||||
def cache_file_path(self):
|
||||
key = code_hash(
|
||||
",".join(
|
||||
[
|
||||
f"{k.fn.cache_key}{k.size_hints!r}{k.triton_meta!r}"
|
||||
for k in self.kernels
|
||||
]
|
||||
)
|
||||
)
|
||||
_, _, path = get_path(key, "picked_kernel")
|
||||
return pathlib.Path(path)
|
||||
|
||||
def load_cache(self):
|
||||
assert self.picked_kernel is None
|
||||
path = self.cache_file_path()
|
||||
if path.exists():
|
||||
with path.open() as fd:
|
||||
# pyrefly: ignore [bad-assignment]
|
||||
self.picked_kernel = int(fd.read())
|
||||
# pyrefly: ignore [unsupported-operation]
|
||||
assert self.picked_kernel >= 0 and self.picked_kernel < len(
|
||||
self._kernels
|
||||
)
|
||||
log.debug(
|
||||
"Load picked kernel %d from cache file %s", self.picked_kernel, path
|
||||
)
|
||||
|
||||
def store_cache(self):
|
||||
assert self.picked_kernel is not None
|
||||
path = self.cache_file_path()
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
write_atomic(path, str(self.picked_kernel))
|
||||
log.debug("Store picked kernel %d to cache file %s", self.picked_kernel, path)
|
||||
|
||||
@property
|
||||
def kernels(self):
|
||||
"""
|
||||
Read results from future.
|
||||
|
||||
This should be called after parallel compilation is done.
|
||||
In case you call this before compilation is done,
|
||||
it may slow down the parallel compilation.
|
||||
"""
|
||||
for i, kernel in enumerate(self._kernels):
|
||||
if isinstance(kernel, CodeCacheFuture):
|
||||
self._kernels[i] = kernel.result()
|
||||
|
||||
return self._kernels
|
||||
|
||||
def benchmark_sub_kernels(self, *args, **kwargs):
|
||||
"""
|
||||
Benchmark all the sub kernels and return the execution time
|
||||
(in milliseconds) for each of time.
|
||||
|
||||
Unit test may mock this method to force a specific kernel to
|
||||
be picked.
|
||||
"""
|
||||
|
||||
def wrap_fn(kernel, index):
|
||||
def inner():
|
||||
filtered_args = self._get_filtered_args(args, index)
|
||||
args_clone, kwargs_clone = kernel.clone_args(*filtered_args, **kwargs)
|
||||
return kernel.run(*args_clone, **kwargs_clone)
|
||||
|
||||
return inner
|
||||
|
||||
return [
|
||||
benchmarker.benchmark(
|
||||
wrap_fn(kernel, index),
|
||||
# Currently the kernel type must be a CachingAutotuner
|
||||
device=kernel.device_props.type,
|
||||
rep=40,
|
||||
)
|
||||
for index, kernel in enumerate(self.kernels)
|
||||
]
|
||||
|
||||
def _get_filtered_args(self, args, index):
|
||||
"""
|
||||
We pass in all arguments to all kernels into the MultiKernelCall
|
||||
so when invoking a particular kernel we need to filter to only the
|
||||
arguments for that specific kernel.
|
||||
"""
|
||||
|
||||
# This is sometimes invoked at runtime where V.graph is
|
||||
# a NullHandler
|
||||
if hasattr(V.graph, "cpp_wrapper") and V.graph.cpp_wrapper:
|
||||
# for cpp-wrapper, we should not filter the args since
|
||||
# we already have chosen a single kernel and arg set.
|
||||
return args
|
||||
return [item for s in self.arg_index[index] for item in args[s]]
|
||||
|
||||
# record_choice and lookup_choice are helper functions for cpp-wrapper
|
||||
# codegen. The first pass use record_choice to keep the choice and
|
||||
# the second pass do lookup by calling lookup_choice.
|
||||
#
|
||||
# An alternative that reused the multi-kernel cache does not work well
|
||||
# since during codegen of the second pass, it's very hard to know the
|
||||
# path for the cache file. Also reading the cache file need do some IO
|
||||
# which can be slower.
|
||||
@staticmethod
|
||||
def record_choice(multi_kernel_name: str, picked_kernel_name: str):
|
||||
"""
|
||||
Record the multi-kernel choice for cpp-wrapper after autotuning
|
||||
|
||||
We should do nothing if this function is not called during codegen.
|
||||
"""
|
||||
from torch._inductor.graph import GraphLowering
|
||||
|
||||
if not isinstance(V.graph, GraphLowering):
|
||||
return
|
||||
|
||||
if not V.graph.record_multi_kernel_choice:
|
||||
return
|
||||
|
||||
V.graph.multi_kernel_to_choice[multi_kernel_name] = picked_kernel_name
|
||||
|
||||
@staticmethod
|
||||
def lookup_choice(multi_kernel_name: str) -> str:
|
||||
# this should always been done during cpp-wrapper codegen
|
||||
assert (
|
||||
V.graph.record_multi_kernel_choice
|
||||
and multi_kernel_name in V.graph.multi_kernel_to_choice
|
||||
)
|
||||
# there should be no miss
|
||||
return V.graph.multi_kernel_to_choice[multi_kernel_name]
|
||||
|
||||
def run(self, *args, **kwargs):
|
||||
if self.picked_kernel is None:
|
||||
timings = self.benchmark_sub_kernels(*args, **kwargs)
|
||||
self.picked_kernel = timings.index(min(timings))
|
||||
k0 = self.kernels[0]
|
||||
log.debug(
|
||||
"pick %dth sub-kernel in %s. Size hints %s. Reduction hint %s. Timings %s",
|
||||
self.picked_kernel,
|
||||
[k.inductor_meta.get("kernel_name") for k in self.kernels],
|
||||
k0.size_hints,
|
||||
k0.inductor_meta.get("reduction_hint"),
|
||||
timings,
|
||||
)
|
||||
get_metric_table("persistent_red_perf").add_row(
|
||||
functools.partial(self._metrics_table_row, timings)
|
||||
)
|
||||
|
||||
if not self.disable_cache:
|
||||
self.store_cache()
|
||||
|
||||
if not self._recorded:
|
||||
self._recorded = True
|
||||
picked_kernel_name = self.kernels[self.picked_kernel].inductor_meta.get(
|
||||
"kernel_name"
|
||||
)
|
||||
assert picked_kernel_name is not None
|
||||
self.record_choice(self.multi_kernel_name, picked_kernel_name)
|
||||
|
||||
run = self.kernels[self.picked_kernel].run # type: ignore[method-assign]
|
||||
filtered_args = self._get_filtered_args(args, self.picked_kernel)
|
||||
run(*filtered_args, **kwargs)
|
||||
|
||||
def _metrics_table_row(self, timings):
|
||||
def get_kernel_path(k):
|
||||
return k.fn.fn.__code__.co_filename
|
||||
|
||||
k0 = self.kernels[0]
|
||||
row = {
|
||||
"size_hints": k0.size_hints,
|
||||
"reduction_hint": k0.inductor_meta.get("reduction_hint"),
|
||||
}
|
||||
max_kernels = 4
|
||||
assert len(timings) <= max_kernels
|
||||
for i in range(max_kernels):
|
||||
if i < len(self.kernels):
|
||||
row[f"kernel{i}_path"] = get_kernel_path(self.kernels[i])
|
||||
row[f"kernel{i}_latency"] = timings[i]
|
||||
else:
|
||||
row[f"kernel{i}_path"] = ""
|
||||
row[f"kernel{i}_latency"] = ""
|
||||
return row
|
||||
|
||||
|
||||
class SizeHintMultiKernel(MultiKernel):
|
||||
"""
|
||||
Version of multi-kernel that generates kernels based on specified size hints.
|
||||
Currently only performs 1-d search over hints; doesn't perform combinatorial n-d search
|
||||
if n > 1 dynamic dimensions are specified.
|
||||
|
||||
e.g. matmul([s0, s1], [s1, s2]) with size-hints [64, 256] only generates 2 kernels,
|
||||
based on tuning shapes ([64, 64], [64, 64]) and ([256, 256], [256, 256])
|
||||
"""
|
||||
|
||||
def __init__(self, kernels):
|
||||
assert isinstance(kernels, dict) and len(kernels) >= 1
|
||||
|
||||
self.kernels, self.kernel_shape_keys = [], []
|
||||
for shape_key, kernel in kernels.items():
|
||||
self.kernels.append(kernel)
|
||||
self.kernel_shape_keys.append(shape_key)
|
||||
self.kernel_name = V.graph.wrapper_code.multi_kernel_state.define_kernel(
|
||||
self.kernels, self.kernel_shape_keys
|
||||
)
|
||||
|
||||
# need this since some code in inductor check if the kernel object has an args
|
||||
# attribute to decide if it's a non-null kernel.
|
||||
self.args = object()
|
||||
|
||||
|
||||
class SizeHintMultiKernelCall(MultiKernelCall):
|
||||
"""
|
||||
Runtime class for size-hint multi-kernels.
|
||||
Instead of having a plain list of kernels to benchmark over, keys them by input & output shapes,
|
||||
and optionally perform shape-based selection. The pre-generated kernel is chosen based on the shape keys,
|
||||
with the heuristic being log2 l1 distance between the pre-generated / runtime input & output shapes.
|
||||
"""
|
||||
|
||||
def __init__(self, multi_kernel_name, kernels, arg_index):
|
||||
super().__init__(multi_kernel_name, list(kernels.values()), arg_index)
|
||||
self._kernel_hints = list(kernels.keys())
|
||||
|
||||
# Caches results for unique shapes.
|
||||
self._shape_cache = {}
|
||||
|
||||
def _get_shape_cache_key(self, *args, **kwargs):
|
||||
"""
|
||||
Generate a cache key based on tensor shapes for shape-specialized dispatch.
|
||||
"""
|
||||
shapes = []
|
||||
for arg in args:
|
||||
if hasattr(arg, "shape"):
|
||||
shapes.append(tuple(arg.shape))
|
||||
return tuple(shapes)
|
||||
|
||||
def _get_cached_shape_choice(self, cache_key):
|
||||
"""
|
||||
Get cached kernel choice for a specific shape.
|
||||
"""
|
||||
return self._shape_cache.get(cache_key)
|
||||
|
||||
def _cache_shape_choice(self, cache_key, kernel_idx):
|
||||
"""
|
||||
Cache kernel choice for a specific shape.
|
||||
"""
|
||||
self._shape_cache[cache_key] = kernel_idx
|
||||
|
||||
def _dist_heuristic(self, k1, k2):
|
||||
"""
|
||||
log2 L1 distance heuristic for kernel selection.
|
||||
"""
|
||||
|
||||
def dist(x, y):
|
||||
lx = math.log2(x) if x > 0 else -1
|
||||
ly = math.log2(y) if y > 0 else -1
|
||||
return abs(lx - ly)
|
||||
|
||||
out = 0
|
||||
for s1, s2 in zip(k1, k2):
|
||||
out += sum(dist(x, y) for x, y in zip(s1, s2))
|
||||
return out
|
||||
|
||||
def run(self, *args, **kwargs):
|
||||
cache_key = self._get_shape_cache_key(*args, **kwargs)
|
||||
cached_choice = self._get_cached_shape_choice(cache_key)
|
||||
if cached_choice is not None:
|
||||
self.picked_kernel = cached_choice
|
||||
log.debug(
|
||||
"using cached shape-specialized choice %dth sub-kernel in %s. Cache key: %s",
|
||||
self.picked_kernel,
|
||||
[k.inductor_meta.get("kernel_name") for k in self.kernels],
|
||||
cache_key,
|
||||
)
|
||||
else:
|
||||
self._select_kernel_by_shape(*args, **kwargs)
|
||||
|
||||
if not self._recorded:
|
||||
self._recorded = True
|
||||
picked_kernel_name = self.kernels[self.picked_kernel].inductor_meta.get(
|
||||
"kernel_name"
|
||||
)
|
||||
assert picked_kernel_name is not None
|
||||
self.record_choice(self.multi_kernel_name, picked_kernel_name)
|
||||
|
||||
run = self.kernels[self.picked_kernel].run # type: ignore[method-assign]
|
||||
filtered_args = self._get_filtered_args(args, self.picked_kernel)
|
||||
run(*filtered_args, **kwargs)
|
||||
|
||||
def _select_kernel_by_shape(self, *args, **kwargs):
|
||||
"""
|
||||
Benchmark kernels for a particular shape and return the
|
||||
best kernel for this shape.
|
||||
"""
|
||||
shape_key = self._get_shape_cache_key(*args, **kwargs)
|
||||
dists = [
|
||||
self._dist_heuristic(shape_key, key) if key is not None else 2**62
|
||||
for key in self._kernel_hints
|
||||
]
|
||||
# pyrefly: ignore [bad-assignment]
|
||||
self.picked_kernel = dists.index(min(dists))
|
||||
self._cache_shape_choice(shape_key, self.picked_kernel)
|
||||
+17
@@ -0,0 +1,17 @@
|
||||
# mypy: allow-untyped-defs
|
||||
from .nv_universal_gemm import (
|
||||
add_nv_universal_gemm_choices,
|
||||
add_nv_universal_grouped_gemm_choices,
|
||||
add_nv_universal_scaled_gemm_choices,
|
||||
GemmVariant,
|
||||
NVUniversalGemmCaller,
|
||||
)
|
||||
|
||||
|
||||
__all__ = [
|
||||
"GemmVariant",
|
||||
"NVUniversalGemmCaller",
|
||||
"add_nv_universal_gemm_choices",
|
||||
"add_nv_universal_grouped_gemm_choices",
|
||||
"add_nv_universal_scaled_gemm_choices",
|
||||
]
|
||||
+96
@@ -0,0 +1,96 @@
|
||||
# mypy: allow-untyped-defs
|
||||
"""
|
||||
Global kernel cache for NVIDIA Universal GEMM.
|
||||
|
||||
This module provides a lazy-initialized cache for cutlass_api kernels,
|
||||
avoiding expensive manifest scans on every kernel lookup.
|
||||
|
||||
The first call to get_kernel_by_name() loads all kernels from cutlass_api
|
||||
(~10 seconds) and builds a name->kernel dict. Subsequent calls use the
|
||||
dict for O(1) lookup (~0.1 μs).
|
||||
"""
|
||||
|
||||
import logging
|
||||
from collections.abc import Callable
|
||||
from typing import Any
|
||||
|
||||
|
||||
log = logging.getLogger(__name__)
|
||||
|
||||
# Global cache: kernel_name -> kernel object
|
||||
_kernel_by_name_cache: dict[str, Any] | None = None
|
||||
|
||||
|
||||
def _build_kernel_cache() -> dict[str, Any]:
|
||||
"""Build the kernel name -> kernel object cache."""
|
||||
import cutlass_api
|
||||
|
||||
log.debug("Building NVGEMM kernel cache (this may take a few seconds)...")
|
||||
|
||||
try:
|
||||
from torch._inductor.kernel.vendored_templates.cutedsl import ( # noqa: F401
|
||||
wrappers,
|
||||
)
|
||||
except ImportError:
|
||||
log.debug("Vendored kernel wrappers not available")
|
||||
|
||||
all_kernels = cutlass_api.get_kernels()
|
||||
cache = {k.metadata.kernel_name: k for k in all_kernels}
|
||||
log.debug("NVGEMM kernel cache built: %d kernels", len(cache))
|
||||
return cache
|
||||
|
||||
|
||||
def get_compatible_kernels(
|
||||
args: Any,
|
||||
cc: int,
|
||||
metadata_filter: Callable[[Any], bool] | None = None,
|
||||
) -> list[Any]:
|
||||
"""Get kernels compatible with the given arguments from the cache."""
|
||||
global _kernel_by_name_cache
|
||||
|
||||
if _kernel_by_name_cache is None:
|
||||
_kernel_by_name_cache = _build_kernel_cache()
|
||||
|
||||
compatible = []
|
||||
for kernel in _kernel_by_name_cache.values():
|
||||
if kernel.metadata.min_cc > cc:
|
||||
continue
|
||||
|
||||
if metadata_filter is not None and not metadata_filter(kernel.metadata):
|
||||
continue
|
||||
|
||||
status = kernel.supports(args)
|
||||
if status.error is not None:
|
||||
continue
|
||||
compatible.append(kernel)
|
||||
|
||||
log.debug(
|
||||
"Found %d compatible kernels from cache of %d total",
|
||||
len(compatible),
|
||||
len(_kernel_by_name_cache),
|
||||
)
|
||||
return compatible
|
||||
|
||||
|
||||
def get_kernel_by_name(kernel_name: str) -> Any:
|
||||
"""Get a cutlass_api kernel by name using the global cache."""
|
||||
global _kernel_by_name_cache
|
||||
|
||||
if _kernel_by_name_cache is None:
|
||||
_kernel_by_name_cache = _build_kernel_cache()
|
||||
|
||||
return _kernel_by_name_cache.get(kernel_name)
|
||||
|
||||
|
||||
def ensure_cache_initialized() -> None:
|
||||
"""Ensure the kernel cache is initialized."""
|
||||
global _kernel_by_name_cache
|
||||
|
||||
if _kernel_by_name_cache is None:
|
||||
_kernel_by_name_cache = _build_kernel_cache()
|
||||
|
||||
|
||||
def clear_cache() -> None:
|
||||
"""Clear the kernel cache."""
|
||||
global _kernel_by_name_cache
|
||||
_kernel_by_name_cache = None
|
||||
+555
@@ -0,0 +1,555 @@
|
||||
# mypy: allow-untyped-defs
|
||||
"""
|
||||
NVIDIA Universal GEMM (NVGEMM) backend for PyTorch Inductor.
|
||||
|
||||
This module provides integration with the cutlass_api library to enable
|
||||
high-performance GEMM kernels for NVIDIA GPUs.
|
||||
"""
|
||||
|
||||
import itertools
|
||||
from enum import auto, Enum
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
from torch._inductor import config
|
||||
from torch._inductor.autotune_process import (
|
||||
BenchmarkRequest,
|
||||
GPUDeviceBenchmarkMixin,
|
||||
TensorMeta,
|
||||
)
|
||||
from torch._inductor.codegen.cuda.cuda_env import get_cuda_arch
|
||||
from torch._inductor.codegen.nv_universal_gemm.nv_universal_gemm_utils import (
|
||||
to_cutlass_scale_mode,
|
||||
)
|
||||
from torch._inductor.ir import Buffer, ChoiceCaller, Layout, TensorBox
|
||||
from torch._inductor.kernel_inputs import MMKernelInputs
|
||||
from torch._inductor.template_heuristics.nv_universal_gemm import get_nvgemm_heuristics
|
||||
from torch._inductor.utils import ensure_nv_universal_gemm_available
|
||||
from torch._logging import getArtifactLogger
|
||||
|
||||
|
||||
log = getArtifactLogger(__name__, "output_code")
|
||||
|
||||
|
||||
class GemmVariant(Enum):
|
||||
"""
|
||||
Enum for different GEMM operation types supported by NVIDIA Universal GEMM.
|
||||
"""
|
||||
|
||||
GEMM = auto()
|
||||
|
||||
GROUPED_GEMM = auto()
|
||||
|
||||
SCALED_GEMM = auto()
|
||||
|
||||
@property
|
||||
def op_name(self) -> str:
|
||||
"""Return the operation name for logging and naming."""
|
||||
if self == GemmVariant.GROUPED_GEMM:
|
||||
return "nv_universal_grouped_gemm"
|
||||
if self == GemmVariant.SCALED_GEMM:
|
||||
return "nv_universal_scaled_gemm"
|
||||
return "nv_universal_gemm"
|
||||
|
||||
@property
|
||||
def arguments_class_name(self) -> str:
|
||||
"""Return the cutlass_api arguments class name."""
|
||||
if self == GemmVariant.GROUPED_GEMM:
|
||||
return "GroupedGemmArguments"
|
||||
return "GemmArguments"
|
||||
|
||||
|
||||
class NVUniversalGemmBenchmarkRequest(GPUDeviceBenchmarkMixin, BenchmarkRequest):
|
||||
"""Benchmark request for NVIDIA Universal GEMM kernels."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
kernel_name: str,
|
||||
input_tensor_meta: TensorMeta | list[TensorMeta],
|
||||
output_tensor_meta: TensorMeta | list[TensorMeta],
|
||||
kernel, # cutlass_api.Kernel object
|
||||
accumulator_type: torch.dtype,
|
||||
variant: GemmVariant,
|
||||
workspace_size: int = 0,
|
||||
scale_type_a: Any | None = None,
|
||||
scale_type_b: Any | None = None,
|
||||
swizzle_type_a: Any | None = None,
|
||||
swizzle_type_b: Any | None = None,
|
||||
) -> None:
|
||||
super().__init__(kernel_name, input_tensor_meta, output_tensor_meta, ())
|
||||
self.kernel = kernel
|
||||
self.accumulator_type = accumulator_type
|
||||
self._compiled_artifact = None
|
||||
self._workspace: torch.Tensor | None = None
|
||||
self.workspace_size = workspace_size
|
||||
self.variant = variant
|
||||
self.scale_type_a = scale_type_a
|
||||
self.scale_type_b = scale_type_b
|
||||
self.swizzle_type_a = swizzle_type_a
|
||||
self.swizzle_type_b = swizzle_type_b
|
||||
|
||||
def benchmark(
|
||||
self,
|
||||
*input_tensors: torch.Tensor,
|
||||
out: torch.Tensor | None = None,
|
||||
) -> float:
|
||||
"""Benchmark the NVIDIA Universal GEMM kernel.
|
||||
|
||||
Override the base class to always create tensors from input_tensor_meta.
|
||||
This is necessary because input_nodes may be ReinterpretViews that share
|
||||
the same underlying buffer name. The autotuning framework deduplicates
|
||||
inputs by name (in AlgorithmSelectorCache.get_inputs()), resulting in
|
||||
fewer tensors than expected. By always creating from input_tensor_meta,
|
||||
we ensure each input gets its own tensor with the correct size/stride/offset
|
||||
from the view's layout.
|
||||
"""
|
||||
# Always create tensors from input_tensor_meta, ignoring passed-in tensors
|
||||
input_tensors = tuple(x.to_tensor() for x in self.input_tensor_meta)
|
||||
if out is None:
|
||||
out = self.output_tensor_meta.to_tensor()
|
||||
|
||||
fn = self.make_run_fn(*input_tensors, out=out)
|
||||
return self.do_bench(fn, *input_tensors, out=out)
|
||||
|
||||
def make_run_fn(self, *input_tensors: torch.Tensor, out: torch.Tensor):
|
||||
"""Create a function to run the NVIDIA Universal GEMM kernel."""
|
||||
import cutlass_api
|
||||
|
||||
args = self._create_gemm_arguments(cutlass_api, input_tensors, out)
|
||||
|
||||
if self._compiled_artifact is None:
|
||||
self._compiled_artifact = self.kernel.compile(args)
|
||||
artifact = self._compiled_artifact
|
||||
kernel = self.kernel
|
||||
|
||||
# Allocate workspace if needed
|
||||
if self.workspace_size > 0:
|
||||
self._workspace = torch.empty(
|
||||
self.workspace_size, device=out.device, dtype=torch.int8
|
||||
)
|
||||
else:
|
||||
self._workspace = None
|
||||
|
||||
workspace = self._workspace
|
||||
|
||||
def run_kernel():
|
||||
stream = torch.cuda.current_stream()
|
||||
kernel.run(
|
||||
args,
|
||||
artifact,
|
||||
stream=stream,
|
||||
workspace=workspace,
|
||||
assume_supported_args=True,
|
||||
)
|
||||
|
||||
return run_kernel
|
||||
|
||||
def _create_gemm_arguments(self, cutlass_api, input_tensors, out):
|
||||
"""Create the appropriate GemmArguments based on variant."""
|
||||
if self.variant == GemmVariant.GROUPED_GEMM:
|
||||
a, b, offsets = input_tensors
|
||||
return cutlass_api.arguments.GroupedGemmArguments(
|
||||
a,
|
||||
b,
|
||||
out,
|
||||
accumulator_type=self.accumulator_type,
|
||||
offsets=offsets,
|
||||
)
|
||||
elif self.variant == GemmVariant.SCALED_GEMM:
|
||||
from cutlass_api.arguments import ScaledTensor
|
||||
|
||||
scale_mode_a, swizzle_mode_a = to_cutlass_scale_mode(
|
||||
self.scale_type_a, self.swizzle_type_a
|
||||
)
|
||||
scale_mode_b, swizzle_mode_b = to_cutlass_scale_mode(
|
||||
self.scale_type_b, self.swizzle_type_b
|
||||
)
|
||||
|
||||
a, b, scale_a, scale_b = input_tensors
|
||||
scaled_a = ScaledTensor(a, scale_a, scale_mode_a, swizzle_mode_a)
|
||||
scaled_b = ScaledTensor(b, scale_b, scale_mode_b, swizzle_mode_b)
|
||||
return cutlass_api.arguments.GemmArguments(
|
||||
scaled_a,
|
||||
scaled_b,
|
||||
out,
|
||||
accumulator_type=self.accumulator_type,
|
||||
)
|
||||
else:
|
||||
a, b = input_tensors
|
||||
return cutlass_api.arguments.GemmArguments(
|
||||
a,
|
||||
b,
|
||||
out,
|
||||
accumulator_type=self.accumulator_type,
|
||||
)
|
||||
|
||||
def cleanup_run_fn(self) -> None:
|
||||
self._workspace = None
|
||||
|
||||
|
||||
class NVUniversalGemmCaller(ChoiceCaller):
|
||||
"""
|
||||
ChoiceCaller for NVIDIA Universal GEMM kernels.
|
||||
|
||||
Wraps a cutlass_api kernel and integrates with Inductor's autotuning.
|
||||
"""
|
||||
|
||||
index_counter = itertools.count()
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
name: str,
|
||||
input_nodes: list[Buffer],
|
||||
layout: Layout,
|
||||
kernel, # cutlass_api.Kernel object
|
||||
accumulator_type: torch.dtype,
|
||||
variant: GemmVariant,
|
||||
workspace_size: int = 0,
|
||||
scale_type_a: Any | None = None,
|
||||
scale_type_b: Any | None = None,
|
||||
swizzle_type_a: Any | None = None,
|
||||
swizzle_type_b: Any | None = None,
|
||||
) -> None:
|
||||
super().__init__(
|
||||
name=name,
|
||||
input_nodes=input_nodes,
|
||||
layout=layout,
|
||||
description=f"{variant.op_name} {kernel.metadata.kernel_name}",
|
||||
)
|
||||
self.kernel = kernel
|
||||
self.accumulator_type = accumulator_type
|
||||
self.workspace_size = workspace_size
|
||||
self.variant = variant
|
||||
self.scale_type_a = scale_type_a
|
||||
self.scale_type_b = scale_type_b
|
||||
self.swizzle_type_a = swizzle_type_a
|
||||
self.swizzle_type_b = swizzle_type_b
|
||||
|
||||
output_buffer = Buffer(name=f"{variant.op_name}_out", layout=layout)
|
||||
|
||||
self.bmreq = NVUniversalGemmBenchmarkRequest(
|
||||
kernel_name=name,
|
||||
input_tensor_meta=TensorMeta.from_irnodes(input_nodes),
|
||||
output_tensor_meta=TensorMeta.from_irnodes(output_buffer),
|
||||
kernel=kernel,
|
||||
accumulator_type=accumulator_type,
|
||||
workspace_size=workspace_size,
|
||||
variant=variant,
|
||||
scale_type_a=scale_type_a,
|
||||
scale_type_b=scale_type_b,
|
||||
swizzle_type_a=swizzle_type_a,
|
||||
swizzle_type_b=swizzle_type_b,
|
||||
)
|
||||
|
||||
def __str__(self) -> str:
|
||||
return f"NVUniversalGemmCaller({self.kernel.metadata.kernel_name})"
|
||||
|
||||
def benchmark(self, *args, out) -> float:
|
||||
return self.bmreq.benchmark(*args, out=out)
|
||||
|
||||
def output_node(self) -> TensorBox:
|
||||
from torch._inductor.ir import NVUniversalGemmBuffer
|
||||
|
||||
buffer = NVUniversalGemmBuffer(
|
||||
layout=self.layout,
|
||||
inputs=self.input_nodes,
|
||||
kernel=self.kernel,
|
||||
accumulator_type=self.accumulator_type,
|
||||
workspace_size=self.workspace_size,
|
||||
variant=self.variant,
|
||||
scale_type_a=self.scale_type_a,
|
||||
scale_type_b=self.scale_type_b,
|
||||
swizzle_type_a=self.swizzle_type_a,
|
||||
swizzle_type_b=self.swizzle_type_b,
|
||||
)
|
||||
# Pass KTC annotation to the buffer for encoding
|
||||
if "ktc" in self.annotations:
|
||||
buffer.annotations["ktc"] = self.annotations["ktc"]
|
||||
return TensorBox.create(buffer)
|
||||
|
||||
def call_name(self) -> str:
|
||||
return self.name
|
||||
|
||||
def to_callable(self):
|
||||
return self.bmreq.make_run_fn
|
||||
|
||||
def hash_key(self) -> str:
|
||||
return f"{self.variant.op_name}_{self.kernel.metadata.kernel_name}"
|
||||
|
||||
def info_dict(self) -> dict[str, Any]:
|
||||
return {
|
||||
"name": self.name,
|
||||
"backend": self.variant.op_name,
|
||||
"kernel_name": self.kernel.metadata.kernel_name,
|
||||
}
|
||||
|
||||
|
||||
def _create_dummy_tensor_from_layout(layout: Layout) -> torch.Tensor | None:
|
||||
"""
|
||||
Create a FakeTensor from a Layout for kernel filtering.
|
||||
|
||||
Uses Layout.get_example() which creates FakeTensors within V.fake_mode,
|
||||
avoiding real CUDA memory allocation. cutlass_api only needs shape/stride/dtype
|
||||
metadata for its supports() checks.
|
||||
"""
|
||||
try:
|
||||
return layout.get_example()
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
|
||||
def _exclude_efc_kernels(metadata) -> bool:
|
||||
"""
|
||||
Filter out EFC kernels.
|
||||
|
||||
EFC kernels support custom epilogue operations but have additional overhead.
|
||||
Since NVGEMM doesn't support epilogue fusion yet (see nv_universal_gemm_scheduling.py),
|
||||
we use non-EFC kernels which are equivalent for identity epilogue (out = acc).
|
||||
TODO(nikhilap): Remove this filter once NVGEMM supports epilogue fusion.
|
||||
"""
|
||||
return "EFC" not in metadata.kernel_class.__name__
|
||||
|
||||
|
||||
def _add_nv_gemm_choices_impl(
|
||||
choices: list[ChoiceCaller],
|
||||
layout: Layout,
|
||||
input_nodes: list[Buffer],
|
||||
variant: GemmVariant,
|
||||
accumulator_type: torch.dtype,
|
||||
mm_inputs: MMKernelInputs | None = None,
|
||||
scale_type_a: Any | None = None,
|
||||
scale_type_b: Any | None = None,
|
||||
swizzle_type_a: Any | None = None,
|
||||
swizzle_type_b: Any | None = None,
|
||||
) -> None:
|
||||
"""
|
||||
Unified implementation for adding NVIDIA Universal GEMM choices.
|
||||
|
||||
Args:
|
||||
choices: List to append ChoiceCaller objects to
|
||||
layout: Output layout
|
||||
input_nodes: Input tensor nodes
|
||||
variant: The GEMM variant (determines behavior)
|
||||
accumulator_type: Accumulator dtype
|
||||
mm_inputs: Optional MMKernelInputs for heuristics
|
||||
scale_type_a: ScalingType for A (required for SCALED_GEMM)
|
||||
scale_type_b: ScalingType for B (required for SCALED_GEMM)
|
||||
swizzle_type_a: SwizzleType for A (required for SCALED_GEMM)
|
||||
swizzle_type_b: SwizzleType for B (required for SCALED_GEMM)
|
||||
"""
|
||||
import cutlass_api
|
||||
|
||||
from torch._inductor.codegen.nv_universal_gemm.kernel_cache import (
|
||||
get_compatible_kernels,
|
||||
)
|
||||
|
||||
# Create dummy tensors for cutlass_api's supports() checks
|
||||
dummy_tensors = [
|
||||
_create_dummy_tensor_from_layout(node.get_layout()) for node in input_nodes
|
||||
]
|
||||
out_tensor = _create_dummy_tensor_from_layout(layout)
|
||||
|
||||
if any(t is None for t in dummy_tensors) or out_tensor is None:
|
||||
log.debug("Failed to create dummy tensors for %s", variant.op_name)
|
||||
return
|
||||
|
||||
if variant == GemmVariant.GROUPED_GEMM:
|
||||
a_tensor, b_tensor, offs_tensor = dummy_tensors
|
||||
assert b_tensor is not None
|
||||
args = cutlass_api.arguments.GroupedGemmArguments(
|
||||
a_tensor,
|
||||
b_tensor,
|
||||
out_tensor,
|
||||
accumulator_type=accumulator_type,
|
||||
offsets=offs_tensor,
|
||||
)
|
||||
elif variant == GemmVariant.SCALED_GEMM:
|
||||
from cutlass_api.arguments import ScaledTensor
|
||||
|
||||
scale_mode_a, swizzle_mode_a = to_cutlass_scale_mode(
|
||||
scale_type_a, swizzle_type_a
|
||||
)
|
||||
scale_mode_b, swizzle_mode_b = to_cutlass_scale_mode(
|
||||
scale_type_b, swizzle_type_b
|
||||
)
|
||||
if scale_mode_a is None or scale_mode_b is None:
|
||||
return
|
||||
|
||||
a_tensor, b_tensor, scale_a_tensor, scale_b_tensor = dummy_tensors
|
||||
scaled_a = ScaledTensor(
|
||||
a_tensor,
|
||||
scale_a_tensor,
|
||||
scale_mode_a,
|
||||
swizzle_mode_a,
|
||||
)
|
||||
scaled_b = ScaledTensor(
|
||||
b_tensor,
|
||||
scale_b_tensor,
|
||||
scale_mode_b,
|
||||
swizzle_mode_b,
|
||||
)
|
||||
args = cutlass_api.arguments.GemmArguments(
|
||||
scaled_a,
|
||||
scaled_b,
|
||||
out_tensor,
|
||||
accumulator_type=accumulator_type,
|
||||
)
|
||||
else:
|
||||
a_tensor, b_tensor = dummy_tensors
|
||||
args = cutlass_api.arguments.GemmArguments(
|
||||
a_tensor,
|
||||
b_tensor,
|
||||
out_tensor,
|
||||
accumulator_type=accumulator_type,
|
||||
)
|
||||
|
||||
cc = get_cuda_arch()
|
||||
if cc is None:
|
||||
log.debug("Failed to get CUDA arch")
|
||||
return
|
||||
cc_int = int(cc)
|
||||
|
||||
kernels = get_compatible_kernels(args, cc_int, metadata_filter=_exclude_efc_kernels)
|
||||
if not kernels:
|
||||
log.debug("No compatible %s kernels found", variant.op_name)
|
||||
return
|
||||
|
||||
max_configs = config.nvgemm_max_profiling_configs or len(kernels)
|
||||
if variant in (GemmVariant.GEMM, GemmVariant.SCALED_GEMM) and mm_inputs is not None:
|
||||
heuristics = get_nvgemm_heuristics()
|
||||
kernels = heuristics.filter_kernels(
|
||||
kernels, mm_inputs, max_configs, accumulator_type
|
||||
)
|
||||
else:
|
||||
# TODO(nikhilap): Enable heuristics for grouped GEMM
|
||||
# when nvMatmulHeuristics adds support
|
||||
kernels = kernels[:max_configs]
|
||||
|
||||
# Add callers for each kernel
|
||||
num_added = 0
|
||||
for kernel in kernels:
|
||||
name = f"{variant.op_name}_{next(NVUniversalGemmCaller.index_counter)}"
|
||||
workspace_size = kernel.get_workspace_size(args)
|
||||
try:
|
||||
caller = NVUniversalGemmCaller(
|
||||
name=name,
|
||||
input_nodes=input_nodes,
|
||||
layout=layout,
|
||||
kernel=kernel,
|
||||
accumulator_type=accumulator_type,
|
||||
workspace_size=workspace_size,
|
||||
variant=variant,
|
||||
scale_type_a=scale_type_a,
|
||||
scale_type_b=scale_type_b,
|
||||
swizzle_type_a=swizzle_type_a,
|
||||
swizzle_type_b=swizzle_type_b,
|
||||
)
|
||||
choices.append(caller)
|
||||
num_added += 1
|
||||
except Exception:
|
||||
log.debug("Failed to create %s choice", variant.op_name, exc_info=True)
|
||||
|
||||
log.debug("Added %d %s choices", num_added, variant.op_name)
|
||||
|
||||
|
||||
def add_nv_universal_gemm_choices(
|
||||
choices: list[ChoiceCaller],
|
||||
layout: Layout,
|
||||
inputs: MMKernelInputs,
|
||||
accumulator_type: torch.dtype | None = None,
|
||||
) -> None:
|
||||
"""
|
||||
Add NVIDIA Universal GEMM kernels to the autotune choices.
|
||||
|
||||
Thin wrapper around _add_nv_gemm_choices_impl for regular GEMM.
|
||||
"""
|
||||
if not ensure_nv_universal_gemm_available():
|
||||
log.debug("cutlass_api not available, skipping NVIDIA Universal GEMM choices")
|
||||
return
|
||||
|
||||
_add_nv_gemm_choices_impl(
|
||||
choices=choices,
|
||||
layout=layout,
|
||||
input_nodes=inputs.nodes(),
|
||||
variant=GemmVariant.GEMM,
|
||||
accumulator_type=accumulator_type or torch.float32,
|
||||
mm_inputs=inputs,
|
||||
)
|
||||
|
||||
|
||||
def add_nv_universal_grouped_gemm_choices(
|
||||
choices: list[ChoiceCaller],
|
||||
layout: Layout,
|
||||
input_nodes: list[Buffer],
|
||||
accumulator_type: torch.dtype | None = None,
|
||||
) -> None:
|
||||
"""
|
||||
Add NVIDIA Universal Grouped GEMM kernels to the autotune choices.
|
||||
|
||||
Thin wrapper around _add_nv_gemm_choices_impl for grouped GEMM.
|
||||
|
||||
For grouped GEMM (contiguous offset variant):
|
||||
- A is (TotalM, K) with problems stacked along M
|
||||
- B is (G, K, N) where B[i] is the weight for problem i
|
||||
- offsets is (G,) marking where each problem ends in A
|
||||
- Output is (TotalM, N)
|
||||
"""
|
||||
if not ensure_nv_universal_gemm_available():
|
||||
log.debug(
|
||||
"cutlass_api not available, skipping NVIDIA Universal Grouped GEMM choices"
|
||||
)
|
||||
return
|
||||
|
||||
_add_nv_gemm_choices_impl(
|
||||
choices=choices,
|
||||
layout=layout,
|
||||
input_nodes=input_nodes,
|
||||
variant=GemmVariant.GROUPED_GEMM,
|
||||
accumulator_type=accumulator_type or torch.float32,
|
||||
)
|
||||
|
||||
|
||||
def add_nv_universal_scaled_gemm_choices(
|
||||
choices: list[ChoiceCaller],
|
||||
layout: Layout,
|
||||
input_nodes: list[Buffer],
|
||||
accumulator_type: torch.dtype | None = None,
|
||||
kernel_inputs: MMKernelInputs | None = None,
|
||||
) -> None:
|
||||
"""
|
||||
Add NVIDIA Universal Scaled GEMM (FP8) kernels to the autotune choices.
|
||||
|
||||
The scaling type is inferred from the input shapes/dtypes.
|
||||
If the scaling mode is unsupported by NVGEMM, this function returns without
|
||||
adding any choices.
|
||||
"""
|
||||
if not ensure_nv_universal_gemm_available():
|
||||
return
|
||||
|
||||
from torch._inductor.utils import infer_scale_swizzle_ir
|
||||
|
||||
if len(input_nodes) < 4:
|
||||
return
|
||||
|
||||
mat_a, mat_b, scale_a, scale_b = input_nodes[:4]
|
||||
|
||||
scale_type_a, swizzle_type_a = infer_scale_swizzle_ir(mat_a, scale_a)
|
||||
scale_type_b, swizzle_type_b = infer_scale_swizzle_ir(
|
||||
mat_b, scale_b, transpose=True
|
||||
)
|
||||
|
||||
if scale_type_a is None or scale_type_b is None:
|
||||
return
|
||||
|
||||
_add_nv_gemm_choices_impl(
|
||||
choices=choices,
|
||||
layout=layout,
|
||||
input_nodes=input_nodes,
|
||||
variant=GemmVariant.SCALED_GEMM,
|
||||
accumulator_type=accumulator_type or torch.float32,
|
||||
mm_inputs=kernel_inputs,
|
||||
scale_type_a=scale_type_a,
|
||||
scale_type_b=scale_type_b,
|
||||
swizzle_type_a=swizzle_type_a,
|
||||
swizzle_type_b=swizzle_type_b,
|
||||
)
|
||||
+293
@@ -0,0 +1,293 @@
|
||||
# mypy: allow-untyped-defs
|
||||
"""
|
||||
NVIDIA Universal GEMM kernel code generation.
|
||||
|
||||
This module generates Python code that calls cutlass_api to execute GEMM operations.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from typing import Any, TYPE_CHECKING
|
||||
|
||||
from torch._inductor.codegen.common import (
|
||||
IndentedBuffer,
|
||||
Kernel,
|
||||
WorkspaceArg,
|
||||
WorkspaceZeroMode,
|
||||
)
|
||||
from torch._inductor.codegen.cutedsl.cutedsl_op_overrides import CuteDSLOpOverrides
|
||||
from torch._inductor.codegen.nv_universal_gemm.nv_universal_gemm_utils import (
|
||||
to_cutlass_scale_mode,
|
||||
)
|
||||
from torch._inductor.ir import (
|
||||
BaseView,
|
||||
Buffer,
|
||||
ExternKernel,
|
||||
MutableBox,
|
||||
ReinterpretView,
|
||||
)
|
||||
from torch._inductor.virtualized import V
|
||||
from torch.utils._ordered_set import OrderedSet
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from torch._inductor.codegen.nv_universal_gemm.nv_universal_gemm import GemmVariant
|
||||
|
||||
|
||||
log = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class NVUniversalGemmKernelWrapper:
|
||||
"""Wrapper to provide .run() interface for NVIDIA Universal GEMM kernels."""
|
||||
|
||||
def __init__(self, kernel_fn, kernel_path: str | None = None):
|
||||
self.kernel_fn = kernel_fn
|
||||
self.kernel_path = kernel_path
|
||||
|
||||
def run(self, *args, stream=None, **kwargs):
|
||||
"""Execute the NVIDIA Universal GEMM kernel."""
|
||||
return self.kernel_fn(*args, stream=stream, **kwargs)
|
||||
|
||||
|
||||
class NVUniversalGemmKernel(Kernel):
|
||||
"""
|
||||
Kernel implementation for NVIDIA Universal GEMM.
|
||||
|
||||
Generates Python code that calls cutlass_api to execute GEMM operations.
|
||||
Unlike CuteDSL which uses Jinja templates, this generates simpler direct
|
||||
Python code.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
kernel_name: str,
|
||||
input_nodes: list[Buffer],
|
||||
output_node: Buffer,
|
||||
kernel_metadata: dict[str, Any],
|
||||
accumulator_type: Any,
|
||||
variant: GemmVariant,
|
||||
workspace_size: int = 0,
|
||||
scale_type_a: Any | None = None,
|
||||
scale_type_b: Any | None = None,
|
||||
swizzle_type_a: Any | None = None,
|
||||
swizzle_type_b: Any | None = None,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self.kernel_name = kernel_name
|
||||
self.input_nodes = input_nodes
|
||||
self.output_node = output_node
|
||||
self.kernel_metadata = kernel_metadata
|
||||
self.accumulator_type = accumulator_type
|
||||
self.workspace_size = workspace_size
|
||||
self.variant = variant
|
||||
self.scale_type_a = scale_type_a
|
||||
self.scale_type_b = scale_type_b
|
||||
self.swizzle_type_a = swizzle_type_a
|
||||
self.swizzle_type_b = swizzle_type_b
|
||||
|
||||
self._template_input_args: list[tuple[str, Buffer]] = []
|
||||
self._seen_input_args: OrderedSet[str] = OrderedSet()
|
||||
|
||||
for i, input_node in enumerate(input_nodes):
|
||||
param_name = f"in_ptr{i}"
|
||||
self._template_input_args.append((param_name, input_node))
|
||||
self._seen_input_args.add(param_name)
|
||||
|
||||
def render(self) -> str:
|
||||
"""
|
||||
Render the NVIDIA Universal GEMM kernel code as a Python source string.
|
||||
|
||||
Generates Python code that:
|
||||
1. Looks up the cutlass_api kernel by name from the manifest (cached in
|
||||
_nv_universal_gemm_kernel_cache to avoid repeated manifest searches)
|
||||
2. Creates GemmArguments with the input/output tensors and accumulator type
|
||||
3. Compiles the kernel for the specific tensor shapes/dtypes (cached in
|
||||
_nv_universal_gemm_artifact_cache keyed by (shape, dtype) tuple)
|
||||
4. Runs the kernel with the compiled artifact and CUDA stream
|
||||
|
||||
The caching strategy ensures:
|
||||
- Kernel lookup happens once per unique kernel name
|
||||
- Compilation happens once per unique (shape, dtype) combination
|
||||
- Runtime execution is just the kernel.run() call with cached artifact
|
||||
|
||||
Returns:
|
||||
Python source code string to be written to a .py file and loaded
|
||||
via async_compile.nv_universal_gemm()
|
||||
"""
|
||||
from torch._inductor.codegen.nv_universal_gemm.nv_universal_gemm import (
|
||||
GemmVariant,
|
||||
)
|
||||
|
||||
kernel_name_str = self.kernel_metadata["kernel_name"]
|
||||
is_grouped = self.variant == GemmVariant.GROUPED_GEMM
|
||||
is_scaled = self.variant == GemmVariant.SCALED_GEMM
|
||||
|
||||
acc_dtype_str = CuteDSLOpOverrides.TORCH_TO_CUTE_DTYPE.get(
|
||||
self.accumulator_type, "cutlass.Float32"
|
||||
)
|
||||
|
||||
input_params = [f"in_ptr{i}" for i, _ in enumerate(self.input_nodes)]
|
||||
input_params.append("out_ptr0")
|
||||
if self.workspace_size > 0:
|
||||
input_params.append("workspace")
|
||||
input_params.append("stream=None")
|
||||
params_str = ", ".join(input_params)
|
||||
|
||||
workspace_arg = "workspace" if self.workspace_size > 0 else "None"
|
||||
|
||||
var_prefix = self.variant.op_name.upper()
|
||||
cache_var = f"_{var_prefix}_compiled_cache"
|
||||
kernel_name_var = f"_{var_prefix}_KERNEL_NAME"
|
||||
|
||||
extra_imports = ""
|
||||
if is_scaled:
|
||||
extra_imports = """from cutlass_api.arguments import ScaledTensor
|
||||
from cutlass_api.library import ScaleMode, ScaleSwizzleMode"""
|
||||
|
||||
# Variant-specific code generation:
|
||||
# - cache_key_code: expression for cache key
|
||||
# - create_args_code: code to create Arguments object
|
||||
if is_grouped:
|
||||
cache_key_code = "(in_ptr0.shape, in_ptr0.dtype, in_ptr1.shape, in_ptr1.dtype, in_ptr2.shape)"
|
||||
create_args_code = f"""args = cutlass_api.arguments.GroupedGemmArguments(
|
||||
in_ptr0,
|
||||
in_ptr1,
|
||||
out_ptr0,
|
||||
accumulator_type={acc_dtype_str},
|
||||
offsets=in_ptr2,
|
||||
)"""
|
||||
elif is_scaled:
|
||||
scale_mode_a, swizzle_mode_a = to_cutlass_scale_mode(
|
||||
self.scale_type_a, self.swizzle_type_a
|
||||
)
|
||||
scale_mode_b, swizzle_mode_b = to_cutlass_scale_mode(
|
||||
self.scale_type_b, self.swizzle_type_b
|
||||
)
|
||||
scale_mode_a_str = scale_mode_a.name if scale_mode_a else ""
|
||||
scale_mode_b_str = scale_mode_b.name if scale_mode_b else ""
|
||||
swizzle_mode_a_str = swizzle_mode_a.name if swizzle_mode_a else ""
|
||||
swizzle_mode_b_str = swizzle_mode_b.name if swizzle_mode_b else ""
|
||||
cache_key_code = "(in_ptr0.shape, in_ptr0.dtype, in_ptr1.shape, in_ptr1.dtype, in_ptr2.shape, in_ptr3.shape)"
|
||||
create_args_code = f"""scaled_a = ScaledTensor(
|
||||
in_ptr0, in_ptr2, ScaleMode.{scale_mode_a_str}, ScaleSwizzleMode.{swizzle_mode_a_str}
|
||||
)
|
||||
scaled_b = ScaledTensor(
|
||||
in_ptr1, in_ptr3, ScaleMode.{scale_mode_b_str}, ScaleSwizzleMode.{swizzle_mode_b_str}
|
||||
)
|
||||
args = cutlass_api.arguments.GemmArguments(
|
||||
scaled_a,
|
||||
scaled_b,
|
||||
out_ptr0,
|
||||
accumulator_type={acc_dtype_str},
|
||||
)"""
|
||||
else:
|
||||
cache_key_code = (
|
||||
"(in_ptr0.shape, in_ptr0.dtype, in_ptr1.shape, in_ptr1.dtype)"
|
||||
)
|
||||
create_args_code = f"""args = cutlass_api.arguments.GemmArguments(
|
||||
in_ptr0,
|
||||
in_ptr1,
|
||||
out_ptr0,
|
||||
accumulator_type={acc_dtype_str},
|
||||
)"""
|
||||
|
||||
code = IndentedBuffer()
|
||||
code.splice(
|
||||
f"""
|
||||
import cutlass
|
||||
import cutlass_api
|
||||
from torch._inductor.codegen.nv_universal_gemm.kernel_cache import get_kernel_by_name
|
||||
{extra_imports}
|
||||
|
||||
{kernel_name_var} = "{kernel_name_str}"
|
||||
# Maps (shape, dtype, shape, dtype, ...) -> compiled kernel artifact
|
||||
{cache_var} = {{}}
|
||||
|
||||
def {self.kernel_name}_main({params_str}):
|
||||
global {cache_var}
|
||||
|
||||
kernel = get_kernel_by_name({kernel_name_var})
|
||||
if kernel is None:
|
||||
raise RuntimeError(f"Could not find kernel: {{{kernel_name_var}}}")
|
||||
|
||||
{create_args_code}
|
||||
|
||||
cache_key = {cache_key_code}
|
||||
artifact = {cache_var}.get(cache_key)
|
||||
if artifact is None:
|
||||
artifact = kernel.compile(args)
|
||||
{cache_var}[cache_key] = artifact
|
||||
|
||||
kernel.run(args, artifact, stream=stream, workspace={workspace_arg}, assume_supported_args=True)
|
||||
"""
|
||||
)
|
||||
|
||||
return code.getvalue()
|
||||
|
||||
def _get_reinterpret_view(self, node) -> ReinterpretView | None:
|
||||
"""Extract or convert to ReinterpretView from a node, handling all views."""
|
||||
while isinstance(node, MutableBox):
|
||||
node = node.data
|
||||
if isinstance(node, BaseView):
|
||||
return ExternKernel.convert_to_reinterpret_view(node)
|
||||
return None
|
||||
|
||||
def call_kernel(self, name: str, node=None):
|
||||
"""
|
||||
Generate the kernel call in the wrapper code.
|
||||
|
||||
Similar to CuteDSLTemplateKernel.call_kernel but simplified for NVIDIA Universal GEMM.
|
||||
"""
|
||||
wrapper = V.graph.wrapper_code
|
||||
|
||||
call_args: list[str] = []
|
||||
arg_types: list[Any] = []
|
||||
raw_args: list[Buffer | ReinterpretView | None] = []
|
||||
|
||||
for _, input_node in self._template_input_args:
|
||||
reinterpret_view = self._get_reinterpret_view(input_node)
|
||||
if reinterpret_view is not None:
|
||||
call_args.append(reinterpret_view.codegen_reference())
|
||||
# Pass the ReinterpretView as raw_arg so autotune_at_compile_time
|
||||
# can use it to generate example tensors
|
||||
raw_args.append(reinterpret_view)
|
||||
else:
|
||||
call_args.append(input_node.get_name())
|
||||
raw_args.append(input_node)
|
||||
arg_types.append(V.graph.get_dtype(input_node.get_name()))
|
||||
|
||||
output_name = self.output_node.get_name()
|
||||
call_args.append(output_name)
|
||||
arg_types.append(V.graph.get_dtype(output_name))
|
||||
raw_args.append(None) # Output buffer is findable by name
|
||||
|
||||
# Allocate workspace if needed
|
||||
ws: WorkspaceArg | None = None
|
||||
if self.workspace_size > 0:
|
||||
ws = WorkspaceArg(
|
||||
count=self.workspace_size,
|
||||
device=V.graph.get_current_device_or_throw(),
|
||||
zero_mode=WorkspaceZeroMode.UNINITIALIZED,
|
||||
outer_name=WorkspaceArg.unique_name(),
|
||||
)
|
||||
wrapper.generate_workspace_allocation(ws)
|
||||
call_args.append(ws.outer_name)
|
||||
arg_types.append(ws.dtype)
|
||||
raw_args.append(None)
|
||||
|
||||
# Generate the kernel call using triton=True for Python-based kernels
|
||||
# Pass raw_keys as None list to match raw_args length
|
||||
# TODO(nikhilap) We don't use autotune_args like the Triton path
|
||||
wrapper.generate_kernel_call(
|
||||
name,
|
||||
call_args,
|
||||
triton=True,
|
||||
arg_types=arg_types,
|
||||
raw_args=raw_args,
|
||||
raw_keys=[None] * len(raw_args),
|
||||
)
|
||||
|
||||
# Deallocate workspace after kernel call
|
||||
if ws is not None:
|
||||
wrapper.generate_workspace_deallocation(ws)
|
||||
+153
@@ -0,0 +1,153 @@
|
||||
# mypy: allow-untyped-defs
|
||||
"""
|
||||
NVIDIA Universal GEMM scheduling for PyTorch Inductor.
|
||||
"""
|
||||
|
||||
import hashlib
|
||||
import logging
|
||||
from collections.abc import Sequence
|
||||
from typing import cast
|
||||
|
||||
from torch._inductor.utils import (
|
||||
get_fused_kernel_name,
|
||||
get_kernel_metadata,
|
||||
Placeholder,
|
||||
)
|
||||
from torch.utils._ordered_set import OrderedSet
|
||||
|
||||
from ... import config
|
||||
from ...codecache import code_hash, get_path
|
||||
from ...ir import NVUniversalGemmBuffer
|
||||
from ...scheduler import (
|
||||
BaseSchedulerNode,
|
||||
BaseScheduling,
|
||||
FusedSchedulerNode,
|
||||
SchedulerNode,
|
||||
)
|
||||
from ...virtualized import V
|
||||
from ..common import BackendFeature, IndentedBuffer
|
||||
|
||||
|
||||
log = logging.getLogger(__name__)
|
||||
|
||||
MAIN_SUFFIX = "main"
|
||||
|
||||
|
||||
class NVUniversalGemmScheduling(BaseScheduling):
|
||||
"""
|
||||
Scheduling implementation for NVIDIA Universal GEMM kernels.
|
||||
|
||||
This class is intended to be used in combination with other schedulers,
|
||||
and delegated to by CUDACombinedScheduling.
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def get_backend_features(cls, device) -> OrderedSet[BackendFeature]:
|
||||
return OrderedSet()
|
||||
|
||||
@staticmethod
|
||||
def is_nv_universal_gemm_template(node: BaseSchedulerNode) -> bool:
|
||||
"""Check if a node is a NVIDIA Universal GEMM template."""
|
||||
return isinstance(node, SchedulerNode) and isinstance(
|
||||
node.node, NVUniversalGemmBuffer
|
||||
)
|
||||
|
||||
def is_nv_universal_gemm_fused_template(self, node: BaseSchedulerNode) -> bool:
|
||||
"""Check if a node is a fused NVIDIA Universal GEMM template."""
|
||||
return isinstance(
|
||||
node, FusedSchedulerNode
|
||||
) and self.is_nv_universal_gemm_template(node)
|
||||
|
||||
def can_fuse_vertical(
|
||||
self, node1: BaseSchedulerNode, node2: BaseSchedulerNode
|
||||
) -> bool:
|
||||
# NVIDIA Universal GEMM templates don't support vertical fusion yet
|
||||
return False
|
||||
|
||||
def can_fuse_horizontal(
|
||||
self, node1: BaseSchedulerNode, node2: BaseSchedulerNode
|
||||
) -> bool:
|
||||
# NVIDIA Universal GEMM templates don't support horizontal fusion yet
|
||||
return False
|
||||
|
||||
def define_kernel(self, src_code: str, node_schedule) -> str:
|
||||
"""
|
||||
Define a NVIDIA Universal GEMM kernel by writing source code and generating wrapper.
|
||||
|
||||
Based on CuteDSLScheduling.define_kernel.
|
||||
"""
|
||||
wrapper = V.graph.wrapper_code
|
||||
|
||||
# Use the string as the key for caching
|
||||
if src_code in wrapper.src_to_kernel:
|
||||
return wrapper.src_to_kernel[src_code]
|
||||
|
||||
fused_name = (
|
||||
get_fused_kernel_name(node_schedule, config.triton.descriptive_names)
|
||||
if config.triton.descriptive_names
|
||||
else ""
|
||||
)
|
||||
|
||||
kernel_hash = hashlib.sha256(src_code.encode("utf-8")).hexdigest()[:8]
|
||||
if fused_name == "fused":
|
||||
kernel_name = f"nv_universal_gemm_{kernel_hash}"
|
||||
else:
|
||||
kernel_name = f"nv_universal_gemm_{fused_name}_{kernel_hash}"
|
||||
|
||||
wrapper.src_to_kernel[src_code] = kernel_name
|
||||
|
||||
src_code = src_code.replace(str(Placeholder.KERNEL_NAME), kernel_name)
|
||||
|
||||
_, _, kernel_path = get_path(code_hash(src_code), "py")
|
||||
|
||||
compile_wrapper = IndentedBuffer()
|
||||
compile_wrapper.writeline(
|
||||
f"async_compile.nv_universal_gemm({kernel_name!r}, r'''"
|
||||
)
|
||||
compile_wrapper.splice(src_code, strip=True)
|
||||
compile_wrapper.writeline("''')")
|
||||
|
||||
metadata_comment = f"# kernel path: {kernel_path}"
|
||||
origins, detailed_origins = get_kernel_metadata(node_schedule, wrapper)
|
||||
metadata_comment += "\n" + origins + "\n" + detailed_origins
|
||||
wrapper.define_kernel(kernel_name, compile_wrapper.getvalue(), metadata_comment)
|
||||
|
||||
return kernel_name
|
||||
|
||||
def codegen_template(
|
||||
self,
|
||||
template_node: BaseSchedulerNode,
|
||||
epilogue_nodes: Sequence[BaseSchedulerNode],
|
||||
prologue_nodes: Sequence[BaseSchedulerNode],
|
||||
):
|
||||
"""
|
||||
Codegen a NVIDIA Universal GEMM template. Currently doesn't support fusion.
|
||||
"""
|
||||
assert self.is_nv_universal_gemm_template(template_node), (
|
||||
"Template node passed to NVUniversalGemmScheduling.codegen_template must be a "
|
||||
"SchedulerNode that wraps a NVUniversalGemmBuffer"
|
||||
)
|
||||
# TODO: add support for fusion when needed
|
||||
assert not epilogue_nodes, (
|
||||
"NVIDIA Universal GEMM doesn't support epilogue fusion yet"
|
||||
)
|
||||
assert not prologue_nodes, (
|
||||
"NVIDIA Universal GEMM doesn't support prologue fusion yet"
|
||||
)
|
||||
|
||||
template_node = cast(SchedulerNode, template_node)
|
||||
ctb: NVUniversalGemmBuffer = cast(NVUniversalGemmBuffer, template_node.node)
|
||||
|
||||
assert ctb.make_kernel_render is not None
|
||||
kernel, render = ctb.make_kernel_render(ctb)
|
||||
template_node.mark_run()
|
||||
src_code = render()
|
||||
|
||||
with V.set_kernel_handler(kernel):
|
||||
node_schedule = [template_node]
|
||||
kernel_name = self.define_kernel(src_code, node_schedule)
|
||||
|
||||
self.codegen_comment(node_schedule, kernel_name)
|
||||
kernel.call_kernel(kernel_name, ctb)
|
||||
V.graph.removed_buffers |= kernel.removed_buffers
|
||||
self.free_buffers_in_scheduler()
|
||||
+42
@@ -0,0 +1,42 @@
|
||||
# mypy: allow-untyped-defs
|
||||
"""
|
||||
Utility functions for NVIDIA Universal GEMM.
|
||||
"""
|
||||
|
||||
from typing import Any
|
||||
|
||||
from torch.nn.functional import ScalingType, SwizzleType
|
||||
|
||||
|
||||
def to_cutlass_scale_mode(
|
||||
scale_type: Any, swizzle_type: Any
|
||||
) -> tuple[Any | None, Any | None]:
|
||||
"""
|
||||
Map PyTorch ScalingType/SwizzleType to cutlass_api ScaleMode/ScaleSwizzleMode.
|
||||
|
||||
Args:
|
||||
scale_type: ScalingType from torch.nn.functional
|
||||
swizzle_type: SwizzleType from torch.nn.functional
|
||||
|
||||
Returns:
|
||||
Tuple of (ScaleMode, ScaleSwizzleMode) from cutlass_api.library,
|
||||
or (None, None) if the types are not supported.
|
||||
|
||||
The returned enum objects can be used directly with cutlass_api, or their
|
||||
.name attribute can be used for codegen (e.g., scale_mode.name -> "Blockwise1x32").
|
||||
|
||||
NOTE:
|
||||
Currently on Blackwell (SM100), NVGEMM only supports MXFP8 scaling modes.
|
||||
Update this mapping when additional scaling modes are added.
|
||||
"""
|
||||
from cutlass_api.library import ScaleMode, ScaleSwizzleMode
|
||||
|
||||
scale_mode_map = {
|
||||
ScalingType.BlockWise1x32: ScaleMode.Blockwise1x32,
|
||||
ScalingType.BlockWise1x16: ScaleMode.Blockwise1x16,
|
||||
}
|
||||
swizzle_mode_map = {
|
||||
SwizzleType.SWIZZLE_32_4_4: ScaleSwizzleMode.Swizzle32x4x4,
|
||||
SwizzleType.NO_SWIZZLE: ScaleSwizzleMode.SwizzleNone,
|
||||
}
|
||||
return scale_mode_map.get(scale_type), swizzle_mode_map.get(swizzle_type)
|
||||
File diff suppressed because it is too large
Load Diff
+33
@@ -0,0 +1,33 @@
|
||||
from typing_extensions import override
|
||||
|
||||
from torch._inductor import ir
|
||||
|
||||
from .wrapper import PythonWrapperCodegen
|
||||
|
||||
|
||||
class PythonWrapperMtia(PythonWrapperCodegen):
|
||||
"""
|
||||
A thin wrapper of PythonWrapperCodegen with MTIA specific logic
|
||||
"""
|
||||
|
||||
@override
|
||||
def write_header(self) -> None:
|
||||
super().write_header()
|
||||
|
||||
# MITA specific imports
|
||||
self.imports.splice("import mtia.host_runtime.torch_mtia.dynamic_library")
|
||||
|
||||
@override
|
||||
@staticmethod
|
||||
def create(
|
||||
is_subgraph: bool,
|
||||
subgraph_name: str | None,
|
||||
parent_wrapper: PythonWrapperCodegen | None,
|
||||
partition_signatures: ir.GraphPartitionSignature | None = None,
|
||||
) -> PythonWrapperCodegen:
|
||||
if is_subgraph:
|
||||
# Delegate to the parent class to handle the case of subgraph
|
||||
return PythonWrapperCodegen.create(
|
||||
is_subgraph, subgraph_name, parent_wrapper, partition_signatures
|
||||
)
|
||||
return PythonWrapperMtia()
|
||||
+627
@@ -0,0 +1,627 @@
|
||||
# mypy: allow-untyped-defs
|
||||
import copy
|
||||
import logging
|
||||
import random
|
||||
from typing import Any
|
||||
from typing_extensions import override
|
||||
|
||||
from torch._inductor.virtualized import V
|
||||
|
||||
from .rocm_template import ArgInfo
|
||||
|
||||
|
||||
try:
|
||||
import ck4inductor # type: ignore[import]
|
||||
except ImportError:
|
||||
ck4inductor = None
|
||||
|
||||
if ck4inductor is not None:
|
||||
from ck4inductor.grouped_conv_fwd.gen_instances import ( # type: ignore[import]
|
||||
gen_conv_ops_library,
|
||||
)
|
||||
from ck4inductor.grouped_conv_fwd.op import ( # type: ignore[import] # noqa: TCH002
|
||||
CKGroupedConvFwdOp,
|
||||
)
|
||||
else:
|
||||
|
||||
def gen_conv_ops_library():
|
||||
return []
|
||||
|
||||
|
||||
from torch._inductor import config
|
||||
from torch._inductor.codegen.rocm.ck_template import CKTemplate
|
||||
from torch._inductor.codegen.rocm.rocm_kernel import ROCmTemplateKernel
|
||||
from torch._inductor.utils import IndentedBuffer
|
||||
|
||||
|
||||
log = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def torch_layout_to_ck_layouts(torch_layout):
|
||||
# logically, torch tensors are always NCHW,
|
||||
# and channels-last memory layout is visible in the strides
|
||||
if V.graph.sizevars.statically_known_equals(torch_layout.stride[-1], 1):
|
||||
# when input or output is NCHW
|
||||
# NB: torch.conv2d result is always NCHW
|
||||
return ["NGCHW", "GKCYX", "NGKHW"]
|
||||
elif V.graph.sizevars.statically_known_equals(torch_layout.stride[-3], 1):
|
||||
# when input or output or weight is channels-last
|
||||
return ["NHWGC", "GKYXC", "NHWGK"]
|
||||
else:
|
||||
return None
|
||||
|
||||
|
||||
def torch_layout_to_ck_input_layout(torch_layout):
|
||||
if V.graph.sizevars.statically_known_equals(torch_layout.stride[-1], 1):
|
||||
return "NGCHW"
|
||||
elif V.graph.sizevars.statically_known_equals(torch_layout.stride[-3], 1):
|
||||
return "NHWGC"
|
||||
else:
|
||||
return None
|
||||
|
||||
|
||||
def torch_layout_to_ck_weight_layout(torch_layout):
|
||||
if V.graph.sizevars.statically_known_equals(torch_layout.stride[-1], 1):
|
||||
return "GKCYX"
|
||||
elif V.graph.sizevars.statically_known_equals(torch_layout.stride[-3], 1):
|
||||
return "GKYXC"
|
||||
else:
|
||||
return None
|
||||
|
||||
|
||||
def torch_layout_to_ck_output_layout(torch_layout):
|
||||
if V.graph.sizevars.statically_known_equals(torch_layout.stride[-1], 1):
|
||||
return "NGKHW"
|
||||
elif V.graph.sizevars.statically_known_equals(torch_layout.stride[-3], 1):
|
||||
return "NHWGK"
|
||||
else:
|
||||
return None
|
||||
|
||||
|
||||
class CKGroupedConvFwdTemplate(CKTemplate):
|
||||
conv_template = r"""
|
||||
{{headers}}
|
||||
{{globals}}
|
||||
{{instance_definition}}
|
||||
extern "C" {
|
||||
PT_EXPORT {{kernel_definition}} {
|
||||
auto conv = {{instance_type}} {};
|
||||
auto invoker = conv.MakeInvoker();
|
||||
|
||||
using ck::index_t;
|
||||
|
||||
constexpr index_t NumDTensor = {{n_d_tensors}};
|
||||
constexpr index_t NDimSpatial = {{n_dim_spatial}};
|
||||
const std::vector<index_t> FilterSize = { FilterSize_0, FilterSize_1 };
|
||||
const std::vector<index_t> InputSize = { InputSize_0, InputSize_1 };
|
||||
const std::vector<index_t> ConvolutionStrides = { ConvolutionStrides_0, ConvolutionStrides_1 };
|
||||
const std::vector<index_t> Dilations = { Dilations_0, Dilations_1 };
|
||||
const std::vector<index_t> LeftPads = { LeftPads_0, LeftPads_1 };
|
||||
const std::vector<index_t> RightPads = { RightPads_0, RightPads_1 };
|
||||
|
||||
|
||||
auto conv_param = ck::utils::conv::ConvParam {
|
||||
NDimSpatial,
|
||||
GroupCount,
|
||||
NBatch,
|
||||
NOutChannels,
|
||||
NInChannels,
|
||||
FilterSize,
|
||||
InputSize,
|
||||
ConvolutionStrides,
|
||||
Dilations,
|
||||
LeftPads,
|
||||
RightPads,
|
||||
};
|
||||
|
||||
using InLayout = ck::tensor_layout::convolution::{{input_layout}};
|
||||
using WeiLayout = ck::tensor_layout::convolution::{{weight_layout}};
|
||||
using OutLayout = ck::tensor_layout::convolution::{{output_layout}};
|
||||
|
||||
const auto in_g_n_c_wis_desc =
|
||||
ck::utils::conv::make_input_host_tensor_descriptor_g_n_c_wis_packed<InLayout>(conv_param);
|
||||
const auto wei_g_k_c_xs_desc =
|
||||
ck::utils::conv::make_weight_host_tensor_descriptor_g_k_c_xs_packed<WeiLayout>(conv_param);
|
||||
const auto out_g_n_k_wos_desc =
|
||||
ck::utils::conv::make_output_host_tensor_descriptor_g_n_k_wos_packed<OutLayout>(conv_param);
|
||||
|
||||
const void* p_a = input;
|
||||
const void* p_b = weight;
|
||||
const std::array<const void*, NumDTensor> p_ds;
|
||||
void* p_e = output;
|
||||
std::array<index_t, NDimSpatial + 3> a_g_n_c_wis_lengths;
|
||||
std::array<index_t, NDimSpatial + 3> a_g_n_c_wis_strides;
|
||||
std::array<index_t, NDimSpatial + 3> b_g_k_c_xs_lengths;
|
||||
std::array<index_t, NDimSpatial + 3> b_g_k_c_xs_strides;
|
||||
std::array<std::array<index_t, NDimSpatial + 3>, NumDTensor> ds_g_n_k_wos_lengths;
|
||||
std::array<std::array<index_t, NDimSpatial + 3>, NumDTensor> ds_g_n_k_wos_strides;
|
||||
std::array<index_t, NDimSpatial + 3> e_g_n_k_wos_lengths;
|
||||
std::array<index_t, NDimSpatial + 3> e_g_n_k_wos_strides;
|
||||
std::array<index_t, NDimSpatial> conv_filter_strides;
|
||||
std::array<index_t, NDimSpatial> conv_filter_dilations;
|
||||
std::array<index_t, NDimSpatial> input_left_pads;
|
||||
std::array<index_t, NDimSpatial> input_right_pads;
|
||||
const auto a_element_op = PassThrough {};
|
||||
const auto b_element_op = PassThrough {};
|
||||
const auto cde_element_op = PassThrough {};
|
||||
|
||||
auto copy = [](auto& x, auto& y) { ck::ranges::copy(x, y.begin()); };
|
||||
|
||||
copy(in_g_n_c_wis_desc.GetLengths(), a_g_n_c_wis_lengths);
|
||||
copy(in_g_n_c_wis_desc.GetStrides(), a_g_n_c_wis_strides);
|
||||
copy(wei_g_k_c_xs_desc.GetLengths(), b_g_k_c_xs_lengths);
|
||||
copy(wei_g_k_c_xs_desc.GetStrides(), b_g_k_c_xs_strides);
|
||||
copy(out_g_n_k_wos_desc.GetLengths(), e_g_n_k_wos_lengths);
|
||||
copy(out_g_n_k_wos_desc.GetStrides(), e_g_n_k_wos_strides);
|
||||
copy(conv_param.conv_filter_strides_, conv_filter_strides);
|
||||
copy(conv_param.conv_filter_dilations_, conv_filter_dilations);
|
||||
copy(conv_param.input_left_pads_, input_left_pads);
|
||||
copy(conv_param.input_right_pads_, input_right_pads);
|
||||
|
||||
auto argument = conv.MakeArgument(
|
||||
p_a,
|
||||
p_b,
|
||||
p_ds,
|
||||
p_e,
|
||||
a_g_n_c_wis_lengths,
|
||||
a_g_n_c_wis_strides,
|
||||
b_g_k_c_xs_lengths,
|
||||
b_g_k_c_xs_strides,
|
||||
ds_g_n_k_wos_lengths,
|
||||
ds_g_n_k_wos_strides,
|
||||
e_g_n_k_wos_lengths,
|
||||
e_g_n_k_wos_strides,
|
||||
conv_filter_strides,
|
||||
conv_filter_dilations,
|
||||
input_left_pads,
|
||||
input_right_pads,
|
||||
a_element_op,
|
||||
b_element_op,
|
||||
cde_element_op
|
||||
);
|
||||
if (!conv.IsSupportedArgument(argument)) {
|
||||
// we do our best to statically avoid this case in `filter_op`
|
||||
std::cerr << "invalid argument for conv instance " << conv.GetTypeString() << std::endl;
|
||||
argument.Print();
|
||||
return -23;
|
||||
}
|
||||
if (workspace_size) {
|
||||
*workspace_size = conv.GetWorkSpaceSize(&argument);
|
||||
return 0;
|
||||
}
|
||||
|
||||
if (p_a == nullptr) {
|
||||
std::cerr << "p_a is nullptr" << std::endl;
|
||||
return -1;
|
||||
}
|
||||
if (p_b == nullptr) {
|
||||
std::cerr << "p_b is nullptr" << std::endl;
|
||||
return -1;
|
||||
}
|
||||
if (p_e == nullptr) {
|
||||
std::cerr << "p_e is nullptr" << std::endl;
|
||||
return -1;
|
||||
}
|
||||
|
||||
// when debugging, do time kernel to serialize launches
|
||||
auto stream_config = StreamConfig{stream, /* time kernel */ false, /* log level */ 0};
|
||||
|
||||
if (workspace != nullptr) {
|
||||
conv.SetWorkSpacePointer(&argument, workspace, stream_config);
|
||||
}
|
||||
|
||||
// run the kernel
|
||||
float elapsed_time = invoker.Run(argument, stream_config);
|
||||
return 0;
|
||||
} // kernel definition
|
||||
} // extern C
|
||||
|
||||
#ifdef GENERATE_CK_STANDALONE_RUNNER
|
||||
int main(int argc, char** argv) {
|
||||
(void) argc;
|
||||
(void) argv;
|
||||
return 0;
|
||||
}
|
||||
#endif // GENERATE_CK_STANDALONE_RUNNER
|
||||
"""
|
||||
|
||||
def globals(self) -> IndentedBuffer:
|
||||
res = super().globals()
|
||||
res.splice(
|
||||
"""
|
||||
// CK conv globals
|
||||
|
||||
using NWC = ck::tensor_layout::convolution::NWC;
|
||||
using NHWC = ck::tensor_layout::convolution::NHWC;
|
||||
using NDHWC = ck::tensor_layout::convolution::NDHWC;
|
||||
|
||||
using KXC = ck::tensor_layout::convolution::KXC;
|
||||
using KYXC = ck::tensor_layout::convolution::KYXC;
|
||||
using KZYXC = ck::tensor_layout::convolution::KZYXC;
|
||||
|
||||
using NWK = ck::tensor_layout::convolution::NWK;
|
||||
using NHWK = ck::tensor_layout::convolution::NHWK;
|
||||
using NDHWK = ck::tensor_layout::convolution::NDHWK;
|
||||
|
||||
using GNWC = ck::tensor_layout::convolution::GNWC;
|
||||
using GNHWC = ck::tensor_layout::convolution::GNHWC;
|
||||
using GNDHWC = ck::tensor_layout::convolution::GNDHWC;
|
||||
|
||||
using GKXC = ck::tensor_layout::convolution::GKXC;
|
||||
using GKYXC = ck::tensor_layout::convolution::GKYXC;
|
||||
using GKZYXC = ck::tensor_layout::convolution::GKZYXC;
|
||||
|
||||
using GKCX = ck::tensor_layout::convolution::GKCX;
|
||||
using GKCYX = ck::tensor_layout::convolution::GKCYX;
|
||||
using GKCZYX = ck::tensor_layout::convolution::GKCZYX;
|
||||
|
||||
using GNWK = ck::tensor_layout::convolution::GNWK;
|
||||
using GNHWK = ck::tensor_layout::convolution::GNHWK;
|
||||
using GNDHWK = ck::tensor_layout::convolution::GNDHWK;
|
||||
|
||||
using NGKW = ck::tensor_layout::convolution::NGKW;
|
||||
using NGKHW = ck::tensor_layout::convolution::NGKHW;
|
||||
using NGKDHW = ck::tensor_layout::convolution::NGKDHW;
|
||||
|
||||
using NWGC = ck::tensor_layout::convolution::NWGC;
|
||||
using NHWGC = ck::tensor_layout::convolution::NHWGC;
|
||||
using NDHWGC = ck::tensor_layout::convolution::NDHWGC;
|
||||
|
||||
using KXGC = ck::tensor_layout::convolution::KXGC;
|
||||
using KYXGC = ck::tensor_layout::convolution::KYXGC;
|
||||
using KZYXGC = ck::tensor_layout::convolution::KZYXGC;
|
||||
|
||||
using NWGK = ck::tensor_layout::convolution::NWGK;
|
||||
using NHWGK = ck::tensor_layout::convolution::NHWGK;
|
||||
using NDHWGK = ck::tensor_layout::convolution::NDHWGK;
|
||||
|
||||
using NGCW = ck::tensor_layout::convolution::NGCW;
|
||||
using NGCHW = ck::tensor_layout::convolution::NGCHW;
|
||||
using NGCDHW = ck::tensor_layout::convolution::NGCDHW;
|
||||
|
||||
using G_K = ck::tensor_layout::convolution::G_K;
|
||||
|
||||
using BlockGemmPipelineScheduler = ck::BlockGemmPipelineScheduler;
|
||||
using GemmSpecialization = ck::tensor_operation::device::GemmSpecialization;
|
||||
using BlockGemmPipelineVersion = ck::BlockGemmPipelineVersion;
|
||||
|
||||
using ConvolutionForwardSpecialization = ck::tensor_operation::device::ConvolutionForwardSpecialization;
|
||||
|
||||
using OutElementOp = PassThrough;
|
||||
|
||||
namespace ck {
|
||||
namespace utils {
|
||||
namespace conv {
|
||||
|
||||
ConvParam::ConvParam(ck::index_t n_dim,
|
||||
ck::index_t group_count,
|
||||
ck::index_t n_batch,
|
||||
ck::index_t n_out_channels,
|
||||
ck::index_t n_in_channels,
|
||||
const std::vector<ck::index_t>& filters_len,
|
||||
const std::vector<ck::index_t>& input_len,
|
||||
const std::vector<ck::index_t>& strides,
|
||||
const std::vector<ck::index_t>& dilations,
|
||||
const std::vector<ck::index_t>& left_pads,
|
||||
const std::vector<ck::index_t>& right_pads)
|
||||
: num_dim_spatial_(static_cast<ck::long_index_t>(n_dim)),
|
||||
G_(static_cast<ck::long_index_t>(group_count)),
|
||||
N_(static_cast<ck::long_index_t>(n_batch)),
|
||||
K_(static_cast<ck::long_index_t>(n_out_channels)),
|
||||
C_(static_cast<ck::long_index_t>(n_in_channels)),
|
||||
filter_spatial_lengths_(num_dim_spatial_),
|
||||
input_spatial_lengths_(num_dim_spatial_),
|
||||
output_spatial_lengths_(num_dim_spatial_),
|
||||
conv_filter_strides_(num_dim_spatial_),
|
||||
conv_filter_dilations_(num_dim_spatial_),
|
||||
input_left_pads_(num_dim_spatial_),
|
||||
input_right_pads_(num_dim_spatial_)
|
||||
{
|
||||
if(static_cast<ck::index_t>(filter_spatial_lengths_.size()) != num_dim_spatial_ ||
|
||||
static_cast<ck::index_t>(input_spatial_lengths_.size()) != num_dim_spatial_ ||
|
||||
static_cast<ck::index_t>(conv_filter_strides_.size()) != num_dim_spatial_ ||
|
||||
static_cast<ck::index_t>(conv_filter_dilations_.size()) != num_dim_spatial_ ||
|
||||
static_cast<ck::index_t>(input_left_pads_.size()) != num_dim_spatial_ ||
|
||||
static_cast<ck::index_t>(input_right_pads_.size()) != num_dim_spatial_)
|
||||
{
|
||||
throw(
|
||||
std::runtime_error("ConvParam::ConvParam: "
|
||||
"parameter size is different from number of declared dimensions!"));
|
||||
}
|
||||
|
||||
for(ck::index_t i = 0; i < num_dim_spatial_; ++i)
|
||||
{
|
||||
filter_spatial_lengths_[i] = static_cast<ck::long_index_t>(filters_len[i]);
|
||||
input_spatial_lengths_[i] = static_cast<ck::long_index_t>(input_len[i]);
|
||||
conv_filter_strides_[i] = static_cast<ck::long_index_t>(strides[i]);
|
||||
conv_filter_dilations_[i] = static_cast<ck::long_index_t>(dilations[i]);
|
||||
input_left_pads_[i] = static_cast<ck::long_index_t>(left_pads[i]);
|
||||
input_right_pads_[i] = static_cast<ck::long_index_t>(right_pads[i]);
|
||||
|
||||
// XEff = (X - 1) * conv_dilation_w + 1;
|
||||
// Wo = (Wi + in_left_pad_w + in_right_pad_w - XEff) / conv_stride_w + 1;
|
||||
const ck::long_index_t x_eff =
|
||||
(filter_spatial_lengths_[i] - 1) * conv_filter_dilations_[i] + 1;
|
||||
|
||||
output_spatial_lengths_[i] =
|
||||
(input_spatial_lengths_[i] + input_left_pads_[i] + input_right_pads_[i] - x_eff) /
|
||||
conv_filter_strides_[i] +
|
||||
1;
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace conv
|
||||
} // namespace utils
|
||||
} // namespace ck
|
||||
|
||||
const std::vector<std::size_t>& HostTensorDescriptor::GetLengths() const { return mLens; }
|
||||
const std::vector<std::size_t>& HostTensorDescriptor::GetStrides() const { return mStrides; }
|
||||
std::size_t HostTensorDescriptor::GetNumOfDimension() const { return mLens.size(); }
|
||||
void HostTensorDescriptor::CalculateStrides() {
|
||||
mStrides.clear();
|
||||
mStrides.resize(mLens.size(), 0);
|
||||
if(mStrides.empty())
|
||||
return;
|
||||
|
||||
mStrides.back() = 1;
|
||||
std::partial_sum(
|
||||
mLens.rbegin(), mLens.rend() - 1, mStrides.rbegin() + 1, std::multiplies<std::size_t>());
|
||||
}
|
||||
"""
|
||||
)
|
||||
return res
|
||||
|
||||
def header(self) -> IndentedBuffer:
|
||||
res = super().header()
|
||||
res.splice(
|
||||
"""
|
||||
// CK conv headers
|
||||
|
||||
#include "ck/tensor_operation/gpu/device/impl/device_grouped_conv_fwd_multiple_abd_xdl_cshuffle_v3.hpp"
|
||||
#include "ck/tensor_operation/gpu/device/convolution_forward_specialization.hpp"
|
||||
#include "ck/tensor_operation/gpu/device/gemm_specialization.hpp"
|
||||
|
||||
#include "ck/library/utility/convolution_parameter.hpp"
|
||||
#include "ck/library/utility/convolution_host_tensor_descriptor_helper.hpp"
|
||||
"""
|
||||
)
|
||||
return res
|
||||
|
||||
@staticmethod
|
||||
def add_ck_conv_choices(
|
||||
choices,
|
||||
layout,
|
||||
input_nodes,
|
||||
*,
|
||||
stride,
|
||||
padding,
|
||||
dilation,
|
||||
groups,
|
||||
n_spatial_dimensions,
|
||||
):
|
||||
template = CKGroupedConvFwdTemplate(
|
||||
input_nodes,
|
||||
layout,
|
||||
stride=stride,
|
||||
padding=padding,
|
||||
dilation=dilation,
|
||||
groups=groups,
|
||||
n_spatial_dimensions=n_spatial_dimensions,
|
||||
)
|
||||
ops = template.gen_ops()
|
||||
for op in ops:
|
||||
template.maybe_append_choice(
|
||||
choices,
|
||||
op=op,
|
||||
)
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
input_nodes,
|
||||
layout,
|
||||
*,
|
||||
stride,
|
||||
padding,
|
||||
dilation,
|
||||
groups,
|
||||
n_spatial_dimensions,
|
||||
):
|
||||
super().__init__(
|
||||
"ck_conv_template",
|
||||
input_nodes,
|
||||
layout,
|
||||
)
|
||||
self.stride = stride
|
||||
self.padding = padding
|
||||
self.dilation = dilation
|
||||
self.groups = groups
|
||||
self.n_spatial_dimensions = n_spatial_dimensions
|
||||
|
||||
def filter_op(self, op: "CKGroupedConvFwdOp"): # type: ignore[name-defined]
|
||||
metas = [
|
||||
T.get_layout()
|
||||
for T in [*self.input_nodes, self.output_node]
|
||||
if T is not None
|
||||
]
|
||||
X_meta = metas[0]
|
||||
W_meta = metas[1]
|
||||
Y_meta = metas[-1]
|
||||
# disable the instance if dtypes don't match
|
||||
if op.a_element_dtype != self._TORCH_DTYPE_TO_CK[X_meta.dtype]:
|
||||
return None
|
||||
if op.b_element_dtype != self._TORCH_DTYPE_TO_CK[W_meta.dtype]:
|
||||
return None
|
||||
if op.e_element_dtype != self._TORCH_DTYPE_TO_CK[Y_meta.dtype]:
|
||||
return None
|
||||
# disable the instance if layouts don't match
|
||||
if op.a_layout != torch_layout_to_ck_input_layout(X_meta):
|
||||
return None
|
||||
if op.b_layout != torch_layout_to_ck_weight_layout(W_meta):
|
||||
return None
|
||||
if op.e_layout != torch_layout_to_ck_output_layout(Y_meta):
|
||||
return None
|
||||
# disable the instance if number of spatial dimensions doesn't match
|
||||
if op.n_dim_spatial != self.n_spatial_dimensions:
|
||||
return None
|
||||
# disable 1x1 and odd-channels conv specializations for now
|
||||
if "Default" not in op.conv_forward_specialization:
|
||||
return None
|
||||
return op
|
||||
|
||||
def gen_ops(self):
|
||||
unfiltered_instances = gen_conv_ops_library()
|
||||
|
||||
filtered_instances = list(
|
||||
filter(lambda op: self.filter_op(op), unfiltered_instances)
|
||||
)
|
||||
# NB: when using a fixed list order, most likely we will pick the subset of instances
|
||||
# which are very similar to each other. Randomizing the choice seems to solve this.
|
||||
random.seed(-11)
|
||||
chosen_instances = (
|
||||
random.sample(
|
||||
filtered_instances,
|
||||
min(len(filtered_instances), config.rocm.ck_max_profiling_configs),
|
||||
)
|
||||
if config.rocm.ck_max_profiling_configs
|
||||
else filtered_instances
|
||||
)
|
||||
log.debug(
|
||||
"generated %d ck instances after filter: %s",
|
||||
len(chosen_instances),
|
||||
chosen_instances,
|
||||
)
|
||||
return chosen_instances
|
||||
|
||||
def emit_ck_instance(self, op: "CKGroupedConvFwdOp") -> tuple[str, str]: # type: ignore[name-defined]
|
||||
# The Jinja template for generating a C++ type alias *definition* for a Universal GEMM instance
|
||||
template_definition = r"""
|
||||
// Gemm operator {{operation_name}}
|
||||
using Operation_{{operation_name}} =
|
||||
ck::tensor_operation::device::DeviceGroupedConvFwdMultipleABD_Xdl_CShuffle_V3<
|
||||
{{template_params}}>;
|
||||
|
||||
"""
|
||||
# The Jinja template for generating a C++ type alias *usage* for a Universal GEMM instance
|
||||
template_type = r"""
|
||||
Operation_{{operation_name}}
|
||||
"""
|
||||
template_params = []
|
||||
for field_name, field_value in op.dict_items():
|
||||
if isinstance(field_value, tuple):
|
||||
tuple_elements = ", ".join(map(str, iter(field_value)))
|
||||
if "ds" in field_name: # element type and layout for bias
|
||||
arg = f"/* {field_name} */ Tuple<{tuple_elements}>"
|
||||
else: # tile shape
|
||||
arg = f"/* {field_name} */ S<{tuple_elements}>"
|
||||
# pyrefly: ignore [bad-argument-type]
|
||||
template_params.append(arg)
|
||||
else:
|
||||
if field_value is not None:
|
||||
# pyrefly: ignore [bad-argument-type]
|
||||
template_params.append(f"/* {field_name} */ {field_value}")
|
||||
return self._template_from_string(template_definition).render(
|
||||
operation_name=op.name(),
|
||||
template_params=(",\n" + 12 * " ").join(template_params),
|
||||
), self._template_from_string(template_type).render(operation_name=op.name())
|
||||
|
||||
def render( # type: ignore[override]
|
||||
self,
|
||||
kernel: ROCmTemplateKernel,
|
||||
op: "CKGroupedConvFwdOp", # type: ignore[name-defined]
|
||||
**kwargs,
|
||||
) -> str:
|
||||
template_buffer_node = kwargs.get("template_buffer_node")
|
||||
if template_buffer_node is not None:
|
||||
self.output_node = template_buffer_node
|
||||
X, W = self.input_nodes[0], self.input_nodes[1]
|
||||
Y = self.output_node
|
||||
Bias = self.input_nodes[2] if 3 == len(self.input_nodes) else None
|
||||
|
||||
op = copy.deepcopy(op)
|
||||
|
||||
instance_definition, instance_type = self.emit_ck_instance(op)
|
||||
|
||||
size_arg_strs = [
|
||||
"GroupCount",
|
||||
"NBatch",
|
||||
"NOutChannels",
|
||||
"NInChannels",
|
||||
"FilterSize_0",
|
||||
"FilterSize_1",
|
||||
"InputSize_0",
|
||||
"InputSize_1",
|
||||
"ConvolutionStrides_0",
|
||||
"ConvolutionStrides_1",
|
||||
"Dilations_0",
|
||||
"Dilations_1",
|
||||
"LeftPads_0",
|
||||
"LeftPads_1",
|
||||
"RightPads_0",
|
||||
"RightPads_1",
|
||||
]
|
||||
|
||||
return self._template_from_string(self.conv_template).render(
|
||||
headers=self.header().getvalue(),
|
||||
globals=self.globals().getvalue(),
|
||||
instance_definition=instance_definition,
|
||||
instance_type=instance_type,
|
||||
kernel_definition=kernel.def_kernel(
|
||||
inputs=[X, W, Bias] if Bias is not None else [X, W],
|
||||
outputs=[Y],
|
||||
names_str="input, weight, bias, output"
|
||||
if Bias is not None
|
||||
else "input, weight, output",
|
||||
size_args=[f"int32_t {arg}" for arg in size_arg_strs],
|
||||
),
|
||||
n_d_tensors=1 if Bias is not None else 0,
|
||||
n_dim_spatial=self.n_spatial_dimensions,
|
||||
input_layout=op.a_layout,
|
||||
weight_layout=op.b_layout,
|
||||
output_layout=op.e_layout,
|
||||
)
|
||||
|
||||
def size_args(self):
|
||||
x, w = self.input_nodes[0], self.input_nodes[1]
|
||||
y = self.output_node
|
||||
|
||||
group_count = self.groups
|
||||
n_batch = x.shape[0] # type: ignore[index]
|
||||
n_out_channels = y.shape[1] # type: ignore[index]
|
||||
n_in_channels = x.shape[1] # type: ignore[index]
|
||||
|
||||
filter_size_0, filter_size_1 = w.shape[2:4] # type: ignore[index]
|
||||
input_size_0, input_size_1 = x.shape[2:4] # type: ignore[index]
|
||||
convolution_strides_0, convolution_strides_1 = self.stride
|
||||
dilations_0, dilations_1 = self.dilation
|
||||
left_pads_0, left_pads_1 = self.padding
|
||||
right_pads_0, right_pads_1 = self.padding
|
||||
|
||||
return (
|
||||
group_count,
|
||||
n_batch,
|
||||
n_out_channels,
|
||||
n_in_channels,
|
||||
filter_size_0,
|
||||
filter_size_1,
|
||||
input_size_0,
|
||||
input_size_1,
|
||||
convolution_strides_0,
|
||||
convolution_strides_1,
|
||||
dilations_0,
|
||||
dilations_1,
|
||||
left_pads_0,
|
||||
left_pads_1,
|
||||
right_pads_0,
|
||||
right_pads_1,
|
||||
)
|
||||
|
||||
@override
|
||||
def get_runtime_arg_info(self) -> list[ArgInfo]:
|
||||
return []
|
||||
|
||||
@override
|
||||
def get_runtime_arg_values(self, **kwargs: Any) -> list[Any]:
|
||||
"""
|
||||
Helper method to retrieve runtime args from generate kwargs
|
||||
"""
|
||||
return []
|
||||
@@ -0,0 +1,110 @@
|
||||
from typing import Any
|
||||
from typing_extensions import override
|
||||
|
||||
import torch
|
||||
from torch._inductor.codegen.rocm.rocm_template import ROCmTemplate
|
||||
from torch._inductor.ir import IRNode
|
||||
from torch._inductor.utils import IndentedBuffer
|
||||
|
||||
from .rocm_template import ArgInfo
|
||||
|
||||
|
||||
class CKTemplate(ROCmTemplate):
|
||||
"""
|
||||
Base class for generating CK templates, has common, i.e. non-gemm-specific, code generation logic
|
||||
"""
|
||||
|
||||
_TORCH_DTYPE_TO_CK = {
|
||||
torch.float32: "F32",
|
||||
torch.float64: "F64",
|
||||
torch.float16: "F16",
|
||||
torch.bfloat16: "BF16",
|
||||
torch.int32: "I32",
|
||||
torch.int8: "I8",
|
||||
torch.float8_e4m3fnuz: "F8", # gfx94
|
||||
torch.float8_e4m3fn: "F8", # gfx95
|
||||
torch.float8_e5m2fnuz: "BF8", # gfx94
|
||||
torch.float8_e5m2: "BF8", # gfx95
|
||||
}
|
||||
|
||||
def header(self) -> IndentedBuffer:
|
||||
res = super().header()
|
||||
res.splice(
|
||||
"""
|
||||
// CK headers
|
||||
|
||||
#ifdef DEBUG_LOG
|
||||
#define DEBUG_LOG_TMP DEBUG_LOG
|
||||
#undef DEBUG_LOG
|
||||
#else
|
||||
#define DEBUG_LOG_TMP 0
|
||||
#endif
|
||||
#include "ck/ck.hpp"
|
||||
#undef DEBUG_LOG
|
||||
#define DEBUG_LOG DEBUG_LOG_TMP
|
||||
|
||||
#include "ck/utility/data_type.hpp"
|
||||
#include "ck/library/utility/check_err.hpp"
|
||||
#include "ck/library/utility/device_memory.hpp"
|
||||
#include "ck/library/utility/fill.hpp"
|
||||
#include "ck/library/utility/host_tensor.hpp"
|
||||
#include "ck/library/utility/host_tensor_generator.hpp"
|
||||
#include "ck/library/utility/literals.hpp"
|
||||
"""
|
||||
)
|
||||
return res
|
||||
|
||||
def globals(self) -> IndentedBuffer:
|
||||
res = super().globals()
|
||||
res.splice(
|
||||
"""
|
||||
// CK globals
|
||||
|
||||
template <ck::index_t... Is>
|
||||
using S = ck::Sequence<Is...>;
|
||||
|
||||
template<typename... Ts>
|
||||
using Tuple = ck::Tuple<Ts...>;
|
||||
|
||||
using PassThrough = ck::tensor_operation::element_wise::PassThrough;
|
||||
using Bilinear = ck::tensor_operation::element_wise::Bilinear;
|
||||
using Scale = ck::tensor_operation::element_wise::Scale;
|
||||
using ScaleAdd = ck::tensor_operation::element_wise::ScaleAdd;
|
||||
using MultiplyMultiply = ck::tensor_operation::element_wise::MultiplyMultiply;
|
||||
|
||||
// see "composable_kernel/include/ck/utility/data_type.hpp"
|
||||
using F8 = ck::f8_t;
|
||||
using BF8 = ck::bf8_t;
|
||||
using F16 = ck::half_t;
|
||||
using F32 = float;
|
||||
// using F64 = double;
|
||||
using BF16 = ck::bhalf_t;
|
||||
// using I32 = int32_t;
|
||||
// using I8 = int8_t;
|
||||
// using I4 = ck::int4_t;
|
||||
|
||||
#if DEBUG_LOG
|
||||
static constexpr auto kDEBUG_LOG = 1;
|
||||
#else
|
||||
static constexpr auto kDEBUG_LOG = 0;
|
||||
#endif
|
||||
"""
|
||||
)
|
||||
return res
|
||||
|
||||
def torch_type_to_ck(self, node: IRNode, ptr: str) -> str:
|
||||
if node is None:
|
||||
return ptr
|
||||
else:
|
||||
return f"({self._TORCH_DTYPE_TO_CK.get(node.get_dtype())}*)({ptr})"
|
||||
|
||||
@override
|
||||
def get_runtime_arg_info(self) -> list[ArgInfo]:
|
||||
return [ArgInfo("kBatch", "int32_t")]
|
||||
|
||||
@override
|
||||
def get_runtime_arg_values(self, **kwargs: Any) -> list[Any]:
|
||||
"""
|
||||
Helper method to retrieve runtime args from generate kwargs
|
||||
"""
|
||||
return [kwargs[arg.name] for arg in self.get_runtime_arg_info()]
|
||||
+58
@@ -0,0 +1,58 @@
|
||||
import torch
|
||||
from torch._inductor.codegen.rocm.rocm_template import ROCmTemplate
|
||||
from torch._inductor.ir import IRNode
|
||||
from torch._inductor.utils import IndentedBuffer
|
||||
|
||||
|
||||
class CKTileTemplate(ROCmTemplate):
|
||||
"""
|
||||
Base class for generating CK templates, has common, i.e. non-gemm-specific, code generation logic
|
||||
"""
|
||||
|
||||
_TORCH_DTYPE_TO_CK = {
|
||||
torch.float32: "F32",
|
||||
torch.float64: "F64",
|
||||
torch.float16: "F16",
|
||||
torch.bfloat16: "BF16",
|
||||
torch.int32: "I32",
|
||||
torch.int8: "I8",
|
||||
torch.float8_e4m3fnuz: "F8", # gfx94
|
||||
torch.float8_e4m3fn: "F8", # gfx95
|
||||
torch.float8_e5m2fnuz: "BF8", # gfx94
|
||||
torch.float8_e5m2: "BF8", # gfx95
|
||||
}
|
||||
|
||||
ck_dtype_to_size = {
|
||||
"FP16": 2,
|
||||
"BF16": 2,
|
||||
}
|
||||
|
||||
def header(self) -> IndentedBuffer:
|
||||
res = super().header()
|
||||
res.splice(
|
||||
"""
|
||||
// CK headers
|
||||
#include "ck_tile/core.hpp"
|
||||
|
||||
"""
|
||||
)
|
||||
return res
|
||||
|
||||
def globals(self) -> IndentedBuffer:
|
||||
res = super().globals()
|
||||
res.splice(
|
||||
"""
|
||||
using F8 = ck_tile::fp8_t;
|
||||
using BF8 = ck_tile::bf8_t;
|
||||
using F16 = ck_tile::half_t;
|
||||
using F32 = float;
|
||||
using BF16 = ck_tile::bfloat16_t;
|
||||
"""
|
||||
)
|
||||
return res
|
||||
|
||||
def torch_type_to_ck(self, node: IRNode, ptr: str) -> str:
|
||||
if node is None:
|
||||
return ptr
|
||||
else:
|
||||
return f"({self._TORCH_DTYPE_TO_CK.get(node.get_dtype())}*)({ptr})"
|
||||
+979
@@ -0,0 +1,979 @@
|
||||
# mypy: allow-untyped-defs, disable-error-code="attr-defined, valid-type"
|
||||
import functools
|
||||
import logging
|
||||
import random
|
||||
from dataclasses import asdict, dataclass
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
from torch._inductor import config
|
||||
from torch._inductor.codegen.rocm.ck_tile_template import CKTileTemplate
|
||||
from torch._inductor.codegen.rocm.rocm_kernel import ROCmTemplateKernel
|
||||
from torch._inductor.codegen.rocm.rocm_template import ArgInfo
|
||||
from torch._inductor.ir import Buffer, Layout
|
||||
from torch.utils._ordered_set import OrderedSet
|
||||
|
||||
from ...utils import IndentedBuffer
|
||||
|
||||
|
||||
log = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def is_static_int(number):
|
||||
import sympy
|
||||
|
||||
return isinstance(number, (int, sympy.Integer))
|
||||
|
||||
|
||||
def torch_layout_to_ck_layout(torch_layout):
|
||||
if torch_layout.stride[-1] == 1:
|
||||
return "Row"
|
||||
elif torch_layout.stride[-2] == 1:
|
||||
return "Col"
|
||||
else:
|
||||
return None
|
||||
|
||||
|
||||
@dataclass
|
||||
class CKTileGemmOperation:
|
||||
layout_a: str
|
||||
layout_b: str
|
||||
layout_c: str
|
||||
|
||||
datatype_a: str
|
||||
datatype_b: str
|
||||
datatype_c: str
|
||||
|
||||
tile_m: int
|
||||
tile_n: int
|
||||
tile_k: int
|
||||
|
||||
warp_m: int
|
||||
warp_n: int
|
||||
warp_k: int
|
||||
|
||||
warp_tile_m: int
|
||||
warp_tile_n: int
|
||||
warp_tile_k: int
|
||||
|
||||
m_is_padded: str
|
||||
n_is_padded: str
|
||||
k_is_padded: str
|
||||
|
||||
pipeline: str
|
||||
scheduler: str
|
||||
epilogue: str
|
||||
|
||||
def layout_repr(self):
|
||||
return f"{self.layout_a[0]}{self.layout_b[0]}{self.layout_c[0]}"
|
||||
|
||||
def dtype_repr(self):
|
||||
return f"{self.datatype_a}{self.datatype_b}{self.datatype_c}"
|
||||
|
||||
def tile_sizes(self):
|
||||
return "_".join(
|
||||
[
|
||||
f"{self.tile_m}{self.tile_n}{self.tile_k}",
|
||||
f"{self.warp_m}{self.warp_n}{self.warp_k}",
|
||||
f"{self.warp_tile_m}{self.warp_tile_n}{self.warp_tile_k}",
|
||||
]
|
||||
)
|
||||
|
||||
def name(self):
|
||||
return "ck_tile_gemm_universal_" + "_".join(
|
||||
[
|
||||
f"{self.layout_repr()}",
|
||||
f"{self.dtype_repr()}",
|
||||
f"{self.tile_sizes()}",
|
||||
f"{self.pipeline}",
|
||||
f"{self.scheduler}",
|
||||
f"{self.epilogue}",
|
||||
]
|
||||
)
|
||||
|
||||
def dict_items(self):
|
||||
return asdict(self).items()
|
||||
|
||||
|
||||
@functools.cache
|
||||
def ops():
|
||||
"""
|
||||
Generate the supported instance dataclasses
|
||||
"""
|
||||
import itertools
|
||||
|
||||
compute_v3_instances = [
|
||||
CKTileGemmOperation(
|
||||
layout_a=layout_a,
|
||||
layout_b=layout_b,
|
||||
layout_c=layout_c,
|
||||
datatype_a=datatype_a,
|
||||
datatype_b=datatype_b,
|
||||
datatype_c=datatype_c,
|
||||
tile_m=tile_m,
|
||||
tile_n=tile_n,
|
||||
tile_k=tile_k,
|
||||
warp_m=warp_m,
|
||||
warp_n=warp_n,
|
||||
warp_k=warp_k,
|
||||
warp_tile_m=warp_tile_m,
|
||||
warp_tile_n=warp_tile_n,
|
||||
warp_tile_k=warp_tile_k,
|
||||
m_is_padded=m_is_padded,
|
||||
n_is_padded=n_is_padded,
|
||||
k_is_padded=k_is_padded,
|
||||
pipeline="CompV3",
|
||||
scheduler="Intrawave",
|
||||
epilogue=epilogue,
|
||||
)
|
||||
for (layout_a, layout_b, layout_c) in [
|
||||
("Row", "Row", "Row"),
|
||||
("Row", "Col", "Row"),
|
||||
]
|
||||
for (datatype_a, datatype_b, datatype_c) in [("FP16",) * 3, ("BF16",) * 3]
|
||||
for (tile_m, tile_n, tile_k) in [(256, 256, 32), (256, 256, 64)]
|
||||
for (warp_m, warp_n, warp_k) in [(2, 2, 1)]
|
||||
for (warp_tile_m, warp_tile_n, warp_tile_k) in [(32, 32, 16)]
|
||||
for m_is_padded in ["true", "false"]
|
||||
for n_is_padded in ["true", "false"]
|
||||
for k_is_padded in ["true", "false"]
|
||||
for epilogue in ["Default", "CShuffle"]
|
||||
]
|
||||
|
||||
compute_v4_instances = [
|
||||
CKTileGemmOperation(
|
||||
layout_a=layout_a,
|
||||
layout_b=layout_b,
|
||||
layout_c=layout_c,
|
||||
datatype_a=datatype_a,
|
||||
datatype_b=datatype_b,
|
||||
datatype_c=datatype_c,
|
||||
tile_m=tile_m,
|
||||
tile_n=tile_n,
|
||||
tile_k=tile_k,
|
||||
warp_m=warp_m,
|
||||
warp_n=warp_n,
|
||||
warp_k=warp_k,
|
||||
warp_tile_m=warp_tile_m,
|
||||
warp_tile_n=warp_tile_n,
|
||||
warp_tile_k=warp_tile_k,
|
||||
m_is_padded=m_is_padded,
|
||||
n_is_padded=n_is_padded,
|
||||
k_is_padded=k_is_padded,
|
||||
pipeline="CompV4",
|
||||
scheduler="Intrawave",
|
||||
epilogue=epilogue,
|
||||
)
|
||||
for (layout_a, layout_b, layout_c) in [
|
||||
("Row", "Row", "Row"),
|
||||
("Row", "Col", "Row"),
|
||||
]
|
||||
for (datatype_a, datatype_b, datatype_c) in [("FP16",) * 3, ("BF16",) * 3]
|
||||
for (tile_m, tile_n, tile_k) in [
|
||||
(256, 256, 32)
|
||||
] # half the tile size since it has double buffering
|
||||
for (warp_m, warp_n, warp_k) in [(2, 2, 1)]
|
||||
for (warp_tile_m, warp_tile_n, warp_tile_k) in [(32, 32, 16)]
|
||||
for m_is_padded in ["true", "false"]
|
||||
for n_is_padded in ["true", "false"]
|
||||
for k_is_padded in ["true", "false"]
|
||||
for epilogue in ["Default", "CShuffle"]
|
||||
]
|
||||
|
||||
mem_instances = [
|
||||
CKTileGemmOperation(
|
||||
layout_a=layout_a,
|
||||
layout_b=layout_b,
|
||||
layout_c=layout_c,
|
||||
datatype_a=datatype_a,
|
||||
datatype_b=datatype_b,
|
||||
datatype_c=datatype_c,
|
||||
tile_m=tile_m,
|
||||
tile_n=tile_n,
|
||||
tile_k=tile_k,
|
||||
warp_m=warp_m,
|
||||
warp_n=warp_n,
|
||||
warp_k=warp_k,
|
||||
warp_tile_m=warp_tile_m,
|
||||
warp_tile_n=warp_tile_n,
|
||||
warp_tile_k=warp_tile_k,
|
||||
m_is_padded=m_is_padded,
|
||||
n_is_padded=n_is_padded,
|
||||
k_is_padded=k_is_padded,
|
||||
pipeline="Mem",
|
||||
scheduler=scheduler,
|
||||
epilogue=epilogue,
|
||||
)
|
||||
for (layout_a, layout_b, layout_c) in [
|
||||
("Row", "Row", "Row"),
|
||||
("Row", "Col", "Row"),
|
||||
]
|
||||
for (datatype_a, datatype_b, datatype_c) in [("FP16",) * 3, ("BF16",) * 3]
|
||||
for (tile_m, tile_n, tile_k) in [(256, 256, 32), (256, 256, 64)]
|
||||
for (warp_m, warp_n, warp_k) in [(2, 2, 1)]
|
||||
for (warp_tile_m, warp_tile_n, warp_tile_k) in [(32, 32, 16)]
|
||||
for m_is_padded in ["true", "false"]
|
||||
for n_is_padded in ["true", "false"]
|
||||
for k_is_padded in ["true", "false"]
|
||||
for scheduler in ["Intrawave", "Interwave"]
|
||||
for epilogue in ["Default", "CShuffle"]
|
||||
]
|
||||
|
||||
return list(
|
||||
itertools.chain(compute_v3_instances, compute_v4_instances, mem_instances)
|
||||
)
|
||||
|
||||
|
||||
class CKTileGemmTemplate(CKTileTemplate):
|
||||
"""
|
||||
This class is used for rendering CK-Tile Universal GEMM kernels
|
||||
"""
|
||||
|
||||
gemm_template = r"""{{version_comment}}
|
||||
{{headers}}
|
||||
{{globals}}
|
||||
{{instance_definition}}
|
||||
extern "C" {
|
||||
PT_EXPORT {{kernel_definition}} {
|
||||
|
||||
using {{instance_namespace}}::BaseGemmPipeline;
|
||||
using {{instance_namespace}}::TilePartitioner;
|
||||
|
||||
constexpr auto TileK = {{instance_namespace}}::TileK;
|
||||
constexpr auto kPrefetchStages = BaseGemmPipeline::PrefetchStages;
|
||||
|
||||
const auto BiasTerms = std::array<const void*, 0> ();
|
||||
const auto BiasStrides = std::array<int32_t, 0> ();
|
||||
|
||||
auto kargs = ck_tile::UniversalGemmKernelArgs<> {
|
||||
{X},
|
||||
{W},
|
||||
BiasTerms,
|
||||
Y,
|
||||
M,
|
||||
N,
|
||||
K,
|
||||
{LDA},
|
||||
{LDB},
|
||||
BiasStrides,
|
||||
LDC,
|
||||
kBatch
|
||||
};
|
||||
|
||||
if (workspace_size) {
|
||||
*workspace_size = 0;
|
||||
return 0;
|
||||
}
|
||||
|
||||
// run the kernel
|
||||
const auto dispatch = [&](const auto has_hot_loop_, const auto tail_number_) constexpr {
|
||||
using Kernel = {{instance_namespace}}::Kernel<has_hot_loop_.value, tail_number_.value>;
|
||||
|
||||
if (!Kernel::IsSupportedArgument(kargs)) {
|
||||
// we do our best to statically avoid this case in `filter_op`
|
||||
throw std::runtime_error("invalid argument");
|
||||
}
|
||||
auto stream_config = ck_tile::stream_config{stream};
|
||||
auto grid_size = Kernel::GridSize(M, N, kBatch);
|
||||
constexpr auto block_size = Kernel::BlockSize();
|
||||
constexpr auto lds_bytes = 0;
|
||||
constexpr auto kBlockPerCU = 1;
|
||||
auto gemm = ck_tile::make_kernel<block_size.x, kBlockPerCU>(Kernel{}, grid_size, block_size, lds_bytes, kargs);
|
||||
float elapsed_time = ck_tile::launch_kernel(stream_config, gemm);
|
||||
};
|
||||
|
||||
const ck_tile::index_t k_grain = kBatch * TileK;
|
||||
const ck_tile::index_t K_split = (K + k_grain - 1) / k_grain * TileK;
|
||||
const ck_tile::index_t num_loop = TilePartitioner::GetLoopNum(K_split);
|
||||
const bool has_hot_loop = BaseGemmPipeline::BlockHasHotloop(num_loop);
|
||||
const ck_tile::TailNumber tail_num = BaseGemmPipeline::GetBlockLoopTailNum(num_loop);
|
||||
|
||||
{{rendered_dispatch}}
|
||||
|
||||
return 0;
|
||||
} // kernel definition
|
||||
} // extern C
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
input_nodes: list[Buffer],
|
||||
layout: Layout,
|
||||
) -> None:
|
||||
super().__init__(
|
||||
"ck_tile_gemm_template",
|
||||
input_nodes=input_nodes,
|
||||
layout=layout,
|
||||
)
|
||||
|
||||
def header(self) -> IndentedBuffer:
|
||||
res = super().header()
|
||||
res.splice(
|
||||
"""
|
||||
// CK GEMM header(s)
|
||||
|
||||
#include "ck_tile/ops/gemm.hpp"
|
||||
#include "ck_tile/ops/epilogue.hpp"
|
||||
"""
|
||||
)
|
||||
return res
|
||||
|
||||
def globals(self) -> IndentedBuffer:
|
||||
res = super().globals()
|
||||
res.splice(
|
||||
"""
|
||||
// CK GEMM globals
|
||||
|
||||
using Row = ck_tile::tensor_layout::gemm::RowMajor;
|
||||
using Col = ck_tile::tensor_layout::gemm::ColumnMajor;
|
||||
|
||||
template <ck_tile::index_t PrefetchStages, typename Dispatcher>
|
||||
void dispatch_memory_pipeline_hot_loop(const ck_tile::TailNumber tail_num, Dispatcher dispatch)
|
||||
{
|
||||
if(tail_num == ck_tile::TailNumber::One)
|
||||
{
|
||||
dispatch(ck_tile::bool_constant<true>{},
|
||||
ck_tile::integral_constant<ck_tile::TailNumber, ck_tile::TailNumber::One>{});
|
||||
}
|
||||
else if(tail_num == ck_tile::TailNumber::Full)
|
||||
{
|
||||
dispatch(ck_tile::bool_constant<true>{},
|
||||
ck_tile::integral_constant<ck_tile::TailNumber, ck_tile::TailNumber::Full>{});
|
||||
}
|
||||
|
||||
if constexpr(PrefetchStages > 2)
|
||||
{
|
||||
if(tail_num == ck_tile::TailNumber::Two)
|
||||
{
|
||||
dispatch(ck_tile::bool_constant<true>{},
|
||||
ck_tile::integral_constant<ck_tile::TailNumber, ck_tile::TailNumber::Two>{});
|
||||
}
|
||||
}
|
||||
if constexpr(PrefetchStages > 3)
|
||||
{
|
||||
if(tail_num == ck_tile::TailNumber::Three)
|
||||
{
|
||||
dispatch(ck_tile::bool_constant<true>{},
|
||||
ck_tile::integral_constant<ck_tile::TailNumber, ck_tile::TailNumber::Three>{});
|
||||
}
|
||||
}
|
||||
if constexpr(PrefetchStages > 4)
|
||||
{
|
||||
if(tail_num == ck_tile::TailNumber::Four)
|
||||
{
|
||||
dispatch(ck_tile::bool_constant<true>{},
|
||||
ck_tile::integral_constant<ck_tile::TailNumber, ck_tile::TailNumber::Four>{});
|
||||
}
|
||||
}
|
||||
if constexpr(PrefetchStages > 5)
|
||||
{
|
||||
if(tail_num == ck_tile::TailNumber::Five)
|
||||
{
|
||||
dispatch(ck_tile::bool_constant<true>{},
|
||||
ck_tile::integral_constant<ck_tile::TailNumber, ck_tile::TailNumber::Five>{});
|
||||
}
|
||||
}
|
||||
if constexpr(PrefetchStages > 6)
|
||||
{
|
||||
if(tail_num == ck_tile::TailNumber::Six)
|
||||
{
|
||||
dispatch(ck_tile::bool_constant<true>{},
|
||||
ck_tile::integral_constant<ck_tile::TailNumber, ck_tile::TailNumber::Six>{});
|
||||
}
|
||||
}
|
||||
if constexpr(PrefetchStages > 7)
|
||||
{
|
||||
if(tail_num == ck_tile::TailNumber::Seven)
|
||||
{
|
||||
dispatch(ck_tile::bool_constant<true>{},
|
||||
ck_tile::integral_constant<ck_tile::TailNumber, ck_tile::TailNumber::Seven>{});
|
||||
}
|
||||
}
|
||||
}
|
||||
"""
|
||||
)
|
||||
return res
|
||||
|
||||
def check_dtypes(self, op: "CKTileGemmOperation"):
|
||||
X_dtype, W_dtype, out_dtype = [
|
||||
T.get_layout().dtype for T in [*self.input_nodes, self.output_node]
|
||||
]
|
||||
if op.datatype_a != self._TORCH_DTYPE_TO_CK[X_dtype]:
|
||||
return False
|
||||
if op.datatype_b != self._TORCH_DTYPE_TO_CK[W_dtype]:
|
||||
return False
|
||||
if op.datatype_c != self._TORCH_DTYPE_TO_CK[out_dtype]:
|
||||
return False
|
||||
return True
|
||||
|
||||
def check_layouts(self, op: "CKTileGemmOperation"):
|
||||
X_layout, W_layout, out_layout = [
|
||||
torch_layout_to_ck_layout(T.get_layout())
|
||||
for T in [*self.input_nodes, self.output_node]
|
||||
]
|
||||
if op.layout_a != X_layout:
|
||||
return False
|
||||
if op.layout_b != W_layout:
|
||||
return False
|
||||
if op.layout_c != out_layout:
|
||||
return False
|
||||
return True
|
||||
|
||||
def get_gemm_problem_size(self):
|
||||
X_size, W_size = [T.get_layout().size for T in [*self.input_nodes]]
|
||||
|
||||
M, K = X_size
|
||||
_, N = W_size
|
||||
|
||||
return M, N, K
|
||||
|
||||
def check_block_tiles(self, op: "CKTileGemmOperation"):
|
||||
"""
|
||||
The contiguous dimension of a tensor must be divisible by the block tile size
|
||||
This helper function enforces it for the inputs and the output.
|
||||
"""
|
||||
M, N, K = self.get_gemm_problem_size()
|
||||
|
||||
def check(dim_size, tile_size, is_padded):
|
||||
if (
|
||||
is_static_int(dim_size)
|
||||
and dim_size % tile_size != 0
|
||||
and is_padded == "false"
|
||||
):
|
||||
return False
|
||||
return True
|
||||
|
||||
if op.layout_a == "Row":
|
||||
# handle in kBatch check
|
||||
return True
|
||||
elif op.layout_a == "Col":
|
||||
if not check(M, op.tile_m, op.m_is_padded):
|
||||
return False
|
||||
else:
|
||||
raise AssertionError(f"Invalid layout {op.layout_a=}")
|
||||
|
||||
if op.layout_b == "Row":
|
||||
if not check(N, op.tile_n, op.n_is_padded):
|
||||
return False
|
||||
elif op.layout_b == "Col":
|
||||
# handle in kBatch check
|
||||
return True
|
||||
else:
|
||||
raise AssertionError(f"Invalid {op.layout_b=}")
|
||||
|
||||
if op.layout_c == "Row":
|
||||
if not check(N, op.tile_n, op.n_is_padded):
|
||||
return False
|
||||
elif op.layout_c == "Col":
|
||||
if not check(M, op.tile_m, op.m_is_padded):
|
||||
return False
|
||||
else:
|
||||
raise AssertionError(f"Invalid layout {op.layout_c=}")
|
||||
|
||||
return True
|
||||
|
||||
def check_alignments(self, op: "CKTileGemmOperation"):
|
||||
"""
|
||||
The contiguous dimension of a tensor must be divisible by the vector load size.
|
||||
"""
|
||||
M, N, K = self.get_gemm_problem_size()
|
||||
|
||||
def max_alignment(contiguous_elements_per_tile, elements_per_thread, ck_dtype):
|
||||
for vector_load_bytes in (16, 8, 4, 2, 1):
|
||||
alignment = vector_load_bytes // self.ck_dtype_to_size[ck_dtype]
|
||||
if (
|
||||
alignment > 0
|
||||
and contiguous_elements_per_tile % alignment == 0
|
||||
and elements_per_thread % alignment == 0
|
||||
):
|
||||
return alignment
|
||||
|
||||
threads_per_block = (
|
||||
op.warp_m * op.warp_n * op.warp_k * self.gfx9_threads_per_warp
|
||||
)
|
||||
a_elements_per_thread = op.tile_m * op.tile_k / threads_per_block
|
||||
b_elements_per_thread = op.tile_n * op.tile_k / threads_per_block
|
||||
|
||||
if op.layout_a == "Row":
|
||||
# K is contiguous tensor dimension
|
||||
a_max_vector_size = max_alignment(
|
||||
op.tile_k, a_elements_per_thread, op.datatype_a
|
||||
)
|
||||
if is_static_int(K) and K % a_max_vector_size != 0:
|
||||
return False
|
||||
elif op.layout_a == "Col":
|
||||
# M is contiguous tensor dimension
|
||||
a_max_vector_size = max_alignment(
|
||||
op.tile_m, a_elements_per_thread, op.datatype_a
|
||||
)
|
||||
if is_static_int(M) and M % a_max_vector_size != 0:
|
||||
return False
|
||||
else:
|
||||
raise AssertionError(f"Invalid layout {op.layout_a=}")
|
||||
|
||||
if op.layout_b == "Row":
|
||||
# N is contiguous tensor dimension
|
||||
b_max_vector_size = max_alignment(
|
||||
op.tile_n, b_elements_per_thread, op.datatype_b
|
||||
)
|
||||
if is_static_int(N) and N % b_max_vector_size != 0:
|
||||
return False
|
||||
elif op.layout_b == "Col":
|
||||
# K is contiguous tensor dimension
|
||||
b_max_vector_size = max_alignment(
|
||||
op.tile_k, b_elements_per_thread, op.datatype_b
|
||||
)
|
||||
if is_static_int(K) and K % b_max_vector_size != 0:
|
||||
return False
|
||||
else:
|
||||
raise AssertionError(f"Invalid layout {op.layout_b=}")
|
||||
|
||||
# the `default` epilogue writes C to memory by 1 tensor element
|
||||
# (divisibility check not necessary)
|
||||
# the `cshuffle` epilogue writes C to memory by 16 bytes
|
||||
# (so the contiguous C dimension size must be divisible by the number of tensor elements in 16 bytes)
|
||||
if op.epilogue == "CShuffle":
|
||||
if (
|
||||
op.layout_c == "Row"
|
||||
and is_static_int(N)
|
||||
and N % (16 / self.ck_dtype_to_size[op.datatype_c]) != 0
|
||||
):
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
def check_warp_tiles(self, op: "CKTileGemmOperation"):
|
||||
if op.tile_m % (op.warp_m * op.warp_tile_m) != 0:
|
||||
return False
|
||||
if op.tile_n % (op.warp_n * op.warp_tile_n) != 0:
|
||||
return False
|
||||
if op.tile_k % (op.warp_k * op.warp_tile_k) != 0:
|
||||
return False
|
||||
return True
|
||||
|
||||
def check_block_tile_size(self, op: "CKTileGemmOperation"):
|
||||
# assuming LDS size is 64KB
|
||||
if op.pipeline == "CompV4":
|
||||
max_block_tile_size = 2**15
|
||||
else:
|
||||
max_block_tile_size = 2**16
|
||||
|
||||
block_tile_size = (
|
||||
self.ck_dtype_to_size[op.datatype_a] * op.tile_m * op.tile_k
|
||||
+ self.ck_dtype_to_size[op.datatype_b] * op.tile_n * op.tile_k
|
||||
)
|
||||
if block_tile_size > max_block_tile_size:
|
||||
return False
|
||||
return True
|
||||
|
||||
def filter_op(self, op: "CKTileGemmOperation"):
|
||||
"""
|
||||
Determines whether a given op definition is suitable for the current
|
||||
input / output of the operation that this template implements.
|
||||
|
||||
Filter is based on inputs' dtype, layout and statically inferred size.
|
||||
|
||||
Returns None if the op is not suitable, otherwise returns the op to be used.
|
||||
"""
|
||||
if not self.check_dtypes(op):
|
||||
return None
|
||||
if not self.check_layouts(op):
|
||||
return None
|
||||
if not self.check_block_tiles(op):
|
||||
return None
|
||||
if not self.check_alignments(op):
|
||||
return None
|
||||
|
||||
return op
|
||||
|
||||
def emit_ck_instance(self, op: "CKTileGemmOperation"):
|
||||
"""
|
||||
This method is used to generate code which defines the type alias for the generated kernel class
|
||||
"""
|
||||
template_definition = r"""
|
||||
// Gemm operator {{operation_name}}
|
||||
|
||||
namespace {{operation_name}} {
|
||||
// block tile
|
||||
constexpr int32_t TileM = {{tile_m}};
|
||||
constexpr int32_t TileN = {{tile_n}};
|
||||
constexpr int32_t TileK = {{tile_k}};
|
||||
// warps per block
|
||||
constexpr int32_t WarpM = {{warp_m}};
|
||||
constexpr int32_t WarpN = {{warp_n}};
|
||||
constexpr int32_t WarpK = {{warp_k}};
|
||||
// xdl tile
|
||||
constexpr int32_t WarpTileM = {{warp_tile_m}};
|
||||
constexpr int32_t WarpTileN = {{warp_tile_n}};
|
||||
constexpr int32_t WarpTileK = {{warp_tile_k}};
|
||||
|
||||
constexpr bool kPadM = {{m_is_padded}};
|
||||
constexpr bool kPadN = {{n_is_padded}};
|
||||
constexpr bool kPadK = {{k_is_padded}};
|
||||
|
||||
using ALayout = {{layout_a}};
|
||||
using BLayout = {{layout_b}};
|
||||
using CLayout = {{layout_c}};
|
||||
|
||||
using ADataType = {{datatype_a}};
|
||||
using BDataType = {{datatype_b}};
|
||||
using CDataType = {{datatype_c}};
|
||||
using AccDataType = F32;
|
||||
|
||||
constexpr bool permuteA = false;
|
||||
constexpr bool permuteB = false;
|
||||
constexpr bool DoubleSmemBuffer = {{has_double_smem_buffer}};
|
||||
constexpr bool TransposeC = false;
|
||||
|
||||
constexpr int kBlockPerCu = 1;
|
||||
constexpr ck_tile::index_t TilePartitionerGroupNum = 8;
|
||||
constexpr ck_tile::index_t TilePartitionerM01 = 4;
|
||||
|
||||
using GemmShape =
|
||||
ck_tile::TileGemmShape<ck_tile::sequence<TileM, TileN, TileK>,
|
||||
ck_tile::sequence<WarpM, WarpN, WarpK>,
|
||||
ck_tile::sequence<WarpTileM, WarpTileN, WarpTileK>,
|
||||
permuteA,
|
||||
permuteB>;
|
||||
|
||||
using TilePartitioner =
|
||||
ck_tile::GemmSpatiallyLocalTilePartitioner<GemmShape,
|
||||
TilePartitionerGroupNum,
|
||||
TilePartitionerM01>;
|
||||
|
||||
using Traits =
|
||||
ck_tile::TileGemmTraits<kPadM, kPadN, kPadK, ALayout, BLayout, CLayout>;
|
||||
|
||||
using GemmUniversalTraits =
|
||||
ck_tile::TileGemmUniversalTraits<kPadM, kPadN, kPadK, DoubleSmemBuffer,
|
||||
ALayout, BLayout, CLayout, TransposeC>;
|
||||
|
||||
using GemmPipelineProblem =
|
||||
ck_tile::GemmPipelineProblem<ADataType, BDataType, AccDataType, GemmShape, Traits>;
|
||||
|
||||
{{rendered_scheduler}}
|
||||
|
||||
template<bool has_hot_loop_v, ck_tile::TailNumber tail_number_v>
|
||||
using UniversalGemmProblem =
|
||||
ck_tile::UniversalGemmPipelineProblem<ADataType,
|
||||
BDataType,
|
||||
AccDataType,
|
||||
GemmShape,
|
||||
GemmUniversalTraits,
|
||||
scheduler,
|
||||
has_hot_loop_v,
|
||||
tail_number_v>;
|
||||
|
||||
{{rendered_pipeline}}
|
||||
|
||||
{{rendered_epilogue}}
|
||||
|
||||
template<bool has_hot_loop_v, ck_tile::TailNumber tail_number_v>
|
||||
using Kernel = ck_tile::GemmKernel<TilePartitioner, GemmPipeline<has_hot_loop_v, tail_number_v>, GemmEpilogue>;
|
||||
}
|
||||
|
||||
"""
|
||||
|
||||
def render_epilogue(epilogue_type):
|
||||
if epilogue_type == "Default":
|
||||
return r"""
|
||||
using EpilogueProblem = ck_tile::DefaultGemm2DEpilogueProblem<ADataType,
|
||||
BDataType,
|
||||
AccDataType,
|
||||
CDataType,
|
||||
CLayout,
|
||||
kPadM,
|
||||
kPadN,
|
||||
WarpTileM,
|
||||
WarpTileN,
|
||||
WarpTileK,
|
||||
TransposeC>;
|
||||
using GemmEpilogue = ck_tile::DefaultGemm2DEpilogue<EpilogueProblem>;
|
||||
"""
|
||||
elif epilogue_type == "CShuffle":
|
||||
return r"""
|
||||
constexpr auto kMemoryOperation = ck_tile::memory_operation_enum::set;
|
||||
using DsDataType = ck_tile::tuple<>; // no bias terms for vanilla GEMM
|
||||
using DsLayout = ck_tile::tuple<>;
|
||||
constexpr auto ELayout = CLayout;
|
||||
using CDEElementWise = ck_tile::element_wise::PassThrough; // no-op
|
||||
using EpilogueProblem = ck_tile::CShuffleEpilogueProblem<ADataType,
|
||||
BDataType,
|
||||
DsDataType,
|
||||
AccDataType,
|
||||
CDataType,
|
||||
DsLayout,
|
||||
ELayout,
|
||||
CDEElementWise,
|
||||
GemmPipelineProblem::kBlockSize,
|
||||
TileM,
|
||||
TileN,
|
||||
WarpM,
|
||||
WarpN,
|
||||
WarpTileM,
|
||||
WarpTileN,
|
||||
WarpTileK,
|
||||
TransposeC,
|
||||
kMemoryOperation>;
|
||||
|
||||
using GemmEpilogue = ck_tile::CShuffleEpilogue<EpilogueProblem>;
|
||||
"""
|
||||
else:
|
||||
raise AssertionError("Epilogue must be set")
|
||||
|
||||
def render_pipeline(pipeline_type):
|
||||
return rf"""
|
||||
using BaseGemmPipeline = ck_tile::BaseGemmPipelineAgBgCr{pipeline_type}<GemmPipelineProblem>;
|
||||
|
||||
template<bool has_hot_loop_v, ck_tile::TailNumber tail_number_v>
|
||||
using GemmPipeline = ck_tile::GemmPipelineAgBgCr{pipeline_type}<UniversalGemmProblem<has_hot_loop_v, tail_number_v>>;
|
||||
"""
|
||||
|
||||
def render_scheduler(scheduler_type):
|
||||
return rf"""
|
||||
constexpr auto scheduler = ck_tile::GemmPipelineScheduler::{scheduler_type};
|
||||
"""
|
||||
|
||||
rendered_definition = self._template_from_string(template_definition).render(
|
||||
operation_name=op.name(),
|
||||
**asdict(op),
|
||||
rendered_scheduler=render_scheduler(op.scheduler),
|
||||
rendered_pipeline=render_pipeline(op.pipeline),
|
||||
rendered_epilogue=render_epilogue(op.epilogue),
|
||||
has_double_smem_buffer=("true" if op.pipeline == "CompV4" else "false"),
|
||||
)
|
||||
return rendered_definition
|
||||
|
||||
def render( # type: ignore[override]
|
||||
self, kernel: ROCmTemplateKernel, op: "CKTileGemmOperation", **kwargs
|
||||
) -> str:
|
||||
"""
|
||||
The primary entry point for the code rendering process used in this template.
|
||||
"""
|
||||
epilogue_nodes = kwargs.get("epilogue_nodes")
|
||||
assert epilogue_nodes is None or 0 == len(epilogue_nodes)
|
||||
template_buffer_node = kwargs.get("template_buffer_node")
|
||||
if template_buffer_node is not None:
|
||||
self.output_node = template_buffer_node
|
||||
assert 2 == len(self.input_nodes)
|
||||
X, W = self.input_nodes
|
||||
Y = self.output_node
|
||||
|
||||
instance_definition = self.emit_ck_instance(op)
|
||||
|
||||
version_comment = rf"""/**
|
||||
* Generated code for CK inductor backend
|
||||
* See {type(self).__module__}.{type(self).__qualname__}
|
||||
*
|
||||
* Template instance {op}
|
||||
*
|
||||
* {torch.__version__=}
|
||||
* torch.version.git_version={getattr(torch.version, "git_version", "None")}
|
||||
*/
|
||||
"""
|
||||
|
||||
def render_dispatch(pipeline_type, op_name):
|
||||
switch_tailnum_template = r"""
|
||||
switch (tail_num) {
|
||||
{% for tail_num in valid_tailnums %}
|
||||
case ck_tile::TailNumber::{{tail_num}}:
|
||||
dispatch({{has_hot_loop}},
|
||||
ck_tile::integral_constant<ck_tile::TailNumber, ck_tile::TailNumber::{{tail_num}}>{});
|
||||
break;
|
||||
{% endfor %}
|
||||
default:
|
||||
std::ostringstream err;
|
||||
err << "Unsupported dispatch: "
|
||||
<< "Pipeline: " << "{{pipeline}}"
|
||||
<< "Prefetch stages: " << kPrefetchStages
|
||||
<< "Tail num: " << tail_num;
|
||||
throw std::runtime_error(err.str());
|
||||
} // switch tail_num
|
||||
"""
|
||||
dispatch_template = r"""
|
||||
if (has_hot_loop) {
|
||||
{{rendered_with_hot_loop}}
|
||||
}
|
||||
else { // has_hot_loop == false
|
||||
{{rendered_without_hot_loop}}
|
||||
} // if has_hot_loop
|
||||
"""
|
||||
if pipeline_type == "CompV3":
|
||||
return self._template_from_string(dispatch_template).render(
|
||||
rendered_with_hot_loop=self._template_from_string(
|
||||
switch_tailnum_template
|
||||
).render(
|
||||
has_hot_loop="ck_tile::integral_constant<bool, true>{}",
|
||||
valid_tailnums=("Full", "Odd", "Even"),
|
||||
pipeline=pipeline_type,
|
||||
),
|
||||
rendered_without_hot_loop=self._template_from_string(
|
||||
switch_tailnum_template
|
||||
).render(
|
||||
has_hot_loop="ck_tile::integral_constant<bool, false>{}",
|
||||
valid_tailnums=("Full", "Odd", "Even"),
|
||||
pipeline=pipeline_type,
|
||||
),
|
||||
)
|
||||
elif pipeline_type == "Mem":
|
||||
return self._template_from_string(dispatch_template).render(
|
||||
rendered_with_hot_loop="dispatch_memory_pipeline_hot_loop<kPrefetchStages>(tail_num, dispatch);",
|
||||
rendered_without_hot_loop=self._template_from_string(
|
||||
switch_tailnum_template
|
||||
).render(
|
||||
has_hot_loop="ck_tile::integral_constant<bool, false>{}",
|
||||
valid_tailnums=("Full", "Odd", "Even"),
|
||||
pipeline=pipeline_type,
|
||||
),
|
||||
)
|
||||
elif pipeline_type == "CompV4":
|
||||
return self._template_from_string(dispatch_template).render(
|
||||
rendered_with_hot_loop=self._template_from_string(
|
||||
switch_tailnum_template
|
||||
).render(
|
||||
has_hot_loop="ck_tile::integral_constant<bool, true>{}",
|
||||
valid_tailnums=("Two", "Three"),
|
||||
pipeline=pipeline_type,
|
||||
),
|
||||
rendered_without_hot_loop=self._template_from_string(
|
||||
switch_tailnum_template
|
||||
).render(
|
||||
has_hot_loop="ck_tile::integral_constant<bool, false>{}",
|
||||
valid_tailnums=("Full", "Odd", "Even"),
|
||||
pipeline=pipeline_type,
|
||||
),
|
||||
)
|
||||
else:
|
||||
raise AssertionError(f"Pipeline {pipeline_type} is not supported")
|
||||
|
||||
return self._template_from_string(self.gemm_template).render(
|
||||
headers=self.header().getvalue(),
|
||||
globals=self.globals().getvalue(),
|
||||
instance_definition=instance_definition,
|
||||
kernel_definition=kernel.def_kernel(
|
||||
inputs=[X, W], # type: ignore[list-item]
|
||||
outputs=[Y],
|
||||
names_str="X, W, Y",
|
||||
size_args=[
|
||||
f"int32_t {arg}" for arg in ["M", "N", "K", "LDA", "LDB", "LDC"]
|
||||
],
|
||||
),
|
||||
instance_namespace=op.name(),
|
||||
version_comment=version_comment,
|
||||
rendered_dispatch=render_dispatch(op.pipeline, op.name()),
|
||||
)
|
||||
|
||||
def gen_ops(self):
|
||||
"""
|
||||
Creates a list of `CKTileGemmOperation` instances that match the GEMM operation this template represents.
|
||||
The instances are guaranteed to have the correct layout, dtype and dimension padding for the GEMM input arguments.
|
||||
|
||||
An instance may invalidate the GEMM configuration at runtime.
|
||||
Such instances will be assigned +inf runtime by the autotune process.
|
||||
"""
|
||||
instances = ops()
|
||||
if not instances:
|
||||
raise AssertionError(
|
||||
"No Composable Kernel Universal GEMM instances found. "
|
||||
"Please check if the library is installed."
|
||||
)
|
||||
filtered_instances = list(filter(self.filter_op, instances))
|
||||
# NB: when using a fixed list order, most likely we will pick the subset of instances
|
||||
# which are very similar to each other. Randomizing the choice seems to solve this.
|
||||
random.seed(-11)
|
||||
chosen_instances = (
|
||||
random.sample(
|
||||
filtered_instances,
|
||||
min(len(filtered_instances), config.rocm.ck_tile_max_profiling_configs),
|
||||
)
|
||||
if config.rocm.ck_tile_max_profiling_configs
|
||||
else filtered_instances
|
||||
)
|
||||
log.debug(
|
||||
"generated %d ck instances after sample: %s",
|
||||
len(chosen_instances),
|
||||
chosen_instances,
|
||||
)
|
||||
return chosen_instances
|
||||
|
||||
@staticmethod
|
||||
def add_choices(
|
||||
choices,
|
||||
layout,
|
||||
input_nodes,
|
||||
):
|
||||
"""
|
||||
Add Composable Kernel Universal GEMM instance choices to the auto-tuning list.
|
||||
"""
|
||||
template = CKTileGemmTemplate(
|
||||
input_nodes,
|
||||
layout,
|
||||
)
|
||||
ops = template.gen_ops()
|
||||
for op in ops:
|
||||
for k_batch in template.k_batch_choices(op):
|
||||
template.maybe_append_choice(
|
||||
choices,
|
||||
op=op,
|
||||
kBatch=k_batch,
|
||||
)
|
||||
|
||||
def k_batch_choices(self, op: "CKTileGemmOperation") -> tuple[int, ...]:
|
||||
"""
|
||||
Returns a list of k_batch choices for the template.
|
||||
"""
|
||||
default_choices = (1, 2, 4, 8, 16, 32)
|
||||
|
||||
def check(dim_size, tile_size, is_padded):
|
||||
if (
|
||||
is_static_int(dim_size)
|
||||
and dim_size % tile_size != 0
|
||||
and is_padded == "false"
|
||||
):
|
||||
return False
|
||||
return True
|
||||
|
||||
_, _, K, _, _, _ = self.size_args()
|
||||
if op.layout_a == "Row" or op.layout_b == "Col":
|
||||
choices = tuple(
|
||||
filter(
|
||||
lambda k_batch: check(K, op.tile_k * k_batch, op.k_is_padded),
|
||||
default_choices,
|
||||
)
|
||||
)
|
||||
else:
|
||||
choices = default_choices
|
||||
|
||||
if op.epilogue == "Default":
|
||||
choices = (1,)
|
||||
|
||||
return choices
|
||||
|
||||
def size_args(self):
|
||||
"""
|
||||
Sizes and strides to be used for the kernel call
|
||||
"""
|
||||
X = self.input_nodes[0]
|
||||
W = self.input_nodes[1]
|
||||
Y = self.output_node
|
||||
|
||||
M = X.get_size()[0]
|
||||
K = X.get_size()[1]
|
||||
N = W.get_size()[1]
|
||||
LDA = X.get_stride()[0 if X.get_stride()[1] == 1 else 1]
|
||||
LDB = W.get_stride()[0 if W.get_stride()[1] == 1 else 1]
|
||||
LDC = Y.get_stride()[0 if Y.get_stride()[1] == 1 else 1]
|
||||
|
||||
return M, N, K, LDA, LDB, LDC
|
||||
|
||||
def get_runtime_arg_info(self) -> list[ArgInfo]:
|
||||
return [ArgInfo("kBatch", "int32_t")]
|
||||
|
||||
def get_runtime_arg_values(self, **kwargs: Any) -> list[Any]:
|
||||
# maybe_append_choice kwarg for k_batch must match the name of the argument
|
||||
arg_names = OrderedSet([arg.name for arg in self.get_runtime_arg_info()])
|
||||
if not arg_names.issubset(kwargs):
|
||||
raise ValueError(
|
||||
"Missing runtime arguments: " + ", ".join(arg_names - kwargs.keys())
|
||||
)
|
||||
return [kwargs[k] for k in arg_names]
|
||||
+1018
File diff suppressed because it is too large
Load Diff
+152
@@ -0,0 +1,152 @@
|
||||
# mypy: allow-untyped-defs
|
||||
import logging
|
||||
import os
|
||||
|
||||
from torch._inductor import config
|
||||
from torch._inductor.utils import is_linux, try_import_ck_lib
|
||||
|
||||
|
||||
log = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def _rocm_include_paths(dst_file_ext: str) -> list[str]:
|
||||
from torch.utils import cpp_extension
|
||||
|
||||
rocm_include = (
|
||||
os.path.join(config.rocm.rocm_home, "include")
|
||||
if config.rocm.rocm_home
|
||||
else cpp_extension._join_rocm_home("include")
|
||||
)
|
||||
|
||||
if config.is_fbcode():
|
||||
from libfb.py import parutil
|
||||
|
||||
ck_path = parutil.get_dir_path("composable-kernel-headers")
|
||||
else:
|
||||
if not config.rocm.ck_dir:
|
||||
ck_dir, _, _, _ = try_import_ck_lib()
|
||||
if not ck_dir:
|
||||
log.warning("Unspecified Composable Kernel directory")
|
||||
config.rocm.ck_dir = ck_dir
|
||||
ck_path = config.rocm.ck_dir or cpp_extension._join_rocm_home(
|
||||
"composable_kernel"
|
||||
)
|
||||
|
||||
log.debug("Using ck path %s", ck_path)
|
||||
|
||||
ck_include = os.path.join(ck_path, "include")
|
||||
ck_library_include = os.path.join(ck_path, "library", "include")
|
||||
|
||||
# CK has to take priority over ROCm include paths
|
||||
# Since CK is potentially more up-to-date
|
||||
paths = [
|
||||
os.path.realpath(p) for p in (ck_include, ck_library_include, rocm_include)
|
||||
]
|
||||
if dst_file_ext == "exe":
|
||||
ck_utility_include = os.path.join(ck_path, "library", "src", "utility")
|
||||
paths.append(os.path.realpath(ck_utility_include))
|
||||
return paths
|
||||
|
||||
|
||||
def _rocm_lib_options(dst_file_ext: str) -> list[str]:
|
||||
from torch.utils import cpp_extension
|
||||
|
||||
rocm_lib_dir = (
|
||||
os.path.join(config.rocm.rocm_home, "lib")
|
||||
if config.rocm.rocm_home
|
||||
else cpp_extension._join_rocm_home("lib")
|
||||
)
|
||||
hip_lib_dir = (
|
||||
os.path.join(config.rocm.rocm_home, "hip", "lib")
|
||||
if config.rocm.rocm_home
|
||||
else cpp_extension._join_rocm_home("hip", "lib")
|
||||
)
|
||||
|
||||
opts = [
|
||||
"-include __clang_hip_runtime_wrapper.h",
|
||||
f"-L{os.path.realpath(rocm_lib_dir)}",
|
||||
f"-L{os.path.realpath(hip_lib_dir)}",
|
||||
"-lamdhip64",
|
||||
]
|
||||
if dst_file_ext == "exe":
|
||||
opts += ["-lpthread", "-lstdc++"]
|
||||
return opts
|
||||
|
||||
|
||||
def _rocm_compiler_options() -> list[str]:
|
||||
arch_list = config.rocm.arch or ["native"]
|
||||
gpu_arch_flags = [f"--offload-arch={arch}" for arch in arch_list]
|
||||
opts = [
|
||||
config.rocm.compile_opt_level,
|
||||
"-x",
|
||||
"hip",
|
||||
"-std=c++20",
|
||||
*gpu_arch_flags,
|
||||
"-fno-gpu-rdc",
|
||||
"-fPIC",
|
||||
"-fvisibility=hidden",
|
||||
"-mllvm",
|
||||
"-amdgpu-early-inline-all=true",
|
||||
"-mllvm",
|
||||
"-amdgpu-function-calls=false",
|
||||
"-mllvm",
|
||||
"-enable-post-misched=0",
|
||||
]
|
||||
if config.rocm.is_debug:
|
||||
opts += ["-DDEBUG_LOG=1", "-g"]
|
||||
if config.rocm.save_temps:
|
||||
opts += ["--save-temps=obj"]
|
||||
if config.rocm.print_kernel_resource_usage:
|
||||
opts += ["-Rpass-analysis=kernel-resource-usage"]
|
||||
if config.rocm.flush_denormals:
|
||||
opts += ["-fgpu-flush-denormals-to-zero"]
|
||||
if config.rocm.use_fast_math:
|
||||
opts += ["-ffast-math"]
|
||||
return opts
|
||||
|
||||
|
||||
def rocm_compiler() -> str | None:
|
||||
if is_linux():
|
||||
if config.rocm.rocm_home:
|
||||
return os.path.realpath(
|
||||
os.path.join(config.rocm.rocm_home, "llvm", "bin", "clang")
|
||||
)
|
||||
try:
|
||||
from torch.utils import cpp_extension
|
||||
|
||||
return os.path.realpath(
|
||||
cpp_extension._join_rocm_home("llvm", "bin", "clang")
|
||||
)
|
||||
except OSError:
|
||||
# neither config.rocm.rocm_home nor env variable ROCM_HOME are set
|
||||
return "clang"
|
||||
return None
|
||||
|
||||
|
||||
def rocm_compile_command(
|
||||
src_files: list[str],
|
||||
dst_file: str,
|
||||
dst_file_ext: str,
|
||||
extra_args: list[str] | None = None,
|
||||
) -> str:
|
||||
include_paths = _rocm_include_paths(dst_file_ext)
|
||||
lib_options = _rocm_lib_options(dst_file_ext)
|
||||
compiler_options = _rocm_compiler_options()
|
||||
compiler = rocm_compiler()
|
||||
options = (
|
||||
compiler_options
|
||||
+ (extra_args or [])
|
||||
+ [f"-I{path}" for path in include_paths]
|
||||
+ lib_options
|
||||
)
|
||||
src_file = " ".join(src_files)
|
||||
# supported extensions: .o, .so, .exe
|
||||
if dst_file_ext == "o":
|
||||
options.append("-c")
|
||||
elif dst_file_ext == "so":
|
||||
options.append("-shared")
|
||||
elif dst_file_ext == "exe":
|
||||
options.append("-DGENERATE_CK_STANDALONE_RUNNER")
|
||||
else:
|
||||
raise NotImplementedError(f"Unsupported output file suffix {dst_file_ext}!")
|
||||
return f"{compiler} {' '.join(options)} -o {dst_file} {src_file}"
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user