[AutoSchedule] Sparse dense tuning support with custom sketch rule

[jcf94]

No description provided.

[ANSHUMAN87]

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.

[tkonolige]

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)?

[jcf94]

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.

[tkonolige]

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).

[merrymercy]

minor style comments

[merrymercy]

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

[jcf94]

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.

[tkonolige]

I just want to echo my earlier concerns about special casing sparse inputs instead of having a generic mechanism for detecting special task inputs.

[tkonolige]

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).

[jcf94]

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.

[jcf94]

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.

[tkonolige]

Could we associate the lookup mechanism with @register_workload? It would at least be extensible then.

[jcf94]

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.

[jcf94]

@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.

comments are addressed