Add enable_options and disable_options to fd.execute

[Priya2698]

This PR adds enable_options and disable_options to fd.execute to allow setting them through the python frontend in lieu of the environment variables.
This work will be used to allow enabling nvfuser matmul codegen from within Thunder.

Inspired by @jacobhinkle's PR #1905!

Tracking Issue: #3022

[Priya2698]

!build

[jacobhinkle]

Thanks for taking the initiative with this! We will also need to update FusionExecutorCache to be option-aware so that we don't re-use kernels compiled with different options. That was part of the idea with #2077: separating options (features) into ones that affect each level of caching or not.

[Priya2698]

Thanks for taking the initiative with this! We will also need to update FusionExecutorCache to be option-aware so that we don't re-use kernels compiled with different options. That was part of the idea with #2077: separating options (features) into ones that affect each level of caching or not.

Got it -- we can add that support separately, what do you think? That will allow development on #3022. IIUC, we are currently missing this for the environment variables based approach as well so we may be reusing kernels which used different env variables?

[jacobhinkle]

IIUC, we are currently missing this for the environment variables based approach as well so we may be reusing kernels which used different env variables?

I guess this is true: a process can change its env vars at any time and we will not recognize that change after the first time we look up an option. However that is off-label use. If we are providing a kwarg to execute then a user would probably assume it's fully supported.

I wonder if we could do something quickly like use a separate FusionCache for each set of options we receive in the frontend. If no options are provided we'd use the default. That could be done entirely in python.

[Priya2698]

I wonder if we could do something quickly like use a separate FusionCache for each set of options we receive in the frontend. If no options are provided we'd use the default. That could be done entirely in python.

What do you mean by using separate FusionCache?

A fusion cache is initialized here:

Fuser/csrc/python_frontend/fusion_definition.cpp

Lines 58 to 67 in f6975f3

	FusionDefinition::FusionDefinition(std::optional<size_t> id, size_t max_length)
	: FusionState(),
	max_length_(max_length),
	fusion_id_(id),
	fusion_cache_(FusionCache::get()),
	trie_node_(nullptr),
	prev_fusion_(nullptr),
	user_sched_(nullptr),
	ops(this),
	sched(this) {}

. The options are passed in the execute method so they are not available when the fusion definition is initialized.

I am still going through the FusionCache files.

[jacobhinkle]

. The options are passed in the execute method so they are not available when the fusion definition is initialized.

Hmmm. Yeah I guess you're right. FusionCache is not the right place to do this. Instead we might do it by having a separate FusionExecutorCache might work:

Fuser/csrc/python_frontend/fusion_definition.cpp

Lines 387 to 390 in 5db18de

	if (outputs.empty()) {
	outputs = scheds->auto_gen_schedules->runFusionWithInputs(
	inputs, std::nullopt, selected_device);
	}

Here scheds->auto_gen_schedules is a unique_ptr<FusionExecutorCache>. This is fine as a default but when features are supplied, we should create a new cache and use it instead I think.

[Priya2698]

As discussed offline, moving forward with renaming the kwargs to _enable/_disable_options. I added a warning about the possibility of kernel reuse so that we are aware of this when utilizing this feature.

[Priya2698]

!build

[Priya2698]

!build

[jacobhinkle]

LGTM. Thanks for making it clear that this is a somewhat unsafe option at least until we figure out the caching problem.

[Priya2698]

!test

[Priya2698]

!test --serde

[Priya2698]

!test