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:
[email protected]
