add custom_op executor to the default executor list

Signed-off-by: Masaki Kozuki <mkozuki@nvidia.com>

Author: crcrpar

thunder/__init__.py

@@ -199,9 +199,11 @@
 apex_executor: None | extend.Executor = extend.get_executor("apex")
 nvfuser_executor: None | extend.Executor = extend.get_executor("nvfuser")
 pytorch_executor: None | extend.Executor = extend.get_executor("torch")
+custom_op_executor: None | extend.Executor = extend.get_executor("custom_op")
 
-# Default executor list is [cudnn -> sdpa -> torchcompile_cat -> torchcompile_xentropy -> nvfuser -> torch -> python]
+# Default executor list is [cudnn -> sdpa -> torchcompile_cat -> torchcompile_xentropy -> nvfuser -> custom_op -> torch -> python]
 # Note that add_default_executor inserts executor at start of list, hence the reverse order below.
+add_default_executor(custom_op_executor)
 if nvfuser_executor:
     add_default_executor(nvfuser_executor)
 

thunder/extend/__init__.py

@@ -533,6 +533,7 @@ def get_all_executors() -> tuple[Executor, ...]:
         apexex,
         cudnn_layernormex,
         cudnnex,
+        custom_op_ex,
         nvfuserex,
         pythonex,
         sdpaex,

thunder/tests/test_torch_library_custom_op.py

@@ -263,58 +263,3 @@ def mul_translator(a, b, c=None, *, fd, lc_to_nv_map):
             if bsym.sym.name == f"{_symbol.name}_backward" and bsym.sym.executor is custom_op_ex:
                 bsym_custom_ex_bsym_found = True
         assert bsym_custom_ex_bsym_found
-
-
-def test_custom_op_executor_cleanup():
-    """Test that custom_op executor is properly removed from default executors after deregistration.
-    
-    This is a regression test for the issue where the custom_op executor would remain in the
-    default executors list after all custom ops were deregistered, causing failures in tests
-    that check the expected executor list.
-    
-    The original issue manifested when running test_recipes.py after test_torch_library_custom_op.py.
-    The test_recipes tests use get_expected_executors() which filters thunder.get_default_executors(),
-    and they expect that only executors actually used by the model are present. When custom_op
-    executor wasn't properly cleaned up, it would remain in the default executors list even though
-    no custom ops were registered, causing assertions like:
-        assert ex.name in [el.name for el in cd.executors_list]
-    to fail because 'custom_op' was in get_expected_executors() but not in cd.executors_list.
-    """
-    import thunder
-    
-    # Define a test custom op
-    @torch.library.custom_op("test_cleanup::mul", mutates_args=())
-    def cleanup_mul(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
-        return a * b
-
-    @torch.library.register_kernel("test_cleanup::mul", "cpu")
-    def _(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
-        return a * b
-
-    @torch.library.register_fake("test_cleanup::mul")
-    def _(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
-        return torch.empty_like(a)
-    
-    # Get initial state (custom_op should not be in default executors)
-    initial_executors = [ex.name for ex in thunder.get_default_executors()]
-    assert "custom_op" not in initial_executors, "custom_op should not be in default executors initially"
-    
-    # Simulate what happens in test_torch_library_custom_op tests
-    # Register the custom op (this adds custom_op_ex to default executors)
-    symbol = _register_custom_op(cleanup_mul)
-    executors_after_register = [ex.name for ex in thunder.get_default_executors()]
-    assert "custom_op" in executors_after_register, "custom_op should be added after registration"
-    
-    # Simulate the cleanup that happens in the autouse fixture
-    # Without the fix, this would NOT remove custom_op_ex from default executors
-    _deregister_custom_op(cleanup_mul)
-    executors_after_deregister = [ex.name for ex in thunder.get_default_executors()]
-    
-    # This is the critical assertion that would fail with the bug:
-    # After deregistration, custom_op should be removed from default executors
-    assert "custom_op" not in executors_after_deregister, \
-        "custom_op should be removed from default executors when no custom ops remain"
-    
-    # Verify we're back to the initial state, which is what test_recipes.py expects
-    assert executors_after_deregister == initial_executors, \
-        "Should return to initial executor state after deregistration"

thunder/torch/custom_op.py

@@ -332,9 +332,6 @@ def _register_custom_op(custom_op: CustomOpDef) -> Symbol:
     .. note::
         This feature is experimental and subject to change.
     """
-    from thunder.extend import add_executor_lists
-    from thunder.extend import get_default_executors
-    from thunder.extend import set_default_executors
     from thunder.executors.torchex import _always_executable
     from thunder.executors.custom_op_ex import custom_op_ex
     from thunder.torch import register_function
@@ -412,12 +409,6 @@ def _register_custom_op(custom_op: CustomOpDef) -> Symbol:
         backward_op = custom_op_ex.register_operator(bwd_fn_name, meta=backward_meta, fn=backward_impl)
         register_backward(symbol.id)(backward_op)
 
-    # NOTE: `thunder.extend.add_default_executor` basically does `lst.insert(ex, 0)`.
-    if custom_op_ex not in get_default_executors():
-        default_executors = get_default_executors()
-        new_default_executors = add_executor_lists(default_executors, [custom_op_ex])
-        set_default_executors(new_default_executors)
-
     _CUSTOM_OP_TO_TORCHFN_AND_SYMBOL[custom_op] = ((torch_opoverload, torch_opoverload_packet), symbol)
 
     return symbol
@@ -450,12 +441,6 @@ def _deregister_custom_op(custom_op: CustomOpDef) -> None:
 
     del _CUSTOM_OP_TO_TORCHFN_AND_SYMBOL[custom_op]
 
-    # Remove custom_op_ex from default executors if no custom ops remain
-    if not _CUSTOM_OP_TO_TORCHFN_AND_SYMBOL:
-        from thunder.extend import remove_default_executor
-
-        remove_default_executor(custom_op_ex)
-
 
 def _register_nvfuser_translator(
     symbol: Symbol,