You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@tvm.apache.org by GitBox <gi...@apache.org> on 2021/01/20 08:18:01 UTC
[GitHub] [tvm] jcf94 opened a new pull request #7313: [AutoSchedule] Sparse dense tuning support with custom sketch rule
jcf94 opened a new pull request #7313:
URL: https://github.com/apache/tvm/pull/7313
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] merrymercy edited a comment on pull request #7313: [AutoSchedule] Sparse dense tuning support with custom sketch rule
Posted by GitBox <gi...@apache.org>.
merrymercy edited a comment on pull request #7313:
URL: https://github.com/apache/tvm/pull/7313#issuecomment-789393341
1. Update the type of `SearchTaskNode::task_inputs`. Change it from `Map<String, runtime::NDArray>` to `Array<String>`, so we only need to store nd arrays in one place.
2. Remove `SearchTask.AddTaskInput` interface to make `SearchTask` immutable. We do not have the need to dynamically update task inputs, so we can provide all arguments to the constructors.
3. Make sure we can use the same interface to support the use case where we want to match the special buffers by name
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] tkonolige commented on a change in pull request #7313: [AutoSchedule] Sparse dense tuning support with custom sketch rule
Posted by GitBox <gi...@apache.org>.
tkonolige commented on a change in pull request #7313:
URL: https://github.com/apache/tvm/pull/7313#discussion_r566970128
##########
File path: python/tvm/auto_scheduler/measure.py
##########
@@ -719,6 +722,87 @@ def local_builder_build(inputs, timeout, n_parallel, build_func="default", verbo
return results
+def _process_sparse_input(args):
+ sparse_prefix = sparse_data = sparse_indices = sparse_indptr = None
+
+ def _process_inputs(input_tensors, M, N, prefix_init):
+ nonlocal sparse_prefix
+ nonlocal sparse_data
+ nonlocal sparse_indices
+ nonlocal sparse_indptr
+
+ assert len(input_tensors) == 4
+ unsure_tensors = list(input_tensors)
+ # Get the Dense data
+ dense_data = None
+ for tensor in unsure_tensors:
+ if len(tensor.shape) == 2:
+ assert dense_data is None
+ dense_data = tensor
+ assert M == dense_data.shape[0]
+ K = dense_data.shape[1]
+ unsure_tensors.remove(dense_data)
+
+ # Get the Sparse data
+ sparse_data = None
+ for tensor in unsure_tensors:
+ if len(tensor.shape) == 3:
+ assert sparse_data is None
+ sparse_data = tensor
+ block_size, BS_R, BS_C = sparse_data.shape
+ unsure_tensors.remove(sparse_data)
+
+ # Get the Sparse indptr & indices
+ sparse_indices = None
+ for tensor in unsure_tensors:
+ assert len(tensor.shape) == 1
+ if tensor.shape[0] == block_size:
+ assert sparse_indices is None
+ sparse_indices = tensor
+ unsure_tensors.remove(sparse_indices)
+ assert len(unsure_tensors) == 1
+ sparse_indptr = unsure_tensors[0]
+
+ # Generate the sparse_prefix
+ density = 1.0
+ for i in sparse_data.shape:
+ density *= i
+ density /= (K * N)
+ density = density.value
+ sparse_prefix = "%s_%d_%d_%d_%d_%d_%.2f_" % (
Review comment:
You could hash the `indptr` and `indices` arrays as these determine the structure. Alternatively you could hash the number of nonzeros per row.
It would be interesting to study if tuning performs the same independent of structure (but for the same sparsity).
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] jcf94 commented on a change in pull request #7313: [AutoSchedule] Sparse dense tuning support with custom sketch rule
Posted by GitBox <gi...@apache.org>.
jcf94 commented on a change in pull request #7313:
URL: https://github.com/apache/tvm/pull/7313#discussion_r586979630
##########
File path: python/tvm/auto_scheduler/measure.py
##########
@@ -719,6 +720,45 @@ def local_builder_build(inputs, timeout, n_parallel, build_func="default", verbo
return results
+def _prepare_input_map(args):
+ """This function deals with special task inputs.
+
+ Parameters
+ ----------
+ args : List[Tensor]
+ Input/output Tensor of a TVM subgraph.
+
+ Returns
+ -------
+ A Dict[Tensor, str] that maps the input Tensor to a buffer name.
+
+ Note
+ ----
+ The buffer name is specially designed, and these buffer should be provided in
+ `SearchTask(..., task_inputs={...})`.
+ """
+ # pylint: disable=import-outside-toplevel
+ from tvm import topi # lazily import to avoid recursive dependency
+
+ # A dict that maps the input tensor arg to a buffer name
+ tensor_input_map = {}
+
+ # Case 0: Check placeholder name
+ for arg in args:
+ if isinstance(arg.op, tvm.te.PlaceholderOp):
+ if arg.op.name != "placeholder":
+ tensor_input_map[arg] = arg.op.name
+
+ # Case 1: Check sparse op
+ sparse_input_map = topi.nn.sparse.try_get_sparse_input(args)
Review comment:
By the way, my colleague is going to add Ansor support for sparse_conv2d. We'll add extra check to this entry first, and see if there's any better way to merge them.
##########
File path: python/tvm/auto_scheduler/measure.py
##########
@@ -719,6 +720,45 @@ def local_builder_build(inputs, timeout, n_parallel, build_func="default", verbo
return results
+def _prepare_input_map(args):
+ """This function deals with special task inputs.
+
+ Parameters
+ ----------
+ args : List[Tensor]
+ Input/output Tensor of a TVM subgraph.
+
+ Returns
+ -------
+ A Dict[Tensor, str] that maps the input Tensor to a buffer name.
+
+ Note
+ ----
+ The buffer name is specially designed, and these buffer should be provided in
+ `SearchTask(..., task_inputs={...})`.
+ """
+ # pylint: disable=import-outside-toplevel
+ from tvm import topi # lazily import to avoid recursive dependency
+
+ # A dict that maps the input tensor arg to a buffer name
+ tensor_input_map = {}
+
+ # Case 0: Check placeholder name
+ for arg in args:
+ if isinstance(arg.op, tvm.te.PlaceholderOp):
+ if arg.op.name != "placeholder":
+ tensor_input_map[arg] = arg.op.name
+
+ # Case 1: Check sparse op
+ sparse_input_map = topi.nn.sparse.try_get_sparse_input(args)
Review comment:
Yeah, I've also had some discussions in our weekly sync while didn't figure out any better solutions.
There're several reasons:
1. Different ops have different requirements over specific inputs;
2. While the problems is in a subgraph generated in Relay Integration, the placeholder are all the same, we can not differentiate them from tag, name or any other way, even the order of inputs are not guaranteed.
Current approach is to merge all specific inputs checking to this function, at least they have a same entry here. For the other ops, you have to add their own check functions below.
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] merrymercy commented on pull request #7313: [AutoSchedule] Sparse dense tuning support with custom sketch rule
Posted by GitBox <gi...@apache.org>.
merrymercy commented on pull request #7313:
URL: https://github.com/apache/tvm/pull/7313#issuecomment-789393341
1. Update the type of `SearchTaskNode::task_inputs`. Change it from `Map<String, runtime::NDArray>` to `Array<String>`, so we only need to store nd arrays in one place.
2. Remove `SearchTask.AddTaskInput` interface to make `SearchTask` immutable. We do not have the need to dynamically update task inputs, so we can provide all arguments to the constructors.
3. Make sure we can use the same interface to do the support the most simple use case where we want to match the special buffers by name
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] tkonolige commented on a change in pull request #7313: [AutoSchedule] Sparse dense tuning support with custom sketch rule
Posted by GitBox <gi...@apache.org>.
tkonolige commented on a change in pull request #7313:
URL: https://github.com/apache/tvm/pull/7313#discussion_r586586361
##########
File path: python/tvm/auto_scheduler/measure.py
##########
@@ -719,6 +720,45 @@ def local_builder_build(inputs, timeout, n_parallel, build_func="default", verbo
return results
+def _prepare_input_map(args):
+ """This function deals with special task inputs.
+
+ Parameters
+ ----------
+ args : List[Tensor]
+ Input/output Tensor of a TVM subgraph.
+
+ Returns
+ -------
+ A Dict[Tensor, str] that maps the input Tensor to a buffer name.
+
+ Note
+ ----
+ The buffer name is specially designed, and these buffer should be provided in
+ `SearchTask(..., task_inputs={...})`.
+ """
+ # pylint: disable=import-outside-toplevel
+ from tvm import topi # lazily import to avoid recursive dependency
+
+ # A dict that maps the input tensor arg to a buffer name
+ tensor_input_map = {}
+
+ # Case 0: Check placeholder name
+ for arg in args:
+ if isinstance(arg.op, tvm.te.PlaceholderOp):
+ if arg.op.name != "placeholder":
+ tensor_input_map[arg] = arg.op.name
+
+ # Case 1: Check sparse op
+ sparse_input_map = topi.nn.sparse.try_get_sparse_input(args)
Review comment:
I think I asked this before, but can we have a more general mechanism than checking only for sparse. There are other use cases that require specific input (sorting, scatter).
##########
File path: tests/python/unittest/test_auto_scheduler_search_task.py
##########
@@ -0,0 +1,211 @@
+# 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 test_auto_scheduler_common import (
+ matmul_auto_scheduler_test,
+ zero_rank_compute_auto_scheduler_test,
+ zero_rank_reduce_auto_scheduler_test,
+)
+import multiprocessing
Review comment:
Double unused import of multiprocessing.
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] jcf94 commented on a change in pull request #7313: [AutoSchedule] Sparse dense tuning support with custom sketch rule
Posted by GitBox <gi...@apache.org>.
jcf94 commented on a change in pull request #7313:
URL: https://github.com/apache/tvm/pull/7313#discussion_r586192659
##########
File path: tutorials/auto_scheduler/tune_sparse_x86.py
##########
@@ -0,0 +1,331 @@
+# 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.
+"""
+Auto-scheduling Sparse Matrix Multiplication for CPU by Custom Sketch Rule
+==========================================================================
+**Author**: `Chengfan Jia <https://github.com/jcf94/>`_
+
+This is a tutorial on how to use the auto-scheduler to tune a sparse matrix multiplication for
+CPUs.
+
+Auto-scheduler is designed to explore the schedule with best performance for a given computation
+declaration automatically. While sometimes, we may have a demand to try some special ops which may
+not been well supported by auto-scheduler's default search policy. Auto-scheduler currently allows
+user to provide a CustomSketch to cover these cases.
+
+We use sparse matrix multiplication as an example in this tutorial.
+
+Note that this tutorial will not run on Windows or recent versions of macOS. To
+get it to run, you will need to wrap the body of this tutorial in a :code:`if
+__name__ == "__main__":` block.
+"""
+
+import os
+import itertools
+
+import numpy as np
+import tvm
+from tvm import te, auto_scheduler, topi
+from tvm.auto_scheduler import _ffi_api
+from tvm.topi.utils import get_const_tuple
+
+import scipy.sparse as sp
+
+######################################################################
+# Define the computation
+# ^^^^^^^^^^^^^^^^^^^^^^
+# To begin with, let us define the computation of a sparse matmul with several relu and bias add.
+# The function should return the list of input/output tensors.
+# From these tensors, the auto-scheduler can get the whole computational graph.
+
+# We use this function to generate a random bsr matrix
+def random_bsr_matrix(M, N, BS_R, BS_C, density, dtype):
+ import itertools
+
+ Y = np.zeros((M, N), dtype=dtype)
+ assert M % BS_R == 0
+ assert N % BS_C == 0
+ nnz = int(density * M * N)
+ num_blocks = int(nnz / (BS_R * BS_C)) + 1
+ candidate_blocks = np.asarray(list(itertools.product(range(0, M, BS_R), range(0, N, BS_C))))
+ assert candidate_blocks.shape[0] == M // BS_R * N // BS_C
+ chosen_blocks = candidate_blocks[
+ np.random.choice(candidate_blocks.shape[0], size=num_blocks, replace=False)
+ ]
+ for i in range(len(chosen_blocks)):
+ r, c = chosen_blocks[i]
+ Y[r : r + BS_R, c : c + BS_C] = np.random.randn(BS_R, BS_C)
+ s = sp.bsr_matrix(Y, blocksize=(BS_R, BS_C))
+ assert s.data.shape == (num_blocks, BS_R, BS_C)
+ assert s.indices.shape == (num_blocks,)
+ assert s.indptr.shape == (M // BS_R + 1,)
+ return s
+
+@auto_scheduler.register_workload
+def sparse_dense(M, N, K, w_data_shape, w_indices_shape, w_indptr_shape, dtype):
+ X = te.placeholder(shape=(M, K), dtype=dtype)
+ W_data = te.placeholder(shape=w_data_shape, dtype=dtype)
+ W_indices = te.placeholder(shape=w_indices_shape, dtype="int32")
+ W_indptr = te.placeholder(shape=w_indptr_shape, dtype="int32")
+ B = te.placeholder(shape=(M, N), dtype=dtype)
+
+ out = topi.nn.sparse_dense(
+ topi.nn.relu(X), W_data, W_indices, W_indptr
+ )
+ out = te.compute((M, N), lambda i, j: out[i, j] + B[i, j], name="BiasAdd")
+ out = topi.nn.relu(out)
+
+ return [X, W_data, W_indices, W_indptr, B, out]
+
+######################################################################
+# Special step for sparse workload
+# ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+# During schedule tuning, auto-scheduler will use random inputs to measure the performance of a
+# generated schedule. While we cannot directly use a random array as the input of a sparse op, for
+# the "indices" and "indptr" array are meaningful for the computation.
+#
+# To solve this problem, we register these as special buffers, and load them when process program
+# measuring.
+# See the :any:`auto_scheduler.measure` code for more details.
+
+# Define the basic shapes of this sparse computation
+M = K = N = 512
+BS_R = 16
+BS_C = 1
+density = 0.6
+
+# Generate the test data with numpy
+X_np = np.random.randn(M, K).astype("float32")
+X_np = np.maximum(np.zeros((M, K), dtype="float32"), X_np) # Relu
+W_sp_np = random_bsr_matrix(N, K, BS_R, BS_C, density=density, dtype="float32")
+W_np = W_sp_np.todense()
+Y_np = X_np @ W_np.T # Process the matrix multiplication
+B_np = np.random.randn(M, N).astype("float32")
+Y_np = Y_np + B_np # Bias add
+Y_np = np.maximum(np.zeros((M, N), dtype="float32"), Y_np) # Relu
+
+# Register the sparse data to special buffer
+prefix = "sparse_dense_bsr_%d_%d_%d_%d_%d_%.2f_" % (M, N, K, BS_R, BS_C, density)
+auto_scheduler.measure.register_special_buffer(prefix + "W_data", W_sp_np.data)
+auto_scheduler.measure.register_special_buffer(prefix + "W_indices", W_sp_np.indices)
+auto_scheduler.measure.register_special_buffer(prefix + "W_indptr", W_sp_np.indptr)
+
+######################################################################
+# Create the search task
+# ^^^^^^^^^^^^^^^^^^^^^^
+# We then create a search task with M=N=K=512 and dtype="float32"
+# If your machine supports avx instructions, you can
+#
+# - replace "llvm" below with "llvm -mcpu=core-avx2" to enable AVX2
+# - replace "llvm" below with "llvm -mcpu=skylake-avx512" to enable AVX-512
+
+target = tvm.target.Target("llvm")
+
+task = tvm.auto_scheduler.SearchTask(
+ func=sparse_dense,
+ args=(
+ M, N, K,
+ W_sp_np.data.shape,
+ W_sp_np.indices.shape,
+ W_sp_np.indptr.shape,
+ "float32"
+ ),
+ target=target
+)
+
+# Inspect the computational graph
+print("Computational DAG:")
+print(task.compute_dag)
+
+######################################################################
+# Write the custom sketch for sparse dense op
+# ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+# Before tuning, we will need to define the CustomSketchRule for the sparse dense op.
+#
+# CustomSketchRule consists of two parts: the condition function and the apply function.
+#
+# - condition function: describe when to use this sketch rule. For example, we can match the op
+# by their name or tag.
+# - apply function: describe how to generate the initial sketch. Auto-scheduler provides a set of
+# loop state APIs.
+
+def meet_condition_func(search_policy, state, stage_id):
+ state = auto_scheduler.loop_state.State(state, search_policy.search_task.compute_dag)
+ if state.stages[stage_id].op.tag in [
+ "sparse_dense_sp_rhs_bsrmm", "sparse_dense_sp_rhs_bsrmm_block"
+ ]:
+ return auto_scheduler.PreloadCustomSketchRule.APPLY_AND_SKIP_REST
+ else:
+ return auto_scheduler.PreloadCustomSketchRule.PASS
+
+def apply_func(search_policy, state, stage_id):
+ ret = []
+ s0 = auto_scheduler.loop_state.State(state, search_policy.search_task.compute_dag)
+ if s0.stages[stage_id].op.tag == "sparse_dense_sp_rhs_bsrmm_block":
+ return [s0.state_object, stage_id - 1]
+
+ sparse_dense = s0.stages[stage_id].op
+ sparse_dense_block = s0.stages[stage_id - 1].op
+ assert sparse_dense.tag == "sparse_dense_sp_rhs_bsrmm"
+ assert sparse_dense_block.tag == "sparse_dense_sp_rhs_bsrmm_block"
+
+ # Set the default consumer of compute block
+ consumer = sparse_dense
+
+ # If sparse dense has a single elementwise consumer
+ # We can compute inline the sparse_dense output stage
+ consumers = _ffi_api.SearchPolicyUtilsGetConsumers(
+ search_policy.search_task, s0.state_object, stage_id
+ )
+ if len(consumers) == 1:
+ consumer_id = int(consumers.items()[0][0])
+ if _ffi_api.SearchPolicyUtilsIsElementwiseMatch(
+ search_policy.search_task, s0.state_object, stage_id, consumer_id
+ ):
+ consumer = s0.stages[consumer_id].op
+ s0.compute_inline(sparse_dense)
+
+ i, nb_j, j, row_offset, c = s0[sparse_dense_block].iters
+ m, n = s0[consumer].iters
+ i0, i1, i2 = s0.split(sparse_dense_block, i, [None, None])
+ m0, m1 = s0.follow_split(consumer, m, len(s0.transform_steps) - 1, 1)
+ j0, j1 = s0.split(sparse_dense_block, nb_j, [None])
+ n0, n1 = s0.follow_split(consumer, n, len(s0.transform_steps) - 1, 1)
+ s0.reorder(sparse_dense_block, [i0, j0, i1, j1, row_offset, i2, j, c])
+ s0.reorder(consumer, [m0, n0, m1, n1])
+ s0.compute_at(sparse_dense_block, consumer, n0)
+
+ ret.append([s0.state_object, stage_id - 2])
+
+ return ret
+
+######################################################################
+# Next, we set parameters for the auto-scheduler.
+#
+# * :code:`num_measure_trials` is the number of measurement trials we can use during the search.
+# We only make 10 trials in this tutorial for a fast demonstration. In practice, 1000 is a
+# good value for the search to converge. You can do more trials according to your time budget.
+# * In addition, we use :code:`RecordToFile` to dump measurement records into a file `matmul.json`.
+# The measurement records can be used to query the history best, resume the search,
+# and do more analyses later.
+# * see :any:`auto_scheduler.TuningOptions` for more parameters
+# * Here, we need to create a :code:`auto_scheduler.SketchPolicy` object, and add the custom sketch
+# rule as a `init_search_callbacks`.
+
+log_file = "sparse_dense.json"
+tune_option = auto_scheduler.TuningOptions(
+ num_measure_trials=10,
+ measure_callbacks=[auto_scheduler.RecordToFile(log_file)],
+ verbose=2,
+)
+
+search_policy = auto_scheduler.SketchPolicy(
+ task,
+ program_cost_model=auto_scheduler.XGBModel(),
+ init_search_callbacks=[
+ auto_scheduler.PreloadCustomSketchRule(meet_condition_func, apply_func, "SparseDense")
+ ]
+)
+
+######################################################################
+# Run the search
+# ^^^^^^^^^^^^^^
+# Now we get all inputs ready.
+# We can kick off the search and let the auto-scheduler do its magic.
+# After some measurement trials, we can load the best schedule from the log
+# file and apply it.
+
+# Run auto-tuning (search)
+task.tune(tune_option, search_policy)
Review comment:
Ok, I think you're right...
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] comaniac commented on a change in pull request #7313: [AutoSchedule] Sparse dense tuning support with custom sketch rule
Posted by GitBox <gi...@apache.org>.
comaniac commented on a change in pull request #7313:
URL: https://github.com/apache/tvm/pull/7313#discussion_r587976799
##########
File path: python/tvm/auto_scheduler/search_task.py
##########
@@ -157,6 +164,149 @@ def __init__(
)
+# The map stores special registered buffer for measurement
+# This can be used for sparse workloads when we cannot use random tensors for measurment.
+# {
+# "workload_key_0": {
+# "task_input_0": Tensor(...),
+# "task_input_1": Tensor(...)
+# },
+# "workload_key_1": {
+# "task_input_2": Tensor(...),
+# "task_input_3": Tensor(...)
+# },
+# ...
+# }
+TASK_INPUT_BUFFER_TABLE = {}
+
+
+def _save_buffer_to_file(buffer_name, buffer_data):
+ """Save the current Tensor buffer to a numpy file.
+
+ File name will be: {buffer_name}.{buffer_shape}_{buffer_data_type}
+ """
+ np_data = buffer_data.asnumpy()
+
+ buffer_name += "."
+ for i in np_data.shape:
+ buffer_name += "%d_" % (i)
+ buffer_name += "%s" % (np_data.dtype)
+
+ np_data.tofile(buffer_name, " ")
+
+
+def _try_load_buffer_from_file(buffer_name):
+ """Try to load buffer from a numpy file, if not found, return None.
+
+ File name has a same format as `_save_buffer_to_file`.
+ """
+ filelist = os.listdir()
+
+ for file in filelist:
+ if file.startswith(buffer_name) and file.count("."):
+ meta_info = file.split(".")[-1].split("_")
+ shape = [int(i) for i in meta_info[:-1]]
+ dtype = meta_info[-1]
+ buffer_data = np.fromfile(file, dtype=dtype, sep=" ")
+ buffer_data = buffer_data.reshape(shape)
+ return ndarray.array(buffer_data)
+
+ return None
+
+
+def register_task_input_buffer(
+ workload_key,
+ input_name,
+ input_data,
+ overwrite=False,
+ save_to_file=False,
+):
+ """Register special buffer for measurement.
+
+ Parameters
+ ----------
+ workload_key : str
+ The workload key of the SearchTask.
+
+ input_name : str
+ The name of input buffer.
+
+ input_data : tvm.nd.NDArray
+ The input Tensor data.
+
+ overwrite : bool = False
+ Whether overwrite the data if a name has already in the global table.
+
+ save_to_file : bool = False
+ Whether record this buffer to a local file. This can be reused to continue the last tuning
+ process.
+ """
+ global TASK_INPUT_BUFFER_TABLE
+
+ if workload_key not in TASK_INPUT_BUFFER_TABLE:
+ TASK_INPUT_BUFFER_TABLE[workload_key] = {}
+ input_table = TASK_INPUT_BUFFER_TABLE[workload_key]
+
+ if not overwrite:
+ if input_name not in input_table.keys():
+ # Try to load buffer data from local file
+ tensor_from_file = _try_load_buffer_from_file(input_name)
+ if tensor_from_file:
+ input_table[input_name] = tensor_from_file
+
+ if input_name in input_table.keys():
Review comment:
Ah I see. That makes sense.
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] tkonolige commented on a change in pull request #7313: [AutoSchedule] Sparse dense tuning support with custom sketch rule
Posted by GitBox <gi...@apache.org>.
tkonolige commented on a change in pull request #7313:
URL: https://github.com/apache/tvm/pull/7313#discussion_r587703710
##########
File path: python/tvm/auto_scheduler/measure.py
##########
@@ -719,6 +720,45 @@ def local_builder_build(inputs, timeout, n_parallel, build_func="default", verbo
return results
+def _prepare_input_map(args):
+ """This function deals with special task inputs.
+
+ Parameters
+ ----------
+ args : List[Tensor]
+ Input/output Tensor of a TVM subgraph.
+
+ Returns
+ -------
+ A Dict[Tensor, str] that maps the input Tensor to a buffer name.
+
+ Note
+ ----
+ The buffer name is specially designed, and these buffer should be provided in
+ `SearchTask(..., task_inputs={...})`.
+ """
+ # pylint: disable=import-outside-toplevel
+ from tvm import topi # lazily import to avoid recursive dependency
+
+ # A dict that maps the input tensor arg to a buffer name
+ tensor_input_map = {}
+
+ # Case 0: Check placeholder name
+ for arg in args:
+ if isinstance(arg.op, tvm.te.PlaceholderOp):
+ if arg.op.name != "placeholder":
+ tensor_input_map[arg] = arg.op.name
+
+ # Case 1: Check sparse op
+ sparse_input_map = topi.nn.sparse.try_get_sparse_input(args)
Review comment:
Could we associate the lookup mechanism with `@register_workload`? It would at least be extensible then.
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] jcf94 commented on a change in pull request #7313: [AutoSchedule] Sparse dense tuning support with custom sketch rule
Posted by GitBox <gi...@apache.org>.
jcf94 commented on a change in pull request #7313:
URL: https://github.com/apache/tvm/pull/7313#discussion_r587975519
##########
File path: python/tvm/auto_scheduler/search_task.py
##########
@@ -157,6 +164,149 @@ def __init__(
)
+# The map stores special registered buffer for measurement
+# This can be used for sparse workloads when we cannot use random tensors for measurment.
+# {
+# "workload_key_0": {
+# "task_input_0": Tensor(...),
+# "task_input_1": Tensor(...)
+# },
+# "workload_key_1": {
+# "task_input_2": Tensor(...),
+# "task_input_3": Tensor(...)
+# },
+# ...
+# }
+TASK_INPUT_BUFFER_TABLE = {}
+
+
+def _save_buffer_to_file(buffer_name, buffer_data):
+ """Save the current Tensor buffer to a numpy file.
+
+ File name will be: {buffer_name}.{buffer_shape}_{buffer_data_type}
+ """
+ np_data = buffer_data.asnumpy()
+
+ buffer_name += "."
+ for i in np_data.shape:
+ buffer_name += "%d_" % (i)
+ buffer_name += "%s" % (np_data.dtype)
+
+ np_data.tofile(buffer_name, " ")
+
+
+def _try_load_buffer_from_file(buffer_name):
+ """Try to load buffer from a numpy file, if not found, return None.
+
+ File name has a same format as `_save_buffer_to_file`.
+ """
+ filelist = os.listdir()
+
+ for file in filelist:
+ if file.startswith(buffer_name) and file.count("."):
+ meta_info = file.split(".")[-1].split("_")
+ shape = [int(i) for i in meta_info[:-1]]
+ dtype = meta_info[-1]
+ buffer_data = np.fromfile(file, dtype=dtype, sep=" ")
+ buffer_data = buffer_data.reshape(shape)
+ return ndarray.array(buffer_data)
+
+ return None
+
+
+def register_task_input_buffer(
+ workload_key,
+ input_name,
+ input_data,
+ overwrite=False,
+ save_to_file=False,
+):
+ """Register special buffer for measurement.
+
+ Parameters
+ ----------
+ workload_key : str
+ The workload key of the SearchTask.
+
+ input_name : str
+ The name of input buffer.
+
+ input_data : tvm.nd.NDArray
+ The input Tensor data.
+
+ overwrite : bool = False
+ Whether overwrite the data if a name has already in the global table.
+
+ save_to_file : bool = False
+ Whether record this buffer to a local file. This can be reused to continue the last tuning
+ process.
+ """
+ global TASK_INPUT_BUFFER_TABLE
+
+ if workload_key not in TASK_INPUT_BUFFER_TABLE:
+ TASK_INPUT_BUFFER_TABLE[workload_key] = {}
+ input_table = TASK_INPUT_BUFFER_TABLE[workload_key]
+
+ if not overwrite:
+ if input_name not in input_table.keys():
+ # Try to load buffer data from local file
+ tensor_from_file = _try_load_buffer_from_file(input_name)
+ if tensor_from_file:
+ input_table[input_name] = tensor_from_file
+
+ if input_name in input_table.keys():
Review comment:
> I suppose missing inputs are more common, as most use cases still rely on random inputs?
Oh, I mean if this SearchTask has task_inpus, while Ansor does not map them all here, then we can raise a warning. It would be fine if a SearchTask does not have any task_inpus.
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] jcf94 commented on a change in pull request #7313: [AutoSchedule] Sparse dense tuning support with custom sketch rule
Posted by GitBox <gi...@apache.org>.
jcf94 commented on a change in pull request #7313:
URL: https://github.com/apache/tvm/pull/7313#discussion_r566526531
##########
File path: python/tvm/auto_scheduler/measure.py
##########
@@ -719,6 +722,87 @@ def local_builder_build(inputs, timeout, n_parallel, build_func="default", verbo
return results
+def _process_sparse_input(args):
+ sparse_prefix = sparse_data = sparse_indices = sparse_indptr = None
+
+ def _process_inputs(input_tensors, M, N, prefix_init):
+ nonlocal sparse_prefix
+ nonlocal sparse_data
+ nonlocal sparse_indices
+ nonlocal sparse_indptr
+
+ assert len(input_tensors) == 4
+ unsure_tensors = list(input_tensors)
+ # Get the Dense data
+ dense_data = None
+ for tensor in unsure_tensors:
+ if len(tensor.shape) == 2:
+ assert dense_data is None
+ dense_data = tensor
+ assert M == dense_data.shape[0]
+ K = dense_data.shape[1]
+ unsure_tensors.remove(dense_data)
+
+ # Get the Sparse data
+ sparse_data = None
+ for tensor in unsure_tensors:
+ if len(tensor.shape) == 3:
+ assert sparse_data is None
+ sparse_data = tensor
+ block_size, BS_R, BS_C = sparse_data.shape
+ unsure_tensors.remove(sparse_data)
+
+ # Get the Sparse indptr & indices
+ sparse_indices = None
+ for tensor in unsure_tensors:
+ assert len(tensor.shape) == 1
+ if tensor.shape[0] == block_size:
+ assert sparse_indices is None
+ sparse_indices = tensor
+ unsure_tensors.remove(sparse_indices)
+ assert len(unsure_tensors) == 1
+ sparse_indptr = unsure_tensors[0]
+
+ # Generate the sparse_prefix
+ density = 1.0
+ for i in sparse_data.shape:
+ density *= i
+ density /= (K * N)
+ density = density.value
+ sparse_prefix = "%s_%d_%d_%d_%d_%d_%.2f_" % (
Review comment:
Though in my test a schedule seems to have similar performance with different random sparse data, I think that may still be a potential problem. Unfortunately, I have not figured out any better solution.
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] jcf94 commented on pull request #7313: [AutoSchedule] Sparse dense tuning support with custom sketch rule
Posted by GitBox <gi...@apache.org>.
jcf94 commented on pull request #7313:
URL: https://github.com/apache/tvm/pull/7313#issuecomment-789492629
> According to our offline discussion,
>
> 1. Update the type of `SearchTaskNode::task_inputs`. Change it from `Map<String, runtime::NDArray>` to `Array<String>`, so we only need to store nd arrays in one place. We can query it from the global table in `measure.py`
> 2. Remove `SearchTask.AddTaskInput` interface to make `SearchTask` immutable. We do not have the need to dynamically update task inputs, so we can provide all arguments to the constructors.
> 3. Make sure we can use the same interface to support the use case where we want to match the special buffers by name
@comaniac @merrymercy Comments all addressed:
1/2: Removed the `add_task_input` API, and only provide them in constructor. Now, SearchTask only keeps the name of each special buffer.
3: Add a extra case in `measure.py:_prepare_input_map` to check the placeholder name, as well as a unit test in `test_auto_scheduler_measure.py:test_measure_special_inputs_map_by_name`.
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] jcf94 commented on a change in pull request #7313: [AutoSchedule] Sparse dense tuning support with custom sketch rule
Posted by GitBox <gi...@apache.org>.
jcf94 commented on a change in pull request #7313:
URL: https://github.com/apache/tvm/pull/7313#discussion_r569407977
##########
File path: python/tvm/auto_scheduler/measure.py
##########
@@ -1132,3 +1249,44 @@ def rpc_runner_run(
print("")
return results
+
+
+# The map stores special registered buffer for measurement
+# This can be used for sparse workloads when we cannot use random tensors for measurment.
+global special_buffer_table
+special_buffer_table = {}
Review comment:
Please have another look~ 😄
I've moved the input data to serialize in the SearchTask. (Even though we may still have to get a global table in the search_task.py)
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] jcf94 commented on a change in pull request #7313: [AutoSchedule] Sparse dense tuning support with custom sketch rule
Posted by GitBox <gi...@apache.org>.
jcf94 commented on a change in pull request #7313:
URL: https://github.com/apache/tvm/pull/7313#discussion_r570068598
##########
File path: python/tvm/auto_scheduler/measure.py
##########
@@ -943,18 +1047,36 @@ def _timed_rpc_run(
if error_no == 0:
try:
- args = [ndarray.empty(get_const_tuple(x.shape), x.dtype, ctx) for x in build_res.args]
try:
random_fill = remote.get_function("tvm.contrib.random.random_fill")
except AttributeError:
raise AttributeError(
"Please make sure USE_RANDOM is ON in the config.cmake " "on the remote devices"
)
- for arg in args:
- random_fill(arg)
- ctx.sync()
+ # Check sparse op
+ sparse_prefix, sparse_data, sparse_indices, sparse_indptr = \
+ _process_sparse_input(build_res.args)
+ if sparse_prefix:
Review comment:
Refactored this part.
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] jcf94 commented on a change in pull request #7313: [AutoSchedule] Sparse dense tuning support with custom sketch rule
Posted by GitBox <gi...@apache.org>.
jcf94 commented on a change in pull request #7313:
URL: https://github.com/apache/tvm/pull/7313#discussion_r570068598
##########
File path: python/tvm/auto_scheduler/measure.py
##########
@@ -943,18 +1047,36 @@ def _timed_rpc_run(
if error_no == 0:
try:
- args = [ndarray.empty(get_const_tuple(x.shape), x.dtype, ctx) for x in build_res.args]
try:
random_fill = remote.get_function("tvm.contrib.random.random_fill")
except AttributeError:
raise AttributeError(
"Please make sure USE_RANDOM is ON in the config.cmake " "on the remote devices"
)
- for arg in args:
- random_fill(arg)
- ctx.sync()
+ # Check sparse op
+ sparse_prefix, sparse_data, sparse_indices, sparse_indptr = \
+ _process_sparse_input(build_res.args)
+ if sparse_prefix:
Review comment:
Refactored this part.
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] comaniac commented on a change in pull request #7313: [AutoSchedule] Sparse dense tuning support with custom sketch rule
Posted by GitBox <gi...@apache.org>.
comaniac commented on a change in pull request #7313:
URL: https://github.com/apache/tvm/pull/7313#discussion_r586185313
##########
File path: tutorials/auto_scheduler/tune_sparse_x86.py
##########
@@ -0,0 +1,331 @@
+# 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.
+"""
+Auto-scheduling Sparse Matrix Multiplication for CPU by Custom Sketch Rule
+==========================================================================
+**Author**: `Chengfan Jia <https://github.com/jcf94/>`_
+
+This is a tutorial on how to use the auto-scheduler to tune a sparse matrix multiplication for
+CPUs.
+
+Auto-scheduler is designed to explore the schedule with best performance for a given computation
+declaration automatically. While sometimes, we may have a demand to try some special ops which may
+not been well supported by auto-scheduler's default search policy. Auto-scheduler currently allows
+user to provide a CustomSketch to cover these cases.
+
+We use sparse matrix multiplication as an example in this tutorial.
+
+Note that this tutorial will not run on Windows or recent versions of macOS. To
+get it to run, you will need to wrap the body of this tutorial in a :code:`if
+__name__ == "__main__":` block.
+"""
+
+import os
+import itertools
+
+import numpy as np
+import tvm
+from tvm import te, auto_scheduler, topi
+from tvm.auto_scheduler import _ffi_api
+from tvm.topi.utils import get_const_tuple
+
+import scipy.sparse as sp
+
+######################################################################
+# Define the computation
+# ^^^^^^^^^^^^^^^^^^^^^^
+# To begin with, let us define the computation of a sparse matmul with several relu and bias add.
+# The function should return the list of input/output tensors.
+# From these tensors, the auto-scheduler can get the whole computational graph.
+
+# We use this function to generate a random bsr matrix
+def random_bsr_matrix(M, N, BS_R, BS_C, density, dtype):
+ import itertools
+
+ Y = np.zeros((M, N), dtype=dtype)
+ assert M % BS_R == 0
+ assert N % BS_C == 0
+ nnz = int(density * M * N)
+ num_blocks = int(nnz / (BS_R * BS_C)) + 1
+ candidate_blocks = np.asarray(list(itertools.product(range(0, M, BS_R), range(0, N, BS_C))))
+ assert candidate_blocks.shape[0] == M // BS_R * N // BS_C
+ chosen_blocks = candidate_blocks[
+ np.random.choice(candidate_blocks.shape[0], size=num_blocks, replace=False)
+ ]
+ for i in range(len(chosen_blocks)):
+ r, c = chosen_blocks[i]
+ Y[r : r + BS_R, c : c + BS_C] = np.random.randn(BS_R, BS_C)
+ s = sp.bsr_matrix(Y, blocksize=(BS_R, BS_C))
+ assert s.data.shape == (num_blocks, BS_R, BS_C)
+ assert s.indices.shape == (num_blocks,)
+ assert s.indptr.shape == (M // BS_R + 1,)
+ return s
+
+@auto_scheduler.register_workload
+def sparse_dense(M, N, K, w_data_shape, w_indices_shape, w_indptr_shape, dtype):
+ X = te.placeholder(shape=(M, K), dtype=dtype)
+ W_data = te.placeholder(shape=w_data_shape, dtype=dtype)
+ W_indices = te.placeholder(shape=w_indices_shape, dtype="int32")
+ W_indptr = te.placeholder(shape=w_indptr_shape, dtype="int32")
+ B = te.placeholder(shape=(M, N), dtype=dtype)
+
+ out = topi.nn.sparse_dense(
+ topi.nn.relu(X), W_data, W_indices, W_indptr
+ )
+ out = te.compute((M, N), lambda i, j: out[i, j] + B[i, j], name="BiasAdd")
+ out = topi.nn.relu(out)
+
+ return [X, W_data, W_indices, W_indptr, B, out]
+
+######################################################################
+# Special step for sparse workload
+# ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+# During schedule tuning, auto-scheduler will use random inputs to measure the performance of a
+# generated schedule. While we cannot directly use a random array as the input of a sparse op, for
+# the "indices" and "indptr" array are meaningful for the computation.
+#
+# To solve this problem, we register these as special buffers, and load them when process program
+# measuring.
+# See the :any:`auto_scheduler.measure` code for more details.
+
+# Define the basic shapes of this sparse computation
+M = K = N = 512
+BS_R = 16
+BS_C = 1
+density = 0.6
+
+# Generate the test data with numpy
+X_np = np.random.randn(M, K).astype("float32")
+X_np = np.maximum(np.zeros((M, K), dtype="float32"), X_np) # Relu
+W_sp_np = random_bsr_matrix(N, K, BS_R, BS_C, density=density, dtype="float32")
+W_np = W_sp_np.todense()
+Y_np = X_np @ W_np.T # Process the matrix multiplication
+B_np = np.random.randn(M, N).astype("float32")
+Y_np = Y_np + B_np # Bias add
+Y_np = np.maximum(np.zeros((M, N), dtype="float32"), Y_np) # Relu
+
+# Register the sparse data to special buffer
+prefix = "sparse_dense_bsr_%d_%d_%d_%d_%d_%.2f_" % (M, N, K, BS_R, BS_C, density)
+auto_scheduler.measure.register_special_buffer(prefix + "W_data", W_sp_np.data)
+auto_scheduler.measure.register_special_buffer(prefix + "W_indices", W_sp_np.indices)
+auto_scheduler.measure.register_special_buffer(prefix + "W_indptr", W_sp_np.indptr)
+
+######################################################################
+# Create the search task
+# ^^^^^^^^^^^^^^^^^^^^^^
+# We then create a search task with M=N=K=512 and dtype="float32"
+# If your machine supports avx instructions, you can
+#
+# - replace "llvm" below with "llvm -mcpu=core-avx2" to enable AVX2
+# - replace "llvm" below with "llvm -mcpu=skylake-avx512" to enable AVX-512
+
+target = tvm.target.Target("llvm")
+
+task = tvm.auto_scheduler.SearchTask(
+ func=sparse_dense,
+ args=(
+ M, N, K,
+ W_sp_np.data.shape,
+ W_sp_np.indices.shape,
+ W_sp_np.indptr.shape,
+ "float32"
+ ),
+ target=target
+)
+
+# Inspect the computational graph
+print("Computational DAG:")
+print(task.compute_dag)
+
+######################################################################
+# Write the custom sketch for sparse dense op
+# ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+# Before tuning, we will need to define the CustomSketchRule for the sparse dense op.
+#
+# CustomSketchRule consists of two parts: the condition function and the apply function.
+#
+# - condition function: describe when to use this sketch rule. For example, we can match the op
+# by their name or tag.
+# - apply function: describe how to generate the initial sketch. Auto-scheduler provides a set of
+# loop state APIs.
+
+def meet_condition_func(search_policy, state, stage_id):
+ state = auto_scheduler.loop_state.State(state, search_policy.search_task.compute_dag)
+ if state.stages[stage_id].op.tag in [
+ "sparse_dense_sp_rhs_bsrmm", "sparse_dense_sp_rhs_bsrmm_block"
+ ]:
+ return auto_scheduler.PreloadCustomSketchRule.APPLY_AND_SKIP_REST
+ else:
+ return auto_scheduler.PreloadCustomSketchRule.PASS
+
+def apply_func(search_policy, state, stage_id):
+ ret = []
+ s0 = auto_scheduler.loop_state.State(state, search_policy.search_task.compute_dag)
+ if s0.stages[stage_id].op.tag == "sparse_dense_sp_rhs_bsrmm_block":
+ return [s0.state_object, stage_id - 1]
+
+ sparse_dense = s0.stages[stage_id].op
+ sparse_dense_block = s0.stages[stage_id - 1].op
+ assert sparse_dense.tag == "sparse_dense_sp_rhs_bsrmm"
+ assert sparse_dense_block.tag == "sparse_dense_sp_rhs_bsrmm_block"
+
+ # Set the default consumer of compute block
+ consumer = sparse_dense
+
+ # If sparse dense has a single elementwise consumer
+ # We can compute inline the sparse_dense output stage
+ consumers = _ffi_api.SearchPolicyUtilsGetConsumers(
+ search_policy.search_task, s0.state_object, stage_id
+ )
+ if len(consumers) == 1:
+ consumer_id = int(consumers.items()[0][0])
+ if _ffi_api.SearchPolicyUtilsIsElementwiseMatch(
+ search_policy.search_task, s0.state_object, stage_id, consumer_id
+ ):
+ consumer = s0.stages[consumer_id].op
+ s0.compute_inline(sparse_dense)
+
+ i, nb_j, j, row_offset, c = s0[sparse_dense_block].iters
+ m, n = s0[consumer].iters
+ i0, i1, i2 = s0.split(sparse_dense_block, i, [None, None])
+ m0, m1 = s0.follow_split(consumer, m, len(s0.transform_steps) - 1, 1)
+ j0, j1 = s0.split(sparse_dense_block, nb_j, [None])
+ n0, n1 = s0.follow_split(consumer, n, len(s0.transform_steps) - 1, 1)
+ s0.reorder(sparse_dense_block, [i0, j0, i1, j1, row_offset, i2, j, c])
+ s0.reorder(consumer, [m0, n0, m1, n1])
+ s0.compute_at(sparse_dense_block, consumer, n0)
+
+ ret.append([s0.state_object, stage_id - 2])
+
+ return ret
+
+######################################################################
+# Next, we set parameters for the auto-scheduler.
+#
+# * :code:`num_measure_trials` is the number of measurement trials we can use during the search.
+# We only make 10 trials in this tutorial for a fast demonstration. In practice, 1000 is a
+# good value for the search to converge. You can do more trials according to your time budget.
+# * In addition, we use :code:`RecordToFile` to dump measurement records into a file `matmul.json`.
+# The measurement records can be used to query the history best, resume the search,
+# and do more analyses later.
+# * see :any:`auto_scheduler.TuningOptions` for more parameters
+# * Here, we need to create a :code:`auto_scheduler.SketchPolicy` object, and add the custom sketch
+# rule as a `init_search_callbacks`.
+
+log_file = "sparse_dense.json"
+tune_option = auto_scheduler.TuningOptions(
+ num_measure_trials=10,
+ measure_callbacks=[auto_scheduler.RecordToFile(log_file)],
+ verbose=2,
+)
+
+search_policy = auto_scheduler.SketchPolicy(
+ task,
+ program_cost_model=auto_scheduler.XGBModel(),
+ init_search_callbacks=[
+ auto_scheduler.PreloadCustomSketchRule(meet_condition_func, apply_func, "SparseDense")
+ ]
+)
+
+######################################################################
+# Run the search
+# ^^^^^^^^^^^^^^
+# Now we get all inputs ready.
+# We can kick off the search and let the auto-scheduler do its magic.
+# After some measurement trials, we can load the best schedule from the log
+# file and apply it.
+
+# Run auto-tuning (search)
+task.tune(tune_option, search_policy)
Review comment:
Well...the failed case in #7548 also runs two trials. We should avoid potential CI flaky as possible.
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] comaniac commented on a change in pull request #7313: [AutoSchedule] Sparse dense tuning support with custom sketch rule
Posted by GitBox <gi...@apache.org>.
comaniac commented on a change in pull request #7313:
URL: https://github.com/apache/tvm/pull/7313#discussion_r587680749
##########
File path: python/tvm/auto_scheduler/measure.py
##########
@@ -719,6 +720,45 @@ def local_builder_build(inputs, timeout, n_parallel, build_func="default", verbo
return results
+def _prepare_input_map(args):
+ """This function deals with special task inputs.
Review comment:
Better to say more about what is special inputs.
##########
File path: python/tvm/auto_scheduler/measure.py
##########
@@ -719,6 +720,45 @@ def local_builder_build(inputs, timeout, n_parallel, build_func="default", verbo
return results
+def _prepare_input_map(args):
+ """This function deals with special task inputs.
+
+ Parameters
+ ----------
+ args : List[Tensor]
+ Input/output Tensor of a TVM subgraph.
+
+ Returns
+ -------
+ A Dict[Tensor, str] that maps the input Tensor to a buffer name.
+
+ Note
+ ----
Review comment:
```suggestion
Notes
-----
```
##########
File path: include/tvm/auto_scheduler/search_task.h
##########
@@ -120,6 +121,8 @@ class SearchTaskNode : public Object {
HardwareParams hardware_params;
/*! \brief The layout rewrite option used for measuring programs. */
LayoutRewriteOption layout_rewrite_option;
+ /*! \brief Names of some user defined input data used in program measuring. */
+ Array<String> task_inputs;
Review comment:
Per functionality, `task_input_names` would be better.
##########
File path: python/tvm/auto_scheduler/search_task.py
##########
@@ -157,6 +164,149 @@ def __init__(
)
+# The map stores special registered buffer for measurement
+# This can be used for sparse workloads when we cannot use random tensors for measurment.
+# {
+# "workload_key_0": {
+# "task_input_0": Tensor(...),
+# "task_input_1": Tensor(...)
+# },
+# "workload_key_1": {
+# "task_input_2": Tensor(...),
+# "task_input_3": Tensor(...)
+# },
+# ...
+# }
+TASK_INPUT_BUFFER_TABLE = {}
+
+
+def _save_buffer_to_file(buffer_name, buffer_data):
+ """Save the current Tensor buffer to a numpy file.
+
+ File name will be: {buffer_name}.{buffer_shape}_{buffer_data_type}
+ """
+ np_data = buffer_data.asnumpy()
+
+ buffer_name += "."
+ for i in np_data.shape:
+ buffer_name += "%d_" % (i)
+ buffer_name += "%s" % (np_data.dtype)
+
+ np_data.tofile(buffer_name, " ")
+
+
+def _try_load_buffer_from_file(buffer_name):
+ """Try to load buffer from a numpy file, if not found, return None.
+
+ File name has a same format as `_save_buffer_to_file`.
+ """
+ filelist = os.listdir()
+
+ for file in filelist:
+ if file.startswith(buffer_name) and file.count("."):
+ meta_info = file.split(".")[-1].split("_")
+ shape = [int(i) for i in meta_info[:-1]]
+ dtype = meta_info[-1]
+ buffer_data = np.fromfile(file, dtype=dtype, sep=" ")
+ buffer_data = buffer_data.reshape(shape)
+ return ndarray.array(buffer_data)
+
+ return None
+
+
+def register_task_input_buffer(
+ workload_key,
+ input_name,
+ input_data,
+ overwrite=False,
+ save_to_file=False,
+):
+ """Register special buffer for measurement.
+
+ Parameters
+ ----------
+ workload_key : str
+ The workload key of the SearchTask.
+
+ input_name : str
+ The name of input buffer.
+
+ input_data : tvm.nd.NDArray
+ The input Tensor data.
+
+ overwrite : bool = False
+ Whether overwrite the data if a name has already in the global table.
+
+ save_to_file : bool = False
+ Whether record this buffer to a local file. This can be reused to continue the last tuning
+ process.
+ """
Review comment:
Returns?
##########
File path: python/tvm/auto_scheduler/search_task.py
##########
@@ -157,6 +164,149 @@ def __init__(
)
+# The map stores special registered buffer for measurement
+# This can be used for sparse workloads when we cannot use random tensors for measurment.
+# {
+# "workload_key_0": {
+# "task_input_0": Tensor(...),
+# "task_input_1": Tensor(...)
+# },
+# "workload_key_1": {
+# "task_input_2": Tensor(...),
+# "task_input_3": Tensor(...)
+# },
+# ...
+# }
+TASK_INPUT_BUFFER_TABLE = {}
+
+
+def _save_buffer_to_file(buffer_name, buffer_data):
+ """Save the current Tensor buffer to a numpy file.
+
+ File name will be: {buffer_name}.{buffer_shape}_{buffer_data_type}
+ """
+ np_data = buffer_data.asnumpy()
+
+ buffer_name += "."
+ for i in np_data.shape:
+ buffer_name += "%d_" % (i)
+ buffer_name += "%s" % (np_data.dtype)
+
+ np_data.tofile(buffer_name, " ")
+
+
+def _try_load_buffer_from_file(buffer_name):
+ """Try to load buffer from a numpy file, if not found, return None.
+
+ File name has a same format as `_save_buffer_to_file`.
+ """
+ filelist = os.listdir()
+
+ for file in filelist:
+ if file.startswith(buffer_name) and file.count("."):
+ meta_info = file.split(".")[-1].split("_")
+ shape = [int(i) for i in meta_info[:-1]]
+ dtype = meta_info[-1]
+ buffer_data = np.fromfile(file, dtype=dtype, sep=" ")
+ buffer_data = buffer_data.reshape(shape)
+ return ndarray.array(buffer_data)
+
+ return None
+
+
+def register_task_input_buffer(
+ workload_key,
+ input_name,
+ input_data,
+ overwrite=False,
+ save_to_file=False,
+):
+ """Register special buffer for measurement.
+
+ Parameters
+ ----------
+ workload_key : str
+ The workload key of the SearchTask.
+
+ input_name : str
+ The name of input buffer.
+
+ input_data : tvm.nd.NDArray
+ The input Tensor data.
+
+ overwrite : bool = False
+ Whether overwrite the data if a name has already in the global table.
+
+ save_to_file : bool = False
+ Whether record this buffer to a local file. This can be reused to continue the last tuning
+ process.
+ """
+ global TASK_INPUT_BUFFER_TABLE
+
+ if workload_key not in TASK_INPUT_BUFFER_TABLE:
+ TASK_INPUT_BUFFER_TABLE[workload_key] = {}
+ input_table = TASK_INPUT_BUFFER_TABLE[workload_key]
+
+ if not overwrite:
+ if input_name not in input_table.keys():
+ # Try to load buffer data from local file
+ tensor_from_file = _try_load_buffer_from_file(input_name)
+ if tensor_from_file:
+ input_table[input_name] = tensor_from_file
+
+ if input_name in input_table.keys():
+ logger.warning(
+ "Tensor %s exists in TASK_INPUT_BUFFER_TABLE, %s",
+ input_name,
+ "set overwrite to True or this Tensor will not be registered",
+ )
+ return input_table[input_name]
+
+ input_table[input_name] = input_data
+ if save_to_file:
+ _save_buffer_to_file(input_name, input_data)
+ return input_data
+
+
+@tvm._ffi.register_func("auto_scheduler.search_task.get_task_input_buffer")
+def get_task_input_buffer(workload_key, input_name):
+ """Get special buffer for measurement.
+
+ The buffers are registered by `register_task_input_buffer`.
+
+ Parameters
+ ----------
+ workload_key : str
+ The workload key of the SearchTask.
+
+ input_name : str
+ The name of input buffer.
+
+ Returns
+ -------
+ The registered input buffer.
Review comment:
type?
##########
File path: python/tvm/auto_scheduler/search_task.py
##########
@@ -185,6 +335,16 @@ class SearchTask(Object):
The NO_REWRITE and INSERT_TRANSFORM_STAGE are expected to be used when tuning a standalone
op, and the REWRITE_FOR_PRE_TRANSFORMED is expected to be used when tuning ops inside a
network.
+ task_inputs : Union[Dict[str, tvm.nd.NDArray], List[str]]
+ A dict maps the input names to input tensors or a list of input names.
+ Some special Tensor used as inputs in program measuring. Usually we do not need to care
+ about it, but for special workloads like Sparse computation the Sparse Tensor input are
+ meaningful that we cannot use random input directly.
+ task_inputs_overwrite : bool = False
+ Whether overwrite the data if a name has already in the global table.
+ task_inputs_save_to_file : bool = False
+ Whether record this buffer to a local file. This can be reused to continue the last
+ tuning process.
Review comment:
ditto
##########
File path: python/tvm/auto_scheduler/search_task.py
##########
@@ -157,6 +164,149 @@ def __init__(
)
+# The map stores special registered buffer for measurement
+# This can be used for sparse workloads when we cannot use random tensors for measurment.
Review comment:
```suggestion
# The map stores special registered buffer for measurement.
# This can be used for sparse workloads when we cannot use random tensors for measurment.
```
##########
File path: python/tvm/auto_scheduler/search_task.py
##########
@@ -157,6 +164,149 @@ def __init__(
)
+# The map stores special registered buffer for measurement
+# This can be used for sparse workloads when we cannot use random tensors for measurment.
+# {
+# "workload_key_0": {
+# "task_input_0": Tensor(...),
+# "task_input_1": Tensor(...)
+# },
+# "workload_key_1": {
+# "task_input_2": Tensor(...),
+# "task_input_3": Tensor(...)
+# },
+# ...
+# }
+TASK_INPUT_BUFFER_TABLE = {}
+
+
+def _save_buffer_to_file(buffer_name, buffer_data):
+ """Save the current Tensor buffer to a numpy file.
+
+ File name will be: {buffer_name}.{buffer_shape}_{buffer_data_type}
Review comment:
Would that be better to have the file extension (i.e., `npy`), so it becomes `{buffer_name}_{buffer_shape}_{buffer_data_type}.npy`.
##########
File path: python/tvm/auto_scheduler/search_task.py
##########
@@ -157,6 +164,149 @@ def __init__(
)
+# The map stores special registered buffer for measurement
+# This can be used for sparse workloads when we cannot use random tensors for measurment.
+# {
+# "workload_key_0": {
+# "task_input_0": Tensor(...),
+# "task_input_1": Tensor(...)
+# },
+# "workload_key_1": {
+# "task_input_2": Tensor(...),
+# "task_input_3": Tensor(...)
+# },
+# ...
+# }
+TASK_INPUT_BUFFER_TABLE = {}
+
+
+def _save_buffer_to_file(buffer_name, buffer_data):
+ """Save the current Tensor buffer to a numpy file.
+
+ File name will be: {buffer_name}.{buffer_shape}_{buffer_data_type}
+ """
+ np_data = buffer_data.asnumpy()
+
+ buffer_name += "."
+ for i in np_data.shape:
+ buffer_name += "%d_" % (i)
+ buffer_name += "%s" % (np_data.dtype)
+
+ np_data.tofile(buffer_name, " ")
+
+
+def _try_load_buffer_from_file(buffer_name):
+ """Try to load buffer from a numpy file, if not found, return None.
+
+ File name has a same format as `_save_buffer_to_file`.
+ """
+ filelist = os.listdir()
+
+ for file in filelist:
+ if file.startswith(buffer_name) and file.count("."):
+ meta_info = file.split(".")[-1].split("_")
+ shape = [int(i) for i in meta_info[:-1]]
+ dtype = meta_info[-1]
+ buffer_data = np.fromfile(file, dtype=dtype, sep=" ")
+ buffer_data = buffer_data.reshape(shape)
+ return ndarray.array(buffer_data)
+
+ return None
+
+
+def register_task_input_buffer(
+ workload_key,
+ input_name,
+ input_data,
+ overwrite=False,
+ save_to_file=False,
+):
+ """Register special buffer for measurement.
+
+ Parameters
+ ----------
+ workload_key : str
+ The workload key of the SearchTask.
+
+ input_name : str
+ The name of input buffer.
+
+ input_data : tvm.nd.NDArray
+ The input Tensor data.
+
+ overwrite : bool = False
+ Whether overwrite the data if a name has already in the global table.
+
+ save_to_file : bool = False
+ Whether record this buffer to a local file. This can be reused to continue the last tuning
+ process.
Review comment:
```suggestion
Whether to save the data to a local file as well. This can be reused to resume the last tuning
process.
```
##########
File path: python/tvm/auto_scheduler/search_task.py
##########
@@ -157,6 +164,149 @@ def __init__(
)
+# The map stores special registered buffer for measurement
+# This can be used for sparse workloads when we cannot use random tensors for measurment.
+# {
+# "workload_key_0": {
+# "task_input_0": Tensor(...),
+# "task_input_1": Tensor(...)
+# },
+# "workload_key_1": {
+# "task_input_2": Tensor(...),
+# "task_input_3": Tensor(...)
+# },
+# ...
+# }
+TASK_INPUT_BUFFER_TABLE = {}
+
+
+def _save_buffer_to_file(buffer_name, buffer_data):
+ """Save the current Tensor buffer to a numpy file.
+
+ File name will be: {buffer_name}.{buffer_shape}_{buffer_data_type}
+ """
+ np_data = buffer_data.asnumpy()
+
+ buffer_name += "."
+ for i in np_data.shape:
+ buffer_name += "%d_" % (i)
+ buffer_name += "%s" % (np_data.dtype)
+
+ np_data.tofile(buffer_name, " ")
+
+
+def _try_load_buffer_from_file(buffer_name):
+ """Try to load buffer from a numpy file, if not found, return None.
+
+ File name has a same format as `_save_buffer_to_file`.
+ """
+ filelist = os.listdir()
+
+ for file in filelist:
+ if file.startswith(buffer_name) and file.count("."):
+ meta_info = file.split(".")[-1].split("_")
+ shape = [int(i) for i in meta_info[:-1]]
+ dtype = meta_info[-1]
+ buffer_data = np.fromfile(file, dtype=dtype, sep=" ")
+ buffer_data = buffer_data.reshape(shape)
+ return ndarray.array(buffer_data)
+
+ return None
+
+
+def register_task_input_buffer(
+ workload_key,
+ input_name,
+ input_data,
+ overwrite=False,
+ save_to_file=False,
+):
+ """Register special buffer for measurement.
+
+ Parameters
+ ----------
+ workload_key : str
+ The workload key of the SearchTask.
+
+ input_name : str
+ The name of input buffer.
+
+ input_data : tvm.nd.NDArray
+ The input Tensor data.
+
+ overwrite : bool = False
+ Whether overwrite the data if a name has already in the global table.
+
+ save_to_file : bool = False
+ Whether record this buffer to a local file. This can be reused to continue the last tuning
+ process.
+ """
+ global TASK_INPUT_BUFFER_TABLE
+
+ if workload_key not in TASK_INPUT_BUFFER_TABLE:
+ TASK_INPUT_BUFFER_TABLE[workload_key] = {}
+ input_table = TASK_INPUT_BUFFER_TABLE[workload_key]
+
+ if not overwrite:
+ if input_name not in input_table.keys():
+ # Try to load buffer data from local file
+ tensor_from_file = _try_load_buffer_from_file(input_name)
+ if tensor_from_file:
+ input_table[input_name] = tensor_from_file
+
+ if input_name in input_table.keys():
Review comment:
```suggestion
else:
```
btw should we have a message saying what buffer is loaded? otherwise if users suppose a buffer has been loaded but actually not (like the user accidently removed the file), then the tuning results may be useless.
##########
File path: python/tvm/auto_scheduler/search_task.py
##########
@@ -185,6 +335,16 @@ class SearchTask(Object):
The NO_REWRITE and INSERT_TRANSFORM_STAGE are expected to be used when tuning a standalone
op, and the REWRITE_FOR_PRE_TRANSFORMED is expected to be used when tuning ops inside a
network.
+ task_inputs : Union[Dict[str, tvm.nd.NDArray], List[str]]
+ A dict maps the input names to input tensors or a list of input names.
+ Some special Tensor used as inputs in program measuring. Usually we do not need to care
+ about it, but for special workloads like Sparse computation the Sparse Tensor input are
+ meaningful that we cannot use random input directly.
+ task_inputs_overwrite : bool = False
+ Whether overwrite the data if a name has already in the global table.
Review comment:
```suggestion
Whether to overwrite the data if a name has already in the global table.
```
##########
File path: python/tvm/auto_scheduler/search_task.py
##########
@@ -157,6 +164,149 @@ def __init__(
)
+# The map stores special registered buffer for measurement
+# This can be used for sparse workloads when we cannot use random tensors for measurment.
+# {
+# "workload_key_0": {
+# "task_input_0": Tensor(...),
+# "task_input_1": Tensor(...)
+# },
+# "workload_key_1": {
+# "task_input_2": Tensor(...),
+# "task_input_3": Tensor(...)
+# },
+# ...
+# }
+TASK_INPUT_BUFFER_TABLE = {}
+
+
+def _save_buffer_to_file(buffer_name, buffer_data):
+ """Save the current Tensor buffer to a numpy file.
+
+ File name will be: {buffer_name}.{buffer_shape}_{buffer_data_type}
+ """
+ np_data = buffer_data.asnumpy()
+
+ buffer_name += "."
+ for i in np_data.shape:
+ buffer_name += "%d_" % (i)
+ buffer_name += "%s" % (np_data.dtype)
+
+ np_data.tofile(buffer_name, " ")
+
+
+def _try_load_buffer_from_file(buffer_name):
+ """Try to load buffer from a numpy file, if not found, return None.
+
+ File name has a same format as `_save_buffer_to_file`.
+ """
+ filelist = os.listdir()
+
+ for file in filelist:
+ if file.startswith(buffer_name) and file.count("."):
+ meta_info = file.split(".")[-1].split("_")
+ shape = [int(i) for i in meta_info[:-1]]
+ dtype = meta_info[-1]
+ buffer_data = np.fromfile(file, dtype=dtype, sep=" ")
+ buffer_data = buffer_data.reshape(shape)
+ return ndarray.array(buffer_data)
+
+ return None
+
+
+def register_task_input_buffer(
+ workload_key,
+ input_name,
+ input_data,
+ overwrite=False,
+ save_to_file=False,
+):
+ """Register special buffer for measurement.
+
+ Parameters
+ ----------
+ workload_key : str
+ The workload key of the SearchTask.
+
+ input_name : str
+ The name of input buffer.
+
+ input_data : tvm.nd.NDArray
+ The input Tensor data.
+
+ overwrite : bool = False
+ Whether overwrite the data if a name has already in the global table.
Review comment:
```suggestion
Whether to overwrite the data if a name has already registered.
```
##########
File path: python/tvm/auto_scheduler/measure.py
##########
@@ -758,11 +802,25 @@ def _timed_eval_func(
if error_no == 0:
try:
- args = [ndarray.empty(get_const_tuple(x.shape), x.dtype, ctx) for x in build_res.args]
random_fill = tvm.get_global_func("tvm.contrib.random.random_fill", True)
assert random_fill, "Please make sure USE_RANDOM is ON in the config.cmake"
- for arg in args:
- random_fill(arg)
+
+ tensor_input_map = _prepare_input_map(build_res.args) if task_inputs else {}
+ args = []
+ for arg in build_res.args:
+ if arg in tensor_input_map:
+ tensor_name = tensor_input_map[arg]
+ if tensor_name in task_inputs:
+ args.append(get_task_input_buffer(inp.task.workload_key, tensor_name))
+ else:
+ raise ValueError(
+ "%s not found in task_inputs, " % (tensor_name)
+ + "should provide with SearchTask.AddTaskInput()"
Review comment:
Didn't find `AddTaskInput`?
##########
File path: python/tvm/auto_scheduler/measure.py
##########
@@ -943,18 +1005,30 @@ def _timed_rpc_run(
if error_no == 0:
try:
- args = [ndarray.empty(get_const_tuple(x.shape), x.dtype, ctx) for x in build_res.args]
- try:
- random_fill = remote.get_function("tvm.contrib.random.random_fill")
- except AttributeError:
- raise AttributeError(
- "Please make sure USE_RANDOM is ON in the config.cmake " "on the remote devices"
- )
- for arg in args:
- random_fill(arg)
+ random_fill = remote.get_function("tvm.contrib.random.random_fill")
+ assert (
+ random_fill
+ ), "Please make sure USE_RANDOM is ON in the config.cmake on the remote devices"
+
+ tensor_input_map = _prepare_input_map(build_res.args) if task_inputs else {}
+ args = []
+ for arg in build_res.args:
+ if arg in tensor_input_map:
+ tensor_name = tensor_input_map[arg]
+ if tensor_name in task_inputs:
+ args.append(get_task_input_buffer(inp.task.workload_key, tensor_name))
+ else:
+ raise ValueError(
+ "%s not found in task_inputs, " % (tensor_name)
+ + "should provide with SearchTask.AddTaskInput()"
Review comment:
ditto
##########
File path: python/tvm/auto_scheduler/measure.py
##########
@@ -719,6 +720,45 @@ def local_builder_build(inputs, timeout, n_parallel, build_func="default", verbo
return results
+def _prepare_input_map(args):
+ """This function deals with special task inputs.
+
+ Parameters
+ ----------
+ args : List[Tensor]
+ Input/output Tensor of a TVM subgraph.
+
+ Returns
+ -------
+ A Dict[Tensor, str] that maps the input Tensor to a buffer name.
Review comment:
```suggestion
Dict[Tensor, str] :
Map from the input Tensor to its buffer name.
```
##########
File path: python/tvm/topi/nn/sparse.py
##########
@@ -356,3 +359,110 @@ def sparse_dense_alter_layout(_attrs, _inputs, _tinfos, _out_type):
Unlike other TOPI functions, this function operates on both graph level and operator level.
"""
return None
+
+
+def try_get_sparse_input(args):
+ """Analyze the input data from the given args.
+
+ Parameters
+ ----------
+ args : List[Tensor]
+ Input/output Tensor of a TVM subgraph.
+
+ Returns
+ -------
+ A Dict[Tensor, str] that maps the input Tensor to a buffer name.
+
+ Note
+ ----
+ The buffer name is specially designed, and these buffer should be provided in
+ `SearchTask(..., task_inputs={...})`.
Review comment:
ditto
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] jcf94 commented on a change in pull request #7313: [AutoSchedule] Sparse dense tuning support with custom sketch rule
Posted by GitBox <gi...@apache.org>.
jcf94 commented on a change in pull request #7313:
URL: https://github.com/apache/tvm/pull/7313#discussion_r587962953
##########
File path: python/tvm/auto_scheduler/search_task.py
##########
@@ -157,6 +164,149 @@ def __init__(
)
+# The map stores special registered buffer for measurement
+# This can be used for sparse workloads when we cannot use random tensors for measurment.
+# {
+# "workload_key_0": {
+# "task_input_0": Tensor(...),
+# "task_input_1": Tensor(...)
+# },
+# "workload_key_1": {
+# "task_input_2": Tensor(...),
+# "task_input_3": Tensor(...)
+# },
+# ...
+# }
+TASK_INPUT_BUFFER_TABLE = {}
+
+
+def _save_buffer_to_file(buffer_name, buffer_data):
+ """Save the current Tensor buffer to a numpy file.
+
+ File name will be: {buffer_name}.{buffer_shape}_{buffer_data_type}
Review comment:
Great. But we need a "." to seperate the buffer_name out, maybe update to:
`{buffer_name}.{buffer_shape}_{buffer_data_type}.npy` ?
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] comaniac commented on a change in pull request #7313: [AutoSchedule] Sparse dense tuning support with custom sketch rule
Posted by GitBox <gi...@apache.org>.
comaniac commented on a change in pull request #7313:
URL: https://github.com/apache/tvm/pull/7313#discussion_r561459545
##########
File path: tutorials/auto_scheduler/tune_sparse_x86.py
##########
@@ -0,0 +1,331 @@
+# 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.
+"""
+Auto-scheduling Sparse Matrix Multiplication for CPU by Custom Sketch Rule
+==========================================================================
Review comment:
```suggestion
Auto-scheduling Sparse Matrix Multiplication on CPU with Custom Sketch Rule
===========================================================================
```
##########
File path: tutorials/auto_scheduler/tune_sparse_x86.py
##########
@@ -0,0 +1,331 @@
+# 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.
+"""
+Auto-scheduling Sparse Matrix Multiplication for CPU by Custom Sketch Rule
+==========================================================================
+**Author**: `Chengfan Jia <https://github.com/jcf94/>`_
+
+This is a tutorial on how to use the auto-scheduler to tune a sparse matrix multiplication for
+CPUs.
+
+Auto-scheduler is designed to explore the schedule with best performance for a given computation
+declaration automatically. While sometimes, we may have a demand to try some special ops which may
+not been well supported by auto-scheduler's default search policy. Auto-scheduler currently allows
+user to provide a CustomSketch to cover these cases.
+
+We use sparse matrix multiplication as an example in this tutorial.
+
+Note that this tutorial will not run on Windows or recent versions of macOS. To
+get it to run, you will need to wrap the body of this tutorial in a :code:`if
+__name__ == "__main__":` block.
+"""
+
+import os
+import itertools
+
+import numpy as np
+import tvm
+from tvm import te, auto_scheduler, topi
+from tvm.auto_scheduler import _ffi_api
+from tvm.topi.utils import get_const_tuple
+
+import scipy.sparse as sp
+
+######################################################################
+# Define the computation
+# ^^^^^^^^^^^^^^^^^^^^^^
+# To begin with, let us define the computation of a sparse matmul with several relu and bias add.
+# The function should return the list of input/output tensors.
+# From these tensors, the auto-scheduler can get the whole computational graph.
+
+# We use this function to generate a random bsr matrix
+def random_bsr_matrix(M, N, BS_R, BS_C, density, dtype):
+ import itertools
+
+ Y = np.zeros((M, N), dtype=dtype)
+ assert M % BS_R == 0
+ assert N % BS_C == 0
+ nnz = int(density * M * N)
+ num_blocks = int(nnz / (BS_R * BS_C)) + 1
+ candidate_blocks = np.asarray(list(itertools.product(range(0, M, BS_R), range(0, N, BS_C))))
+ assert candidate_blocks.shape[0] == M // BS_R * N // BS_C
+ chosen_blocks = candidate_blocks[
+ np.random.choice(candidate_blocks.shape[0], size=num_blocks, replace=False)
+ ]
+ for i in range(len(chosen_blocks)):
+ r, c = chosen_blocks[i]
+ Y[r : r + BS_R, c : c + BS_C] = np.random.randn(BS_R, BS_C)
+ s = sp.bsr_matrix(Y, blocksize=(BS_R, BS_C))
+ assert s.data.shape == (num_blocks, BS_R, BS_C)
+ assert s.indices.shape == (num_blocks,)
+ assert s.indptr.shape == (M // BS_R + 1,)
+ return s
+
+@auto_scheduler.register_workload
+def sparse_dense(M, N, K, w_data_shape, w_indices_shape, w_indptr_shape, dtype):
+ X = te.placeholder(shape=(M, K), dtype=dtype)
+ W_data = te.placeholder(shape=w_data_shape, dtype=dtype)
+ W_indices = te.placeholder(shape=w_indices_shape, dtype="int32")
+ W_indptr = te.placeholder(shape=w_indptr_shape, dtype="int32")
+ B = te.placeholder(shape=(M, N), dtype=dtype)
+
+ out = topi.nn.sparse_dense(
+ topi.nn.relu(X), W_data, W_indices, W_indptr
+ )
+ out = te.compute((M, N), lambda i, j: out[i, j] + B[i, j], name="BiasAdd")
+ out = topi.nn.relu(out)
+
+ return [X, W_data, W_indices, W_indptr, B, out]
+
+######################################################################
+# Special step for sparse workload
+# ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+# During schedule tuning, auto-scheduler will use random inputs to measure the performance of a
+# generated schedule. While we cannot directly use a random array as the input of a sparse op, for
+# the "indices" and "indptr" array are meaningful for the computation.
+#
+# To solve this problem, we register these as special buffers, and load them when process program
+# measuring.
+# See the :any:`auto_scheduler.measure` code for more details.
Review comment:
I feel this is too ad hoc. Can we just expose the input buffers in general? For example, Relay graph runtime uses `set_input` to accept data, maybe we can have a similar API in `task` instead of `measure`? This is more reasonable because `measure_ctx` can actually be used by all tasks.
##########
File path: tutorials/auto_scheduler/tune_sparse_x86.py
##########
@@ -0,0 +1,331 @@
+# 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.
+"""
+Auto-scheduling Sparse Matrix Multiplication for CPU by Custom Sketch Rule
+==========================================================================
+**Author**: `Chengfan Jia <https://github.com/jcf94/>`_
+
+This is a tutorial on how to use the auto-scheduler to tune a sparse matrix multiplication for
+CPUs.
+
+Auto-scheduler is designed to explore the schedule with best performance for a given computation
+declaration automatically. While sometimes, we may have a demand to try some special ops which may
+not been well supported by auto-scheduler's default search policy. Auto-scheduler currently allows
+user to provide a CustomSketch to cover these cases.
+
+We use sparse matrix multiplication as an example in this tutorial.
+
+Note that this tutorial will not run on Windows or recent versions of macOS. To
+get it to run, you will need to wrap the body of this tutorial in a :code:`if
+__name__ == "__main__":` block.
+"""
+
+import os
+import itertools
+
+import numpy as np
+import tvm
+from tvm import te, auto_scheduler, topi
+from tvm.auto_scheduler import _ffi_api
+from tvm.topi.utils import get_const_tuple
+
+import scipy.sparse as sp
+
+######################################################################
+# Define the computation
+# ^^^^^^^^^^^^^^^^^^^^^^
+# To begin with, let us define the computation of a sparse matmul with several relu and bias add.
+# The function should return the list of input/output tensors.
+# From these tensors, the auto-scheduler can get the whole computational graph.
+
+# We use this function to generate a random bsr matrix
+def random_bsr_matrix(M, N, BS_R, BS_C, density, dtype):
+ import itertools
+
+ Y = np.zeros((M, N), dtype=dtype)
+ assert M % BS_R == 0
+ assert N % BS_C == 0
+ nnz = int(density * M * N)
+ num_blocks = int(nnz / (BS_R * BS_C)) + 1
+ candidate_blocks = np.asarray(list(itertools.product(range(0, M, BS_R), range(0, N, BS_C))))
+ assert candidate_blocks.shape[0] == M // BS_R * N // BS_C
+ chosen_blocks = candidate_blocks[
+ np.random.choice(candidate_blocks.shape[0], size=num_blocks, replace=False)
+ ]
+ for i in range(len(chosen_blocks)):
+ r, c = chosen_blocks[i]
+ Y[r : r + BS_R, c : c + BS_C] = np.random.randn(BS_R, BS_C)
+ s = sp.bsr_matrix(Y, blocksize=(BS_R, BS_C))
+ assert s.data.shape == (num_blocks, BS_R, BS_C)
+ assert s.indices.shape == (num_blocks,)
+ assert s.indptr.shape == (M // BS_R + 1,)
+ return s
+
+@auto_scheduler.register_workload
+def sparse_dense(M, N, K, w_data_shape, w_indices_shape, w_indptr_shape, dtype):
+ X = te.placeholder(shape=(M, K), dtype=dtype)
+ W_data = te.placeholder(shape=w_data_shape, dtype=dtype)
+ W_indices = te.placeholder(shape=w_indices_shape, dtype="int32")
+ W_indptr = te.placeholder(shape=w_indptr_shape, dtype="int32")
+ B = te.placeholder(shape=(M, N), dtype=dtype)
+
+ out = topi.nn.sparse_dense(
+ topi.nn.relu(X), W_data, W_indices, W_indptr
+ )
+ out = te.compute((M, N), lambda i, j: out[i, j] + B[i, j], name="BiasAdd")
+ out = topi.nn.relu(out)
+
+ return [X, W_data, W_indices, W_indptr, B, out]
+
+######################################################################
+# Special step for sparse workload
+# ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+# During schedule tuning, auto-scheduler will use random inputs to measure the performance of a
+# generated schedule. While we cannot directly use a random array as the input of a sparse op, for
+# the "indices" and "indptr" array are meaningful for the computation.
+#
+# To solve this problem, we register these as special buffers, and load them when process program
+# measuring.
+# See the :any:`auto_scheduler.measure` code for more details.
+
+# Define the basic shapes of this sparse computation
+M = K = N = 512
+BS_R = 16
+BS_C = 1
+density = 0.6
+
+# Generate the test data with numpy
+X_np = np.random.randn(M, K).astype("float32")
+X_np = np.maximum(np.zeros((M, K), dtype="float32"), X_np) # Relu
+W_sp_np = random_bsr_matrix(N, K, BS_R, BS_C, density=density, dtype="float32")
+W_np = W_sp_np.todense()
+Y_np = X_np @ W_np.T # Process the matrix multiplication
+B_np = np.random.randn(M, N).astype("float32")
+Y_np = Y_np + B_np # Bias add
+Y_np = np.maximum(np.zeros((M, N), dtype="float32"), Y_np) # Relu
+
+# Register the sparse data to special buffer
+prefix = "sparse_dense_bsr_%d_%d_%d_%d_%d_%.2f_" % (M, N, K, BS_R, BS_C, density)
+auto_scheduler.measure.register_special_buffer(prefix + "W_data", W_sp_np.data)
+auto_scheduler.measure.register_special_buffer(prefix + "W_indices", W_sp_np.indices)
+auto_scheduler.measure.register_special_buffer(prefix + "W_indptr", W_sp_np.indptr)
+
+######################################################################
+# Create the search task
+# ^^^^^^^^^^^^^^^^^^^^^^
+# We then create a search task with M=N=K=512 and dtype="float32"
+# If your machine supports avx instructions, you can
+#
+# - replace "llvm" below with "llvm -mcpu=core-avx2" to enable AVX2
+# - replace "llvm" below with "llvm -mcpu=skylake-avx512" to enable AVX-512
+
+target = tvm.target.Target("llvm")
+
+task = tvm.auto_scheduler.SearchTask(
+ func=sparse_dense,
+ args=(
+ M, N, K,
+ W_sp_np.data.shape,
+ W_sp_np.indices.shape,
+ W_sp_np.indptr.shape,
+ "float32"
+ ),
+ target=target
+)
+
+# Inspect the computational graph
+print("Computational DAG:")
+print(task.compute_dag)
+
+######################################################################
+# Write the custom sketch for sparse dense op
+# ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+# Before tuning, we will need to define the CustomSketchRule for the sparse dense op.
+#
+# CustomSketchRule consists of two parts: the condition function and the apply function.
+#
+# - condition function: describe when to use this sketch rule. For example, we can match the op
+# by their name or tag.
+# - apply function: describe how to generate the initial sketch. Auto-scheduler provides a set of
+# loop state APIs.
+
+def meet_condition_func(search_policy, state, stage_id):
+ state = auto_scheduler.loop_state.State(state, search_policy.search_task.compute_dag)
+ if state.stages[stage_id].op.tag in [
+ "sparse_dense_sp_rhs_bsrmm", "sparse_dense_sp_rhs_bsrmm_block"
+ ]:
+ return auto_scheduler.PreloadCustomSketchRule.APPLY_AND_SKIP_REST
+ else:
+ return auto_scheduler.PreloadCustomSketchRule.PASS
+
+def apply_func(search_policy, state, stage_id):
+ ret = []
+ s0 = auto_scheduler.loop_state.State(state, search_policy.search_task.compute_dag)
+ if s0.stages[stage_id].op.tag == "sparse_dense_sp_rhs_bsrmm_block":
+ return [s0.state_object, stage_id - 1]
+
+ sparse_dense = s0.stages[stage_id].op
+ sparse_dense_block = s0.stages[stage_id - 1].op
+ assert sparse_dense.tag == "sparse_dense_sp_rhs_bsrmm"
+ assert sparse_dense_block.tag == "sparse_dense_sp_rhs_bsrmm_block"
+
+ # Set the default consumer of compute block
+ consumer = sparse_dense
+
+ # If sparse dense has a single elementwise consumer
+ # We can compute inline the sparse_dense output stage
+ consumers = _ffi_api.SearchPolicyUtilsGetConsumers(
+ search_policy.search_task, s0.state_object, stage_id
+ )
+ if len(consumers) == 1:
+ consumer_id = int(consumers.items()[0][0])
+ if _ffi_api.SearchPolicyUtilsIsElementwiseMatch(
+ search_policy.search_task, s0.state_object, stage_id, consumer_id
+ ):
+ consumer = s0.stages[consumer_id].op
+ s0.compute_inline(sparse_dense)
+
+ i, nb_j, j, row_offset, c = s0[sparse_dense_block].iters
+ m, n = s0[consumer].iters
+ i0, i1, i2 = s0.split(sparse_dense_block, i, [None, None])
+ m0, m1 = s0.follow_split(consumer, m, len(s0.transform_steps) - 1, 1)
+ j0, j1 = s0.split(sparse_dense_block, nb_j, [None])
+ n0, n1 = s0.follow_split(consumer, n, len(s0.transform_steps) - 1, 1)
+ s0.reorder(sparse_dense_block, [i0, j0, i1, j1, row_offset, i2, j, c])
+ s0.reorder(consumer, [m0, n0, m1, n1])
+ s0.compute_at(sparse_dense_block, consumer, n0)
+
+ ret.append([s0.state_object, stage_id - 2])
+
+ return ret
+
+######################################################################
+# Next, we set parameters for the auto-scheduler.
Review comment:
```suggestion
# Next, we set parameters for the auto-scheduler with the custom sketch plugged in.
```
##########
File path: tutorials/auto_scheduler/tune_sparse_x86.py
##########
@@ -0,0 +1,331 @@
+# 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.
+"""
+Auto-scheduling Sparse Matrix Multiplication for CPU by Custom Sketch Rule
+==========================================================================
+**Author**: `Chengfan Jia <https://github.com/jcf94/>`_
+
+This is a tutorial on how to use the auto-scheduler to tune a sparse matrix multiplication for
+CPUs.
+
+Auto-scheduler is designed to explore the schedule with best performance for a given computation
+declaration automatically. While sometimes, we may have a demand to try some special ops which may
+not been well supported by auto-scheduler's default search policy. Auto-scheduler currently allows
+user to provide a CustomSketch to cover these cases.
+
+We use sparse matrix multiplication as an example in this tutorial.
+
+Note that this tutorial will not run on Windows or recent versions of macOS. To
+get it to run, you will need to wrap the body of this tutorial in a :code:`if
+__name__ == "__main__":` block.
+"""
+
+import os
+import itertools
+
+import numpy as np
+import tvm
+from tvm import te, auto_scheduler, topi
+from tvm.auto_scheduler import _ffi_api
+from tvm.topi.utils import get_const_tuple
+
+import scipy.sparse as sp
+
+######################################################################
+# Define the computation
+# ^^^^^^^^^^^^^^^^^^^^^^
+# To begin with, let us define the computation of a sparse matmul with several relu and bias add.
+# The function should return the list of input/output tensors.
+# From these tensors, the auto-scheduler can get the whole computational graph.
+
+# We use this function to generate a random bsr matrix
+def random_bsr_matrix(M, N, BS_R, BS_C, density, dtype):
+ import itertools
+
+ Y = np.zeros((M, N), dtype=dtype)
+ assert M % BS_R == 0
+ assert N % BS_C == 0
+ nnz = int(density * M * N)
+ num_blocks = int(nnz / (BS_R * BS_C)) + 1
+ candidate_blocks = np.asarray(list(itertools.product(range(0, M, BS_R), range(0, N, BS_C))))
+ assert candidate_blocks.shape[0] == M // BS_R * N // BS_C
+ chosen_blocks = candidate_blocks[
+ np.random.choice(candidate_blocks.shape[0], size=num_blocks, replace=False)
+ ]
+ for i in range(len(chosen_blocks)):
+ r, c = chosen_blocks[i]
+ Y[r : r + BS_R, c : c + BS_C] = np.random.randn(BS_R, BS_C)
+ s = sp.bsr_matrix(Y, blocksize=(BS_R, BS_C))
+ assert s.data.shape == (num_blocks, BS_R, BS_C)
+ assert s.indices.shape == (num_blocks,)
+ assert s.indptr.shape == (M // BS_R + 1,)
+ return s
+
+@auto_scheduler.register_workload
+def sparse_dense(M, N, K, w_data_shape, w_indices_shape, w_indptr_shape, dtype):
+ X = te.placeholder(shape=(M, K), dtype=dtype)
+ W_data = te.placeholder(shape=w_data_shape, dtype=dtype)
+ W_indices = te.placeholder(shape=w_indices_shape, dtype="int32")
+ W_indptr = te.placeholder(shape=w_indptr_shape, dtype="int32")
+ B = te.placeholder(shape=(M, N), dtype=dtype)
+
+ out = topi.nn.sparse_dense(
+ topi.nn.relu(X), W_data, W_indices, W_indptr
+ )
+ out = te.compute((M, N), lambda i, j: out[i, j] + B[i, j], name="BiasAdd")
+ out = topi.nn.relu(out)
+
+ return [X, W_data, W_indices, W_indptr, B, out]
+
+######################################################################
+# Special step for sparse workload
+# ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+# During schedule tuning, auto-scheduler will use random inputs to measure the performance of a
+# generated schedule. While we cannot directly use a random array as the input of a sparse op, for
+# the "indices" and "indptr" array are meaningful for the computation.
+#
+# To solve this problem, we register these as special buffers, and load them when process program
+# measuring.
+# See the :any:`auto_scheduler.measure` code for more details.
+
+# Define the basic shapes of this sparse computation
+M = K = N = 512
+BS_R = 16
+BS_C = 1
+density = 0.6
+
+# Generate the test data with numpy
+X_np = np.random.randn(M, K).astype("float32")
+X_np = np.maximum(np.zeros((M, K), dtype="float32"), X_np) # Relu
+W_sp_np = random_bsr_matrix(N, K, BS_R, BS_C, density=density, dtype="float32")
+W_np = W_sp_np.todense()
+Y_np = X_np @ W_np.T # Process the matrix multiplication
+B_np = np.random.randn(M, N).astype("float32")
+Y_np = Y_np + B_np # Bias add
+Y_np = np.maximum(np.zeros((M, N), dtype="float32"), Y_np) # Relu
+
+# Register the sparse data to special buffer
+prefix = "sparse_dense_bsr_%d_%d_%d_%d_%d_%.2f_" % (M, N, K, BS_R, BS_C, density)
+auto_scheduler.measure.register_special_buffer(prefix + "W_data", W_sp_np.data)
+auto_scheduler.measure.register_special_buffer(prefix + "W_indices", W_sp_np.indices)
+auto_scheduler.measure.register_special_buffer(prefix + "W_indptr", W_sp_np.indptr)
+
+######################################################################
+# Create the search task
+# ^^^^^^^^^^^^^^^^^^^^^^
+# We then create a search task with M=N=K=512 and dtype="float32"
+# If your machine supports avx instructions, you can
+#
+# - replace "llvm" below with "llvm -mcpu=core-avx2" to enable AVX2
+# - replace "llvm" below with "llvm -mcpu=skylake-avx512" to enable AVX-512
+
+target = tvm.target.Target("llvm")
+
+task = tvm.auto_scheduler.SearchTask(
+ func=sparse_dense,
+ args=(
+ M, N, K,
+ W_sp_np.data.shape,
+ W_sp_np.indices.shape,
+ W_sp_np.indptr.shape,
+ "float32"
+ ),
+ target=target
+)
+
+# Inspect the computational graph
+print("Computational DAG:")
+print(task.compute_dag)
+
+######################################################################
+# Write the custom sketch for sparse dense op
+# ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+# Before tuning, we will need to define the CustomSketchRule for the sparse dense op.
+#
+# CustomSketchRule consists of two parts: the condition function and the apply function.
+#
+# - condition function: describe when to use this sketch rule. For example, we can match the op
+# by their name or tag.
+# - apply function: describe how to generate the initial sketch. Auto-scheduler provides a set of
+# loop state APIs.
Review comment:
```suggestion
# - apply function: describe how to generate the initial sketch. You can implement it using auto-scheduler provided loop state APIs.
```
##########
File path: tutorials/auto_scheduler/tune_sparse_x86.py
##########
@@ -0,0 +1,331 @@
+# 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.
+"""
+Auto-scheduling Sparse Matrix Multiplication for CPU by Custom Sketch Rule
+==========================================================================
+**Author**: `Chengfan Jia <https://github.com/jcf94/>`_
+
+This is a tutorial on how to use the auto-scheduler to tune a sparse matrix multiplication for
+CPUs.
+
+Auto-scheduler is designed to explore the schedule with best performance for a given computation
+declaration automatically. While sometimes, we may have a demand to try some special ops which may
+not been well supported by auto-scheduler's default search policy. Auto-scheduler currently allows
+user to provide a CustomSketch to cover these cases.
+
+We use sparse matrix multiplication as an example in this tutorial.
Review comment:
```suggestion
We use sparse matrix multiplication as an example in this tutorial to demonstrate how to implement and plug a custom sketch rule to the auto-scheduler search policy.
```
##########
File path: tutorials/auto_scheduler/tune_sparse_x86.py
##########
@@ -0,0 +1,331 @@
+# 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.
+"""
+Auto-scheduling Sparse Matrix Multiplication for CPU by Custom Sketch Rule
+==========================================================================
+**Author**: `Chengfan Jia <https://github.com/jcf94/>`_
+
+This is a tutorial on how to use the auto-scheduler to tune a sparse matrix multiplication for
+CPUs.
+
+Auto-scheduler is designed to explore the schedule with best performance for a given computation
+declaration automatically. While sometimes, we may have a demand to try some special ops which may
+not been well supported by auto-scheduler's default search policy. Auto-scheduler currently allows
Review comment:
```suggestion
not been well-supported by auto-scheduler's default sketch rules and result in poor performance. Fortunately, auto-scheduler currently allows
```
##########
File path: tutorials/auto_scheduler/tune_sparse_x86.py
##########
@@ -0,0 +1,331 @@
+# 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.
+"""
+Auto-scheduling Sparse Matrix Multiplication for CPU by Custom Sketch Rule
+==========================================================================
+**Author**: `Chengfan Jia <https://github.com/jcf94/>`_
+
+This is a tutorial on how to use the auto-scheduler to tune a sparse matrix multiplication for
+CPUs.
+
+Auto-scheduler is designed to explore the schedule with best performance for a given computation
+declaration automatically. While sometimes, we may have a demand to try some special ops which may
+not been well supported by auto-scheduler's default search policy. Auto-scheduler currently allows
+user to provide a CustomSketch to cover these cases.
+
+We use sparse matrix multiplication as an example in this tutorial.
+
+Note that this tutorial will not run on Windows or recent versions of macOS. To
+get it to run, you will need to wrap the body of this tutorial in a :code:`if
+__name__ == "__main__":` block.
+"""
+
+import os
+import itertools
+
+import numpy as np
+import tvm
+from tvm import te, auto_scheduler, topi
+from tvm.auto_scheduler import _ffi_api
+from tvm.topi.utils import get_const_tuple
+
+import scipy.sparse as sp
+
+######################################################################
+# Define the computation
+# ^^^^^^^^^^^^^^^^^^^^^^
+# To begin with, let us define the computation of a sparse matmul with several relu and bias add.
+# The function should return the list of input/output tensors.
+# From these tensors, the auto-scheduler can get the whole computational graph.
+
+# We use this function to generate a random bsr matrix
+def random_bsr_matrix(M, N, BS_R, BS_C, density, dtype):
+ import itertools
+
+ Y = np.zeros((M, N), dtype=dtype)
+ assert M % BS_R == 0
+ assert N % BS_C == 0
+ nnz = int(density * M * N)
+ num_blocks = int(nnz / (BS_R * BS_C)) + 1
+ candidate_blocks = np.asarray(list(itertools.product(range(0, M, BS_R), range(0, N, BS_C))))
+ assert candidate_blocks.shape[0] == M // BS_R * N // BS_C
+ chosen_blocks = candidate_blocks[
+ np.random.choice(candidate_blocks.shape[0], size=num_blocks, replace=False)
+ ]
+ for i in range(len(chosen_blocks)):
+ r, c = chosen_blocks[i]
+ Y[r : r + BS_R, c : c + BS_C] = np.random.randn(BS_R, BS_C)
+ s = sp.bsr_matrix(Y, blocksize=(BS_R, BS_C))
+ assert s.data.shape == (num_blocks, BS_R, BS_C)
+ assert s.indices.shape == (num_blocks,)
+ assert s.indptr.shape == (M // BS_R + 1,)
+ return s
+
+@auto_scheduler.register_workload
+def sparse_dense(M, N, K, w_data_shape, w_indices_shape, w_indptr_shape, dtype):
+ X = te.placeholder(shape=(M, K), dtype=dtype)
+ W_data = te.placeholder(shape=w_data_shape, dtype=dtype)
+ W_indices = te.placeholder(shape=w_indices_shape, dtype="int32")
+ W_indptr = te.placeholder(shape=w_indptr_shape, dtype="int32")
+ B = te.placeholder(shape=(M, N), dtype=dtype)
+
+ out = topi.nn.sparse_dense(
+ topi.nn.relu(X), W_data, W_indices, W_indptr
+ )
+ out = te.compute((M, N), lambda i, j: out[i, j] + B[i, j], name="BiasAdd")
+ out = topi.nn.relu(out)
+
+ return [X, W_data, W_indices, W_indptr, B, out]
+
+######################################################################
+# Special step for sparse workload
+# ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+# During schedule tuning, auto-scheduler will use random inputs to measure the performance of a
+# generated schedule. While we cannot directly use a random array as the input of a sparse op, for
+# the "indices" and "indptr" array are meaningful for the computation.
+#
+# To solve this problem, we register these as special buffers, and load them when process program
+# measuring.
+# See the :any:`auto_scheduler.measure` code for more details.
+
+# Define the basic shapes of this sparse computation
+M = K = N = 512
+BS_R = 16
+BS_C = 1
+density = 0.6
+
+# Generate the test data with numpy
+X_np = np.random.randn(M, K).astype("float32")
+X_np = np.maximum(np.zeros((M, K), dtype="float32"), X_np) # Relu
+W_sp_np = random_bsr_matrix(N, K, BS_R, BS_C, density=density, dtype="float32")
+W_np = W_sp_np.todense()
+Y_np = X_np @ W_np.T # Process the matrix multiplication
+B_np = np.random.randn(M, N).astype("float32")
+Y_np = Y_np + B_np # Bias add
+Y_np = np.maximum(np.zeros((M, N), dtype="float32"), Y_np) # Relu
+
+# Register the sparse data to special buffer
+prefix = "sparse_dense_bsr_%d_%d_%d_%d_%d_%.2f_" % (M, N, K, BS_R, BS_C, density)
+auto_scheduler.measure.register_special_buffer(prefix + "W_data", W_sp_np.data)
+auto_scheduler.measure.register_special_buffer(prefix + "W_indices", W_sp_np.indices)
+auto_scheduler.measure.register_special_buffer(prefix + "W_indptr", W_sp_np.indptr)
+
+######################################################################
+# Create the search task
+# ^^^^^^^^^^^^^^^^^^^^^^
+# We then create a search task with M=N=K=512 and dtype="float32"
+# If your machine supports avx instructions, you can
+#
+# - replace "llvm" below with "llvm -mcpu=core-avx2" to enable AVX2
+# - replace "llvm" below with "llvm -mcpu=skylake-avx512" to enable AVX-512
+
+target = tvm.target.Target("llvm")
+
+task = tvm.auto_scheduler.SearchTask(
+ func=sparse_dense,
+ args=(
+ M, N, K,
+ W_sp_np.data.shape,
+ W_sp_np.indices.shape,
+ W_sp_np.indptr.shape,
+ "float32"
+ ),
+ target=target
+)
+
+# Inspect the computational graph
+print("Computational DAG:")
+print(task.compute_dag)
+
+######################################################################
+# Write the custom sketch for sparse dense op
+# ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+# Before tuning, we will need to define the CustomSketchRule for the sparse dense op.
+#
+# CustomSketchRule consists of two parts: the condition function and the apply function.
+#
+# - condition function: describe when to use this sketch rule. For example, we can match the op
+# by their name or tag.
+# - apply function: describe how to generate the initial sketch. Auto-scheduler provides a set of
+# loop state APIs.
+
+def meet_condition_func(search_policy, state, stage_id):
+ state = auto_scheduler.loop_state.State(state, search_policy.search_task.compute_dag)
+ if state.stages[stage_id].op.tag in [
+ "sparse_dense_sp_rhs_bsrmm", "sparse_dense_sp_rhs_bsrmm_block"
+ ]:
+ return auto_scheduler.PreloadCustomSketchRule.APPLY_AND_SKIP_REST
+ else:
+ return auto_scheduler.PreloadCustomSketchRule.PASS
+
+def apply_func(search_policy, state, stage_id):
+ ret = []
+ s0 = auto_scheduler.loop_state.State(state, search_policy.search_task.compute_dag)
+ if s0.stages[stage_id].op.tag == "sparse_dense_sp_rhs_bsrmm_block":
+ return [s0.state_object, stage_id - 1]
+
+ sparse_dense = s0.stages[stage_id].op
+ sparse_dense_block = s0.stages[stage_id - 1].op
+ assert sparse_dense.tag == "sparse_dense_sp_rhs_bsrmm"
+ assert sparse_dense_block.tag == "sparse_dense_sp_rhs_bsrmm_block"
+
+ # Set the default consumer of compute block
+ consumer = sparse_dense
+
+ # If sparse dense has a single elementwise consumer
+ # We can compute inline the sparse_dense output stage
+ consumers = _ffi_api.SearchPolicyUtilsGetConsumers(
+ search_policy.search_task, s0.state_object, stage_id
+ )
+ if len(consumers) == 1:
+ consumer_id = int(consumers.items()[0][0])
+ if _ffi_api.SearchPolicyUtilsIsElementwiseMatch(
+ search_policy.search_task, s0.state_object, stage_id, consumer_id
+ ):
+ consumer = s0.stages[consumer_id].op
+ s0.compute_inline(sparse_dense)
+
+ i, nb_j, j, row_offset, c = s0[sparse_dense_block].iters
+ m, n = s0[consumer].iters
+ i0, i1, i2 = s0.split(sparse_dense_block, i, [None, None])
+ m0, m1 = s0.follow_split(consumer, m, len(s0.transform_steps) - 1, 1)
+ j0, j1 = s0.split(sparse_dense_block, nb_j, [None])
+ n0, n1 = s0.follow_split(consumer, n, len(s0.transform_steps) - 1, 1)
+ s0.reorder(sparse_dense_block, [i0, j0, i1, j1, row_offset, i2, j, c])
+ s0.reorder(consumer, [m0, n0, m1, n1])
+ s0.compute_at(sparse_dense_block, consumer, n0)
+
+ ret.append([s0.state_object, stage_id - 2])
+
+ return ret
+
+######################################################################
+# Next, we set parameters for the auto-scheduler.
+#
+# * :code:`num_measure_trials` is the number of measurement trials we can use during the search.
+# We only make 10 trials in this tutorial for a fast demonstration. In practice, 1000 is a
+# good value for the search to converge. You can do more trials according to your time budget.
+# * In addition, we use :code:`RecordToFile` to dump measurement records into a file `matmul.json`.
+# The measurement records can be used to query the history best, resume the search,
+# and do more analyses later.
+# * see :any:`auto_scheduler.TuningOptions` for more parameters
+# * Here, we need to create a :code:`auto_scheduler.SketchPolicy` object, and add the custom sketch
+# rule as a `init_search_callbacks`.
+
+log_file = "sparse_dense.json"
+tune_option = auto_scheduler.TuningOptions(
+ num_measure_trials=10,
+ measure_callbacks=[auto_scheduler.RecordToFile(log_file)],
+ verbose=2,
+)
+
+search_policy = auto_scheduler.SketchPolicy(
+ task,
+ program_cost_model=auto_scheduler.XGBModel(),
+ init_search_callbacks=[
+ auto_scheduler.PreloadCustomSketchRule(meet_condition_func, apply_func, "SparseDense")
+ ]
+)
+
+######################################################################
+# Run the search
+# ^^^^^^^^^^^^^^
+# Now we get all inputs ready.
+# We can kick off the search and let the auto-scheduler do its magic.
+# After some measurement trials, we can load the best schedule from the log
+# file and apply it.
+
+# Run auto-tuning (search)
+task.tune(tune_option, search_policy)
+# Apply the best schedule
+sch, args = task.apply_best(log_file)
+
+######################################################################
+# We can lower the schedule to see the IR after auto-scheduling.
+# The auto-scheduler correctly performs optimizations including multi-level tiling,
+# layout transformation, parallelization, vectorization, unrolling, and operator fusion.
+
+print("Lowered TIR:")
+print(tvm.lower(sch, args, simple_mode=True))
+
+######################################################################
+# Check correctness and evaluate performance
+# ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+# We build the binary and check its correctness and performance.
+
+func = tvm.build(sch, args, target)
+
+ctx = tvm.cpu()
+
+X_tvm = tvm.nd.array(X_np, ctx=ctx)
+W_data_tvm = tvm.nd.array(W_sp_np.data, ctx=ctx)
+W_indices_tvm = tvm.nd.array(W_sp_np.indices, ctx=ctx)
+W_indptr_tvm = tvm.nd.array(W_sp_np.indptr, ctx=ctx)
+B_tvm = tvm.nd.array(B_np, ctx=ctx)
+Y_tvm = tvm.nd.empty(Y_np.shape, ctx=ctx)
+
+func(X_tvm, W_data_tvm, W_indices_tvm, W_indptr_tvm, B_tvm, Y_tvm)
+
+# Check results
+tvm.testing.assert_allclose(Y_np, Y_tvm.asnumpy(), atol=1e-4, rtol=1e-4)
+
+# Evaluate execution time.
+evaluator = func.time_evaluator(func.entry_name, ctx, min_repeat_ms=500)
+print(
+ "Execution time of this operator: %.3f ms"
+ % (np.median(evaluator(X_tvm, W_data_tvm, W_indices_tvm, W_indptr_tvm, B_tvm, Y_tvm).results) * 1000)
+)
+
+######################################################################
+# Using the record file
+# ^^^^^^^^^^^^^^^^^^^^^
+# During the search, all measurement records are dumped into the record
+# file "matmul.json". The measurement records can be used to re-apply search results,
+# resume the search, and perform other analyses.
+
+######################################################################
+# Here is an example where we load the best schedule from a file,
+# and print the equivalent python schedule API. This can be used for
+# debugging and learning the behavior of the auto-scheduler.
+
+print("Equivalent python schedule:")
+print(task.print_best(log_file))
+
+######################################################################
+# A more complicated example is to resume the search.
+# In this case, we need to create the search policy and cost model by ourselves
+# and resume the status of search policy and cost model with the log file.
+# In the example below we resume the status and do more 5 trials.
+
+
+def resume_search(task, log_file):
+ print("Resume search:")
+ cost_model = auto_scheduler.XGBModel()
+ cost_model.update_from_file(log_file)
+ search_policy = auto_scheduler.SketchPolicy(
+ task, cost_model, init_search_callbacks=[
+ auto_scheduler.PreloadMeasuredStates(log_file),
+ auto_scheduler.PreloadCustomSketchRule(meet_condition_func, apply_func, "SparseDense")
+ ]
+ )
+ tune_option = auto_scheduler.TuningOptions(
+ num_measure_trials=5, measure_callbacks=[auto_scheduler.RecordToFile(log_file)]
+ )
+ task.tune(tune_option, search_policy=search_policy)
+
+
+resume_search(task, log_file)
Review comment:
I think you can simply refer to other tutorials to skip this part. This tutorial is more advance so it should be fine to assume most readers are more or less familiar with auto-scheduler already.
##########
File path: tutorials/auto_scheduler/tune_sparse_x86.py
##########
@@ -0,0 +1,331 @@
+# 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.
+"""
+Auto-scheduling Sparse Matrix Multiplication for CPU by Custom Sketch Rule
+==========================================================================
+**Author**: `Chengfan Jia <https://github.com/jcf94/>`_
+
+This is a tutorial on how to use the auto-scheduler to tune a sparse matrix multiplication for
+CPUs.
+
+Auto-scheduler is designed to explore the schedule with best performance for a given computation
+declaration automatically. While sometimes, we may have a demand to try some special ops which may
+not been well supported by auto-scheduler's default search policy. Auto-scheduler currently allows
+user to provide a CustomSketch to cover these cases.
+
+We use sparse matrix multiplication as an example in this tutorial.
+
+Note that this tutorial will not run on Windows or recent versions of macOS. To
+get it to run, you will need to wrap the body of this tutorial in a :code:`if
+__name__ == "__main__":` block.
+"""
+
+import os
+import itertools
+
+import numpy as np
+import tvm
+from tvm import te, auto_scheduler, topi
+from tvm.auto_scheduler import _ffi_api
+from tvm.topi.utils import get_const_tuple
+
+import scipy.sparse as sp
+
+######################################################################
+# Define the computation
+# ^^^^^^^^^^^^^^^^^^^^^^
+# To begin with, let us define the computation of a sparse matmul with several relu and bias add.
+# The function should return the list of input/output tensors.
+# From these tensors, the auto-scheduler can get the whole computational graph.
+
+# We use this function to generate a random bsr matrix
+def random_bsr_matrix(M, N, BS_R, BS_C, density, dtype):
+ import itertools
+
+ Y = np.zeros((M, N), dtype=dtype)
+ assert M % BS_R == 0
+ assert N % BS_C == 0
+ nnz = int(density * M * N)
+ num_blocks = int(nnz / (BS_R * BS_C)) + 1
+ candidate_blocks = np.asarray(list(itertools.product(range(0, M, BS_R), range(0, N, BS_C))))
+ assert candidate_blocks.shape[0] == M // BS_R * N // BS_C
+ chosen_blocks = candidate_blocks[
+ np.random.choice(candidate_blocks.shape[0], size=num_blocks, replace=False)
+ ]
+ for i in range(len(chosen_blocks)):
+ r, c = chosen_blocks[i]
+ Y[r : r + BS_R, c : c + BS_C] = np.random.randn(BS_R, BS_C)
+ s = sp.bsr_matrix(Y, blocksize=(BS_R, BS_C))
+ assert s.data.shape == (num_blocks, BS_R, BS_C)
+ assert s.indices.shape == (num_blocks,)
+ assert s.indptr.shape == (M // BS_R + 1,)
+ return s
+
+@auto_scheduler.register_workload
+def sparse_dense(M, N, K, w_data_shape, w_indices_shape, w_indptr_shape, dtype):
+ X = te.placeholder(shape=(M, K), dtype=dtype)
+ W_data = te.placeholder(shape=w_data_shape, dtype=dtype)
+ W_indices = te.placeholder(shape=w_indices_shape, dtype="int32")
+ W_indptr = te.placeholder(shape=w_indptr_shape, dtype="int32")
+ B = te.placeholder(shape=(M, N), dtype=dtype)
+
+ out = topi.nn.sparse_dense(
+ topi.nn.relu(X), W_data, W_indices, W_indptr
+ )
+ out = te.compute((M, N), lambda i, j: out[i, j] + B[i, j], name="BiasAdd")
+ out = topi.nn.relu(out)
+
+ return [X, W_data, W_indices, W_indptr, B, out]
+
+######################################################################
+# Special step for sparse workload
+# ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+# During schedule tuning, auto-scheduler will use random inputs to measure the performance of a
+# generated schedule. While we cannot directly use a random array as the input of a sparse op, for
+# the "indices" and "indptr" array are meaningful for the computation.
+#
+# To solve this problem, we register these as special buffers, and load them when process program
+# measuring.
+# See the :any:`auto_scheduler.measure` code for more details.
+
+# Define the basic shapes of this sparse computation
+M = K = N = 512
+BS_R = 16
+BS_C = 1
+density = 0.6
+
+# Generate the test data with numpy
+X_np = np.random.randn(M, K).astype("float32")
+X_np = np.maximum(np.zeros((M, K), dtype="float32"), X_np) # Relu
+W_sp_np = random_bsr_matrix(N, K, BS_R, BS_C, density=density, dtype="float32")
+W_np = W_sp_np.todense()
+Y_np = X_np @ W_np.T # Process the matrix multiplication
+B_np = np.random.randn(M, N).astype("float32")
+Y_np = Y_np + B_np # Bias add
+Y_np = np.maximum(np.zeros((M, N), dtype="float32"), Y_np) # Relu
+
+# Register the sparse data to special buffer
+prefix = "sparse_dense_bsr_%d_%d_%d_%d_%d_%.2f_" % (M, N, K, BS_R, BS_C, density)
+auto_scheduler.measure.register_special_buffer(prefix + "W_data", W_sp_np.data)
+auto_scheduler.measure.register_special_buffer(prefix + "W_indices", W_sp_np.indices)
+auto_scheduler.measure.register_special_buffer(prefix + "W_indptr", W_sp_np.indptr)
+
+######################################################################
+# Create the search task
+# ^^^^^^^^^^^^^^^^^^^^^^
+# We then create a search task with M=N=K=512 and dtype="float32"
+# If your machine supports avx instructions, you can
+#
+# - replace "llvm" below with "llvm -mcpu=core-avx2" to enable AVX2
+# - replace "llvm" below with "llvm -mcpu=skylake-avx512" to enable AVX-512
+
+target = tvm.target.Target("llvm")
+
+task = tvm.auto_scheduler.SearchTask(
+ func=sparse_dense,
+ args=(
+ M, N, K,
+ W_sp_np.data.shape,
+ W_sp_np.indices.shape,
+ W_sp_np.indptr.shape,
+ "float32"
+ ),
+ target=target
+)
+
+# Inspect the computational graph
+print("Computational DAG:")
+print(task.compute_dag)
+
+######################################################################
+# Write the custom sketch for sparse dense op
+# ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+# Before tuning, we will need to define the CustomSketchRule for the sparse dense op.
+#
+# CustomSketchRule consists of two parts: the condition function and the apply function.
+#
+# - condition function: describe when to use this sketch rule. For example, we can match the op
+# by their name or tag.
+# - apply function: describe how to generate the initial sketch. Auto-scheduler provides a set of
+# loop state APIs.
+
+def meet_condition_func(search_policy, state, stage_id):
+ state = auto_scheduler.loop_state.State(state, search_policy.search_task.compute_dag)
+ if state.stages[stage_id].op.tag in [
+ "sparse_dense_sp_rhs_bsrmm", "sparse_dense_sp_rhs_bsrmm_block"
+ ]:
+ return auto_scheduler.PreloadCustomSketchRule.APPLY_AND_SKIP_REST
+ else:
+ return auto_scheduler.PreloadCustomSketchRule.PASS
+
+def apply_func(search_policy, state, stage_id):
+ ret = []
+ s0 = auto_scheduler.loop_state.State(state, search_policy.search_task.compute_dag)
+ if s0.stages[stage_id].op.tag == "sparse_dense_sp_rhs_bsrmm_block":
+ return [s0.state_object, stage_id - 1]
+
+ sparse_dense = s0.stages[stage_id].op
+ sparse_dense_block = s0.stages[stage_id - 1].op
+ assert sparse_dense.tag == "sparse_dense_sp_rhs_bsrmm"
+ assert sparse_dense_block.tag == "sparse_dense_sp_rhs_bsrmm_block"
+
+ # Set the default consumer of compute block
+ consumer = sparse_dense
+
+ # If sparse dense has a single elementwise consumer
+ # We can compute inline the sparse_dense output stage
+ consumers = _ffi_api.SearchPolicyUtilsGetConsumers(
+ search_policy.search_task, s0.state_object, stage_id
+ )
+ if len(consumers) == 1:
+ consumer_id = int(consumers.items()[0][0])
+ if _ffi_api.SearchPolicyUtilsIsElementwiseMatch(
+ search_policy.search_task, s0.state_object, stage_id, consumer_id
+ ):
+ consumer = s0.stages[consumer_id].op
+ s0.compute_inline(sparse_dense)
+
+ i, nb_j, j, row_offset, c = s0[sparse_dense_block].iters
+ m, n = s0[consumer].iters
+ i0, i1, i2 = s0.split(sparse_dense_block, i, [None, None])
+ m0, m1 = s0.follow_split(consumer, m, len(s0.transform_steps) - 1, 1)
+ j0, j1 = s0.split(sparse_dense_block, nb_j, [None])
+ n0, n1 = s0.follow_split(consumer, n, len(s0.transform_steps) - 1, 1)
+ s0.reorder(sparse_dense_block, [i0, j0, i1, j1, row_offset, i2, j, c])
+ s0.reorder(consumer, [m0, n0, m1, n1])
+ s0.compute_at(sparse_dense_block, consumer, n0)
+
+ ret.append([s0.state_object, stage_id - 2])
+
+ return ret
+
+######################################################################
+# Next, we set parameters for the auto-scheduler.
+#
+# * :code:`num_measure_trials` is the number of measurement trials we can use during the search.
+# We only make 10 trials in this tutorial for a fast demonstration. In practice, 1000 is a
+# good value for the search to converge. You can do more trials according to your time budget.
+# * In addition, we use :code:`RecordToFile` to dump measurement records into a file `matmul.json`.
+# The measurement records can be used to query the history best, resume the search,
+# and do more analyses later.
+# * see :any:`auto_scheduler.TuningOptions` for more parameters
+# * Here, we need to create a :code:`auto_scheduler.SketchPolicy` object, and add the custom sketch
+# rule as a `init_search_callbacks`.
+
+log_file = "sparse_dense.json"
+tune_option = auto_scheduler.TuningOptions(
+ num_measure_trials=10,
+ measure_callbacks=[auto_scheduler.RecordToFile(log_file)],
+ verbose=2,
+)
+
+search_policy = auto_scheduler.SketchPolicy(
+ task,
+ program_cost_model=auto_scheduler.XGBModel(),
+ init_search_callbacks=[
+ auto_scheduler.PreloadCustomSketchRule(meet_condition_func, apply_func, "SparseDense")
+ ]
+)
+
+######################################################################
+# Run the search
+# ^^^^^^^^^^^^^^
+# Now we get all inputs ready.
+# We can kick off the search and let the auto-scheduler do its magic.
+# After some measurement trials, we can load the best schedule from the log
+# file and apply it.
+
+# Run auto-tuning (search)
+task.tune(tune_option, search_policy)
Review comment:
Do not run the tuning in the tutorial. We should comment out this line and commit a pre-tuned log to `ci_logs`.
##########
File path: tutorials/auto_scheduler/tune_sparse_x86.py
##########
@@ -0,0 +1,331 @@
+# 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.
+"""
+Auto-scheduling Sparse Matrix Multiplication for CPU by Custom Sketch Rule
+==========================================================================
+**Author**: `Chengfan Jia <https://github.com/jcf94/>`_
+
+This is a tutorial on how to use the auto-scheduler to tune a sparse matrix multiplication for
+CPUs.
+
+Auto-scheduler is designed to explore the schedule with best performance for a given computation
+declaration automatically. While sometimes, we may have a demand to try some special ops which may
+not been well supported by auto-scheduler's default search policy. Auto-scheduler currently allows
+user to provide a CustomSketch to cover these cases.
+
+We use sparse matrix multiplication as an example in this tutorial.
+
+Note that this tutorial will not run on Windows or recent versions of macOS. To
+get it to run, you will need to wrap the body of this tutorial in a :code:`if
+__name__ == "__main__":` block.
+"""
+
+import os
+import itertools
+
+import numpy as np
+import tvm
+from tvm import te, auto_scheduler, topi
+from tvm.auto_scheduler import _ffi_api
+from tvm.topi.utils import get_const_tuple
+
+import scipy.sparse as sp
+
+######################################################################
+# Define the computation
+# ^^^^^^^^^^^^^^^^^^^^^^
+# To begin with, let us define the computation of a sparse matmul with several relu and bias add.
+# The function should return the list of input/output tensors.
+# From these tensors, the auto-scheduler can get the whole computational graph.
+
+# We use this function to generate a random bsr matrix
+def random_bsr_matrix(M, N, BS_R, BS_C, density, dtype):
+ import itertools
+
+ Y = np.zeros((M, N), dtype=dtype)
+ assert M % BS_R == 0
+ assert N % BS_C == 0
+ nnz = int(density * M * N)
+ num_blocks = int(nnz / (BS_R * BS_C)) + 1
+ candidate_blocks = np.asarray(list(itertools.product(range(0, M, BS_R), range(0, N, BS_C))))
+ assert candidate_blocks.shape[0] == M // BS_R * N // BS_C
+ chosen_blocks = candidate_blocks[
+ np.random.choice(candidate_blocks.shape[0], size=num_blocks, replace=False)
+ ]
+ for i in range(len(chosen_blocks)):
+ r, c = chosen_blocks[i]
+ Y[r : r + BS_R, c : c + BS_C] = np.random.randn(BS_R, BS_C)
+ s = sp.bsr_matrix(Y, blocksize=(BS_R, BS_C))
+ assert s.data.shape == (num_blocks, BS_R, BS_C)
+ assert s.indices.shape == (num_blocks,)
+ assert s.indptr.shape == (M // BS_R + 1,)
+ return s
+
+@auto_scheduler.register_workload
+def sparse_dense(M, N, K, w_data_shape, w_indices_shape, w_indptr_shape, dtype):
+ X = te.placeholder(shape=(M, K), dtype=dtype)
+ W_data = te.placeholder(shape=w_data_shape, dtype=dtype)
+ W_indices = te.placeholder(shape=w_indices_shape, dtype="int32")
+ W_indptr = te.placeholder(shape=w_indptr_shape, dtype="int32")
+ B = te.placeholder(shape=(M, N), dtype=dtype)
+
+ out = topi.nn.sparse_dense(
+ topi.nn.relu(X), W_data, W_indices, W_indptr
+ )
+ out = te.compute((M, N), lambda i, j: out[i, j] + B[i, j], name="BiasAdd")
+ out = topi.nn.relu(out)
+
+ return [X, W_data, W_indices, W_indptr, B, out]
+
+######################################################################
+# Special step for sparse workload
+# ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+# During schedule tuning, auto-scheduler will use random inputs to measure the performance of a
+# generated schedule. While we cannot directly use a random array as the input of a sparse op, for
+# the "indices" and "indptr" array are meaningful for the computation.
+#
+# To solve this problem, we register these as special buffers, and load them when process program
+# measuring.
+# See the :any:`auto_scheduler.measure` code for more details.
+
+# Define the basic shapes of this sparse computation
+M = K = N = 512
+BS_R = 16
+BS_C = 1
+density = 0.6
+
+# Generate the test data with numpy
+X_np = np.random.randn(M, K).astype("float32")
+X_np = np.maximum(np.zeros((M, K), dtype="float32"), X_np) # Relu
+W_sp_np = random_bsr_matrix(N, K, BS_R, BS_C, density=density, dtype="float32")
+W_np = W_sp_np.todense()
+Y_np = X_np @ W_np.T # Process the matrix multiplication
+B_np = np.random.randn(M, N).astype("float32")
+Y_np = Y_np + B_np # Bias add
+Y_np = np.maximum(np.zeros((M, N), dtype="float32"), Y_np) # Relu
+
+# Register the sparse data to special buffer
+prefix = "sparse_dense_bsr_%d_%d_%d_%d_%d_%.2f_" % (M, N, K, BS_R, BS_C, density)
+auto_scheduler.measure.register_special_buffer(prefix + "W_data", W_sp_np.data)
+auto_scheduler.measure.register_special_buffer(prefix + "W_indices", W_sp_np.indices)
+auto_scheduler.measure.register_special_buffer(prefix + "W_indptr", W_sp_np.indptr)
+
+######################################################################
+# Create the search task
+# ^^^^^^^^^^^^^^^^^^^^^^
+# We then create a search task with M=N=K=512 and dtype="float32"
+# If your machine supports avx instructions, you can
+#
+# - replace "llvm" below with "llvm -mcpu=core-avx2" to enable AVX2
+# - replace "llvm" below with "llvm -mcpu=skylake-avx512" to enable AVX-512
+
+target = tvm.target.Target("llvm")
+
+task = tvm.auto_scheduler.SearchTask(
+ func=sparse_dense,
+ args=(
+ M, N, K,
+ W_sp_np.data.shape,
+ W_sp_np.indices.shape,
+ W_sp_np.indptr.shape,
+ "float32"
+ ),
+ target=target
+)
+
+# Inspect the computational graph
+print("Computational DAG:")
+print(task.compute_dag)
+
+######################################################################
+# Write the custom sketch for sparse dense op
+# ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+# Before tuning, we will need to define the CustomSketchRule for the sparse dense op.
+#
+# CustomSketchRule consists of two parts: the condition function and the apply function.
+#
+# - condition function: describe when to use this sketch rule. For example, we can match the op
+# by their name or tag.
Review comment:
```suggestion
# - condition function: describe when to apply this sketch rule. For example, we can only apply the rule
# to the sparse ops by matching their name and tag.
```
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] merrymercy commented on a change in pull request #7313: [AutoSchedule] Sparse dense tuning support with custom sketch rule
Posted by GitBox <gi...@apache.org>.
merrymercy commented on a change in pull request #7313:
URL: https://github.com/apache/tvm/pull/7313#discussion_r586075755
##########
File path: python/tvm/auto_scheduler/search_task.py
##########
@@ -157,6 +164,149 @@ def __init__(
)
+# The map stores special registered buffer for measurement
+# This can be used for sparse workloads when we cannot use random tensors for measurment.
+# {
+# "workload_key_0": {
+# "task_input_0": Tensor(...),
+# "task_input_1": Tensor(...)
+# },
+# "workload_key_1": {
+# "task_input_2": Tensor(...),
+# "task_input_3": Tensor(...)
+# },
+# ...
+# }
+TASK_INPUT_BUFFER_TABLE = {}
+
+
+def _save_buffer_to_file(buffer_name, buffer_data):
+ """Save the current Tensor buffer to a numpy file.
+
+ File name will be: {buffer_name}.{buffer_shape}_{buffer_data_type}
+ """
+ np_data = buffer_data.asnumpy()
+
+ buffer_name += "."
+ for i in np_data.shape:
+ buffer_name += "%d_" % (i)
+ buffer_name += "%s" % (np_data.dtype)
+
+ np_data.tofile(buffer_name, " ")
+
+
+def _try_load_buffer_from_file(buffer_name):
+ """Try to load buffer from a numpy file, if not found, return None.
+
+ File name has a same format as `_save_buffer_to_file`.
+ """
+ filelist = os.listdir()
+
+ for file in filelist:
+ if file.startswith(buffer_name) and file.count("."):
+ meta_info = file.split(".")[-1].split("_")
+ shape = [int(i) for i in meta_info[:-1]]
+ dtype = meta_info[-1]
+ buffer_data = np.fromfile(file, dtype=dtype, sep=" ")
+ buffer_data = buffer_data.reshape(shape)
+ return ndarray.array(buffer_data)
+
+ return None
+
+
+def register_task_input_buffer(
+ workload_key,
+ input_name,
+ input_data,
+ overwrite=False,
+ save_to_file=False,
+):
+ """Register special buffer for measurement.
+
+ Parameters
+ ----------
+ workload_key : str
+ The workload key of the SearchTask.
+
+ input_name : str
+ The name of input buffer.
+
+ input_data : Tensor
Review comment:
```suggestion
input_data : tvm.nd.NDArray
```
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] jcf94 commented on a change in pull request #7313: [AutoSchedule] Sparse dense tuning support with custom sketch rule
Posted by GitBox <gi...@apache.org>.
jcf94 commented on a change in pull request #7313:
URL: https://github.com/apache/tvm/pull/7313#discussion_r586942876
##########
File path: tests/python/unittest/test_auto_scheduler_search_task.py
##########
@@ -0,0 +1,211 @@
+# 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 test_auto_scheduler_common import (
+ matmul_auto_scheduler_test,
+ zero_rank_compute_auto_scheduler_test,
+ zero_rank_reduce_auto_scheduler_test,
+)
+import multiprocessing
Review comment:
Good catch! Why should pylint/flake8 didn't find this bug .....
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] merrymercy commented on a change in pull request #7313: [AutoSchedule] Sparse dense tuning support with custom sketch rule
Posted by GitBox <gi...@apache.org>.
merrymercy commented on a change in pull request #7313:
URL: https://github.com/apache/tvm/pull/7313#discussion_r585920940
##########
File path: tutorials/auto_scheduler/tune_sparse_x86.py
##########
@@ -0,0 +1,293 @@
+# 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.
+"""
+Auto-scheduling Sparse Matrix Multiplication on CPU with Custom Sketch Rule
+===========================================================================
+**Author**: `Chengfan Jia <https://github.com/jcf94/>`_
+
+This is a tutorial on how to use the auto-scheduler to tune a sparse matrix multiplication for
+CPUs.
+
+Auto-scheduler is designed to explore the schedule with best performance for a given computation
+declaration automatically. While sometimes, we may have a demand to try some special ops which may
+not been well-supported by auto-scheduler's default sketch rules and result in poor performance.
+Fortunately, auto-scheduler currently allows user to provide a CustomSketch to cover these cases.
+
+We use sparse matrix multiplication as an example in this tutorial to demonstrate how to implement
+and plug a custom sketch rule to the auto-scheduler search policy.
Review comment:
```suggestion
and plug a custom sketch rule to the auto-scheduler's search policy.
```
##########
File path: tutorials/auto_scheduler/tune_sparse_x86.py
##########
@@ -0,0 +1,293 @@
+# 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.
+"""
+Auto-scheduling Sparse Matrix Multiplication on CPU with Custom Sketch Rule
+===========================================================================
+**Author**: `Chengfan Jia <https://github.com/jcf94/>`_
+
+This is a tutorial on how to use the auto-scheduler to tune a sparse matrix multiplication for
+CPUs.
+
+Auto-scheduler is designed to explore the schedule with best performance for a given computation
+declaration automatically. While sometimes, we may have a demand to try some special ops which may
+not been well-supported by auto-scheduler's default sketch rules and result in poor performance.
+Fortunately, auto-scheduler currently allows user to provide a CustomSketch to cover these cases.
+
+We use sparse matrix multiplication as an example in this tutorial to demonstrate how to implement
+and plug a custom sketch rule to the auto-scheduler search policy.
+
+Note that this tutorial will not run on Windows or recent versions of macOS. To
+get it to run, you will need to wrap the body of this tutorial in a :code:`if
+__name__ == "__main__":` block.
+"""
+
+import os
+import itertools
+
+import numpy as np
+import tvm
+from tvm import te, auto_scheduler, runtime, topi
+from tvm.auto_scheduler import _ffi_api
+from tvm.topi.utils import get_const_tuple
+
+import scipy.sparse as sp
+
+######################################################################
+# Define the computation
+# ^^^^^^^^^^^^^^^^^^^^^^
+# To begin with, let us define the computation of a sparse matmul with several relu and bias add.
+# The function should return the list of input/output tensors.
+# From these tensors, the auto-scheduler can get the whole computational graph.
+
+# We use this function to generate a random bsr matrix
+def random_bsr_matrix(M, N, BS_R, BS_C, density, dtype):
+ import itertools
+
+ Y = np.zeros((M, N), dtype=dtype)
+ assert M % BS_R == 0
+ assert N % BS_C == 0
+ nnz = int(density * M * N)
+ num_blocks = int(nnz / (BS_R * BS_C)) + 1
+ candidate_blocks = np.asarray(list(itertools.product(range(0, M, BS_R), range(0, N, BS_C))))
+ assert candidate_blocks.shape[0] == M // BS_R * N // BS_C
+ chosen_blocks = candidate_blocks[
+ np.random.choice(candidate_blocks.shape[0], size=num_blocks, replace=False)
+ ]
+ for i in range(len(chosen_blocks)):
+ r, c = chosen_blocks[i]
+ Y[r : r + BS_R, c : c + BS_C] = np.random.randn(BS_R, BS_C)
+ s = sp.bsr_matrix(Y, blocksize=(BS_R, BS_C))
+ assert s.data.shape == (num_blocks, BS_R, BS_C)
+ assert s.indices.shape == (num_blocks,)
+ assert s.indptr.shape == (M // BS_R + 1,)
+ return s
+
+
+@auto_scheduler.register_workload
+def sparse_dense(M, N, K, w_data_shape, w_indices_shape, w_indptr_shape, dtype):
+ X = te.placeholder(shape=(M, K), dtype=dtype)
+ W_data = te.placeholder(shape=w_data_shape, dtype=dtype)
+ W_indices = te.placeholder(shape=w_indices_shape, dtype="int32")
+ W_indptr = te.placeholder(shape=w_indptr_shape, dtype="int32")
+ B = te.placeholder(shape=(M, N), dtype=dtype)
+
+ out = topi.nn.sparse_dense(topi.nn.relu(X), W_data, W_indices, W_indptr)
+ out = te.compute((M, N), lambda i, j: out[i, j] + B[i, j], name="BiasAdd")
+ out = topi.nn.relu(out)
+
+ return [X, W_data, W_indices, W_indptr, B, out]
+
+
+######################################################################
+# Special step for sparse workload
+# ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+# During schedule tuning, auto-scheduler will use random inputs to measure the performance of a
+# generated schedule. While we cannot directly use a random array as the input of a sparse op, for
+# the "indices" and "indptr" array are meaningful for the computation.
+#
+# To solve this problem, we register these as special buffers, and load them when process program
+# measuring.
+# See the :any:`auto_scheduler.measure` code for more details.
Review comment:
```suggestion
# See the `tvm.auto_scheduler.measure.py` for more details.
```
The original hyperlink won't work. We can just use the filename directly.
##########
File path: src/auto_scheduler/feature.cc
##########
@@ -1468,7 +1468,7 @@ void GetPerStoreFeaturesFromMeasurePairs(const Array<MeasureInput>& inputs,
Array<te::Tensor> tensors = (*workload_key_to_tensors)(workload_key);
task = SearchTask(ComputeDAG(tensors), workload_key, inputs[i]->task->target,
inputs[i]->task->target_host, inputs[i]->task->hardware_params,
- inputs[i]->task->layout_rewrite_option);
+ inputs[i]->task->layout_rewrite_option, {});
Review comment:
```suggestion
inputs[i]->task->layout_rewrite_option, inputs[i]->task->task_inputs);
```
Should we use this?
##########
File path: python/tvm/topi/nn/sparse.py
##########
@@ -356,3 +359,107 @@ def sparse_dense_alter_layout(_attrs, _inputs, _tinfos, _out_type):
Unlike other TOPI functions, this function operates on both graph level and operator level.
"""
return None
+
+
+def try_get_sparse_input(args):
+ """Analise the input data from the given args.
Review comment:
```suggestion
"""Analyze the input data from the given args.
```
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] tkonolige commented on a change in pull request #7313: [AutoSchedule] Sparse dense tuning support with custom sketch rule
Posted by GitBox <gi...@apache.org>.
tkonolige commented on a change in pull request #7313:
URL: https://github.com/apache/tvm/pull/7313#discussion_r566970128
##########
File path: python/tvm/auto_scheduler/measure.py
##########
@@ -719,6 +722,87 @@ def local_builder_build(inputs, timeout, n_parallel, build_func="default", verbo
return results
+def _process_sparse_input(args):
+ sparse_prefix = sparse_data = sparse_indices = sparse_indptr = None
+
+ def _process_inputs(input_tensors, M, N, prefix_init):
+ nonlocal sparse_prefix
+ nonlocal sparse_data
+ nonlocal sparse_indices
+ nonlocal sparse_indptr
+
+ assert len(input_tensors) == 4
+ unsure_tensors = list(input_tensors)
+ # Get the Dense data
+ dense_data = None
+ for tensor in unsure_tensors:
+ if len(tensor.shape) == 2:
+ assert dense_data is None
+ dense_data = tensor
+ assert M == dense_data.shape[0]
+ K = dense_data.shape[1]
+ unsure_tensors.remove(dense_data)
+
+ # Get the Sparse data
+ sparse_data = None
+ for tensor in unsure_tensors:
+ if len(tensor.shape) == 3:
+ assert sparse_data is None
+ sparse_data = tensor
+ block_size, BS_R, BS_C = sparse_data.shape
+ unsure_tensors.remove(sparse_data)
+
+ # Get the Sparse indptr & indices
+ sparse_indices = None
+ for tensor in unsure_tensors:
+ assert len(tensor.shape) == 1
+ if tensor.shape[0] == block_size:
+ assert sparse_indices is None
+ sparse_indices = tensor
+ unsure_tensors.remove(sparse_indices)
+ assert len(unsure_tensors) == 1
+ sparse_indptr = unsure_tensors[0]
+
+ # Generate the sparse_prefix
+ density = 1.0
+ for i in sparse_data.shape:
+ density *= i
+ density /= (K * N)
+ density = density.value
+ sparse_prefix = "%s_%d_%d_%d_%d_%d_%.2f_" % (
Review comment:
You could hash the `indptr` and `indices` arrays as these determine the structure.
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] ANSHUMAN87 commented on pull request #7313: [AutoSchedule] Sparse dense tuning support with custom sketch rule
Posted by GitBox <gi...@apache.org>.
ANSHUMAN87 commented on pull request #7313:
URL: https://github.com/apache/tvm/pull/7313#issuecomment-763723572
Thanks @jcf94 for the PR!
May be once the PR is ready, it would be really great if you can share the stats of sparse_dense Op with and without Ansor.
Really excited to see those.
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] jcf94 commented on a change in pull request #7313: [AutoSchedule] Sparse dense tuning support with custom sketch rule
Posted by GitBox <gi...@apache.org>.
jcf94 commented on a change in pull request #7313:
URL: https://github.com/apache/tvm/pull/7313#discussion_r565944345
##########
File path: tutorials/auto_scheduler/tune_sparse_x86.py
##########
@@ -0,0 +1,331 @@
+# 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.
+"""
+Auto-scheduling Sparse Matrix Multiplication for CPU by Custom Sketch Rule
+==========================================================================
+**Author**: `Chengfan Jia <https://github.com/jcf94/>`_
+
+This is a tutorial on how to use the auto-scheduler to tune a sparse matrix multiplication for
+CPUs.
+
+Auto-scheduler is designed to explore the schedule with best performance for a given computation
+declaration automatically. While sometimes, we may have a demand to try some special ops which may
+not been well supported by auto-scheduler's default search policy. Auto-scheduler currently allows
+user to provide a CustomSketch to cover these cases.
+
+We use sparse matrix multiplication as an example in this tutorial.
+
+Note that this tutorial will not run on Windows or recent versions of macOS. To
+get it to run, you will need to wrap the body of this tutorial in a :code:`if
+__name__ == "__main__":` block.
+"""
+
+import os
+import itertools
+
+import numpy as np
+import tvm
+from tvm import te, auto_scheduler, topi
+from tvm.auto_scheduler import _ffi_api
+from tvm.topi.utils import get_const_tuple
+
+import scipy.sparse as sp
+
+######################################################################
+# Define the computation
+# ^^^^^^^^^^^^^^^^^^^^^^
+# To begin with, let us define the computation of a sparse matmul with several relu and bias add.
+# The function should return the list of input/output tensors.
+# From these tensors, the auto-scheduler can get the whole computational graph.
+
+# We use this function to generate a random bsr matrix
+def random_bsr_matrix(M, N, BS_R, BS_C, density, dtype):
+ import itertools
+
+ Y = np.zeros((M, N), dtype=dtype)
+ assert M % BS_R == 0
+ assert N % BS_C == 0
+ nnz = int(density * M * N)
+ num_blocks = int(nnz / (BS_R * BS_C)) + 1
+ candidate_blocks = np.asarray(list(itertools.product(range(0, M, BS_R), range(0, N, BS_C))))
+ assert candidate_blocks.shape[0] == M // BS_R * N // BS_C
+ chosen_blocks = candidate_blocks[
+ np.random.choice(candidate_blocks.shape[0], size=num_blocks, replace=False)
+ ]
+ for i in range(len(chosen_blocks)):
+ r, c = chosen_blocks[i]
+ Y[r : r + BS_R, c : c + BS_C] = np.random.randn(BS_R, BS_C)
+ s = sp.bsr_matrix(Y, blocksize=(BS_R, BS_C))
+ assert s.data.shape == (num_blocks, BS_R, BS_C)
+ assert s.indices.shape == (num_blocks,)
+ assert s.indptr.shape == (M // BS_R + 1,)
+ return s
+
+@auto_scheduler.register_workload
+def sparse_dense(M, N, K, w_data_shape, w_indices_shape, w_indptr_shape, dtype):
+ X = te.placeholder(shape=(M, K), dtype=dtype)
+ W_data = te.placeholder(shape=w_data_shape, dtype=dtype)
+ W_indices = te.placeholder(shape=w_indices_shape, dtype="int32")
+ W_indptr = te.placeholder(shape=w_indptr_shape, dtype="int32")
+ B = te.placeholder(shape=(M, N), dtype=dtype)
+
+ out = topi.nn.sparse_dense(
+ topi.nn.relu(X), W_data, W_indices, W_indptr
+ )
+ out = te.compute((M, N), lambda i, j: out[i, j] + B[i, j], name="BiasAdd")
+ out = topi.nn.relu(out)
+
+ return [X, W_data, W_indices, W_indptr, B, out]
+
+######################################################################
+# Special step for sparse workload
+# ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+# During schedule tuning, auto-scheduler will use random inputs to measure the performance of a
+# generated schedule. While we cannot directly use a random array as the input of a sparse op, for
+# the "indices" and "indptr" array are meaningful for the computation.
+#
+# To solve this problem, we register these as special buffers, and load them when process program
+# measuring.
+# See the :any:`auto_scheduler.measure` code for more details.
Review comment:
Good idea! I'm now trying adding a `add_measure_input` API to SearchTask.
Then we will need to serialize/unserialize the added data buffer ... Is it fine to do so?
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] jcf94 commented on a change in pull request #7313: [AutoSchedule] Sparse dense tuning support with custom sketch rule
Posted by GitBox <gi...@apache.org>.
jcf94 commented on a change in pull request #7313:
URL: https://github.com/apache/tvm/pull/7313#discussion_r586937376
##########
File path: python/tvm/auto_scheduler/measure.py
##########
@@ -719,6 +720,45 @@ def local_builder_build(inputs, timeout, n_parallel, build_func="default", verbo
return results
+def _prepare_input_map(args):
+ """This function deals with special task inputs.
+
+ Parameters
+ ----------
+ args : List[Tensor]
+ Input/output Tensor of a TVM subgraph.
+
+ Returns
+ -------
+ A Dict[Tensor, str] that maps the input Tensor to a buffer name.
+
+ Note
+ ----
+ The buffer name is specially designed, and these buffer should be provided in
+ `SearchTask(..., task_inputs={...})`.
+ """
+ # pylint: disable=import-outside-toplevel
+ from tvm import topi # lazily import to avoid recursive dependency
+
+ # A dict that maps the input tensor arg to a buffer name
+ tensor_input_map = {}
+
+ # Case 0: Check placeholder name
+ for arg in args:
+ if isinstance(arg.op, tvm.te.PlaceholderOp):
+ if arg.op.name != "placeholder":
+ tensor_input_map[arg] = arg.op.name
+
+ # Case 1: Check sparse op
+ sparse_input_map = topi.nn.sparse.try_get_sparse_input(args)
Review comment:
Yeah, we've also had some discussions while didn't figure out any better solutions.
There're several reasons:
1. Different ops have different requirements over specific inputs;
2. While the problems is in a subgraph generated in Relay Integration, the placeholder are all the same, we can not differentiate them from tag, name or any other way, even the order of inputs are not guaranteed.
Currently approach is to merge all specific inputs checking to this function, at least they have a same entry here. For the other ops, you have to add their own check functions below.
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] merrymercy edited a comment on pull request #7313: [AutoSchedule] Sparse dense tuning support with custom sketch rule
Posted by GitBox <gi...@apache.org>.
merrymercy edited a comment on pull request #7313:
URL: https://github.com/apache/tvm/pull/7313#issuecomment-789393341
According to our offline discussion,
1. Update the type of `SearchTaskNode::task_inputs`. Change it from `Map<String, runtime::NDArray>` to `Array<String>`, so we only need to store nd arrays in one place. We can query it from the global table in `measure.py`
2. Remove `SearchTask.AddTaskInput` interface to make `SearchTask` immutable. We do not have the need to dynamically update task inputs, so we can provide all arguments to the constructors.
3. Make sure we can use the same interface to support the use case where we want to match the special buffers by name
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] jcf94 commented on pull request #7313: [AutoSchedule] Sparse dense tuning support with custom sketch rule
Posted by GitBox <gi...@apache.org>.
jcf94 commented on pull request #7313:
URL: https://github.com/apache/tvm/pull/7313#issuecomment-791121946
@merrymercy @comaniac @tkonolige Thanks! Comments has all been addressed.
Additionally, add a `@auto_scheduler.register_task_input_check_func`, now we can add extra input check functions more easily.
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] antinucleon commented on a change in pull request #7313: [AutoSchedule] Sparse dense tuning support with custom sketch rule
Posted by GitBox <gi...@apache.org>.
antinucleon commented on a change in pull request #7313:
URL: https://github.com/apache/tvm/pull/7313#discussion_r560760923
##########
File path: python/tvm/auto_scheduler/measure.py
##########
@@ -1132,3 +1249,44 @@ def rpc_runner_run(
print("")
return results
+
+
+# The map stores special registered buffer for measurement
+# This can be used for sparse workloads when we cannot use random tensors for measurment.
+global special_buffer_table
+special_buffer_table = {}
Review comment:
I am against to use global variable. In my personal [test branch](https://github.com/antinucleon/tvm/tree/metal), I use pickle to dump all buffers to `tmpdir/workload_key.pkl` to avoid global variable. In setting with ARM64 Python + Multiprocessing, global variable behavior is not same to what we expect on Linux.
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] jcf94 commented on a change in pull request #7313: [AutoSchedule] Sparse dense tuning support with custom sketch rule
Posted by GitBox <gi...@apache.org>.
jcf94 commented on a change in pull request #7313:
URL: https://github.com/apache/tvm/pull/7313#discussion_r586167940
##########
File path: tutorials/auto_scheduler/tune_sparse_x86.py
##########
@@ -0,0 +1,331 @@
+# 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.
+"""
+Auto-scheduling Sparse Matrix Multiplication for CPU by Custom Sketch Rule
+==========================================================================
+**Author**: `Chengfan Jia <https://github.com/jcf94/>`_
+
+This is a tutorial on how to use the auto-scheduler to tune a sparse matrix multiplication for
+CPUs.
+
+Auto-scheduler is designed to explore the schedule with best performance for a given computation
+declaration automatically. While sometimes, we may have a demand to try some special ops which may
+not been well supported by auto-scheduler's default search policy. Auto-scheduler currently allows
+user to provide a CustomSketch to cover these cases.
+
+We use sparse matrix multiplication as an example in this tutorial.
+
+Note that this tutorial will not run on Windows or recent versions of macOS. To
+get it to run, you will need to wrap the body of this tutorial in a :code:`if
+__name__ == "__main__":` block.
+"""
+
+import os
+import itertools
+
+import numpy as np
+import tvm
+from tvm import te, auto_scheduler, topi
+from tvm.auto_scheduler import _ffi_api
+from tvm.topi.utils import get_const_tuple
+
+import scipy.sparse as sp
+
+######################################################################
+# Define the computation
+# ^^^^^^^^^^^^^^^^^^^^^^
+# To begin with, let us define the computation of a sparse matmul with several relu and bias add.
+# The function should return the list of input/output tensors.
+# From these tensors, the auto-scheduler can get the whole computational graph.
+
+# We use this function to generate a random bsr matrix
+def random_bsr_matrix(M, N, BS_R, BS_C, density, dtype):
+ import itertools
+
+ Y = np.zeros((M, N), dtype=dtype)
+ assert M % BS_R == 0
+ assert N % BS_C == 0
+ nnz = int(density * M * N)
+ num_blocks = int(nnz / (BS_R * BS_C)) + 1
+ candidate_blocks = np.asarray(list(itertools.product(range(0, M, BS_R), range(0, N, BS_C))))
+ assert candidate_blocks.shape[0] == M // BS_R * N // BS_C
+ chosen_blocks = candidate_blocks[
+ np.random.choice(candidate_blocks.shape[0], size=num_blocks, replace=False)
+ ]
+ for i in range(len(chosen_blocks)):
+ r, c = chosen_blocks[i]
+ Y[r : r + BS_R, c : c + BS_C] = np.random.randn(BS_R, BS_C)
+ s = sp.bsr_matrix(Y, blocksize=(BS_R, BS_C))
+ assert s.data.shape == (num_blocks, BS_R, BS_C)
+ assert s.indices.shape == (num_blocks,)
+ assert s.indptr.shape == (M // BS_R + 1,)
+ return s
+
+@auto_scheduler.register_workload
+def sparse_dense(M, N, K, w_data_shape, w_indices_shape, w_indptr_shape, dtype):
+ X = te.placeholder(shape=(M, K), dtype=dtype)
+ W_data = te.placeholder(shape=w_data_shape, dtype=dtype)
+ W_indices = te.placeholder(shape=w_indices_shape, dtype="int32")
+ W_indptr = te.placeholder(shape=w_indptr_shape, dtype="int32")
+ B = te.placeholder(shape=(M, N), dtype=dtype)
+
+ out = topi.nn.sparse_dense(
+ topi.nn.relu(X), W_data, W_indices, W_indptr
+ )
+ out = te.compute((M, N), lambda i, j: out[i, j] + B[i, j], name="BiasAdd")
+ out = topi.nn.relu(out)
+
+ return [X, W_data, W_indices, W_indptr, B, out]
+
+######################################################################
+# Special step for sparse workload
+# ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+# During schedule tuning, auto-scheduler will use random inputs to measure the performance of a
+# generated schedule. While we cannot directly use a random array as the input of a sparse op, for
+# the "indices" and "indptr" array are meaningful for the computation.
+#
+# To solve this problem, we register these as special buffers, and load them when process program
+# measuring.
+# See the :any:`auto_scheduler.measure` code for more details.
+
+# Define the basic shapes of this sparse computation
+M = K = N = 512
+BS_R = 16
+BS_C = 1
+density = 0.6
+
+# Generate the test data with numpy
+X_np = np.random.randn(M, K).astype("float32")
+X_np = np.maximum(np.zeros((M, K), dtype="float32"), X_np) # Relu
+W_sp_np = random_bsr_matrix(N, K, BS_R, BS_C, density=density, dtype="float32")
+W_np = W_sp_np.todense()
+Y_np = X_np @ W_np.T # Process the matrix multiplication
+B_np = np.random.randn(M, N).astype("float32")
+Y_np = Y_np + B_np # Bias add
+Y_np = np.maximum(np.zeros((M, N), dtype="float32"), Y_np) # Relu
+
+# Register the sparse data to special buffer
+prefix = "sparse_dense_bsr_%d_%d_%d_%d_%d_%.2f_" % (M, N, K, BS_R, BS_C, density)
+auto_scheduler.measure.register_special_buffer(prefix + "W_data", W_sp_np.data)
+auto_scheduler.measure.register_special_buffer(prefix + "W_indices", W_sp_np.indices)
+auto_scheduler.measure.register_special_buffer(prefix + "W_indptr", W_sp_np.indptr)
+
+######################################################################
+# Create the search task
+# ^^^^^^^^^^^^^^^^^^^^^^
+# We then create a search task with M=N=K=512 and dtype="float32"
+# If your machine supports avx instructions, you can
+#
+# - replace "llvm" below with "llvm -mcpu=core-avx2" to enable AVX2
+# - replace "llvm" below with "llvm -mcpu=skylake-avx512" to enable AVX-512
+
+target = tvm.target.Target("llvm")
+
+task = tvm.auto_scheduler.SearchTask(
+ func=sparse_dense,
+ args=(
+ M, N, K,
+ W_sp_np.data.shape,
+ W_sp_np.indices.shape,
+ W_sp_np.indptr.shape,
+ "float32"
+ ),
+ target=target
+)
+
+# Inspect the computational graph
+print("Computational DAG:")
+print(task.compute_dag)
+
+######################################################################
+# Write the custom sketch for sparse dense op
+# ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+# Before tuning, we will need to define the CustomSketchRule for the sparse dense op.
+#
+# CustomSketchRule consists of two parts: the condition function and the apply function.
+#
+# - condition function: describe when to use this sketch rule. For example, we can match the op
+# by their name or tag.
+# - apply function: describe how to generate the initial sketch. Auto-scheduler provides a set of
+# loop state APIs.
+
+def meet_condition_func(search_policy, state, stage_id):
+ state = auto_scheduler.loop_state.State(state, search_policy.search_task.compute_dag)
+ if state.stages[stage_id].op.tag in [
+ "sparse_dense_sp_rhs_bsrmm", "sparse_dense_sp_rhs_bsrmm_block"
+ ]:
+ return auto_scheduler.PreloadCustomSketchRule.APPLY_AND_SKIP_REST
+ else:
+ return auto_scheduler.PreloadCustomSketchRule.PASS
+
+def apply_func(search_policy, state, stage_id):
+ ret = []
+ s0 = auto_scheduler.loop_state.State(state, search_policy.search_task.compute_dag)
+ if s0.stages[stage_id].op.tag == "sparse_dense_sp_rhs_bsrmm_block":
+ return [s0.state_object, stage_id - 1]
+
+ sparse_dense = s0.stages[stage_id].op
+ sparse_dense_block = s0.stages[stage_id - 1].op
+ assert sparse_dense.tag == "sparse_dense_sp_rhs_bsrmm"
+ assert sparse_dense_block.tag == "sparse_dense_sp_rhs_bsrmm_block"
+
+ # Set the default consumer of compute block
+ consumer = sparse_dense
+
+ # If sparse dense has a single elementwise consumer
+ # We can compute inline the sparse_dense output stage
+ consumers = _ffi_api.SearchPolicyUtilsGetConsumers(
+ search_policy.search_task, s0.state_object, stage_id
+ )
+ if len(consumers) == 1:
+ consumer_id = int(consumers.items()[0][0])
+ if _ffi_api.SearchPolicyUtilsIsElementwiseMatch(
+ search_policy.search_task, s0.state_object, stage_id, consumer_id
+ ):
+ consumer = s0.stages[consumer_id].op
+ s0.compute_inline(sparse_dense)
+
+ i, nb_j, j, row_offset, c = s0[sparse_dense_block].iters
+ m, n = s0[consumer].iters
+ i0, i1, i2 = s0.split(sparse_dense_block, i, [None, None])
+ m0, m1 = s0.follow_split(consumer, m, len(s0.transform_steps) - 1, 1)
+ j0, j1 = s0.split(sparse_dense_block, nb_j, [None])
+ n0, n1 = s0.follow_split(consumer, n, len(s0.transform_steps) - 1, 1)
+ s0.reorder(sparse_dense_block, [i0, j0, i1, j1, row_offset, i2, j, c])
+ s0.reorder(consumer, [m0, n0, m1, n1])
+ s0.compute_at(sparse_dense_block, consumer, n0)
+
+ ret.append([s0.state_object, stage_id - 2])
+
+ return ret
+
+######################################################################
+# Next, we set parameters for the auto-scheduler.
+#
+# * :code:`num_measure_trials` is the number of measurement trials we can use during the search.
+# We only make 10 trials in this tutorial for a fast demonstration. In practice, 1000 is a
+# good value for the search to converge. You can do more trials according to your time budget.
+# * In addition, we use :code:`RecordToFile` to dump measurement records into a file `matmul.json`.
+# The measurement records can be used to query the history best, resume the search,
+# and do more analyses later.
+# * see :any:`auto_scheduler.TuningOptions` for more parameters
+# * Here, we need to create a :code:`auto_scheduler.SketchPolicy` object, and add the custom sketch
+# rule as a `init_search_callbacks`.
+
+log_file = "sparse_dense.json"
+tune_option = auto_scheduler.TuningOptions(
+ num_measure_trials=10,
+ measure_callbacks=[auto_scheduler.RecordToFile(log_file)],
+ verbose=2,
+)
+
+search_policy = auto_scheduler.SketchPolicy(
+ task,
+ program_cost_model=auto_scheduler.XGBModel(),
+ init_search_callbacks=[
+ auto_scheduler.PreloadCustomSketchRule(meet_condition_func, apply_func, "SparseDense")
+ ]
+)
+
+######################################################################
+# Run the search
+# ^^^^^^^^^^^^^^
+# Now we get all inputs ready.
+# We can kick off the search and let the auto-scheduler do its magic.
+# After some measurement trials, we can load the best schedule from the log
+# file and apply it.
+
+# Run auto-tuning (search)
+task.tune(tune_option, search_policy)
Review comment:
This task is very small ... Almost cost no time for only measure 2 states.
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] merrymercy edited a comment on pull request #7313: [AutoSchedule] Sparse dense tuning support with custom sketch rule
Posted by GitBox <gi...@apache.org>.
merrymercy edited a comment on pull request #7313:
URL: https://github.com/apache/tvm/pull/7313#issuecomment-789393341
1. Update the type of `SearchTaskNode::task_inputs`. Change it from `Map<String, runtime::NDArray>` to `Array<String>`, so we only need to store nd arrays in one place. We can query it from the global table in `measure.py`
2. Remove `SearchTask.AddTaskInput` interface to make `SearchTask` immutable. We do not have the need to dynamically update task inputs, so we can provide all arguments to the constructors.
3. Make sure we can use the same interface to support the use case where we want to match the special buffers by name
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] jcf94 commented on a change in pull request #7313: [AutoSchedule] Sparse dense tuning support with custom sketch rule
Posted by GitBox <gi...@apache.org>.
jcf94 commented on a change in pull request #7313:
URL: https://github.com/apache/tvm/pull/7313#discussion_r587996929
##########
File path: python/tvm/auto_scheduler/measure.py
##########
@@ -719,6 +720,45 @@ def local_builder_build(inputs, timeout, n_parallel, build_func="default", verbo
return results
+def _prepare_input_map(args):
+ """This function deals with special task inputs.
+
+ Parameters
+ ----------
+ args : List[Tensor]
+ Input/output Tensor of a TVM subgraph.
+
+ Returns
+ -------
+ A Dict[Tensor, str] that maps the input Tensor to a buffer name.
+
+ Note
+ ----
+ The buffer name is specially designed, and these buffer should be provided in
+ `SearchTask(..., task_inputs={...})`.
+ """
+ # pylint: disable=import-outside-toplevel
+ from tvm import topi # lazily import to avoid recursive dependency
+
+ # A dict that maps the input tensor arg to a buffer name
+ tensor_input_map = {}
+
+ # Case 0: Check placeholder name
+ for arg in args:
+ if isinstance(arg.op, tvm.te.PlaceholderOp):
+ if arg.op.name != "placeholder":
+ tensor_input_map[arg] = arg.op.name
+
+ # Case 1: Check sparse op
+ sparse_input_map = topi.nn.sparse.try_get_sparse_input(args)
Review comment:
> Could we associate the lookup mechanism with `@register_workload`? It would at least be extensible then.
Thanks! This is a pretty good idea, I'll have a try.
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] tkonolige commented on a change in pull request #7313: [AutoSchedule] Sparse dense tuning support with custom sketch rule
Posted by GitBox <gi...@apache.org>.
tkonolige commented on a change in pull request #7313:
URL: https://github.com/apache/tvm/pull/7313#discussion_r566479973
##########
File path: python/tvm/auto_scheduler/measure.py
##########
@@ -943,18 +1047,36 @@ def _timed_rpc_run(
if error_no == 0:
try:
- args = [ndarray.empty(get_const_tuple(x.shape), x.dtype, ctx) for x in build_res.args]
try:
random_fill = remote.get_function("tvm.contrib.random.random_fill")
except AttributeError:
raise AttributeError(
"Please make sure USE_RANDOM is ON in the config.cmake " "on the remote devices"
)
- for arg in args:
- random_fill(arg)
- ctx.sync()
+ # Check sparse op
+ sparse_prefix, sparse_data, sparse_indices, sparse_indptr = \
+ _process_sparse_input(build_res.args)
+ if sparse_prefix:
+ args = []
+ for arg in build_res.args:
+ if arg == sparse_data:
+ args.append(ndarray.array(get_special_buffer(sparse_prefix+"W_data"), ctx))
+ elif arg == sparse_indices:
+ args.append(ndarray.array(get_special_buffer(sparse_prefix+"W_indices"), ctx))
+ elif arg == sparse_indptr:
+ args.append(ndarray.array(get_special_buffer(sparse_prefix+"W_indptr"), ctx))
+ else:
+ empty_array = ndarray.empty(get_const_tuple(arg.shape), arg.dtype, ctx)
+ random_fill(empty_array)
+ args.append(empty_array)
+ else:
+ args = [ndarray.empty(get_const_tuple(x.shape), x.dtype, ctx) for x in build_res.args]
+ for arg in args:
+ random_fill(arg)
+ ctx.sync()
Review comment:
This code is duplicated above. Maybe we can unify it.
##########
File path: python/tvm/auto_scheduler/measure.py
##########
@@ -943,18 +1047,36 @@ def _timed_rpc_run(
if error_no == 0:
try:
- args = [ndarray.empty(get_const_tuple(x.shape), x.dtype, ctx) for x in build_res.args]
try:
random_fill = remote.get_function("tvm.contrib.random.random_fill")
except AttributeError:
raise AttributeError(
"Please make sure USE_RANDOM is ON in the config.cmake " "on the remote devices"
)
- for arg in args:
- random_fill(arg)
- ctx.sync()
+ # Check sparse op
+ sparse_prefix, sparse_data, sparse_indices, sparse_indptr = \
+ _process_sparse_input(build_res.args)
+ if sparse_prefix:
Review comment:
We are special casing for sparse here, but there are other situations where we need non-random input data. For example, sorting and scatter. Maybe we could make a more general approach?
##########
File path: python/tvm/auto_scheduler/measure.py
##########
@@ -719,6 +722,87 @@ def local_builder_build(inputs, timeout, n_parallel, build_func="default", verbo
return results
+def _process_sparse_input(args):
+ sparse_prefix = sparse_data = sparse_indices = sparse_indptr = None
+
+ def _process_inputs(input_tensors, M, N, prefix_init):
+ nonlocal sparse_prefix
+ nonlocal sparse_data
+ nonlocal sparse_indices
+ nonlocal sparse_indptr
+
+ assert len(input_tensors) == 4
+ unsure_tensors = list(input_tensors)
+ # Get the Dense data
+ dense_data = None
+ for tensor in unsure_tensors:
+ if len(tensor.shape) == 2:
+ assert dense_data is None
+ dense_data = tensor
+ assert M == dense_data.shape[0]
+ K = dense_data.shape[1]
+ unsure_tensors.remove(dense_data)
+
+ # Get the Sparse data
+ sparse_data = None
+ for tensor in unsure_tensors:
+ if len(tensor.shape) == 3:
+ assert sparse_data is None
+ sparse_data = tensor
+ block_size, BS_R, BS_C = sparse_data.shape
+ unsure_tensors.remove(sparse_data)
+
+ # Get the Sparse indptr & indices
+ sparse_indices = None
+ for tensor in unsure_tensors:
+ assert len(tensor.shape) == 1
+ if tensor.shape[0] == block_size:
+ assert sparse_indices is None
+ sparse_indices = tensor
+ unsure_tensors.remove(sparse_indices)
+ assert len(unsure_tensors) == 1
+ sparse_indptr = unsure_tensors[0]
+
+ # Generate the sparse_prefix
+ density = 1.0
+ for i in sparse_data.shape:
+ density *= i
+ density /= (K * N)
+ density = density.value
+ sparse_prefix = "%s_%d_%d_%d_%d_%d_%.2f_" % (
Review comment:
We could run into the case that two matrices have the same `sparse_prefix`, but different non-zero structure. Will this cause issues? What if one of the matrices has one nonzero per row and the other has one dense row (while maintaining the same sparsity)?
##########
File path: python/tvm/auto_scheduler/measure.py
##########
@@ -1132,3 +1254,44 @@ def rpc_runner_run(
print("")
return results
+
+
+# The map stores special registered buffer for measurement
+# This can be used for sparse workloads when we cannot use random tensors for measurment.
+global special_buffer_table
+special_buffer_table = {}
+
+def register_special_buffer(tensor_name, data):
+ """Register special buffer for measurement
+ This can be used for sparse workloads when we cannot use random tensors for measurment.
+ """
+ if tensor_name in special_buffer_table.keys():
+ return True
+
+ if os.path.isfile(tensor_name):
+ print("Load ", tensor_name)
+ if tensor_name.startswith("sparse_dense_bsr"):
+ if tensor_name.endswith("data"):
+ data = np.fromfile(tensor_name, dtype="float32", sep=" ")
Review comment:
This hard-coding of dtype seems like it will fail if the sparse matrix does not use float32
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] jcf94 commented on a change in pull request #7313: [AutoSchedule] Sparse dense tuning support with custom sketch rule
Posted by GitBox <gi...@apache.org>.
jcf94 commented on a change in pull request #7313:
URL: https://github.com/apache/tvm/pull/7313#discussion_r587971202
##########
File path: python/tvm/auto_scheduler/search_task.py
##########
@@ -157,6 +164,149 @@ def __init__(
)
+# The map stores special registered buffer for measurement
+# This can be used for sparse workloads when we cannot use random tensors for measurment.
+# {
+# "workload_key_0": {
+# "task_input_0": Tensor(...),
+# "task_input_1": Tensor(...)
+# },
+# "workload_key_1": {
+# "task_input_2": Tensor(...),
+# "task_input_3": Tensor(...)
+# },
+# ...
+# }
+TASK_INPUT_BUFFER_TABLE = {}
+
+
+def _save_buffer_to_file(buffer_name, buffer_data):
+ """Save the current Tensor buffer to a numpy file.
+
+ File name will be: {buffer_name}.{buffer_shape}_{buffer_data_type}
+ """
+ np_data = buffer_data.asnumpy()
+
+ buffer_name += "."
+ for i in np_data.shape:
+ buffer_name += "%d_" % (i)
+ buffer_name += "%s" % (np_data.dtype)
+
+ np_data.tofile(buffer_name, " ")
+
+
+def _try_load_buffer_from_file(buffer_name):
+ """Try to load buffer from a numpy file, if not found, return None.
+
+ File name has a same format as `_save_buffer_to_file`.
+ """
+ filelist = os.listdir()
+
+ for file in filelist:
+ if file.startswith(buffer_name) and file.count("."):
+ meta_info = file.split(".")[-1].split("_")
+ shape = [int(i) for i in meta_info[:-1]]
+ dtype = meta_info[-1]
+ buffer_data = np.fromfile(file, dtype=dtype, sep=" ")
+ buffer_data = buffer_data.reshape(shape)
+ return ndarray.array(buffer_data)
+
+ return None
+
+
+def register_task_input_buffer(
+ workload_key,
+ input_name,
+ input_data,
+ overwrite=False,
+ save_to_file=False,
+):
+ """Register special buffer for measurement.
+
+ Parameters
+ ----------
+ workload_key : str
+ The workload key of the SearchTask.
+
+ input_name : str
+ The name of input buffer.
+
+ input_data : tvm.nd.NDArray
+ The input Tensor data.
+
+ overwrite : bool = False
+ Whether overwrite the data if a name has already in the global table.
+
+ save_to_file : bool = False
+ Whether record this buffer to a local file. This can be reused to continue the last tuning
+ process.
+ """
+ global TASK_INPUT_BUFFER_TABLE
+
+ if workload_key not in TASK_INPUT_BUFFER_TABLE:
+ TASK_INPUT_BUFFER_TABLE[workload_key] = {}
+ input_table = TASK_INPUT_BUFFER_TABLE[workload_key]
+
+ if not overwrite:
+ if input_name not in input_table.keys():
+ # Try to load buffer data from local file
+ tensor_from_file = _try_load_buffer_from_file(input_name)
+ if tensor_from_file:
+ input_table[input_name] = tensor_from_file
+
+ if input_name in input_table.keys():
Review comment:
I actually have thought about this, but if we print at each measure the output information will be a mess.
Or we can raise a warning if the inputs are missing here.
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] comaniac commented on a change in pull request #7313: [AutoSchedule] Sparse dense tuning support with custom sketch rule
Posted by GitBox <gi...@apache.org>.
comaniac commented on a change in pull request #7313:
URL: https://github.com/apache/tvm/pull/7313#discussion_r587973222
##########
File path: python/tvm/auto_scheduler/search_task.py
##########
@@ -157,6 +164,149 @@ def __init__(
)
+# The map stores special registered buffer for measurement
+# This can be used for sparse workloads when we cannot use random tensors for measurment.
+# {
+# "workload_key_0": {
+# "task_input_0": Tensor(...),
+# "task_input_1": Tensor(...)
+# },
+# "workload_key_1": {
+# "task_input_2": Tensor(...),
+# "task_input_3": Tensor(...)
+# },
+# ...
+# }
+TASK_INPUT_BUFFER_TABLE = {}
+
+
+def _save_buffer_to_file(buffer_name, buffer_data):
+ """Save the current Tensor buffer to a numpy file.
+
+ File name will be: {buffer_name}.{buffer_shape}_{buffer_data_type}
+ """
+ np_data = buffer_data.asnumpy()
+
+ buffer_name += "."
+ for i in np_data.shape:
+ buffer_name += "%d_" % (i)
+ buffer_name += "%s" % (np_data.dtype)
+
+ np_data.tofile(buffer_name, " ")
+
+
+def _try_load_buffer_from_file(buffer_name):
+ """Try to load buffer from a numpy file, if not found, return None.
+
+ File name has a same format as `_save_buffer_to_file`.
+ """
+ filelist = os.listdir()
+
+ for file in filelist:
+ if file.startswith(buffer_name) and file.count("."):
+ meta_info = file.split(".")[-1].split("_")
+ shape = [int(i) for i in meta_info[:-1]]
+ dtype = meta_info[-1]
+ buffer_data = np.fromfile(file, dtype=dtype, sep=" ")
+ buffer_data = buffer_data.reshape(shape)
+ return ndarray.array(buffer_data)
+
+ return None
+
+
+def register_task_input_buffer(
+ workload_key,
+ input_name,
+ input_data,
+ overwrite=False,
+ save_to_file=False,
+):
+ """Register special buffer for measurement.
+
+ Parameters
+ ----------
+ workload_key : str
+ The workload key of the SearchTask.
+
+ input_name : str
+ The name of input buffer.
+
+ input_data : tvm.nd.NDArray
+ The input Tensor data.
+
+ overwrite : bool = False
+ Whether overwrite the data if a name has already in the global table.
+
+ save_to_file : bool = False
+ Whether record this buffer to a local file. This can be reused to continue the last tuning
+ process.
+ """
+ global TASK_INPUT_BUFFER_TABLE
+
+ if workload_key not in TASK_INPUT_BUFFER_TABLE:
+ TASK_INPUT_BUFFER_TABLE[workload_key] = {}
+ input_table = TASK_INPUT_BUFFER_TABLE[workload_key]
+
+ if not overwrite:
+ if input_name not in input_table.keys():
+ # Try to load buffer data from local file
+ tensor_from_file = _try_load_buffer_from_file(input_name)
+ if tensor_from_file:
+ input_table[input_name] = tensor_from_file
+
+ if input_name in input_table.keys():
Review comment:
I suppose missing inputs are more common, as most use cases still rely on random inputs?
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] merrymercy merged pull request #7313: [AutoSchedule] Sparse dense tuning support with custom sketch rule
Posted by GitBox <gi...@apache.org>.
merrymercy merged pull request #7313:
URL: https://github.com/apache/tvm/pull/7313
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] comaniac commented on a change in pull request #7313: [AutoSchedule] Sparse dense tuning support with custom sketch rule
Posted by GitBox <gi...@apache.org>.
comaniac commented on a change in pull request #7313:
URL: https://github.com/apache/tvm/pull/7313#discussion_r587963714
##########
File path: python/tvm/auto_scheduler/search_task.py
##########
@@ -157,6 +164,149 @@ def __init__(
)
+# The map stores special registered buffer for measurement
+# This can be used for sparse workloads when we cannot use random tensors for measurment.
+# {
+# "workload_key_0": {
+# "task_input_0": Tensor(...),
+# "task_input_1": Tensor(...)
+# },
+# "workload_key_1": {
+# "task_input_2": Tensor(...),
+# "task_input_3": Tensor(...)
+# },
+# ...
+# }
+TASK_INPUT_BUFFER_TABLE = {}
+
+
+def _save_buffer_to_file(buffer_name, buffer_data):
+ """Save the current Tensor buffer to a numpy file.
+
+ File name will be: {buffer_name}.{buffer_shape}_{buffer_data_type}
Review comment:
Ah I see. Yeah then we can go with the one you proposed.
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] comaniac commented on a change in pull request #7313: [AutoSchedule] Sparse dense tuning support with custom sketch rule
Posted by GitBox <gi...@apache.org>.
comaniac commented on a change in pull request #7313:
URL: https://github.com/apache/tvm/pull/7313#discussion_r561459545
##########
File path: tutorials/auto_scheduler/tune_sparse_x86.py
##########
@@ -0,0 +1,331 @@
+# 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.
+"""
+Auto-scheduling Sparse Matrix Multiplication for CPU by Custom Sketch Rule
+==========================================================================
Review comment:
```suggestion
Auto-scheduling Sparse Matrix Multiplication on CPU with Custom Sketch Rule
===========================================================================
```
##########
File path: tutorials/auto_scheduler/tune_sparse_x86.py
##########
@@ -0,0 +1,331 @@
+# 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.
+"""
+Auto-scheduling Sparse Matrix Multiplication for CPU by Custom Sketch Rule
+==========================================================================
+**Author**: `Chengfan Jia <https://github.com/jcf94/>`_
+
+This is a tutorial on how to use the auto-scheduler to tune a sparse matrix multiplication for
+CPUs.
+
+Auto-scheduler is designed to explore the schedule with best performance for a given computation
+declaration automatically. While sometimes, we may have a demand to try some special ops which may
+not been well supported by auto-scheduler's default search policy. Auto-scheduler currently allows
+user to provide a CustomSketch to cover these cases.
+
+We use sparse matrix multiplication as an example in this tutorial.
+
+Note that this tutorial will not run on Windows or recent versions of macOS. To
+get it to run, you will need to wrap the body of this tutorial in a :code:`if
+__name__ == "__main__":` block.
+"""
+
+import os
+import itertools
+
+import numpy as np
+import tvm
+from tvm import te, auto_scheduler, topi
+from tvm.auto_scheduler import _ffi_api
+from tvm.topi.utils import get_const_tuple
+
+import scipy.sparse as sp
+
+######################################################################
+# Define the computation
+# ^^^^^^^^^^^^^^^^^^^^^^
+# To begin with, let us define the computation of a sparse matmul with several relu and bias add.
+# The function should return the list of input/output tensors.
+# From these tensors, the auto-scheduler can get the whole computational graph.
+
+# We use this function to generate a random bsr matrix
+def random_bsr_matrix(M, N, BS_R, BS_C, density, dtype):
+ import itertools
+
+ Y = np.zeros((M, N), dtype=dtype)
+ assert M % BS_R == 0
+ assert N % BS_C == 0
+ nnz = int(density * M * N)
+ num_blocks = int(nnz / (BS_R * BS_C)) + 1
+ candidate_blocks = np.asarray(list(itertools.product(range(0, M, BS_R), range(0, N, BS_C))))
+ assert candidate_blocks.shape[0] == M // BS_R * N // BS_C
+ chosen_blocks = candidate_blocks[
+ np.random.choice(candidate_blocks.shape[0], size=num_blocks, replace=False)
+ ]
+ for i in range(len(chosen_blocks)):
+ r, c = chosen_blocks[i]
+ Y[r : r + BS_R, c : c + BS_C] = np.random.randn(BS_R, BS_C)
+ s = sp.bsr_matrix(Y, blocksize=(BS_R, BS_C))
+ assert s.data.shape == (num_blocks, BS_R, BS_C)
+ assert s.indices.shape == (num_blocks,)
+ assert s.indptr.shape == (M // BS_R + 1,)
+ return s
+
+@auto_scheduler.register_workload
+def sparse_dense(M, N, K, w_data_shape, w_indices_shape, w_indptr_shape, dtype):
+ X = te.placeholder(shape=(M, K), dtype=dtype)
+ W_data = te.placeholder(shape=w_data_shape, dtype=dtype)
+ W_indices = te.placeholder(shape=w_indices_shape, dtype="int32")
+ W_indptr = te.placeholder(shape=w_indptr_shape, dtype="int32")
+ B = te.placeholder(shape=(M, N), dtype=dtype)
+
+ out = topi.nn.sparse_dense(
+ topi.nn.relu(X), W_data, W_indices, W_indptr
+ )
+ out = te.compute((M, N), lambda i, j: out[i, j] + B[i, j], name="BiasAdd")
+ out = topi.nn.relu(out)
+
+ return [X, W_data, W_indices, W_indptr, B, out]
+
+######################################################################
+# Special step for sparse workload
+# ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+# During schedule tuning, auto-scheduler will use random inputs to measure the performance of a
+# generated schedule. While we cannot directly use a random array as the input of a sparse op, for
+# the "indices" and "indptr" array are meaningful for the computation.
+#
+# To solve this problem, we register these as special buffers, and load them when process program
+# measuring.
+# See the :any:`auto_scheduler.measure` code for more details.
Review comment:
I feel this is too ad hoc. Can we just expose the input buffers in general? For example, Relay graph runtime uses `set_input` to accept data, maybe we can have a similar API in `task` instead of `measure`? This is more reasonable because `measure_ctx` can actually be used by all tasks.
##########
File path: tutorials/auto_scheduler/tune_sparse_x86.py
##########
@@ -0,0 +1,331 @@
+# 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.
+"""
+Auto-scheduling Sparse Matrix Multiplication for CPU by Custom Sketch Rule
+==========================================================================
+**Author**: `Chengfan Jia <https://github.com/jcf94/>`_
+
+This is a tutorial on how to use the auto-scheduler to tune a sparse matrix multiplication for
+CPUs.
+
+Auto-scheduler is designed to explore the schedule with best performance for a given computation
+declaration automatically. While sometimes, we may have a demand to try some special ops which may
+not been well supported by auto-scheduler's default search policy. Auto-scheduler currently allows
+user to provide a CustomSketch to cover these cases.
+
+We use sparse matrix multiplication as an example in this tutorial.
+
+Note that this tutorial will not run on Windows or recent versions of macOS. To
+get it to run, you will need to wrap the body of this tutorial in a :code:`if
+__name__ == "__main__":` block.
+"""
+
+import os
+import itertools
+
+import numpy as np
+import tvm
+from tvm import te, auto_scheduler, topi
+from tvm.auto_scheduler import _ffi_api
+from tvm.topi.utils import get_const_tuple
+
+import scipy.sparse as sp
+
+######################################################################
+# Define the computation
+# ^^^^^^^^^^^^^^^^^^^^^^
+# To begin with, let us define the computation of a sparse matmul with several relu and bias add.
+# The function should return the list of input/output tensors.
+# From these tensors, the auto-scheduler can get the whole computational graph.
+
+# We use this function to generate a random bsr matrix
+def random_bsr_matrix(M, N, BS_R, BS_C, density, dtype):
+ import itertools
+
+ Y = np.zeros((M, N), dtype=dtype)
+ assert M % BS_R == 0
+ assert N % BS_C == 0
+ nnz = int(density * M * N)
+ num_blocks = int(nnz / (BS_R * BS_C)) + 1
+ candidate_blocks = np.asarray(list(itertools.product(range(0, M, BS_R), range(0, N, BS_C))))
+ assert candidate_blocks.shape[0] == M // BS_R * N // BS_C
+ chosen_blocks = candidate_blocks[
+ np.random.choice(candidate_blocks.shape[0], size=num_blocks, replace=False)
+ ]
+ for i in range(len(chosen_blocks)):
+ r, c = chosen_blocks[i]
+ Y[r : r + BS_R, c : c + BS_C] = np.random.randn(BS_R, BS_C)
+ s = sp.bsr_matrix(Y, blocksize=(BS_R, BS_C))
+ assert s.data.shape == (num_blocks, BS_R, BS_C)
+ assert s.indices.shape == (num_blocks,)
+ assert s.indptr.shape == (M // BS_R + 1,)
+ return s
+
+@auto_scheduler.register_workload
+def sparse_dense(M, N, K, w_data_shape, w_indices_shape, w_indptr_shape, dtype):
+ X = te.placeholder(shape=(M, K), dtype=dtype)
+ W_data = te.placeholder(shape=w_data_shape, dtype=dtype)
+ W_indices = te.placeholder(shape=w_indices_shape, dtype="int32")
+ W_indptr = te.placeholder(shape=w_indptr_shape, dtype="int32")
+ B = te.placeholder(shape=(M, N), dtype=dtype)
+
+ out = topi.nn.sparse_dense(
+ topi.nn.relu(X), W_data, W_indices, W_indptr
+ )
+ out = te.compute((M, N), lambda i, j: out[i, j] + B[i, j], name="BiasAdd")
+ out = topi.nn.relu(out)
+
+ return [X, W_data, W_indices, W_indptr, B, out]
+
+######################################################################
+# Special step for sparse workload
+# ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+# During schedule tuning, auto-scheduler will use random inputs to measure the performance of a
+# generated schedule. While we cannot directly use a random array as the input of a sparse op, for
+# the "indices" and "indptr" array are meaningful for the computation.
+#
+# To solve this problem, we register these as special buffers, and load them when process program
+# measuring.
+# See the :any:`auto_scheduler.measure` code for more details.
+
+# Define the basic shapes of this sparse computation
+M = K = N = 512
+BS_R = 16
+BS_C = 1
+density = 0.6
+
+# Generate the test data with numpy
+X_np = np.random.randn(M, K).astype("float32")
+X_np = np.maximum(np.zeros((M, K), dtype="float32"), X_np) # Relu
+W_sp_np = random_bsr_matrix(N, K, BS_R, BS_C, density=density, dtype="float32")
+W_np = W_sp_np.todense()
+Y_np = X_np @ W_np.T # Process the matrix multiplication
+B_np = np.random.randn(M, N).astype("float32")
+Y_np = Y_np + B_np # Bias add
+Y_np = np.maximum(np.zeros((M, N), dtype="float32"), Y_np) # Relu
+
+# Register the sparse data to special buffer
+prefix = "sparse_dense_bsr_%d_%d_%d_%d_%d_%.2f_" % (M, N, K, BS_R, BS_C, density)
+auto_scheduler.measure.register_special_buffer(prefix + "W_data", W_sp_np.data)
+auto_scheduler.measure.register_special_buffer(prefix + "W_indices", W_sp_np.indices)
+auto_scheduler.measure.register_special_buffer(prefix + "W_indptr", W_sp_np.indptr)
+
+######################################################################
+# Create the search task
+# ^^^^^^^^^^^^^^^^^^^^^^
+# We then create a search task with M=N=K=512 and dtype="float32"
+# If your machine supports avx instructions, you can
+#
+# - replace "llvm" below with "llvm -mcpu=core-avx2" to enable AVX2
+# - replace "llvm" below with "llvm -mcpu=skylake-avx512" to enable AVX-512
+
+target = tvm.target.Target("llvm")
+
+task = tvm.auto_scheduler.SearchTask(
+ func=sparse_dense,
+ args=(
+ M, N, K,
+ W_sp_np.data.shape,
+ W_sp_np.indices.shape,
+ W_sp_np.indptr.shape,
+ "float32"
+ ),
+ target=target
+)
+
+# Inspect the computational graph
+print("Computational DAG:")
+print(task.compute_dag)
+
+######################################################################
+# Write the custom sketch for sparse dense op
+# ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+# Before tuning, we will need to define the CustomSketchRule for the sparse dense op.
+#
+# CustomSketchRule consists of two parts: the condition function and the apply function.
+#
+# - condition function: describe when to use this sketch rule. For example, we can match the op
+# by their name or tag.
+# - apply function: describe how to generate the initial sketch. Auto-scheduler provides a set of
+# loop state APIs.
+
+def meet_condition_func(search_policy, state, stage_id):
+ state = auto_scheduler.loop_state.State(state, search_policy.search_task.compute_dag)
+ if state.stages[stage_id].op.tag in [
+ "sparse_dense_sp_rhs_bsrmm", "sparse_dense_sp_rhs_bsrmm_block"
+ ]:
+ return auto_scheduler.PreloadCustomSketchRule.APPLY_AND_SKIP_REST
+ else:
+ return auto_scheduler.PreloadCustomSketchRule.PASS
+
+def apply_func(search_policy, state, stage_id):
+ ret = []
+ s0 = auto_scheduler.loop_state.State(state, search_policy.search_task.compute_dag)
+ if s0.stages[stage_id].op.tag == "sparse_dense_sp_rhs_bsrmm_block":
+ return [s0.state_object, stage_id - 1]
+
+ sparse_dense = s0.stages[stage_id].op
+ sparse_dense_block = s0.stages[stage_id - 1].op
+ assert sparse_dense.tag == "sparse_dense_sp_rhs_bsrmm"
+ assert sparse_dense_block.tag == "sparse_dense_sp_rhs_bsrmm_block"
+
+ # Set the default consumer of compute block
+ consumer = sparse_dense
+
+ # If sparse dense has a single elementwise consumer
+ # We can compute inline the sparse_dense output stage
+ consumers = _ffi_api.SearchPolicyUtilsGetConsumers(
+ search_policy.search_task, s0.state_object, stage_id
+ )
+ if len(consumers) == 1:
+ consumer_id = int(consumers.items()[0][0])
+ if _ffi_api.SearchPolicyUtilsIsElementwiseMatch(
+ search_policy.search_task, s0.state_object, stage_id, consumer_id
+ ):
+ consumer = s0.stages[consumer_id].op
+ s0.compute_inline(sparse_dense)
+
+ i, nb_j, j, row_offset, c = s0[sparse_dense_block].iters
+ m, n = s0[consumer].iters
+ i0, i1, i2 = s0.split(sparse_dense_block, i, [None, None])
+ m0, m1 = s0.follow_split(consumer, m, len(s0.transform_steps) - 1, 1)
+ j0, j1 = s0.split(sparse_dense_block, nb_j, [None])
+ n0, n1 = s0.follow_split(consumer, n, len(s0.transform_steps) - 1, 1)
+ s0.reorder(sparse_dense_block, [i0, j0, i1, j1, row_offset, i2, j, c])
+ s0.reorder(consumer, [m0, n0, m1, n1])
+ s0.compute_at(sparse_dense_block, consumer, n0)
+
+ ret.append([s0.state_object, stage_id - 2])
+
+ return ret
+
+######################################################################
+# Next, we set parameters for the auto-scheduler.
Review comment:
```suggestion
# Next, we set parameters for the auto-scheduler with the custom sketch plugged in.
```
##########
File path: tutorials/auto_scheduler/tune_sparse_x86.py
##########
@@ -0,0 +1,331 @@
+# 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.
+"""
+Auto-scheduling Sparse Matrix Multiplication for CPU by Custom Sketch Rule
+==========================================================================
+**Author**: `Chengfan Jia <https://github.com/jcf94/>`_
+
+This is a tutorial on how to use the auto-scheduler to tune a sparse matrix multiplication for
+CPUs.
+
+Auto-scheduler is designed to explore the schedule with best performance for a given computation
+declaration automatically. While sometimes, we may have a demand to try some special ops which may
+not been well supported by auto-scheduler's default search policy. Auto-scheduler currently allows
+user to provide a CustomSketch to cover these cases.
+
+We use sparse matrix multiplication as an example in this tutorial.
+
+Note that this tutorial will not run on Windows or recent versions of macOS. To
+get it to run, you will need to wrap the body of this tutorial in a :code:`if
+__name__ == "__main__":` block.
+"""
+
+import os
+import itertools
+
+import numpy as np
+import tvm
+from tvm import te, auto_scheduler, topi
+from tvm.auto_scheduler import _ffi_api
+from tvm.topi.utils import get_const_tuple
+
+import scipy.sparse as sp
+
+######################################################################
+# Define the computation
+# ^^^^^^^^^^^^^^^^^^^^^^
+# To begin with, let us define the computation of a sparse matmul with several relu and bias add.
+# The function should return the list of input/output tensors.
+# From these tensors, the auto-scheduler can get the whole computational graph.
+
+# We use this function to generate a random bsr matrix
+def random_bsr_matrix(M, N, BS_R, BS_C, density, dtype):
+ import itertools
+
+ Y = np.zeros((M, N), dtype=dtype)
+ assert M % BS_R == 0
+ assert N % BS_C == 0
+ nnz = int(density * M * N)
+ num_blocks = int(nnz / (BS_R * BS_C)) + 1
+ candidate_blocks = np.asarray(list(itertools.product(range(0, M, BS_R), range(0, N, BS_C))))
+ assert candidate_blocks.shape[0] == M // BS_R * N // BS_C
+ chosen_blocks = candidate_blocks[
+ np.random.choice(candidate_blocks.shape[0], size=num_blocks, replace=False)
+ ]
+ for i in range(len(chosen_blocks)):
+ r, c = chosen_blocks[i]
+ Y[r : r + BS_R, c : c + BS_C] = np.random.randn(BS_R, BS_C)
+ s = sp.bsr_matrix(Y, blocksize=(BS_R, BS_C))
+ assert s.data.shape == (num_blocks, BS_R, BS_C)
+ assert s.indices.shape == (num_blocks,)
+ assert s.indptr.shape == (M // BS_R + 1,)
+ return s
+
+@auto_scheduler.register_workload
+def sparse_dense(M, N, K, w_data_shape, w_indices_shape, w_indptr_shape, dtype):
+ X = te.placeholder(shape=(M, K), dtype=dtype)
+ W_data = te.placeholder(shape=w_data_shape, dtype=dtype)
+ W_indices = te.placeholder(shape=w_indices_shape, dtype="int32")
+ W_indptr = te.placeholder(shape=w_indptr_shape, dtype="int32")
+ B = te.placeholder(shape=(M, N), dtype=dtype)
+
+ out = topi.nn.sparse_dense(
+ topi.nn.relu(X), W_data, W_indices, W_indptr
+ )
+ out = te.compute((M, N), lambda i, j: out[i, j] + B[i, j], name="BiasAdd")
+ out = topi.nn.relu(out)
+
+ return [X, W_data, W_indices, W_indptr, B, out]
+
+######################################################################
+# Special step for sparse workload
+# ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+# During schedule tuning, auto-scheduler will use random inputs to measure the performance of a
+# generated schedule. While we cannot directly use a random array as the input of a sparse op, for
+# the "indices" and "indptr" array are meaningful for the computation.
+#
+# To solve this problem, we register these as special buffers, and load them when process program
+# measuring.
+# See the :any:`auto_scheduler.measure` code for more details.
+
+# Define the basic shapes of this sparse computation
+M = K = N = 512
+BS_R = 16
+BS_C = 1
+density = 0.6
+
+# Generate the test data with numpy
+X_np = np.random.randn(M, K).astype("float32")
+X_np = np.maximum(np.zeros((M, K), dtype="float32"), X_np) # Relu
+W_sp_np = random_bsr_matrix(N, K, BS_R, BS_C, density=density, dtype="float32")
+W_np = W_sp_np.todense()
+Y_np = X_np @ W_np.T # Process the matrix multiplication
+B_np = np.random.randn(M, N).astype("float32")
+Y_np = Y_np + B_np # Bias add
+Y_np = np.maximum(np.zeros((M, N), dtype="float32"), Y_np) # Relu
+
+# Register the sparse data to special buffer
+prefix = "sparse_dense_bsr_%d_%d_%d_%d_%d_%.2f_" % (M, N, K, BS_R, BS_C, density)
+auto_scheduler.measure.register_special_buffer(prefix + "W_data", W_sp_np.data)
+auto_scheduler.measure.register_special_buffer(prefix + "W_indices", W_sp_np.indices)
+auto_scheduler.measure.register_special_buffer(prefix + "W_indptr", W_sp_np.indptr)
+
+######################################################################
+# Create the search task
+# ^^^^^^^^^^^^^^^^^^^^^^
+# We then create a search task with M=N=K=512 and dtype="float32"
+# If your machine supports avx instructions, you can
+#
+# - replace "llvm" below with "llvm -mcpu=core-avx2" to enable AVX2
+# - replace "llvm" below with "llvm -mcpu=skylake-avx512" to enable AVX-512
+
+target = tvm.target.Target("llvm")
+
+task = tvm.auto_scheduler.SearchTask(
+ func=sparse_dense,
+ args=(
+ M, N, K,
+ W_sp_np.data.shape,
+ W_sp_np.indices.shape,
+ W_sp_np.indptr.shape,
+ "float32"
+ ),
+ target=target
+)
+
+# Inspect the computational graph
+print("Computational DAG:")
+print(task.compute_dag)
+
+######################################################################
+# Write the custom sketch for sparse dense op
+# ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+# Before tuning, we will need to define the CustomSketchRule for the sparse dense op.
+#
+# CustomSketchRule consists of two parts: the condition function and the apply function.
+#
+# - condition function: describe when to use this sketch rule. For example, we can match the op
+# by their name or tag.
+# - apply function: describe how to generate the initial sketch. Auto-scheduler provides a set of
+# loop state APIs.
Review comment:
```suggestion
# - apply function: describe how to generate the initial sketch. You can implement it using auto-scheduler provided loop state APIs.
```
##########
File path: tutorials/auto_scheduler/tune_sparse_x86.py
##########
@@ -0,0 +1,331 @@
+# 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.
+"""
+Auto-scheduling Sparse Matrix Multiplication for CPU by Custom Sketch Rule
+==========================================================================
+**Author**: `Chengfan Jia <https://github.com/jcf94/>`_
+
+This is a tutorial on how to use the auto-scheduler to tune a sparse matrix multiplication for
+CPUs.
+
+Auto-scheduler is designed to explore the schedule with best performance for a given computation
+declaration automatically. While sometimes, we may have a demand to try some special ops which may
+not been well supported by auto-scheduler's default search policy. Auto-scheduler currently allows
+user to provide a CustomSketch to cover these cases.
+
+We use sparse matrix multiplication as an example in this tutorial.
Review comment:
```suggestion
We use sparse matrix multiplication as an example in this tutorial to demonstrate how to implement and plug a custom sketch rule to the auto-scheduler search policy.
```
##########
File path: tutorials/auto_scheduler/tune_sparse_x86.py
##########
@@ -0,0 +1,331 @@
+# 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.
+"""
+Auto-scheduling Sparse Matrix Multiplication for CPU by Custom Sketch Rule
+==========================================================================
+**Author**: `Chengfan Jia <https://github.com/jcf94/>`_
+
+This is a tutorial on how to use the auto-scheduler to tune a sparse matrix multiplication for
+CPUs.
+
+Auto-scheduler is designed to explore the schedule with best performance for a given computation
+declaration automatically. While sometimes, we may have a demand to try some special ops which may
+not been well supported by auto-scheduler's default search policy. Auto-scheduler currently allows
Review comment:
```suggestion
not been well-supported by auto-scheduler's default sketch rules and result in poor performance. Fortunately, auto-scheduler currently allows
```
##########
File path: tutorials/auto_scheduler/tune_sparse_x86.py
##########
@@ -0,0 +1,331 @@
+# 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.
+"""
+Auto-scheduling Sparse Matrix Multiplication for CPU by Custom Sketch Rule
+==========================================================================
+**Author**: `Chengfan Jia <https://github.com/jcf94/>`_
+
+This is a tutorial on how to use the auto-scheduler to tune a sparse matrix multiplication for
+CPUs.
+
+Auto-scheduler is designed to explore the schedule with best performance for a given computation
+declaration automatically. While sometimes, we may have a demand to try some special ops which may
+not been well supported by auto-scheduler's default search policy. Auto-scheduler currently allows
+user to provide a CustomSketch to cover these cases.
+
+We use sparse matrix multiplication as an example in this tutorial.
+
+Note that this tutorial will not run on Windows or recent versions of macOS. To
+get it to run, you will need to wrap the body of this tutorial in a :code:`if
+__name__ == "__main__":` block.
+"""
+
+import os
+import itertools
+
+import numpy as np
+import tvm
+from tvm import te, auto_scheduler, topi
+from tvm.auto_scheduler import _ffi_api
+from tvm.topi.utils import get_const_tuple
+
+import scipy.sparse as sp
+
+######################################################################
+# Define the computation
+# ^^^^^^^^^^^^^^^^^^^^^^
+# To begin with, let us define the computation of a sparse matmul with several relu and bias add.
+# The function should return the list of input/output tensors.
+# From these tensors, the auto-scheduler can get the whole computational graph.
+
+# We use this function to generate a random bsr matrix
+def random_bsr_matrix(M, N, BS_R, BS_C, density, dtype):
+ import itertools
+
+ Y = np.zeros((M, N), dtype=dtype)
+ assert M % BS_R == 0
+ assert N % BS_C == 0
+ nnz = int(density * M * N)
+ num_blocks = int(nnz / (BS_R * BS_C)) + 1
+ candidate_blocks = np.asarray(list(itertools.product(range(0, M, BS_R), range(0, N, BS_C))))
+ assert candidate_blocks.shape[0] == M // BS_R * N // BS_C
+ chosen_blocks = candidate_blocks[
+ np.random.choice(candidate_blocks.shape[0], size=num_blocks, replace=False)
+ ]
+ for i in range(len(chosen_blocks)):
+ r, c = chosen_blocks[i]
+ Y[r : r + BS_R, c : c + BS_C] = np.random.randn(BS_R, BS_C)
+ s = sp.bsr_matrix(Y, blocksize=(BS_R, BS_C))
+ assert s.data.shape == (num_blocks, BS_R, BS_C)
+ assert s.indices.shape == (num_blocks,)
+ assert s.indptr.shape == (M // BS_R + 1,)
+ return s
+
+@auto_scheduler.register_workload
+def sparse_dense(M, N, K, w_data_shape, w_indices_shape, w_indptr_shape, dtype):
+ X = te.placeholder(shape=(M, K), dtype=dtype)
+ W_data = te.placeholder(shape=w_data_shape, dtype=dtype)
+ W_indices = te.placeholder(shape=w_indices_shape, dtype="int32")
+ W_indptr = te.placeholder(shape=w_indptr_shape, dtype="int32")
+ B = te.placeholder(shape=(M, N), dtype=dtype)
+
+ out = topi.nn.sparse_dense(
+ topi.nn.relu(X), W_data, W_indices, W_indptr
+ )
+ out = te.compute((M, N), lambda i, j: out[i, j] + B[i, j], name="BiasAdd")
+ out = topi.nn.relu(out)
+
+ return [X, W_data, W_indices, W_indptr, B, out]
+
+######################################################################
+# Special step for sparse workload
+# ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+# During schedule tuning, auto-scheduler will use random inputs to measure the performance of a
+# generated schedule. While we cannot directly use a random array as the input of a sparse op, for
+# the "indices" and "indptr" array are meaningful for the computation.
+#
+# To solve this problem, we register these as special buffers, and load them when process program
+# measuring.
+# See the :any:`auto_scheduler.measure` code for more details.
+
+# Define the basic shapes of this sparse computation
+M = K = N = 512
+BS_R = 16
+BS_C = 1
+density = 0.6
+
+# Generate the test data with numpy
+X_np = np.random.randn(M, K).astype("float32")
+X_np = np.maximum(np.zeros((M, K), dtype="float32"), X_np) # Relu
+W_sp_np = random_bsr_matrix(N, K, BS_R, BS_C, density=density, dtype="float32")
+W_np = W_sp_np.todense()
+Y_np = X_np @ W_np.T # Process the matrix multiplication
+B_np = np.random.randn(M, N).astype("float32")
+Y_np = Y_np + B_np # Bias add
+Y_np = np.maximum(np.zeros((M, N), dtype="float32"), Y_np) # Relu
+
+# Register the sparse data to special buffer
+prefix = "sparse_dense_bsr_%d_%d_%d_%d_%d_%.2f_" % (M, N, K, BS_R, BS_C, density)
+auto_scheduler.measure.register_special_buffer(prefix + "W_data", W_sp_np.data)
+auto_scheduler.measure.register_special_buffer(prefix + "W_indices", W_sp_np.indices)
+auto_scheduler.measure.register_special_buffer(prefix + "W_indptr", W_sp_np.indptr)
+
+######################################################################
+# Create the search task
+# ^^^^^^^^^^^^^^^^^^^^^^
+# We then create a search task with M=N=K=512 and dtype="float32"
+# If your machine supports avx instructions, you can
+#
+# - replace "llvm" below with "llvm -mcpu=core-avx2" to enable AVX2
+# - replace "llvm" below with "llvm -mcpu=skylake-avx512" to enable AVX-512
+
+target = tvm.target.Target("llvm")
+
+task = tvm.auto_scheduler.SearchTask(
+ func=sparse_dense,
+ args=(
+ M, N, K,
+ W_sp_np.data.shape,
+ W_sp_np.indices.shape,
+ W_sp_np.indptr.shape,
+ "float32"
+ ),
+ target=target
+)
+
+# Inspect the computational graph
+print("Computational DAG:")
+print(task.compute_dag)
+
+######################################################################
+# Write the custom sketch for sparse dense op
+# ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+# Before tuning, we will need to define the CustomSketchRule for the sparse dense op.
+#
+# CustomSketchRule consists of two parts: the condition function and the apply function.
+#
+# - condition function: describe when to use this sketch rule. For example, we can match the op
+# by their name or tag.
+# - apply function: describe how to generate the initial sketch. Auto-scheduler provides a set of
+# loop state APIs.
+
+def meet_condition_func(search_policy, state, stage_id):
+ state = auto_scheduler.loop_state.State(state, search_policy.search_task.compute_dag)
+ if state.stages[stage_id].op.tag in [
+ "sparse_dense_sp_rhs_bsrmm", "sparse_dense_sp_rhs_bsrmm_block"
+ ]:
+ return auto_scheduler.PreloadCustomSketchRule.APPLY_AND_SKIP_REST
+ else:
+ return auto_scheduler.PreloadCustomSketchRule.PASS
+
+def apply_func(search_policy, state, stage_id):
+ ret = []
+ s0 = auto_scheduler.loop_state.State(state, search_policy.search_task.compute_dag)
+ if s0.stages[stage_id].op.tag == "sparse_dense_sp_rhs_bsrmm_block":
+ return [s0.state_object, stage_id - 1]
+
+ sparse_dense = s0.stages[stage_id].op
+ sparse_dense_block = s0.stages[stage_id - 1].op
+ assert sparse_dense.tag == "sparse_dense_sp_rhs_bsrmm"
+ assert sparse_dense_block.tag == "sparse_dense_sp_rhs_bsrmm_block"
+
+ # Set the default consumer of compute block
+ consumer = sparse_dense
+
+ # If sparse dense has a single elementwise consumer
+ # We can compute inline the sparse_dense output stage
+ consumers = _ffi_api.SearchPolicyUtilsGetConsumers(
+ search_policy.search_task, s0.state_object, stage_id
+ )
+ if len(consumers) == 1:
+ consumer_id = int(consumers.items()[0][0])
+ if _ffi_api.SearchPolicyUtilsIsElementwiseMatch(
+ search_policy.search_task, s0.state_object, stage_id, consumer_id
+ ):
+ consumer = s0.stages[consumer_id].op
+ s0.compute_inline(sparse_dense)
+
+ i, nb_j, j, row_offset, c = s0[sparse_dense_block].iters
+ m, n = s0[consumer].iters
+ i0, i1, i2 = s0.split(sparse_dense_block, i, [None, None])
+ m0, m1 = s0.follow_split(consumer, m, len(s0.transform_steps) - 1, 1)
+ j0, j1 = s0.split(sparse_dense_block, nb_j, [None])
+ n0, n1 = s0.follow_split(consumer, n, len(s0.transform_steps) - 1, 1)
+ s0.reorder(sparse_dense_block, [i0, j0, i1, j1, row_offset, i2, j, c])
+ s0.reorder(consumer, [m0, n0, m1, n1])
+ s0.compute_at(sparse_dense_block, consumer, n0)
+
+ ret.append([s0.state_object, stage_id - 2])
+
+ return ret
+
+######################################################################
+# Next, we set parameters for the auto-scheduler.
+#
+# * :code:`num_measure_trials` is the number of measurement trials we can use during the search.
+# We only make 10 trials in this tutorial for a fast demonstration. In practice, 1000 is a
+# good value for the search to converge. You can do more trials according to your time budget.
+# * In addition, we use :code:`RecordToFile` to dump measurement records into a file `matmul.json`.
+# The measurement records can be used to query the history best, resume the search,
+# and do more analyses later.
+# * see :any:`auto_scheduler.TuningOptions` for more parameters
+# * Here, we need to create a :code:`auto_scheduler.SketchPolicy` object, and add the custom sketch
+# rule as a `init_search_callbacks`.
+
+log_file = "sparse_dense.json"
+tune_option = auto_scheduler.TuningOptions(
+ num_measure_trials=10,
+ measure_callbacks=[auto_scheduler.RecordToFile(log_file)],
+ verbose=2,
+)
+
+search_policy = auto_scheduler.SketchPolicy(
+ task,
+ program_cost_model=auto_scheduler.XGBModel(),
+ init_search_callbacks=[
+ auto_scheduler.PreloadCustomSketchRule(meet_condition_func, apply_func, "SparseDense")
+ ]
+)
+
+######################################################################
+# Run the search
+# ^^^^^^^^^^^^^^
+# Now we get all inputs ready.
+# We can kick off the search and let the auto-scheduler do its magic.
+# After some measurement trials, we can load the best schedule from the log
+# file and apply it.
+
+# Run auto-tuning (search)
+task.tune(tune_option, search_policy)
+# Apply the best schedule
+sch, args = task.apply_best(log_file)
+
+######################################################################
+# We can lower the schedule to see the IR after auto-scheduling.
+# The auto-scheduler correctly performs optimizations including multi-level tiling,
+# layout transformation, parallelization, vectorization, unrolling, and operator fusion.
+
+print("Lowered TIR:")
+print(tvm.lower(sch, args, simple_mode=True))
+
+######################################################################
+# Check correctness and evaluate performance
+# ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+# We build the binary and check its correctness and performance.
+
+func = tvm.build(sch, args, target)
+
+ctx = tvm.cpu()
+
+X_tvm = tvm.nd.array(X_np, ctx=ctx)
+W_data_tvm = tvm.nd.array(W_sp_np.data, ctx=ctx)
+W_indices_tvm = tvm.nd.array(W_sp_np.indices, ctx=ctx)
+W_indptr_tvm = tvm.nd.array(W_sp_np.indptr, ctx=ctx)
+B_tvm = tvm.nd.array(B_np, ctx=ctx)
+Y_tvm = tvm.nd.empty(Y_np.shape, ctx=ctx)
+
+func(X_tvm, W_data_tvm, W_indices_tvm, W_indptr_tvm, B_tvm, Y_tvm)
+
+# Check results
+tvm.testing.assert_allclose(Y_np, Y_tvm.asnumpy(), atol=1e-4, rtol=1e-4)
+
+# Evaluate execution time.
+evaluator = func.time_evaluator(func.entry_name, ctx, min_repeat_ms=500)
+print(
+ "Execution time of this operator: %.3f ms"
+ % (np.median(evaluator(X_tvm, W_data_tvm, W_indices_tvm, W_indptr_tvm, B_tvm, Y_tvm).results) * 1000)
+)
+
+######################################################################
+# Using the record file
+# ^^^^^^^^^^^^^^^^^^^^^
+# During the search, all measurement records are dumped into the record
+# file "matmul.json". The measurement records can be used to re-apply search results,
+# resume the search, and perform other analyses.
+
+######################################################################
+# Here is an example where we load the best schedule from a file,
+# and print the equivalent python schedule API. This can be used for
+# debugging and learning the behavior of the auto-scheduler.
+
+print("Equivalent python schedule:")
+print(task.print_best(log_file))
+
+######################################################################
+# A more complicated example is to resume the search.
+# In this case, we need to create the search policy and cost model by ourselves
+# and resume the status of search policy and cost model with the log file.
+# In the example below we resume the status and do more 5 trials.
+
+
+def resume_search(task, log_file):
+ print("Resume search:")
+ cost_model = auto_scheduler.XGBModel()
+ cost_model.update_from_file(log_file)
+ search_policy = auto_scheduler.SketchPolicy(
+ task, cost_model, init_search_callbacks=[
+ auto_scheduler.PreloadMeasuredStates(log_file),
+ auto_scheduler.PreloadCustomSketchRule(meet_condition_func, apply_func, "SparseDense")
+ ]
+ )
+ tune_option = auto_scheduler.TuningOptions(
+ num_measure_trials=5, measure_callbacks=[auto_scheduler.RecordToFile(log_file)]
+ )
+ task.tune(tune_option, search_policy=search_policy)
+
+
+resume_search(task, log_file)
Review comment:
I think you can simply refer to other tutorials to skip this part. This tutorial is more advance so it should be fine to assume most readers are more or less familiar with auto-scheduler already.
##########
File path: tutorials/auto_scheduler/tune_sparse_x86.py
##########
@@ -0,0 +1,331 @@
+# 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.
+"""
+Auto-scheduling Sparse Matrix Multiplication for CPU by Custom Sketch Rule
+==========================================================================
+**Author**: `Chengfan Jia <https://github.com/jcf94/>`_
+
+This is a tutorial on how to use the auto-scheduler to tune a sparse matrix multiplication for
+CPUs.
+
+Auto-scheduler is designed to explore the schedule with best performance for a given computation
+declaration automatically. While sometimes, we may have a demand to try some special ops which may
+not been well supported by auto-scheduler's default search policy. Auto-scheduler currently allows
+user to provide a CustomSketch to cover these cases.
+
+We use sparse matrix multiplication as an example in this tutorial.
+
+Note that this tutorial will not run on Windows or recent versions of macOS. To
+get it to run, you will need to wrap the body of this tutorial in a :code:`if
+__name__ == "__main__":` block.
+"""
+
+import os
+import itertools
+
+import numpy as np
+import tvm
+from tvm import te, auto_scheduler, topi
+from tvm.auto_scheduler import _ffi_api
+from tvm.topi.utils import get_const_tuple
+
+import scipy.sparse as sp
+
+######################################################################
+# Define the computation
+# ^^^^^^^^^^^^^^^^^^^^^^
+# To begin with, let us define the computation of a sparse matmul with several relu and bias add.
+# The function should return the list of input/output tensors.
+# From these tensors, the auto-scheduler can get the whole computational graph.
+
+# We use this function to generate a random bsr matrix
+def random_bsr_matrix(M, N, BS_R, BS_C, density, dtype):
+ import itertools
+
+ Y = np.zeros((M, N), dtype=dtype)
+ assert M % BS_R == 0
+ assert N % BS_C == 0
+ nnz = int(density * M * N)
+ num_blocks = int(nnz / (BS_R * BS_C)) + 1
+ candidate_blocks = np.asarray(list(itertools.product(range(0, M, BS_R), range(0, N, BS_C))))
+ assert candidate_blocks.shape[0] == M // BS_R * N // BS_C
+ chosen_blocks = candidate_blocks[
+ np.random.choice(candidate_blocks.shape[0], size=num_blocks, replace=False)
+ ]
+ for i in range(len(chosen_blocks)):
+ r, c = chosen_blocks[i]
+ Y[r : r + BS_R, c : c + BS_C] = np.random.randn(BS_R, BS_C)
+ s = sp.bsr_matrix(Y, blocksize=(BS_R, BS_C))
+ assert s.data.shape == (num_blocks, BS_R, BS_C)
+ assert s.indices.shape == (num_blocks,)
+ assert s.indptr.shape == (M // BS_R + 1,)
+ return s
+
+@auto_scheduler.register_workload
+def sparse_dense(M, N, K, w_data_shape, w_indices_shape, w_indptr_shape, dtype):
+ X = te.placeholder(shape=(M, K), dtype=dtype)
+ W_data = te.placeholder(shape=w_data_shape, dtype=dtype)
+ W_indices = te.placeholder(shape=w_indices_shape, dtype="int32")
+ W_indptr = te.placeholder(shape=w_indptr_shape, dtype="int32")
+ B = te.placeholder(shape=(M, N), dtype=dtype)
+
+ out = topi.nn.sparse_dense(
+ topi.nn.relu(X), W_data, W_indices, W_indptr
+ )
+ out = te.compute((M, N), lambda i, j: out[i, j] + B[i, j], name="BiasAdd")
+ out = topi.nn.relu(out)
+
+ return [X, W_data, W_indices, W_indptr, B, out]
+
+######################################################################
+# Special step for sparse workload
+# ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+# During schedule tuning, auto-scheduler will use random inputs to measure the performance of a
+# generated schedule. While we cannot directly use a random array as the input of a sparse op, for
+# the "indices" and "indptr" array are meaningful for the computation.
+#
+# To solve this problem, we register these as special buffers, and load them when process program
+# measuring.
+# See the :any:`auto_scheduler.measure` code for more details.
+
+# Define the basic shapes of this sparse computation
+M = K = N = 512
+BS_R = 16
+BS_C = 1
+density = 0.6
+
+# Generate the test data with numpy
+X_np = np.random.randn(M, K).astype("float32")
+X_np = np.maximum(np.zeros((M, K), dtype="float32"), X_np) # Relu
+W_sp_np = random_bsr_matrix(N, K, BS_R, BS_C, density=density, dtype="float32")
+W_np = W_sp_np.todense()
+Y_np = X_np @ W_np.T # Process the matrix multiplication
+B_np = np.random.randn(M, N).astype("float32")
+Y_np = Y_np + B_np # Bias add
+Y_np = np.maximum(np.zeros((M, N), dtype="float32"), Y_np) # Relu
+
+# Register the sparse data to special buffer
+prefix = "sparse_dense_bsr_%d_%d_%d_%d_%d_%.2f_" % (M, N, K, BS_R, BS_C, density)
+auto_scheduler.measure.register_special_buffer(prefix + "W_data", W_sp_np.data)
+auto_scheduler.measure.register_special_buffer(prefix + "W_indices", W_sp_np.indices)
+auto_scheduler.measure.register_special_buffer(prefix + "W_indptr", W_sp_np.indptr)
+
+######################################################################
+# Create the search task
+# ^^^^^^^^^^^^^^^^^^^^^^
+# We then create a search task with M=N=K=512 and dtype="float32"
+# If your machine supports avx instructions, you can
+#
+# - replace "llvm" below with "llvm -mcpu=core-avx2" to enable AVX2
+# - replace "llvm" below with "llvm -mcpu=skylake-avx512" to enable AVX-512
+
+target = tvm.target.Target("llvm")
+
+task = tvm.auto_scheduler.SearchTask(
+ func=sparse_dense,
+ args=(
+ M, N, K,
+ W_sp_np.data.shape,
+ W_sp_np.indices.shape,
+ W_sp_np.indptr.shape,
+ "float32"
+ ),
+ target=target
+)
+
+# Inspect the computational graph
+print("Computational DAG:")
+print(task.compute_dag)
+
+######################################################################
+# Write the custom sketch for sparse dense op
+# ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+# Before tuning, we will need to define the CustomSketchRule for the sparse dense op.
+#
+# CustomSketchRule consists of two parts: the condition function and the apply function.
+#
+# - condition function: describe when to use this sketch rule. For example, we can match the op
+# by their name or tag.
+# - apply function: describe how to generate the initial sketch. Auto-scheduler provides a set of
+# loop state APIs.
+
+def meet_condition_func(search_policy, state, stage_id):
+ state = auto_scheduler.loop_state.State(state, search_policy.search_task.compute_dag)
+ if state.stages[stage_id].op.tag in [
+ "sparse_dense_sp_rhs_bsrmm", "sparse_dense_sp_rhs_bsrmm_block"
+ ]:
+ return auto_scheduler.PreloadCustomSketchRule.APPLY_AND_SKIP_REST
+ else:
+ return auto_scheduler.PreloadCustomSketchRule.PASS
+
+def apply_func(search_policy, state, stage_id):
+ ret = []
+ s0 = auto_scheduler.loop_state.State(state, search_policy.search_task.compute_dag)
+ if s0.stages[stage_id].op.tag == "sparse_dense_sp_rhs_bsrmm_block":
+ return [s0.state_object, stage_id - 1]
+
+ sparse_dense = s0.stages[stage_id].op
+ sparse_dense_block = s0.stages[stage_id - 1].op
+ assert sparse_dense.tag == "sparse_dense_sp_rhs_bsrmm"
+ assert sparse_dense_block.tag == "sparse_dense_sp_rhs_bsrmm_block"
+
+ # Set the default consumer of compute block
+ consumer = sparse_dense
+
+ # If sparse dense has a single elementwise consumer
+ # We can compute inline the sparse_dense output stage
+ consumers = _ffi_api.SearchPolicyUtilsGetConsumers(
+ search_policy.search_task, s0.state_object, stage_id
+ )
+ if len(consumers) == 1:
+ consumer_id = int(consumers.items()[0][0])
+ if _ffi_api.SearchPolicyUtilsIsElementwiseMatch(
+ search_policy.search_task, s0.state_object, stage_id, consumer_id
+ ):
+ consumer = s0.stages[consumer_id].op
+ s0.compute_inline(sparse_dense)
+
+ i, nb_j, j, row_offset, c = s0[sparse_dense_block].iters
+ m, n = s0[consumer].iters
+ i0, i1, i2 = s0.split(sparse_dense_block, i, [None, None])
+ m0, m1 = s0.follow_split(consumer, m, len(s0.transform_steps) - 1, 1)
+ j0, j1 = s0.split(sparse_dense_block, nb_j, [None])
+ n0, n1 = s0.follow_split(consumer, n, len(s0.transform_steps) - 1, 1)
+ s0.reorder(sparse_dense_block, [i0, j0, i1, j1, row_offset, i2, j, c])
+ s0.reorder(consumer, [m0, n0, m1, n1])
+ s0.compute_at(sparse_dense_block, consumer, n0)
+
+ ret.append([s0.state_object, stage_id - 2])
+
+ return ret
+
+######################################################################
+# Next, we set parameters for the auto-scheduler.
+#
+# * :code:`num_measure_trials` is the number of measurement trials we can use during the search.
+# We only make 10 trials in this tutorial for a fast demonstration. In practice, 1000 is a
+# good value for the search to converge. You can do more trials according to your time budget.
+# * In addition, we use :code:`RecordToFile` to dump measurement records into a file `matmul.json`.
+# The measurement records can be used to query the history best, resume the search,
+# and do more analyses later.
+# * see :any:`auto_scheduler.TuningOptions` for more parameters
+# * Here, we need to create a :code:`auto_scheduler.SketchPolicy` object, and add the custom sketch
+# rule as a `init_search_callbacks`.
+
+log_file = "sparse_dense.json"
+tune_option = auto_scheduler.TuningOptions(
+ num_measure_trials=10,
+ measure_callbacks=[auto_scheduler.RecordToFile(log_file)],
+ verbose=2,
+)
+
+search_policy = auto_scheduler.SketchPolicy(
+ task,
+ program_cost_model=auto_scheduler.XGBModel(),
+ init_search_callbacks=[
+ auto_scheduler.PreloadCustomSketchRule(meet_condition_func, apply_func, "SparseDense")
+ ]
+)
+
+######################################################################
+# Run the search
+# ^^^^^^^^^^^^^^
+# Now we get all inputs ready.
+# We can kick off the search and let the auto-scheduler do its magic.
+# After some measurement trials, we can load the best schedule from the log
+# file and apply it.
+
+# Run auto-tuning (search)
+task.tune(tune_option, search_policy)
Review comment:
Do not run the tuning in the tutorial. We should comment out this line and commit a pre-tuned log to `ci_logs`.
##########
File path: tutorials/auto_scheduler/tune_sparse_x86.py
##########
@@ -0,0 +1,331 @@
+# 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.
+"""
+Auto-scheduling Sparse Matrix Multiplication for CPU by Custom Sketch Rule
+==========================================================================
+**Author**: `Chengfan Jia <https://github.com/jcf94/>`_
+
+This is a tutorial on how to use the auto-scheduler to tune a sparse matrix multiplication for
+CPUs.
+
+Auto-scheduler is designed to explore the schedule with best performance for a given computation
+declaration automatically. While sometimes, we may have a demand to try some special ops which may
+not been well supported by auto-scheduler's default search policy. Auto-scheduler currently allows
+user to provide a CustomSketch to cover these cases.
+
+We use sparse matrix multiplication as an example in this tutorial.
+
+Note that this tutorial will not run on Windows or recent versions of macOS. To
+get it to run, you will need to wrap the body of this tutorial in a :code:`if
+__name__ == "__main__":` block.
+"""
+
+import os
+import itertools
+
+import numpy as np
+import tvm
+from tvm import te, auto_scheduler, topi
+from tvm.auto_scheduler import _ffi_api
+from tvm.topi.utils import get_const_tuple
+
+import scipy.sparse as sp
+
+######################################################################
+# Define the computation
+# ^^^^^^^^^^^^^^^^^^^^^^
+# To begin with, let us define the computation of a sparse matmul with several relu and bias add.
+# The function should return the list of input/output tensors.
+# From these tensors, the auto-scheduler can get the whole computational graph.
+
+# We use this function to generate a random bsr matrix
+def random_bsr_matrix(M, N, BS_R, BS_C, density, dtype):
+ import itertools
+
+ Y = np.zeros((M, N), dtype=dtype)
+ assert M % BS_R == 0
+ assert N % BS_C == 0
+ nnz = int(density * M * N)
+ num_blocks = int(nnz / (BS_R * BS_C)) + 1
+ candidate_blocks = np.asarray(list(itertools.product(range(0, M, BS_R), range(0, N, BS_C))))
+ assert candidate_blocks.shape[0] == M // BS_R * N // BS_C
+ chosen_blocks = candidate_blocks[
+ np.random.choice(candidate_blocks.shape[0], size=num_blocks, replace=False)
+ ]
+ for i in range(len(chosen_blocks)):
+ r, c = chosen_blocks[i]
+ Y[r : r + BS_R, c : c + BS_C] = np.random.randn(BS_R, BS_C)
+ s = sp.bsr_matrix(Y, blocksize=(BS_R, BS_C))
+ assert s.data.shape == (num_blocks, BS_R, BS_C)
+ assert s.indices.shape == (num_blocks,)
+ assert s.indptr.shape == (M // BS_R + 1,)
+ return s
+
+@auto_scheduler.register_workload
+def sparse_dense(M, N, K, w_data_shape, w_indices_shape, w_indptr_shape, dtype):
+ X = te.placeholder(shape=(M, K), dtype=dtype)
+ W_data = te.placeholder(shape=w_data_shape, dtype=dtype)
+ W_indices = te.placeholder(shape=w_indices_shape, dtype="int32")
+ W_indptr = te.placeholder(shape=w_indptr_shape, dtype="int32")
+ B = te.placeholder(shape=(M, N), dtype=dtype)
+
+ out = topi.nn.sparse_dense(
+ topi.nn.relu(X), W_data, W_indices, W_indptr
+ )
+ out = te.compute((M, N), lambda i, j: out[i, j] + B[i, j], name="BiasAdd")
+ out = topi.nn.relu(out)
+
+ return [X, W_data, W_indices, W_indptr, B, out]
+
+######################################################################
+# Special step for sparse workload
+# ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+# During schedule tuning, auto-scheduler will use random inputs to measure the performance of a
+# generated schedule. While we cannot directly use a random array as the input of a sparse op, for
+# the "indices" and "indptr" array are meaningful for the computation.
+#
+# To solve this problem, we register these as special buffers, and load them when process program
+# measuring.
+# See the :any:`auto_scheduler.measure` code for more details.
+
+# Define the basic shapes of this sparse computation
+M = K = N = 512
+BS_R = 16
+BS_C = 1
+density = 0.6
+
+# Generate the test data with numpy
+X_np = np.random.randn(M, K).astype("float32")
+X_np = np.maximum(np.zeros((M, K), dtype="float32"), X_np) # Relu
+W_sp_np = random_bsr_matrix(N, K, BS_R, BS_C, density=density, dtype="float32")
+W_np = W_sp_np.todense()
+Y_np = X_np @ W_np.T # Process the matrix multiplication
+B_np = np.random.randn(M, N).astype("float32")
+Y_np = Y_np + B_np # Bias add
+Y_np = np.maximum(np.zeros((M, N), dtype="float32"), Y_np) # Relu
+
+# Register the sparse data to special buffer
+prefix = "sparse_dense_bsr_%d_%d_%d_%d_%d_%.2f_" % (M, N, K, BS_R, BS_C, density)
+auto_scheduler.measure.register_special_buffer(prefix + "W_data", W_sp_np.data)
+auto_scheduler.measure.register_special_buffer(prefix + "W_indices", W_sp_np.indices)
+auto_scheduler.measure.register_special_buffer(prefix + "W_indptr", W_sp_np.indptr)
+
+######################################################################
+# Create the search task
+# ^^^^^^^^^^^^^^^^^^^^^^
+# We then create a search task with M=N=K=512 and dtype="float32"
+# If your machine supports avx instructions, you can
+#
+# - replace "llvm" below with "llvm -mcpu=core-avx2" to enable AVX2
+# - replace "llvm" below with "llvm -mcpu=skylake-avx512" to enable AVX-512
+
+target = tvm.target.Target("llvm")
+
+task = tvm.auto_scheduler.SearchTask(
+ func=sparse_dense,
+ args=(
+ M, N, K,
+ W_sp_np.data.shape,
+ W_sp_np.indices.shape,
+ W_sp_np.indptr.shape,
+ "float32"
+ ),
+ target=target
+)
+
+# Inspect the computational graph
+print("Computational DAG:")
+print(task.compute_dag)
+
+######################################################################
+# Write the custom sketch for sparse dense op
+# ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+# Before tuning, we will need to define the CustomSketchRule for the sparse dense op.
+#
+# CustomSketchRule consists of two parts: the condition function and the apply function.
+#
+# - condition function: describe when to use this sketch rule. For example, we can match the op
+# by their name or tag.
Review comment:
```suggestion
# - condition function: describe when to apply this sketch rule. For example, we can only apply the rule
# to the sparse ops by matching their name and tag.
```
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] jcf94 commented on a change in pull request #7313: [AutoSchedule] Sparse dense tuning support with custom sketch rule
Posted by GitBox <gi...@apache.org>.
jcf94 commented on a change in pull request #7313:
URL: https://github.com/apache/tvm/pull/7313#discussion_r586937376
##########
File path: python/tvm/auto_scheduler/measure.py
##########
@@ -719,6 +720,45 @@ def local_builder_build(inputs, timeout, n_parallel, build_func="default", verbo
return results
+def _prepare_input_map(args):
+ """This function deals with special task inputs.
+
+ Parameters
+ ----------
+ args : List[Tensor]
+ Input/output Tensor of a TVM subgraph.
+
+ Returns
+ -------
+ A Dict[Tensor, str] that maps the input Tensor to a buffer name.
+
+ Note
+ ----
+ The buffer name is specially designed, and these buffer should be provided in
+ `SearchTask(..., task_inputs={...})`.
+ """
+ # pylint: disable=import-outside-toplevel
+ from tvm import topi # lazily import to avoid recursive dependency
+
+ # A dict that maps the input tensor arg to a buffer name
+ tensor_input_map = {}
+
+ # Case 0: Check placeholder name
+ for arg in args:
+ if isinstance(arg.op, tvm.te.PlaceholderOp):
+ if arg.op.name != "placeholder":
+ tensor_input_map[arg] = arg.op.name
+
+ # Case 1: Check sparse op
+ sparse_input_map = topi.nn.sparse.try_get_sparse_input(args)
Review comment:
Yeah, I've also had some discussions in our weekly sync while didn't figure out any better solutions.
There're several reasons:
1. Different ops have different requirements over specific inputs;
2. While the problems is in a subgraph generated in Relay Integration, the placeholder are all the same, we can not differentiate them from tag, name or any other way, even the order of inputs are not guaranteed.
Currently approach is to merge all specific inputs checking to this function, at least they have a same entry here. For the other ops, you have to add their own check functions below.
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] comaniac commented on a change in pull request #7313: [AutoSchedule] Sparse dense tuning support with custom sketch rule
Posted by GitBox <gi...@apache.org>.
comaniac commented on a change in pull request #7313:
URL: https://github.com/apache/tvm/pull/7313#discussion_r566325374
##########
File path: tutorials/auto_scheduler/tune_sparse_x86.py
##########
@@ -0,0 +1,331 @@
+# 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.
+"""
+Auto-scheduling Sparse Matrix Multiplication for CPU by Custom Sketch Rule
+==========================================================================
+**Author**: `Chengfan Jia <https://github.com/jcf94/>`_
+
+This is a tutorial on how to use the auto-scheduler to tune a sparse matrix multiplication for
+CPUs.
+
+Auto-scheduler is designed to explore the schedule with best performance for a given computation
+declaration automatically. While sometimes, we may have a demand to try some special ops which may
+not been well supported by auto-scheduler's default search policy. Auto-scheduler currently allows
+user to provide a CustomSketch to cover these cases.
+
+We use sparse matrix multiplication as an example in this tutorial.
+
+Note that this tutorial will not run on Windows or recent versions of macOS. To
+get it to run, you will need to wrap the body of this tutorial in a :code:`if
+__name__ == "__main__":` block.
+"""
+
+import os
+import itertools
+
+import numpy as np
+import tvm
+from tvm import te, auto_scheduler, topi
+from tvm.auto_scheduler import _ffi_api
+from tvm.topi.utils import get_const_tuple
+
+import scipy.sparse as sp
+
+######################################################################
+# Define the computation
+# ^^^^^^^^^^^^^^^^^^^^^^
+# To begin with, let us define the computation of a sparse matmul with several relu and bias add.
+# The function should return the list of input/output tensors.
+# From these tensors, the auto-scheduler can get the whole computational graph.
+
+# We use this function to generate a random bsr matrix
+def random_bsr_matrix(M, N, BS_R, BS_C, density, dtype):
+ import itertools
+
+ Y = np.zeros((M, N), dtype=dtype)
+ assert M % BS_R == 0
+ assert N % BS_C == 0
+ nnz = int(density * M * N)
+ num_blocks = int(nnz / (BS_R * BS_C)) + 1
+ candidate_blocks = np.asarray(list(itertools.product(range(0, M, BS_R), range(0, N, BS_C))))
+ assert candidate_blocks.shape[0] == M // BS_R * N // BS_C
+ chosen_blocks = candidate_blocks[
+ np.random.choice(candidate_blocks.shape[0], size=num_blocks, replace=False)
+ ]
+ for i in range(len(chosen_blocks)):
+ r, c = chosen_blocks[i]
+ Y[r : r + BS_R, c : c + BS_C] = np.random.randn(BS_R, BS_C)
+ s = sp.bsr_matrix(Y, blocksize=(BS_R, BS_C))
+ assert s.data.shape == (num_blocks, BS_R, BS_C)
+ assert s.indices.shape == (num_blocks,)
+ assert s.indptr.shape == (M // BS_R + 1,)
+ return s
+
+@auto_scheduler.register_workload
+def sparse_dense(M, N, K, w_data_shape, w_indices_shape, w_indptr_shape, dtype):
+ X = te.placeholder(shape=(M, K), dtype=dtype)
+ W_data = te.placeholder(shape=w_data_shape, dtype=dtype)
+ W_indices = te.placeholder(shape=w_indices_shape, dtype="int32")
+ W_indptr = te.placeholder(shape=w_indptr_shape, dtype="int32")
+ B = te.placeholder(shape=(M, N), dtype=dtype)
+
+ out = topi.nn.sparse_dense(
+ topi.nn.relu(X), W_data, W_indices, W_indptr
+ )
+ out = te.compute((M, N), lambda i, j: out[i, j] + B[i, j], name="BiasAdd")
+ out = topi.nn.relu(out)
+
+ return [X, W_data, W_indices, W_indptr, B, out]
+
+######################################################################
+# Special step for sparse workload
+# ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+# During schedule tuning, auto-scheduler will use random inputs to measure the performance of a
+# generated schedule. While we cannot directly use a random array as the input of a sparse op, for
+# the "indices" and "indptr" array are meaningful for the computation.
+#
+# To solve this problem, we register these as special buffers, and load them when process program
+# measuring.
+# See the :any:`auto_scheduler.measure` code for more details.
Review comment:
Sounds good to me.
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org