[BUG | QST] CPU overhead regression with TVM-FFI

Which component has the problem?

CuTe DSL

Bug Report

Describe the bug
In v4.6 update, the following is changed in python/CuTeDSL/cutlass/cutlass_dsl/cutlass.py:

-                if kwargs_wrapper_spec.kwonly_names or kwargs_wrapper_spec.arg_defaults:
+                # Route through the kwargs-capable compiled class whenever
+                # the signature has kwonly/defaults OR any positional arg
+                # carries a dataclass instance (detected by the spec
+                # converter, including Union[...] over dataclasses). The
+                # kwargs wrapper is the only place tvm-ffi's
+                # ``map_dataclass_to_tuple`` unpack hook fires.
+                if (
+                    kwargs_wrapper_spec.kwonly_names
+                    or kwargs_wrapper_spec.arg_defaults
+                    or kwargs_wrapper_spec.arg_names
+                    or map_dataclass_to_tuple
+                ):

So basically if a @cute.jit function satisfies one of the following:

• kwonly_names: have keyword only arguments (indicated by *)
• arg_defaults: have default arguments
• arg_names: have any argument that is not a keyword only argument
• map_dataclass_to_tuple: have dataclass arguments

With these restrictions (especially kwonly_names and arg_names), basically only a CuTeDSL function with no arguments can go to the else branch and become a TVMFFIJitCompiledFunction which is a subclass of tvm_ffi.Function so they get to be called in C++ and bypass python completely. For CuTeDSL functions with arguments, they now have to go through the TVMFFIJitCompiledFunctionWithKwargs path.

This will slightly slow down calling a CuTeDSL function from python due to extra python layers from in TVMFFIJitCompiledFunctionWithKwargs, which isn't a big deal since python is slow anyway. However, if you register the CuTeDSL function via tvm-ffi and call it from C++, the overhead becomes obvious because you will leave C++ and enter python (TVMFFIJitCompiledFunctionWithKwargs.__call__) and then enter the compiled CuTeDSL function. This makes the CPU overhead regress compared to pre-4.6 version since there wasn't going through python with GIL and argument wrapper stuff.

I also wonder what's the motivation behind this. My guess is this makes a @cute.jit function behave more like a python function where it can use kwargs instead of can only be called using positional args, but meanwhile when we do want to call it with positional args (in C++) this only makes it slower.

Currently there is a hack to fix this, which is instead of registering the compiled CuTeDSL function, we can register its __tvm_ffi_object__ which would be a tvm_ffi.Function object, and we need to guarantee the args order is exactly the same in C++. I feel it's a bit hacky and fragile as cutlass evolves.

Another proposition of mine is we can tell users (probably in docs) that if they define CuTeDSL functions in this format: def func(a, b, c, /) with nothing after / and no default values, then we can optimize this case into a TVMFFIJitCompiledFunction because we know nobody will use kwargs in this case. I think if this fix is OK I can open a PR on it.

Not sure what's the opinion of the CUTLASS people?

Steps/Code to reproduce bug
I think my explanation should be clear enough so I don't think I need an example to reproduce.
You can just follow the example in https://github.com/NVIDIA/cutlass/pull/3289 and measure the CPU time to notice the difference.

Expected behavior
Should be faster?

Environment details (please complete the following information):
This regressed after v4.6 update

Additional context
Add any other context about the problem here.

Start in python/CuTeDSL/cutlass/cutlass_dsl/cutlass.py and trace TVMFFIJitCompiledFunction versus TVMFFIJitCompiledFunctionWithKwargs. Reproduce the CPU timing with the example from pull request 3289, comparing positional C++ calls and Python calls. Done means the regression is addressed without losing the reported keyword-argument behavior.