You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@tvm.apache.org by tq...@apache.org on 2025/02/14 22:09:36 UTC
(tvm) 05/05: remove topi schedule and keep compute
This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch refactor
in repository https://gitbox.apache.org/repos/asf/tvm.git
commit b97c804270749610adeb8911e3668679211ec156
Author: tqchen <ti...@gmail.com>
AuthorDate: Fri Feb 14 16:56:12 2025 -0500
remove topi schedule and keep compute
---
CMakeLists.txt | 1 +
python/tvm/contrib/relay_viz/__init__.py | 113 --
python/tvm/contrib/relay_viz/dot.py | 221 ---
python/tvm/contrib/relay_viz/interface.py | 326 ----
python/tvm/contrib/relay_viz/terminal.py | 239 ---
python/tvm/script/parser/tir/parser.py | 2 +-
.../test_auto_scheduler_compute_dag.py | 157 --
.../test_auto_scheduler_cost_model.py | 84 -
.../test_auto_scheduler_evolutionary_search.py | 122 --
.../auto_scheduler/test_auto_scheduler_feature.py | 305 ----
.../test_auto_scheduler_layout_rewrite.py | 190 ---
.../test_auto_scheduler_loop_state.py | 522 ------
.../auto_scheduler/test_auto_scheduler_measure.py | 427 -----
.../test_auto_scheduler_search_policy.py | 220 ---
.../test_auto_scheduler_search_task.py | 216 ---
.../test_auto_scheduler_sketch_generation.py | 454 -----
.../test_auto_scheduler_task_scheduler.py | 145 --
tests/python/autotvm/test_autotvm_database.py | 124 --
.../autotvm/test_autotvm_dispatch_context.py | 51 -
tests/python/autotvm/test_autotvm_droplet_tuner.py | 57 -
tests/python/autotvm/test_autotvm_feature.py | 136 --
.../python/autotvm/test_autotvm_flop_calculator.py | 173 --
tests/python/autotvm/test_autotvm_ga_tuner.py | 89 -
.../autotvm/test_autotvm_graph_tuner_core.py | 792 ---------
.../autotvm/test_autotvm_graph_tuner_utils.py | 180 --
tests/python/autotvm/test_autotvm_index_tuner.py | 131 --
tests/python/autotvm/test_autotvm_measure.py | 94 --
tests/python/autotvm/test_autotvm_record.py | 163 --
tests/python/autotvm/test_autotvm_space.py | 263 ---
tests/python/autotvm/test_autotvm_xgboost_model.py | 98 --
.../python/codegen/test_target_codegen_aarch64.py | 264 +--
tests/python/codegen/test_target_codegen_blob.py | 43 +-
tests/python/codegen/test_target_codegen_cuda.py | 62 -
.../python/codegen/test_target_codegen_hexagon.py | 1 -
tests/python/codegen/test_target_codegen_llvm.py | 7 +-
tests/python/codegen/test_target_codegen_opencl.py | 194 +--
tests/python/codegen/test_target_codegen_vulkan.py | 55 +-
tests/python/codegen/test_target_codegen_x86.py | 43 -
.../codegen/test_target_texture_codegen_opencl.py | 1778 --------------------
tests/python/driver/tvmc/conftest.py | 302 ----
tests/python/driver/tvmc/test_autoscheduler.py | 107 --
tests/python/driver/tvmc/test_autotuner.py | 309 ----
tests/python/driver/tvmc/test_command_line.py | 332 ----
tests/python/driver/tvmc/test_compiler.py | 527 ------
tests/python/driver/tvmc/test_composite_target.py | 63 -
tests/python/driver/tvmc/test_frontends.py | 459 -----
tests/python/driver/tvmc/test_model.py | 88 -
tests/python/driver/tvmc/test_parse_config_file.py | 154 --
tests/python/driver/tvmc/test_pass_config.py | 163 --
tests/python/driver/tvmc/test_pass_list.py | 36 -
tests/python/driver/tvmc/test_registry_options.py | 81 -
tests/python/driver/tvmc/test_runner.py | 189 ---
tests/python/driver/tvmc/test_shape_parser.py | 102 --
tests/python/driver/tvmc/test_target.py | 131 --
tests/python/driver/tvmc/test_target_options.py | 145 --
tests/python/driver/tvmc/test_tracker.py | 49 -
tests/python/driver/tvmc/test_transform.py | 253 ---
tests/python/driver/tvmc/test_workspace_pools.py | 405 -----
tests/python/integration/__init__.py | 17 -
tests/python/integration/test_arm_aprofile.py | 81 -
tests/python/integration/test_auto_tensorize.py | 462 -----
tests/python/integration/test_dot.py | 62 -
tests/python/integration/test_ewise.py | 356 ----
tests/python/integration/test_ewise_fpga.py | 111 --
tests/python/integration/test_gemm.py | 122 --
tests/python/integration/test_legacy_tuning.py | 380 -----
tests/python/integration/test_lower.py | 407 -----
tests/python/integration/test_reduce.py | 677 --------
tests/python/integration/test_scan.py | 78 -
tests/python/integration/test_tuning.py | 120 --
tests/python/integration/test_winograd_nnpack.py | 186 --
tests/python/ir/test_ir_type.py | 9 -
tests/python/ir/test_roundtrip_runtime_module.py | 76 -
tests/python/runtime/test_runtime_heterogeneous.py | 444 -----
tests/python/runtime/test_runtime_module_load.py | 2 +-
tests/python/target/test_arm_target.py | 22 -
tests/python/topi/common.py | 31 -
tests/python/topi/test_fifo_buffer.py | 208 ---
tests/python/topi/test_topi_argwhere.py | 80 -
tests/python/topi/test_topi_basic.py | 54 -
tests/python/topi/test_topi_batch_matmul.py | 164 --
.../topi/test_topi_batch_matmul_tensorcore.py | 114 --
tests/python/topi/test_topi_batch_norm.py | 136 --
tests/python/topi/test_topi_batch_to_space_nd.py | 70 -
tests/python/topi/test_topi_bitserial_conv2d.py | 158 --
.../python/topi/test_topi_bitserial_conv2d_rasp.py | 121 --
tests/python/topi/test_topi_bitserial_dense.py | 90 -
tests/python/topi/test_topi_bnn.py | 75 -
tests/python/topi/test_topi_broadcast.py | 449 -----
tests/python/topi/test_topi_clip.py | 75 -
tests/python/topi/test_topi_conv1d.py | 199 ---
.../python/topi/test_topi_conv1d_transpose_ncw.py | 121 --
tests/python/topi/test_topi_conv2d_NCHWc.py | 290 ----
tests/python/topi/test_topi_conv2d_hwcn.py | 120 --
.../topi/test_topi_conv2d_hwnc_tensorcore.py | 210 ---
tests/python/topi/test_topi_conv2d_int8.py | 623 -------
tests/python/topi/test_topi_conv2d_nchw.py | 354 ----
tests/python/topi/test_topi_conv2d_nhwc.py | 269 ---
.../python/topi/test_topi_conv2d_nhwc_pack_int8.py | 84 -
.../topi/test_topi_conv2d_nhwc_tensorcore.py | 154 --
.../python/topi/test_topi_conv2d_nhwc_winograd.py | 178 --
.../python/topi/test_topi_conv2d_tensordot_opts.py | 437 -----
.../python/topi/test_topi_conv2d_transpose_nchw.py | 127 --
tests/python/topi/test_topi_conv2d_winograd.py | 248 ---
tests/python/topi/test_topi_conv3d_ncdhw.py | 225 ---
tests/python/topi/test_topi_conv3d_ndhwc.py | 112 --
.../topi/test_topi_conv3d_ndhwc_tensorcore.py | 152 --
.../topi/test_topi_conv3d_transpose_ncdhw.py | 150 --
tests/python/topi/test_topi_conv3d_winograd.py | 180 --
tests/python/topi/test_topi_correlation.py | 97 --
tests/python/topi/test_topi_deformable_conv2d.py | 214 ---
tests/python/topi/test_topi_dense.py | 187 --
tests/python/topi/test_topi_dense_tensorcore.py | 144 --
tests/python/topi/test_topi_depth_to_space.py | 85 -
tests/python/topi/test_topi_depthwise_conv2d.py | 441 -----
.../topi/test_topi_depthwise_conv2d_back_input.py | 158 --
.../topi/test_topi_depthwise_conv2d_back_weight.py | 142 --
tests/python/topi/test_topi_dft.py | 88 -
tests/python/topi/test_topi_dilate.py | 60 -
tests/python/topi/test_topi_einsum.py | 118 --
.../topi/test_topi_group_conv1d_transpose_ncw.py | 110 --
tests/python/topi/test_topi_group_conv2d.py | 658 --------
.../topi/test_topi_group_conv2d_NCHWc_int8.py | 170 --
.../topi/test_topi_group_conv2d_transpose.py | 165 --
.../topi/test_topi_group_conv3d_transpose_ncdhw.py | 109 --
tests/python/topi/test_topi_group_norm.py | 67 -
tests/python/topi/test_topi_image.py | 350 ----
tests/python/topi/test_topi_instance_norm.py | 64 -
tests/python/topi/test_topi_layer_norm.py | 63 -
tests/python/topi/test_topi_loss.py | 75 -
tests/python/topi/test_topi_lrn.py | 72 -
tests/python/topi/test_topi_lstm.py | 161 --
tests/python/topi/test_topi_math.py | 250 ---
tests/python/topi/test_topi_matmul.py | 183 --
tests/python/topi/test_topi_pooling.py | 812 ---------
tests/python/topi/test_topi_prng.py | 240 ---
tests/python/topi/test_topi_reduce.py | 275 ---
tests/python/topi/test_topi_relu.py | 116 --
tests/python/topi/test_topi_reorg.py | 75 -
tests/python/topi/test_topi_rms_norm.py | 66 -
tests/python/topi/test_topi_scan.py | 144 --
tests/python/topi/test_topi_scatter.py | 93 -
tests/python/topi/test_topi_searchsorted.py | 93 -
tests/python/topi/test_topi_softmax.py | 110 --
tests/python/topi/test_topi_sort.py | 162 --
tests/python/topi/test_topi_space_to_batch_nd.py | 72 -
tests/python/topi/test_topi_space_to_depth.py | 82 -
tests/python/topi/test_topi_sparse.py | 619 -------
tests/python/topi/test_topi_tensor.py | 137 --
tests/python/topi/test_topi_transform.py | 1416 ----------------
tests/python/topi/test_topi_unique.py | 114 --
tests/python/topi/test_topi_upsampling.py | 312 ----
tests/python/topi/test_topi_util.py | 38 -
tests/python/topi/test_topi_vision.py | 736 --------
tests/scripts/task_python_hexagon.sh | 3 +
155 files changed, 13 insertions(+), 32040 deletions(-)
diff --git a/CMakeLists.txt b/CMakeLists.txt
index f43eab264d..409d1a716a 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -829,6 +829,7 @@ if(NOT DEFINED ENV{CONDA_BUILD})
WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}/python
)
add_dependencies(tvm_cython tvm)
+ message("Add Cython build into the default build step")
endif()
# Installation rules
diff --git a/python/tvm/contrib/relay_viz/__init__.py b/python/tvm/contrib/relay_viz/__init__.py
deleted file mode 100644
index fb4dac226d..0000000000
--- a/python/tvm/contrib/relay_viz/__init__.py
+++ /dev/null
@@ -1,113 +0,0 @@
-# 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.
-"""Relay IR Visualizer"""
-from typing import Dict
-import tvm
-from tvm import relay
-from .interface import (
- Plotter,
- VizGraph,
- VizParser,
-)
-from .terminal import (
- TermPlotter,
- TermVizParser,
-)
-from .dot import (
- DotPlotter,
- DotVizParser,
-)
-
-
-class RelayVisualizer:
- """Relay IR Visualizer
-
- Parameters
- ----------
- relay_mod: tvm.IRModule
- Relay IR module.
- relay_param: None | Dict[str, tvm.runtime.NDArray]
- Relay parameter dictionary. Default `None`.
- plotter: Plotter
- An instance of class inheriting from Plotter interface.
- Default is an instance of `terminal.TermPlotter`.
- parser: VizParser
- An instance of class inheriting from VizParser interface.
- Default is an instance of `terminal.TermVizParser`.
- """
-
- def __init__(
- self,
- relay_mod: tvm.IRModule,
- relay_param: Dict[str, tvm.runtime.NDArray] = None,
- plotter: Plotter = None,
- parser: VizParser = None,
- ):
- self._plotter = plotter if plotter is not None else TermPlotter()
- self._relay_param = relay_param if relay_param is not None else {}
- self._parser = parser if parser is not None else TermVizParser()
-
- global_vars = relay_mod.get_global_vars()
- graph_names = []
- # If we have main function, put it to the first.
- # Then main function can be shown on the top.
- for gv_node in global_vars:
- if gv_node.name_hint == "main":
- graph_names.insert(0, gv_node.name_hint)
- else:
- graph_names.append(gv_node.name_hint)
-
- node_to_id = {}
- # callback to generate an unique string-ID for nodes.
- # node_count_offset ensure each node ID is still unique across subgraph.
- node_count_offset = 0
-
- def traverse_expr(node):
- if node in node_to_id:
- return
- node_to_id[node] = str(len(node_to_id) + node_count_offset)
-
- for name in graph_names:
- node_count_offset += len(node_to_id)
- node_to_id.clear()
- relay.analysis.post_order_visit(relay_mod[name], traverse_expr)
- graph = self._plotter.create_graph(name)
- self._add_nodes(graph, node_to_id)
-
- def _add_nodes(self, graph: VizGraph, node_to_id: Dict[relay.Expr, str]):
- """add nodes and to the graph.
-
- Parameters
- ----------
- graph : VizGraph
- a VizGraph for nodes to be added to.
-
- node_to_id : Dict[relay.expr, str]
- a mapping from nodes to an unique ID.
-
- relay_param : Dict[str, tvm.runtime.NDarray]
- relay parameter dictionary.
- """
- for node in node_to_id:
- viz_node, viz_edges = self._parser.get_node_edges(node, self._relay_param, node_to_id)
- if viz_node is not None:
- graph.node(viz_node)
- for edge in viz_edges:
- graph.edge(edge)
-
- def render(self, filename: str = None) -> None:
- self._plotter.render(filename=filename)
diff --git a/python/tvm/contrib/relay_viz/dot.py b/python/tvm/contrib/relay_viz/dot.py
deleted file mode 100644
index a9e98189a8..0000000000
--- a/python/tvm/contrib/relay_viz/dot.py
+++ /dev/null
@@ -1,221 +0,0 @@
-# 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.
-"""Visualize Relay IR by Graphviz DOT language."""
-
-from typing import (
- Any,
- Callable,
- Dict,
-)
-from .interface import (
- DefaultVizParser,
- Plotter,
- VizEdge,
- VizGraph,
- VizNode,
-)
-
-try:
- import graphviz
-except ImportError:
- # add "from None" to silence
- # "During handling of the above exception, another exception occurred"
- raise ImportError(
- "The graphviz package is required for DOT renderer. "
- "Please install it first. For example, pip3 install graphviz"
- ) from None
-
-DotVizParser = DefaultVizParser
-
-
-class DotGraph(VizGraph):
- """DOT graph for relay IR.
-
- See also :py:class:`tvm.contrib.relay_viz.dot.DotPlotter`
-
- Parameters
- ----------
- name: str
- name of this graph.
- graph_attr: Optional[Dict[str, str]]
- key-value pairs for the graph.
- node_attr: Optional[Dict[str, str]]
- key-value pairs for all nodes.
- edge_attr: Optional[Dict[str, str]]
- key-value pairs for all edges.
- get_node_attr: Optional[Callable[[VizNode], Dict[str, str]]]
- A callable returning attributes for the node.
- """
-
- def __init__(
- self,
- name: str,
- graph_attr: Dict[str, str] = None,
- node_attr: Dict[str, str] = None,
- edge_attr: Dict[str, str] = None,
- get_node_attr: Callable[[VizNode], Dict[str, str]] = None,
- ):
- self._name = name
- self._get_node_attr = self._default_get_node_attr
- if get_node_attr is not None:
- self._get_node_attr = get_node_attr
-
- # graphviz recognizes the subgraph as a cluster subgraph
- # by the name starting with "cluster" (all lowercase)
- self._digraph = graphviz.Digraph(
- name=f"cluster_{self._name}",
- graph_attr=graph_attr,
- node_attr=node_attr,
- edge_attr=edge_attr,
- )
- self._digraph.attr(label=self._name)
-
- def node(self, viz_node: VizNode) -> None:
- """Add a node to the underlying graph.
- Nodes in a Relay IR Module are expected to be added in the post-order.
-
- Parameters
- ----------
- viz_node : VizNode
- A `VizNode` instance.
- """
- self._digraph.node(
- viz_node.identity,
- f"{viz_node.type_name}\n{viz_node.detail}",
- **self._get_node_attr(viz_node),
- )
-
- def edge(self, viz_edge: VizEdge) -> None:
- """Add an edge to the underlying graph.
-
- Parameters
- ----------
- viz_edge : VizEdge
- A `VizEdge` instance.
- """
- self._digraph.edge(viz_edge.start, viz_edge.end)
-
- @property
- def digraph(self):
- return self._digraph
-
- @staticmethod
- def _default_get_node_attr(node: VizNode):
- if "Var" in node.type_name:
- return {"shape": "ellipse"}
- return {"shape": "box"}
-
-
-class DotPlotter(Plotter):
- """DOT language graph plotter
-
- The plotter accepts various graphviz attributes for graphs, nodes, and edges.
- Please refer to https://graphviz.org/doc/info/attrs.html for available attributes.
-
- Parameters
- ----------
- graph_attr: Optional[Dict[str, str]]
- key-value pairs for all graphs.
- node_attr: Optional[Dict[str, str]]
- key-value pairs for all nodes.
- edge_attr: Optional[Dict[str, str]]
- key-value pairs for all edges.
- get_node_attr: Optional[Callable[[VizNode], Dict[str, str]]]
- A callable returning attributes for a specific node.
- render_kwargs: Optional[Dict[str, Any]]
- keyword arguments directly passed to `graphviz.Digraph.render()`.
-
- Examples
- --------
-
- .. code-block:: python
-
- from tvm.contrib import relay_viz
- from tvm.relay.testing import resnet
-
- mod, param = resnet.get_workload(num_layers=18)
- # graphviz attributes
- graph_attr = {"color": "red"}
- node_attr = {"color": "blue"}
- edge_attr = {"color": "black"}
-
- # VizNode is passed to the callback.
- # We want to color NCHW conv2d nodes. Also give Var a different shape.
- def get_node_attr(node):
- if "nn.conv2d" in node.type_name and "NCHW" in node.detail:
- return {
- "fillcolor": "green",
- "style": "filled",
- "shape": "box",
- }
- if "Var" in node.type_name:
- return {"shape": "ellipse"}
- return {"shape": "box"}
-
- # Create plotter and pass it to viz. Then render the graph.
- dot_plotter = relay_viz.DotPlotter(
- graph_attr=graph_attr,
- node_attr=node_attr,
- edge_attr=edge_attr,
- get_node_attr=get_node_attr)
-
- viz = relay_viz.RelayVisualizer(
- mod,
- relay_param=param,
- plotter=dot_plotter,
- parser=relay_viz.DotVizParser())
- viz.render("hello")
- """
-
- def __init__(
- self,
- graph_attr: Dict[str, str] = None,
- node_attr: Dict[str, str] = None,
- edge_attr: Dict[str, str] = None,
- get_node_attr: Callable[[VizNode], Dict[str, str]] = None,
- render_kwargs: Dict[str, Any] = None,
- ):
- self._name_to_graph = {}
- self._graph_attr = graph_attr
- self._node_attr = node_attr
- self._edge_attr = edge_attr
- self._get_node_attr = get_node_attr
-
- self._render_kwargs = {} if render_kwargs is None else render_kwargs
-
- def create_graph(self, name):
- self._name_to_graph[name] = DotGraph(
- name, self._graph_attr, self._node_attr, self._edge_attr, self._get_node_attr
- )
- return self._name_to_graph[name]
-
- def render(self, filename: str = None):
- """render the graph generated from the Relay IR module.
-
- This function is a thin wrapper of `graphviz.Digraph.render()`.
- """
- # Create or update the filename
- if filename is not None:
- self._render_kwargs["filename"] = filename
- # default cleanup
- if "cleanup" not in self._render_kwargs:
- self._render_kwargs["cleanup"] = True
-
- root_graph = graphviz.Digraph()
- for graph in self._name_to_graph.values():
- root_graph.subgraph(graph.digraph)
- root_graph.render(**self._render_kwargs)
diff --git a/python/tvm/contrib/relay_viz/interface.py b/python/tvm/contrib/relay_viz/interface.py
deleted file mode 100644
index 8df188fcf4..0000000000
--- a/python/tvm/contrib/relay_viz/interface.py
+++ /dev/null
@@ -1,326 +0,0 @@
-# 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.
-"""Abstract class used by :py:class:`tvm.contrib.relay_viz.RelayVisualizer`."""
-import abc
-from typing import (
- Dict,
- Union,
- Tuple,
- List,
-)
-
-import tvm
-from tvm import relay
-
-UNKNOWN_TYPE = "unknown"
-
-
-class VizNode:
- """VizNode carry node information for `VizGraph` interface.
-
- Parameters
- ----------
- node_id: str
- Unique identifier for this node.
- node_type: str
- Type of this node.
- node_detail: str
- Any supplement for this node such as attributes.
- """
-
- def __init__(self, node_id: str, node_type: str, node_detail: str):
- self._id = node_id
- self._type = node_type
- self._detail = node_detail
-
- @property
- def identity(self) -> str:
- return self._id
-
- @property
- def type_name(self) -> str:
- return self._type
-
- @property
- def detail(self) -> str:
- return self._detail
-
- def __repr__(self) -> str:
- detail = self._detail.replace("\n", ", ")
- return f"VizNode(identity: {self._id}, type_name: {self._type}, detail: {detail}"
-
-
-class VizEdge:
- """VizEdge connect two `VizNode`.
-
- Parameters
- ----------
- start_node: str
- The identifier of the node starting the edge.
- end_node: str
- The identifier of the node ending the edge.
- """
-
- def __init__(self, start_node: str, end_node: str):
- self._start_node = start_node
- self._end_node = end_node
-
- @property
- def start(self) -> str:
- return self._start_node
-
- @property
- def end(self) -> str:
- return self._end_node
-
-
-class VizParser(abc.ABC):
- """VizParser parses out a VizNode and VizEdges from a `relay.Expr`."""
-
- @abc.abstractmethod
- def get_node_edges(
- self,
- node: relay.Expr,
- relay_param: Dict[str, tvm.runtime.NDArray],
- node_to_id: Dict[relay.Expr, str],
- ) -> Tuple[Union[VizNode, None], List[VizEdge]]:
- """Get VizNode and VizEdges for a `relay.Expr`.
-
- Parameters
- ----------
- node : relay.Expr
- relay.Expr which will be parsed and generate a node and edges.
-
- relay_param: Dict[str, tvm.runtime.NDArray]
- relay parameters dictionary.
-
- node_to_id : Dict[relay.Expr, str]
- This is a mapping from relay.Expr to a unique id, generated by `RelayVisualizer`.
-
- Returns
- -------
- rv1 : Union[VizNode, None]
- VizNode represent the relay.Expr. If the relay.Expr is not intended to introduce a node
- to the graph, return None.
-
- rv2 : List[VizEdge]
- a list of VizEdges to describe the connectivity of the relay.Expr.
- Can be empty list to indicate no connectivity.
- """
-
-
-class VizGraph(abc.ABC):
- """Abstract class for graph, which is composed of nodes and edges."""
-
- @abc.abstractmethod
- def node(self, viz_node: VizNode) -> None:
- """Add a node to the underlying graph.
- Nodes in a Relay IR Module are expected to be added in the post-order.
-
- Parameters
- ----------
- viz_node : VizNode
- A `VizNode` instance.
- """
-
- @abc.abstractmethod
- def edge(self, viz_edge: VizEdge) -> None:
- """Add an edge to the underlying graph.
-
- Parameters
- ----------
- viz_edge : VizEdge
- A `VizEdge` instance.
- """
-
-
-class DefaultVizParser(VizParser):
- """DefaultVizParser provde a set of logics to parse a various relay types.
- These logics are inspired and heavily based on
- `visualize` function in https://tvm.apache.org/2020/07/14/bert-pytorch-tvm
- """
-
- def get_node_edges(
- self,
- node: relay.Expr,
- relay_param: Dict[str, tvm.runtime.NDArray],
- node_to_id: Dict[relay.Expr, str],
- ) -> Tuple[Union[VizNode, None], List[VizEdge]]:
- if isinstance(node, relay.Function):
- return self._function(node, node_to_id)
- if isinstance(node, relay.expr.Call):
- return self._call(node, node_to_id)
- if isinstance(node, relay.expr.Var):
- return self._var(node, relay_param, node_to_id)
- if isinstance(node, relay.expr.Tuple):
- return self._tuple(node, node_to_id)
- if isinstance(node, relay.expr.TupleGetItem):
- return self._tuple_get_item(node, node_to_id)
- if isinstance(node, relay.expr.Constant):
- return self._constant(node, node_to_id)
- # GlobalVar possibly mean another global relay function,
- # which is expected to in "Graph" level, not in "Node" level.
- if isinstance(node, (relay.expr.GlobalVar, tvm.ir.Op)):
- return None, []
-
- viz_node = VizNode(node_to_id[node], UNKNOWN_TYPE, f"don't know how to parse {type(node)}")
- viz_edges = []
- return viz_node, viz_edges
-
- def _var(
- self,
- node: relay.Expr,
- relay_param: Dict[str, tvm.runtime.NDArray],
- node_to_id: Dict[relay.Expr, str],
- ) -> Tuple[Union[VizNode, None], List[VizEdge]]:
- """Render rule for a relay var node"""
-
- node_id = node_to_id[node]
- name_hint = node.name_hint
- node_detail = f"name_hint: {name_hint}"
- node_type = "Var(Param)" if name_hint in relay_param else "Var(Input)"
-
- if node.type_annotation is not None:
- if hasattr(node.type_annotation, "shape"):
- shape = tuple(map(int, node.type_annotation.shape))
- dtype = node.type_annotation.dtype
- node_detail = f"{node_detail}\nshape: {shape}\ndtype: {dtype}"
- else:
- node_detail = f"{node_detail}\ntype_annotation: {node.type_annotation}"
-
- # only node
- viz_node = VizNode(node_id, node_type, node_detail)
- viz_edges = []
- return viz_node, viz_edges
-
- def _function(
- self,
- node: relay.Expr,
- node_to_id: Dict[relay.Expr, str],
- ) -> Tuple[Union[VizNode, None], List[VizEdge]]:
- """Render rule for a relay function node"""
- func_attrs = node.attrs
- node_details = [f"{k}: {func_attrs.get_str(k)}" for k in func_attrs.keys()]
- # "Composite" might from relay.transform.MergeComposite
- if "Composite" in func_attrs.keys():
- name = func_attrs["Composite"]
- else:
- name = ""
-
- node_id = node_to_id[node]
-
- # Body -> FunctionNode
- viz_node = VizNode(node_id, f"Func {name}", "\n".join(node_details))
- viz_edges = [VizEdge(node_to_id[node.body], node_id)]
- return viz_node, viz_edges
-
- def _call(
- self,
- node: relay.Expr,
- node_to_id: Dict[relay.Expr, str],
- ) -> Tuple[Union[VizNode, None], List[VizEdge]]:
- """Render rule for a relay call node"""
- node_id = node_to_id[node]
- op_name = UNKNOWN_TYPE
- node_detail = []
- if isinstance(node.op, tvm.ir.Op):
- op_name = node.op.name
- if node.attrs:
- node_detail = [f"{k}: {node.attrs.get_str(k)}" for k in node.attrs.keys()]
- elif isinstance(node.op, relay.Function):
- func_attrs = node.op.attrs
- op_name = "Anonymous Func"
- node_detail = [f"{k}: {func_attrs.get_str(k)}" for k in func_attrs.keys()]
- # "Composite" might from relay.transform.MergeComposite
- if "Composite" in func_attrs.keys():
- op_name = func_attrs["Composite"]
- elif isinstance(node.op, relay.GlobalVar):
- op_name = "GlobalVar"
- node_detail = [f"GlobalVar.name_hint: {node.op.name_hint}"]
- else:
- op_name = str(type(node.op)).split(".")[-1].split("'")[0]
-
- # Arguments -> CallNode
- viz_node = VizNode(node_id, f"Call {op_name}", "\n".join(node_detail))
- args = [node_to_id[arg] for arg in node.args]
- viz_edges = [VizEdge(arg, node_id) for arg in args]
- return viz_node, viz_edges
-
- def _tuple(
- self,
- node: relay.Expr,
- node_to_id: Dict[relay.Expr, str],
- ) -> Tuple[Union[VizNode, None], List[VizEdge]]:
- node_id = node_to_id[node]
-
- # Fields -> TupleNode
- viz_node = VizNode(node_id, "Tuple", "")
- viz_edges = [VizEdge(node_to_id[field], node_id) for field in node.fields]
- return viz_node, viz_edges
-
- def _tuple_get_item(
- self,
- node: relay.Expr,
- node_to_id: Dict[relay.Expr, str],
- ) -> Tuple[Union[VizNode, None], List[VizEdge]]:
- node_id = node_to_id[node]
-
- # Tuple -> TupleGetItemNode
- viz_node = VizNode(node_id, "TupleGetItem", f"idx: {node.index}")
- viz_edges = [VizEdge(node_to_id[node.tuple_value], node_id)]
- return viz_node, viz_edges
-
- def _constant(
- self,
- node: relay.Expr,
- node_to_id: Dict[relay.Expr, str],
- ) -> Tuple[Union[VizNode, None], List[VizEdge]]:
- node_id = node_to_id[node]
- node_detail = f"shape: {node.data.shape}, dtype: {node.data.dtype}"
-
- # only node
- viz_node = VizNode(node_id, "Const", node_detail)
- viz_edges = []
- return viz_node, viz_edges
-
-
-class Plotter(abc.ABC):
- """Plotter can render a collection of Graph interfaces to a file."""
-
- @abc.abstractmethod
- def create_graph(self, name: str) -> VizGraph:
- """Create a VizGraph
-
- Parameters
- ----------
- name : str
- the name of the graph
-
- Return
- ------
- rv1: an instance of class inheriting from VizGraph interface.
- """
-
- @abc.abstractmethod
- def render(self, filename: str) -> None:
- """Render the graph as a file.
-
- Parameters
- ----------
- filename : str
- see the definition of implemented class.
- """
diff --git a/python/tvm/contrib/relay_viz/terminal.py b/python/tvm/contrib/relay_viz/terminal.py
deleted file mode 100644
index f137bbf9d4..0000000000
--- a/python/tvm/contrib/relay_viz/terminal.py
+++ /dev/null
@@ -1,239 +0,0 @@
-# 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.
-"""Visualize Relay IR in AST text-form."""
-
-from collections import deque
-from typing import (
- Dict,
- Union,
- Tuple,
- List,
-)
-import tvm
-from tvm import relay
-from .interface import (
- DefaultVizParser,
- Plotter,
- VizEdge,
- VizGraph,
- VizNode,
- VizParser,
-)
-
-
-class TermVizParser(VizParser):
- """`TermVizParser` parse nodes and edges for `TermPlotter`."""
-
- def __init__(self):
- self._default_parser = DefaultVizParser()
-
- def get_node_edges(
- self,
- node: relay.Expr,
- relay_param: Dict[str, tvm.runtime.NDArray],
- node_to_id: Dict[relay.Expr, str],
- ) -> Tuple[Union[VizNode, None], List[VizEdge]]:
- """Parse a node and edges from a relay.Expr."""
- if isinstance(node, relay.Call):
- return self._call(node, node_to_id)
- if isinstance(node, relay.Let):
- return self._let(node, node_to_id)
- if isinstance(node, relay.GlobalVar):
- return self._global_var(node, node_to_id)
- if isinstance(node, relay.If):
- return self._if(node, node_to_id)
- if isinstance(node, tvm.ir.Op):
- return self._op(node, node_to_id)
- if isinstance(node, relay.Function):
- return self._function(node, node_to_id)
-
- # Leverage logics from default parser.
- return self._default_parser.get_node_edges(node, relay_param, node_to_id)
-
- def _call(self, node, node_to_id):
- node_id = node_to_id[node]
- viz_node = VizNode(node_id, "Call", "")
- viz_edges = [VizEdge(node_to_id[node.op], node_id)]
- for arg in node.args:
- arg_id = node_to_id[arg]
- viz_edges.append(VizEdge(arg_id, node_id))
- return viz_node, viz_edges
-
- def _let(self, node, node_to_id):
- node_id = node_to_id[node]
- viz_node = VizNode(node_id, "Let", "(var, val, body)")
- viz_edges = [
- VizEdge(node_to_id[node.var], node_id),
- VizEdge(node_to_id[node.value], node_id),
- VizEdge(node_to_id[node.body], node_id),
- ]
- return viz_node, viz_edges
-
- def _global_var(self, node, node_to_id):
- node_id = node_to_id[node]
- viz_node = VizNode(node_id, "GlobalVar", node.name_hint)
- viz_edges = []
- return viz_node, viz_edges
-
- def _if(self, node, node_to_id):
- node_id = node_to_id[node]
- viz_node = VizNode(node_id, "If", "(cond, true, false)")
- viz_edges = [
- VizEdge(node_to_id[node.cond], node_id),
- VizEdge(node_to_id[node.true_branch], node_id),
- VizEdge(node_to_id[node.false_branch], node_id),
- ]
- return viz_node, viz_edges
-
- def _op(self, node, node_to_id):
- node_id = node_to_id[node]
- op_name = node.name
- viz_node = VizNode(node_id, op_name, "")
- viz_edges = []
- return viz_node, viz_edges
-
- def _function(self, node, node_to_id):
- node_id = node_to_id[node]
- viz_node = VizNode(node_id, "Func", str(node.params))
- viz_edges = [VizEdge(node_to_id[node.body], node_id)]
- return viz_node, viz_edges
-
-
-class TermNode:
- """TermNode is aimed to generate text more suitable for terminal visualization."""
-
- def __init__(self, viz_node: VizNode):
- self.type = viz_node.type_name
- # We don't want too many lines in a terminal.
- self.other_info = viz_node.detail.replace("\n", ", ")
-
-
-class TermGraph(VizGraph):
- """Terminal graph for a relay IR Module
-
- Parameters
- ----------
- name: str
- name of this graph.
- """
-
- def __init__(self, name: str):
- self._name = name
- # A graph in adjacency list form.
- # The key is source node, and the value is a list of destination nodes.
- self._graph = {}
- # a hash table for quick searching.
- self._id_to_term_node = {}
- # node_id in reversed post order
- # That mean, root is the first node.
- self._node_id_rpo = deque()
-
- def node(self, viz_node: VizNode) -> None:
- """Add a node to the underlying graph.
- Nodes in a Relay IR Module are expected to be added in the post-order.
-
- Parameters
- ----------
- viz_node : VizNode
- A `VizNode` instance.
- """
-
- self._node_id_rpo.appendleft(viz_node.identity)
-
- if viz_node.identity not in self._graph:
- # Add the node into the graph.
- self._graph[viz_node.identity] = []
-
- # Create TermNode from VizNode
- node = TermNode(viz_node)
- self._id_to_term_node[viz_node.identity] = node
-
- def edge(self, viz_edge: VizEdge) -> None:
- """Add an edge to the terminal graph.
-
- Parameters
- ----------
- viz_edge : VizEdge
- A `VizEdge` instance.
- """
- # Take CallNode as an example, instead of "arguments point to CallNode",
- # we want "CallNode points to arguments" in ast-dump form.
- #
- # The direction of edge is typically controlled by the implemented VizParser.
- # Reverse start/end here simply because we leverage default parser implementation.
- if viz_edge.end in self._graph:
- self._graph[viz_edge.end].append(viz_edge.start)
- else:
- self._graph[viz_edge.end] = [viz_edge.start]
-
- def render(self) -> str:
- """Draw a terminal graph
-
- Returns
- -------
- rv1: str
- text representing a graph.
- """
- lines = []
- seen_node = set()
-
- def gen_line(indent, n_id):
- if (indent, n_id) in seen_node:
- return
- seen_node.add((indent, n_id))
-
- conn_symbol = ["|--", "`--"]
- last = len(self._graph[n_id]) - 1
- for i, next_n_id in enumerate(self._graph[n_id]):
- node = self._id_to_term_node[next_n_id]
- lines.append(
- f"{indent}{conn_symbol[1 if i==last else 0]}{node.type} {node.other_info}"
- )
- next_indent = indent
- # increase indent for the next level.
- next_indent += " " if (i == last) else "| "
- gen_line(next_indent, next_n_id)
-
- first_node_id = self._node_id_rpo[0]
- first_node = self._id_to_term_node[first_node_id]
- lines.append(f"@{self._name}({first_node.other_info})")
- gen_line("", first_node_id)
-
- return "\n".join(lines)
-
-
-class TermPlotter(Plotter):
- """Terminal plotter"""
-
- def __init__(self):
- self._name_to_graph = {}
-
- def create_graph(self, name):
- self._name_to_graph[name] = TermGraph(name)
- return self._name_to_graph[name]
-
- def render(self, filename):
- """If filename is None, print to stdio. Otherwise, write to the filename."""
- lines = []
- for name in self._name_to_graph:
- text_graph = self._name_to_graph[name].render()
- lines.append(text_graph)
- if filename is None:
- print("\n".join(lines))
- else:
- with open(filename, "w") as out_file:
- out_file.write("\n".join(lines))
diff --git a/python/tvm/script/parser/tir/parser.py b/python/tvm/script/parser/tir/parser.py
index 3107354ac3..9cc3e785fe 100644
--- a/python/tvm/script/parser/tir/parser.py
+++ b/python/tvm/script/parser/tir/parser.py
@@ -454,7 +454,7 @@ def visit_expr_stmt(self: Parser, node: doc.Expr) -> None:
T.evaluate(res)
elif isinstance(res, (int, bool)):
T.evaluate(tvm.tir.const(res))
- elif isinstance(res, (tvm.relay.Call, tvm.relax.Call)) and not res.args:
+ elif isinstance(res, tvm.relax.Call) and not res.args:
# Using GlobalVar.__call__ with no arguments is ambiguous, as
# each IR has a different function Call representation. If
# this occurs, convert to the TIR representation.
diff --git a/tests/python/auto_scheduler/test_auto_scheduler_compute_dag.py b/tests/python/auto_scheduler/test_auto_scheduler_compute_dag.py
deleted file mode 100644
index d3b618d675..0000000000
--- a/tests/python/auto_scheduler/test_auto_scheduler_compute_dag.py
+++ /dev/null
@@ -1,157 +0,0 @@
-# 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.
-
-"""Test ComputeDAG (replay, infer bound)"""
-import json
-import pickle
-
-import tvm
-from tvm import topi
-from tvm import auto_scheduler, te
-
-from tvm.testing.auto_scheduler import (
- get_tiled_matmul,
- invalid_compute_definition,
- matmul_auto_scheduler_test,
- parallel_matmul_auto_scheduler_test,
-)
-
-
-def test_apply_steps():
- dag, s = get_tiled_matmul()
- dag.print_python_code_from_state(s)
- sch, tensors = dag.apply_steps_from_state(s)
- tvm.lower(sch, tensors, simple_mode=True)
-
-
-def test_infer_bound():
- dag, s = get_tiled_matmul()
- s = dag.infer_bound_from_state(s)
-
-
-def test_estimate_flop():
- N = 512
- A, B, C = matmul_auto_scheduler_test(N, N, N)
- dag = auto_scheduler.ComputeDAG([A, B, C])
- assert abs(dag.flop_ct - 2 * N**3) < 0.5
-
- D = topi.nn.relu(C)
- dag = auto_scheduler.ComputeDAG([A, B, D])
- assert abs(dag.flop_ct - (2 * N**3 + N * N)) < 0.5
-
- # should not count the comparison operations in padding
- E = topi.nn.pad(C, [1, 1])
- dag = auto_scheduler.ComputeDAG([A, B, E])
- assert abs(dag.flop_ct - 2 * N**3) < 0.5
-
- F = te.compute((N, N), lambda i, j: E[i, j], name="F", attrs={"FLOP": 1234})
- dag = auto_scheduler.ComputeDAG([A, B, F])
- assert abs(dag.flop_ct - (2 * N**3 + 1234)) < 0.5
-
- A = te.placeholder((N, N), dtype="float32", name="A")
- F = te.compute((N, N), lambda i, j: te.if_then_else(A[i, j] > 0, A[i, j], 0))
- dag = auto_scheduler.ComputeDAG([A, F])
- assert abs(dag.flop_ct - N**2) < 0.5
-
-
-def test_stage_order():
- """Test if the stage order is preserved when recovering a DAG."""
- N = 512
- A, B, C, D, E = parallel_matmul_auto_scheduler_test(N)
- sch = te.create_schedule([D.op, E.op])
- (D_local,) = sch.cache_write([D], "local")
- (E_local,) = sch.cache_write([E], "local")
- sch.cache_read(A, "shared", [D_local])
- sch.cache_read(B, "shared", [D_local])
- sch.cache_read(A, "shared", [E_local])
- sch.cache_read(C, "shared", [E_local])
-
- dag = auto_scheduler.ComputeDAG(sch)
- stage_ops_1 = dag.get_init_state().stage_ops
-
- # 3 placeholder, 4 x.shared, 2 {D,E}.local, 2 {D,E} compute
- assert len(stage_ops_1) == 11
-
- # Cache read stage should follow the source stage
- for idx, op in enumerate(stage_ops_1):
- if op.name == "A":
- assert (
- stage_ops_1[idx + 1].name == "A.d.shared"
- and stage_ops_1[idx + 2].name == "A.shared"
- )
- elif op.name in ["B", "C"]:
- assert stage_ops_1[idx + 1].name == "%s.shared" % op.name
-
- # Apply the same schedule to Ansor state and it should have the same stage order
- dag = auto_scheduler.ComputeDAG([A, B, C, D, E])
- state = dag.get_init_state()
-
- D_local = state.cache_write(D, "local")
- E_local = state.cache_write(E, "local")
- state.cache_read(A, "shared", [D_local])
- state.cache_read(B, "shared", [D_local])
- state.cache_read(A, "shared", [E_local])
- state.cache_read(C, "shared", [E_local])
-
- stage_ops_2 = state.stage_ops
- assert len(stage_ops_1) == len(stage_ops_2)
-
- # Cache read stage should follow the source stage
- for op1, op2 in zip(stage_ops_1, stage_ops_2):
- assert op1.name == op2.name
-
- # Serialize and deserialize the ComputeDAG constructed by a list of tensor ops.
- loaded_dag = pickle.loads(pickle.dumps(dag))
- assert str(loaded_dag.get_init_state()) == str(dag.get_init_state())
- assert len(loaded_dag.get_init_state().stage_ops) == len(dag.get_init_state().stage_ops)
-
- # Serialize and deserialize the search task. Note that we intentionally skip hardware_params
- # to test if the default one is serialized along with other attributes as well.
- task = auto_scheduler.SearchTask(
- compute_dag=dag, workload_key=json.dumps(("test-key",)), target=tvm.target.Target("llvm")
- )
-
- task2 = pickle.loads(pickle.dumps(task))
- assert '["test-key"]' in auto_scheduler.workload_registry.WORKLOAD_FUNC_REGISTRY
- assert str(task.compute_dag.get_init_state()) == str(task2.compute_dag.get_init_state())
- assert len(task.compute_dag.get_init_state().stage_ops) == len(
- task2.compute_dag.get_init_state().stage_ops
- )
- assert task.workload_key == task2.workload_key
- assert str(task.target) == str(task2.target)
- assert task.hardware_params.num_cores == task2.hardware_params.num_cores
- assert task.hardware_params.vector_unit_bytes == task2.hardware_params.vector_unit_bytes
- assert task.hardware_params.cache_line_bytes == task2.hardware_params.cache_line_bytes
-
-
-def test_invalid_compute_dag():
- failed = False
- try:
- A, B = invalid_compute_definition()
- auto_scheduler.ComputeDAG([A, B])
- except tvm.TVMError:
- failed = True
-
- assert failed
-
-
-if __name__ == "__main__":
- test_apply_steps()
- test_infer_bound()
- test_estimate_flop()
- test_stage_order()
- test_invalid_compute_dag()
diff --git a/tests/python/auto_scheduler/test_auto_scheduler_cost_model.py b/tests/python/auto_scheduler/test_auto_scheduler_cost_model.py
deleted file mode 100644
index 50e3ceb6f5..0000000000
--- a/tests/python/auto_scheduler/test_auto_scheduler_cost_model.py
+++ /dev/null
@@ -1,84 +0,0 @@
-# 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.
-
-"""Test cost models"""
-
-import tempfile
-
-import numpy as np
-
-import tvm
-from tvm import auto_scheduler
-
-from tvm.testing.auto_scheduler import matmul_auto_scheduler_test
-
-
-def get_sample_records(number):
- """Generate a list of random MeasureInput and MeasureResult pairs"""
- N = 128
- task = auto_scheduler.SearchTask(func=matmul_auto_scheduler_test, args=(N, N, N), target="llvm")
- policy = auto_scheduler.SketchPolicy(task, verbose=0)
- states = policy.sample_initial_population()[:number]
-
- inputs = [auto_scheduler.MeasureInput(task, s) for s in states]
- results = [
- auto_scheduler.MeasureResult([np.random.uniform(0.5, 1.0)], 0, "", 0.1, 0)
- for _ in range(len(inputs))
- ]
-
- return task, inputs, results
-
-
-def test_random_model():
- task, inputs, results = get_sample_records(50)
-
- model = auto_scheduler.RandomModel()
- model.update(inputs, results)
- scores = model.predict(task, [x.state for x in inputs])
- assert len(scores) == len(inputs)
-
-
-def test_xgb_model():
- task, inputs, results = get_sample_records(50)
-
- model = auto_scheduler.XGBModel(num_warmup_sample=-1)
- model.update(inputs, results)
- preds = model.predict(task, [x.state for x in inputs])
- assert len(preds) == len(inputs)
-
- costs = [np.mean([x.value for x in res.costs]) for res in results]
- throughputs = np.min(costs) / costs
-
- # test regression quality
- rmse = np.sqrt(np.mean([np.square(pred - label) for pred, label in zip(preds, throughputs)]))
- assert rmse <= 0.3
-
- # test loading a record file
- tmpdir = tvm.contrib.utils.tempdir()
- tmpfile = tmpdir.relpath("test1")
- auto_scheduler.save_records(tmpfile, inputs, results)
- model.update_from_file(tmpfile)
-
- # test model serialization
- tmpfile = tmpdir.relpath("test2")
- model.save(tmpfile)
- model.load(tmpfile)
-
-
-if __name__ == "__main__":
- test_random_model()
- test_xgb_model()
diff --git a/tests/python/auto_scheduler/test_auto_scheduler_evolutionary_search.py b/tests/python/auto_scheduler/test_auto_scheduler_evolutionary_search.py
deleted file mode 100644
index 93853b4e7c..0000000000
--- a/tests/python/auto_scheduler/test_auto_scheduler_evolutionary_search.py
+++ /dev/null
@@ -1,122 +0,0 @@
-# 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.
-""" Test evolutionary search. """
-
-import tvm
-import pytest
-from tvm.testing.auto_scheduler import matmul_auto_scheduler_test
-from tvm import auto_scheduler, te
-from tvm.auto_scheduler.cost_model.cost_model import PythonBasedModel
-
-
-def test_mutate_tile_size():
- """
- The test case initializes evo search with a batch of "bad" states and check whether
- the search algorithm can find "good" states by mutating the "bad" states.
-
- This unit test has been tested with 1,000 runs with no failures, meaning that
- the failure rate is less than 0.1%.
- """
-
- class MockCostModel(PythonBasedModel):
- """A mock cost model that rates 1 only for the states with tile_k=2."""
-
- @staticmethod
- def is_good_state(state):
- for line in str(state).split("\n"):
- if line.find("k.1") != -1 and line.find("(0,2)") != -1:
- return True
- return False
-
- def predict(self, task, states):
- scores = []
- for state in states:
- scores.append(1 if self.is_good_state(state) else 0)
- return scores
-
- task = auto_scheduler.SearchTask(
- func=matmul_auto_scheduler_test, args=(10, 10, 4), target=tvm.target.Target("llvm")
- )
- policy = auto_scheduler.SketchPolicy(task, program_cost_model=MockCostModel(), verbose=0)
- states = policy.sample_initial_population()[:50]
-
- bad_states = []
- for state in states:
- if not MockCostModel.is_good_state(state):
- bad_states.append(state)
-
- new_states = policy.evolutionary_search(bad_states, 50)
- found = False
- for state in new_states:
- if MockCostModel.is_good_state(state):
- found = True
- break
- assert found
-
-
-@pytest.mark.skip(reason="See https://github.com/apache/tvm/issues/11440")
-def test_mutate_parallel():
- """
- The test case initializes evo search with a batch of "bad" states and check whether
- the search algorithm can find "good" states by mutating the "bad" states.
- """
-
- class MockCostModel(PythonBasedModel):
- @staticmethod
- def is_good_state(state):
- for line in str(state).split("\n"):
- if (
- line.find("parallel i.0@ (0") != -1
- or line.find("parallel i.0@j.0@ (0") != -1
- or line.find("parallel i.0@j.0@i.1@ (0") != -1
- ):
- return True
- return False
-
- def predict(self, task, states):
- scores = []
- for state in states:
- scores.append(1 if self.is_good_state(state) else 0)
- return scores
-
- task = auto_scheduler.SearchTask(
- func=matmul_auto_scheduler_test, args=(1024, 1024, 1024), target="llvm"
- )
- policy = auto_scheduler.SketchPolicy(task, program_cost_model=MockCostModel(), verbose=0)
-
- found = False
- retry_ct = 0
- while retry_ct < 10 and not found:
- states = policy.sample_initial_population()[:100]
- bad_states = []
- for state in states:
- if not MockCostModel.is_good_state(state):
- bad_states.append(state)
-
- new_states = policy.evolutionary_search(bad_states, 50)
- for state in new_states:
- if MockCostModel.is_good_state(state):
- found = True
- break
- retry_ct += 1
-
- assert found
-
-
-if __name__ == "__main__":
- test_mutate_tile_size()
- test_mutate_parallel()
diff --git a/tests/python/auto_scheduler/test_auto_scheduler_feature.py b/tests/python/auto_scheduler/test_auto_scheduler_feature.py
deleted file mode 100644
index c8edebfd3b..0000000000
--- a/tests/python/auto_scheduler/test_auto_scheduler_feature.py
+++ /dev/null
@@ -1,305 +0,0 @@
-# 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.
-
-"""Test feature extraction"""
-
-import math
-import tempfile
-
-import tvm
-from tvm import te, auto_scheduler, relay
-from tvm.script import tir as T
-
-from tvm.testing.auto_scheduler import matmul_auto_scheduler_test
-
-
-def fequal(a, b):
- return math.fabs(a - b) < 1e-6
-
-
-def test_cpu_matmul():
- dag = auto_scheduler.ComputeDAG(matmul_auto_scheduler_test(512, 512, 512))
- s = dag.get_init_state()
- C = s.stage_ops[2]
-
- i, j, k = s[C].iters
- io, ii = s.split(C, i, [16])
- jo, ji = s.split(C, j, [8])
- s.reorder(C, [io, jo, k, ji, ii])
- s.vectorize(C, ji)
- s.parallel(C, io)
- s.parallel(C, jo)
- s.unroll(C, k)
-
- target = tvm.target.Target("llvm")
- task = auto_scheduler.SearchTask(compute_dag=dag, workload_key="test", target=target)
- names = auto_scheduler.feature.get_per_store_feature_names()
- fea = auto_scheduler.feature.get_per_store_features_from_states([s], task)[0]
-
- stage_0 = fea[0]
- assert len(stage_0) == len(names), "%d vs %d" % (len(stage_0), len(names))
- fea_dict = {}
- for name, value in zip(names, stage_0):
- fea_dict[name] = value
-
- for name in ["B0", "B1", "B2"]:
- if fequal(fea_dict[name + ".acc_type.kReadWrite"], 1.0):
- c_name = name
- if fequal(fea_dict[name + ".acc_type.kRead"], 1.0):
- if fequal(fea_dict[name + ".stride"], 0.0):
- b_name = name
- else:
- a_name = name
-
- """
- lowered IR:
-
- Placeholder: A, B
- parallel i.0 (0,32)
- parallel j.0 (0,64)
- unroll k (0,512)
- vectorize j.1 (0,8)
- for i.1 (0,16)
- C...] = A[...] * B[...]
- """
-
- # check touched memory in bytes, touched unique memory in bytes, reuse distance, etc.
- assert fequal(fea_dict[c_name + ".bytes"], math.log2(512**3 * 4 + 1))
- assert fequal(fea_dict[b_name + ".unique_bytes"], math.log2(512**2 * 4 + 1))
- assert fequal(fea_dict[c_name + ".reuse_dis_iter"], math.log2(8 * 16 + 1))
- assert fequal(fea_dict[c_name + ".reuse_dis_bytes"], math.log2((8 * 16 + 8 + 16) * 4 + 1))
- assert fequal(fea_dict[c_name + ".reuse_ct"], math.log2(512 + 1))
-
- # check annotations
- assert fequal(fea_dict["unroll_num"], math.log2(1 + 1))
- # assert fequal(fea_dict["unroll_type.kPosInnerReduce"], 1.0)
- assert fequal(fea_dict["vec_num"], math.log2(1 + 1))
- assert fequal(fea_dict["parallel_num"], math.log2(2 + 1))
- assert fequal(fea_dict["parallel_prod"], math.log2((512 * 512 / 16 / 8) + 1))
-
-
-def test_cpu_fusion():
- def fusion_test(N, M):
- A = te.placeholder((N, M), name="A")
- B = te.compute((N, M), lambda i, j: A[i][j], name="B")
- C = te.compute((N, M), lambda i, j: B[i][j], name="C")
- return [A, B, C]
-
- dag = auto_scheduler.ComputeDAG(fusion_test(64, 32))
- s = dag.get_init_state()
- s.compute_at(1, 2, s.stages[2].iters[1])
-
- target = tvm.target.Target("llvm")
- task = auto_scheduler.SearchTask(compute_dag=dag, workload_key="test", target=target)
- names = auto_scheduler.feature.get_per_store_feature_names()
- fea = auto_scheduler.feature.get_per_store_features_from_states([s], task)[0]
-
- """
- lowered IR:
-
- Placeholder: A
- for i (0,64)
- for j (0,32)
- for ii (1)
- for jj (1)
- B[...] = A[...]
- C[...] = B[...]
- """
-
- # check reuse distance and reuse type after fusion
- found = False
- for stage_fea in fea:
- for i, (name, value) in enumerate(zip(names, stage_fea)):
- if "reuse_type.kSerialMultipleReadWrite" in name and value > 0.5:
- # reuse distance in #iter
- assert fequal(stage_fea[i + 2], 1.0)
- # reuse distance in bytes
- assert fequal(stage_fea[i + 3], math.log2(16 + 1))
- found = True
- assert found
-
-
-def test_gpu_feature():
- # Use records to build a complicated GPU program
- json_records = "\n".join(
- (
- """{"i": [["[\\"matmul_auto_scheduler_test\\", 512, 512, 512]", "cuda"], [[], [["CHW", 2, "local"], ["SP", 2, 0, 512, [1, 16, 32, 1], 1], ["SP", 2, 5, 512, [4, 1, 1, 16], 1], ["SP", 2, 10, 512, [1, 2], 1], ["RE", 2, [0, 5, 1, 6, 2, 7, 10, 11, 3, 8, 12, 4, 9]], ["FSP", 3, 0, 1, 3], ["FSP", 3, 4, 2, 3], ["RE", 3, [0, 4, 1, 5, 2, 6, 3, 7]], ["FU", 2, [0, 1]], ["FU", 3, [0, 1]], ["FU", 2, [1, 2]], ["FU", 3, [1, 2]], ["FU", 2, [2, 3]], ["FU", 3, [2, 3]], ["CA", 2, 3, 2], ["CHR", 1 [...]
- )
- )
-
- # load states
- with tempfile.NamedTemporaryFile(mode="w") as f:
- f.write(json_records)
- f.flush()
- inputs, _ = auto_scheduler.RecordReader(f.name).read_lines()
-
- inp = inputs[0]
- task = auto_scheduler.SearchTask(
- workload_key=inp.task.workload_key,
- target=inp.task.target,
- hardware_params=auto_scheduler.HardwareParams(
- 100000, 16, 64, 1 << 30, 1 << 30, 1 << 30, 1 << 30, 1 << 30
- ),
- )
-
- state = task.compute_dag.infer_bound_from_state(inputs[0].state)
- fea = auto_scheduler.feature.get_per_store_features_from_states([state], task)[0]
- names = auto_scheduler.feature.get_per_store_feature_names()
-
- # build feature dict
- fea_dicts = []
- for i in range(len(fea)):
- tmp_dict = {}
- for j in range(len(names)):
- tmp_dict[names[j]] = fea[i][j]
- fea_dicts.append(tmp_dict)
-
- """
- lowered IR:
-
- Placeholder: A, B
- blockIdx.x i.0@j.0@ (0,8)
- vthread i.1@j.1@ (0,4)
- threadIdx.x i.2@j.2@ (0,16)
- C.local auto_unroll: 1024
- for k.0 (0,256)
- for ax0@ax1@.0 (0,8)
- threadIdx.x ax0@ax1@.1 (0,16)
- B.shared = ...
- for ax0@ax1@.0 (0,64)
- threadIdx.x ax0@ax1@.1 (0,16)
- A.shared = ...
- for i_c.3 (0,32)
- for k.2 (0,2)
- for j_c.4 (0,16)
- C.local = ...
- for i.3 (0,32)
- for j.3 (0,16)
- C = ...
- """
-
- # check gpu-related features
- assert fequal(fea_dicts[0]["blockIdx_x_len"], math.log2(8 + 1))
- assert fequal(fea_dicts[0]["vthread_len"], math.log2(4 + 1))
- assert fequal(fea_dicts[1]["threadIdx_x_len"], math.log2(16 + 1))
- assert fequal(fea_dicts[0]["threadIdx_y_len"], math.log2(1 + 1))
- assert fequal(fea_dicts[2]["blockIdx_z_len"], math.log2(1 + 1))
- assert fequal(fea_dicts[0]["is_gpu"], 1.0)
-
-
-@T.prim_func
-def tir_matmul(
- A: T.Buffer((256, 256), "float32"),
- B: T.Buffer((256, 256), "float32"),
- C: T.Buffer((256, 256), "float32"),
-) -> None:
- # function attr dict
- T.func_attr({"from_legacy_te_schedule": True, "global_symbol": "main", "tir.noalias": True})
- A_flat = T.Buffer([16384], dtype="float32", data=A.data)
- B_flat = T.Buffer([16384], dtype="float32", data=B.data)
- C_flat = T.Buffer([16384], dtype="float32", data=C.data)
- # body
- for x, y in T.grid(128, 128):
- C_flat[x * 128 + y] = T.float32(0)
- for k in T.serial(128):
- C_flat[x * 128 + y] = C_flat[x * 128 + y] + A_flat[x * 128 + k] * B_flat[y * 128 + k]
-
-
-def test_primfunc_without_lowering():
- features = auto_scheduler.feature.named_features_from_primfunc(tir_matmul)
- assert features["float_mad"].shape == (1,)
- # featurization does not handle multiple-add right now, so they are split out
- assert abs(features["float_addsub"][0] - 128 * 128 * 128) < 10
- assert abs(features["float_mul"][0] - 128 * 128 * 128) < 10
- for i in range(0, 3):
- assert abs(features[f"B{i}.unique_bytes"][0] - 128 * 128 * 4) < 10 # 4 bytes per float32
-
-
-def test_primfunc_lowered():
- # Lower tir function so all passes get applied
- f = tvm.lower(tir_matmul)
- features = auto_scheduler.feature.named_features_from_primfunc(f["main"])
- assert features["float_mad"].shape == (1,)
- # featurization does not handle multiple-add right now, so they are split out
- assert abs(features["float_addsub"][0] - 128 * 128 * 128) < 10
- assert abs(features["float_mul"][0] - 128 * 128 * 128) < 10
- for i in range(0, 3):
- assert abs(features[f"B{i}.unique_bytes"][0] - 128 * 128 * 4) < 10 # 4 bytes per float32
-
-
-def test_dense_lowered():
- a = relay.var("a", relay.TensorType((128, 128), "float32"))
- b = relay.var("b", relay.TensorType((128, 128), "float32"))
- c = relay.nn.dense(a, b)
- mod = tvm.IRModule.from_expr(relay.Function([a, b], c))
- target = "llvm"
- comp = relay.vm.VMCompiler()
- mod, params = comp.optimize(mod, params={}, target=target)
- for name, func in mod.functions.items():
- if name.name_hint != "main":
- break
- features = auto_scheduler.feature.named_features_from_primfunc(func)
- # featurization does not handle multiple-add right now, so they are split out
- assert features["float_addsub"].sum() >= 128 * 128 * 128
- assert features["float_mul"].sum() >= 128 * 128 * 128
- total_bytes_loaded = 0
- for i in range(0, 4):
- total_bytes_loaded += features[f"B{i}.unique_bytes"].sum()
- assert total_bytes_loaded > 2 * 128 * 128 * 4 # 4 bytes per float32
-
-
-@T.prim_func
-def negative_extent(A: T.Buffer((1,), "float32")):
- for j in range(0, -1):
- A[j] = A[j] + 1.0
-
-
-def test_negative_extent():
- features = auto_scheduler.feature.named_features_from_primfunc(negative_extent)
- assert features["B0.unique_bytes"] == 0
-
-
-@T.prim_func
-def zero_dim(
- p2: T.Buffer((), "float32"),
- T_cast: T.Buffer((T.int64(1), T.int64(768)), "int8"),
-):
- # function attr dict
- T.func_attr(
- {
- "tir.noalias": True,
- "Primitive": 1,
- }
- )
- # buffer definition
- T_cast_1 = T.buffer_decl([T.int64(768)], dtype="int8", data=T_cast.data)
- p2_1 = T.buffer_decl([1], dtype="float32", data=p2.data)
- # body
- for i0_i1_fused in T.serial(768):
- T_cast_1[i0_i1_fused] = p2_1[0]
-
-
-def test_zero_dim():
- features = auto_scheduler.feature.named_features_from_primfunc(zero_dim)
- assert features["B1.stride"] == 1
- assert features["B0.stride"] == 1
-
-
-if __name__ == "__main__":
- test_cpu_matmul()
- test_cpu_fusion()
- test_gpu_feature()
diff --git a/tests/python/auto_scheduler/test_auto_scheduler_layout_rewrite.py b/tests/python/auto_scheduler/test_auto_scheduler_layout_rewrite.py
deleted file mode 100644
index 39673fad24..0000000000
--- a/tests/python/auto_scheduler/test_auto_scheduler_layout_rewrite.py
+++ /dev/null
@@ -1,190 +0,0 @@
-# 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.
-
-"""Test AutoScheduler Layout Rewrite"""
-import tempfile
-import numpy as np
-
-import pytest
-
-import tvm
-import tvm.testing
-from tvm import topi
-from tvm import auto_scheduler, te
-
-from tvm.testing.auto_scheduler import get_tiled_matmul, matmul_auto_scheduler_test
-
-
-def test_apply_steps_with_layout_rewrite():
- dag, s = get_tiled_matmul()
- _, bufs = dag.apply_steps_from_state(s)
- assert bufs[1].shape[0] == 512
- assert bufs[1].shape[1] == 512
- _, bufs = dag.apply_steps_from_state(
- s, layout_rewrite=auto_scheduler.LayoutRewriteOption.REWRITE_FOR_PRE_TRANSFORMED
- )
- assert bufs[1].shape[0] == 4
- assert bufs[1].shape[1] == 8
- assert bufs[1].shape[2] == 4
- assert bufs[1].shape[3] == 4
- assert bufs[1].shape[4] == 512
- _, bufs = dag.apply_steps_from_state(
- s, layout_rewrite=auto_scheduler.LayoutRewriteOption.INSERT_TRANSFORM_STAGE
- )
- assert bufs[1].shape[0] == 512
- assert bufs[1].shape[1] == 512
-
-
-def test_apply_steps_with_layout_rewrite_corner_case():
- A, B, C = matmul_auto_scheduler_test(1, 1, 1)
- dag = auto_scheduler.ComputeDAG([A, B, C])
-
- s = dag.get_init_state()
-
- s.compute_root(C)
- i_j_fused = s.fuse(C, [s[C].iters[0], s[C].iters[1]])
- s.parallel(C, i_j_fused)
-
- _, bufs = dag.apply_steps_from_state(
- s, layout_rewrite=auto_scheduler.LayoutRewriteOption.REWRITE_FOR_PRE_TRANSFORMED
- )
-
-
-@tvm.testing.requires_llvm
-def test_correctness_layout_rewrite_rewrite_for_preTransformed():
- N = 16
- target = tvm.target.Target("llvm")
- task = auto_scheduler.SearchTask(func=matmul_auto_scheduler_test, args=(N, N, N), target=target)
- dag = task.compute_dag
-
- with tempfile.NamedTemporaryFile() as fp:
- log_file = fp.name
-
- search_policy = auto_scheduler.SketchPolicy(task)
-
- measure_ctx = auto_scheduler.LocalRPCMeasureContext()
- tuning_options = auto_scheduler.TuningOptions(
- num_measure_trials=100,
- runner=measure_ctx.runner,
- verbose=2,
- early_stopping=1,
- measure_callbacks=[auto_scheduler.RecordToFile(log_file)],
- )
- task.tune(tuning_options, search_policy=search_policy)
- inp, _ = auto_scheduler.load_best_record(log_file, task.workload_key, target)
- s, bufs = dag.apply_steps_from_state(
- inp.state, layout_rewrite=auto_scheduler.LayoutRewriteOption.REWRITE_FOR_PRE_TRANSFORMED
- )
- s_ref, bufs_ref = dag.apply_steps_from_state(inp.state)
- np_args = [np.random.randn(*topi.get_const_tuple(x.shape)).astype(x.dtype) for x in bufs]
- np_args_ref = [np.array(x) for x in np_args]
-
- weight = np_args_ref[1]
- # infer shape for the rewritten layout
- if len(weight.shape) >= 6:
- # For cpu tile structure SSRSRS
- base = len(weight.shape) - 6
- red_dim = weight.shape[2 + base] * weight.shape[4 + base]
- out_dim = weight.shape[3 + base] * weight.shape[5 + base]
- for i in range(base + 2):
- out_dim *= weight.shape[i]
- new_order = (
- [
- 2 + base,
- 4 + base,
- ]
- + list(range(base + 2))
- + [
- 3 + base,
- 5 + base,
- ]
- )
- np_args_ref[1] = np_args_ref[1].transpose(new_order)
- np_args_ref[1] = np_args_ref[1].reshape((red_dim, out_dim))
-
- func = tvm.build(s, bufs, target=target)
- func_ref = tvm.build(s_ref, bufs_ref, target=target)
-
- dev = tvm.device(str(target))
- dev_ref = tvm.cpu()
-
- args = [tvm.nd.array(x, device=dev) for x in np_args]
- args_ref = [tvm.nd.array(x, device=dev_ref) for x in np_args_ref]
- dev.sync()
-
- func(*args)
- func_ref(*args_ref)
- dev.sync()
-
- tvm.testing.assert_allclose(args[0].numpy(), args_ref[0].numpy(), atol=1e-3, rtol=1e-3)
- tvm.testing.assert_allclose(args[2].numpy(), args_ref[2].numpy(), atol=1e-3, rtol=1e-3)
- del measure_ctx
-
-
-@tvm.testing.requires_llvm
-def test_correctness_layout_rewrite_insert_transform_stage():
- N = 128
- target = tvm.target.Target("llvm")
- task = auto_scheduler.SearchTask(func=matmul_auto_scheduler_test, args=(N, N, N), target=target)
- dag = task.compute_dag
-
- with tempfile.NamedTemporaryFile() as fp:
- log_file = fp.name
-
- search_policy = auto_scheduler.SketchPolicy(task)
-
- measure_ctx = auto_scheduler.LocalRPCMeasureContext()
- tuning_options = auto_scheduler.TuningOptions(
- num_measure_trials=2,
- runner=measure_ctx.runner,
- verbose=1,
- measure_callbacks=[auto_scheduler.RecordToFile(log_file)],
- )
- task.tune(tuning_options, search_policy=search_policy)
- inp, _ = auto_scheduler.load_best_record(log_file, task.workload_key, target)
- s, bufs = dag.apply_steps_from_state(
- inp.state, layout_rewrite=auto_scheduler.LayoutRewriteOption.INSERT_TRANSFORM_STAGE
- )
-
- s_ref, bufs_ref = dag.apply_steps_from_state(inp.state)
- np_args = [np.random.randn(*topi.get_const_tuple(x.shape)).astype(x.dtype) for x in bufs]
-
- func = tvm.build(s, bufs, target=target)
- func_ref = tvm.build(s_ref, bufs_ref, target=target)
-
- dev = tvm.device(str(target))
- dev_ref = tvm.cpu()
-
- args = [tvm.nd.array(x, device=dev) for x in np_args]
- args_ref = [tvm.nd.array(x, device=dev_ref) for x in np_args]
- dev.sync()
-
- func(*args)
- func_ref(*args_ref)
- dev.sync()
-
- tvm.testing.assert_allclose(args[0].numpy(), args_ref[0].numpy(), atol=1e-3, rtol=1e-3)
- tvm.testing.assert_allclose(args[1].numpy(), args_ref[1].numpy(), atol=1e-3, rtol=1e-3)
- tvm.testing.assert_allclose(args[2].numpy(), args_ref[2].numpy(), atol=1e-3, rtol=1e-3)
- del measure_ctx
-
-
-if __name__ == "__main__":
- test_apply_steps_with_layout_rewrite()
- test_apply_steps_with_layout_rewrite_corner_case()
- test_correctness_layout_rewrite_rewrite_for_preTransformed()
- test_correctness_layout_rewrite_insert_transform_stage()
diff --git a/tests/python/auto_scheduler/test_auto_scheduler_loop_state.py b/tests/python/auto_scheduler/test_auto_scheduler_loop_state.py
deleted file mode 100644
index 0965ed9efb..0000000000
--- a/tests/python/auto_scheduler/test_auto_scheduler_loop_state.py
+++ /dev/null
@@ -1,522 +0,0 @@
-# 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.
-
-"""Test loop state and schedule primitives"""
-
-import numpy as np
-
-import tvm
-from tvm import auto_scheduler, te
-from tvm import topi
-
-from tvm.testing.auto_scheduler import (
- matmul_auto_scheduler_test,
- conv2d_nchw_bn_relu_auto_scheduler_test,
-)
-
-
-def test_split_fuse_reorder_annotation():
- A, B, C = matmul_auto_scheduler_test(N=512, M=512, K=512)
- dag = auto_scheduler.ComputeDAG([A, B, C])
- s0 = dag.get_init_state()
- i, j, k = s0[C].iters
-
- assert i.range.extent == 512
-
- io, ii = s0.split(C, i, [16])
- assert s0[C].iters[0] == io
- assert s0[C].iters[1] == ii
- assert io.range.extent == 32
- assert ii.range.extent == 16
-
- jo, ji = s0.split(C, j, [8])
- assert jo.range.extent == 64
- assert ji.range.extent == 8
-
- s0.reorder(C, [io, jo, k, ji, ii])
- assert s0[C].iters[2].range.extent == 512
-
- fused_it = s0.fuse(C, [io, jo])
- assert fused_it.range.extent == 2048
-
- s1 = dag.get_init_state()
- i, j, _ = s1[C].iters
- i1, i2, i3 = s1.split(C, i, [8, 2])
- j1, j2, j3 = s1.split(C, j, [32, 8], False)
- assert s1[C].iters[0].range.extent == 32
- assert s1[C].iters[1].range.extent == 8
- assert s1[C].iters[2].range.extent == 2
- assert s1[C].iters[3].range.extent == 32
- assert s1[C].iters[4].range.extent == 8
- assert s1[C].iters[5].range.extent == 2
-
- res = s1.bind(C, i1, "blockIdx.x")
- assert res == s1[C].iters[0]
- assert res.annotation == auto_scheduler.loop_state.State.ANNOTATION_TRANS_TABLE["blockIdx.x"]
-
- res = s1.bind(C, i2, "vthread")
- assert res == s1[C].iters[1]
- assert res.annotation == auto_scheduler.loop_state.State.ANNOTATION_TRANS_TABLE["vthread"]
-
- res = s1.bind(C, i3, "threadIdx.y")
- assert res == s1[C].iters[2]
- assert res.annotation == auto_scheduler.loop_state.State.ANNOTATION_TRANS_TABLE["threadIdx.y"]
-
- res = s1.parallel(C, j1)
- assert res == s1[C].iters[3]
- assert res.annotation == auto_scheduler.loop_state.State.ANNOTATION_TRANS_TABLE["parallel"]
-
- res = s1.unroll(C, j2)
- assert res == s1[C].iters[4]
- assert res.annotation == auto_scheduler.loop_state.State.ANNOTATION_TRANS_TABLE["unroll"]
-
- res = s1.vectorize(C, j3)
- assert res == s1[C].iters[5]
- assert res.annotation == auto_scheduler.loop_state.State.ANNOTATION_TRANS_TABLE["vectorize"]
-
-
-def test_compute_at_root_inline():
- dag = auto_scheduler.ComputeDAG(
- conv2d_nchw_bn_relu_auto_scheduler_test(
- N=1, H=224, W=224, CI=3, CO=64, kernel_size=7, strides=2, padding=3
- )
- )
- s0 = dag.get_init_state()
-
- # data, padding, kernel = 0, 1, 2
- conv = s0.stage_ops[3]
- # bias = 4
- bias_add = s0.stage_ops[5]
- # bn_scale = 6
- bn_mul = s0.stage_ops[7]
- # bn_offset = 8
- bn_add = s0.stage_ops[9]
- relu = s0.stage_ops[10]
-
- s0.compute_inline(bn_add)
- assert s0[bn_add].compute_at == 1
-
- s0.compute_inline(bn_mul)
- assert s0[bn_mul].compute_at == 1
-
- s0.compute_inline(bias_add)
- assert s0[bias_add].compute_at == 1
-
- assert s0[conv].iters[0].range.extent == 1
- assert s0[conv].iters[1].range.extent == 64
- assert s0[conv].iters[2].range.extent == 112
- assert s0[conv].iters[3].range.extent == 112
- assert s0[conv].iters[4].range.extent == 3
- assert s0[conv].iters[5].range.extent == 7
- assert s0[conv].iters[6].range.extent == 7
- s0.compute_at(conv, relu, s0[relu].iters[2])
- assert s0[conv].compute_at == 2
- s0 = dag.infer_bound_from_state(s0)
- assert s0[conv].iters[0].range.extent == 1
- assert s0[conv].iters[1].range.extent == 1
- assert s0[conv].iters[2].range.extent == 1
- assert s0[conv].iters[3].range.extent == 112
- assert s0[conv].iters[4].range.extent == 3
- assert s0[conv].iters[5].range.extent == 7
- assert s0[conv].iters[6].range.extent == 7
-
- s0.compute_root(bn_mul)
- assert s0[bn_mul].compute_at == 0
-
- s0.compute_root(conv)
- assert s0[conv].compute_at == 0
- s0 = dag.infer_bound_from_state(s0)
- assert s0[conv].iters[0].range.extent == 1
- assert s0[conv].iters[1].range.extent == 64
- assert s0[conv].iters[2].range.extent == 112
- assert s0[conv].iters[3].range.extent == 112
- assert s0[conv].iters[4].range.extent == 3
- assert s0[conv].iters[5].range.extent == 7
- assert s0[conv].iters[6].range.extent == 7
-
-
-def test_cache_read_write():
- N, H, W, CO, CI, KH, KW, strides, padding = 4, 7, 7, 512, 512, 3, 3, (1, 1), (1, 1)
-
- data = te.placeholder((N, CI, H, W), name="Data")
- kernel_data = te.placeholder((CO, CI, KH, KW), name="Kernel_data")
- k0, k1 = te.compute(
- kernel_data.shape,
- lambda *i: (kernel_data(*i) + 1, kernel_data(*i) / 2),
- name="Kernel_split",
- )
- kernel = te.compute(kernel_data.shape, lambda *i: k0(*i) + k1(*i), name="Kernel")
- conv = topi.nn.conv2d_nchw(data, kernel, strides, padding, dilation=1)
- relu = topi.nn.relu(conv)
- add = topi.add(data, relu)
-
- dag = auto_scheduler.ComputeDAG([data, kernel_data, add])
- s0 = dag.get_init_state()
-
- pad_temp = s0.stage_ops[1]
- kernel_split = s0.stage_ops[3]
-
- # 0: init state
- ori_its = s0[add].iters
- its = s0.split(add, s0[add].iters[0], [2])
- s0.reorder(add, [its[0], ori_its[1], its[1], ori_its[2], ori_its[3]])
- s0.compute_inline(relu)
-
- # 1: simple cache_write with compute_at
- conv_global = s0.cache_write(conv, "global")
- s0.compute_at(conv_global, conv, s0[conv].iters[3])
-
- # 2: simple cache_read with compute_at
- kernel_global = s0.cache_read(kernel, "global", [conv_global])
- s0.compute_at(kernel_global, conv_global, s0[conv_global].iters[4])
- """
- Placeholder: Data, Kernel_data
- for i0 (0,4)
- for i1 (0,512)
- for i2 (0,9)
- for i3 (0,9)
- pad_temp = ...
- for i0 (0,512)
- for i1 (0,512)
- for i2 (0,3)
- for i3 (0,3)
- Kernel_split = ...
- for i0 (0,512)
- for i1 (0,512)
- for i2 (0,3)
- for i3 (0,3)
- Kernel = ...
- for nn (0,4)
- for ff (0,512)
- for yy (0,7)
- for xx (0,7)
- for nn_c (None)
- for ff_c (None)
- for yy_c (None)
- for xx_c (None)
- for rc (None)
- for ax0 (None)
- for ax1 (None)
- for ax2 (None)
- for ax3 (None)
- Kernel.global = ...
- for ry (None)
- for rx (None)
- compute.global = ...
- compute = ...
- for ax0.0 (0,2)
- for ax1 (0,512)
- for ax0.1 (0,2)
- for ax2 (0,7)
- for ax3 (0,7)
- T_add = ...
- """
- s1 = dag.infer_bound_from_state(s0)
- assert s1[conv].iters[0].range.extent == 4
- assert s1[conv].iters[1].range.extent == 512
- assert s1[conv].iters[2].range.extent == 7
- assert s1[conv].iters[3].range.extent == 7
- assert s1[kernel_global].iters[0].range.extent == 1
- assert s1[kernel_global].iters[1].range.extent == 1
- assert s1[kernel_global].iters[2].range.extent == 3
- assert s1[kernel_global].iters[3].range.extent == 3
- assert s1[conv_global].iters[0].range.extent == 1
- assert s1[conv_global].iters[1].range.extent == 1
- assert s1[conv_global].iters[2].range.extent == 1
- assert s1[conv_global].iters[3].range.extent == 1
- assert s1[conv_global].iters[4].range.extent == 512
- assert s1[conv_global].iters[5].range.extent == 3
- assert s1[conv_global].iters[6].range.extent == 3
-
- # 3: two level cache_read with compute_at
- # preparing for GPU's shared memory & local memory
- pad_temp_global = s0.cache_read(pad_temp, "global", [conv_global])
- pad_temp_shared = s0.cache_read(pad_temp_global, "shared", [conv_global])
- s0.compute_at(pad_temp_global, conv_global, s0[conv_global].iters[2])
- s0.compute_at(pad_temp_shared, conv_global, s0[conv_global].iters[4])
-
- # 4: cache_read with multi readers
- # This stage cannot be compute at to its consumer
- s0.cache_read(data, "global", [pad_temp, add])
- """
- Placeholder: Data, Kernel_data
- for ax0 (0,4)
- for ax1 (0,512)
- for ax2 (0,7)
- for ax3 (0,7)
- Data.global = ...
- for i0 (0,4)
- for i1 (0,512)
- for i2 (0,9)
- for i3 (0,9)
- pad_temp = ...
- for i0 (0,512)
- for i1 (0,512)
- for i2 (0,3)
- for i3 (0,3)
- Kernel_split = ...
- for i0 (0,512)
- for i1 (0,512)
- for i2 (0,3)
- for i3 (0,3)
- Kernel = ...
- for nn (0,4)
- for ff (0,512)
- for yy (0,7)
- for xx (0,7)
- for nn_c (None)
- for ff_c (None)
- for yy_c (None)
- for ax0 (None)
- for ax1 (None)
- for ax2 (None)
- for ax3 (None)
- pad_temp.global = ...
- for xx_c (None)
- for rc (None)
- for ax0 (None)
- for ax1 (None)
- for ax2 (None)
- for ax3 (None)
- Kernel.global = ...
- for ax0 (None)
- for ax1 (None)
- for ax2 (None)
- for ax3 (None)
- pad_temp.global.shared = ...
- for ry (None)
- for rx (None)
- compute.global = ...
- compute = ...
- for ax0.0 (0,2)
- for ax1 (0,512)
- for ax0.1 (0,2)
- for ax2 (0,7)
- for ax3 (0,7)
- T_add = ...
- """
- s1 = dag.infer_bound_from_state(s0)
- assert s1[conv].iters[0].range.extent == 4
- assert s1[conv].iters[1].range.extent == 512
- assert s1[conv].iters[2].range.extent == 7
- assert s1[conv].iters[3].range.extent == 7
- assert s1[kernel_global].iters[0].range.extent == 1
- assert s1[kernel_global].iters[1].range.extent == 1
- assert s1[kernel_global].iters[2].range.extent == 3
- assert s1[kernel_global].iters[3].range.extent == 3
- assert s1[conv_global].iters[0].range.extent == 1
- assert s1[conv_global].iters[1].range.extent == 1
- assert s1[conv_global].iters[2].range.extent == 1
- assert s1[conv_global].iters[3].range.extent == 1
- assert s1[conv_global].iters[4].range.extent == 512
- assert s1[conv_global].iters[5].range.extent == 3
- assert s1[conv_global].iters[6].range.extent == 3
- assert s1[pad_temp_global].iters[0].range.extent == 1
- assert s1[pad_temp_global].iters[1].range.extent == 512
- assert s1[pad_temp_global].iters[2].range.extent == 3
- assert s1[pad_temp_global].iters[3].range.extent == 3
- assert s1[pad_temp_shared].iters[0].range.extent == 1
- assert s1[pad_temp_shared].iters[1].range.extent == 1
- assert s1[pad_temp_shared].iters[2].range.extent == 3
- assert s1[pad_temp_shared].iters[3].range.extent == 3
-
- # 5: cache_write with multi outputs
- # TVM's cache_write actually has a bug with this case:
- #
- # After schedule.cache_write, TVM generate one new stage:
- # From: kernel_data -> kernel_split -> kernel
- # To: kernel_data -> kernel_split_global -> kernel_split -> kernel
- #
- # But with topo sort analyse, we get:
- # // kernel_data -> kernel_split_global -> kernel_split -> kernel
- # \ /
- # ----------------> kernel_split ---------------->
- #
- # TODO(jcf94): Seems there's bug with the input/output tensor. Such multi outputs case
- # should be unusual, so we make some hack on DoCacheWrite. This should be fixed later.
- kernel_split_global = s0.cache_write(kernel_split, "global")
- """
- Placeholder: Data, Kernel_data
- for ax0 (0,4)
- for ax1 (0,512)
- for ax2 (0,7)
- for ax3 (0,7)
- Data.global = ...
- for i0 (0,4)
- for i1 (0,512)
- for i2 (0,9)
- for i3 (0,9)
- pad_temp = ...
- for i0_c (0,512)
- for i1_c (0,512)
- for i2_c (0,3)
- for i3_c (0,3)
- Kernel_split.global = ...
- for i0 (0,512)
- for i1 (0,512)
- for i2 (0,3)
- for i3 (0,3)
- Kernel_split = ...
- (******* Bug here, there should not be two kernel_split stage *******)
- for i0 (0,512)
- for i1 (0,512)
- for i2 (0,3)
- for i3 (0,3)
- Kernel_split = ...
- (******* Bug here, there should not be two kernel_split stage *******)
- for i0 (0,512)
- for i1 (0,512)
- for i2 (0,3)
- for i3 (0,3)
- Kernel = ...
- for nn (0,4)
- for ff (0,512)
- for yy (0,7)
- for xx (0,7)
- for nn_c (None)
- for ff_c (None)
- for yy_c (None)
- for ax0 (None)
- for ax1 (None)
- for ax2 (None)
- for ax3 (None)
- pad_temp.global = ...
- for xx_c (None)
- for rc (None)
- for ax0 (None)
- for ax1 (None)
- for ax2 (None)
- for ax3 (None)
- Kernel.global = ...
- for ax0 (None)
- for ax1 (None)
- for ax2 (None)
- for ax3 (None)
- pad_temp.global.shared = ...
- for ry (None)
- for rx (None)
- compute.global = ...
- compute = ...
- for ax0.0 (0,2)
- for ax1 (0,512)
- for ax0.1 (0,2)
- for ax2 (0,7)
- for ax3 (0,7)
- T_add = ...
- """
- assert len(s0[kernel_split].iters) == len(s0[kernel_split_global].iters)
- for it0, it1 in zip(s0[kernel_split].iters, s0[kernel_split_global].iters):
- assert it0.range == it1.range
-
-
-def test_follow_split_follow_fused_split():
- A, B, C = matmul_auto_scheduler_test(512, 512, 512)
- dag = auto_scheduler.ComputeDAG([A, B, C])
- s0 = dag.get_init_state()
-
- C_global = s0.cache_write(C, "global")
- its0 = s0.split(C, s0[C].iters[0], [4, 2, 8, 4], True)
- split_step0 = len(s0.transform_steps) - 1
- for level in range(1, 6):
- tmp = s0.copy()
- tmp.follow_split(C_global, tmp[C_global].iters[0], split_step0, level)
- for i in range(0, level):
- assert tmp[C].iters[i].range.extent == tmp[C_global].iters[i].range.extent
-
- its1 = s0.split(C, s0[C].iters[5], [2, 2, 4, 8])
- split_step1 = len(s0.transform_steps) - 1
- its = []
- for i0, i1 in zip(its0, its1):
- its.append(i0)
- its.append(i1)
- s0.reorder(C, its)
- for i in range(0, 5):
- s0.fuse(C, [s0[C].iters[i], s0[C].iters[i + 1]])
-
- for level in range(0, 4):
- tmp = s0.copy()
- tmp.follow_fused_split(
- C_global, tmp[C_global].iters[0], [split_step0, split_step1], level, False
- )
- assert tmp[C].iters[level + 1].range.extent == tmp[C_global].iters[0].range.extent
-
- for level in range(0, 4):
- tmp = s0.copy()
- tmp.follow_fused_split(
- C_global, tmp[C_global].iters[0], [split_step0, split_step1], level, True
- )
- assert tmp[C].iters[level + 1].range.extent == tmp[C_global].iters[1].range.extent
-
-
-def test_rfactor():
- A, B, C = matmul_auto_scheduler_test(8, 8, 512)
- dag = auto_scheduler.ComputeDAG([A, B, C])
- s0 = dag.get_init_state()
-
- ko, ki = s0.split(C, s0[C].iters[2], [16])
-
- s1 = s0.copy()
- C_r = s1.rfactor(C, ko, 2)
- """
- Placeholder: A, B
- for i (0,8)
- for j (0,8)
- for k_o (0,32)
- for k_i (0,16)
- C.rf = ...
- for ax0 (0,8)
- for ax1 (0,8)
- for k_o_v (0,32)
- C.repl = ...
- """
- assert s1[C_r].iters[0].range.extent == 8
- assert s1[C_r].iters[1].range.extent == 8
- assert s1[C_r].iters[2].range.extent == 32
- assert s1[C_r].iters[3].range.extent == 16
- assert s1[C].iters[0].range.extent == 8
- assert s1[C].iters[1].range.extent == 8
- assert s1[C].iters[2].range.extent == 32
-
- s2 = s0.copy()
- C_r = s2.rfactor(C, ki, 2)
- """
- Placeholder: A, B
- for i (0,8)
- for j (0,8)
- for k_i (0,16)
- for k_o (0,32)
- C.rf = ...
- for ax0 (0,8)
- for ax1 (0,8)
- for k_i_v (0,16)
- C.repl = ...
- """
- assert s2[C_r].iters[0].range.extent == 8
- assert s2[C_r].iters[1].range.extent == 8
- assert s2[C_r].iters[2].range.extent == 16
- assert s2[C_r].iters[3].range.extent == 32
- assert s2[C].iters[0].range.extent == 8
- assert s2[C].iters[1].range.extent == 8
- assert s2[C].iters[2].range.extent == 16
-
-
-if __name__ == "__main__":
- test_split_fuse_reorder_annotation()
- test_compute_at_root_inline()
- test_cache_read_write()
- test_follow_split_follow_fused_split()
- test_rfactor()
diff --git a/tests/python/auto_scheduler/test_auto_scheduler_measure.py b/tests/python/auto_scheduler/test_auto_scheduler_measure.py
deleted file mode 100644
index 3fd5f97dd8..0000000000
--- a/tests/python/auto_scheduler/test_auto_scheduler_measure.py
+++ /dev/null
@@ -1,427 +0,0 @@
-# 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.
-
-""" Test measurement and log serialization. """
-import json
-
-import multiprocessing
-import numpy as np
-import tvm
-from tvm import topi
-from tvm import te, auto_scheduler
-import tempfile
-import tvm.testing
-import pickle
-from tvm.testing.auto_scheduler import matmul_auto_scheduler_test
-from tvm.auto_scheduler import workload_registry
-
-
-def record_common(dag, s):
- target = tvm.target.Target("llvm")
- task = auto_scheduler.SearchTask(compute_dag=dag, workload_key="test", target=target)
-
- inp = auto_scheduler.measure.MeasureInput(task, s)
- res = auto_scheduler.measure.MeasureResult([0.1], 0, "", 0.2, 1)
-
- # Test in-memory record processing.
- record_str = auto_scheduler.measure_record.dump_record_to_string(inp, res)
- r_inp, r_res = auto_scheduler.measure_record.load_record_from_string(record_str)
- # Only check the workload_key for simplification.
- assert inp.task.workload_key == r_inp.task.workload_key
- assert str(res) == str(r_res)
-
- # Test file-based record processing.
- with tempfile.NamedTemporaryFile() as fp:
- auto_scheduler.save_records(fp.name, [inp], [res])
-
- log_reader = auto_scheduler.RecordReader(fp.name)
- inputs, _ = log_reader.read_lines()
- assert len(inputs) == 1
-
- s1 = dag.infer_bound_from_state(s)
- s2 = dag.infer_bound_from_state(inputs[0].state)
-
- assert s1 == s2
- assert not (s1 == dag.get_init_state())
-
-
-def test_record_split_reorder_fuse_annotation():
- if not tvm.testing.device_enabled("llvm"):
- return
-
- A = te.placeholder((512, 512), name="A")
- B = te.placeholder((512, 512), name="B")
- k = te.reduce_axis((0, 512), name="k")
- C = te.compute((512, 512), lambda i, j: te.sum(A[i][k] * B[k][j], axis=[k]), name="C")
-
- dag = auto_scheduler.ComputeDAG([A, B, C])
- s = dag.get_init_state()
-
- # Split
- its0 = s.split(C, s[C].iters[0], [4, 8, 8])
- its1 = s.split(C, s[C].iters[4], [8, 4, 4])
- # Reorder
- s.reorder(
- C, [its0[0], its1[0], its0[1], its1[1], its0[2], its1[2], its0[3], s[C].iters[8], its1[3]]
- )
- # Fuse
- s.fuse(C, [s[C].iters[0], s[C].iters[1], s[C].iters[2]])
- # Parallel
- s.parallel(C, s[C].iters[0])
- # Thread bind(The blockIdx & threadIdx are used in GPU, just for record testing here)
- s.bind(C, s[C].iters[1], "blockIdx.x")
- s.bind(C, s[C].iters[2], "threadIdx.z")
- s.bind(C, s[C].iters[3], "vthread")
- # Unroll
- s.unroll(C, s[C].iters[4])
- # Vectorize
- s.vectorize(C, s[C].iters[6])
-
- record_common(dag, s)
-
-
-def test_record_compute_at_root_inline_cache_read_write():
- if not tvm.testing.device_enabled("llvm"):
- return
-
- A = te.placeholder((512, 512), name="A")
- AA = topi.nn.relu(A)
- B = te.placeholder((512, 512), name="B")
- k = te.reduce_axis((0, 512), name="k")
- C = te.compute((512, 512), lambda i, j: te.sum(AA[i][k] * B[k][j], axis=[k]), name="C")
-
- dag = auto_scheduler.ComputeDAG([A, B, C])
- s = dag.get_init_state()
-
- # Cache Write
- C_shared = s.cache_write(C, "shared")
- # Compute At
- s.compute_at(C_shared, C, s[C].iters[0])
- # Cache Read
- B_global = s.cache_read(B, "global", [C_shared])
- s.compute_at(B_global, C_shared, s[C_shared].iters[2])
- # Compute Inline
- s.compute_inline(AA)
- # Compute Root
- s.compute_root(C_shared)
-
- record_common(dag, s)
-
-
-def test_record_follow_split_follow_fused_split():
- if not tvm.testing.device_enabled("llvm"):
- return
-
- A = te.placeholder((512, 512), name="A")
- B = te.placeholder((512, 512), name="B")
- k = te.reduce_axis((0, 512), name="k")
- C = te.compute((512, 512), lambda i, j: te.sum(A[i][k] * B[k][j], axis=[k]), name="C")
- D = topi.nn.relu(C)
- E = topi.nn.relu(D)
-
- dag = auto_scheduler.ComputeDAG([A, B, E])
- s = dag.get_init_state()
-
- # Follow Split
- s.split(C, s[C].iters[0], [4, 2, 8, 4], True)
- split_step0 = len(s.transform_steps) - 1
- s.follow_split(C, s[C].iters[5], split_step0, 4)
- # Follow Fused Split
- its0 = s.split(E, s[E].iters[0], [4, 2, 8, 4], True)
- split_step1 = len(s.transform_steps) - 1
- its1 = s.split(E, s[E].iters[5], [2, 4, 2, 4], True)
- split_step2 = len(s.transform_steps) - 1
- its = []
- for i0, i1 in zip(its0, its1):
- its.append(i0)
- its.append(i1)
- for i in range(0, 5):
- s.fuse(E, [s[E].iters[i], s[E].iters[i + 1]])
- s.follow_fused_split(D, s[D].iters[0], [split_step1, split_step2], 2, True)
-
- record_common(dag, s)
-
-
-def test_record_pragma_storage_align_rfactor():
- if not tvm.testing.device_enabled("llvm"):
- return
-
- A = te.placeholder((512, 512), name="A")
- B = te.placeholder((512, 512), name="B")
- k = te.reduce_axis((0, 512), name="k")
- C = te.compute((512, 512), lambda i, j: te.sum(A[i][k] * B[k][j], axis=[k]), name="C")
-
- dag = auto_scheduler.ComputeDAG([A, B, C])
- s = dag.get_init_state()
-
- # Rfactor
- ko, _ = s.split(C, s[C].iters[2], [16])
- s.rfactor(C, ko, 2)
- # Pragma
- s.pragma(C, s[C].iters[0], "auto_unroll_max_step$64")
- # StorageAlign
- s.storage_align(C, s[C].iters[-1], 8, 4)
-
- record_common(dag, s)
-
-
-def test_recover_measure_input():
- task = auto_scheduler.SearchTask(
- func=matmul_auto_scheduler_test, args=(512, 512, 512), target="llvm"
- )
-
- inp = auto_scheduler.measure.MeasureInput(task, task.compute_dag.init_state)
- res = auto_scheduler.measure.MeasureResult([0.1], 0, "", 0.2, 1)
-
- with tempfile.NamedTemporaryFile() as fp:
- auto_scheduler.save_records(fp.name, [inp], [res])
-
- log_reader = auto_scheduler.RecordReader(fp.name)
- inputs, _ = log_reader.read_lines()
- assert len(inputs) == 1
-
- raw_inp = inputs[0]
-
- correct_inp = auto_scheduler.measure.recover_measure_input(raw_inp)
- assert str(correct_inp.task.compute_dag) == str(inp.task.compute_dag)
-
- correct_inp = auto_scheduler.measure.recover_measure_input(raw_inp, rebuild_state=True)
- assert str(correct_inp.state) == str(inp.state)
-
-
-def test_workload_dis_factor():
- calc = auto_scheduler.utils.calc_workload_dis_factor
- decode = auto_scheduler.utils.decode_workload_key
-
- # Identical
- target_wkl_key = json.dumps(
- ["func1", [8, 3, 224, 224], [32, 3, 3, 3], [0, 0], [1, 1], "float32"]
- )
- assert calc(decode(target_wkl_key), decode(target_wkl_key)) == 1
-
- # Compatible with a factor
- wkl_key = json.dumps(["func1", [1, 3, 112, 112], [32, 3, 3, 3], [0, 0], [1, 1], "float32"])
- assert calc(decode(target_wkl_key), decode(wkl_key)) == 8 * 2 * 2
-
- # Incompatible argument with zeros
- wkl_key = json.dumps(["func1", [8, 3, 224, 224], [32, 3, 3, 3], [1, 1], [1, 1], "float32"])
- assert calc(decode(target_wkl_key), decode(wkl_key)) == float("inf")
- wkl_key = json.dumps(["func1", [8, 3, 224, 224], [32, 3, 3, 3], [0, 0], [0, 0], "float32"])
- assert calc(decode(target_wkl_key), decode(wkl_key)) == float("inf")
-
- # Incompatible non-integter argument
- wkl_key = json.dumps(["func1", [8, 3, 224, 224], [32, 3, 3, 3], [0, 0], [1, 1], "int8"])
- assert calc(decode(target_wkl_key), decode(wkl_key)) == float("inf")
-
- # Incompatible function
- wkl_key = json.dumps(["func2", [8, 3, 224, 224], [32, 3, 3, 3], [0, 0], [1, 1], "float32"])
- assert calc(decode(target_wkl_key), decode(wkl_key)) == float("inf")
-
- # Incompatible due to non-dividable factor
- wkl_key = json.dumps(["func1", [8, 3, 223, 223], [32, 3, 3, 3], [0, 0], [1, 1], "float32"])
- assert calc(decode(target_wkl_key), decode(wkl_key)) == float("inf")
-
-
-def test_measure_local_builder_runner():
- if not tvm.testing.device_enabled("llvm"):
- return
-
- task = auto_scheduler.SearchTask(
- func=matmul_auto_scheduler_test, args=(512, 512, 512), target="llvm"
- )
-
- for enable_cpu_cache_flush in [True, False]:
- minp = auto_scheduler.MeasureInput(task, task.compute_dag.init_state)
- local_builder = auto_scheduler.LocalBuilder()
- local_runner = auto_scheduler.LocalRunner(
- timeout=60, enable_cpu_cache_flush=enable_cpu_cache_flush
- )
-
- bress = local_builder.build([minp])
- assert bress[0].error_no == 0
- mress = local_runner.run([minp], bress)
- assert mress[0].error_no == 0
-
-
-def test_dag_measure_local_builder_runner():
- if not tvm.testing.device_enabled("llvm"):
- return
-
- A = te.placeholder((512, 512), name="A")
- B = te.placeholder((512, 512), name="B")
- k = te.reduce_axis((0, 512), name="k")
- C = te.compute((512, 512), lambda i, j: te.sum(A[i][k] * B[k][j], axis=[k]), name="C")
- D = topi.nn.relu(C)
- E = topi.nn.relu(D)
-
- tensors = [A, B, E]
- dag = auto_scheduler.ComputeDAG(tensors)
- key = workload_registry.register_workload_tensors(dag.workload_key(), tensors)
- transfer_data = workload_registry.serialize_workload_registry_entry(key)
- f_data = pickle.dumps(transfer_data)
- f_new = pickle.loads(f_data)
- del workload_registry.WORKLOAD_FUNC_REGISTRY[key]
- workload_registry.deserialize_workload_registry_entry(f_new)
-
- target = tvm.target.Target("llvm")
- task = auto_scheduler.SearchTask(compute_dag=dag, workload_key=key, target=target)
-
- for enable_cpu_cache_flush in [True, False]:
- minp = auto_scheduler.MeasureInput(task, task.compute_dag.init_state)
- local_builder = auto_scheduler.LocalBuilder()
- local_runner = auto_scheduler.LocalRunner(
- timeout=60, enable_cpu_cache_flush=enable_cpu_cache_flush
- )
-
- bress = local_builder.build([minp])
- assert bress[0].error_no == 0
- mress = local_runner.run([minp], bress)
- assert mress[0].error_no == 0
-
-
-def test_workload_serialization():
- key = tvm.auto_scheduler.utils.get_func_name(matmul_auto_scheduler_test)
- transfer_data = workload_registry.serialize_workload_registry_entry(key)
- f_data = pickle.dumps(transfer_data)
- f_new = pickle.loads(f_data)
- del workload_registry.WORKLOAD_FUNC_REGISTRY[key]
- workload_registry.deserialize_workload_registry_entry(f_new)
-
-
-def test_measure_local_builder_rpc_runner():
- if not tvm.testing.device_enabled("llvm"):
- return
-
- task = auto_scheduler.SearchTask(
- func=matmul_auto_scheduler_test, args=(512, 512, 512), target="llvm"
- )
-
- for enable_cpu_cache_flush in [True, False]:
- minp = auto_scheduler.MeasureInput(task, task.compute_dag.init_state)
- local_builder = auto_scheduler.LocalBuilder()
- measure_ctx = auto_scheduler.LocalRPCMeasureContext(
- timeout=60, enable_cpu_cache_flush=enable_cpu_cache_flush
- )
- rpc_runner = measure_ctx.runner
-
- bress = local_builder.build([minp])
- assert bress[0].error_no == 0
- mress = rpc_runner.run([minp], bress)
- assert mress[0].error_no == 0
-
- del measure_ctx
-
-
-def measure_local_builder_rpc_runner_spawn():
- assert multiprocessing.get_start_method(False) == "spawn"
- test_measure_local_builder_rpc_runner()
-
-
-@tvm.testing.requires_llvm
-def test_measure_local_builder_rpc_runner_spawn():
- ctx = multiprocessing.get_context("spawn")
- p = ctx.Process(target=measure_local_builder_rpc_runner_spawn)
- p.start()
- p.join()
-
-
-@tvm.testing.requires_llvm
-def test_measure_target_host():
- task = auto_scheduler.SearchTask(
- func=matmul_auto_scheduler_test,
- args=(512, 512, 512),
- target=tvm.target.Target("llvm", "llvm -mtriple=aarch64-linux-gnu"),
- )
-
- inp = auto_scheduler.measure.MeasureInput(task, task.compute_dag.init_state)
- res = auto_scheduler.measure.MeasureResult([0.1], 0, "", 0.2, 1)
-
- with tempfile.NamedTemporaryFile() as fp:
- auto_scheduler.save_records(fp.name, [inp], [res])
-
- log_reader = auto_scheduler.RecordReader(fp.name)
- inputs, _ = log_reader.read_lines()
- assert len(inputs) == 1
-
- raw_inp = inputs[0]
-
- recovered_inp = auto_scheduler.measure.recover_measure_input(raw_inp)
- assert str(recovered_inp.task.target.host) == str(inp.task.target.host)
-
-
-@tvm.testing.requires_llvm
-def test_measure_special_inputs_map_by_name_local_runner():
- @auto_scheduler.register_workload
- def foo():
- X = te.placeholder(shape=[10], dtype="int32")
- Index = te.placeholder(shape=[1], dtype="int32", name="Index")
- Y = te.compute((1,), lambda i: X[Index[i]])
- return [X, Index, Y]
-
- # This workload cannot use random input for the `Index` input
- task = auto_scheduler.SearchTask(
- func=foo,
- target="llvm",
- task_inputs={
- "Index": tvm.nd.array(np.array([5], dtype="int32")),
- },
- )
-
- minp = auto_scheduler.MeasureInput(task, task.compute_dag.init_state)
- local_builder = auto_scheduler.LocalBuilder()
- local_runner = auto_scheduler.LocalRunner(timeout=10)
-
- bress = local_builder.build([minp])
- assert bress[0].error_no == 0
- mress = local_runner.run([minp], bress)
- assert mress[0].error_no == 0
-
-
-@tvm.testing.requires_llvm
-def test_measure_special_inputs_map_by_name_rpc_runner():
- @auto_scheduler.register_workload
- def foo():
- X = te.placeholder(shape=[10], dtype="int32")
- Index = te.placeholder(shape=[1], dtype="int32", name="Index")
- Y = te.compute((1,), lambda i: X[Index[i]])
- return [X, Index, Y]
-
- # This workload cannot use random input for the `Index` input
- task = auto_scheduler.SearchTask(
- func=foo,
- target="llvm",
- task_inputs={
- "Index": tvm.nd.array(np.array([5], dtype="int32")),
- },
- )
-
- for enable_cpu_cache_flush in [True, False]:
- minp = auto_scheduler.MeasureInput(task, task.compute_dag.init_state)
- local_builder = auto_scheduler.LocalBuilder()
- measure_ctx = auto_scheduler.LocalRPCMeasureContext(
- timeout=60, enable_cpu_cache_flush=enable_cpu_cache_flush
- )
- rpc_runner = measure_ctx.runner
-
- bress = local_builder.build([minp])
- assert bress[0].error_no == 0
- mress = rpc_runner.run([minp], bress)
- assert mress[0].error_no == 0
-
-
-if __name__ == "__main__":
- tvm.testing.main()
diff --git a/tests/python/auto_scheduler/test_auto_scheduler_search_policy.py b/tests/python/auto_scheduler/test_auto_scheduler_search_policy.py
deleted file mode 100644
index a9f6596a85..0000000000
--- a/tests/python/auto_scheduler/test_auto_scheduler_search_policy.py
+++ /dev/null
@@ -1,220 +0,0 @@
-# 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.
-
-"""Test search policy"""
-
-import random
-import multiprocessing
-import numpy as np
-import tempfile
-
-import tvm
-import tvm.testing
-from tvm import auto_scheduler
-from tvm.auto_scheduler.utils import get_const_tuple
-
-from tvm.testing.auto_scheduler import (
- matmul_auto_scheduler_test,
- zero_rank_compute_auto_scheduler_test,
- zero_rank_reduce_auto_scheduler_test,
-)
-import multiprocessing
-
-
-class CustomMeasureCallback(auto_scheduler.measure.PythonBasedMeasureCallback):
- """A simple Python-based callback for testing."""
-
- def callback(self, policy, inputs, results):
- assert isinstance(policy, auto_scheduler.search_policy.SearchPolicy)
- for inp, res in zip(inputs, results):
- assert isinstance(inp, auto_scheduler.MeasureInput)
- assert isinstance(res, auto_scheduler.MeasureResult)
-
-
-def search_common(
- task=None,
- target="llvm",
- search_policy="sketch",
- runner="local",
- num_measure_trials=100,
- cost_model=auto_scheduler.RandomModel(),
- init_search_callbacks=None,
-):
- if task is None:
- task = auto_scheduler.SearchTask(
- func=matmul_auto_scheduler_test, args=(64, 64, 64), target=target
- )
- target = task.target
-
- print("Test search policy '%s' for '%s'" % (search_policy, target))
-
- with tempfile.NamedTemporaryFile() as fp:
- log_file = fp.name
-
- init_search_callbacks = init_search_callbacks or []
- init_search_callbacks.append(auto_scheduler.PreloadMeasuredStates(log_file))
-
- if search_policy == "empty":
- search_policy = auto_scheduler.EmptyPolicy(task)
- elif search_policy == "sketch":
- search_policy = auto_scheduler.SketchPolicy(
- task, program_cost_model=cost_model, init_search_callbacks=init_search_callbacks
- )
- else:
- raise ValueError("Invalid policy: " + search_policy)
-
- # Tune
- tuning_options = auto_scheduler.TuningOptions(
- num_measure_trials=num_measure_trials,
- num_measures_per_round=2,
- early_stopping=1,
- runner=runner,
- measure_callbacks=[auto_scheduler.RecordToFile(log_file), CustomMeasureCallback()],
- )
- task.tune(tuning_options=tuning_options, search_policy=search_policy)
-
- # Compile with the best schedule
- sch, args = task.apply_best(log_file)
- mod = tvm.build(sch, args, target)
-
- # Compile with naive schedule for correctness check
- sch, args = task.compute_dag.apply_steps_from_state(task.compute_dag.init_state)
- mod_ref = tvm.build(sch, args, "llvm")
-
- ctx = tvm.device(str(target), 0)
- np_arrays = [np.random.uniform(size=get_const_tuple(x.shape)).astype(x.dtype) for x in args]
-
- tvm_arrays = [tvm.nd.array(x, ctx) for x in np_arrays]
- mod(*tvm_arrays)
- actual = [x.numpy() for x in tvm_arrays]
-
- tvm_arrays = [tvm.nd.array(x) for x in np_arrays]
- mod_ref(*tvm_arrays)
- expected = [x.numpy() for x in tvm_arrays]
-
- for x, y in zip(actual, expected):
- tvm.testing.assert_allclose(x, y, rtol=1e-5)
-
-
-@tvm.testing.requires_llvm
-def test_workload_registry_empty_policy():
- search_common(search_policy="empty", num_measure_trials=2)
-
- N = 64
- target = "llvm"
- search_common(
- task=auto_scheduler.SearchTask(
- func="matmul_auto_scheduler_test", args=(N, N, N), target=target
- ),
- num_measure_trials=2,
- search_policy="empty",
- )
- search_common(
- task=auto_scheduler.SearchTask(
- func="matmul_auto_scheduler_test_rename_1", args=(N, N, N), target=target
- ),
- num_measure_trials=2,
- search_policy="empty",
- )
-
-
-@tvm.testing.requires_llvm
-def test_sketch_search_policy_basic():
- search_common()
-
-
-def sketch_search_policy_basic_spawn():
- assert multiprocessing.get_start_method(False) == "spawn"
- test_sketch_search_policy_basic()
-
-
-@tvm.testing.requires_llvm
-def test_sketch_search_policy_basic_spawn():
- ctx = multiprocessing.get_context("spawn")
- p = ctx.Process(target=sketch_search_policy_basic_spawn)
- p.start()
- p.join()
-
-
-@tvm.testing.requires_llvm
-def test_sketch_search_policy_xgbmodel():
- search_common(cost_model=auto_scheduler.XGBModel())
-
-
-@tvm.testing.requires_cuda
-def test_sketch_search_policy_cuda_rpc_runner():
- measure_ctx = auto_scheduler.LocalRPCMeasureContext()
- search_common(target="cuda", runner=measure_ctx.runner)
-
-
-@tvm.testing.requires_cuda
-def test_sketch_search_policy_cuda_xgbmodel_rpc_runner():
- measure_ctx = auto_scheduler.LocalRPCMeasureContext()
- search_common(target="cuda", runner=measure_ctx.runner, cost_model=auto_scheduler.XGBModel())
-
-
-@tvm.testing.requires_llvm
-@tvm.testing.requires_cuda
-def test_sketch_search_policy_zero_rank():
- measure_ctx = auto_scheduler.LocalRPCMeasureContext()
- for target in ["llvm", "cuda"]:
- task = auto_scheduler.SearchTask(
- func=zero_rank_compute_auto_scheduler_test, args=(10,), target=target
- )
- search_common(task, runner=measure_ctx.runner)
-
- task = auto_scheduler.SearchTask(
- func=zero_rank_reduce_auto_scheduler_test, args=(10,), target=target
- )
- search_common(task, runner=measure_ctx.runner)
-
-
-@tvm.testing.requires_llvm
-def test_sketch_search_policy_custom_sketch():
- def meet_condition_func(search_policy, state, stage_id):
- return auto_scheduler.PreloadCustomSketchRule.APPLY_AND_SKIP_REST
-
- def apply_func(search_policy, state, stage_id):
- ret = []
- state = auto_scheduler.loop_state.State(state, search_policy.search_task.compute_dag)
- C = state.stage_ops[2]
-
- ret.append([state.state_object, -1])
-
- s1 = state.copy()
- i, _, _ = s1[C].iters
- s1.split(C, i, [8])
- ret.append([s1.state_object, -1])
- return ret
-
- search_common(
- cost_model=auto_scheduler.XGBModel(),
- init_search_callbacks=[
- auto_scheduler.PreloadCustomSketchRule(meet_condition_func, apply_func)
- ],
- )
-
-
-if __name__ == "__main__":
- test_workload_registry_empty_policy()
- test_sketch_search_policy_basic()
- test_sketch_search_policy_basic_spawn()
- test_sketch_search_policy_xgbmodel()
- test_sketch_search_policy_cuda_rpc_runner()
- test_sketch_search_policy_cuda_xgbmodel_rpc_runner()
- test_sketch_search_policy_zero_rank()
- test_sketch_search_policy_custom_sketch()
diff --git a/tests/python/auto_scheduler/test_auto_scheduler_search_task.py b/tests/python/auto_scheduler/test_auto_scheduler_search_task.py
deleted file mode 100644
index 7c5441e818..0000000000
--- a/tests/python/auto_scheduler/test_auto_scheduler_search_task.py
+++ /dev/null
@@ -1,216 +0,0 @@
-# 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.
-
-"""Test search policy"""
-
-import numpy as np
-import tempfile
-
-import tvm
-import tvm.testing
-from tvm import auto_scheduler
-from tvm.auto_scheduler.utils import get_const_tuple
-from tvm.testing.auto_scheduler import (
- matmul_auto_scheduler_test,
- zero_rank_compute_auto_scheduler_test,
- zero_rank_reduce_auto_scheduler_test,
-)
-
-
-def test_search_task_add_task_input():
- auto_scheduler.search_task.TASK_INPUT_BUFFER_TABLE.clear()
- N = 64
- target = "llvm"
- test_input_0 = tvm.runtime.ndarray.empty((64, 64))
- test_input_1 = tvm.runtime.ndarray.empty((10, 20))
- test_input_2 = tvm.runtime.ndarray.empty((30, 40, 50))
- task = auto_scheduler.SearchTask(
- func="matmul_auto_scheduler_test",
- args=(N, N, N),
- target=target,
- task_inputs={
- "test_input_0": test_input_0,
- "test_input_1": test_input_1,
- "test_input_2": test_input_2,
- },
- task_inputs_overwrite=True,
- )
-
- assert len(task.task_input_names) == 3
- assert task.task_input_names[0] == "test_input_0"
- assert task.task_input_names[1] == "test_input_1"
- assert task.task_input_names[2] == "test_input_2"
-
-
-def test_search_task_record():
- auto_scheduler.search_task.TASK_INPUT_BUFFER_TABLE.clear()
- N = 64
- target = "llvm"
-
- # Log with no task input
- task = auto_scheduler.SearchTask(
- func="matmul_auto_scheduler_test", args=(N, N, N), target=target
- )
- task_record = auto_scheduler._ffi_api.SerializeSearchTask(task)
- new_task = auto_scheduler._ffi_api.DeserializeSearchTask(task_record)
- # TODO(jcf94): Check the compute dag & hardware parameter
- assert task.workload_key == new_task.workload_key
- assert str(task.target) == str(new_task.target)
- assert str(task.target.host) == str(new_task.target.host)
- assert task.layout_rewrite_option == new_task.layout_rewrite_option
-
- # Log with 1 task input
- test_input_0 = tvm.runtime.ndarray.empty((64, 64))
- task = auto_scheduler.SearchTask(
- func="matmul_auto_scheduler_test",
- args=(N, N, N),
- target=target,
- task_inputs={"test_input_0": test_input_0},
- task_inputs_overwrite=True,
- )
- task_record = auto_scheduler._ffi_api.SerializeSearchTask(task)
- new_task = auto_scheduler._ffi_api.DeserializeSearchTask(task_record)
- assert task.workload_key == new_task.workload_key
- assert str(task.target) == str(new_task.target)
- assert str(task.target.host) == str(new_task.target.host)
- assert task.layout_rewrite_option == new_task.layout_rewrite_option
- assert len(new_task.task_input_names) == 1
- assert new_task.task_input_names[0] == "test_input_0"
-
- # Log with multiple task inputs
- test_input_1 = tvm.runtime.ndarray.empty((64, 64))
- task = auto_scheduler.SearchTask(
- func="matmul_auto_scheduler_test",
- args=(N, N, N),
- target=target,
- task_inputs={
- "test_input_0": test_input_0,
- "test_input_1": test_input_1,
- },
- task_inputs_overwrite=True,
- )
- task_record = auto_scheduler._ffi_api.SerializeSearchTask(task)
- new_task = auto_scheduler._ffi_api.DeserializeSearchTask(task_record)
- assert task.workload_key == new_task.workload_key
- assert str(task.target) == str(new_task.target)
- assert str(task.target.host) == str(new_task.target.host)
- assert task.layout_rewrite_option == new_task.layout_rewrite_option
- assert len(new_task.task_input_names) == 2
- assert new_task.task_input_names[0] == "test_input_0"
- assert new_task.task_input_names[1] == "test_input_1"
-
- # Log with version 0.5
- v5_log = (
- """["[\\\"matmul_auto_scheduler_test\\\", 64, 64, 64]", """
- f'"{str(tvm.target.Target(target))}"'
- """, [6, 64, 64, 0, 0, 0, 0, 0], "", 1]"""
- )
- new_task = auto_scheduler._ffi_api.DeserializeSearchTask(v5_log)
- assert task.workload_key == new_task.workload_key
- assert str(task.target) == str(new_task.target)
- assert str(task.target.host) == str(new_task.target.host)
- assert task.layout_rewrite_option == new_task.layout_rewrite_option
- assert len(new_task.task_input_names) == 0
-
-
-def test_recover_measure_input_with_task_input():
- auto_scheduler.search_task.TASK_INPUT_BUFFER_TABLE.clear()
- target = "llvm"
-
- # Since this file is tests for search_task, we only check the search_task here
-
- # Log with no task input
- task = auto_scheduler.SearchTask(
- func=matmul_auto_scheduler_test, args=(512, 512, 512), target=target
- )
- inp = auto_scheduler.measure.MeasureInput(task, task.compute_dag.init_state)
- res = auto_scheduler.measure.MeasureResult([0.1], 0, "", 0.2, 1)
- measure_record = auto_scheduler.measure_record.dump_record_to_string(inp, res)
- measure_log = auto_scheduler.measure_record.load_record_from_string(measure_record)
- new_task = measure_log[0].task
- assert task.workload_key == new_task.workload_key
- assert str(task.target) == str(new_task.target)
- assert str(task.target.host) == str(new_task.target.host)
- assert task.layout_rewrite_option == new_task.layout_rewrite_option
-
- # Log with 1 task input
- test_input_0 = tvm.runtime.ndarray.empty((64, 64))
- task = auto_scheduler.SearchTask(
- func=matmul_auto_scheduler_test,
- args=(512, 512, 512),
- target=target,
- task_inputs={
- "test_input_0": test_input_0,
- },
- task_inputs_overwrite=True,
- )
- inp = auto_scheduler.measure.MeasureInput(task, task.compute_dag.init_state)
- res = auto_scheduler.measure.MeasureResult([0.1], 0, "", 0.2, 1)
- measure_record = auto_scheduler.measure_record.dump_record_to_string(inp, res)
- measure_log = auto_scheduler.measure_record.load_record_from_string(measure_record)
- new_task = measure_log[0].task
- assert task.workload_key == new_task.workload_key
- assert str(task.target) == str(new_task.target)
- assert str(task.target.host) == str(new_task.target.host)
- assert task.layout_rewrite_option == new_task.layout_rewrite_option
- assert len(new_task.task_input_names) == 1
- assert new_task.task_input_names[0] == "test_input_0"
-
- # Log with multiple task inputs
- test_input_1 = tvm.runtime.ndarray.empty((64, 64))
- task = auto_scheduler.SearchTask(
- func=matmul_auto_scheduler_test,
- args=(512, 512, 512),
- target=target,
- task_inputs={
- "test_input_0": test_input_0,
- "test_input_1": test_input_1,
- },
- task_inputs_overwrite=True,
- )
- inp = auto_scheduler.measure.MeasureInput(task, task.compute_dag.init_state)
- res = auto_scheduler.measure.MeasureResult([0.1], 0, "", 0.2, 1)
- measure_record = auto_scheduler.measure_record.dump_record_to_string(inp, res)
- measure_log = auto_scheduler.measure_record.load_record_from_string(measure_record)
- new_task = measure_log[0].task
- assert task.workload_key == new_task.workload_key
- assert str(task.target) == str(new_task.target)
- assert str(task.target.host) == str(new_task.target.host)
- assert task.layout_rewrite_option == new_task.layout_rewrite_option
- assert len(new_task.task_input_names) == 2
- assert new_task.task_input_names[0] == "test_input_0"
- assert new_task.task_input_names[1] == "test_input_1"
-
- # Log with version 0.5
- v5_log = (
- """{"i": [["[\\\"matmul_auto_scheduler_test\\\", 512, 512, 512]", """
- f'"{str(tvm.target.Target(target))}"'
- """, [6, 64, 64, 0, 0, 0, 0, 0], "", 1], [[], []]], "r": [[0.1], 0, 0.2, 1], "v": "v0.6"}"""
- )
- measure_log = auto_scheduler.measure_record.load_record_from_string(v5_log)
- new_task = measure_log[0].task
- assert task.workload_key == new_task.workload_key
- assert str(task.target) == str(new_task.target)
- assert str(task.target.host) == str(new_task.target.host)
- assert task.layout_rewrite_option == new_task.layout_rewrite_option
- assert len(new_task.task_input_names) == 0
-
-
-if __name__ == "__main__":
- test_search_task_add_task_input()
- test_search_task_record()
- test_recover_measure_input_with_task_input()
diff --git a/tests/python/auto_scheduler/test_auto_scheduler_sketch_generation.py b/tests/python/auto_scheduler/test_auto_scheduler_sketch_generation.py
deleted file mode 100644
index a3f63a3849..0000000000
--- a/tests/python/auto_scheduler/test_auto_scheduler_sketch_generation.py
+++ /dev/null
@@ -1,454 +0,0 @@
-# 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.
-
-""" Test sketch generation. """
-
-import sys
-import tvm
-import tvm.testing
-
-import pytest
-
-from tvm import te, auto_scheduler
-from tvm.auto_scheduler import _ffi_api
-from tvm.auto_scheduler.loop_state import Stage
-
-from tvm.testing.auto_scheduler import (
- matmul_auto_scheduler_test,
- double_matmul_auto_scheduler_test,
- conv2d_nchw_bn_relu_auto_scheduler_test,
- max_pool2d_auto_scheduler_test,
- min_nm_auto_scheduler_test,
- softmax_nm_auto_scheduler_test,
- softmax_abcd_auto_scheduler_test,
- conv2d_winograd_nhwc_auto_scheduler_test,
- zero_rank_reduce_auto_scheduler_test,
-)
-
-
-def generate_sketches(
- workload_func, args, target, print_for_debug=False, init_search_callbacks=None
-):
- # NOTE: test_cpu_matmul_sketch and test_cpu_max_pool2d_sketch assume 4 cores to trigger all
- # possible sketch generations.
- task = auto_scheduler.SearchTask(
- func=workload_func,
- args=args,
- target=target,
- hardware_params=auto_scheduler.HardwareParams(num_cores=4, target=target),
- )
- policy = auto_scheduler.SketchPolicy(
- task, verbose=0, init_search_callbacks=init_search_callbacks
- )
- return policy.generate_sketches(print_for_debug)
-
-
-def assert_compute_at_condition(stage, condition):
- assert stage.compute_at == Stage.COMPUTE_AT_TRANS_TABLE[condition]
-
-
-def assert_is_tiled(stage):
- assert _ffi_api.SearchPolicyUtilsIsTiled(stage)
-
-
-def assert_is_not_tiled(stage):
- assert not _ffi_api.SearchPolicyUtilsIsTiled(stage)
-
-
-def assert_has_cache_write(state, stage_id):
- assert _ffi_api.SearchPolicyUtilsHasCacheWriteStage(state, stage_id)
-
-
-def assert_has_cache_read(state, stage_id):
- assert _ffi_api.SearchPolicyUtilsHasCacheReadStage(state, stage_id)
-
-
-def assert_has_rfactor(state, stage_id):
- assert _ffi_api.SearchPolicyUtilsHasRfactorStage(state, stage_id)
-
-
-def assert_has_cross_thread_reduction(state, stage_id):
- assert _ffi_api.SearchPolicyUtilsHasCrossThreadReduction(state, stage_id)
-
-
-def test_cpu_matmul_sketch():
- sketches = generate_sketches(matmul_auto_scheduler_test, (512, 512, 512), "llvm")
- """ 3 multi-level tiling sketches
- No.0 : Multi-level tiling
- No.1 : Multi-level tiling with cache write on position 0
- No.2 : Multi-level tiling with cache write on position 1
- """
- assert len(sketches) == 3
- # Sketch 0
- assert_is_tiled(sketches[0].stages[2])
- # Sketch 1
- assert_is_tiled(sketches[1].stages[2])
- assert_has_cache_write(sketches[1], 2)
- assert_compute_at_condition(sketches[1].stages[2], "iter")
- # Sketch 2
- assert_is_tiled(sketches[2].stages[2])
- assert_has_cache_write(sketches[2], 2)
- assert_compute_at_condition(sketches[2].stages[2], "iter")
- assert sketches[1] != sketches[2]
-
- sketches = generate_sketches(matmul_auto_scheduler_test, (8, 8, 512), "llvm")
- """ 2 rfactor sketches + 3 multi-level tiling sketches
- No.0 : Rfactor with factor position 0
- No.1 : Rfactor with factor position 1
- No.2 : Multi-level tiling
- No.3 : Multi-level tiling with cache write on position 0
- No.4 : Multi-level tiling with cache write on position 1
- """
- assert len(sketches) == 5
- # Sketch 0
- assert_has_rfactor(sketches[0], 2)
- # Sketch 1
- assert_has_rfactor(sketches[1], 2)
- assert sketches[0] != sketches[1]
- # Sketch 2
- assert_is_tiled(sketches[2].stages[2])
- # Sketch 3
- assert_is_tiled(sketches[3].stages[2])
- assert_has_cache_write(sketches[3], 2)
- assert_compute_at_condition(sketches[3].stages[2], "iter")
- # Sketch 4
- assert_is_tiled(sketches[4].stages[2])
- assert_has_cache_write(sketches[4], 2)
- assert_compute_at_condition(sketches[4].stages[2], "iter")
- assert sketches[3] != sketches[4]
-
- sketches = generate_sketches(double_matmul_auto_scheduler_test, (512,), "llvm")
- """ 3 multi-level tiling sketches for one matmul, so 3 * 3 = 9 sketches in total """
- assert len(sketches) == 9
- assert_is_tiled(sketches[8].stages[5])
-
-
-def test_cpu_conv2d_bn_relu_sketch():
- sketches = generate_sketches(
- conv2d_nchw_bn_relu_auto_scheduler_test, (1, 56, 56, 512, 512, 3, 1, 1), "llvm"
- )
- """ 3 multi-level tiling sketches
- No.0 : Conv2d multi-level tiling with fusion on position 0
- No.1 : Conv2d multi-level tiling with fusion on position 1
- No.2 : Conv2d multi-level tiling without fusion
- """
- assert len(sketches) == 3
- # Sketch 0
- assert_is_not_tiled(sketches[0].stages[1])
- assert_is_tiled(sketches[0].stages[3])
- assert_compute_at_condition(sketches[0].stages[3], "iter")
- assert_compute_at_condition(sketches[0].stages[5], "inlined")
- assert_compute_at_condition(sketches[0].stages[7], "inlined")
- assert_compute_at_condition(sketches[0].stages[9], "inlined")
- assert_is_tiled(sketches[0].stages[10])
- # Sketch 1
- assert_is_not_tiled(sketches[1].stages[1])
- assert_is_tiled(sketches[1].stages[3])
- assert_compute_at_condition(sketches[1].stages[3], "iter")
- assert_compute_at_condition(sketches[1].stages[5], "inlined")
- assert_compute_at_condition(sketches[1].stages[7], "inlined")
- assert_compute_at_condition(sketches[1].stages[9], "inlined")
- assert_is_tiled(sketches[1].stages[10])
- # Sketch 2
- assert_is_not_tiled(sketches[2].stages[1])
- assert_is_tiled(sketches[2].stages[3])
- assert_compute_at_condition(sketches[2].stages[3], "root")
- assert_compute_at_condition(sketches[2].stages[5], "inlined")
- assert_compute_at_condition(sketches[2].stages[7], "inlined")
- assert_compute_at_condition(sketches[2].stages[9], "inlined")
- assert_is_not_tiled(sketches[2].stages[10])
-
-
-def test_cpu_max_pool2d_sketch():
- sketches = generate_sketches(max_pool2d_auto_scheduler_test, (1, 56, 56, 512, 1), "llvm")
- """ 1 default sketch """
- assert len(sketches) == 1
- # Sketch 0
- assert len(sketches[0].transform_steps) == 0
-
-
-def test_cpu_min_sketch():
- sketches = generate_sketches(min_nm_auto_scheduler_test, (10, 1024), "llvm")
- """ 2 rfactor sketches + 1 default sketch
- No.0 : Rfactor with factor position 0
- No.1 : Rfactor with factor position 1
- No.2 : Default sketch
- """
- assert len(sketches) == 3
- # Sketch 0
- assert_has_rfactor(sketches[0], 1)
- # Sketch 1
- assert_has_rfactor(sketches[1], 1)
- assert sketches[0] != sketches[1]
- # Sketch 2
- assert len(sketches[2].transform_steps) == 0
-
-
-def test_cpu_softmax_sketch():
- sketches = generate_sketches(softmax_nm_auto_scheduler_test, (1, 1024), "llvm")
- """ (2 rfactor sketches + 1 default sketch) * (2 rfactor sketches + 1 default sketch) """
- assert len(sketches) == (3 * 3)
- for i in range(0, 3):
- for j in range(0, 3):
- sketch = sketches[i * 3 + j]
- if j in [0, 1]:
- assert_has_rfactor(sketch, 1)
- if i in [0, 1]:
- assert_has_rfactor(sketch, 4 if j in [0, 1] else 3)
- assert len(sketches[8].transform_steps) == 0
-
- sketches = generate_sketches(softmax_abcd_auto_scheduler_test, (1, 12, 128, 128), "llvm")
- """ (2 rfactor sketches + 1 default sketch) * (2 rfactor sketches + 1 default sketch) """
- assert len(sketches) == (3 * 3)
- for i in range(0, 3):
- for j in range(0, 3):
- sketch = sketches[i * 3 + j]
- if j in [0, 1]:
- assert_has_rfactor(sketch, 1)
- if i in [0, 1]:
- assert_has_rfactor(sketch, 4 if j in [0, 1] else 3)
- assert len(sketches[8].transform_steps) == 0
-
-
-def test_cpu_conv2d_winograd_sketch():
- sketches = generate_sketches(
- conv2d_winograd_nhwc_auto_scheduler_test, (1, 28, 28, 128, 128, 3, 1, 1), "llvm"
- )
- """ 3 multi-level tiling sketches
- No.0 : Bgemm multi-level tiling
- No.1 : Bgemm multi-level tiling with cache write on position 0
- No.2 : Bgemm multi-level tiling with cache write on position 1
- """
- assert len(sketches) == 3
- # Sketch 0
- assert_is_not_tiled(sketches[0].stages[1])
- assert_is_not_tiled(sketches[0].stages[2])
- assert_compute_at_condition(sketches[0].stages[3], "inlined")
- assert_is_tiled(sketches[0].stages[4])
- assert_is_tiled(sketches[0].stages[6])
- assert_compute_at_condition(sketches[0].stages[7], "inlined")
- assert_is_tiled(sketches[0].stages[8])
- assert_is_not_tiled(sketches[0].stages[9])
- # Sketch 1
- assert_is_not_tiled(sketches[1].stages[1])
- assert_is_not_tiled(sketches[1].stages[2])
- assert_compute_at_condition(sketches[1].stages[3], "inlined")
- assert_is_tiled(sketches[1].stages[4])
- assert_is_tiled(sketches[1].stages[6])
- assert_has_cache_write(sketches[1], 6)
- assert_compute_at_condition(sketches[1].stages[6], "iter")
- assert_compute_at_condition(sketches[1].stages[8], "inlined")
- assert_is_tiled(sketches[1].stages[9])
- assert_is_not_tiled(sketches[1].stages[10])
- # Sketch 2
- assert_is_not_tiled(sketches[2].stages[1])
- assert_is_not_tiled(sketches[2].stages[2])
- assert_compute_at_condition(sketches[2].stages[3], "inlined")
- assert_is_tiled(sketches[2].stages[4])
- assert_is_tiled(sketches[2].stages[6])
- assert_has_cache_write(sketches[2], 6)
- assert_compute_at_condition(sketches[2].stages[6], "iter")
- assert_compute_at_condition(sketches[2].stages[8], "inlined")
- assert_is_tiled(sketches[2].stages[9])
- assert_is_not_tiled(sketches[2].stages[10])
- assert sketches[1] != sketches[2]
-
-
-def test_cpu_zero_rank_sketch():
- sketches = generate_sketches(zero_rank_reduce_auto_scheduler_test, (128,), "llvm")
- """ 2 rfactor sketches + 1 multi-level tiling sketches """
- assert len(sketches) == 3
-
-
-def test_cpu_custom_sketch():
- def meet_condition_func(search_policy, state, stage_id):
- return auto_scheduler.PreloadCustomSketchRule.APPLY_AND_SKIP_REST
-
- def apply_func(search_policy, state, stage_id):
- ret = []
- state = auto_scheduler.loop_state.State(state, search_policy.search_task.compute_dag)
- C = state.stage_ops[2]
-
- ret.append([state.state_object, -1])
-
- s1 = state.copy()
- i, _, _ = s1[C].iters
- s1.split(C, i, [8, 2])
- ret.append([s1.state_object, -1])
- return ret
-
- sketches = generate_sketches(
- matmul_auto_scheduler_test,
- (512, 512, 512),
- "llvm",
- init_search_callbacks=[
- auto_scheduler.PreloadCustomSketchRule(meet_condition_func, apply_func)
- ],
- )
- assert len(sketches) == 2
- assert sketches[0].stages[2].iters[0].range.extent == 512
- assert sketches[0].stages[2].iters[1].range.extent == 512
- assert sketches[0].stages[2].iters[2].range.extent == 512
- assert sketches[1].stages[2].iters[0].range.extent == 32
- assert sketches[1].stages[2].iters[1].range.extent == 8
- assert sketches[1].stages[2].iters[2].range.extent == 2
- assert sketches[1].stages[2].iters[3].range.extent == 512
- assert sketches[1].stages[2].iters[4].range.extent == 512
-
-
-@tvm.testing.requires_cuda
-def test_cuda_matmul_sketch():
- sketches = generate_sketches(matmul_auto_scheduler_test, (512, 512, 512), "cuda")
- """ 1 multi-level tiling sketch """
- assert len(sketches) == 1
- assert_has_cache_read(sketches[0], 0)
- assert_compute_at_condition(sketches[0].stages[1], "iter")
- assert_has_cache_read(sketches[0], 2)
- assert_compute_at_condition(sketches[0].stages[3], "iter")
- assert_has_cache_write(sketches[0], 4)
- assert_is_tiled(sketches[0].stages[4])
- assert_compute_at_condition(sketches[0].stages[4], "iter")
- assert_is_tiled(sketches[0].stages[5])
-
- sketches = generate_sketches(matmul_auto_scheduler_test, (8, 8, 1024), "cuda")
- """ 1 cross thread reuction sketch + 1 multi-level tiling sketch """
- assert len(sketches) == 2
- # Sketch 0
- assert_has_cross_thread_reduction(sketches[0], 2)
- # Sketch 1
- assert_has_cache_read(sketches[1], 0)
- assert_compute_at_condition(sketches[1].stages[1], "iter")
- assert_has_cache_read(sketches[1], 2)
- assert_compute_at_condition(sketches[1].stages[3], "iter")
- assert_has_cache_write(sketches[1], 4)
- assert_is_tiled(sketches[1].stages[4])
- assert_compute_at_condition(sketches[1].stages[4], "iter")
- assert_is_tiled(sketches[1].stages[5])
-
- sketches = generate_sketches(double_matmul_auto_scheduler_test, (512,), "cuda")
- """ 1 multi-level tiling sketch for one matmul, so 1 x 1 = 1 sketch in total """
- assert len(sketches) == 1
- assert_compute_at_condition(sketches[0].stages[5], "root")
- assert_compute_at_condition(sketches[0].stages[6], "iter")
-
-
-@tvm.testing.requires_cuda
-def test_cuda_conv2d_bn_relu_sketch():
- sketches = generate_sketches(
- conv2d_nchw_bn_relu_auto_scheduler_test, (1, 56, 56, 512, 512, 3, 1, 1), "cuda"
- )
- """ 1 multi-level tiling sketch """
- assert len(sketches) == 1
- assert_has_cache_read(sketches[0], 1)
- assert_compute_at_condition(sketches[0].stages[1], "inlined")
- assert_compute_at_condition(sketches[0].stages[2], "iter")
- assert_has_cache_read(sketches[0], 3)
- assert_compute_at_condition(sketches[0].stages[4], "iter")
- assert_is_tiled(sketches[0].stages[5])
- assert_compute_at_condition(sketches[0].stages[5], "iter")
- assert_compute_at_condition(sketches[0].stages[7], "inlined")
- assert_compute_at_condition(sketches[0].stages[9], "inlined")
- assert_compute_at_condition(sketches[0].stages[11], "inlined")
- assert_is_tiled(sketches[0].stages[12])
-
-
-@tvm.testing.requires_cuda
-def test_cuda_max_pool2d_sketch():
- sketches = generate_sketches(max_pool2d_auto_scheduler_test, (1, 56, 56, 512, 0), "cuda")
- """ 1 default sketch """
- assert len(sketches) == 1
- assert len(sketches[0].transform_steps) == 0
-
-
-@tvm.testing.requires_cuda
-def test_cuda_min_sketch():
- sketches = generate_sketches(min_nm_auto_scheduler_test, (10, 1024), "cuda")
- """ 1 cross thread reuction sketch + 1 default sketch """
- assert len(sketches) == 2
- # Sketch 0
- assert_has_cross_thread_reduction(sketches[0], 1)
- # Sketch 1
- assert len(sketches[1].transform_steps) == 0
-
-
-@tvm.testing.requires_cuda
-def test_cuda_softmax_sketch():
- sketches = generate_sketches(softmax_nm_auto_scheduler_test, (2, 1024), "cuda")
- """ (1 cross thread reuction sketch + 1 default sketch) * (1 cross thread reuction sketch + 1 default sketch) """
- assert len(sketches) == (2 * 2)
- # Sketch 0
- assert_has_cross_thread_reduction(sketches[0], 1)
- assert_compute_at_condition(sketches[3].stages[2], "inlined")
- assert_has_cross_thread_reduction(sketches[0], 3)
- # Sketch 1
- assert_compute_at_condition(sketches[3].stages[2], "inlined")
- assert_has_cross_thread_reduction(sketches[1], 3)
- # Sketch 2
- assert_has_cross_thread_reduction(sketches[2], 1)
- assert_compute_at_condition(sketches[3].stages[2], "inlined")
- # Sketch 3
- assert_compute_at_condition(sketches[3].stages[2], "inlined")
-
- sketches = generate_sketches(softmax_abcd_auto_scheduler_test, (1, 12, 128, 128), "cuda")
- """ (1 cross thread reuction sketch + 1 default sketch) * (1 cross thread reuction sketch + 1 default sketch) """
- assert len(sketches) == (2 * 2)
- # Sketch 0
- assert_has_cross_thread_reduction(sketches[0], 1)
- assert_compute_at_condition(sketches[3].stages[2], "inlined")
- assert_has_cross_thread_reduction(sketches[0], 3)
- # Sketch 1
- assert_compute_at_condition(sketches[3].stages[2], "inlined")
- assert_has_cross_thread_reduction(sketches[1], 3)
- # Sketch 2
- assert_has_cross_thread_reduction(sketches[2], 1)
- assert_compute_at_condition(sketches[3].stages[2], "inlined")
- # Sketch 3
- assert_compute_at_condition(sketches[3].stages[2], "inlined")
-
-
-@tvm.testing.requires_cuda
-def test_cuda_conv2d_winograd_sketch():
- sketches = generate_sketches(
- conv2d_winograd_nhwc_auto_scheduler_test, (1, 28, 28, 128, 128, 3, 1, 1), "cuda"
- )
- """ 1 multi-level tiling sketch """
- assert len(sketches) == 1
- assert_compute_at_condition(sketches[0].stages[1], "inlined")
- assert_compute_at_condition(sketches[0].stages[2], "iter")
- assert_compute_at_condition(sketches[0].stages[3], "inlined")
- assert_is_tiled(sketches[0].stages[4])
- assert_has_cache_read(sketches[0], 4)
- assert_compute_at_condition(sketches[0].stages[5], "iter")
- assert_has_cache_read(sketches[0], 6)
- assert_compute_at_condition(sketches[0].stages[7], "iter")
- assert_is_tiled(sketches[0].stages[8])
- assert_compute_at_condition(sketches[0].stages[8], "iter")
- assert_has_cache_write(sketches[0], 8)
- assert_compute_at_condition(sketches[0].stages[9], "root")
- assert_is_tiled(sketches[0].stages[11])
- assert_is_not_tiled(sketches[0].stages[12])
-
-
-@tvm.testing.requires_cuda
-def test_cuda_zero_rank_sketch():
- sketches = generate_sketches(zero_rank_reduce_auto_scheduler_test, (128,), "cuda")
- """ 1 cross thread reuction sketch + 1 multi-level tiling sketch """
- assert len(sketches) == 2
-
-
-if __name__ == "__main__":
- tvm.testing.main()
diff --git a/tests/python/auto_scheduler/test_auto_scheduler_task_scheduler.py b/tests/python/auto_scheduler/test_auto_scheduler_task_scheduler.py
deleted file mode 100644
index 66e1ba9d68..0000000000
--- a/tests/python/auto_scheduler/test_auto_scheduler_task_scheduler.py
+++ /dev/null
@@ -1,145 +0,0 @@
-# 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.
-""" Test task scheduler """
-
-import tempfile
-
-import multiprocessing
-import numpy as np
-
-import tvm
-import tvm.testing
-from tvm import auto_scheduler
-
-from tvm.testing.auto_scheduler import matmul_auto_scheduler_test
-
-
-@tvm.testing.requires_llvm
-def test_task_scheduler_round_robin():
- tasks = []
- for n in [2, 4, 8]:
- tasks.append(
- auto_scheduler.SearchTask(
- func=matmul_auto_scheduler_test, args=(n, n, n), target="llvm"
- )
- )
-
- with tempfile.NamedTemporaryFile() as fp:
- log_file = fp.name
- num_trials_per_task = 2
-
- # Tune all tasks
- measure_ctx = auto_scheduler.LocalRPCMeasureContext()
- tune_option = auto_scheduler.TuningOptions(
- num_measure_trials=num_trials_per_task * len(tasks),
- runner=measure_ctx.runner,
- num_measures_per_round=1,
- measure_callbacks=[auto_scheduler.RecordToFile(log_file)],
- )
- task_scheduler = auto_scheduler.TaskScheduler(tasks, strategy="round-robin", callbacks=[])
- task_scheduler.tune(tune_option, search_policy="sketch.random")
-
- # Check the result of round robin
- counters = {}
- for task in tasks:
- counters[task.workload_key] = 0
-
- for inp, _ in auto_scheduler.load_records(log_file):
- counters[inp.task.workload_key] += 1
-
- for task in tasks:
- assert counters[task.workload_key] == num_trials_per_task
-
- # test continuous tuning (restoring the status)
- task_scheduler = auto_scheduler.TaskScheduler(
- tasks, strategy="round-robin", load_log_file=log_file, callbacks=[]
- )
- tune_option = auto_scheduler.TuningOptions(
- num_measure_trials=len(tasks),
- num_measures_per_round=1,
- )
- task_scheduler.tune(tune_option, search_policy="sketch.random")
- del measure_ctx
-
-
-@tvm.testing.requires_llvm
-def task_scheduler_round_robin_spawn():
- assert multiprocessing.get_start_method(False) == "spawn"
- test_task_scheduler_round_robin()
-
-
-@tvm.testing.requires_llvm
-def test_task_scheduler_round_robin_spawn():
- ctx = multiprocessing.get_context("spawn")
- p = ctx.Process(target=task_scheduler_round_robin_spawn)
- p.start()
- p.join()
-
-
-@tvm.testing.requires_llvm
-def test_task_scheduler_gradient():
- tasks = []
- for n in [2, 4]:
- tasks.append(
- auto_scheduler.SearchTask(
- func=matmul_auto_scheduler_test, args=(n, n, n), target="llvm"
- )
- )
-
- def objective_func(costs):
- return 1e5 * costs[0]
-
- with tempfile.NamedTemporaryFile() as fp:
- log_file = fp.name
-
- n_trials = 5
-
- # Tune all tasks
- measure_ctx = auto_scheduler.LocalRPCMeasureContext()
- tune_option = auto_scheduler.TuningOptions(
- num_measure_trials=n_trials,
- runner=measure_ctx.runner,
- num_measures_per_round=1,
- measure_callbacks=[auto_scheduler.RecordToFile(log_file)],
- )
- task_scheduler = auto_scheduler.TaskScheduler(
- tasks, objective_func=objective_func, callbacks=[]
- )
-
- # Forcely rewrite the initial values.
- # This can make this test more stable on the slow CI machines
- task_scheduler.best_costs = np.array([1e2, 1e-8])
-
- task_scheduler.tune(tune_option, search_policy="sketch.random")
-
- # Check the allocation results
- counters = {}
- for task in tasks:
- counters[task.workload_key] = 0
-
- for inp, _ in auto_scheduler.load_records(log_file):
- counters[inp.task.workload_key] += 1
-
- assert counters[tasks[0].workload_key] == n_trials - 1
- assert counters[tasks[1].workload_key] == 1
- del measure_ctx
-
-
-if __name__ == "__main__":
- test_task_scheduler_round_robin()
- test_task_scheduler_round_robin_spawn()
- test_task_scheduler_gradient()
diff --git a/tests/python/autotvm/test_autotvm_database.py b/tests/python/autotvm/test_autotvm_database.py
deleted file mode 100644
index d598002281..0000000000
--- a/tests/python/autotvm/test_autotvm_database.py
+++ /dev/null
@@ -1,124 +0,0 @@
-# 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.
-"""Test database"""
-import copy
-import logging
-
-from tvm.autotvm import database
-from tvm.autotvm.record import encode, MeasureResult
-
-from tvm.testing.autotvm import get_sample_records
-
-
-def test_save_load():
- logging.info("test basic db load/save ...")
- records = get_sample_records(3)
- inp1, res1 = records[0]
- inp2, res2 = records[1]
- inp3, _ = records[2]
-
- _db = database.DummyDatabase()
- _db.flush()
- _db.save(inp1, res1)
- _db.save(inp2, res2)
-
- load1 = _db.load(inp1)
- load2 = _db.load(inp2)
- load3 = _db.load(inp3)
- assert load1 == res1
- assert load2 == res2
- assert load3 is None
- assert load1 != load2
-
-
-TRIAL_LIMIT = 2
-
-
-def test_db_hash():
- logging.info("test db hash check ...")
- inp1, res1 = get_sample_records(1)[0]
- inp2 = copy.deepcopy(inp1)
- inp1.config.code_hash = "cafecafe"
- inp2.config.code_hash = "dbffdbff"
- res2l = list(tuple(res1))
-
- # set timestamp
- res2l[-1] = -1
- res2 = MeasureResult(*res2l)
- _db = database.DummyDatabase()
- _db.flush()
- _db.save(inp1, res1, extend=True)
- _db.save(inp2, res2, extend=True)
-
- load1 = _db.load(inp1)
- load2 = _db.load(inp2)
- assert load1 != load2
- assert load1.timestamp != -1
- assert load2.timestamp == -1
-
-
-def test_db_latest_all():
- logging.info("test db load w/ multiple results ...")
- inp1, res1 = get_sample_records(1)[0]
- lis1 = list(tuple(res1))
- lis2 = list(tuple(res1))
- lis3 = list(tuple(res1))
-
- # set timestamp
- lis1[-1] = 0.0
- lis2[-1] = 1.1
- lis3[-1] = 9999.9999
- res1 = MeasureResult(*lis1)
- res2 = MeasureResult(*lis2)
- res3 = MeasureResult(*lis3)
-
- _db = database.DummyDatabase()
- _db.flush()
- _db.save(inp1, res1, extend=True)
- load1 = _db.load(inp1)
- assert load1.timestamp == 0.0
- _db.save(inp1, res2, extend=True)
- load2 = _db.load(inp1)
- assert load2.timestamp == 1.1
- _db.save(inp1, res3, extend=True)
- load3 = _db.load(inp1)
- assert load3.timestamp == 9999.9999
-
- load4 = _db.load(inp1, get_all=True)
- assert encode(inp1, load4[0]) == encode(inp1, res1)
- assert encode(inp1, load4[1]) == encode(inp1, res2)
- assert encode(inp1, load4[2]) == encode(inp1, res3)
-
-
-def test_db_filter():
- logging.info("test db filter ...")
- records = get_sample_records(5)
- _db = database.DummyDatabase()
- _db.flush()
- for inp, result in records:
- _db.save(inp, result)
-
- records = _db.filter(lambda inp, ress: any(r.costs[0] <= 2 for r in ress))
- assert len(records) == 2
-
-
-if __name__ == "__main__":
- logging.basicConfig(level=logging.INFO)
- test_save_load()
- test_db_hash()
- test_db_latest_all()
- test_db_filter()
diff --git a/tests/python/autotvm/test_autotvm_dispatch_context.py b/tests/python/autotvm/test_autotvm_dispatch_context.py
deleted file mode 100644
index ba75992128..0000000000
--- a/tests/python/autotvm/test_autotvm_dispatch_context.py
+++ /dev/null
@@ -1,51 +0,0 @@
-# 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.
-"""Test dispatcher.
-The dispatcher can choose which template to use according
-to the parameters of workload"""
-
-from tvm import autotvm
-import tvm
-
-
-@autotvm.template("testing/dispatch_fallback")
-def simple_template(a, b):
- cfg = autotvm.get_config()
- assert cfg.is_fallback
-
-
-def test_fallback():
- simple_template(2, 3)
-
-
-def test_tophub_kinds_match():
- def verify_arm_cpu(target):
- best_by_targetkey = autotvm.tophub.context(target).best_by_targetkey
- assert len(best_by_targetkey)
- found_arm_cpu = False
- for a, _ in best_by_targetkey:
- if "arm_cpu" in a:
- found_arm_cpu = True
- break
- assert found_arm_cpu
-
- verify_arm_cpu("llvm -device=arm_cpu -mtriple=aarch64-linux-gnu -mattr=+neon,+v8.2a,+dotprod")
- verify_arm_cpu("llvm -model=snapdragon835 -mtriple=arm64-linux-android -mattr=+neon")
-
-
-if __name__ == "__main__":
- test_fallback()
diff --git a/tests/python/autotvm/test_autotvm_droplet_tuner.py b/tests/python/autotvm/test_autotvm_droplet_tuner.py
deleted file mode 100644
index 10ffa18784..0000000000
--- a/tests/python/autotvm/test_autotvm_droplet_tuner.py
+++ /dev/null
@@ -1,57 +0,0 @@
-# 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.
-"""Test droplet algorithm tuner"""
-
-from tvm.testing.autotvm import DummyRunner, get_sample_task, get_sample_records
-from tvm import autotvm
-
-
-def test_tuner():
- """Test Droplet Tuner"""
-
- task, _ = get_sample_task()
- measure_option = autotvm.measure_option(builder=autotvm.LocalBuilder(), runner=DummyRunner())
-
- # When no range index, range_length should be the length of config space
- tuner = autotvm.tuner.DropletTuner(task)
- assert len(tuner.best_choice) == 3
- assert tuner.execution == 1
- assert tuner.batch == 16
- assert tuner.total_execution == max(tuner.dims)
- assert tuner.step == 1
-
-
-def test_multi_filter():
- # Test with multi-filter
- task, _ = get_sample_task()
- task.config_space.multi_filter(
- filter=lambda entity: 0 <= (entity["tile_x"].size[1] * entity["tile_y"].size[1]) < 1024
- )
- tuner = autotvm.tuner.DropletTuner(task)
- valid_indexes = list(
- filter(lambda idx: tuner.space.is_index_valid(idx), range(tuner.space.range_length))
- )
- assert tuner.visited.issubset(valid_indexes)
-
- measure_option = autotvm.measure_option(builder=autotvm.LocalBuilder(), runner=DummyRunner())
- tuner.tune(n_trial=8, measure_option=measure_option)
- assert tuner.visited.issubset(valid_indexes)
-
-
-if __name__ == "__main__":
- test_tuner()
- test_multi_filter()
diff --git a/tests/python/autotvm/test_autotvm_feature.py b/tests/python/autotvm/test_autotvm_feature.py
deleted file mode 100644
index 26268e5599..0000000000
--- a/tests/python/autotvm/test_autotvm_feature.py
+++ /dev/null
@@ -1,136 +0,0 @@
-# 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.
-"""Test feature extraction"""
-
-import numpy as np
-
-import tvm
-from tvm import te
-from tvm.autotvm import feature
-
-
-def test_iter_feature_gemm():
- N = 128
-
- k = te.reduce_axis((0, N), "k")
- A = te.placeholder((N, N), name="A")
- B = te.placeholder((N, N), name="B")
- C = te.compute(A.shape, lambda y, x: te.sum(A[y, k] * B[k, x], axis=k), name="C")
-
- s = te.create_schedule(C.op)
-
- feas = feature.get_itervar_feature(s, [A, B, C], take_log=False)
-
- expected = [
- {
- "_attr_": [128, 1, 128, 2097152, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0],
- "A_0": [128, -1, 16384, 128, 0, 0],
- "B_0": [0, -1, 16384, 128, 0, 0],
- "C_0": [128, -1, 16384, 128, 0, 0],
- "C_1": [128, -1, 16384, 128, 0, 0],
- },
- {
- "_attr_": [128, 2, 16384, 16384, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0],
- "A_0": [0, -1, 128, 128, 0, 0],
- "B_0": [1, -1, 16384, 1, 0, 0],
- "C_0": [1, -1, 128, 128, 0, 0],
- "C_1": [1, -1, 128, 128, 0, 0],
- },
- {
- "_attr_": [128, 3, 2097152, 128, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0],
- "A_0": [1, -1, 128, 1, 0, 0],
- "B_0": [128, -1, 128, 1, 0, 0],
- "C_1": [0, -1, 1, 128, 0, 0],
- "C_2": [0, -1, 1, 128, 0, 0],
- },
- ]
-
- for ans, row in zip(expected, feas):
- for pair in row:
- if pair[0] not in ans:
- continue
- assert ans[pair[0]] == pair[1:], "%s: %s vs %s" % (pair[0], ans[pair[0]], pair[1:])
-
-
-def test_curve_feature_gemm():
- N = 128
-
- k = te.reduce_axis((0, N), "k")
- A = te.placeholder((N, N), name="A")
- B = te.placeholder((N, N), name="B")
- C = te.compute(A.shape, lambda y, x: te.sum(A[y, k] * B[k, x], axis=k), name="C")
-
- s = te.create_schedule(C.op)
-
- feas = feature.get_buffer_curve_sample_flatten(s, [A, B, C], sample_n=30)
- # sample_n * #buffers * #curves * 2 numbers per curve
- assert len(feas) == 30 * 3 * 4 * 2
-
-
-def test_feature_shape():
- """test the dimensions of flatten feature are the same"""
-
- N = 1024
- n_sample = 100
-
- def get_gemm_feature(target):
- k = te.reduce_axis((0, N), "k")
- A = te.placeholder((N, N), name="A")
- B = te.placeholder((N, N), name="B")
- C = te.compute(A.shape, lambda y, x: te.sum(A[y, k] * B[k, x], axis=k), name="C")
-
- s = te.create_schedule(C.op)
-
- y, x = s[C].op.axis
- axes = list(s[C].tile(y, x, 8, 8)) + [k]
- perm = np.random.permutation(5)
- axes = [axes[x] for x in perm]
- s[C].reorder(*axes)
-
- if "gpu" in target.keys:
- pick = []
- # filter out reduction axis
- for i in range(len(perm)):
- if perm[i] != 4:
- pick.append(axes[i])
- s[C].bind(pick[0], te.thread_axis("blockIdx.x"))
- s[C].bind(pick[1], te.thread_axis("vthread"))
- s[C].bind(pick[2], te.thread_axis("threadIdx.y"))
-
- with target:
- feas = feature.get_itervar_feature(s, [A, B, C])
- feas = feature.flatten_itervar_feature(feas)
- return feas
-
- targets = [
- tvm.target.cuda(),
- tvm.target.mali(),
- tvm.target.arm_cpu(),
- ]
-
- for target in targets:
- dim = len(get_gemm_feature(target))
- for i in range(n_sample):
- assert dim == len(get_gemm_feature(target)), (
- "dimensions of feature do not match" " for different configurations"
- )
-
-
-if __name__ == "__main__":
- test_iter_feature_gemm()
- test_curve_feature_gemm()
- test_feature_shape()
diff --git a/tests/python/autotvm/test_autotvm_flop_calculator.py b/tests/python/autotvm/test_autotvm_flop_calculator.py
deleted file mode 100644
index e28beaf987..0000000000
--- a/tests/python/autotvm/test_autotvm_flop_calculator.py
+++ /dev/null
@@ -1,173 +0,0 @@
-# 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.
-"""Test flop calculation"""
-
-import tvm
-from tvm import te
-import numpy as np
-
-from tvm.autotvm.task.task import compute_flop
-
-
-def random_dtypes():
- """Return pair of (input, accumulator) dtypes"""
- candidates = [("float32", "float32"), ("float16", "float32"), ("int8", "int32")]
- return candidates[np.random.choice(len(candidates))]
-
-
-def test_conv():
- for i in range(5):
- N, H, W, CO, CI, KH, KW = [np.random.randint(10, 32) for _ in range(7)]
- (input_dtype, acc_dtype) = random_dtypes()
- D = te.placeholder((N, CI, H, W), dtype=input_dtype)
- K = te.placeholder((CO, CI, KH, KW), dtype=input_dtype)
-
- KH = min(H, KH)
- KW = min(W, KW)
-
- ci = te.reduce_axis((0, CI))
- kh = te.reduce_axis((0, KH))
- kw = te.reduce_axis((0, KW))
-
- OH = (H - KH) + 1
- OW = (W - KW) + 1
-
- C = te.compute(
- (N, CO, OH, OW),
- lambda n, co, h, w: te.sum(
- D[n][ci][h][w].astype(acc_dtype) * K[co][ci][h][w].astype(acc_dtype),
- axis=[ci, kh, kw],
- ),
- )
-
- s = te.create_schedule([C.op])
-
- assert compute_flop(s) == 2 * N * CO * OH * OW * CI * KH * KW
-
-
-def test_pack_gemm():
- for i in range(5):
- N, L, M = [np.random.randint(10, 128) * 4 for _ in range(3)]
- (input_dtype, acc_dtype) = random_dtypes()
- A = te.placeholder((N, L), dtype=input_dtype)
- B = te.placeholder((M, L), dtype=input_dtype)
- k = te.reduce_axis((0, L))
-
- bn = 4
- idxd = tvm.tir.indexdiv
- idxm = tvm.tir.indexmod
-
- A_pack = te.compute((N // bn, L, bn), lambda i, j, k: A[i * bn + k][j])
- B_pack = te.compute((M // bn, L, bn), lambda i, j, k: B[i * bn + k][j])
- C_pack = te.compute(
- (N // bn, M // bn, bn, bn),
- lambda i, j, ii, jj: te.sum(
- A_pack[i, k, ii].astype(acc_dtype) * B_pack[j, k, jj].astype(acc_dtype), axis=[k]
- ),
- )
- C = te.compute(
- (N, M), lambda i, j: C_pack[idxd(i, bn)][idxd(j, bn)][idxm(i, bn)][idxm(j, bn)]
- )
-
- s = te.create_schedule([C.op])
- assert compute_flop(s) == 2 * N * L * M
-
-
-def test_outer_dot():
- for i in range(5):
- N, M = [np.random.randint(10, 128) * 4 for _ in range(2)]
- (input_dtype, acc_dtype) = random_dtypes()
- A = te.placeholder((N,), dtype=input_dtype)
- B = te.placeholder((M,), dtype=input_dtype)
-
- C = te.compute((N, M), lambda i, j: A[i].astype(acc_dtype) * B[j].astype(acc_dtype))
-
- s = te.create_schedule([C.op])
- assert compute_flop(s) == N * M
-
-
-def test_max_pool():
- for i in range(5):
- N, H, W, CO, CI, KH, KW = [np.random.randint(10, 32) for _ in range(7)]
- (input_dtype, _) = random_dtypes()
- D = te.placeholder((N, CI, H, W), dtype=input_dtype)
-
- KH = min(H, KH)
- KW = min(W, KW)
-
- kh = te.reduce_axis((0, KH))
- kw = te.reduce_axis((0, KW))
-
- OH = (H - KH) + 1
- OW = (W - KW) + 1
-
- C = te.compute(
- (N, CO, OH, OW), lambda n, co, h, w: tvm.te.max(D[n][co][h + kh][w + kw], axis=[kh, kw])
- )
-
- s = te.create_schedule([C.op])
-
- assert compute_flop(s) == N * CO * OH * OW * KH * KW
-
-
-def test_average_pool():
- for i in range(5):
- N, H, W, CO, CI, KH, KW = [np.random.randint(10, 32) for _ in range(7)]
- (input_dtype, acc_dtype) = random_dtypes()
- D = te.placeholder((N, CI, H, W), dtype=input_dtype)
-
- KH = min(H, KH)
- KW = min(W, KW)
-
- kh = te.reduce_axis((0, KH))
- kw = te.reduce_axis((0, KW))
-
- OH = (H - KH) + 1
- OW = (W - KW) + 1
-
- C = te.compute(
- (N, CO, OH, OW),
- lambda n, co, h, w: te.sum(
- te.div(D[n][co][h + kh][w + kw].astype(acc_dtype), (KW * KH)), axis=[kh, kw]
- ),
- )
-
- s = te.create_schedule([C.op])
-
- assert compute_flop(s) == 2 * N * CO * OH * OW * KH * KW
-
-
-def test_move():
- """No float number operation in simple move. So the estimator should raise an error"""
- N = 1024
-
- A = te.placeholder((N,))
- C = te.compute((N,), lambda i: A[i])
- s = te.create_schedule([C.op])
-
- try:
- compute_flop(s)
- assert False
- except RuntimeError:
- pass
-
-
-if __name__ == "__main__":
- test_conv()
- test_pack_gemm()
- test_outer_dot()
- test_move()
diff --git a/tests/python/autotvm/test_autotvm_ga_tuner.py b/tests/python/autotvm/test_autotvm_ga_tuner.py
deleted file mode 100644
index 625c6c66b6..0000000000
--- a/tests/python/autotvm/test_autotvm_ga_tuner.py
+++ /dev/null
@@ -1,89 +0,0 @@
-# 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.
-"""Test genetic algorithm tuner"""
-
-from tvm.testing.autotvm import DummyRunner, get_sample_task
-from tvm import autotvm
-
-
-def test_ga_tuner():
- """Test GATuner"""
- # Test population size smaller than space size tuning configuration
- task, _ = get_sample_task()
- tuner = autotvm.tuner.GATuner(task, pop_size=32)
- valid_indexes = list(
- filter(lambda idx: tuner.space.is_index_valid(idx), range(tuner.space.range_length))
- )
- assert tuner.visited.issubset(valid_indexes)
- assert tuner.pop_size == len(tuner.visited) == len(tuner.genes)
- assert len(tuner.space) == 64
-
- measure_option = autotvm.measure_option(builder=autotvm.LocalBuilder(), runner=DummyRunner())
- tuner.tune(n_trial=len(tuner.space), measure_option=measure_option)
- assert tuner.visited.issubset(valid_indexes)
-
- # Test population size bigger than space size tuning configuration
- task, _ = get_sample_task()
- tuner = autotvm.tuner.GATuner(task, pop_size=100)
- valid_indexes = list(
- filter(lambda idx: tuner.space.is_index_valid(idx), range(tuner.space.range_length))
- )
- assert tuner.visited.issubset(valid_indexes)
- assert tuner.pop_size == len(tuner.visited) == len(tuner.genes)
- assert len(tuner.space) == 64
-
- measure_option = autotvm.measure_option(builder=autotvm.LocalBuilder(), runner=DummyRunner())
- tuner.tune(n_trial=len(tuner.space), measure_option=measure_option)
- assert tuner.visited.issubset(valid_indexes)
-
- # Test population size smaller than multi-filtered space size tuning configuration
- task, _ = get_sample_task()
- task.config_space.multi_filter(
- filter=lambda entity: 8 <= (entity["tile_x"].size[1] * entity["tile_y"].size[1]) < 1024
- )
- tuner = autotvm.tuner.GATuner(task, pop_size=32)
- valid_indexes = list(
- filter(lambda idx: tuner.space.is_index_valid(idx), range(tuner.space.range_length))
- )
- assert tuner.visited.issubset(valid_indexes)
- assert tuner.pop_size == len(tuner.visited) == len(tuner.genes)
- assert len(tuner.space) == 43
-
- measure_option = autotvm.measure_option(builder=autotvm.LocalBuilder(), runner=DummyRunner())
- tuner.tune(n_trial=len(tuner.space), measure_option=measure_option)
- assert tuner.visited.issubset(valid_indexes)
-
- # Test population size bigger than multi-filtered space size tuning configuration
- task, _ = get_sample_task()
- task.config_space.multi_filter(
- filter=lambda entity: 8 <= (entity["tile_x"].size[1] * entity["tile_y"].size[1]) < 1024
- )
- tuner = autotvm.tuner.GATuner(task, pop_size=100)
- valid_indexes = list(
- filter(lambda idx: tuner.space.is_index_valid(idx), range(tuner.space.range_length))
- )
- assert tuner.visited.issubset(valid_indexes)
- assert tuner.pop_size == len(tuner.visited) == len(tuner.genes)
- assert len(tuner.space) == 43
-
- measure_option = autotvm.measure_option(builder=autotvm.LocalBuilder(), runner=DummyRunner())
- tuner.tune(n_trial=len(tuner.space), measure_option=measure_option)
- assert tuner.visited.issubset(valid_indexes)
-
-
-if __name__ == "__main__":
- test_ga_tuner()
diff --git a/tests/python/autotvm/test_autotvm_graph_tuner_core.py b/tests/python/autotvm/test_autotvm_graph_tuner_core.py
deleted file mode 100644
index e1aff87241..0000000000
--- a/tests/python/autotvm/test_autotvm_graph_tuner_core.py
+++ /dev/null
@@ -1,792 +0,0 @@
-# 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.
-
-# NOTE: We name this test file to start with test_graph_tuner
-# to make it execute after zero_rank tensor test cases. This
-# helps avoid topi arithmetic operator overloading issue:
-# https://github.com/apache/tvm/issues/3240.
-# TODO: restore the file name after this issue is resolved.
-import os
-import copy
-import numpy as np
-import tvm
-from tvm import te
-import tvm.relay.testing
-
-from tvm import autotvm
-from tvm import relay
-from tvm.autotvm.task import ConfigEntity
-from tvm.autotvm.measure import MeasureResult, MeasureInput
-from tvm.autotvm.graph_tuner import DPTuner, PBQPTuner
-
-
-def _create_args(dshape, kshape, strides, padding, dilation, layout, out_layout, dtype, out_dtype):
- data = tvm.te.placeholder(dshape, dtype=dtype)
- kernel = tvm.te.placeholder(kshape, dtype=dtype)
- return autotvm.task.serialize_args(
- [data, kernel, strides, padding, dilation, layout, layout, out_dtype]
- )
-
-
-def _create_data(target, dshape, dtype, layout):
- data = relay.var("data", shape=dshape, dtype=dtype)
- w0 = relay.var("w0_weight")
- conv0 = relay.nn.conv2d(data, w0, channels=16, kernel_size=(3, 3), padding=(1, 1))
- w1 = relay.var("w1_weight")
- conv1 = relay.nn.conv2d(conv0, w1, channels=32, kernel_size=(1, 1))
- w2 = relay.var("w2_weight")
- conv2 = relay.nn.conv2d(conv1, w2, channels=32, kernel_size=(3, 3), padding=(1, 1))
- out = relay.add(conv1, conv2)
- net = relay.Function(relay.analysis.free_vars(out), out)
- mod, params = relay.testing.create_workload(net)
- tasks = autotvm.task.extract_from_program(
- mod["main"], target=target, params=params, ops=(relay.op.get("nn.conv2d"),)
- )
- new_args = [
- _create_args(
- (1, 3, 8, 8), (16, 3, 3, 3), (1, 1), (1, 1, 1, 1), (1, 1), layout, layout, dtype, dtype
- ),
- _create_args(
- (1, 16, 8, 8),
- (32, 16, 1, 1),
- (1, 1),
- (0, 0, 0, 0),
- (1, 1),
- layout,
- layout,
- dtype,
- dtype,
- ),
- _create_args(
- (1, 32, 8, 8),
- (32, 32, 3, 3),
- (1, 1),
- (1, 1, 1, 1),
- (1, 1),
- layout,
- layout,
- dtype,
- dtype,
- ),
- ]
-
- costs = [0.04, 0.012, 0.03]
- config_list = []
- cfg_dict = {
- "index": -1,
- "code_hash": None,
- "entity": [
- ["tile_ic", "sp", [3, 1]],
- ["tile_oc", "sp", [4, 4]],
- ["tile_ow", "sp", [4, 2]],
- ["unroll_kw", "ot", True],
- ],
- }
- config_list.append(ConfigEntity.from_json_dict(cfg_dict))
- cfg_dict = {
- "index": -1,
- "code_hash": None,
- "entity": [
- ["tile_ic", "sp", [2, 8]],
- ["tile_oc", "sp", [1, 32]],
- ["tile_oh", "ot", 1],
- ["tile_ow", "sp", [4, 2]],
- ],
- }
- config_list.append(ConfigEntity.from_json_dict(cfg_dict))
- cfg_dict = {
- "index": -1,
- "code_hash": None,
- "entity": [
- ["tile_ic", "sp", [8, 4]],
- ["tile_oc", "sp", [4, 8]],
- ["tile_ow", "sp", [2, 4]],
- ["unroll_kw", "ot", False],
- ],
- }
- config_list.append(ConfigEntity.from_json_dict(cfg_dict))
-
- records = []
- for args, cost, config, task in zip(new_args, costs, config_list, tasks):
- task.args = args
- ms_input = MeasureInput(target=target, task=task, config=config)
- ms_output = MeasureResult(costs=(cost,), error_no=0, all_cost=-1, timestamp=-1)
- records.append((ms_input, ms_output))
-
- ltf_records = []
- ltf_arg = [te.placeholder((1, 64, 16, 16, 8), dtype=dtype), "NCHW8c", "NCHW512c"]
- ltf_task = autotvm.task.create("layout_transform", ltf_arg, target)
- ms_input = MeasureInput(target=target, task=ltf_task, config=None)
- ms_output = MeasureResult(costs=(1.91224744e-05,), error_no=0, all_cost=-1, timestamp=-1)
- ltf_records.append((ms_input, ms_output))
-
- ltf_keys = []
- ltf_arg = [te.placeholder((1, 4, 8, 8, 4), dtype=dtype), "NCHW4c", "NCHW8c"]
- ltf_wkl = autotvm.task.args_to_workload(ltf_arg, "layout_transform")
- ltf_keys.append(ltf_wkl)
- ltf_arg = [te.placeholder((1, 1, 8, 8, 32), dtype=dtype), "NCHW32c", "NCHW4c"]
- ltf_wkl = autotvm.task.args_to_workload(ltf_arg, "layout_transform")
- ltf_keys.append(ltf_wkl)
- ltf_arg = [te.placeholder((1, 4, 8, 8, 8), dtype=dtype), "NCHW8c", "NCHW32c"]
- ltf_wkl = autotvm.task.args_to_workload(ltf_arg, "layout_transform")
- ltf_keys.append(ltf_wkl)
-
- return net, records, ltf_records, ltf_keys, tasks
-
-
-@tvm.testing.requires_x86
-def test_graph_tuner_layout_transform():
- log_file = "%s/test_tuner.log" % (os.getcwd())
- target = "llvm"
- dshape = (1, 3, 8, 8)
- dtype = "float32"
- layout = "NCHW"
- conv2d = relay.op.get("nn.conv2d")
- target_ops = [conv2d]
-
- g, records, ltf_records, ltf_keys, _ = _create_data(target, dshape, dtype, layout)
- executor = DPTuner(g, {"data": dshape}, records, target_ops, target=target, log_file=log_file)
- executor.benchmark_layout_transform(layout_records=ltf_records, infer_layout=True)
- out = executor._layout_transform_perf_records
-
- num_flops = 0
- total_time = 0
- for record in ltf_records:
- ltf_wkl = record[0].task.workload
- input_shape = ltf_wkl[1][1]
- flops = np.prod(input_shape)
- num_flops += flops
- total_time += record[1].costs[0]
- avg_time = total_time / num_flops
-
- for ltf_workload in out:
- input_shape = ltf_workload[1][1]
- flops = 1
- for i in input_shape:
- flops *= i
- expected_time = flops * avg_time
- out_time = out[ltf_workload][1].costs[0]
- assert (
- expected_time == out_time
- ), "Inferred layout transformation time mismatch for %s: " "expecting %f but got %f" % (
- str(ltf_workload),
- expected_time,
- out_time,
- )
-
-
-@tvm.testing.requires_x86
-def test_graph_tuner_layout_transform_runner():
- log_file = "%s/test_tuner.log" % (os.getcwd())
- target = "llvm"
- dshape = (1, 3, 8, 8)
- dtype = "float32"
- layout = "NCHW"
- conv2d = relay.op.get("nn.conv2d")
- target_ops = [conv2d]
-
- g, records, ltf_records, ltf_keys, _ = _create_data(target, dshape, dtype, layout)
- executor = DPTuner(g, {"data": dshape}, records, target_ops, target=target, log_file=log_file)
- runner = autotvm.LocalRunner(number=100, repeat=1, timeout=10)
- executor.benchmark_layout_transform(
- layout_records=ltf_records, infer_layout=True, runner=runner
- )
- out = executor._layout_transform_perf_records
-
- num_flops = 0
- total_time = 0
- for record in ltf_records:
- ltf_wkl = record[0].task.workload
- input_shape = ltf_wkl[1][1]
- flops = np.prod(input_shape)
- num_flops += flops
- total_time += record[1].costs[0]
- avg_time = total_time / num_flops
-
- for ltf_workload in out:
- input_shape = ltf_workload[1][1]
- flops = 1
- for i in input_shape:
- flops *= i
- expected_time = flops * avg_time
- out_time = out[ltf_workload][1].costs[0]
- assert (
- expected_time == out_time
- ), "Inferred layout transformation time mismatch for %s: " "expecting %f but got %f" % (
- str(ltf_workload),
- expected_time,
- out_time,
- )
-
-
-@tvm.testing.requires_x86
-def test_DPTuner_run():
- log_file = "%s/test_tuner.log" % (os.getcwd())
- target = "llvm"
- dtype = "float32"
- layout = "NCHW"
- dshape = (1, 3, 8, 8)
- conv2d = relay.op.get("nn.conv2d")
- target_ops = [conv2d]
-
- g, records, ltf_records, ltf_keys, tasks = _create_data(target, dshape, dtype, layout)
- mod = tvm.IRModule()
- mod["main"] = g
- costs = [0.02, 0.02, 0.045]
- config_list = []
- cfg_dict = {
- "index": -1,
- "code_hash": None,
- "entity": [
- ["tile_ic", "sp", [1, 3]],
- ["tile_oc", "sp", [2, 8]],
- ["tile_ow", "sp", [4, 2]],
- ["unroll_kw", "ot", True],
- ],
- }
- config_list.append(ConfigEntity.from_json_dict(cfg_dict))
- cfg_dict = {
- "index": -1,
- "code_hash": None,
- "entity": [
- ["tile_ic", "sp", [4, 4]],
- ["tile_oc", "sp", [2, 16]],
- ["tile_oh", "ot", 1],
- ["tile_ow", "sp", [4, 2]],
- ],
- }
- config_list.append(ConfigEntity.from_json_dict(cfg_dict))
- cfg_dict = {
- "index": -1,
- "code_hash": None,
- "entity": [
- ["tile_ic", "sp", [16, 2]],
- ["tile_oc", "sp", [8, 4]],
- ["tile_ow", "sp", [2, 4]],
- ["unroll_kw", "ot", False],
- ],
- }
- config_list.append(ConfigEntity.from_json_dict(cfg_dict))
- for cost, config, task in zip(costs, config_list, tasks):
- ms_input = MeasureInput(target=target, task=task, config=config)
- ms_output = MeasureResult(costs=(cost,), error_no=0, all_cost=-1, timestamp=-1)
- records.append((ms_input, ms_output))
-
- executor = DPTuner(mod, {"data": dshape}, records, target_ops, target, log_file=log_file)
- executor.benchmark_layout_transform(layout_records=ltf_records, infer_layout=True)
- executor.run()
- out = [record[0].config for record in executor.get_optimal_records()]
- expected_out = [records[3][0].config, records[1][0].config, records[2][0].config]
- assert expected_out == out, "Output mismatch: expecting %s but got %s" % (
- str(expected_out),
- str(out),
- )
- assert os.path.isfile(log_file), "No log file with name %s exists." % log_file
-
-
-@tvm.testing.requires_x86
-def test_PBQPTuner_run():
- target = "llvm"
- dtype = "float32"
- layout = "NCHW"
- dshape = (1, 3, 8, 8)
- conv2d = relay.op.get("nn.conv2d")
- target_ops = [conv2d]
-
- g, records, ltf_records, ltf_keys, tasks = _create_data(target, dshape, dtype, layout)
- costs = [0.02, 0.02, 0.045]
- config_list = []
- cfg_dict = {
- "index": -1,
- "code_hash": None,
- "entity": [
- ["tile_ic", "sp", [1, 3]],
- ["tile_oc", "sp", [2, 8]],
- ["tile_ow", "sp", [4, 2]],
- ["unroll_kw", "ot", True],
- ],
- }
- config_list.append(ConfigEntity.from_json_dict(cfg_dict))
- cfg_dict = {
- "index": -1,
- "code_hash": None,
- "entity": [
- ["tile_ic", "sp", [4, 4]],
- ["tile_oc", "sp", [2, 16]],
- ["tile_oh", "ot", 1],
- ["tile_ow", "sp", [4, 2]],
- ],
- }
- config_list.append(ConfigEntity.from_json_dict(cfg_dict))
- cfg_dict = {
- "index": -1,
- "code_hash": None,
- "entity": [
- ["tile_ic", "sp", [16, 2]],
- ["tile_oc", "sp", [8, 4]],
- ["tile_ow", "sp", [2, 4]],
- ["unroll_kw", "ot", False],
- ],
- }
- config_list.append(ConfigEntity.from_json_dict(cfg_dict))
- for cost, config, task in zip(costs, config_list, tasks):
- ms_input = MeasureInput(target=target, task=task, config=config)
- ms_output = MeasureResult(costs=(cost,), error_no=0, all_cost=-1, timestamp=-1)
- records.append((ms_input, ms_output))
-
- executor = PBQPTuner(g, {"data": dshape}, records, target_ops, target)
- executor.benchmark_layout_transform(layout_records=ltf_records, infer_layout=True)
- executor.run()
- out = [record[0].config for record in executor.get_optimal_records()]
- expected_out = [records[3][0].config, records[1][0].config, records[2][0].config]
- assert expected_out == out, "Output mismatch: expecting %s but got %s" % (
- str(expected_out),
- str(out),
- )
-
-
-@tvm.testing.requires_x86
-def test_many_sub_graphs():
- target = "llvm"
- dtype = "float32"
- dshape = (1, 8, 8, 3)
- layout = "NCHW"
- conv2d = relay.op.get("nn.conv2d")
- target_ops = [conv2d]
-
- data = relay.var("data", shape=dshape, dtype=dtype)
- t0 = relay.transpose(data, (0, 3, 1, 2))
- w0 = relay.var("w0_weight")
- conv0 = relay.nn.conv2d(t0, w0, channels=16, kernel_size=(3, 3), padding=(1, 1))
- t1 = relay.transpose(conv0, (0, 2, 3, 1))
- w1 = relay.var("w1_weight")
- t2 = relay.transpose(t1, (0, 3, 1, 2))
- conv1 = relay.nn.conv2d(t2, w1, channels=32, kernel_size=(1, 1))
- t3 = relay.transpose(conv1, (0, 2, 3, 1))
- w2 = relay.var("w2_weight")
- t4 = relay.transpose(t3, (0, 3, 1, 2))
- conv2 = relay.nn.conv2d(t4, w2, channels=32, kernel_size=(3, 3), padding=(1, 1))
- t5 = relay.transpose(conv2, (0, 2, 3, 1))
- out = relay.add(t3, t5)
- net = relay.Function(relay.analysis.free_vars(out), out)
- net, params = relay.testing.create_workload(net)
-
- tasks = autotvm.task.extract_from_program(
- net["main"], target=target, params=params, ops=(conv2d,)
- )
- new_args = [
- _create_args(
- (1, 3, 8, 8), (16, 3, 3, 3), (1, 1), (1, 1, 1, 1), (1, 1), layout, layout, dtype, dtype
- ),
- _create_args(
- (1, 16, 8, 8),
- (32, 16, 1, 1),
- (1, 1),
- (0, 0, 0, 0),
- (1, 1),
- layout,
- layout,
- dtype,
- dtype,
- ),
- _create_args(
- (1, 32, 8, 8),
- (32, 32, 3, 3),
- (1, 1),
- (1, 1, 1, 1),
- (1, 1),
- layout,
- layout,
- dtype,
- dtype,
- ),
- ]
-
- costs = [0.04, 0.012, 0.03, 0.02, 0.02, 0.045]
- config_list = []
- cfg_dict = {
- "index": -1,
- "code_hash": None,
- "entity": [
- ["tile_ic", "sp", [3, 1]],
- ["tile_oc", "sp", [4, 4]],
- ["tile_ow", "sp", [4, 2]],
- ["unroll_kw", "ot", True],
- ],
- }
- config_list.append(ConfigEntity.from_json_dict(cfg_dict))
- cfg_dict = {
- "index": -1,
- "code_hash": None,
- "entity": [
- ["tile_ic", "sp", [2, 8]],
- ["tile_oc", "sp", [1, 32]],
- ["tile_oh", "ot", 1],
- ["tile_ow", "sp", [4, 2]],
- ],
- }
- config_list.append(ConfigEntity.from_json_dict(cfg_dict))
- cfg_dict = {
- "index": -1,
- "code_hash": None,
- "entity": [
- ["tile_ic", "sp", [8, 4]],
- ["tile_oc", "sp", [4, 8]],
- ["tile_ow", "sp", [2, 4]],
- ["unroll_kw", "ot", False],
- ],
- }
- config_list.append(ConfigEntity.from_json_dict(cfg_dict))
- cfg_dict = {
- "index": -1,
- "code_hash": None,
- "entity": [
- ["tile_ic", "sp", [1, 3]],
- ["tile_oc", "sp", [2, 8]],
- ["tile_ow", "sp", [4, 2]],
- ["unroll_kw", "ot", True],
- ],
- }
- config_list.append(ConfigEntity.from_json_dict(cfg_dict))
- cfg_dict = {
- "index": -1,
- "code_hash": None,
- "entity": [
- ["tile_ic", "sp", [4, 4]],
- ["tile_oc", "sp", [2, 16]],
- ["tile_oh", "ot", 1],
- ["tile_ow", "sp", [4, 2]],
- ],
- }
- config_list.append(ConfigEntity.from_json_dict(cfg_dict))
- cfg_dict = {
- "index": -1,
- "code_hash": None,
- "entity": [
- ["tile_ic", "sp", [16, 2]],
- ["tile_oc", "sp", [8, 4]],
- ["tile_ow", "sp", [2, 4]],
- ["unroll_kw", "ot", False],
- ],
- }
- config_list.append(ConfigEntity.from_json_dict(cfg_dict))
-
- records = []
- new_args = new_args + new_args
- tasks = tasks + tasks
- for args, cost, config, task in zip(new_args, costs, config_list, tasks):
- task.args = args
- ms_input = MeasureInput(target=target, task=task, config=config)
- ms_output = MeasureResult(costs=(cost,), error_no=0, all_cost=-1, timestamp=-1)
- records.append((ms_input, ms_output))
-
- ltf_records = []
- ltf_arg = [te.placeholder((1, 64, 16, 16, 8), dtype=dtype), "NCHW8c", "NCHW512c"]
- ltf_task = autotvm.task.create("layout_transform", ltf_arg, target)
- ms_input = MeasureInput(target=target, task=ltf_task, config=None)
- ms_output = MeasureResult(costs=(1.91224744e-05,), error_no=0, all_cost=-1, timestamp=-1)
- ltf_records.append((ms_input, ms_output))
-
- executor = DPTuner(net, {"data": dshape}, records, target_ops, target)
- executor.benchmark_layout_transform(layout_records=ltf_records, infer_layout=True)
- executor.run()
- out = [record[0].config for record in executor.get_optimal_records()]
- expected_out = [records[3][0].config, records[1][0].config, records[2][0].config]
- assert expected_out == out, "Output mismatch: expecting %s but got %s" % (
- str(expected_out),
- str(out),
- )
-
- executor = PBQPTuner(net, {"data": dshape}, records, target_ops, target)
- executor.benchmark_layout_transform(layout_records=ltf_records, infer_layout=True)
- executor.run()
- out = [record[0].config for record in executor.get_optimal_records()]
- expected_out = [records[3][0].config, records[1][0].config, records[2][0].config]
- assert expected_out == out, "Output mismatch: expecting %s but got %s" % (
- str(expected_out),
- str(out),
- )
-
-
-@tvm.testing.requires_x86
-def test_tuple():
- target = "llvm"
- dtype = "float32"
- dshape = (1, 5, 32, 32)
- layout = "NCHW"
- conv2d = relay.op.get("nn.conv2d")
- target_ops = [conv2d]
-
- data = relay.var("data", shape=dshape, dtype=dtype)
- w0 = relay.var("w0_weight")
- conv0 = relay.nn.conv2d(data, w0, channels=2, kernel_size=(3, 3), padding=(1, 1))
- w1 = relay.var("w1_weight")
- conv1 = relay.nn.conv2d(data, w1, channels=3, kernel_size=(3, 3), padding=(1, 1))
- out = relay.concatenate([conv0, conv1], axis=1)
- net = relay.Function(relay.analysis.free_vars(out), out)
- net, params = relay.testing.create_workload(net)
-
- tasks = autotvm.task.extract_from_program(
- net["main"], target=target, params=params, ops=(conv2d,)
- )
- new_args = [
- _create_args(
- (1, 5, 32, 32), (2, 5, 3, 3), (1, 1), (1, 1, 1, 1), (1, 1), layout, layout, dtype, dtype
- ),
- _create_args(
- (1, 5, 32, 32), (3, 5, 3, 3), (1, 1), (1, 1, 1, 1), (1, 1), layout, layout, dtype, dtype
- ),
- ]
- costs = [0.01, 0.012, 0.03, 0.04]
- config_list = []
- cfg_dict = {
- "index": -1,
- "code_hash": None,
- "entity": [
- ["tile_ic", "sp", [1, 5]],
- ["tile_oc", "sp", [1, 2]],
- ["tile_ow", "sp", [4, 8]],
- ["unroll_kw", "ot", True],
- ],
- }
- config_list.append(ConfigEntity.from_json_dict(cfg_dict))
- cfg_dict = {
- "index": -1,
- "code_hash": None,
- "entity": [
- ["tile_ic", "sp", [1, 5]],
- ["tile_oc", "sp", [1, 3]],
- ["tile_ow", "sp", [2, 16]],
- ["unroll_kw", "ot", False],
- ],
- }
- config_list.append(ConfigEntity.from_json_dict(cfg_dict))
- cfg_dict = {
- "index": -1,
- "code_hash": None,
- "entity": [
- ["tile_ic", "sp", [1, 5]],
- ["tile_oc", "sp", [2, 1]],
- ["tile_ow", "sp", [4, 8]],
- ["unroll_kw", "ot", True],
- ],
- }
- config_list.append(ConfigEntity.from_json_dict(cfg_dict))
- cfg_dict = {
- "index": -1,
- "code_hash": None,
- "entity": [
- ["tile_ic", "sp", [1, 5]],
- ["tile_oc", "sp", [3, 1]],
- ["tile_ow", "sp", [2, 16]],
- ["unroll_kw", "ot", False],
- ],
- }
- config_list.append(ConfigEntity.from_json_dict(cfg_dict))
-
- records = []
- new_args = new_args + new_args
- tasks = tasks + tasks
- for args, cost, config, task in zip(new_args, costs, config_list, tasks):
- task.args = args
- ms_input = MeasureInput(target=target, task=task, config=config)
- ms_output = MeasureResult(costs=(cost,), error_no=0, all_cost=-1, timestamp=-1)
- records.append((ms_input, ms_output))
-
- ltf_records = []
- ltf_arg = [te.placeholder((1, 64, 16, 16, 8), dtype=dtype), "NCHW8c", "NCHW512c"]
- ltf_task = autotvm.task.create("layout_transform", ltf_arg, target)
- ms_input = MeasureInput(target=target, task=ltf_task, config=None)
- ms_output = MeasureResult(costs=(1.91224744e-05,), error_no=0, all_cost=-1, timestamp=-1)
- ltf_records.append((ms_input, ms_output))
-
- executor = DPTuner(net, {"data": dshape}, records, target_ops, target)
- executor.benchmark_layout_transform(layout_records=ltf_records, infer_layout=True)
- executor.run()
- out = [record[0].config for record in executor.get_optimal_records()]
- expected_out = [records[2][0].config, records[1][0].config]
- assert expected_out == out, "Output mismatch: expecting %s but got %s" % (
- str(expected_out),
- str(out),
- )
-
- executor = PBQPTuner(net, {"data": dshape}, records, target_ops, target)
- executor.benchmark_layout_transform(layout_records=ltf_records, infer_layout=True)
- executor.run()
- out = [record[0].config for record in executor.get_optimal_records()]
- expected_out = [records[2][0].config, records[1][0].config]
- assert expected_out == out, "Output mismatch: expecting %s but got %s" % (
- str(expected_out),
- str(out),
- )
-
-
-@tvm.testing.requires_x86
-def test_triangle_block():
- target = "llvm"
- dtype = "float32"
- dshape = (1, 3, 8, 8)
- layout = "NCHW"
- conv2d = relay.op.get("nn.conv2d")
- target_ops = [conv2d]
-
- data = relay.var("data", shape=dshape, dtype=dtype)
- w0 = relay.var("w0_weight")
- conv0 = relay.nn.conv2d(data, w0, channels=16, kernel_size=(3, 3), padding=(1, 1))
- w1 = relay.var("w1_weight")
- conv1 = relay.nn.conv2d(conv0, w1, channels=32, kernel_size=(1, 1))
- w2 = relay.var("w2_weight")
- conv2 = relay.nn.conv2d(data, w2, channels=32, kernel_size=(3, 3), padding=(1, 1))
- out = relay.concatenate([conv0, conv1, conv2], axis=1)
- net = relay.Function(relay.analysis.free_vars(out), out)
- net, params = relay.testing.create_workload(net)
-
- tasks = autotvm.task.extract_from_program(
- net["main"], target=target, params=params, ops=(conv2d,)
- )
- new_args = [
- _create_args(
- (1, 3, 8, 8), (16, 3, 3, 3), (1, 1), (1, 1, 1, 1), (1, 1), layout, layout, dtype, dtype
- ),
- _create_args(
- (1, 16, 8, 8),
- (32, 16, 1, 1),
- (1, 1),
- (0, 0, 0, 0),
- (1, 1),
- layout,
- layout,
- dtype,
- dtype,
- ),
- _create_args(
- (1, 3, 8, 8), (32, 3, 3, 3), (1, 1), (1, 1, 1, 1), (1, 1), layout, layout, dtype, dtype
- ),
- ]
- costs = [0.04, 0.012, 0.03, 0.02, 0.02, 0.045]
- config_list = []
- cfg_dict = {
- "index": -1,
- "code_hash": None,
- "entity": [
- ["tile_ic", "sp", [3, 1]],
- ["tile_oc", "sp", [4, 4]],
- ["tile_ow", "sp", [4, 2]],
- ["unroll_kw", "ot", True],
- ],
- }
- config_list.append(ConfigEntity.from_json_dict(cfg_dict))
- cfg_dict = {
- "index": -1,
- "code_hash": None,
- "entity": [
- ["tile_ic", "sp", [2, 8]],
- ["tile_oc", "sp", [1, 32]],
- ["tile_oh", "ot", 1],
- ["tile_ow", "sp", [4, 2]],
- ],
- }
- config_list.append(ConfigEntity.from_json_dict(cfg_dict))
- cfg_dict = {
- "index": -1,
- "code_hash": None,
- "entity": [
- ["tile_ic", "sp", [8, 4]],
- ["tile_oc", "sp", [4, 8]],
- ["tile_ow", "sp", [2, 4]],
- ["unroll_kw", "ot", False],
- ],
- }
- config_list.append(ConfigEntity.from_json_dict(cfg_dict))
- cfg_dict = {
- "index": -1,
- "code_hash": None,
- "entity": [
- ["tile_ic", "sp", [1, 3]],
- ["tile_oc", "sp", [2, 8]],
- ["tile_ow", "sp", [4, 2]],
- ["unroll_kw", "ot", True],
- ],
- }
- config_list.append(ConfigEntity.from_json_dict(cfg_dict))
- cfg_dict = {
- "index": -1,
- "code_hash": None,
- "entity": [
- ["tile_ic", "sp", [4, 4]],
- ["tile_oc", "sp", [2, 16]],
- ["tile_oh", "ot", 1],
- ["tile_ow", "sp", [4, 2]],
- ],
- }
- config_list.append(ConfigEntity.from_json_dict(cfg_dict))
- cfg_dict = {
- "index": -1,
- "code_hash": None,
- "entity": [
- ["tile_ic", "sp", [16, 2]],
- ["tile_oc", "sp", [8, 4]],
- ["tile_ow", "sp", [2, 4]],
- ["unroll_kw", "ot", False],
- ],
- }
- config_list.append(ConfigEntity.from_json_dict(cfg_dict))
-
- records = []
- new_args = new_args + new_args
- tasks = tasks + tasks
- for args, cost, config, task in zip(new_args, costs, config_list, tasks):
- task.args = args
- ms_input = MeasureInput(target=target, task=task, config=config)
- ms_output = MeasureResult(costs=(cost,), error_no=0, all_cost=-1, timestamp=-1)
- records.append((ms_input, ms_output))
-
- ltf_records = []
- ltf_arg = [te.placeholder((1, 64, 16, 16, 8), dtype=dtype), "NCHW8c", "NCHW512c"]
- ltf_task = autotvm.task.create("layout_transform", ltf_arg, target)
- ms_input = MeasureInput(target=target, task=ltf_task, config=None)
- ms_output = MeasureResult(costs=(1.91224744e-05,), error_no=0, all_cost=-1, timestamp=-1)
- ltf_records.append((ms_input, ms_output))
-
- executor = DPTuner(net, {"data": dshape}, records, target_ops, target)
- executor.benchmark_layout_transform(layout_records=ltf_records, infer_layout=True)
- executor.run()
- out = [record[0].config for record in executor.get_optimal_records()]
- expected_out = [records[3][0].config, records[1][0].config, records[2][0].config]
- assert expected_out == out, "Output mismatch: expecting %s but got %s" % (
- str(expected_out),
- str(out),
- )
-
- executor = PBQPTuner(net, {"data": dshape}, records, target_ops, target)
- executor.benchmark_layout_transform(layout_records=ltf_records, infer_layout=True)
- executor.run()
- out = [record[0].config for record in executor.get_optimal_records()]
- expected_out = [records[3][0].config, records[1][0].config, records[2][0].config]
- assert expected_out == out, "Output mismatch: expecting %s but got %s" % (
- str(expected_out),
- str(out),
- )
-
-
-if __name__ == "__main__":
- test_graph_tuner_layout_transform()
- test_DPTuner_run()
- test_PBQPTuner_run()
- test_many_sub_graphs()
- test_tuple()
- test_triangle_block()
diff --git a/tests/python/autotvm/test_autotvm_graph_tuner_utils.py b/tests/python/autotvm/test_autotvm_graph_tuner_utils.py
deleted file mode 100644
index 583bd36684..0000000000
--- a/tests/python/autotvm/test_autotvm_graph_tuner_utils.py
+++ /dev/null
@@ -1,180 +0,0 @@
-# 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.
-
-# NOTE: We name this test file to start with test_graph_tuner
-# to make it execute after zero_rank tensor test cases. This
-# helps avoid topi arithmetic operator overloading issue:
-# https://github.com/apache/tvm/issues/3240
-# TODO: restore the file name after this issue is resolved.
-import pytest
-
-import tvm
-from tvm import te
-
-from tvm import autotvm, relay
-from tvm.relay.testing import synthetic
-from tvm.autotvm.graph_tuner.utils import (
- has_multiple_inputs,
- get_direct_ancestor,
- get_in_nodes,
- get_out_nodes,
- expr2graph,
- bind_inputs,
-)
-from tvm.autotvm.graph_tuner._base import OPT_OUT_OP
-from tvm.autotvm.graph_tuner.utils.traverse_graph import _replace_device_with_tracing
-from tvm.relay.expr import Call, TupleGetItem, Tuple, Var
-
-
-def verify_has_multiple_inputs(node_list, node_idx, input_names, expected_result):
- out = has_multiple_inputs(node_list, node_idx, input_names, OPT_OUT_OP)
- assert out == expected_result, "Output mismatch: expecting checking %s to be %s but got %s." % (
- node_list[node_idx]["op"],
- str(expected_result),
- str(out),
- )
-
-
-def test_has_multiple_inputs():
- data = relay.var("data")
- out1 = data * relay.expr.const(3.0)
- w0 = relay.var("w0")
- out2 = relay.nn.conv2d(data, w0)
- out = relay.add(out1, out2)
- net = relay.Function(relay.analysis.free_vars(out), out)
- net = bind_inputs(net, {"data": (1, 16, 224, 224), "w0": (16, 16, 1, 1)})
- target_ops = [relay.op.get("nn.conv2d")]
- node_list = []
- node_dict = {}
- expr2graph(net, target_ops, node_dict, node_list, tvm.target.Target("llvm"))
- input_names = ["data"]
- verify_has_multiple_inputs(node_list, 2, input_names, False)
- verify_has_multiple_inputs(node_list, 4, input_names, False)
- verify_has_multiple_inputs(node_list, 5, input_names, True)
-
-
-def test_expr2graph():
- mod, _ = synthetic.get_workload()
- node_dict = {}
- node_list = []
- target_ops = [relay.op.get("nn.conv2d")]
- op_name_list = []
-
- def _count_node(node):
- if isinstance(node, Call):
- op_name_list.append(node.op)
- elif isinstance(node, (Var, TupleGetItem, Tuple)):
- op_name_list.append(None)
-
- relay.analysis.post_order_visit(mod["main"], _count_node)
-
- expr2graph(mod["main"], target_ops, node_dict, node_list, tvm.target.Target("llvm"))
- assert len(node_list) == len(op_name_list)
- for i, item in enumerate(zip(op_name_list, node_list)):
- op_name, node = item
- assert op_name == node["op"], "%dth Node operator mismatch: expecting %s but got %s" % (
- i,
- str(op_name),
- str(node["op"]),
- )
-
-
-def test_get_direct_ancestor():
- data = relay.var("data")
- w0 = relay.var("w0")
- out1 = relay.nn.conv2d(data, w0)
- out2 = relay.add(out1, data * relay.expr.const(5.0))
- out3 = out2 + relay.expr.const(2.5)
- w1 = relay.var("w1")
- out = relay.nn.conv2d(out3, w1)
- net = relay.Function(relay.analysis.free_vars(out), out)
- net = bind_inputs(net, {"data": (1, 16, 224, 224), "w0": (16, 16, 1, 1), "w1": (16, 16, 1, 1)})
- target_ops = [relay.op.get("nn.conv2d")]
- node_list = []
- node_dict = {}
- expr2graph(net, target_ops, node_dict, node_list, tvm.target.Target("llvm"))
- visited_dict = {}
- input_names = ["data"]
- out = get_direct_ancestor(node_list, visited_dict, target_ops, 5, input_names)
- assert out == [0], "Output mismatch: expecting [0] but got %s." % str(out)
-
- # non-regression test
- out = relay.add(relay.log(data), relay.sqrt(data))
- net = relay.Function(relay.analysis.free_vars(out), out)
- net = bind_inputs(net, {"data": (1, 16, 224, 224)})
- node_list = []
- node_dict = {}
- expr2graph(net, target_ops, node_dict, node_list, tvm.target.Target("llvm"))
- out = get_direct_ancestor(node_list, visited_dict, target_ops, 3, input_names)
- assert out == [0], "Output mismatch: expecting [0] but got %s." % str(out)
-
-
-def test_get_in_nodes():
- data = relay.var("data")
- w0 = relay.var("w0")
- out1 = relay.nn.conv2d(data, w0)
- out2 = relay.add(out1, data)
- out3 = out2 + relay.expr.const(2.5)
- w1 = relay.var("w1")
- out = relay.nn.conv2d(out3, w1)
- net = relay.Function(relay.analysis.free_vars(out), out)
- net = bind_inputs(net, {"data": (1, 16, 224, 224), "w0": (16, 16, 1, 1), "w1": (16, 16, 1, 1)})
- target_ops = [relay.op.get("nn.conv2d")]
- input_names = ["data"]
- node_list = []
- node_dict = {}
- expr2graph(net, target_ops, node_dict, node_list, tvm.target.Target("llvm"))
- out = get_in_nodes(node_list, target_ops, input_names)
- expected_out = {3: [0], 4: [3, 0], 7: [4]}
- diff_set = set(out) ^ set(expected_out)
- if len(diff_set) != 0:
- raise RuntimeError(
- "Output mismatch: expecting %s but got %s." % (str(expected_out), str(out))
- )
-
-
-def test_get_out_nodes():
- in_nodes_dict = {8: [4], 4: [3, 0], 3: [0]}
- expected_out = {0: [3, 4], 3: [4], 4: [8], 8: []}
- out = get_out_nodes(in_nodes_dict)
- diff_set = set(out) ^ set(expected_out)
- if len(diff_set) != 0:
- raise RuntimeError(
- "Output mismatch: expecting %s but got %s." % (str(expected_out), str(out))
- )
-
-
-def test_target_device_replacement():
- assert _replace_device_with_tracing("cuda") == "cuda -device=tracing"
- assert (
- _replace_device_with_tracing("cuda -device=some_device -libs=cudnn")
- == "cuda -device=tracing -libs=cudnn"
- )
- assert (
- _replace_device_with_tracing("llvm -device=arm_cpu -arg=xxx")
- == "llvm -device=tracing -arg=xxx"
- )
- assert _replace_device_with_tracing("llvm -device=arm_cpu") == "llvm -device=tracing"
- assert _replace_device_with_tracing("llvm -device=abc, def") == "llvm -device=tracing"
-
-
-if __name__ == "__main__":
- test_has_multiple_inputs()
- test_expr2graph()
- test_get_direct_ancestor()
- test_get_in_nodes()
- test_get_out_nodes()
diff --git a/tests/python/autotvm/test_autotvm_index_tuner.py b/tests/python/autotvm/test_autotvm_index_tuner.py
deleted file mode 100644
index 514577f1c9..0000000000
--- a/tests/python/autotvm/test_autotvm_index_tuner.py
+++ /dev/null
@@ -1,131 +0,0 @@
-# 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.
-"""Test index based tuners"""
-
-import multiprocessing
-from tvm.testing.autotvm import DummyRunner, get_sample_task
-from tvm import autotvm
-
-
-def test_grid_search_tuner():
- """Test GridSearchTuner"""
-
- task, _ = get_sample_task()
- measure_option = autotvm.measure_option(builder=autotvm.LocalBuilder(), runner=DummyRunner())
-
- # When no range index, range_length should be the length of config space
- tuner = autotvm.tuner.GridSearchTuner(task)
- assert tuner.begin_idx == 0
- assert tuner.end_idx == 64
- assert tuner.index == 0
- assert tuner.range_length == 64
- assert tuner.visited_max == 64
-
- # With range index, range_length should be the length of the specified range
- tuner = autotvm.tuner.GridSearchTuner(task, range_idx=(8, 15))
- assert tuner.begin_idx == 8
- assert tuner.end_idx == 16
- assert tuner.index == 8
- assert tuner.range_length == 8
- assert tuner.visited_max == 8
-
- # Tuner should only focus on the specified range
- tuner.tune(n_trial=8, measure_option=measure_option)
- assert len(tuner.visited) == 8
- assert not tuner.has_next()
-
- # With multi-filter
- task, _ = get_sample_task()
- task.config_space.multi_filter(
- filter=lambda entity: 32 <= (entity["tile_x"].size[1] * entity["tile_y"].size[1]) < 1024
- )
-
- tuner = autotvm.tuner.GridSearchTuner(task)
- assert tuner.begin_idx == 0
- assert tuner.end_idx == 64
- assert tuner.index == 5
- assert tuner.range_length == 64
- assert tuner.visited_max == 34
-
- # With range index, range_length should be the length of the specified range
- tuner = autotvm.tuner.GridSearchTuner(task, range_idx=(8, 15))
- assert tuner.begin_idx == 8
- assert tuner.end_idx == 16
- assert tuner.index == 12
- assert tuner.range_length == 8
- assert tuner.visited_max == 4
-
- # Tuner should only focus on the specified range
- tuner.tune(n_trial=8, measure_option=measure_option)
- assert len(tuner.visited) == 4
- assert not tuner.has_next()
-
-
-def grid_search_spawn():
- assert multiprocessing.get_spawn_method(False) == "spawn"
- test_grid_search_tuner()
-
-
-def test_grid_search_tuner_spawn():
- ctx = multiprocessing.get_context("spawn")
- p = ctx.Process(target=test_grid_search_tuner)
- p.start()
- p.join()
-
-
-def test_random_tuner():
- """Test RandomTuner"""
-
- task, _ = get_sample_task()
- measure_option = autotvm.measure_option(builder=autotvm.LocalBuilder(), runner=DummyRunner())
-
- tuner = autotvm.tuner.RandomTuner(task, range_idx=(8, 15))
- assert tuner.begin_idx == 8
- assert tuner.end_idx == 16
- assert tuner.range_length == 8
- assert tuner.visited_max == 8
-
- # Tuner should only focus on the specified range and should visit all indices
- tuner.tune(n_trial=8, measure_option=measure_option)
- assert len(tuner.visited) == 8
- assert not tuner.has_next()
- for idx in tuner.visited:
- assert 8 <= idx <= 15
-
- # With multi-filter
- task, _ = get_sample_task()
- task.config_space.multi_filter(
- filter=lambda entity: 32 <= (entity["tile_x"].size[1] * entity["tile_y"].size[1]) < 1024
- )
- tuner = autotvm.tuner.RandomTuner(task, range_idx=(8, 15))
- assert tuner.begin_idx == 8
- assert tuner.end_idx == 16
- assert tuner.range_length == 8
- assert tuner.visited_max == 4
-
- # Tuner should only focus on the specified range and should visit all indices
- tuner.tune(n_trial=8, measure_option=measure_option)
- assert len(tuner.visited) == 4
- assert not tuner.has_next()
- for idx in tuner.visited:
- assert 8 <= idx <= 15
-
-
-if __name__ == "__main__":
- test_grid_search_tuner()
- test_grid_search_tuner_spawn()
- test_random_tuner()
diff --git a/tests/python/autotvm/test_autotvm_measure.py b/tests/python/autotvm/test_autotvm_measure.py
deleted file mode 100644
index 3ef5cbdad6..0000000000
--- a/tests/python/autotvm/test_autotvm_measure.py
+++ /dev/null
@@ -1,94 +0,0 @@
-# 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.
-"""Test builder and runner"""
-import logging
-import multiprocessing
-import concurrent
-
-import numpy as np
-
-import tvm
-from tvm import te
-from tvm.autotvm.measure import executor
-from tvm.testing.autotvm import DummyRunner, bad_matmul, get_sample_task
-from tvm import autotvm
-from tvm.autotvm.measure.measure import MeasureErrorNo, MeasureResult
-from tvm.autotvm import measure
-from inspect import Signature
-
-
-def test_task_tuner_without_measurement():
- """test task and tuner without measurement"""
- task, _ = get_sample_task()
-
- measure_option = autotvm.measure_option(builder=autotvm.LocalBuilder(), runner=DummyRunner())
-
- logging.info("%s", task.config_space)
-
- for tuner_class in [
- autotvm.tuner.RandomTuner,
- autotvm.tuner.GridSearchTuner,
- autotvm.tuner.GATuner,
- autotvm.tuner.XGBTuner,
- ]:
- tuner = tuner_class(task)
- tuner.tune(n_trial=10, measure_option=measure_option)
- assert tuner.best_flops > 1
-
-
-def task_tuner_spawn():
- assert multiprocessing.get_start_method(False) == "spawn"
- test_task_tuner_without_measurement()
-
-
-def test_task_tuner_without_measurement_spawn():
- # Subprocesses inherit the spawn method of their parents
- ctx = multiprocessing.get_context("spawn")
- p = ctx.Process(target=task_tuner_spawn)
- p.start()
- p.join()
-
-
-def test_task_runner_with_ref_input():
- """test runner ref_input without measurement"""
- refinp = [np.random.rand(128, 128) for i in range(3)]
- runner = measure.LocalRunner()
- runner.ref_input = refinp
-
- class DummyExecutor(measure.executor.Executor):
- def __init__(self):
- self.ran_dummy_executor = False
-
- def submit(self, func, *args, **kwargs):
- self.ran_dummy_executor = True
- sig = Signature.from_callable(func)
- assert sig.bind(*args, **kwargs).arguments["ref_input"] == refinp
- dummy_future = concurrent.futures.Future()
- dummy_future.set_result(None)
- return dummy_future
-
- runner.executor = DummyExecutor()
- runner.run([None], [None])
- assert runner.executor.ran_dummy_executor
-
-
-if __name__ == "__main__":
- logging.basicConfig(level=logging.INFO)
-
- test_task_tuner_without_measurement()
- test_task_tuner_without_measurement_spawn()
- test_task_runner_with_ref_input()
diff --git a/tests/python/autotvm/test_autotvm_record.py b/tests/python/autotvm/test_autotvm_record.py
deleted file mode 100644
index 693810d3f9..0000000000
--- a/tests/python/autotvm/test_autotvm_record.py
+++ /dev/null
@@ -1,163 +0,0 @@
-# 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.
-"""test the correctness of dump and load of data log"""
-from io import StringIO
-from os import PathLike
-import time
-
-from tvm.contrib import utils
-
-from tvm import autotvm
-from tvm.autotvm.measure import MeasureInput, MeasureResult, MeasureErrorNo
-from tvm.autotvm.record import encode, decode, ApplyHistoryBest, measure_str_key
-
-from tvm.testing.autotvm import get_sample_task
-
-
-def test_load_dump():
- task, target = get_sample_task()
-
- inp = MeasureInput(target, task, task.config_space.get(0))
- result = MeasureResult(
- (2.0, 2.23, 0.23, 0.123, 0.234, 0.123), MeasureErrorNo.NO_ERROR, 2.3, time.time()
- )
-
- for protocol in ["json", "pickle"]:
- row = encode(inp, result, protocol=protocol)
- inp_2, result_2 = decode(row, protocol=protocol)
-
- assert measure_str_key(inp) == measure_str_key(inp_2), "%s vs %s" % (
- measure_str_key(inp),
- measure_str_key(inp_2),
- )
- assert result.costs == result_2.costs
- assert result.error_no == result_2.error_no
- assert result.timestamp == result_2.timestamp
-
-
-def test_file_io():
- temp = utils.tempdir()
- file_path = temp.relpath("temp.log")
-
- tsk, target = get_sample_task()
- inputs = [MeasureInput(target, tsk, tsk.config_space.get(i)) for i in range(0, 10)]
- results = [MeasureResult((i,), 0, 0, 0) for i in range(0, 10)]
-
- invalid_inp = MeasureInput(target, tsk, tsk.config_space.get(10))
- invalid_res = MeasureResult((10,), 0, 0, 0)
-
- # Erase the entity map to test if it will be ignored when loading back.
- invalid_inp.config._entity_map = {}
-
- with open(file_path, "w") as fo:
- cb = autotvm.callback.log_to_file(fo)
- cb(None, inputs, results)
- cb(None, [invalid_inp], [invalid_res])
-
- ref = zip(inputs, results)
- for x, y in zip(ref, autotvm.record.load_from_file(file_path)):
- assert x[1] == y[1]
-
- # Confirm functionality of multiple file loads
- hist_best = ApplyHistoryBest([file_path, file_path])
- x = hist_best.query(target, tsk.workload)
- assert str(x) == str(inputs[0][2])
-
-
-def test_apply_history_best(tmpdir):
- tsk, target = get_sample_task()
- best = str(tsk.config_space.get(2))
-
- inputs_batch_1 = [MeasureInput(target, tsk, tsk.config_space.get(i)) for i in range(3)]
- results_batch_1 = [MeasureResult((i,), 0, 0, 0) for i in range(1, 3)]
- results_batch_1.append(MeasureResult((0.5,), 0, 2.3, 0))
-
- # Write data out to file
- filepath_batch_1 = tmpdir / "batch_1.log"
- with open(filepath_batch_1, "w") as file:
- autotvm.callback.log_to_file(file)(None, inputs_batch_1, results_batch_1)
-
- # Load best results from Path
- assert isinstance(filepath_batch_1, PathLike)
- hist_best = ApplyHistoryBest(filepath_batch_1)
- assert str(hist_best.query(target, tsk.workload)) == best
-
- # Load best results from str(Path)
- hist_best = ApplyHistoryBest(str(filepath_batch_1))
- assert str(hist_best.query(target, tsk.workload)) == best
-
- # Write data into StringIO buffer
- stringio_batch_1 = StringIO()
- assert isinstance(filepath_batch_1, PathLike)
- callback = autotvm.callback.log_to_file(stringio_batch_1)
- callback(None, inputs_batch_1, results_batch_1)
- stringio_batch_1.seek(0)
-
- # Load best results from strIO
- hist_best = ApplyHistoryBest(stringio_batch_1)
- assert str(hist_best.query(target, tsk.workload)) == best
-
- # Load best result from list of tuples (MeasureInput, MeasureResult)
- hist_best = ApplyHistoryBest(list(zip(inputs_batch_1, results_batch_1)))
- assert str(hist_best.query(target, tsk.workload)) == best
-
- # Same thing, but iterable instead of list (i.e. no subscripting)
- hist_best = ApplyHistoryBest(zip(inputs_batch_1, results_batch_1))
- assert str(hist_best.query(target, tsk.workload)) == best
-
-
-def test_apply_history_best_multiple_batches(tmpdir):
- tsk, target = get_sample_task()
- best = str(tsk.config_space.get(2))
-
- inputs_batch_1 = [MeasureInput(target, tsk, tsk.config_space.get(i)) for i in range(2)]
- results_batch_1 = [MeasureResult((i,), 0, 0, 0) for i in range(1, 3)]
- filepath_batch_1 = tmpdir / "batch_1.log"
- with open(filepath_batch_1, "w") as file:
- autotvm.callback.log_to_file(file)(None, inputs_batch_1, results_batch_1)
-
- inputs_batch_2 = [MeasureInput(target, tsk, tsk.config_space.get(i)) for i in range(2, 4)]
- results_batch_2 = [MeasureResult((0.5,), 0, 0, 0), MeasureResult((3,), 0, 0, 0)]
- filepath_batch_2 = tmpdir / "batch_2.log"
- with open(filepath_batch_2, "w") as file:
- autotvm.callback.log_to_file(file)(None, inputs_batch_2, results_batch_2)
-
- # Check two Path filepaths works
- hist_best = ApplyHistoryBest([filepath_batch_1, filepath_batch_2])
- assert str(hist_best.query(target, tsk.workload)) == best
-
- # Check that an arbitrary Iterable of Paths works
- # Calling zip() on a single list gives a non-subscriptable Iterable
- hist_best = ApplyHistoryBest(zip([filepath_batch_1, filepath_batch_2]))
- assert str(hist_best.query(target, tsk.workload)) == best
-
- # Check that Iterable of Iterable of tuples is correctly merged
- hist_best = ApplyHistoryBest(
- zip(
- [
- zip(inputs_batch_1, results_batch_1),
- zip(inputs_batch_2, results_batch_2),
- ]
- )
- )
- assert str(hist_best.query(target, tsk.workload)) == best
-
-
-if __name__ == "__main__":
- test_load_dump()
- test_apply_history_best()
- test_file_io()
diff --git a/tests/python/autotvm/test_autotvm_space.py b/tests/python/autotvm/test_autotvm_space.py
deleted file mode 100644
index eb783a9f8b..0000000000
--- a/tests/python/autotvm/test_autotvm_space.py
+++ /dev/null
@@ -1,263 +0,0 @@
-# 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.
-"""Test space definition primitives"""
-
-from tvm import te
-from tvm.autotvm.task.space import ConfigSpace, FallbackConfigEntity
-
-
-def gemm_func(cfg, N, filter_y=None, filter_x=None):
- A = te.placeholder((N, N), name="A")
- B = te.placeholder((N, N), name="B")
-
- k = te.reduce_axis((0, N), name="k")
- C = te.compute((N, N), lambda i, j: te.sum(A[i, k] * B[k, j], axis=[k]), name="C")
-
- s = te.create_schedule([C.op])
-
- y, x = s[C].op.axis
-
- cfg.define_split("tile_y", cfg.axis(y), num_outputs=2, filter=filter_y)
- cfg.define_split("tile_x", cfg.axis(x), num_outputs=2, filter=filter_x)
-
- return s, [A, B, C]
-
-
-def test_split():
- cfg = ConfigSpace()
-
- gemm_func(cfg, 128)
- assert cfg.range_length == 64
- assert len(cfg.space_map["tile_y"]) == 8
-
- # test policy
- cfg = ConfigSpace()
- cfg.define_split("tile_x", cfg.axis(256), policy="factors", num_outputs=3)
- assert len(cfg.space_map["tile_x"]) == 45
-
- cfg.define_split("tile_y", cfg.axis(256), policy="power2", num_outputs=3)
- assert len(cfg.space_map["tile_y"]) == 45
-
- cfg.define_split("tile_z", cfg.axis(256), policy="verbose", num_outputs=3)
- assert len(cfg.space_map["tile_z"]) == 45
-
- cfg.define_split("tile_a", cfg.axis(224), policy="factors", num_outputs=3)
- assert len(cfg.space_map["tile_a"]) == 63
-
- cfg.define_split("tile_b", cfg.axis(224), policy="power2", num_outputs=3)
- assert len(cfg.space_map["tile_b"]) == 36
-
- cfg.define_split("tile_c", cfg.axis(224), policy="verbose", num_outputs=3)
- assert len(cfg.space_map["tile_c"]) == 84
-
- # Count the number of non-negative integer solutions of a + b + c + d = n
- def count4(n):
- cnt = 0
- for a in range(0, n + 1):
- for b in range(0, n - a + 1):
- cnt += n - a - b + 1
- return cnt
-
- # test overflow
- n = 25
- cfg = ConfigSpace()
- cfg.define_split("x", cfg.axis(2**n), policy="factors", num_outputs=4)
- # count4(25) is 3276.
- assert len(cfg.space_map["x"]) == count4(n)
-
- # test fallback
- cfg = FallbackConfigEntity()
- cfg.define_split("tile_n", cfg.axis(128), num_outputs=3)
- cfg.fallback_split("tile_n", [-1, 8, 4])
- # verify if define_split override previously manualy defined split params
- cfg.define_split("tile_n", cfg.axis(128), num_outputs=3)
- assert cfg["tile_n"].size == [4, 8, 4]
-
- cfg = FallbackConfigEntity()
- cfg.define_split("tile_n", cfg.axis(49), num_outputs=3)
- cfg.fallback_split("tile_n", [-1, 8, 4])
- assert cfg["tile_n"].size == [7, 7, 1]
-
- cfg = FallbackConfigEntity()
- cfg.define_split("tile_n", cfg.axis(49), num_outputs=3)
- try:
- cfg.fallback_split("tile_n", [-1, 1, 0])
- assert False
- except RuntimeError:
- pass
-
-
-def _raises_exception(f):
- try:
- f()
- except Exception:
- return True
- return False
-
-
-def test_multi_filter():
- # create config without multi_filter
- cfg = ConfigSpace()
- gemm_func(cfg, 128)
- # create config with multi_filter
- cfg_mf = ConfigSpace()
- gemm_func(cfg_mf, 128)
- cfg_mf.multi_filter(
- filter=lambda entity: 32 <= (entity["tile_x"].size[1] * entity["tile_y"].size[1]) < 1024
- )
- # test len
- assert len(cfg) == 64
- assert len(cfg_mf) == 34
- # test range_length
- assert cfg.range_length == 64
- assert cfg_mf.range_length == 64
- # test dims
- assert cfg.dims == [8, 8]
- assert cfg_mf.dims == [8, 8]
- # test is_index_valid
- assert cfg.is_index_valid(0) is True
- assert cfg.is_index_valid(15) is True
- assert cfg_mf.is_index_valid(0) is False
- assert cfg_mf.is_index_valid(15) is True
- # test get
- assert _raises_exception(lambda: cfg.get(0)) is False
- assert _raises_exception(lambda: cfg.get(15)) is False
- assert _raises_exception(lambda: cfg_mf.get(0)) is True
- assert _raises_exception(lambda: cfg_mf.get(15)) is False
- # test subrange_length
- assert cfg.subrange_length(0, 64) == 64
- assert cfg.subrange_length(0, 32) == 32
- assert cfg.subrange_length(16, 32) == 16
- assert cfg.subrange_length(16, 16) == 0
- assert _raises_exception(lambda: cfg.subrange_length(0, 128))
- assert _raises_exception(lambda: cfg.subrange_length(-64, 64))
- assert _raises_exception(lambda: cfg.subrange_length(64, 0))
- assert cfg_mf.subrange_length(0, 64) == 34
- assert cfg_mf.subrange_length(0, 32) == 17
- assert cfg_mf.subrange_length(16, 32) == 10
- assert cfg_mf.subrange_length(16, 16) == 0
- assert _raises_exception(lambda: cfg_mf.subrange_length(0, 128))
- assert _raises_exception(lambda: cfg_mf.subrange_length(-64, 64))
- assert _raises_exception(lambda: cfg_mf.subrange_length(64, 0))
- # test point2knob
- assert cfg.point2knob(0) == [0, 0]
- assert cfg.point2knob(4) == [4, 0]
- assert cfg.point2knob(8) == [0, 1]
- assert cfg.point2knob(12) == [4, 1]
- assert cfg_mf.point2knob(0) == [0, 0]
- assert cfg_mf.point2knob(4) == [4, 0]
- assert cfg_mf.point2knob(8) == [0, 1]
- assert cfg_mf.point2knob(12) == [4, 1]
- # test knob2point
- assert cfg.knob2point([0, 0]) == 0
- assert cfg.knob2point([4, 0]) == 4
- assert cfg.knob2point([0, 1]) == 8
- assert cfg.knob2point([4, 1]) == 12
- assert cfg_mf.knob2point([0, 0]) == 0
- assert cfg_mf.knob2point([4, 0]) == 4
- assert cfg_mf.knob2point([0, 1]) == 8
- assert cfg_mf.knob2point([4, 1]) == 12
- # get_rand_index
- cfg_valid_indexes = list(filter(lambda idx: cfg.is_index_valid(idx), range(cfg.range_length)))
- assert cfg.get_rand_index() in cfg_valid_indexes
- assert cfg.get_rand_index(start=15, end=16) == 15
- assert 10 <= cfg.get_rand_index(start=10, end=20) < 20
- assert cfg.get_rand_index(to_exclude=cfg_valid_indexes[:-1]) == cfg_valid_indexes[-1:][0]
- cfg_mf_valid_indexes = list(
- filter(lambda idx: cfg_mf.is_index_valid(idx), range(cfg_mf.range_length))
- )
- assert cfg_mf.get_rand_index() in cfg_mf_valid_indexes
- assert cfg_mf.get_rand_index(start=15, end=16) == 15
- assert 10 <= cfg_mf.get_rand_index(start=10, end=20) < 20
- assert (
- cfg_mf.get_rand_index(to_exclude=cfg_mf_valid_indexes[:-1]) == cfg_mf_valid_indexes[-1:][0]
- )
- # get_next_index
- assert cfg.get_next_index(0) == 1
- assert cfg.get_next_index(0, 1) == 1
- assert cfg.get_next_index(0, 2) == 2
- assert cfg.get_next_index(0, -1) is None
- assert cfg.get_next_index(0, -2) is None
- assert cfg.get_next_index(63) is None
- assert cfg.get_next_index(63, 1) is None
- assert cfg.get_next_index(63, 2) is None
- assert cfg.get_next_index(63, -1) == 62
- assert cfg.get_next_index(63, -2) == 61
- assert cfg.get_next_index(60, 1, end=63) == 61
- assert cfg.get_next_index(63, -1, start=60) == 62
- assert cfg_mf.get_next_index(0) == 5
- assert cfg_mf.get_next_index(0, 1) == 5
- assert cfg_mf.get_next_index(0, 2) == 6
- assert cfg_mf.get_next_index(0, -1) is None
- assert cfg_mf.get_next_index(0, -2) is None
- assert cfg_mf.get_next_index(63) is None
- assert cfg_mf.get_next_index(63, 1) is None
- assert cfg_mf.get_next_index(63, 2) is None
- assert cfg_mf.get_next_index(63, -1) == 58
- assert cfg_mf.get_next_index(63, -2) == 57
- assert cfg_mf.get_next_index(60, 1, end=63) is None
- assert cfg_mf.get_next_index(63, -1, start=60) is None
- # test sample_ints
- cfg_ints = cfg.sample_ints(5)
- assert len(cfg_ints) == 5
- assert set(cfg_ints).issubset(cfg_valid_indexes)
- cfg_mf_ints = cfg_mf.sample_ints(5)
- assert len(cfg_mf_ints) == 5
- assert set(cfg_mf_ints).issubset(cfg_mf_valid_indexes)
- # test random_walk
- cfg_walk = cfg.random_walk(15)
- assert cfg_walk != 15
- assert cfg_walk in cfg_valid_indexes
- cfg_mf_walk = cfg_mf.random_walk(15)
- assert cfg_mf_walk != 15
- assert cfg_mf_walk in cfg_mf_valid_indexes
-
-
-def test_filter_and_multi_filter():
- # test the order: filter -> multi_filter
- cfg = ConfigSpace()
- gemm_func(cfg, 128, filter_y=lambda y: y.size[-1] < 64)
- # after adding filter
- assert len(cfg) == 48
- assert cfg.range_length == 48
- cfg.multi_filter(
- filter=lambda entity: 32 <= (entity["tile_x"].size[1] * entity["tile_y"].size[1]) < 1024
- )
- # after adding multi_filter
- assert len(cfg) == 27
- assert cfg.range_length == 48
-
- # test the order: multi_filter -> filter
- cfg = ConfigSpace()
- s, (A, B, C) = gemm_func(cfg, 128, filter_y=None)
- cfg.multi_filter(
- filter=lambda entity: 32 <= (entity["tile_x"].size[1] * entity["tile_y"].size[1]) < 1024
- )
- # after adding multi_filter
- assert len(cfg) == 34
- assert cfg.range_length == 64
- y, x = s[C].op.axis
- cfg.define_split("tile_y", cfg.axis(y), num_outputs=2, filter=lambda y: y.size[-1] < 64)
- # after adding filter
- assert len(cfg) == 27
- assert cfg.range_length == 48
-
-
-if __name__ == "__main__":
- test_split()
- test_multi_filter()
- test_filter_and_multi_filter()
diff --git a/tests/python/autotvm/test_autotvm_xgboost_model.py b/tests/python/autotvm/test_autotvm_xgboost_model.py
deleted file mode 100644
index b9f157247e..0000000000
--- a/tests/python/autotvm/test_autotvm_xgboost_model.py
+++ /dev/null
@@ -1,98 +0,0 @@
-# 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.
-import time
-
-import multiprocessing
-import numpy as np
-
-import tvm
-from tvm import te
-from tvm import autotvm
-from tvm.autotvm import MeasureInput, MeasureResult
-from tvm.autotvm.tuner.xgboost_cost_model import XGBoostCostModel
-
-from tvm.testing.autotvm import get_sample_task, get_sample_records
-
-
-def test_fit():
- task, target = get_sample_task()
- records = get_sample_records(n=500)
-
- base_model = XGBoostCostModel(task, feature_type="itervar", loss_type="reg")
- base_model.fit_log(records, plan_size=32)
-
- upper_model = XGBoostCostModel(task, feature_type="itervar", loss_type="reg")
- upper_model.load_basemodel(base_model)
-
- xs = np.arange(10)
- ys = np.arange(10)
-
- upper_model.fit(xs, ys, plan_size=32)
-
- # feature lengths are not guaranteed to always be the same
- upper_model.predict(np.ones(12))
- upper_model.predict(np.ones(8))
-
-
-def fit_spawn():
- assert multiprocessing.get_start_method(False) == "spawn"
- test_fit()
-
-
-def test_fit_spawn():
- # Subprocesses inherit the spawn method of their parents
- ctx = multiprocessing.get_context("spawn")
- p = ctx.Process(target=test_fit)
- p.start()
- p.join()
-
-
-def test_tuner():
- task, target = get_sample_task()
- records = get_sample_records(n=10)
-
- tuner = autotvm.tuner.XGBTuner(task)
- tuner.load_history(records, min_seed_records=10)
- # Confirm that loading history successfully loaded a
- # base_model.
- assert tuner.cost_model.base_model is not None
-
- tuner = autotvm.tuner.XGBTuner(task)
- tuner.load_history(records, min_seed_records=11)
- # Confirm that loading history did not load base_model
- # when not enough records according to `min_seed_records`
- # are provided
- assert tuner.cost_model.base_model is None
-
-
-def test_update():
- task, target = get_sample_task()
- tuner = autotvm.tuner.XGBTuner(task)
- n_records = 5
- records = get_sample_records(n=n_records)
- tuner.update([inp for inp, _ in records], [res for _, res in records])
- assert len(tuner.xs) == n_records
- assert len(tuner.ys) == n_records
- assert len(tuner.visited) == n_records
- assert all(x in tuner.visited for x in tuner.xs)
-
-
-if __name__ == "__main__":
- test_fit()
- test_fit_spawn()
- test_tuner()
- test_update()
diff --git a/tests/python/codegen/test_target_codegen_aarch64.py b/tests/python/codegen/test_target_codegen_aarch64.py
index f596549a10..366198c7de 100644
--- a/tests/python/codegen/test_target_codegen_aarch64.py
+++ b/tests/python/codegen/test_target_codegen_aarch64.py
@@ -25,7 +25,6 @@ import pytest
import tvm
from tvm import te
from tvm.script import tir as T
-from tvm.topi.arm_cpu.pstate_attributes import SMEAttributes
from tvm.target.codegen import llvm_version_major
@@ -498,65 +497,6 @@ def test_codegen_vscale():
assert re.findall(r"llvm.vscale.i32", llvm), "No vscale in generated LLVM."
-@pytest.mark.skipif(
- llvm_version_major() < 16, reason="SME is not supported in earlier versions of LLVM"
-)
-@pytest.mark.parametrize("dtype", ["float32", "float16"])
-def test_matmul_sme(dtype):
- target = "llvm -mtriple=aarch64-linux-gnu -mattr=+v9a,+sme"
-
- def check_correct_assembly(dtype):
- A = te.placeholder((32, 32), dtype=dtype, name="A")
- B = te.placeholder((32, 32), dtype=dtype, name="B")
-
- with tvm.target.Target(target):
- C = tvm.topi.arm_cpu.matmul.compute_matmul_sme(
- A, B, None, "float32", False, dtype == "float16"
- )
- prim_func = te.create_prim_func([A, B, C])
-
- sch = tvm.tir.Schedule(prim_func)
- tvm.topi.arm_cpu.matmul.tir_schedule_matmul_sme(sch)
- prim_func = sch.mod
-
- f = tvm.build(prim_func, target=target)
-
- assembly = f.get_source("asm")
- smstart = re.findall(r"smstart\t(sm|za)", assembly)
- loads = re.findall(r"ld1[whdb]\t{\s?za", assembly)
- mopa = re.findall(
- r"fmopa\tza[0-9].[shdb],( p[0-9]/[zm],)?( p[0-9]/[zm],)? z[0-9].[shdb], z[0-9].[shdb]",
- assembly,
- )
- stores = re.findall(r"st1[whdb]\t{\s?za", assembly)
- smstop = re.findall(r"smstop\t(sm|za)", assembly)
- whilelo = re.findall(r"whilelo\tp[0-9].[shdb]", assembly)
-
- assert len(smstart) > 0
- assert len(loads) > 0
- assert len(mopa) > 0
- assert len(stores) > 0
- assert len(smstop) > 0
- assert len(whilelo) > 0
-
- check_correct_assembly(dtype=dtype)
-
-
-def test_matmul_sme_no_reduction_block():
- @T.prim_func
- def prim_func(a: T.handle, b: T.handle):
- A = T.match_buffer(a, (4,))
- B = T.match_buffer(b, (4,))
- for i in range(3):
- with T.block("block"):
- vi = T.axis.remap("S", [i])
- B[vi] = A[vi]
-
- sch = tvm.tir.Schedule(prim_func)
- with pytest.raises(AssertionError, match="Expected a single gemm reduction block."):
- tvm.topi.arm_cpu.matmul.tir_schedule_matmul_sme(sch)
-
-
@pytest.mark.skipif(
llvm_version_major() < 11, reason="Vscale is not supported in earlier versions of LLVM"
)
@@ -636,206 +576,7 @@ def test_vscale_range_function_attribute(mattr, expect_attr):
), f"Unexpected function attribute vscale_range() was found in generated LLVM IR"
-@pytest.mark.skipif(
- llvm_version_major() < 16, reason="Test requires an LLVM version of at least 16 to target SME"
-)
-@pytest.mark.parametrize(
- "attr_key,attr_value,expected",
- [
- (
- SMEAttributes.STREAMING_MODE,
- SMEAttributes.StreamingModeValues.ENABLED,
- "aarch64_pstate_sm_enabled",
- ),
- (
- SMEAttributes.STREAMING_MODE,
- SMEAttributes.StreamingModeValues.COMPATIBLE,
- "aarch64_pstate_sm_compatible",
- ),
- (SMEAttributes.ZA_STORAGE, SMEAttributes.ZAStorageValues.NEW, "aarch64_pstate_za_new"),
- (
- SMEAttributes.ZA_STORAGE,
- SMEAttributes.ZAStorageValues.SHARED,
- "aarch64_pstate_za_shared",
- ),
- ],
-)
-def test_function_attributes(attr_key, attr_value, expected):
- target = "llvm -mtriple=aarch64-linux-gnu -mattr=+sme"
-
- @T.prim_func
- def prim_func(a: T.handle, c: T.handle):
- T.func_attr({"global_symbol": "main", "tir.noalias": T.bool(True)})
- A = T.match_buffer(a, (16,), "float32")
- C = T.match_buffer(c, (1,), "float32")
-
- with T.block("extern"):
- T.block_attr({attr_key: attr_value})
- for i in range(16):
- C[0] += A[i]
-
- func = tvm.build(prim_func, target=target)
- ll = func.get_source("ll")
-
- # Check that the attribute exists
- attr = re.findall(rf".*{expected}*.", ll)
- assert attr, f"Function attribute {expected} was not found in generated LLVM IR"
-
- # Check this attribute is used on the "compute" function
- func_attr_label = attr[0].split(" ")[1]
- found_compute_func = False
- for match in re.findall(rf".*{func_attr_label}*.", ll):
- if "_compute_" in match:
- found_compute_func = True
-
- assert found_compute_func, (
- f"The attribute {expected} was found to be under the label {func_attr_label}, "
- "but it was not used by the 'compute' scope function."
- )
-
-
-def test_unsupported_function_attribute_type():
- target = "llvm -mtriple=aarch64-linux-gnu -mattr=+sme"
-
- @T.prim_func
- def prim_func(a: T.handle, c: T.handle):
- T.func_attr({"global_symbol": "main", "tir.noalias": T.bool(True)})
- A = T.match_buffer(a, (16,), "float32")
- C = T.match_buffer(c, (1,), "float32")
-
- with T.block("extern"):
- T.block_attr({SMEAttributes.STREAMING_MODE: True})
- with T.block("root"):
- for i in range(16):
- C[0] += A[i]
-
- err_msg = f"Expect {SMEAttributes.STREAMING_MODE} to have a String value but was IntImm"
- with pytest.raises(tvm.error.TVMError, match=err_msg):
- tvm.build(prim_func, target=target)
-
-
-@pytest.mark.parametrize(
- "attr_key,attr_value",
- [
- (SMEAttributes.STREAMING_MODE, SMEAttributes.StreamingModeValues.ENABLED),
- (SMEAttributes.ZA_STORAGE, SMEAttributes.ZAStorageValues.NEW),
- ],
-)
-def test_unsupported_multiple_function_attributes(attr_key, attr_value):
- target = "llvm -mtriple=aarch64-linux-gnu -mattr=+sme"
-
- @T.prim_func
- def prim_func(a: T.handle, c: T.handle):
- A = T.match_buffer(a, (16,), "float32")
- C = T.match_buffer(c, (1,), "float32")
-
- with T.block("root"):
- with T.block("extern"):
- T.block_attr({attr_key: attr_value})
- for i in range(16):
- C[0] += A[i] * 2
- with T.block("extern2"):
- T.block_attr({attr_key: attr_value})
- for i in range(16):
- C[0] += A[i] * 3
-
- err_msg = f"Multiple definitions of {attr_key} attribute found in the function default_function_compute_"
- with pytest.raises(tvm.error.TVMError, match=err_msg):
- tvm.build(prim_func, target=target)
-
-
-@pytest.mark.skipif(
- llvm_version_major() < 15, reason="Test requires an LLVM version of at least 15 to target SVE"
-)
-@pytest.mark.parametrize("dtype", ["float16", "float32"])
-@pytest.mark.parametrize(
- "conv2d_impl",
- [
- (
- tvm.topi.arm_cpu.compute_conv2d_NHWC_hybrid_SVE,
- tvm.topi.arm_cpu.schedule_conv2d_NHWC_hybrid_SVE,
- False,
- ),
- (
- tvm.topi.arm_cpu.compute_conv2d_NHWC_hybrid_SVE,
- tvm.topi.arm_cpu.schedule_conv2d_NHWC_hybrid_TIR,
- True,
- ),
- ],
-)
-def test_conv2d_sve(dtype, conv2d_impl):
- target = "llvm -mtriple=aarch64-linux-gnu -mattr=+sve"
-
- def check_correct_assembly(dtype, compute, schedule, use_tir_schedule):
- A = te.placeholder((1, 32, 32, 3), dtype=dtype, name="A")
- W = te.placeholder((3, 3, 3, 8), dtype=dtype, name="B")
- stride = padding = dilation = 1
- B = compute(A, W, stride, padding, dilation, dtype)
- if use_tir_schedule:
- func = te.create_prim_func([A, W, B])
- sch = schedule(tvm.tir.Schedule(func))
- f = tvm.build(sch.mod["main"], target)
- else:
- s = schedule([B])
- f = tvm.build(s, [A, W, B], target)
- assembly = f.get_source("asm")
-
- loads = re.findall(r"ld1[r]?[q]?[whdb]\t{\s?z", assembly)
- compute_ops = re.findall(
- r"fm(la|ad)\tz\d+.[shdb], (p\d+\/[zm], )?z\d+.[shdb], z\d+.[shdb]",
- assembly,
- )
- stores = re.findall(r"st1[whdb]\t{\s?z", assembly)
-
- assert len(loads) > 0
- assert len(compute_ops) > 0
- assert len(stores) > 0
-
- with tvm.target.Target(target):
- check_correct_assembly(dtype, *conv2d_impl)
-
-
-@pytest.mark.skipif(
- llvm_version_major() < 16, reason="Test requires an LLVM version of at least 16 to target SME"
-)
-@pytest.mark.parametrize("dtype", ["float32"])
-def test_conv2d_sme(dtype):
- target = "llvm -mtriple=aarch64-linux-gnu -mattr=+v9a,+sme"
-
- def check_correct_assembly(dtype):
- A = te.placeholder((1, 32, 32, 3), dtype=dtype, name="A")
- W = te.placeholder((3, 3, 3, 8), dtype=dtype, name="B")
- stride = padding = dilation = 1
-
- B = tvm.topi.arm_cpu.compute_conv2d_NHWC_hybrid_SME(A, W, stride, padding, dilation, dtype)
- func = te.create_prim_func([A, W, B])
- sch = tvm.topi.arm_cpu.schedule_conv2d_NHWC_hybrid_TIR(tvm.tir.Schedule(func))
- f = tvm.build(sch.mod["main"], target)
-
- assembly = f.get_source("asm")
- smstart = re.findall(r"smstart\t(sm|za)", assembly)
- loads = re.findall(r"ld1[whdb]\t{\s?za", assembly)
- mopa = re.findall(
- r"fmopa\tza[0-9].[shdb],( p[0-9]/[zm],)?( p[0-9]/[zm],)? z[0-9].[shdb], z[0-9].[shdb]",
- assembly,
- )
- stores = re.findall(r"st1[whdb]\t{\s?za", assembly)
- smstop = re.findall(r"smstop\t(sm|za)", assembly)
- whilelo = re.findall(r"whilelo\tp[0-9].[shdb]", assembly)
-
- assert len(smstart) > 0
- assert len(loads) > 0
- assert len(mopa) > 0
- assert len(stores) > 0
- assert len(smstop) > 0
- assert len(whilelo) > 0
-
- with tvm.target.Target(target):
- check_correct_assembly(dtype=dtype)
-
-
-@pytest.mark.skipif(
- llvm_version_major() < 11,
+@pytest.mark.skip(
reason="Vscale and get.active.lane.mask are not supported in earlier versions of LLVM",
)
def test_get_active_lane_mask():
@@ -854,8 +595,7 @@ def test_get_active_lane_mask():
assert "get.active.lane.mask" in ll
-@pytest.mark.skipif(
- llvm_version_major() < 11,
+@pytest.mark.skip(
reason="Vscale and get.active.lane.mask are not supported in earlier versions of LLVM",
)
def test_predicated_scalable_buffer():
diff --git a/tests/python/codegen/test_target_codegen_blob.py b/tests/python/codegen/test_target_codegen_blob.py
index 5266f481f5..a61e6e894c 100644
--- a/tests/python/codegen/test_target_codegen_blob.py
+++ b/tests/python/codegen/test_target_codegen_blob.py
@@ -17,53 +17,12 @@
import ctypes
import numpy as np
-from tvm import relay
-import tvm.relay.testing
-from tvm.contrib import graph_executor, cc, utils, popen_pool, tar
+from tvm.contrib import cc, utils, popen_pool, tar
import tvm
import tvm.testing
from tvm.script import ir as I, tir as T
-@tvm.testing.uses_gpu
-def test_synthetic():
- for device in ["llvm", "cuda"]:
- if not tvm.testing.device_enabled(device):
- print("skip because %s is not enabled..." % device)
- return
-
- input_shape = (1, 5, 23, 61)
-
- def verify(data):
- mod, params = relay.testing.synthetic.get_workload(input_shape=input_shape)
- with tvm.transform.PassContext(opt_level=3):
- lib = relay.build_module.build(mod, "llvm", params=params)
- dev = tvm.cpu()
- module = graph_executor.GraphModule(lib["default"](dev))
- module.set_input("data", data)
- module.run()
- out = module.get_output(0).numpy()
- return out
-
- synthetic_mod, synthetic_params = relay.testing.synthetic.get_workload(input_shape=input_shape)
- with tvm.transform.PassContext(opt_level=3):
- synthetic_gpu_lib = relay.build_module.build(synthetic_mod, "cuda", params=synthetic_params)
-
- temp = utils.tempdir()
- path_lib = temp.relpath("deploy_lib.so")
- synthetic_gpu_lib.export_library(path_lib)
-
- loaded_lib = tvm.runtime.load_module(path_lib)
- data = np.random.uniform(-1, 1, size=input_shape).astype("float32")
- dev = tvm.cuda()
- module = graph_executor.GraphModule(loaded_lib["default"](dev))
- module.set_input("data", data)
- module.run()
- out = module.get_output(0).numpy()
-
- tvm.testing.assert_allclose(out, verify(data), atol=1e-5)
-
-
@tvm.testing.uses_gpu
def test_cuda_multi_lib():
# test combining two system lib together
diff --git a/tests/python/codegen/test_target_codegen_cuda.py b/tests/python/codegen/test_target_codegen_cuda.py
index 112d1151fe..764328a49b 100644
--- a/tests/python/codegen/test_target_codegen_cuda.py
+++ b/tests/python/codegen/test_target_codegen_cuda.py
@@ -455,68 +455,6 @@ def test_cuda_const_float_to_half():
np.testing.assert_equal(c.numpy(), a_np > b.value)
-@tvm.testing.requires_gpu
-@tvm.testing.requires_cuda
-def test_cuda_reduction():
- def check(device, dtype, m=32, n=32):
- if not tvm.testing.device_enabled(device):
- print("Skipping", device)
- return
- dev = tvm.device(device, 0)
- a = te.placeholder((m, n), name="a", dtype=dtype)
- b = te.placeholder((m, n), name="b", dtype=dtype)
- c = a + b
- d = a * b
- e = topi.elemwise_sum([c, d])
- g = topi.sum(e)
- with tvm.target.Target(device):
- sg = topi.cuda.schedule_reduce(g)
- func = tvm.build(sg, [a, b, g], device)
- a_np = np.random.uniform(size=(m, n)).astype(a.dtype)
- b_np = np.random.uniform(size=(m, n)).astype(b.dtype)
- g_np = np.sum(np.add(a_np * b_np, a_np + b_np))
- a_nd = tvm.nd.array(a_np, dev)
- b_nd = tvm.nd.array(b_np, dev)
- g_nd = tvm.nd.array(np.zeros(g_np.shape, dtype=g_np.dtype), dev)
- func(a_nd, b_nd, g_nd)
- tvm.testing.assert_allclose(g_nd.numpy(), g_np, rtol=1e-3)
-
- check("cuda", "float32")
- check("rocm", "float32")
- check("cuda", "float16")
-
-
-@tvm.testing.requires_gpu
-@tvm.testing.requires_cuda
-def test_cuda_mix_threaded_and_normal_reduction():
- def check(device, dtype, m=32, n=32):
- if not tvm.testing.device_enabled(device):
- print("Skipping", device)
- return
- dev = tvm.device(device, 0)
- if dtype == "float16" and not have_fp16(dev.compute_version):
- print("Skip because gpu does not have fp16 support")
- return
-
- a = tvm.te.placeholder((m, n), name="a", dtype=dtype)
- b = topi.sum(a)
- with tvm.target.Target(device):
- sb = tvm.te.create_schedule(b.op)
- i, _ = b.op.reduce_axis
- sb[b].bind(i, tvm.te.thread_axis("threadIdx.x"))
- func = tvm.build(sb, [a, b], device)
- a_np = np.random.uniform(size=(m, n)).astype(a.dtype)
- b_np = np.sum(a_np)
- a_nd = tvm.nd.array(a_np, dev)
- b_nd = tvm.nd.array(np.zeros(b_np.shape, dtype=b_np.dtype), dev)
- func(a_nd, b_nd)
- tvm.testing.assert_allclose(b_nd.numpy(), b_np, rtol=1e-3)
-
- check("cuda", "float32")
- check("rocm", "float32")
- check("cuda", "float16")
-
-
@tvm.testing.requires_gpu
@tvm.testing.requires_cuda
def test_cuda_floordiv_with_vectorization():
diff --git a/tests/python/codegen/test_target_codegen_hexagon.py b/tests/python/codegen/test_target_codegen_hexagon.py
index 344c7a9762..c97637f927 100644
--- a/tests/python/codegen/test_target_codegen_hexagon.py
+++ b/tests/python/codegen/test_target_codegen_hexagon.py
@@ -21,7 +21,6 @@ import pytest
import re
import sys
import tvm
-import tvm.relay
import tvm.testing
import tvm.contrib.hexagon as hexagon
diff --git a/tests/python/codegen/test_target_codegen_llvm.py b/tests/python/codegen/test_target_codegen_llvm.py
index e8036467ff..d629d93d36 100644
--- a/tests/python/codegen/test_target_codegen_llvm.py
+++ b/tests/python/codegen/test_target_codegen_llvm.py
@@ -27,7 +27,6 @@ import tvm
import tvm.testing
from tvm import te
from tvm.contrib import clang, utils
-from tvm.relay.backend import Runtime
from tvm.script import tir as T, ir as I
from tvm.target.codegen import llvm_get_intrinsic_name, llvm_lookup_intrinsic_id
@@ -754,12 +753,10 @@ def test_llvm_crt_static_lib():
A = te.placeholder((32,), dtype="bfloat16")
B = te.placeholder((32,), dtype="bfloat16")
d = te.compute((32,), lambda x: A[x] + B[x])
- sch = te.create_schedule(d.op)
+ mod = tvm.IRModule.from_expr(te.create_prim_func([A, B, d]))
module = tvm.build(
- sch,
- [A, B, d],
+ mod.with_attr("system_lib_prefix", ""),
target=tvm.target.Target("llvm"),
- runtime=Runtime("crt", {"system-lib": True}),
)
print(module.get_source())
module.save("test.o")
diff --git a/tests/python/codegen/test_target_codegen_opencl.py b/tests/python/codegen/test_target_codegen_opencl.py
index 9222947ae4..079553665f 100644
--- a/tests/python/codegen/test_target_codegen_opencl.py
+++ b/tests/python/codegen/test_target_codegen_opencl.py
@@ -15,7 +15,7 @@
# specific language governing permissions and limitations
# under the License.
import tvm
-from tvm import te, relay
+from tvm import te
import tvm.testing
import re
import pytest
@@ -234,197 +234,5 @@ def _get_maximum_kernel_args(source):
return max_args
-def _validate_opencl_executors(executor_type, get_model, ref_impl):
- from tvm.contrib import graph_executor
- from tvm.runtime.vm import VirtualMachine
-
- input_dict, model = get_model()
- if executor_type == "ge":
- with tvm.transform.PassContext(opt_level=3):
- lib = relay.build(model, target_host="llvm", target=target)
- ocl_lib = lib.get_lib()
- else:
- module = tvm.IRModule({})
- module["main"] = model
- with tvm.transform.PassContext(opt_level=3):
- lib = relay.vm.compile(module, target=target, target_host="llvm")
- ocl_lib = lib.module.imported_modules[0]
- opencl_modules = list(filter(lambda mod: mod.type_key == "opencl", ocl_lib.imported_modules))
- assembly = opencl_modules[0].get_source()
- with tvm.target.Target(target):
- limit = tvm.target.Target.current().max_function_args
- max_num = _get_maximum_kernel_args(assembly)
- assert max_num <= limit
-
- dev = tvm.cl()
- if executor_type == "ge":
- module = graph_executor.GraphModule(lib["default"](dev))
- module.set_input(**input_dict)
- module.run()
- tvm_out = module.get_output(0)
- else:
- vm = VirtualMachine(lib, dev, "naive")
- data = {}
- for k, v in input_dict.items():
- data[k] = tvm.nd.array(v, dev)
- vm.set_input("main", **data)
- vm.invoke_stateful("main")
- tvm_out = vm.get_outputs()[0]
-
- np_result = ref_impl(list(input_dict.values()))
- np.testing.assert_allclose(tvm_out.asnumpy(), np_result, rtol=1e-2, atol=1e-2)
-
-
-shape_type = tvm.testing.parameter("dynamic", "static")
-executor_type = tvm.testing.parameter("ge", "vm")
-
-
-@tvm.testing.requires_gpu
-@tvm.testing.requires_opencl
-def test_opencl_args_split(executor_type, shape_type):
- def _get_model():
- if shape_type == "dynamic":
- shape = (tvm.tir.Any(), 1, 1, 3)
- else:
- shape = (1, 1, 1, 3)
- shape_np = (1, 1, 1, 3)
- dtype = "float32"
- axis = 1
- tensors_num = 300
- inputs = []
- inputs_np = {}
- for i in range(tensors_num):
- inputs.append(relay.var("p{}".format(i), shape=shape, dtype=dtype))
- inputs_np[f"p{i}"] = np.random.uniform(size=shape_np).astype(dtype)
-
- inp = relay.Tuple(inputs)
- concat = relay.op.concatenate(inp, axis)
- return inputs_np, relay.Function(inputs, concat)
-
- def ref_impl(inputs):
- axis = 1
- return np.concatenate(tuple(inputs), axis=axis)
-
- if executor_type == "ge" and shape_type == "dynamic":
- pytest.skip()
- _validate_opencl_executors(executor_type, _get_model, ref_impl)
-
-
-@tvm.testing.requires_gpu
-@tvm.testing.requires_opencl
-def test_opencl_fuse_max_args(executor_type, shape_type):
- if shape_type == "dynamic":
- shape = (tvm.tir.Any(), 20)
- ops_num = 80
- else:
- shape = (1, 20)
- ops_num = 300
- shape_np = (1, 20)
- dtype = "float32"
-
- def _base_func(name):
- x = relay.var(name, shape=shape)
- y = relay.add(x, relay.const(1, "float32"))
- w = relay.exp(y)
- return x, w
-
- def _get_model():
- inp = []
- inputs_np = {}
- out = []
- for i in range(ops_num):
- x, w = _base_func(f"x{i}")
- inputs_np[f"x{i}"] = np.random.uniform(size=shape_np).astype(dtype)
- inp.append(x)
- out.append(w)
- w = out[0]
- for i in range(len(out) - 1):
- w = relay.add(w, out[i + 1])
- return inputs_np, relay.Function(inp, w)
-
- def ref_impl(inputs):
- w = np.exp(inputs[0] + 1)
- for i in range(len(inputs) - 1):
- w = w + np.exp(inputs[i + 1] + 1)
- return w
-
- if executor_type == "ge" and shape_type == "dynamic":
- pytest.skip()
- _validate_opencl_executors(executor_type, _get_model, ref_impl)
-
-
-@tvm.testing.requires_gpu
-@tvm.testing.requires_opencl
-def test_fuse_concat_max_num_args(executor_type, shape_type):
- """
- In this test, we have an operation with 3 inputs before concat. In the
- SplitArgs we cannot calculate these inputs as inputs to the concat layer,
- because they will be added to the concat after the fusing operation. So
- FuseOps pass should handle this case and stop fusing before the concat
- layer.
-
- The example:
- x y z x y z
- \ | / \ | /
- \ | / \ | /
- where ... where
- | |
- exp exp
- \ /
- \ /
- \-----> concat <-----/
- """
- if shape_type == "dynamic":
- shape = (tvm.tir.Any(), 20)
- ops_num = 80
- else:
- shape = (10, 20)
- ops_num = 300
- shape_np = (10, 20)
- dtype = "float32"
- axis = 1
-
- def _base_func(name):
- x = relay.var(name, shape=shape)
- y = relay.var(f"y{name}", shape=shape)
- z = relay.var(f"z{name}", shape=shape)
- cond = relay.less(x, relay.const(1, "float32"))
- l = relay.add(y, relay.const(1, "float32"))
- r = relay.add(z, relay.const(5, "float32"))
- w = relay.where(cond, l, r)
- w = relay.exp(w)
- return [x, y, z], w
-
- def _get_model():
- inp = []
- out = []
- inputs_np = {}
- for i in range(ops_num):
- inputs, w = _base_func(f"x{i}")
- inputs_np[f"x{i}"] = np.random.uniform(size=shape_np).astype(dtype)
- inputs_np[f"yx{i}"] = np.random.uniform(size=shape_np).astype(dtype)
- inputs_np[f"zx{i}"] = np.random.uniform(size=shape_np).astype(dtype)
- inp.extend(inputs)
- out.append(w)
- t = relay.Tuple(out)
- w = relay.op.concatenate(t, axis)
- return inputs_np, relay.Function(inp, w)
-
- def ref_impl(inputs):
- res = []
- for i in range(0, len(inputs), 3):
- x = inputs[i]
- y = inputs[i + 1]
- z = inputs[i + 2]
- comp = np.where(x < 1, y + 1, z + 5)
- comp = np.exp(comp)
- res.append(comp)
- return np.concatenate(tuple(res), axis=axis)
-
- if executor_type == "ge" and shape_type == "dynamic":
- pytest.skip()
- _validate_opencl_executors(executor_type, _get_model, ref_impl)
-
-
if __name__ == "__main__":
tvm.testing.main()
diff --git a/tests/python/codegen/test_target_codegen_vulkan.py b/tests/python/codegen/test_target_codegen_vulkan.py
index 6973040cb2..2da1023869 100644
--- a/tests/python/codegen/test_target_codegen_vulkan.py
+++ b/tests/python/codegen/test_target_codegen_vulkan.py
@@ -26,7 +26,7 @@ import pytest
import tvm
import tvm.testing
-from tvm import relay, te
+from tvm import te
from tvm.topi.math import cast
from tvm.script import tir as T, ir as I
from tvm.tir import TensorIntrin, IntImm, Cast, Schedule
@@ -233,59 +233,6 @@ def test_vulkan_bool_load(target, dev):
tvm.testing.assert_allclose(b.numpy(), ref)
-def check_mod(target, dev, mod, x_np, res_np):
- res = relay.create_executor("vm", mod=mod, device=dev, target=target).evaluate()(x_np).numpy()
- tvm.testing.assert_allclose(res, res_np, atol=1e-5)
-
-
-def test_sqrt(target, dev):
- # Three 32 bit pushconstants: any_dim, stride, stride
- dtype = "float32"
- x = relay.var("x", shape=(relay.Any(),), dtype=dtype)
- mod = tvm.IRModule()
- mod["main"] = relay.Function([x], relay.sqrt(x))
- x_np = np.random.uniform(size=(10,)).astype(dtype)
- res_np = np.sqrt(x_np)
-
- check_mod(target, dev, mod, x_np, res_np)
-
-
-def test_argsort(target, dev):
- # One 64 bit and one 32 bit constants
- dtype = "int32"
- x = relay.var("x", shape=(relay.Any(),), dtype=dtype)
- mod = tvm.IRModule()
- mod["main"] = relay.Function([x], relay.argsort(x))
- x_np = np.random.randint(0, high=10, size=(10,)).astype(dtype)
- res_np = np.argsort(x_np, kind="stable")
-
- check_mod(target, dev, mod, x_np, res_np)
-
-
-def test_cumsum(target, dev):
- # One 64 bit and one 32 bit constants
- dtype = "int32"
- x = relay.var("x", shape=(relay.Any(),), dtype=dtype)
- mod = tvm.IRModule()
- mod["main"] = relay.Function([x], relay.cumsum(x))
- x_np = np.random.randint(0, high=10, size=(10,)).astype(dtype)
- res_np = np.cumsum(x_np)
-
- check_mod(target, dev, mod, x_np, res_np)
-
-
-@tvm.testing.skip_if_wheel_test
-def test_unique(target, dev):
- dtype = "int32"
- x = relay.var("x", shape=(relay.Any(),), dtype=dtype)
- mod = tvm.IRModule()
- [unique, _, _, num_unique] = relay.unique(x, is_sorted=True)
- mod["main"] = relay.Function([x], relay.op.strided_slice(unique, begin=[0], end=num_unique))
- x_np = np.random.randint(0, high=10, size=(10,)).astype(dtype)
- res_np = np.unique(x_np)
- check_mod(target, dev, mod, x_np, res_np)
-
-
vulkan_parameter_impl = tvm.testing.parameter("push_constants", "ubo")
vulkan_parameter_dtype = tvm.testing.parameter("int32", "float32", "int64")
diff --git a/tests/python/codegen/test_target_codegen_x86.py b/tests/python/codegen/test_target_codegen_x86.py
index 8ff9dbb3dd..a276940050 100644
--- a/tests/python/codegen/test_target_codegen_x86.py
+++ b/tests/python/codegen/test_target_codegen_x86.py
@@ -64,48 +64,5 @@ def test_fp16_to_fp32():
is_32bit = platform.architecture()[0] == "32bit"
-@tvm.testing.requires_llvm
-@pytest.mark.skipif(is_32bit, reason=f"Fails in CI due to architecture mismatch in JIT")
-@pytest.mark.parametrize("feature_string", ["-sse2", "+sse2"])
-def test_fp16_fp32_conversions(feature_string):
- relay_model = textwrap.dedent(
- """
- #[version = "0.0.5"]
- def @main(%inp : Tensor[(3), float32], %cst : Tensor[(3), float32]) {
- %1 = cast(%inp, dtype="float16");
- %2 = cast(%cst, dtype="float16");
- %3 = add(%1, %2);
- %4 = cast(%3, dtype="float32");
- %4
- }
- """
- )
-
- ir_mod = tvm.relay.fromtext(relay_model)
-
- arch = "i386" if machine == "i386" else "x86_64"
- aot_factory = tvm.relay.build(
- ir_mod,
- params={"cst": np.array([1.0, 2.0, 3.0], dtype="float32")},
- target=f"llvm --mtriple={arch} --mattr={feature_string}",
- executor=tvm.relay.backend.Executor(
- "aot", {"interface-api": "packed", "unpacked-api": False}
- ),
- )
-
- mod_name = aot_factory["list_module_names"]()[0]
- executor = aot_factory[mod_name]
- mod = executor(tvm.cpu(0))
-
- inp = tvm.nd.array(np.array([1.1, 2.1, 3.1], dtype="float32"), device=tvm.cpu(0))
-
- mod.get_function("set_input")(0, inp)
- mod.get_function("run")()
- out = mod.get_function("get_output")(0)
-
- expected = np.array([2.1, 4.1, 6.1], dtype="float32")
- np.testing.assert_allclose(out.asnumpy(), expected, rtol=1e-3)
-
-
if __name__ == "__main__":
test_fp16_to_fp32()
diff --git a/tests/python/codegen/test_target_texture_codegen_opencl.py b/tests/python/codegen/test_target_texture_codegen_opencl.py
deleted file mode 100644
index 213ad8d7ba..0000000000
--- a/tests/python/codegen/test_target_texture_codegen_opencl.py
+++ /dev/null
@@ -1,1778 +0,0 @@
-# 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.
-import sys
-
-import numpy as np
-import pytest
-
-import tvm
-import tvm.testing
-from tvm import autotvm
-from tvm import te
-from tvm.topi import testing
-from tvm.topi.utils import get_const_tuple, simplify
-from tvm.topi import nn
-
-
-def compute_plus_one_rank3(shape):
- X = te.placeholder(shape, name="X", dtype="float32")
- Y = te.compute(shape, lambda i, j, k: X[i, j, k] + 1, name="Compute_Y")
- return X, Y
-
-
-def schedule_plus_one_rank3(X, Y):
- s = te.create_schedule(Y.op)
- # Xt = s.cache_read(X, "texture", [Y])
- # Xt = s.cache_read(X, "global", [Y])
- Xt = s.cache_read(X, "global.texture", [Y])
-
- # copy to texture stage
- x, y, c = s[Xt].op.axis
- s[Xt].bind(x, te.thread_axis("blockIdx.x"))
- s[Xt].bind(y, te.thread_axis("threadIdx.x"))
- s[Xt].vectorize(c)
-
- # the compute stage
- x, y, c = s[Y].op.axis
- xo, yo, xi, yi = s[Y].tile(x, y, 4, 4)
- s[Y].bind(xo, te.thread_axis("blockIdx.x"))
- s[Y].bind(yo, te.thread_axis("threadIdx.x"))
- s[Y].vectorize(c)
- return s
-
-
-def compute_plus_one_rank5(shape):
- X = te.placeholder(shape, name="X", dtype="float32")
- Y = te.compute(shape, lambda i, j, k, l, m: X[i, j, k, l, m] + 1, name="Compute_Y")
- return X, Y
-
-
-def schedule_plus_one_rank5(X, Y):
- s = te.create_schedule(Y.op)
- Xt = s.cache_read(X, "global.texture", [Y])
-
- # copy to texture stage
- a, b, c, d, e = s[Xt].op.axis
- abc = s[Xt].fuse(a, b, c)
- s[Xt].bind(abc, te.thread_axis("blockIdx.x"))
- s[Xt].bind(d, te.thread_axis("threadIdx.x"))
- s[Xt].vectorize(e)
-
- # the compute stage
- a, b, c, d, e = s[Y].op.axis
- abc = s[Y].fuse(a, b, c)
- xo, yo, xi, yi = s[Y].tile(abc, d, 4, 4)
- s[Y].bind(xo, te.thread_axis("blockIdx.x"))
- s[Y].bind(yo, te.thread_axis("threadIdx.x"))
- s[Y].vectorize(e)
- return s
-
-
-def compute_matmul(shape):
- A = te.placeholder(shape, name="A", dtype="float32")
- B = te.placeholder(shape, name="B", dtype="float32")
- k = te.reduce_axis((0, shape[1]), name="k")
- C = te.compute(
- (shape[0] * shape[2], shape[0] * shape[2]),
- lambda i, j: te.sum(
- A[i // shape[2], k, i % shape[2]].astype("float32")
- * B[j // shape[2], k, j % shape[2]].astype("float32"),
- axis=[k],
- ),
- name="Compute_MatMul",
- )
- return A, B, C
-
-
-def schedule_matmul(A, B, C, local=False):
- s = te.create_schedule(C.op)
- At = s.cache_read(A, "global.texture", [C])
- Bt = s.cache_read(B, "global.texture", [C])
- if local:
- Al = s.cache_read(At, "local", [C])
- Bl = s.cache_read(Bt, "local", [C])
- Cl = s.cache_write(C, "local")
-
- bx = te.thread_axis("blockIdx.x")
- tx = te.thread_axis("threadIdx.x")
-
- def copy_to_texture(stage):
- _io, _k, _ii = s[stage].op.axis
- s[stage].vectorize(_ii)
- s[stage].bind(_io, bx)
- s[stage].bind(_k, tx)
-
- copy_to_texture(At)
- copy_to_texture(Bt)
-
- # copy to global stage
- _i, _j = s[C].op.axis
- xo, yo, xi, yi = s[C].tile(_i, _j, 4, 4)
- s[C].unroll(xi)
- s[C].vectorize(yi)
- s[C].bind(xo, te.thread_axis("blockIdx.x"))
- s[C].bind(yo, te.thread_axis("threadIdx.x"))
-
- # the compute stage
- s[Cl].compute_at(s[C], yo)
- (_k,) = Cl.op.reduce_axis
- _x, _y = s[Cl].op.axis
- s[Cl].reorder(_k, _x, _y)
- s[Cl].unroll(_x)
- s[Cl].vectorize(_y)
-
- if local:
- s[Al].compute_at(s[Cl], _k)
- s[Al].vectorize(s[Al].op.axis[-1])
- s[Bl].compute_at(s[Cl], _k)
- s[Bl].vectorize(s[Bl].op.axis[-1])
-
- return s
-
-
-def compute_matmul_inner(shape):
- A = te.placeholder(shape, name="A", dtype="float32")
- B = te.placeholder(shape, name="B", dtype="float32")
- k = te.reduce_axis((0, shape[1] * shape[2]), name="k")
- # (M, K) x (N, K)
- # (32, 256) x (32, 256)
- # (32, 64, 4) x (32, 64, 4)
- C = te.compute(
- (shape[0], shape[0]),
- lambda i, j: te.sum(
- A[i, k // shape[2], k % shape[2]].astype("float32")
- * B[j, k // shape[2], k % shape[2]].astype("float32"),
- axis=[k],
- ),
- name="Compute_MatMul",
- )
- return A, B, C
-
-
-def schedule_matmul_inner(A, B, C, local=False):
- s = te.create_schedule(C.op)
- At = s.cache_read(A, "global.texture", [C])
- Bt = s.cache_read(B, "global.texture", [C])
- if local:
- Al = s.cache_read(At, "local", [C])
- Bl = s.cache_read(Bt, "local", [C])
- Cl = s.cache_write(C, "local")
-
- bx = te.thread_axis("blockIdx.x")
- tx = te.thread_axis("threadIdx.x")
-
- def copy_to_texture(stage):
- _i, _ko, _ki = s[stage].op.axis
- s[stage].vectorize(_ki)
- s[stage].bind(_i, bx)
- s[stage].bind(_ko, tx)
-
- copy_to_texture(At)
- copy_to_texture(Bt)
-
- # copy to global stage
- _i, _j = s[C].op.axis
- xo, yo, xi, yi = s[C].tile(_i, _j, 4, 4)
- s[C].unroll(xi)
- s[C].vectorize(yi)
- s[C].bind(xo, te.thread_axis("blockIdx.x"))
- s[C].bind(yo, te.thread_axis("threadIdx.x"))
-
- # the compute stage
- s[Cl].compute_at(s[C], yo)
- (_k,) = Cl.op.reduce_axis
- _x, _y = s[Cl].op.axis
- s[Cl].reorder(_x, _y, _k)
- s[Cl].unroll(_x)
- # TODO(csullivan): consider whether the below error is worth resolving
- # s[Cl].vectorize(_y) # error
-
- if local:
- s[Al].compute_at(s[Cl], _x)
- s[Al].vectorize(s[Al].op.axis[-1])
- s[Bl].compute_at(s[Cl], _x)
- s[Bl].vectorize(s[Bl].op.axis[-1])
-
- return s
-
-
-def compute_matmul_vector_accumulator(shapeA, shapeB):
- # A x B
- # (K/4, M, K%4) x (K, N/4, N%4) = (M, N)
- # (32, 64, 4) x (128, 16, 4) = (64, 64)
- A = te.placeholder(shapeA, name="A", dtype="float32")
- B = te.placeholder(shapeB, name="B", dtype="float32")
- k = te.reduce_axis((0, shapeB[0]), name="k")
- C = te.compute(
- (shapeA[1], shapeB[1] * shapeB[2]),
- lambda i, j: te.sum(
- A[k // shapeA[-1], i, k % shapeA[-1]].astype("float32")
- * B[k, j // shapeB[-1], j % shapeB[-1]].astype("float32"),
- axis=[k],
- ),
- name="Compute_MatMul",
- )
- return A, B, C
-
-
-def schedule_matmul_vector_accumulator(A, B, C, local=False):
- s = te.create_schedule(C.op)
- At = s.cache_read(A, "global.texture", [C])
- Bt = s.cache_read(B, "global.texture", [C])
- if local:
- Al = s.cache_read(At, "local", [C])
- Bl = s.cache_read(Bt, "local", [C])
- Cl = s.cache_write(C, "local")
-
- def copy_to_texture(stage):
- _y, _x, _v = s[stage].op.axis
- # TODO(csullivan): removing this vectorize results in numerical errors, autovectorize
- s[stage].vectorize(_v)
- s[stage].bind(_y, te.thread_axis("blockIdx.x"))
- s[stage].bind(_x, te.thread_axis("threadIdx.x"))
-
- copy_to_texture(At)
- copy_to_texture(Bt)
-
- # copy to global stage
- _i, _j = s[C].op.axis
- xo, yo, xi, yi = s[C].tile(_i, _j, 4, 4)
- s[C].unroll(xi)
- s[C].vectorize(yi)
- s[C].bind(xo, te.thread_axis("blockIdx.x"))
- s[C].bind(yo, te.thread_axis("threadIdx.x"))
-
- # the compute stage
- s[Cl].compute_at(s[C], yo)
- (_k,) = Cl.op.reduce_axis
- _a, _b = s[Cl].op.axis
- _ko, _ki = s[Cl].split(_k, factor=4)
- s[Cl].reorder(_ko, _a, _ki, _b)
- s[Cl].unroll(_ki)
- s[Cl].unroll(_a)
- s[Cl].vectorize(_b)
-
- if local:
- s[Al].compute_at(s[Cl], _a)
- _aa, _ka, _ba = s[Al].op.axis
- # TODO(csullivan)[BEFORE PR]: removing this vectorize command causes a crash. This needs to be autovectorized.
- s[Al].vectorize(_ba)
- s[Bl].compute_at(s[Cl], _ko)
- _ab, _kb, _bb = s[Bl].op.axis
- s[Bl].vectorize(_bb)
- s[Bl].unroll(_ab)
-
- return s
-
-
-def compute_conv2d_1x1_NCHWc_RSCKk(input_shape, filter_shape):
- # conv2d( [N, C, H, W, c] , [1, 1, C, K, k]
- data = te.placeholder(input_shape, name="data", dtype="float32")
- filt = te.placeholder(filter_shape, name="filter", dtype="float32")
- c = te.reduce_axis((0, input_shape[1]), name="C")
- c4 = te.reduce_axis((0, input_shape[-1]), name="c4")
- kh = te.reduce_axis((0, filter_shape[0]), name="kh")
- kw = te.reduce_axis((0, filter_shape[1]), name="kw")
- conv = te.compute(
- (input_shape[0], filter_shape[-2], input_shape[2], input_shape[3], filter_shape[-1]),
- lambda n, ko, i, j, ki: te.sum(
- data[n, c, i, j, c4].astype("float32")
- * filt[kh, kw, c * input_shape[-1] + c4, ko, ki].astype("float32"),
- axis=[kh, kw, c, c4],
- ),
- # name="Compute_conv2d_1x1_NCHWc_RSCKk",
- name="conv2d_1x1",
- )
- return data, filt, conv
-
-
-def schedule_conv2d_1x1_NCHWc_RSCKk(data, filt, conv):
- # inputs: (1, 128//4, 56, 56, 4), (1, 1, 128, 128//4, 4)
- # outputs:
- s = te.create_schedule(conv.op)
- A, B, C = data, filt, conv
- At = s.cache_read(A, "global.texture", [C])
- Bt = s.cache_read(B, "global.texture", [C])
- Al = s.cache_read(At, "local", [C])
- Bl = s.cache_read(Bt, "local", [C])
- Cl = s.cache_write(C, "local")
-
- def copy_to_texture(stage):
- axes = s[stage].op.axis
- fused = s[stage].fuse(*axes[:-1])
- block, thread = s[stage].split(fused, factor=32)
- s[stage].vectorize(axes[-1])
- s[stage].bind(block, te.thread_axis("blockIdx.x"))
- s[stage].bind(thread, te.thread_axis("threadIdx.x"))
-
- copy_to_texture(At)
- copy_to_texture(Bt)
-
- _n, _ko, _h, _w, _ki = s[C].op.axis
- s[C].vectorize(_ki)
- s[C].bind(_n, te.thread_axis("blockIdx.x"))
- s[C].bind(_ko, te.thread_axis("threadIdx.x"))
-
- s[Cl].compute_at(s[C], _w)
- _nl, _kol, _hl, _wl, _kil = s[Cl].op.axis
- _khl, _kwl, _cl, _cl4 = s[Cl].op.reduce_axis
- _clo, _cli = s[Cl].split(_cl, factor=4)
- s[Cl].reorder(_clo, _cli, _cl4, _kil)
- s[Cl].unroll(_cli)
- s[Cl].unroll(_cl4)
- s[Cl].vectorize(_kil)
-
- s[Al].compute_at(s[Cl], _cli)
- s[Al].vectorize(s[Al].op.axis[-1])
- s[Bl].compute_at(s[Cl], _kwl)
- s[Bl].vectorize(s[Bl].op.axis[-1])
-
- return s
-
-
-def compute_conv2d_1x1_WCHNc_CRSKk(input_shape, filter_shape):
- # input_shape = [W, C, H, N, c] -> [W, C, H*N, c]
- # filter_shape = [C, R, S, K, k] -> [C, R*S*K, k]
- # output_shape: [WK, HN, k] -> [W, K, H, N, k]
- data = te.placeholder(input_shape, name="data", dtype="float32")
- filt = te.placeholder(filter_shape, name="filter", dtype="float32")
-
- packed_data = te.compute(
- (input_shape[0], input_shape[1], input_shape[2] * input_shape[3], input_shape[4]),
- lambda i, j, k, l: data[i, j, k // input_shape[3], k % input_shape[3], l],
- name="packed_data",
- )
-
- # Logical transformation of Nd -> 3d tensor
- # CRSKk -> C|RSK|k
- # r = rsk // SK
- # sk = rsk % SK
- # s = sk // K == (rsk % SK) // K == (rsk // K) % S
- # k = sk % K == (rsk % SK) % K == rsk % K
- packed_filter = te.compute(
- (filter_shape[0], filter_shape[1] * filter_shape[2] * filter_shape[3], filter_shape[4]),
- lambda i, j, k: filt[
- i,
- j // (filter_shape[3] * filter_shape[2]),
- (j // filter_shape[3]) % filter_shape[2],
- j % filter_shape[3],
- k,
- ],
- name="packed_filter",
- )
-
- c = te.reduce_axis((0, input_shape[1]), name="C")
- c4 = te.reduce_axis((0, input_shape[-1]), name="c4")
- r = te.reduce_axis((0, filter_shape[1]), name="r")
- s = te.reduce_axis((0, filter_shape[2]), name="s")
-
- conv = te.compute(
- (input_shape[0], filter_shape[3], input_shape[2], input_shape[3], filter_shape[4]),
- lambda w, ko, h, n, ki: te.sum(
- packed_data[w, c, h * input_shape[3] + n, c4].astype("float32")
- * packed_filter[
- c * input_shape[-1] + c4, ((r * filter_shape[2]) + s) * filter_shape[3] + ko, ki
- ].astype("float32"),
- axis=[r, s, c, c4],
- ),
- name="conv2d_1x1",
- )
- return data, filt, packed_data, packed_filter, conv
-
-
-def schedule_conv2d_1x1_WCHNc_CRSKk(data, filt, packed_data, packed_filter, conv):
- # data: [W, C, H*N, c]
- # filter: [C, R*S*K, k]
- # output: [W, K, H, N, k]
-
- # conv2d( [N, C, H, W, c] , [1, 1, C, K, k]
- # inputs: (1, 128//4, 56, 56, 4), (1, 1, 128, 128//4, 4)
-
- # data: (56, 128//4, 56*1, 4) = (56, 32, 56, 4)
- # filt: (128, 1*1*128//4, 4) = (128, 32, 4)
- # conv: (56, 32, 56, 1, 4)
-
- s = te.create_schedule(conv.op)
- cfg = autotvm.get_config()
-
- s[packed_data].compute_inline()
- s[packed_filter].compute_inline()
- A, B, C = packed_data, packed_filter, conv
- At = s.cache_read(A, "global.texture", [C])
- Bt = s.cache_read(B, "global.texture", [C])
- Al = s.cache_read(At, "local", [C])
- Bl = s.cache_read(Bt, "local", [C])
- Cl = s.cache_write(C, "local")
-
- def copy_to_texture(stage):
- axes = s[stage].op.axis
- fused = s[stage].fuse(*axes[:-1])
- block, thread = s[stage].split(fused, factor=32)
- s[stage].vectorize(axes[-1])
- s[stage].bind(block, te.thread_axis("blockIdx.x"))
- s[stage].bind(thread, te.thread_axis("threadIdx.x"))
-
- copy_to_texture(At)
- copy_to_texture(Bt)
-
- _w, _ko, _h, _n, _ki = s[C].op.axis
- kernel_scope, _n = s[C].split(_n, nparts=1)
-
- cfg.define_split("tile_f", _ko, num_outputs=4)
- cfg.define_split("tile_w", _w, num_outputs=4)
- cfg.define_split("tile_h", _h, num_outputs=4)
- cfg.define_knob("auto_unroll_max_step", [0, 512, 1500])
-
- bk, vk, tk, ki = cfg["tile_f"].apply(s, C, _ko)
- bw, vw, tw, wi = cfg["tile_w"].apply(s, C, _w)
- bh, vh, th, hi = cfg["tile_h"].apply(s, C, _h)
- s[C].reorder(bh, _n, vh, th, hi)
- bhn = s[C].fuse(bh, _n)
-
- s[C].bind(bk, te.thread_axis("blockIdx.z"))
- s[C].bind(bhn, te.thread_axis("blockIdx.y"))
- s[C].bind(bw, te.thread_axis("blockIdx.x"))
- s[C].bind(vk, te.thread_axis("vthread"))
- s[C].bind(vh, te.thread_axis("vthread"))
- s[C].bind(vw, te.thread_axis("vthread"))
- s[C].bind(tk, te.thread_axis("threadIdx.z"))
- s[C].bind(th, te.thread_axis("threadIdx.y"))
- s[C].bind(tw, te.thread_axis("threadIdx.x"))
- s[C].reorder(bw, bk, bhn, vw, vk, vh, tw, tk, th, ki, hi, wi, _ki)
- s[C].vectorize(_ki)
-
- # TODO(csullivan): Try uneven workgroup split
- # _wo, _wi = s[C].split(_w, factor=4)
- # #_hno, _hni = s[C].split(_hn, factor=8)
- # #s[C].reorder(_wo, _wi, _ko, _hno, _hni, _ki)
- # s[C].reorder(_wo, _ko, _hn, _ki, _wi)
- # s[C].unroll(_wi)
-
- # # mace:
- # # const int out_ch_blk = get_global_id(0);
- # # const int out_w_blk = get_global_id(1);
- # # const int out_hb = get_global_id(2);
-
- # bx = te.thread_axis("blockIdx.x")
- # by = te.thread_axis("blockIdx.y")
- # bz = te.thread_axis("blockIdx.z")
- # s[C].bind(_ko, bx)
- # s[C].bind(_wo, by)
- # s[C].bind(_hn, bz)
-
- # s[Cl].compute_at(s[C], _hn)
- s[Cl].compute_at(s[C], th)
-
- _wl, _kol, _hl, _nl, _kil = s[Cl].op.axis
- _khl, _kwl, _cl, _cl4 = s[Cl].op.reduce_axis
-
- cfg.define_split("tile_c", _cl, num_outputs=2)
- cfg.define_split("tile_kh", _khl, num_outputs=2)
- cfg.define_split("tile_kw", _kwl, num_outputs=2)
-
- _clo, _cli = cfg["tile_c"].apply(s, Cl, _cl)
- _khlo, _khli = cfg["tile_kh"].apply(s, Cl, _khl)
- _kwlo, _kwli = cfg["tile_kw"].apply(s, Cl, _kwl)
- # s[OL].reorder(rco, ryo, rxo, rci, ryi, rxi, n, f, y, x)
- s[Cl].reorder(_clo, _khlo, _kwlo, _cli, _cl4, _khli, _kwli, _kol, _hl, _nl, _kil, _wl)
- # s[Cl].reorder(_clo, _khlo, _kwlo, _cli, _cl4, _khli, _kwli)
- # s[Cl].reorder(_cl, _cl4, _kil, _wl)
- s[Cl].unroll(_cl4)
- s[Cl].unroll(_wl)
- s[Cl].vectorize(_kil)
-
- _wla, _cla, _hnla, _cl4a = s[Al].op.axis
- s[Al].compute_at(s[Cl], _cli)
- s[Al].vectorize(_cl4a)
- s[Al].unroll(_wla)
-
- _clb, _rskolb, _kilb = s[Bl].op.axis
- s[Bl].compute_at(s[Cl], _cli)
- s[Bl].vectorize(_kilb)
- s[Bl].unroll(_clb)
-
- s[C].pragma(kernel_scope, "auto_unroll_max_step", cfg["auto_unroll_max_step"].val)
-
- WO, K, HO, N, K4 = get_const_tuple(C.shape)
- RSC, _, _ = get_const_tuple(B.shape)
- cfg.add_flop(2 * N * K * K4 * HO * WO * RSC)
-
- return s
-
-
-def compute_conv2d_NCHWc_KCRSk(Input, Filter, stride, padding, dilation, out_dtype=None):
- """Convolution operator in NCHWc layout."""
-
- if out_dtype is None:
- out_dtype = Input.dtype
- assert isinstance(stride, int) or len(stride) == 2
- assert isinstance(dilation, int) or len(dilation) == 2
- if isinstance(stride, int):
- stride_h = stride_w = stride
- else:
- stride_h, stride_w = stride
-
- if isinstance(dilation, int):
- dilation_h = dilation_w = dilation
- else:
- dilation_h, dilation_w = dilation
-
- batch, in_channel_chunk, in_height, in_width, in_channel_block = Input.shape
- num_filter_chunk, channel, kernel_h, kernel_w, num_filter_block = Filter.shape
- # compute the output shape
- dilated_kernel_h = (kernel_h - 1) * dilation_h + 1
- dilated_kernel_w = (kernel_w - 1) * dilation_w + 1
- pad_top, pad_left, pad_down, pad_right = nn.get_pad_tuple(
- padding, (dilated_kernel_h, dilated_kernel_w)
- )
-
- out_height = simplify((in_height - dilated_kernel_h + pad_top + pad_down) // stride_h + 1)
- out_width = simplify((in_width - dilated_kernel_w + pad_left + pad_right) // stride_w + 1)
- # compute graph
- pad_before = [0, 0, pad_top, pad_left, 0]
- pad_after = [0, 0, pad_down, pad_right, 0]
- temp = nn.pad(Input, pad_before, pad_after, name="pad_temp")
-
- rcc = te.reduce_axis((0, in_channel_chunk), name="rc")
- rcb = te.reduce_axis((0, in_channel_block), name="rc")
- ry = te.reduce_axis((0, kernel_h), name="ry")
- rx = te.reduce_axis((0, kernel_w), name="rx")
-
- # NCHWc x KCRSk
- # texture: NCH|W|c
- # texture: K|CRS|k
- # c = crs//RS
- # rs = crs % RS
- # r = rs // W == (crs // S) % R
- # s = rs % W == crs % S
- Filter = te.compute(
- (num_filter_chunk, channel * kernel_h * kernel_w, num_filter_block),
- lambda ffc, crs, ffb: Filter[
- ffc, crs // (kernel_h * kernel_w), (crs // kernel_w) % kernel_h, crs % kernel_w, ffb
- ],
- name="packed_filter",
- )
- return te.compute(
- (batch, num_filter_chunk, out_height, out_width, num_filter_block),
- lambda nn, ffc, yy, xx, ffb: te.sum(
- temp[
- nn, rcc, yy * stride_h + ry * dilation_h, xx * stride_w + rx * dilation_w, rcb
- ].astype(out_dtype)
- * Filter[
- ffc, ((rcc * in_channel_block + rcb) * kernel_h + ry) * kernel_w + rx, ffb
- ].astype(out_dtype),
- axis=[rcc, rcb, ry, rx],
- ),
- tag="conv2d_nchwc_kcrsk_texture",
- )
-
-
-def schedule_conv2d_NCHWc_KCRSk(cfg, s, conv):
- """schedule optimized for batch size = 1"""
-
- ##### space definition begin #####
- n, fc, y, x, fb = s[conv].op.axis
- rcc, rcb, ry, rx = s[conv].op.reduce_axis
- cfg.define_split("tile_fc", fc, num_outputs=4)
- cfg.define_split("tile_y", y, num_outputs=4)
- cfg.define_split("tile_x", x, num_outputs=4)
- cfg.define_split("tile_rcc", rcc, num_outputs=2)
- cfg.define_split("tile_ry", ry, num_outputs=2)
- cfg.define_split("tile_rx", rx, num_outputs=2)
- cfg.define_knob("auto_unroll_max_step", [0, 512, 1500])
-
- pad_data, flattened_kernel = s[conv].op.input_tensors
- kernel = s[flattened_kernel].op.input_tensors[0]
- s[flattened_kernel].compute_inline()
-
- s[pad_data].compute_inline()
- if isinstance(kernel.op, tvm.te.ComputeOp) and "dilate" in kernel.op.tag:
- s[kernel].compute_inline()
- kernel = flattened_kernel
-
- if conv.op in s.outputs:
- output = conv
- OL = s.cache_write(conv, "local")
- else:
- output = s.outputs[0].output(0)
- s[conv].set_scope("local")
- OL = conv
-
- # create cache stage
- AT = s.cache_read(pad_data, "global.texture", [OL])
- WT = s.cache_read(kernel, "global.texture", [OL])
-
- def copy_to_texture(stage):
- axes = s[stage].op.axis
- fused = s[stage].fuse(*axes[:-1])
- block, thread = s[stage].split(fused, factor=32)
- s[stage].vectorize(axes[-1])
- s[stage].bind(block, te.thread_axis("blockIdx.x"))
- s[stage].bind(thread, te.thread_axis("threadIdx.x"))
-
- copy_to_texture(AT)
- copy_to_texture(WT)
-
- AA = s.cache_read(AT, "shared", [OL])
- WW = s.cache_read(WT, "shared", [OL])
-
- # tile and bind spatial axes
- n, fc, y, x, fb = s[output].op.axis
-
- kernel_scope, n = s[output].split(n, nparts=1)
-
- bf, vf, tf, fi = cfg["tile_fc"].apply(s, output, fc)
- by, vy, ty, yi = cfg["tile_y"].apply(s, output, y)
- bx, vx, tx, xi = cfg["tile_x"].apply(s, output, x)
-
- bf = s[output].fuse(n, bf)
- s[output].bind(bf, te.thread_axis("blockIdx.z"))
- s[output].bind(by, te.thread_axis("blockIdx.y"))
- s[output].bind(bx, te.thread_axis("blockIdx.x"))
- s[output].bind(vf, te.thread_axis("vthread"))
- s[output].bind(vy, te.thread_axis("vthread"))
- s[output].bind(vx, te.thread_axis("vthread"))
- s[output].bind(tf, te.thread_axis("threadIdx.z"))
- s[output].bind(ty, te.thread_axis("threadIdx.y"))
- s[output].bind(tx, te.thread_axis("threadIdx.x"))
- s[output].reorder(bf, by, bx, vf, vy, vx, tf, ty, tx, fi, yi, xi, fb)
- s[output].vectorize(fb)
- s[OL].compute_at(s[output], tx)
-
- # tile reduction axes
- n, fc, y, x, fb = s[OL].op.axis
-
- rcc, rcb, ry, rx = s[OL].op.reduce_axis
- rco, rci = cfg["tile_rcc"].apply(s, OL, rcc)
- ryo, ryi = cfg["tile_ry"].apply(s, OL, ry)
- rxo, rxi = cfg["tile_rx"].apply(s, OL, rx)
-
- # TODO(csullivan): check position of rcb
- s[OL].reorder(rco, ryo, rxo, rci, ryi, rxi, rcb, n, fc, y, x, fb)
- s[OL].vectorize(fb)
- s[OL].unroll(rcb)
-
- s[AA].compute_at(s[OL], rxo)
- s[WW].compute_at(s[OL], rxo)
- # cooperative fetching
- for load in [AA, WW]:
- if load == WW:
- n, fyx, v = s[load].op.axis
- fused = s[load].fuse(n, fyx)
- else:
- n, f, y, x, v = s[load].op.axis
- fused = s[load].fuse(n, f, y, x)
- tz, fused = s[load].split(fused, nparts=cfg["tile_fc"].size[2])
- ty, fused = s[load].split(fused, nparts=cfg["tile_y"].size[2])
- tx, fused = s[load].split(fused, nparts=cfg["tile_x"].size[2])
- s[load].bind(tz, te.thread_axis("threadIdx.z"))
- s[load].bind(ty, te.thread_axis("threadIdx.y"))
- s[load].bind(tx, te.thread_axis("threadIdx.x"))
- s[load].vectorize(v)
-
- # unroll
- s[output].pragma(kernel_scope, "auto_unroll_max_step", cfg["auto_unroll_max_step"].val)
-
- N, OCC, OH, OW, OCB = get_const_tuple(output.shape)
- _, ICKHKW, _ = get_const_tuple(kernel.shape)
-
- if isinstance(N, int):
- cfg.add_flop(2 * N * OH * OW * OCC * OCB * ICKHKW)
-
-
-def compute_conv2d_NCHWc_KCRSk_acc32(Input, Filter, stride, padding, dilation, out_dtype=None):
- """Convolution operator in NCHWc layout."""
-
- if out_dtype is None:
- out_dtype = Input.dtype
- assert isinstance(stride, int) or len(stride) == 2
- assert isinstance(dilation, int) or len(dilation) == 2
- if isinstance(stride, int):
- stride_h = stride_w = stride
- else:
- stride_h, stride_w = stride
-
- if isinstance(dilation, int):
- dilation_h = dilation_w = dilation
- else:
- dilation_h, dilation_w = dilation
-
- batch, in_channel_chunk, in_height, in_width, in_channel_block = Input.shape
- num_filter_chunk, channel, kernel_h, kernel_w, num_filter_block = Filter.shape
- # compute the output shape
- dilated_kernel_h = (kernel_h - 1) * dilation_h + 1
- dilated_kernel_w = (kernel_w - 1) * dilation_w + 1
- pad_top, pad_left, pad_down, pad_right = nn.get_pad_tuple(
- padding, (dilated_kernel_h, dilated_kernel_w)
- )
-
- out_height = simplify((in_height - dilated_kernel_h + pad_top + pad_down) // stride_h + 1)
- out_width = simplify((in_width - dilated_kernel_w + pad_left + pad_right) // stride_w + 1)
- # compute graph
- pad_before = [0, 0, pad_top, pad_left, 0]
- pad_after = [0, 0, pad_down, pad_right, 0]
- temp = nn.pad(Input, pad_before, pad_after, name="pad_temp")
-
- rcc = te.reduce_axis((0, in_channel_chunk), name="rc")
- rcb = te.reduce_axis((0, in_channel_block), name="rc")
- ry = te.reduce_axis((0, kernel_h), name="ry")
- rx = te.reduce_axis((0, kernel_w), name="rx")
-
- # NCHWc x KCRSk
- # texture: NCH|W|c
- # texture: K|CRS|k
- # c = crs//RS
- # rs = crs % RS
- # r = rs // W == (crs // S) % R
- # s = rs % W == crs % S
- Filter = te.compute(
- (num_filter_chunk, channel * kernel_h * kernel_w, num_filter_block),
- lambda ffc, crs, ffb: Filter[
- ffc, crs // (kernel_h * kernel_w), (crs // kernel_w) % kernel_h, crs % kernel_w, ffb
- ],
- name="packed_filter",
- )
- conv = te.compute(
- (batch, num_filter_chunk, out_height, out_width, num_filter_block),
- lambda nn, ffc, yy, xx, ffb: te.sum(
- (
- temp[nn, rcc, yy * stride_h + ry * dilation_h, xx * stride_w + rx * dilation_w, rcb]
- * Filter[ffc, ((rcc * in_channel_block + rcb) * kernel_h + ry) * kernel_w + rx, ffb]
- ).astype(out_dtype),
- axis=[rcc, rcb, ry, rx],
- ),
- tag="conv2d_nchwc_kcrsk_texture",
- )
- output = te.compute(conv.shape, lambda n, fc, y, x, fb: conv[n, fc, y, x, fb].astype("float32"))
- return output
-
-
-def schedule_conv2d_NCHWc_KCRSk_acc32(cfg, s, output):
- """schedule optimized for batch size = 1"""
-
- conv = output.op.input_tensors[0]
-
- ##### space definition begin #####
- n, fc, y, x, fb = s[conv].op.axis
- rcc, rcb, ry, rx = s[conv].op.reduce_axis
- cfg.define_split("tile_fc", fc, num_outputs=4)
- cfg.define_split("tile_y", y, num_outputs=4)
- cfg.define_split("tile_x", x, num_outputs=4)
- cfg.define_split("tile_rcc", rcc, num_outputs=2)
- cfg.define_split("tile_ry", ry, num_outputs=2)
- cfg.define_split("tile_rx", rx, num_outputs=2)
- cfg.define_knob("auto_unroll_max_step", [0, 512, 1500])
-
- pad_data, flattened_kernel = s[conv].op.input_tensors
- kernel = s[flattened_kernel].op.input_tensors[0]
- s[flattened_kernel].compute_inline()
-
- s[pad_data].compute_inline()
- if isinstance(kernel.op, tvm.te.ComputeOp) and "dilate" in kernel.op.tag:
- s[kernel].compute_inline()
- kernel = flattened_kernel
-
- if conv.op in s.outputs:
- output = conv
- OL = s.cache_write(conv, "local")
- else:
- output = s.outputs[0].output(0)
- s[conv].set_scope("local")
- OL = conv
-
- # create cache stage
- AT = s.cache_read(pad_data, "global.texture", [OL])
- WT = s.cache_read(kernel, "global.texture", [OL])
-
- def copy_to_texture(stage):
- axes = s[stage].op.axis
- fused = s[stage].fuse(*axes[:-1])
- block, thread = s[stage].split(fused, factor=32)
- s[stage].vectorize(axes[-1])
- s[stage].bind(block, te.thread_axis("blockIdx.x"))
- s[stage].bind(thread, te.thread_axis("threadIdx.x"))
-
- copy_to_texture(AT)
- copy_to_texture(WT)
-
- AA = s.cache_read(AT, "shared", [OL])
- WW = s.cache_read(WT, "shared", [OL])
-
- # tile and bind spatial axes
- n, fc, y, x, fb = s[output].op.axis
-
- kernel_scope, n = s[output].split(n, nparts=1)
-
- bf, vf, tf, fi = cfg["tile_fc"].apply(s, output, fc)
- by, vy, ty, yi = cfg["tile_y"].apply(s, output, y)
- bx, vx, tx, xi = cfg["tile_x"].apply(s, output, x)
-
- bf = s[output].fuse(n, bf)
- s[output].bind(bf, te.thread_axis("blockIdx.z"))
- s[output].bind(by, te.thread_axis("blockIdx.y"))
- s[output].bind(bx, te.thread_axis("blockIdx.x"))
- s[output].bind(vf, te.thread_axis("vthread"))
- s[output].bind(vy, te.thread_axis("vthread"))
- s[output].bind(vx, te.thread_axis("vthread"))
- s[output].bind(tf, te.thread_axis("threadIdx.z"))
- s[output].bind(ty, te.thread_axis("threadIdx.y"))
- s[output].bind(tx, te.thread_axis("threadIdx.x"))
- s[output].reorder(bf, by, bx, vf, vy, vx, tf, ty, tx, fi, yi, xi, fb)
- s[output].vectorize(fb)
-
- s[OL].compute_at(s[output], tx)
-
- # tile reduction axes
- n, fc, y, x, fb = s[OL].op.axis
-
- rcc, rcb, ry, rx = s[OL].op.reduce_axis
- rco, rci = cfg["tile_rcc"].apply(s, OL, rcc)
- ryo, ryi = cfg["tile_ry"].apply(s, OL, ry)
- rxo, rxi = cfg["tile_rx"].apply(s, OL, rx)
-
- # TODO(csullivan): check position of rcb
- s[OL].reorder(rco, ryo, rxo, rci, ryi, rxi, rcb, n, fc, y, x, fb)
- s[OL].vectorize(fb)
- s[OL].unroll(rcb)
-
- s[AA].compute_at(s[OL], rxo)
- s[WW].compute_at(s[OL], rxo)
- # cooperative fetching
- for load in [AA, WW]:
- if load == WW:
- n, fyx, v = s[load].op.axis
- fused = s[load].fuse(n, fyx)
- else:
- n, f, y, x, v = s[load].op.axis
- fused = s[load].fuse(n, f, y, x)
- tz, fused = s[load].split(fused, nparts=cfg["tile_fc"].size[2])
- ty, fused = s[load].split(fused, nparts=cfg["tile_y"].size[2])
- tx, fused = s[load].split(fused, nparts=cfg["tile_x"].size[2])
- s[load].bind(tz, te.thread_axis("threadIdx.z"))
- s[load].bind(ty, te.thread_axis("threadIdx.y"))
- s[load].bind(tx, te.thread_axis("threadIdx.x"))
- s[load].vectorize(v)
-
- # unroll
- s[output].pragma(kernel_scope, "auto_unroll_max_step", cfg["auto_unroll_max_step"].val)
-
- N, OCC, OH, OW, OCB = get_const_tuple(output.shape)
- _, ICKHKW, _ = get_const_tuple(kernel.shape)
-
- if isinstance(N, int):
- cfg.add_flop(2 * N * OH * OW * OCC * OCB * ICKHKW)
-
-
-def compute_depthwise_conv2d_NCHWc_KCRSk_acc32(
- Input, Filter, stride, padding, dilation, out_dtype=None
-):
- """Depthwise convolution operator in NCHWc layout."""
- if out_dtype is None:
- out_dtype = Input.dtype
- assert isinstance(stride, int) or len(stride) == 2
- assert isinstance(dilation, int) or len(dilation) == 2
-
- if isinstance(stride, int):
- stride_h = stride_w = stride
- else:
- stride_h, stride_w = stride
-
- if isinstance(dilation, int):
- dilation_h = dilation_w = dilation
- else:
- dilation_h, dilation_w = dilation
-
- batch, channel_chunk, in_height, in_width, channel_block = Input.shape
- _, channel_multiplier, kernel_h, kernel_w, _ = Filter.shape
-
- # compute the output shape
- dilated_kernel_h = (kernel_h - 1) * dilation_h + 1
- dilated_kernel_w = (kernel_w - 1) * dilation_w + 1
- pad_top, pad_left, pad_down, pad_right = nn.get_pad_tuple(
- padding, (dilated_kernel_h, dilated_kernel_w)
- )
- out_channel_chunk = simplify(channel_chunk * channel_multiplier)
- out_height = simplify((in_height - dilated_kernel_h + pad_top + pad_down) // stride_h + 1)
- out_width = simplify((in_width - dilated_kernel_w + pad_left + pad_right) // stride_w + 1)
- # compute graph
- pad_before = [0, 0, pad_top, pad_left, 0]
- pad_after = [0, 0, pad_down, pad_right, 0]
- temp = nn.pad(Input, pad_before, pad_after, name="pad_temp")
-
- ry = te.reduce_axis((0, kernel_h), name="ry")
- rx = te.reduce_axis((0, kernel_w), name="rx")
-
- # NCHWc x CMRSc = [N,(C//4)M,OH,OW, 4c]
- # NCHWc x CMRS
- # texture: NCH|W|c
- # texture: C|MRS|c
- # output: N
- # m = mrs//RS
- # rs = mrs % RS
- # r = rs // W == (mrs // S) % R
- # s = rs % W == mrs % S
- Filter = te.compute(
- (channel_chunk, channel_multiplier * kernel_h * kernel_w, channel_block),
- lambda ffc, mrs, ffb: Filter[
- ffc, mrs // (kernel_h * kernel_w), (mrs // kernel_w) % kernel_h, mrs % kernel_w, ffb
- ],
- name="packed_filter",
- )
-
- conv = te.compute(
- (batch, out_channel_chunk, out_height, out_width, channel_block),
- lambda nn, ffc, yy, xx, ffb: te.sum(
- (
- temp[
- nn,
- ffc // channel_multiplier,
- yy * stride_h + ry * dilation_h,
- xx * stride_w + rx * dilation_w,
- ffb,
- ]
- * Filter[
- ffc // channel_multiplier,
- ((ffc % channel_multiplier) * kernel_h + ry) * kernel_w + rx,
- ffb,
- ]
- ).astype(out_dtype),
- axis=[ry, rx],
- ),
- tag="depthwise_conv2d_nchwc_kcrsk_texture",
- )
- return te.compute(
- conv.shape, lambda n, ffc, y, x, ffb: conv[n, ffc, y, x, ffb].astype("float32")
- )
-
-
-def schedule_depthwise_conv2d_NCHWc_KCRSk_acc32(cfg, s, output):
- """schedule optimized for batch size = 1"""
-
- conv = output.op.input_tensors[0]
-
- ##### space definition begin #####
- n, fc, y, x, fb = s[conv].op.axis
- ry, rx = s[conv].op.reduce_axis
- cfg.define_split("tile_fc", fc, num_outputs=4)
- cfg.define_split("tile_y", y, num_outputs=4)
- cfg.define_split("tile_x", x, num_outputs=4)
- cfg.define_split("tile_ry", ry, num_outputs=2)
- cfg.define_split("tile_rx", rx, num_outputs=2)
- cfg.define_knob("auto_unroll_max_step", [0, 512, 1500])
-
- pad_data, flattened_kernel = s[conv].op.input_tensors
- kernel = s[flattened_kernel].op.input_tensors[0]
- s[flattened_kernel].compute_inline()
-
- s[pad_data].compute_inline()
- if isinstance(kernel.op, tvm.te.ComputeOp) and "dilate" in kernel.op.tag:
- s[kernel].compute_inline()
- kernel = flattened_kernel
-
- if conv.op in s.outputs:
- output = conv
- OL = s.cache_write(conv, "local")
- else:
- output = s.outputs[0].output(0)
- s[conv].set_scope("local")
- OL = conv
-
- # create cache stage
- AT = s.cache_read(pad_data, "global.texture", [OL])
- WT = s.cache_read(kernel, "global.texture", [OL])
-
- def copy_to_texture(stage):
- axes = s[stage].op.axis
- fused = s[stage].fuse(*axes[:-1])
- block, thread = s[stage].split(fused, factor=32)
- s[stage].vectorize(axes[-1])
- s[stage].bind(block, te.thread_axis("blockIdx.x"))
- s[stage].bind(thread, te.thread_axis("threadIdx.x"))
-
- copy_to_texture(AT)
- copy_to_texture(WT)
-
- AA = s.cache_read(AT, "shared", [OL])
- WW = s.cache_read(WT, "shared", [OL])
-
- # tile and bind spatial axes
- n, fc, y, x, fb = s[output].op.axis
-
- kernel_scope, n = s[output].split(n, nparts=1)
-
- bf, vf, tf, fi = cfg["tile_fc"].apply(s, output, fc)
- by, vy, ty, yi = cfg["tile_y"].apply(s, output, y)
- bx, vx, tx, xi = cfg["tile_x"].apply(s, output, x)
-
- bf = s[output].fuse(n, bf)
- s[output].bind(bf, te.thread_axis("blockIdx.z"))
- s[output].bind(by, te.thread_axis("blockIdx.y"))
- s[output].bind(bx, te.thread_axis("blockIdx.x"))
- s[output].bind(vf, te.thread_axis("vthread"))
- s[output].bind(vy, te.thread_axis("vthread"))
- s[output].bind(vx, te.thread_axis("vthread"))
- s[output].bind(tf, te.thread_axis("threadIdx.z"))
- s[output].bind(ty, te.thread_axis("threadIdx.y"))
- s[output].bind(tx, te.thread_axis("threadIdx.x"))
- s[output].reorder(bf, by, bx, vf, vy, vx, tf, ty, tx, fi, yi, xi, fb)
- s[output].vectorize(fb)
-
- s[OL].compute_at(s[output], tx)
-
- # tile reduction axes
- n, fc, y, x, fb = s[OL].op.axis
-
- ry, rx = s[OL].op.reduce_axis
- ryo, ryi = cfg["tile_ry"].apply(s, OL, ry)
- rxo, rxi = cfg["tile_rx"].apply(s, OL, rx)
-
- s[OL].reorder(ryo, rxo, ryi, rxi, n, fc, y, x, fb)
- s[OL].vectorize(fb)
- # s[OL].unroll()
-
- s[AA].compute_at(s[OL], rxo)
- s[WW].compute_at(s[OL], rxo)
- # cooperative fetching
- for load in [AA, WW]:
- if load == WW:
- n, fyx, v = s[load].op.axis
- fused = s[load].fuse(n, fyx)
- else:
- n, f, y, x, v = s[load].op.axis
- fused = s[load].fuse(n, f, y, x)
- tz, fused = s[load].split(fused, nparts=cfg["tile_fc"].size[2])
- ty, fused = s[load].split(fused, nparts=cfg["tile_y"].size[2])
- tx, fused = s[load].split(fused, nparts=cfg["tile_x"].size[2])
- s[load].bind(tz, te.thread_axis("threadIdx.z"))
- s[load].bind(ty, te.thread_axis("threadIdx.y"))
- s[load].bind(tx, te.thread_axis("threadIdx.x"))
- s[load].vectorize(v)
-
- # unroll
- s[output].pragma(kernel_scope, "auto_unroll_max_step", cfg["auto_unroll_max_step"].val)
-
- N, OCC, OH, OW, OCB = get_const_tuple(output.shape)
- ICC, MKHKW, ICB = get_const_tuple(kernel.shape)
- M = (OCC * OCB) // (ICC * ICB)
- KHKW = MKHKW // M
-
- if isinstance(N, int):
- cfg.add_flop(2 * N * OH * OW * OCC * OCB * KHKW)
-
-
-def scheduler(compute, schedule, *args, **kwargs):
- placeholders = compute(*args)
- s = schedule(*placeholders, **kwargs)
- return s, placeholders
-
-
-def conv2d_1x1_NCHWc_RSCKk(input_shape, filter_shape):
- placeholders = compute_conv2d_1x1_NCHWc_RSCKk(input_shape, filter_shape)
- s = schedule_conv2d_1x1_NCHWc_RSCKk(*placeholders)
- return s, placeholders
-
-
-def conv2d_1x1_WCHNc_CRSKk(input_shape, filter_shape):
- placeholders = compute_conv2d_1x1_WCHNc_CRSKk(input_shape, filter_shape)
- s = schedule_conv2d_1x1_WCHNc_CRSKk(*placeholders)
- return s, (placeholders[0], placeholders[1], placeholders[-1])
-
-
-def conv2d_NCHWc_KCRSk(input_shape, filter_shape):
- data = te.placeholder(input_shape, name="data", dtype="float32")
- filt = te.placeholder(filter_shape, name="filter", dtype="float32")
- conv = compute_conv2d_NCHWc_KCRSk(data, filt, [1, 1], [0, 0], [1, 1], "float32")
- cfg = autotvm.get_config()
- s = te.create_schedule([x.op for x in [conv]])
- schedule_conv2d_NCHWc_KCRSk(cfg, s, conv)
- return s, (data, filt, conv)
-
-
-def conv2d_NCHWc_KCRSk_fp32_acc(input_shape, filter_shape):
- data = te.placeholder(input_shape, name="data", dtype="float32")
- filt = te.placeholder(filter_shape, name="filter", dtype="float32")
- output = compute_conv2d_NCHWc_KCRSk_acc32(data, filt, [1, 1], [0, 0], [1, 1], "float32")
- cfg = autotvm.get_config()
- s = te.create_schedule([x.op for x in [output]])
- schedule_conv2d_NCHWc_KCRSk_acc32(cfg, s, output)
- return s, (data, filt, output)
-
-
-def depthwise_conv2d_NCHWc_KCRSk_acc32(input_shape, filter_shape):
- data = te.placeholder(input_shape, name="data", dtype="float32")
- filt = te.placeholder(filter_shape, name="filter", dtype="float32")
- output = compute_depthwise_conv2d_NCHWc_KCRSk_acc32(
- data, filt, [1, 1], [0, 0], [1, 1], "float32"
- )
- cfg = autotvm.get_config()
- s = te.create_schedule([x.op for x in [output]])
- schedule_depthwise_conv2d_NCHWc_KCRSk_acc32(cfg, s, output)
- return s, (data, filt, output)
-
-
-def ref_convolution(data, kernel, stride, pad):
- import mxnet as mx
-
- groups = 1
- kernel_size = (kernel.shape[2], kernel.shape[3])
- num_filter = kernel.shape[0]
- ref_res = mx.nd.Convolution(
- data=mx.nd.array(data),
- weight=mx.nd.array(kernel),
- bias=None,
- no_bias=True,
- kernel=kernel_size,
- stride=stride,
- pad=pad,
- num_filter=num_filter,
- num_group=groups,
- )
- return ref_res.asnumpy()
-
-
-def ref_depthwise_convolution(data, kernel, stride, pad):
- import mxnet as mx
-
- groups = kernel.shape[0]
- kernel_size = (kernel.shape[2], kernel.shape[3])
- num_filter = kernel.shape[0]
- multiplier = kernel.shape[1]
- ref_res = mx.nd.Convolution(
- data=mx.nd.array(data),
- weight=mx.nd.array(kernel),
- bias=None,
- no_bias=True,
- kernel=kernel_size,
- stride=stride,
- pad=pad,
- num_filter=num_filter,
- num_group=groups,
- )
- return ref_res.asnumpy()
-
-
-def validate(workload, target, dev, input_shapes, *args, **kwargs):
- s, placeholders = workload(*input_shapes, *args, **kwargs)
- func = tvm.driver.build(s, [*placeholders], target=target, name="TestFunction")
-
- args_tvm = []
- args_np = []
- for var in placeholders[:-1]:
- var_np = np.random.uniform(size=[i.value for i in var.shape]).astype(var.dtype)
- args_np.append(var_np)
- args_tvm.append(tvm.nd.array(var_np, dev))
- args_tvm.append(
- tvm.nd.array(
- np.zeros([i.value for i in placeholders[-1].shape], dtype=placeholders[-1].dtype), dev
- )
- )
- func(*args_tvm)
-
- if "plus_one" in workload.__name__:
- np_result = args_np[0] + 1.0
- elif "matmul" in workload.__name__:
- if "inner" in workload.__name__:
- np_result = np.matmul(
- args_np[0].reshape(32, 256), args_np[1].reshape(32, 256).transpose(1, 0)
- )
- elif "accum" in workload.__name__:
- np_result = np.matmul(
- args_np[0].transpose((1, 0, 2)).reshape(64, 128), args_np[1].reshape(128, 64)
- )
- else:
- np_result = np.matmul(
- args_np[0].transpose((0, 2, 1)).reshape(128, 64),
- args_np[1].transpose(1, 0, 2).reshape(64, 128),
- )
- elif "conv2d_1x1_NCHWc_RSCKk" in workload.__name__:
- vec_length = args_np[1].shape[-1]
- # nchwc -> nchw
- args_np[0] = (
- args_np[0]
- .transpose((0, 1, 4, 2, 3))
- .reshape(
- args_np[0].shape[0],
- args_np[0].shape[1] * args_np[0].shape[-1],
- args_np[0].shape[2],
- args_np[0].shape[3],
- )
- )
- # rsckk -> rsck -> kcrs
- args_np[1] = (
- args_np[1]
- .reshape(
- args_np[1].shape[0],
- args_np[1].shape[1],
- args_np[1].shape[2],
- args_np[1].shape[3] * args_np[1].shape[4],
- )
- .transpose((3, 2, 0, 1))
- )
- np_result = testing.conv2d_nchw_python(args_np[0], args_np[1], 1, 0)
- # nkhw -> nkhwk
- np_result = np_result.reshape(
- np_result.shape[0],
- np_result.shape[1] // vec_length,
- vec_length,
- np_result.shape[2],
- np_result.shape[3],
- ).transpose(0, 1, 3, 4, 2)
- elif "conv2d_1x1_WCHNc_CRSKk" in workload.__name__:
- vec_length = args_np[1].shape[-1]
- # wchnc -> nchw
- args_np[0] = (
- args_np[0]
- .transpose((3, 1, 4, 2, 0))
- .reshape(
- args_np[0].shape[3],
- args_np[0].shape[1] * args_np[0].shape[-1],
- args_np[0].shape[2],
- args_np[0].shape[0],
- )
- )
- # crskk -> crsk -> kcrs
- args_np[1] = (
- args_np[1]
- .reshape(
- args_np[1].shape[0],
- args_np[1].shape[1],
- args_np[1].shape[2],
- args_np[1].shape[3] * args_np[1].shape[4],
- )
- .transpose((3, 0, 1, 2))
- )
- np_result = testing.conv2d_nchw_python(args_np[0], args_np[1], 1, 0)
- # nkhw -> nkkhw -> wkhnk
- np_result = np_result.reshape(
- np_result.shape[0],
- np_result.shape[1] // vec_length,
- vec_length,
- np_result.shape[2],
- np_result.shape[3],
- ).transpose(4, 1, 3, 0, 2)
- elif "NCHW_KCRS" in workload.__name__:
- np_result = testing.conv2d_nchw_python(args_np[0], args_np[1], 1, 0)
- elif "NCHWc_KCRSk" in workload.__name__:
- vec_length = args_np[1].shape[-1]
- # nchwc -> nchw
- args_np[0] = (
- args_np[0]
- .transpose((0, 1, 4, 2, 3))
- .reshape(
- args_np[0].shape[0],
- args_np[0].shape[1] * args_np[0].shape[-1],
- args_np[0].shape[2],
- args_np[0].shape[3],
- )
- )
- # kcrsk/cmrsc -> kcrs/cmrs
- args_np[1] = (
- args_np[1]
- .transpose((0, 4, 1, 2, 3))
- .reshape(
- args_np[1].shape[0] * args_np[1].shape[4],
- args_np[1].shape[1],
- args_np[1].shape[2],
- args_np[1].shape[3],
- )
- )
- if "depthwise" in workload.__name__:
- # np_result = testing.depthwise_conv2d_python_nchw(args_np[0], args_np[1], 1, "VALID")
- np_result = ref_depthwise_convolution(args_np[0], args_np[1], [], [])
- else:
- # np_result = testing.conv2d_nchw_python(args_np[0], args_np[1], 1, 0)
- np_result = ref_convolution(args_np[0], args_np[1], [], [])
- # nkhw -> nkhwk
- np_result = np_result.reshape(
- np_result.shape[0],
- np_result.shape[1] // vec_length,
- vec_length,
- np_result.shape[2],
- np_result.shape[3],
- ).transpose(0, 1, 3, 4, 2)
- np.testing.assert_allclose(args_tvm[-1].asnumpy(), np_result, rtol=1e-2, atol=1e-2)
-
-
-class BaseSingleShapeValidator:
- @tvm.testing.parametrize_targets("opencl")
- def test_unary(self, test_func, input_shape, target, dev):
- validate(test_func, target, dev, [input_shape])
-
-
-class TestPlusOneRank3(BaseSingleShapeValidator):
- input_shape = tvm.testing.parameter((32, 32, 4))
-
- def plus_one(input_shape):
- return scheduler(compute_plus_one_rank3, schedule_plus_one_rank3, input_shape)
-
- test_func = tvm.testing.parameter(plus_one)
-
-
-class TestPlusOneRank5(BaseSingleShapeValidator):
- input_shape = tvm.testing.parameter((32, 2, 4, 4, 4))
-
- def plus_one(input_shape):
- return scheduler(compute_plus_one_rank5, schedule_plus_one_rank5, input_shape)
-
- test_func = tvm.testing.parameter(plus_one)
-
-
-class TestMatmul:
- input_shape = tvm.testing.parameter((32, 64, 4))
- local = tvm.testing.parameter(False, True)
-
- def matmul(input_shape, local):
- return scheduler(compute_matmul, schedule_matmul, input_shape, local=local)
-
- def matmul_inner(input_shape, local):
- return scheduler(compute_matmul_inner, schedule_matmul_inner, input_shape, local=local)
-
- test_func = tvm.testing.parameter(matmul, matmul_inner)
-
- @tvm.testing.parametrize_targets("opencl")
- def test_matmul(self, test_func, input_shape, local, target, dev):
- validate(test_func, target, dev, [input_shape], local=local)
-
-
-class TestMatmulVectorAccumulator:
- shapeA = tvm.testing.parameter((32, 64, 4))
- shapeB = tvm.testing.parameter((128, 16, 4))
- local = tvm.testing.parameter(False, True)
-
- def matmul_vector_accumulator(shapeA, shapeB, local):
- return scheduler(
- compute_matmul_vector_accumulator,
- schedule_matmul_vector_accumulator,
- shapeA,
- shapeB,
- local=local,
- )
-
- test_func = tvm.testing.parameter(matmul_vector_accumulator)
-
- @tvm.testing.parametrize_targets("opencl")
- def test_matmul_vec_acc(self, test_func, shapeA, shapeB, local, target, dev):
- validate(test_func, target, dev, [shapeA, shapeB], local=local)
-
-
-class BaseConv2DValidator:
- @tvm.testing.parametrize_targets("opencl")
- def test_conv2d(self, test_func, input_shapes, target, dev):
- validate(test_func, target, dev, input_shapes)
-
-
-class TestConv2dNCHWcRSCKk(BaseConv2DValidator):
- input_shapes = tvm.testing.parameter([(1, 32, 56, 56, 4), (1, 1, 128, 32, 4)])
- test_func = tvm.testing.parameter(conv2d_1x1_NCHWc_RSCKk)
-
-
-class TestConv2dWCHNcCRSKk(BaseConv2DValidator):
- input_shapes = tvm.testing.parameter([(56, 32, 56, 1, 4), (128, 1, 1, 32, 4)])
- test_func = tvm.testing.parameter(conv2d_1x1_WCHNc_CRSKk)
-
-
-@pytest.mark.skip("AttributeError: module 'numpy' has no attribute 'bool' raised from mxnet")
-class TestConv2dNCHWcKCRSk(BaseConv2DValidator):
- input_shapes = tvm.testing.parameter(
- [(1, 32, 56, 56, 4), (32, 128, 1, 1, 4)], [(1, 32, 112, 112, 4), (32, 128, 3, 3, 4)]
- )
- test_func = tvm.testing.parameter(conv2d_NCHWc_KCRSk, conv2d_NCHWc_KCRSk_fp32_acc)
-
-
-@pytest.mark.skip("AttributeError: module 'numpy' has no attribute 'bool' raised from mxnet")
-class TestDepthwiseConv2dNCHWcKCRSk(BaseConv2DValidator):
- input_shapes = tvm.testing.parameter([(1, 24, 257, 257, 4), (24, 1, 3, 3, 4)])
- test_func = tvm.testing.parameter(depthwise_conv2d_NCHWc_KCRSk_acc32)
-
-
-def simple_texture_to_scalar_common(
- target, input_info, output_info, find_patterns, dtype, cast_type
-):
- def _compute():
- p0 = te.placeholder(input_info[1], name="p0", dtype=dtype)
- p0_comp = te.compute(input_info[1], lambda *i: p0(*i), name="p0_comp")
- if len(output_info[1]) == 4 and len(input_info[1]) == 5:
- out = te.compute(
- output_info[1],
- lambda n, c, h, w: p0_comp[n][c // 4][h][w][c % 4].astype(cast_type),
- name="out",
- )
- elif len(output_info[1]) == 5 and len(input_info[1]) == 5:
- out = te.compute(
- output_info[1],
- lambda n, c, h, w, cb: p0_comp[n][c][h][w][cb].astype(cast_type),
- name="out",
- )
- else:
- raise Exception("Impossible case")
- dummy_out = te.compute(output_info[1], lambda *i: out(*i), name="dummy_out")
- return p0, dummy_out
-
- def _schedule(dummy_out):
- from tvm.topi.adreno.utils import bind_data_copy
-
- s = te.create_schedule(dummy_out.op)
- out = s[dummy_out].op.input_tensors[0]
- p0_comp = s[out].op.input_tensors[0]
- s[p0_comp].set_scope(input_info[0])
- bind_data_copy(s[p0_comp])
- s[out].set_scope(output_info[0])
- bind_data_copy(s[out])
- bind_data_copy(s[dummy_out])
- return s
-
- p0, dummy_out = _compute()
- s = _schedule(dummy_out)
-
- fun = tvm.build(s, [p0, dummy_out], target)
- dev = tvm.device(target, 0)
- opencl_source = fun.imported_modules[0].get_source()
- start_idx = 0
- for pattern in find_patterns:
- start_idx = opencl_source.find(pattern, start_idx)
- assert start_idx > -1
-
- input_np = np.random.uniform(size=[i for i in input_info[1]]).astype(dtype)
- input_tvm = tvm.nd.array(input_np, dev)
- c = tvm.nd.empty(output_info[1], dtype, dev)
- # Doesn't run OpenCL code for FP16 because GPUs in CI don't support FP16 inference
- if cast_type == "float32":
- fun(input_tvm, c)
- # For output len == 5 it makes no sense to check the accuracy
- if cast_type == "float32" and len(output_info[1]) == 4:
- np_result = input_np.transpose(0, 2, 3, 1, 4) # NCHW4c -> NHWC4c
- np_result = np.squeeze(np_result, axis=3)
- np_result = np_result.transpose(0, 3, 1, 2) # NHWC -> NCHW
- np.testing.assert_allclose(c.asnumpy(), np_result, rtol=1e-2, atol=1e-2)
-
-
-class TestSimpleTextureToScalarFP16:
- # (input [scope, shape], output [scope, shape], [find_patterns])
- input_info, output_info, find_patterns = tvm.testing.parameters(
- # 1. Texture (NCHW4c) -> Cast(FP16) -> Buffer (NCHW)
- (
- ["global.texture", (1, 1, 40, 40, 4)],
- ["", (1, 4, 40, 40)],
- [
- "float4 v_ = READ_IMAGEF(p0_comp, image_sampler, ((int2)(((convert_int(get_local_id(0))) % 40), ((((convert_int(get_group_id(0))) & 1) * 20) + ((convert_int(get_local_id(0))) / 40)))));",
- "out[(((convert_int(get_group_id(0))) * 800) + (convert_int(get_local_id(0))))] = (convert_half(((float*)&v_)[((convert_int(get_group_id(0))) >> 1)]));",
- ],
- ),
- # 2. Buffer (NCHW4c) -> Cast(FP16) -> Buffer (NCHW)
- (
- ["", (1, 1, 40, 40, 4)],
- ["", (1, 4, 40, 40)],
- [
- "out[(((convert_int(get_group_id(0))) * 800) + (convert_int(get_local_id(0))))] = (convert_half(p0_comp[(((((convert_int(get_group_id(0))) & 1) * 3200) + ((convert_int(get_local_id(0))) * 4)) + ((convert_int(get_group_id(0))) >> 1))]));"
- ],
- ),
- # 3. Texture (NCHW4c) -> Cast(FP16) -> Texture (NCHW4c)
- (
- ["global.texture", (1, 1, 40, 40, 4)],
- ["global.texture", (1, 1, 40, 40, 4)],
- [
- "float4 v_ = READ_IMAGEF(p0_comp, image_sampler, ((int2)(((((convert_int(get_group_id(0))) * 24) + (convert_int(get_local_id(0)))) % 40), ((((convert_int(get_group_id(0))) * 8) + ((convert_int(get_local_id(0))) >> 3)) / 5))));",
- "write_imageh(out, (int2)(((((convert_int(get_group_id(0))) * 24) + (convert_int(get_local_id(0)))) % 40), ((((convert_int(get_group_id(0))) * 8) + ((convert_int(get_local_id(0))) >> 3)) / 5)), (convert_half4(v_)));",
- ],
- ),
- )
- dtype = tvm.testing.parameter("float32")
-
- @tvm.testing.parametrize_targets("opencl")
- def test_simple_texture_to_scalar_fp16(
- self, input_info, output_info, find_patterns, dtype, target
- ):
- simple_texture_to_scalar_common(
- target, input_info, output_info, find_patterns, dtype, "float16"
- )
-
-
-class TestSimpleTextureToScalarFP32:
- # (input [scope, shape], output [scope, shape], [find_patterns])
- input_info, output_info, find_patterns = tvm.testing.parameters(
- # 1. Texture (NCHW4c) -> Buffer (NCHW)
- (
- ["global.texture", (1, 1, 40, 40, 4)],
- ["", (1, 4, 40, 40)],
- [
- "float4 v_ = READ_IMAGEF(p0_comp, image_sampler, ((int2)(((convert_int(get_local_id(0))) % 40), ((((convert_int(get_group_id(0))) & 1) * 20) + ((convert_int(get_local_id(0))) / 40)))));",
- "out[(((convert_int(get_group_id(0))) * 800) + (convert_int(get_local_id(0))))] = ((float*)&v_)[((convert_int(get_group_id(0))) >> 1)];",
- ],
- ),
- # 2. Buffer (NCHW4c) -> Buffer (NCHW)
- (
- ["", (1, 1, 40, 40, 4)],
- ["", (1, 4, 40, 40)],
- [
- "out[(((convert_int(get_group_id(0))) * 800) + (convert_int(get_local_id(0))))] = p0_comp[(((((convert_int(get_group_id(0))) & 1) * 3200) + ((convert_int(get_local_id(0))) * 4)) + ((convert_int(get_group_id(0))) >> 1))];"
- ],
- ),
- )
- dtype = tvm.testing.parameter("float32")
-
- @tvm.testing.parametrize_targets("opencl")
- def test_simple_texture_to_scalar_fp32(
- self, input_info, output_info, find_patterns, dtype, target
- ):
- simple_texture_to_scalar_common(
- target, input_info, output_info, find_patterns, dtype, "float32"
- )
-
-
-def texture_to_scalar_reuse_ssa_common(
- target, input_info, output_info, find_patterns, dtype, cast_type
-):
- def _compute():
- p0 = te.placeholder(input_info[1], name="p0", dtype=dtype)
- p0_comp = te.compute(input_info[1], lambda *i: p0(*i), name="p0_comp")
- if len(output_info[1]) == 4 and len(input_info[1]) == 5:
- out = te.compute(
- output_info[1],
- lambda n, c, h, w: p0_comp[n][c // 4][h][w][c % 4].astype(cast_type),
- name="out",
- )
- out2 = te.compute(
- output_info[1],
- lambda n, c, h, w: out[n][c][h][w]
- + p0_comp[n][c // 4][h][w][c % 4].astype(cast_type),
- name="out",
- )
- elif len(output_info[1]) == 5 and len(input_info[1]) == 5:
- out = te.compute(
- output_info[1],
- lambda n, c, h, w, cb: p0_comp[n][c][h][w][cb].astype(cast_type),
- name="out",
- )
- out2 = te.compute(
- output_info[1],
- lambda n, c, h, w, cb: out[n][c][h][w][cb]
- + p0_comp[n][c][h][w][cb].astype(cast_type),
- name="out",
- )
- else:
- raise Exception("Impossible case")
- out_sum = te.compute(output_info[1], lambda *i: out(*i) + out2(*i), name="out_sum")
- dummy_out = te.compute(output_info[1], lambda *i: out_sum(*i), name="dummy_out")
- return p0, dummy_out
-
- def _schedule(dummy_out):
- from tvm.topi.adreno.utils import bind_data_copy
-
- s = te.create_schedule(dummy_out.op)
- out_sum = s[dummy_out].op.input_tensors[0]
- out, out2 = s[out_sum].op.input_tensors
- p0_comp = s[out].op.input_tensors[0]
- s[p0_comp].set_scope(input_info[0])
- bind_data_copy(s[p0_comp])
- s[out].set_scope(output_info[0])
- s[out2].set_scope(output_info[0])
- s[out2].compute_inline()
- s[out].compute_inline()
- s[out_sum].set_scope(output_info[0])
- bind_data_copy(s[out_sum])
- bind_data_copy(s[dummy_out])
- return s
-
- p0, dummy_out = _compute()
- s = _schedule(dummy_out)
-
- fun = tvm.build(s, [p0, dummy_out], target)
- dev = tvm.device(target, 0)
- opencl_source = fun.imported_modules[0].get_source()
- start_idx = 0
- for pattern in find_patterns:
- start_idx = opencl_source.find(pattern, start_idx)
- assert start_idx > -1
-
- input_np = np.random.uniform(size=[i for i in input_info[1]]).astype(dtype)
- input_tvm = tvm.nd.array(input_np, dev)
- c = tvm.nd.empty(output_info[1], dtype, dev)
- # Doesn't run OpenCL code for FP16 because GPUs in CI don't support FP16 inference
- if cast_type == "float32":
- fun(input_tvm, c)
- # For output len == 5 it makes no sense to check the accuracy
- if cast_type == "float32" and len(output_info[1]) == 4:
- np_result = input_np * 3
- np_result = np_result.transpose(0, 2, 3, 1, 4) # NCHW4c -> NHWC4c
- np_result = np.squeeze(np_result, axis=3)
- np_result = np_result.transpose(0, 3, 1, 2) # NHWC -> NCHW
- np.testing.assert_allclose(c.asnumpy(), np_result, rtol=1e-2, atol=1e-2)
-
-
-class TestTextureToScalarReuseSSAFP16:
- # (input [scope, shape], output [scope, shape], [find_patterns])
- input_info, output_info, find_patterns = tvm.testing.parameters(
- # 1. Texture (NCHW4c) -> Cast(FP16) -> Buffer (NCHW)
- (
- ["global.texture", (1, 1, 40, 40, 4)],
- ["", (1, 4, 40, 40)],
- [
- "float4 v_ = READ_IMAGEF(p0_comp, image_sampler, ((int2)(((convert_int(get_local_id(0))) % 40), ((((convert_int(get_group_id(0))) & 1) * 20) + ((convert_int(get_local_id(0))) / 40)))));",
- "out_sum[(((convert_int(get_group_id(0))) * 800) + (convert_int(get_local_id(0))))] = ((convert_half(((float*)&v_)[((convert_int(get_group_id(0))) >> 1)])) + ((convert_half(((float*)&v_)[((convert_int(get_group_id(0))) >> 1)])) + (convert_half(((float*)&v_)[((convert_int(get_group_id(0))) >> 1)]))));",
- ],
- ),
- # 2. Buffer (NCHW4c) -> Cast(FP16) -> Buffer (NCHW)
- (
- ["", (1, 1, 40, 40, 4)],
- ["", (1, 4, 40, 40)],
- [
- " out_sum[(((convert_int(get_group_id(0))) * 800) + (convert_int(get_local_id(0))))] = ((convert_half(p0_comp[(((((convert_int(get_group_id(0))) & 1) * 3200) + ((convert_int(get_local_id(0))) * 4)) + ((convert_int(get_group_id(0))) >> 1))])) + ((convert_half(p0_comp[(((((convert_int(get_group_id(0))) & 1) * 3200) + ((convert_int(get_local_id(0))) * 4)) + ((convert_int(get_group_id(0))) >> 1))])) + (convert_half(p0_comp[(((((convert_int(get_group_id(0))) & 1) * 3200) + ((c [...]
- ],
- ),
- # 3. Texture (NCHW4c) -> Cast(FP16) -> Texture (NCHW4c)
- (
- ["global.texture", (1, 1, 40, 40, 4)],
- ["global.texture", (1, 1, 40, 40, 4)],
- [
- "float4 v_ = READ_IMAGEF(p0_comp, image_sampler, ((int2)(((((convert_int(get_group_id(0))) * 24) + (convert_int(get_local_id(0)))) % 40), ((((convert_int(get_group_id(0))) * 8) + ((convert_int(get_local_id(0))) >> 3)) / 5))));",
- "write_imageh(out_sum, (int2)(((((convert_int(get_group_id(0))) * 24) + (convert_int(get_local_id(0)))) % 40), ((((convert_int(get_group_id(0))) * 8) + ((convert_int(get_local_id(0))) >> 3)) / 5)), ((convert_half4(v_)) + ((convert_half4(v_)) + (convert_half4(v_)))));",
- ],
- ),
- )
- dtype = tvm.testing.parameter("float32")
-
- @tvm.testing.parametrize_targets("opencl")
- def test_texture_to_scalar_reuse_ssa_fp16(
- self, input_info, output_info, find_patterns, dtype, target
- ):
- texture_to_scalar_reuse_ssa_common(
- target, input_info, output_info, find_patterns, dtype, "float16"
- )
-
-
-class TestTextureToScalarReuseSSAFP32:
- # (input [scope, shape], output [scope, shape], [find_patterns])
- input_info, output_info, find_patterns = tvm.testing.parameters(
- # 1. Texture (NCHW4c) -> Buffer (NCHW)
- (
- ["global.texture", (1, 1, 40, 40, 4)],
- ["", (1, 4, 40, 40)],
- [
- "float4 v_ = READ_IMAGEF(p0_comp, image_sampler, ((int2)(((convert_int(get_local_id(0))) % 40), ((((convert_int(get_group_id(0))) & 1) * 20) + ((convert_int(get_local_id(0))) / 40)))));",
- "out_sum[(((convert_int(get_group_id(0))) * 800) + (convert_int(get_local_id(0))))] = (((float*)&v_)[((convert_int(get_group_id(0))) >> 1)] + (((float*)&v_)[((convert_int(get_group_id(0))) >> 1)] + ((float*)&v_)[((convert_int(get_group_id(0))) >> 1)]));",
- ],
- ),
- # 2. Buffer (NCHW4c) -> Buffer (NCHW)
- (
- ["", (1, 1, 40, 40, 4)],
- ["", (1, 4, 40, 40)],
- [
- "out_sum[(((convert_int(get_group_id(0))) * 800) + (convert_int(get_local_id(0))))] = (p0_comp[(((((convert_int(get_group_id(0))) & 1) * 3200) + ((convert_int(get_local_id(0))) * 4)) + ((convert_int(get_group_id(0))) >> 1))] + (p0_comp[(((((convert_int(get_group_id(0))) & 1) * 3200) + ((convert_int(get_local_id(0))) * 4)) + ((convert_int(get_group_id(0))) >> 1))] + p0_comp[(((((convert_int(get_group_id(0))) & 1) * 3200) + ((convert_int(get_local_id(0))) * 4)) + ((convert_ [...]
- ],
- ),
- )
- dtype = tvm.testing.parameter("float32")
-
- @tvm.testing.parametrize_targets("opencl")
- def test_texture_to_scalar_reuse_ssa_fp32(
- self, input_info, output_info, find_patterns, dtype, target
- ):
- texture_to_scalar_reuse_ssa_common(
- target, input_info, output_info, find_patterns, dtype, "float32"
- )
-
-
-class TestLocalArrayToTexture:
- # 1. conv2d(Texture(NCHW4c), Texture(OIHW4o)) -> local_array[4] -> Texture (NCHW4c)
- input_shape1, input_shape2, output_shape, find_patterns = tvm.testing.parameters(
- (
- (1, 1, 40, 40, 4),
- (2, 4, 3, 3, 4),
- (1, 2, 38, 38, 4),
- [
- "float out_local[4];",
- "float4 v_ = READ_IMAGEF(p1_comp, image_sampler, ((int2)(((((convert_int(get_group_id(0))) * 14) + (convert_int(get_local_id(0)))) % 38), (((((convert_int(get_group_id(0))) * 64) + ((convert_int(get_local_id(0))) >> 1)) % 722) / 19))));",
- "float4 v__1 = READ_IMAGEF(p2_comp, image_sampler, ((int2)(rw, (((((((convert_int(get_group_id(0))) * 32) + ((convert_int(get_local_id(0))) >> 2)) / 361) * 12) + (rcb * 3)) + rh))));",
- "out_local[cb_c] = (out_local[cb_c] + (((float*)&v_)[rcb] * ((float*)&v__1)[cb_c]));",
- "write_imagef(out, (int2)(((((convert_int(get_group_id(0))) * 14) + (convert_int(get_local_id(0)))) % 38), ((((convert_int(get_group_id(0))) * 64) + ((convert_int(get_local_id(0))) >> 1)) / 19)), vload4(0, out_local + 0));",
- ],
- ),
- )
- dtype = tvm.testing.parameter("float32")
-
- @tvm.testing.parametrize_targets("opencl")
- def test_local_array_to_texture(
- self, input_shape1, input_shape2, output_shape, find_patterns, dtype, target
- ):
- def _compute():
- p1 = te.placeholder(input_shape1, name="p1", dtype=dtype)
- p1_comp = te.compute(input_shape1, lambda *i: p1(*i), name="p1_comp")
- p2 = te.placeholder(input_shape2, name="p2", dtype=dtype)
- p2_comp = te.compute(input_shape2, lambda *i: p2(*i), name="p2_comp")
- KH, KW = input_shape2[2], input_shape2[3]
- IC, ICB = input_shape1[1], input_shape1[4]
- rh = te.reduce_axis((0, KH), name="rh")
- rw = te.reduce_axis((0, KW), name="rw")
- rc = te.reduce_axis((0, IC), name="rc")
- rcb = te.reduce_axis((0, ICB), name="rcb")
- out = te.compute(
- output_shape,
- lambda n, c, h, w, cb: te.sum(
- (p1_comp[n, rc, h, w, rcb] * p2_comp[c, rc * ICB + rcb, rh, rw, cb]).astype(
- dtype
- ),
- axis=[rh, rw, rc, rcb],
- ),
- name="out",
- )
- dummy_out = te.compute(output_shape, lambda *i: out(*i), name="dummy_out")
- return p1, p2, dummy_out
-
- def _schedule(dummy_out):
- from tvm.topi.adreno.utils import bind_data_copy
-
- s = te.create_schedule(dummy_out.op)
- out = s[dummy_out].op.input_tensors[0]
- p1_comp, p2_comp = s[out].op.input_tensors
- bind_data_copy(s[p1_comp])
- s[p1_comp].set_scope("global.texture")
- bind_data_copy(s[p2_comp])
- s[p2_comp].set_scope("global.texture")
- OL = s.cache_write(out, "local")
- n, c, h, w, cb = s[out].op.axis
- fused = s[out].fuse(n, c, h, w)
- bx, tx = s[out].split(fused, 128)
- s[out].reorder(bx, tx, cb)
- s[out].vectorize(cb)
- s[out].set_scope("global.texture")
- s[out].bind(bx, te.thread_axis("blockIdx.x"))
- s[out].bind(tx, te.thread_axis("threadIdx.x"))
- s[OL].compute_at(s[out], tx)
- bind_data_copy(s[dummy_out])
- return s
-
- p1, p2, dummy_out = _compute()
- s = _schedule(dummy_out)
-
- fun = tvm.build(s, [p1, p2, dummy_out], target)
- dev = tvm.device(target, 0)
- opencl_source = fun.imported_modules[0].get_source()
- start_idx = 0
- for pattern in find_patterns:
- start_idx = opencl_source.find(pattern, start_idx)
- assert start_idx > -1
-
- input_np1 = np.random.uniform(size=[i for i in input_shape1]).astype(dtype)
- input_np2 = np.random.uniform(size=[i for i in input_shape2]).astype(dtype)
- input_tvm1 = tvm.nd.array(input_np1, dev)
- input_tvm2 = tvm.nd.array(input_np2, dev)
- c = tvm.nd.empty(output_shape, dtype, dev)
- fun(input_tvm1, input_tvm2, c)
-
-
-if __name__ == "__main__":
- tvm.testing.main()
diff --git a/tests/python/driver/tvmc/conftest.py b/tests/python/driver/tvmc/conftest.py
deleted file mode 100644
index 5f10ebb24a..0000000000
--- a/tests/python/driver/tvmc/conftest.py
+++ /dev/null
@@ -1,302 +0,0 @@
-# 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.
-import os
-import pytest
-import tarfile
-import textwrap
-
-import numpy as np
-
-from PIL import Image
-
-import tvm
-from tvm import relay
-from tvm.driver import tvmc
-
-from tvm.contrib.download import download_testdata
-
-# Support functions
-
-
-def download_and_untar(model_url, model_sub_path, temp_dir):
- model_tar_name = os.path.basename(model_url)
- model_path = download_testdata(model_url, model_tar_name, module=["tvmc"])
-
- if model_path.endswith("tgz") or model_path.endswith("gz") or model_path.endswith("tar"):
- tar = tarfile.open(model_path)
- tar.extractall(path=temp_dir)
- tar.close()
-
- return os.path.join(temp_dir, model_sub_path)
-
-
-# PyTest fixtures
-
-
-@pytest.fixture(scope="session")
-def tflite_mobilenet_v1_1_quant(tmpdir_factory):
- base_url = "https://storage.googleapis.com/download.tensorflow.org/models"
- model_url = "mobilenet_v1_2018_08_02/mobilenet_v1_1.0_224_quant.tgz"
- model_file = download_and_untar(
- "{}/{}".format(base_url, model_url),
- "mobilenet_v1_1.0_224_quant.tflite",
- temp_dir=tmpdir_factory.mktemp("data"),
- )
-
- return model_file
-
-
-@pytest.fixture(scope="session")
-def pb_mobilenet_v1_1_quant(tmpdir_factory):
- base_url = "https://storage.googleapis.com/download.tensorflow.org/models"
- model_url = "mobilenet_v1_2018_08_02/mobilenet_v1_1.0_224.tgz"
- model_file = download_and_untar(
- "{}/{}".format(base_url, model_url),
- "mobilenet_v1_1.0_224_frozen.pb",
- temp_dir=tmpdir_factory.mktemp("data"),
- )
-
- return model_file
-
-
-@pytest.fixture(scope="session")
-def keras_resnet50(tmpdir_factory):
- try:
- from tensorflow.keras.applications.resnet50 import ResNet50
- except ImportError:
- # not all environments provide TensorFlow, so skip this fixture
- # if that is that case.
- return ""
-
- model_file_name = "{}/{}".format(tmpdir_factory.mktemp("data"), "resnet50.h5")
- model = ResNet50(include_top=True, weights="imagenet", input_shape=(224, 224, 3), classes=1000)
- model.save(model_file_name)
-
- return model_file_name
-
-
-@pytest.fixture(scope="session")
-def keras_simple(tmpdir_factory):
- try:
- from tensorflow import keras
- except ImportError:
- # not all environments provide TensorFlow, so skip this fixture
- # if that is that case.
- return ""
-
- model_file_name = "{}/{}".format(tmpdir_factory.mktemp("data"), "simple_conv.h5")
- model = keras.Sequential(
- [
- keras.layers.InputLayer(input_shape=[32, 32, 3], batch_size=1),
- keras.layers.Conv2D(8, kernel_size=(3, 3)),
- keras.layers.Flatten(),
- keras.layers.Dense(64),
- ]
- )
- model.save(model_file_name)
-
- return model_file_name
-
-
-@pytest.fixture(scope="session")
-def pytorch_resnet18(tmpdir_factory):
- try:
- import torch
- import torchvision.models as models
- except ImportError:
- # Not all environments provide Pytorch, so skip if that's the case.
- return ""
- model = models.resnet18()
- model_file_name = "{}/{}".format(tmpdir_factory.mktemp("data"), "resnet18.pth")
- # Trace model into torchscript.
- traced_cpu = torch.jit.trace(model, torch.randn(1, 3, 224, 224))
- torch.jit.save(traced_cpu, model_file_name)
-
- return model_file_name
-
-
-@pytest.fixture(scope="session")
-def pytorch_mobilenetv2_quantized(tmpdir_factory):
- try:
- import torch
- import torchvision.models as models
- except ImportError:
- # Not all environments provide Pytorch, so skip if that's the case.
- return ""
- model = models.quantization.mobilenet_v2(quantize=True)
- model_file_name = "{}/{}".format(tmpdir_factory.mktemp("data"), "mobilenet_v2_quantized.pth")
- # Trace model into torchscript.
- traced_cpu = torch.jit.trace(model, torch.randn(1, 3, 224, 224))
- torch.jit.save(traced_cpu, model_file_name)
-
- return model_file_name
-
-
-@pytest.fixture(scope="session")
-def onnx_resnet50():
- base_url = "https://github.com/onnx/models/raw/bd206494e8b6a27b25e5cf7199dbcdbfe9d05d1c/vision/classification/resnet/model"
- file_to_download = "resnet50-v2-7.onnx"
- model_file = download_testdata(
- "{}/{}".format(base_url, file_to_download), file_to_download, module=["tvmc"]
- )
-
- return model_file
-
-
-@pytest.fixture(scope="session")
-def paddle_resnet50(tmpdir_factory):
- base_url = "https://bj.bcebos.com/x2paddle/models"
- model_url = "paddle_resnet50.tar"
- model_file = download_and_untar(
- "{}/{}".format(base_url, model_url),
- "paddle_resnet50/model.pdmodel",
- temp_dir=tmpdir_factory.mktemp("data"),
- )
- return model_file
-
-
-@pytest.fixture(scope="session")
-def onnx_mnist():
- base_url = "https://github.com/onnx/models/raw/bd206494e8b6a27b25e5cf7199dbcdbfe9d05d1c/vision/classification/mnist/model"
- file_to_download = "mnist-1.onnx"
- model_file = download_testdata(
- "{}/{}".format(base_url, file_to_download), file_to_download, module=["tvmc"]
- )
-
- return model_file
-
-
-@pytest.fixture
-def tflite_compile_model(tmpdir_factory):
- """Support function that returns a TFLite compiled module"""
-
- def model_compiler(model_file, **overrides):
- package_path = tmpdir_factory.mktemp("data").join("mock.tar")
- tvmc_model = tvmc.frontends.load_model(model_file)
- args = {"target": "llvm", **overrides}
- return tvmc.compiler.compile_model(tvmc_model, package_path=package_path, **args)
-
- # Returns a TVMCPackage
- return model_compiler
-
-
-@pytest.fixture
-def relay_compile_model(tmpdir_factory):
- """Support function that returns a TFLite compiled module"""
-
- def model_compiler(model_file, shape_dict, **overrides):
- package_path = tmpdir_factory.mktemp("data").join("mock.tar")
- tvmc_model = tvmc.frontends.load_model(
- model_file, model_format="relay", shape_dict=shape_dict
- )
- args = {"target": "llvm", **overrides}
- return tvmc.compiler.compile_model(tvmc_model, package_path=package_path, **args)
-
- # Returns a TVMCPackage
- return model_compiler
-
-
-@pytest.fixture(scope="session")
-def imagenet_cat(tmpdir_factory):
- tmpdir_name = tmpdir_factory.mktemp("data")
- cat_file_name = "imagenet_cat.npz"
-
- cat_url = "https://github.com/dmlc/mxnet.js/blob/main/data/cat.png?raw=true"
- image_path = download_testdata(cat_url, "inputs", module=["tvmc"])
- resized_image = Image.open(image_path).resize((224, 224))
- image_data = np.asarray(resized_image).astype("float32")
- image_data = np.expand_dims(image_data, axis=0)
-
- cat_file_full_path = os.path.join(tmpdir_name, cat_file_name)
- np.savez(cat_file_full_path, input=image_data)
-
- return cat_file_full_path
-
-
-@pytest.fixture(scope="session")
-def tflite_mobilenet_v1_0_25_128(tmpdir_factory):
- base_url = "https://storage.googleapis.com/download.tensorflow.org/models"
- model_url = "mobilenet_v1_2018_02_22/mobilenet_v1_0.25_128.tgz"
- model_file = download_and_untar(
- "{}/{}".format(base_url, model_url),
- "mobilenet_v1_0.25_128.tflite",
- temp_dir=tmpdir_factory.mktemp("data"),
- )
-
- return model_file
-
-
-@pytest.fixture(scope="session")
-def tflite_cnn_s_quantized(tmpdir_factory):
- base_url = "https://github.com/ARM-software/ML-zoo/raw/48a22ee22325d15d2371a6df24eb7d67e21dcc97/models/keyword_spotting/cnn_small/tflite_int8"
- file_to_download = "cnn_s_quantized.tflite"
- model_file = download_testdata(
- "{}/{}".format(base_url, file_to_download), file_to_download, module=["tvmc"]
- )
- return model_file
-
-
-@pytest.fixture(scope="session")
-def relay_text_conv2d(tmpdir_factory):
- file_path = os.path.join(tmpdir_factory.mktemp("model"), "relay.txt")
-
- RELAY_MODEL = textwrap.dedent(
- """\
- #[version = "0.0.5"]
- def @main(%data : Tensor[(1, 3, 64, 64), uint8], %weight : Tensor[(3, 3, 5, 5), int8]) {
- %1 = nn.conv2d(
- %data,
- %weight,
- padding=[2, 2],
- channels=3,
- kernel_size=[5, 5],
- data_layout="NCHW",
- kernel_layout="OIHW",
- out_dtype="int32");
- %2 = cast(nn.max_pool2d(%1, pool_size=[3, 3]), dtype="int8");
- %3 = nn.conv2d(
- %2,
- %weight,
- padding=[2, 2],
- channels=3,
- kernel_size=[5, 5],
- data_layout="NCHW",
- kernel_layout="OIHW",
- out_dtype="int32");
- %4 = nn.max_pool2d(%3, pool_size=[3, 3]);
- %4
- }
- """
- )
-
- with open(file_path, "w") as relay_text:
- relay_text.write(RELAY_MODEL)
- return file_path
-
-
-@pytest.fixture(scope="session")
-def relay_conv2d():
- """
- Simple conv2d Relay implementation.
- """
- dtype = "float32"
-
- x = relay.var("x", shape=(1, 4, 2, 2), dtype=dtype)
- weight = relay.const(np.random.uniform(size=(2, 4, 2, 2)), dtype=dtype)
- x = relay.nn.conv2d(x, weight)
- func = relay.Function(relay.analysis.free_vars(x), x)
- return tvm.IRModule.from_expr(func)
diff --git a/tests/python/driver/tvmc/test_autoscheduler.py b/tests/python/driver/tvmc/test_autoscheduler.py
deleted file mode 100644
index a3a3f52d39..0000000000
--- a/tests/python/driver/tvmc/test_autoscheduler.py
+++ /dev/null
@@ -1,107 +0,0 @@
-# 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.
-import platform
-import pytest
-import os
-
-from os import path
-
-from tvm import auto_scheduler
-from tvm.driver import tvmc
-
-
-def _get_tasks(model):
- tvmc_model = tvmc.frontends.load_model(model)
- tasks, weights = tvmc.autotuner.autoscheduler_get_tuning_tasks(
- tvmc_model.mod, tvmc_model.params, "llvm"
- )
- return (tasks, weights)
-
-
-def _autoscheduler_test_helper(model, tmpdir_name, early_stopping=1, prior_records=None):
- tvmc_model = tvmc.frontends.load_model(model)
- log_file = os.path.join(tmpdir_name, "autoscheduler.json")
-
- hardware_params = auto_scheduler.HardwareParams(num_cores=4, target="llvm")
-
- tvmc.tune(
- tvmc_model,
- target="llvm",
- tuning_records=log_file,
- prior_records=prior_records,
- early_stopping=early_stopping,
- enable_autoscheduler=True,
- trials=2,
- hardware_params=hardware_params,
- )
-
- # testing whether the log file was produced
- assert path.exists(log_file), "autoscheduler log file should exist"
-
- with auto_scheduler.ApplyHistoryBest(log_file) as best:
- assert isinstance(
- best, auto_scheduler.dispatcher.ApplyHistoryBest
- ), "unable to load the best results of tuning"
-
- return log_file
-
-
-def test_get_tuning_tasks(keras_simple):
- pytest.importorskip("tensorflow")
-
- tasks, weights = _get_tasks(keras_simple)
- expected_task_type = auto_scheduler.SearchTask
-
- assert type(tasks) is list
- assert len(tasks) > 0
- assert all([type(x) is expected_task_type for x in tasks]) is True
-
-
-@pytest.mark.skipif(
- platform.machine() == "aarch64",
- reason="Currently failing on AArch64 - see https://github.com/apache/tvm/issues/10673",
-)
-def test_tune_tasks(keras_simple, tmpdir_factory):
- pytest.importorskip("tensorflow")
-
- tmpdir_name = tmpdir_factory.mktemp("data")
- _autoscheduler_test_helper(keras_simple, tmpdir_name)
-
-
-@pytest.mark.skipif(
- platform.machine() == "aarch64",
- reason="Currently failing on AArch64 - see https://github.com/apache/tvm/issues/10673",
-)
-def test_tune_tasks__tuning_records(keras_simple, tmpdir_factory):
- pytest.importorskip("tensorflow")
-
- tmpdir_name = tmpdir_factory.mktemp("data")
- output_log_phase_1 = _autoscheduler_test_helper(keras_simple, tmpdir_name)
-
- # Exercises transfer learning by making sure a previous log exists
- _autoscheduler_test_helper(keras_simple, tmpdir_name, prior_records=output_log_phase_1)
-
-
-@pytest.mark.skipif(
- platform.machine() == "aarch64",
- reason="Currently failing on AArch64 - see https://github.com/apache/tvm/issues/10673",
-)
-def test_tune_tasks__no_early_stopping(keras_simple, tmpdir_factory):
- pytest.importorskip("tensorflow")
-
- tmpdir_name = tmpdir_factory.mktemp("data")
- _autoscheduler_test_helper(keras_simple, tmpdir_name, early_stopping=None)
diff --git a/tests/python/driver/tvmc/test_autotuner.py b/tests/python/driver/tvmc/test_autotuner.py
deleted file mode 100644
index ce5b888f25..0000000000
--- a/tests/python/driver/tvmc/test_autotuner.py
+++ /dev/null
@@ -1,309 +0,0 @@
-# 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.
-import platform
-import pytest
-import os
-
-from unittest import mock
-
-from os import path
-from pathlib import Path
-
-import tvm
-import tvm.testing
-from tvm import autotvm, auto_scheduler
-from tvm.driver import tvmc
-from tvm.driver.tvmc.autotuner import filter_tasks, gen_task_list
-
-
-def _get_tasks(model):
- tvmc_model = tvmc.frontends.load_model(model)
- return tvmc.autotuner.autotvm_get_tuning_tasks(tvmc_model.mod, tvmc_model.params, "llvm")
-
-
-def _get_measure_options():
- return autotvm.measure_option(
- builder=autotvm.LocalBuilder(build_func="default"), runner="local"
- )
-
-
-def _tuner_test_helper(model, tuner_name, tmpdir_name, early_stopping=1, prior_records=None):
- tvmc_model = tvmc.frontends.load_model(model)
- log_file = os.path.join(tmpdir_name, "log_{}.txt".format(tuner_name))
-
- tvmc.tune(
- tvmc_model,
- target="llvm",
- tuning_records=log_file,
- prior_records=prior_records,
- tuner=tuner_name,
- trials=4,
- early_stopping=early_stopping,
- )
-
- # testing whether the log file was produced
- assert path.exists(log_file), "tuning log file should exist"
-
- with autotvm.apply_history_best(log_file) as best:
- assert isinstance(
- best, autotvm.task.dispatcher.ApplyHistoryBest
- ), "unable to load the best results of tuning"
-
- return log_file
-
-
-def test_get_tuning_tasks(onnx_mnist):
- pytest.importorskip("onnx")
-
- sut = _get_tasks(onnx_mnist)
- expected_task_type = autotvm.task.Task
-
- assert type(sut) is list
- assert len(sut) > 0
- assert all([type(x) is expected_task_type for x in sut]) is True
-
-
-@pytest.mark.skipif(
- platform.machine() == "aarch64",
- reason="Currently failing on AArch64 - see https://github.com/apache/tvm/issues/10673",
-)
-def test_tune_tasks__tuner__xgb(onnx_mnist, tmpdir_factory):
- pytest.importorskip("onnx")
-
- tmpdir_name = tmpdir_factory.mktemp("data")
- _tuner_test_helper(onnx_mnist, "xgb", tmpdir_name)
-
-
-def test_tune_tasks__tuner__xgb_knob(onnx_mnist, tmpdir_factory):
- pytest.importorskip("onnx")
-
- tmpdir_name = tmpdir_factory.mktemp("data")
- _tuner_test_helper(onnx_mnist, "xgb_knob", tmpdir_name)
-
-
-def test_tune_tasks__tuner__ga(onnx_mnist, tmpdir_factory):
- pytest.importorskip("onnx")
-
- tmpdir_name = tmpdir_factory.mktemp("data")
- _tuner_test_helper(onnx_mnist, "ga", tmpdir_name)
-
-
-def test_tune_tasks__tuner__random(onnx_mnist, tmpdir_factory):
- pytest.importorskip("onnx")
-
- tmpdir_name = tmpdir_factory.mktemp("data")
- _tuner_test_helper(onnx_mnist, "random", tmpdir_name)
-
-
-def test_tune_tasks__tuner__gridsearch(onnx_mnist, tmpdir_factory):
- pytest.importorskip("onnx")
-
- tmpdir_name = tmpdir_factory.mktemp("data")
- _tuner_test_helper(onnx_mnist, "gridsearch", tmpdir_name)
-
-
-def test_tune_tasks__tuner__gridsearch__tuning_records(onnx_mnist, tmpdir_factory):
- pytest.importorskip("onnx")
-
- tmpdir_name = tmpdir_factory.mktemp("data")
- output_log_phase_1 = _tuner_test_helper(onnx_mnist, "gridsearch", tmpdir_name)
-
- # Exercises transfer learning by making sure a previous log exists
- _tuner_test_helper(onnx_mnist, "gridsearch", tmpdir_name, prior_records=output_log_phase_1)
-
-
-def test_tune_tasks__tuner__ga__empty_tasks(tmpdir_factory):
- pytest.importorskip("onnx")
-
- tmpdir_name = tmpdir_factory.mktemp("data")
- log_file = os.path.join(tmpdir_name, "log_{}.txt".format("ga"))
-
- tvmc.autotuner.tune_tasks(
- tasks=[],
- log_file=log_file,
- measure_option=_get_measure_options(),
- tuner="ga",
- trials=1,
- early_stopping=1,
- )
-
-
-def test_tune_tasks__tuner__xgb__no_early_stopping(onnx_mnist, tmpdir_factory):
- pytest.importorskip("onnx")
-
- tmpdir_name = tmpdir_factory.mktemp("data")
- _tuner_test_helper(onnx_mnist, "xgb", tmpdir_name, early_stopping=None)
-
-
-@pytest.mark.skipif(
- platform.machine() == "aarch64",
- reason="Currently failing on AArch64 - see https://github.com/apache/tvm/issues/10673",
-)
-def test_tune_tasks__tuner__xgb__no_tuning_records(onnx_mnist, tmpdir_factory):
- pytest.importorskip("onnx")
-
- tmpdir_name = tmpdir_factory.mktemp("data")
- _tuner_test_helper(onnx_mnist, "xgb", tmpdir_name, prior_records=None)
-
-
-def test_tune_tasks__invalid_tuner(onnx_mnist, tmpdir_factory):
- pytest.importorskip("onnx")
-
- tasks = _get_tasks(onnx_mnist)
- log_file = os.path.join(tmpdir_factory.mktemp("data"), "log2.txt")
-
- with pytest.raises(tvmc.TVMCException):
- tvmc.autotuner.tune_tasks(tasks, log_file, _get_measure_options(), "invalid_tuner", 1, 1)
-
-
-@mock.patch("tvm.driver.tvmc.autotuner.auto_scheduler.HardwareParams", return_value=None)
-@mock.patch("tvm.driver.tvmc.autotuner.tune_model", return_value=None)
-@mock.patch("tvm.driver.tvmc.frontends.load_model", return_value=None)
-def test_tune_rpc_tracker_parsing(mock_load_model, mock_tune_model, mock_auto_scheduler):
- cli_args = mock.MagicMock()
- cli_args.rpc_tracker = "10.0.0.1:9999"
- # FILE is not used but it's set to a valid value here to avoid it being set
- # by mock to a MagicMock class, which won't pass the checks for valid FILE.
- fake_input_file = "./fake_input_file.tflite"
- Path(fake_input_file).touch()
- cli_args.FILE = fake_input_file
-
- tvmc.autotuner.drive_tune(cli_args)
-
- os.remove(fake_input_file)
-
- mock_tune_model.assert_called_once()
-
- # inspect the mock call, to search for specific arguments
- _, _, kwargs = mock_tune_model.mock_calls[0]
- assert "hostname" in kwargs
- assert "10.0.0.1" == kwargs["hostname"]
- assert "port" in kwargs
- assert 9999 == kwargs["port"]
-
-
-@mock.patch("tvm.transform.PassContext", return_value=tvm.transform.PassContext())
-def test_autotune_pass_context(mock_pc, onnx_mnist, tmpdir_factory):
- """
- Check that the pass context while tuning is as expected.
- """
- pytest.importorskip("onnx")
-
- tmpdir_name = tmpdir_factory.mktemp("data")
- _tuner_test_helper(onnx_mnist, "gridsearch", tmpdir_name)
-
- # AutoTVM overrides the pass context later in the pipeline to disable AlterOpLayout
- assert mock_pc.call_count == 2
- assert mock_pc.call_args_list[0][1]["opt_level"] == 3
-
-
-def test_filter_tasks_valid():
- filter_tasks(list(range(10)), "list") == ([], True)
- filter_tasks(list(range(10)), "help") == ([], True)
- filter_tasks(list(range(10)), "all") == ([0, 1, 2, 3, 4, 5, 6, 7, 8, 9], False)
- filter_tasks(list(range(10)), "5") == ([5], False)
- filter_tasks(list(range(10)), "1-5") == ([1, 2, 3, 4, 5], False)
- filter_tasks(list(range(10)), "-5") == ([0, 1, 2, 3, 4, 5], False)
- filter_tasks(list(range(10)), "6-") == ([6, 7, 8, 9], False)
- filter_tasks(list(range(10)), "0,1-3,all") == ([0, 1, 2, 3, 4, 5, 6, 7, 8, 9], False)
- filter_tasks(list(range(10)), "0,4-5,9,list") == ([0, 4, 5, 9], True)
-
-
-@pytest.mark.parametrize(
- "value,err_msg",
- [
- ("10", "Task index out of range"),
- ("5,10", "Task index out of range"),
- ("1-10", "Right-hand side expression out of range"),
- ("-10", "Right-hand side expression out of range"),
- ("-", "Missing lhs or rhs for range expression"),
- ("-10-", "Malformed range expression"),
- ("--", "Malformed range expression"),
- ],
-)
-def test_filter_tasks_invalid(value, err_msg):
- with pytest.raises(AssertionError, match=err_msg):
- filter_tasks(list(range(10)), value)
-
-
-@pytest.mark.parametrize(
- "enable_autoscheduler,expected",
- [
- (
- False,
- """Available Tasks for tuning:
- 0. Task(func_name=taskA, args=[], kwargs={}, workload=('taskA',)) (len=?)
- 1. Task(func_name=taskBtaskBtaskBtaskBtaskBtaskBtaskBtaskBtaskBtaskBtaskBtaskBtaskBtaskBtaskBtaskBta... (len=?)
- 2. Task(func_name=taskC, args=[], kwargs={}, workload=('taskC',)) (len=?)""",
- ),
- (
- True,
- """Available Tasks for tuning:
- 0. taskA
- 1. taskBtaskBtaskBtaskBtaskBtaskBtaskBtaskBtaskBtaskBtaskBtaskBtaskBtaskBtaskBtaskBtaskBtaskBtaskBta...
- 2. Unnamed""",
- ),
- ],
-)
-def test_print_task_list(enable_autoscheduler, expected):
- if enable_autoscheduler:
- auto_scheduler.search_task.TASK_INPUT_BUFFER_TABLE.clear()
- N = 64
- target = "llvm"
- test_input_0 = tvm.runtime.ndarray.empty((64, 64))
- test_input_1 = tvm.runtime.ndarray.empty((10, 20))
- test_input_2 = tvm.runtime.ndarray.empty((30, 40, 50))
- task_inputs = {
- "test_input_0": test_input_0,
- "test_input_1": test_input_1,
- "test_input_2": test_input_2,
- }
- task1 = auto_scheduler.SearchTask(
- func="matmul_auto_scheduler_test",
- args=(N, N, N),
- target=target,
- task_inputs=task_inputs,
- task_inputs_overwrite=True,
- desc="taskA",
- )
- task2 = auto_scheduler.SearchTask(
- func="matmul_auto_scheduler_test",
- args=(N, N, N),
- target=target,
- task_inputs=task_inputs,
- task_inputs_overwrite=True,
- desc="taskB" * 20, # very long name
- )
- task3 = auto_scheduler.SearchTask(
- func="matmul_auto_scheduler_test",
- args=(N, N, N),
- target=target,
- task_inputs=task_inputs,
- task_inputs_overwrite=True,
- # missing description
- )
- else:
- task1 = autotvm.task.Task("taskA", [])
- task2 = autotvm.task.Task("taskB" * 20, []) # very long name
- task3 = autotvm.task.Task("taskC", [])
- tasks = [task1, task2, task3]
- out = gen_task_list(tasks, enable_autoscheduler)
- assert out == expected
-
-
-if __name__ == "__main__":
- tvm.testing.main()
diff --git a/tests/python/driver/tvmc/test_command_line.py b/tests/python/driver/tvmc/test_command_line.py
deleted file mode 100644
index af6cf8a26f..0000000000
--- a/tests/python/driver/tvmc/test_command_line.py
+++ /dev/null
@@ -1,332 +0,0 @@
-# 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.
-import os
-import platform
-import pytest
-import shutil
-import logging
-import sys
-
-from unittest import mock
-
-import tvm
-from tvm.driver.tvmc.main import _main
-from tvm.driver.tvmc import compiler
-from unittest.mock import MagicMock
-
-
-@pytest.mark.skipif(
- platform.machine() == "aarch64",
- reason="Currently failing on AArch64 - see https://github.com/apache/tvm/issues/10673",
-)
-def test_tvmc_cl_workflow(keras_simple, tmpdir_factory):
- pytest.importorskip("tensorflow")
-
- tmpdir = tmpdir_factory.mktemp("data")
-
- # Test model tuning
- log_path = os.path.join(tmpdir, "keras-autotuner_records.json")
- tuning_str = (
- f"tvmc tune --target llvm --output {log_path} "
- f"--trials 2 --enable-autoscheduler {keras_simple}"
- )
- tuning_args = tuning_str.split(" ")[1:]
- _main(tuning_args)
- assert os.path.exists(log_path)
-
- # Test model compilation
- package_path = os.path.join(tmpdir, "keras-tvm.tar")
- compile_str = (
- f"tvmc compile --target llvm --tuning-records {log_path} "
- f"--output {package_path} {keras_simple}"
- )
- compile_args = compile_str.split(" ")[1:]
- _main(compile_args)
- assert os.path.exists(package_path)
-
- # Test running the model
- output_path = os.path.join(tmpdir, "predictions.npz")
- run_str = f"tvmc run --end-to-end --outputs {output_path} {package_path}"
- run_args = run_str.split(" ")[1:]
- _main(run_args)
- assert os.path.exists(output_path)
-
-
-@pytest.mark.skipif(
- platform.machine() == "aarch64",
- reason="Currently failing on AArch64 - see https://github.com/apache/tvm/issues/10673",
-)
-def test_tvmc_cl_workflow_json_config(keras_simple, tmpdir_factory):
- pytest.importorskip("tensorflow")
- tune_config_file = "tune_config_test"
- tmpdir = tmpdir_factory.mktemp("data")
-
- # Test model tuning
- log_path = os.path.join(tmpdir, "keras-autotuner_records.json")
- tuning_str = (
- f"tvmc tune --config {tune_config_file} --output {log_path} "
- f"--enable-autoscheduler {keras_simple}"
- )
- tuning_args = tuning_str.split(" ")[1:]
- _main(tuning_args)
- assert os.path.exists(log_path)
-
- # Test model compilation
- package_path = os.path.join(tmpdir, "keras-tvm.tar")
- compile_str = (
- f"tvmc compile --tuning-records {log_path} " f"--output {package_path} {keras_simple}"
- )
- compile_args = compile_str.split(" ")[1:]
- _main(compile_args)
- assert os.path.exists(package_path)
-
- # Test running the model
- output_path = os.path.join(tmpdir, "predictions.npz")
- run_str = f"tvmc run --outputs {output_path} {package_path}"
- run_args = run_str.split(" ")[1:]
- _main(run_args)
- assert os.path.exists(output_path)
-
-
-@pytest.fixture
-def missing_file():
- missing_file_name = "missing_file_as_invalid_input.tfite"
- return missing_file_name
-
-
-@pytest.fixture
-def broken_symlink(tmp_path):
- broken_symlink = "broken_symlink_as_invalid_input.tflite"
- os.symlink("non_existing_file", tmp_path / broken_symlink)
- yield broken_symlink
- os.unlink(tmp_path / broken_symlink)
-
-
-@pytest.fixture
-def fake_directory(tmp_path):
- dir_as_invalid = "dir_as_invalid_input.tflite"
- os.mkdir(tmp_path / dir_as_invalid)
- yield dir_as_invalid
- shutil.rmtree(tmp_path / dir_as_invalid)
-
-
-@pytest.mark.parametrize(
- "invalid_input",
- ["missing_file", "broken_symlink", "fake_directory"],
-)
-def test_tvmc_compile_file_check(capsys, invalid_input, request):
- invalid_input = request.getfixturevalue(invalid_input)
- compile_cmd = f"tvmc compile --target 'c' {invalid_input}"
- run_arg = compile_cmd.split(" ")[1:]
-
- _main(run_arg)
-
- captured = capsys.readouterr()
- expected_err = (
- f"Error: Input file '{invalid_input}' doesn't exist, "
- "is a broken symbolic link, or a directory.\n"
- )
- on_assert_error = f"'tvmc compile' failed to check invalid FILE: {invalid_input}"
- assert captured.err == expected_err, on_assert_error
-
-
-@pytest.mark.parametrize(
- "invalid_input",
- ["missing_file", "broken_symlink", "fake_directory"],
-)
-def test_tvmc_tune_file_check(capsys, invalid_input, request):
- invalid_input = request.getfixturevalue(invalid_input)
- tune_cmd = f"tvmc tune --target 'llvm' --output output.json {invalid_input}"
- run_arg = tune_cmd.split(" ")[1:]
-
- _main(run_arg)
-
- captured = capsys.readouterr()
- expected_err = (
- f"Error: Input file '{invalid_input}' doesn't exist, "
- "is a broken symbolic link, or a directory.\n"
- )
- on_assert_error = f"'tvmc tune' failed to check invalid FILE: {invalid_input}"
- assert captured.err == expected_err, on_assert_error
-
-
-@mock.patch("tvm.relay.build", side_effect=tvm.relay.build)
-@mock.patch("tvm.driver.tvmc.model.TVMCPackage.__init__", return_value=None)
-def test_tvmc_workspace_pools_check(mock_pkg, mock_relay, keras_simple, tmpdir_factory):
- pytest.importorskip("tensorflow")
- tmpdir = tmpdir_factory.mktemp("data")
-
- # Test model compilation
- package_path = os.path.join(tmpdir, "keras-tvm.tar")
- compile_str = (
- f"tvmc compile --target=llvm --workspace-pools=sram "
- f"--workspace-pools-targets=sram:llvm "
- f"--output={package_path} {keras_simple}"
- )
- compile_args = compile_str.split(" ")[1:]
- _main(compile_args)
- assert os.path.exists(package_path)
- assert mock_relay.call_count == 1
- assert mock_relay.call_args_list[0][1]["workspace_memory_pools"].pools[0].pool_name == "sram"
-
-
-@pytest.fixture
-def paddle_model(paddle_resnet50):
- # If we can't import "paddle" module, skip testing paddle as the input model.
- if pytest.importorskip("paddle", reason="'paddle' module not installed"):
- return paddle_resnet50
-
-
-@pytest.mark.parametrize(
- "model",
- [
- "paddle_model",
- ],
-)
-# compile_model() can take too long and is tested elsewhere, hence it's mocked below
-@mock.patch.object(compiler, "compile_model")
-# @mock.patch.object(compiler, "compile_model")
-def test_tvmc_compile_input_model(mock_compile_model, tmpdir_factory, model, request):
-
- model = request.getfixturevalue(model)
- output_dir = tmpdir_factory.mktemp("output")
- output_file = output_dir / "model.tar"
-
- compile_cmd = (
- f"tvmc compile --target 'llvm' {model} --model-format paddle --output {output_file}"
- )
- run_arg = compile_cmd.split(" ")[1:]
-
- _main(run_arg)
-
- mock_compile_model.assert_called_once()
-
-
-def test_tvmc_logger(caplog, tmpdir_factory, keras_simple):
- pytest.importorskip("tensorflow")
- tmpdir = tmpdir_factory.mktemp("out")
-
... 23281 lines suppressed ...