[Docs] Improve API autotune, search ?

[paugier]

Where is the problem?

https://docs.modular.com/mojo/programming-manual.html#autotuning-adaptive-compilation

What can we do better?

The example in https://docs.modular.com/mojo/programming-manual.html#autotuning-adaptive-compilation is both very impressive/nice and IMHO really ugly and difficult to follow.

It seems to me that the API (and therefore the example) could be improved with 2 "simple" changes. (I guess there are good reasons not to do that but I'm so shocked that I write my reaction here.)

1. If I understand correctly, using autotune in a function magically changes the type of what is produced by the fn keyword. A explicit decorator would make clear that one doesn't define only one function but a VariadicList[fn_type]

from autotune import autotune, search, forked_functions

@forked_functions
fn buffer_elementwise_add_implementations[
    dt: DType
](lhs: DTypePointer[dt], rhs: DTypePointer[dt], result: DTypePointer[dt], N: Int):
    ...

Of course, forked_functions might not be the best name.

2. The syntax to define best_impl in buffer_elementwise_add is for me (a Python dev) very very weird. I didn't see such thing explained in other parts of the documentation and I don't understand why it cannot be written with something like

fn buffer_elementwise_add[
    dt: DType
](lhs: DTypePointer[dt], rhs: DTypePointer[dt], result: DTypePointer[dt], N: Int):
    # Perform search!
    alias best_impl = search[
        fn_type=fn(DTypePointer[dt], DTypePointer[dt], DTypePointer[dt], Int)->None,
        buffer_elementwise_add_implementations[dt], 
        elementwise_evaluator[dt],
    ]()

    # Call the select implementation
    best_impl(lhs, rhs, result, N)

Anything else?

No response

[scottamain]

@arthurevans Please have a look at this.