viable/strict/1781110191: Optimize mutable custom op version bump dispatch (#186175)

Mutable torch.library custom ops install an ADInplaceOrView shim to bump

the version counter for mutated inputs. The shim already precomputed which

schema arguments were mutable, but it still called utils.fill_defaults on

every invocation. That walked every schema argument before bumping only the

mutated values, so mutable custom op overhead grew with total argument count.

Use the dispatcher-provided args and kwargs directly in the hot path, and skip

omitted mutated arguments instead of materializing or incrementing their default

values. Default values for mutated arguments are expected to be non-Tensors or

fresh values, so there is no useful version counter to bump when the caller

omitted them.

This substantially reduces the overhead reported in #139494, but does not claim

to fully close the remaining gap between mutable and non-mutating custom ops.

The alternative was to optimize utils.fill_defaults itself, but that helper is

general schema normalization and the mutable version-bump path only needs the

small subset of mutated arguments that were actually provided.

Generated by my agent

Benchmark Results:

Before, issue benchmark on CUDA H100, 1000 calls per do_bench function:

• mutate2 = 18.27797737121582


• no_mutate2 = 13.457284654889788


• mutate = 66.9056625366211


• no_mutate = 17.066380310058594

After:

• mutate2 = 16.13766403198242


• no_mutate2 = 13.513686997549874


• mutate = 21.02779197692871


• no_mutate = 17.362232971191407

Test Plan:

• python test/test_custom_ops.py TestCustomOpAPI.test_mutated_version_bump_does_not_fill_all_defaults TestCustomOpAPI.test_mutated_optional_arg_default_none TestCustomOpAPI.test_mutated TestCustomOpAPI.test_custom_op_out_tag


• lintrunner -a
  
  Pull Request resolved: #186175
  
  Approved by: https://github.com/zou3519