<commit_before># Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.
# pylint: disable=unused-variable,invalid-name

"""
Integrate auto_scheduler into relay. It implements the following items:
1. Extract search tasks from a relay program
2. Provide auto-scheduling for all TOPI compute functions
"""

import json
import logging
import threading
from copy import deepcopy

import tvm
from tvm import autotvm, transform
from tvm.ir.transform import PassContext
from tvm.runtime import convert_to_object
from tvm.target import Target
from tvm.te.tensor import ComputeOp, PlaceholderOp, Tensor
from tvm.tir import Reduce
from tvm.tir import expr as _expr

from . import _ffi_api
from .compute_dag import ComputeDAG, LayoutRewriteOption
from .dispatcher import DispatchContext
from .search_task import SearchTask
from .utils import get_const_tuple
from .workload_registry import register_workload_tensors

logger = logging.getLogger("auto_scheduler")


def call_all_topi_funcs(mod, params, target, opt_level=3):
    """Call all TOPI compute to extract auto_scheduler tasks in a Relay program"""
    # pylint: disable=import-outside-toplevel
    from tvm import relay
    from tvm.relay.backend import graph_executor_codegen

    # Turn off AutoTVM config not found warnings
    old_autotvm_silent = autotvm.GLOBAL_SCOPE.silent
    autotvm.GLOBAL_SCOPE.silent = True

    with transform.PassContext(
        opt_level=opt_level,
        config={
            "relay.backend.use_auto_scheduler": True,
            "relay.backend.disable_compile_engine_cache": True,
        },
        disabled_pass={"AutoSchedulerLayoutRewrite"},
    ):
        try:
            # TODO(jwfromm) Remove this once AlterOpLayout bug that mutates
            # source module is fixed. Until then, create a clone.
            mod_clone = deepcopy(mod)
            opt_mod, _ = relay.optimize(mod_clone, target, params)
            grc = graph_executor_codegen.GraphExecutorCodegen(None, target)
            grc.codegen(opt_mod["main"])
        except tvm.TVMError:
            print(
                "Get errors with GraphExecutorCodegen for task extraction. "
                "Fallback to VMCompiler."
            )
            mod_clone = deepcopy(mod)
            compiler = relay.vm.VMCompiler()
            if params:
                compiler.set_params(params)
            mod_clone = (
                tvm.IRModule.from_expr(mod_clone)
                if isinstance(mod_clone, relay.Function)
                else mod_clone
            )
            compiler.lower(mod_clone, target)

    autotvm.GLOBAL_SCOPE.silent = old_autotvm_silent


def extract_tasks(
    mod,
    params,
    target,
    target_host=None,
    hardware_params=None,
    include_simple_tasks=False,
    dump_workload_to_dag_log=None,
    opt_level=3,
):
    """Extract tuning tasks from a relay program.

    Parameters
    ----------
    mod: tvm.IRModule or relay.function.Function
        The module or function to tune
    params: dict of str to numpy array
        The associated parameters of the program
    target: Union[tvm.target.Target, str]
        The compilation target
    target_host: Optional[Union[tvm.target.Target, str]]
        The host compilation target
    hardware_params : Optional[HardwareParams]
        Hardware parameters used for the search tasks
    include_simple_tasks: bool
        Whether to extract simple tasks that do not include complicated ops.
    dump_workload_to_dag_log: Optional[str]
        A file to dump an association between the workload keys and the actual DAG
    opt_level : Optional[int]
        The optimization level of the task extractions.

    Returns
    -------
    tasks: List[SearchTask]
        The tasks in this network
    weights: List[int]
        The weight (i.e. the number of appearance) of extracted tasks
    """
    # pylint: disable=import-outside-toplevel

    target, target_host = Target.check_and_update_host_consist(target, target_host)

    # Run the compiler to collect all TOPI calls during compilation.
    env = TracingEnvironment(
        TracingMode.EXTRACT_TASK if include_simple_tasks else TracingMode.EXTRACT_COMPLEX_TASK_ONLY
    )

    dispatch_ctx = DispatchContext.current
    old_verbose = dispatch_ctx.verbose
    dispatch_ctx.verbose = 0
    with env:
        # Wrap build call in a new thread to avoid the conflict
        # between python's multiprocessing and tvm's thread pool
        build_thread = threading.Thread(
            target=call_all_topi_funcs, args=(mod, params, target, opt_level)
        )
        build_thread.start()
        build_thread.join()
    dispatch_ctx.verbose = old_verbose

    # create search tasks
    tasks = []
    weights = []
    for wkl_key, (weight, func_names) in env.wkl_key_to_weight.items():
        tasks.append(
            SearchTask(
                workload_key=wkl_key,
                target=target,
                hardware_params=hardware_params,
                # When auto scheduler is used in end to end network, try to apply layout rewrite
                # to improve the overall performance
                layout_rewrite_option=LayoutRewriteOption.get_target_default(target, True),
                task_inputs=(
                    env.wkl_key_to_input_names[wkl_key]
                    if wkl_key in env.wkl_key_to_input_names
                    else None
                ),
                task_inputs_save_to_file=True,
                desc=",".join(func_names),
            )
        )
        weights.append(weight)

    if dump_workload_to_dag_log is not None:
        with open(dump_workload_to_dag_log, "w") as f:
            json.dump({task.workload_key: str(task.compute_dag) for task in tasks}, f)

    return tasks, weights


class TracingMode:
    """Two modes for tracing"""

    EXTRACT_TASK = 0  # trace all topi calls to extract tasks
    EXTRACT_COMPLEX_TASK_ONLY = 1  # same as EXTRACT_TASK but ignore the task without complex ops
    PREPARE_LAYOUT_REWRITE = 2  # trace topi calls to prepare layout rewrite


class TracingEnvironment:
    """Global environment for tracing all topi function calls"""

    current = None

    def __init__(self, tracing_mode):
        self.tracing_mode = tracing_mode
        self.relay_disable_build_cache = "false"
        self.func_name_to_wkl_key = {}
        self.wkl_key_to_weight = {}
        self.wkl_key_to_input_names = {}
<commit_msg>[Autoscheduler] Reduce task weight coercion overhead (#8995)<commit_after># Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.
# pylint: disable=unused-variable,invalid-name

"""
Integrate auto_scheduler into relay. It implements the following items:
1. Extract search tasks from a relay program
2. Provide auto-scheduling for all TOPI compute functions
"""

import json
import logging
import threading
from copy import deepcopy

import tvm
from tvm import autotvm, transform
from tvm.ir.transform import PassContext
from tvm.runtime import convert_to_object
from tvm.target import Target
from tvm.te.tensor import ComputeOp, PlaceholderOp, Tensor
from tvm.tir import Reduce
from tvm.tir import expr as _expr

from . import _ffi_api
from .compute_dag import ComputeDAG, LayoutRewriteOption
from .dispatcher import DispatchContext
from .search_task import SearchTask
from .utils import get_const_tuple
from .workload_registry import register_workload_tensors

logger = logging.getLogger("auto_scheduler")


def call_all_topi_funcs(mod, params, target, opt_level=3):
    """Call all TOPI compute to extract auto_scheduler tasks in a Relay program"""
    # pylint: disable=import-outside-toplevel
    from tvm import relay
    from tvm.relay.backend import graph_executor_codegen

    # Turn off AutoTVM config not found warnings
    old_autotvm_silent = autotvm.GLOBAL_SCOPE.silent
    autotvm.GLOBAL_SCOPE.silent = True

    with transform.PassContext(
        opt_level=opt_level,
        config={
            "relay.backend.use_auto_scheduler": True,
            "relay.backend.disable_compile_engine_cache": True,
        },
        disabled_pass={"AutoSchedulerLayoutRewrite"},
    ):
        try:
            # TODO(jwfromm) Remove this once AlterOpLayout bug that mutates
            # source module is fixed. Until then, create a clone.
            mod_clone = deepcopy(mod)
            opt_mod, _ = relay.optimize(mod_clone, target, params)
            grc = graph_executor_codegen.GraphExecutorCodegen(None, target)
            grc.codegen(opt_mod["main"])
        except tvm.TVMError:
            print(
                "Get errors with GraphExecutorCodegen for task extraction. "
                "Fallback to VMCompiler."
            )
            mod_clone = deepcopy(mod)
            compiler = relay.vm.VMCompiler()
            if params:
                compiler.set_params(params)
            mod_clone = (
                tvm.IRModule.from_expr(mod_clone)
                if isinstance(mod_clone, relay.Function)
                else mod_clone
            )
            compiler.lower(mod_clone, target)

    autotvm.GLOBAL_SCOPE.silent = old_autotvm_silent


def extract_tasks(
    mod,
    params,
    target,
    target_host=None,
    hardware_params=None,
    include_simple_tasks=False,
    dump_workload_to_dag_log=None,
    opt_level=3,
):
    """Extract tuning tasks from a relay program.

    Parameters
    ----------
    mod: tvm.IRModule or relay.function.Function
        The module or function to tune
    params: dict of str to numpy array
        The associated parameters of the program
    target: Union[tvm.target.Target, str]
        The compilation target
    target_host: Optional[Union[tvm.target.Target, str]]
        The host compilation target
    hardware_params : Optional[HardwareParams]
        Hardware parameters used for the search tasks
    include_simple_tasks: bool
        Whether to extract simple tasks that do not include complicated ops.
    dump_workload_to_dag_log: Optional[str]
        A file to dump an association between the workload keys and the actual DAG
    opt_level : Optional[int]
        The optimization level of the task extractions.

    Returns
    -------
    tasks: List[SearchTask]
        The tasks in this network
    weights: List[int]
        The weight (i.e. the number of appearance) of extracted tasks
    """
    # pylint: disable=import-outside-toplevel

    target, target_host = Target.check_and_update_host_consist(target, target_host)

    # Run the compiler to collect all TOPI calls during compilation.
    env = TracingEnvironment(
        TracingMode.EXTRACT_TASK if include_simple_tasks else TracingMode.EXTRACT_COMPLEX_TASK_ONLY
    )

    dispatch_ctx = DispatchContext.current
    old_verbose = dispatch_ctx.verbose
    dispatch_ctx.verbose = 0
    with env:
        # Wrap build call in a new thread to avoid the conflict
        # between python's multiprocessing and tvm's thread pool
        build_thread = threading.Thread(
            target=call_all_topi_funcs, args=(mod, params, target, opt_level)
        )
        build_thread.start()
        build_thread.join()
    dispatch_ctx.verbose = old_verbose

    # create search tasks
    tasks = []
    weights = []
    for wkl_key, (weight, func_names) in env.wkl_key_to_weight.items():
        tasks.append(
            SearchTask(
                workload_key=wkl_key,
                target=target,
                hardware_params=hardware_params,
                # When auto scheduler is used in end to end network, try to apply layout rewrite
                # to improve the overall performance
                layout_rewrite_option=LayoutRewriteOption.get_target_default(target, True),
                task_inputs=(
                    env.wkl_key_to_input_names[wkl_key]
                    if wkl_key in env.wkl_key_to_input_names
                    else None
                ),
                task_inputs_save_to_file=True,
                desc=",".join(func_names),
            )
        )
        weights.append(int(weight))

    if dump_workload_to_dag_log is not None:
        with open(dump_workload_to_dag_log, "w") as f:
            json.dump({task.workload_key: str(task.compute_dag) for task in tasks}, f)

    return tasks, weights


class TracingMode:
    """Two modes for tracing"""

    EXTRACT_TASK = 0  # trace all topi calls to extract tasks
    EXTRACT_COMPLEX_TASK_ONLY = 1  # same as EXTRACT_TASK but ignore the task without complex ops
    PREPARE_LAYOUT_REWRITE = 2  # trace topi calls to prepare layout rewrite


class TracingEnvironment:
    """Global environment for tracing all topi function calls"""

    current = None

    def __init__(self, tracing_mode):
        self.tracing_mode = tracing_mode
        self.relay_disable_build_cache = "false"
        self.func_name_to_wkl_key = {}
        self.wkl_key_to_weight = {}
        self.wkl_key_to_input_names = {}

{'fraction_non_alphanumeric': 0.06495953111917388, 'fraction_numerical': 0.0019536701088473346, 'mean_word_length': 2.7491498822913942, 'pattern_counts': {'":': 4, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2971430', 'n_tokens_mistral': 4104, 'n_tokens_neox': 3977, 'n_words': 1256}