[Tuner] Fix context management

[kuhar]

Currently, the MLIR context management in the tuner is functional but weird:

shark-ai/tuner/tuner/libtuner.py

Lines 1055 to 1098 in aaee29a

	try:
	# Strip compilation info before generating td_specs, since the generated
	# td_specs can end up matching against the compilation info from the
	# source mlir.
	mlir_text = strip_compilation_info(path_config.template_mlir)
	mlir_module = dispatch_parser.parse_mlir(mlir_text, tuning_client.tuner_context)
	with tuning_client.tuner_context.mlir_ctx:
	logging.debug("Captured messages from candidate_gen.py:")
	config_specs: list[ir.Module] = candidate_gen.generate_configs_and_td_specs(
	input_module=mlir_module,
	tuner_context=tuning_client.tuner_context,
	limit=args.num_candidates,
	num_subgroups=args.num_subgroups,
	)
	logging.debug("candidate_gen.py ends")
	handle_error(
	condition=(len(config_specs) <= 1), msg="Failed to generate any candidates"
	)
	# Create candidate trackers.
	candidates = []
	for candidate_num, spec in enumerate(config_specs):
	candidates.append(candidate_num)
	# Move the specs to the canonical path_config location.
	spec_path = path_config.specs_dir / path_config.get_candidate_spec_filename(
	candidate_num
	)
	with open(spec_path, "w") as f:
	f.write(str(spec))
	new_candidate = CandidateTracker(
	mlir_path=path_config.template_mlir,
	candidate_id=candidate_num,
	spec_path=spec_path,
	)
	candidate_trackers.append(new_candidate)
	except Exception as e:
	logging.error("An error occurred during candidates generation: %s", str(e))
	# Capture and log debug messages from candidate_gen.py.
	tune_logger = logging.getLogger("tune_with_td")
	for handler in logging.getLogger().handlers:
	if isinstance(handler, logging.FileHandler):
	tune_logger.handlers.append(handler)
	tune_logger.exception("Error in candidate_gen.py:")
	raise

We use a with statement to put the existing mlir context onto the stack, and after that the context outlives the with scope. This is unusual because the with statement is generally destructive in python: https://docs.python.org/3/reference/compound_stmts.html#the-with-statement:~:text=with%20EXPRESSION%20as%20TARGET%3A%0A%20%20%20%20SUITE-,is%20semantically%20equivalent%20to%3A,-manager%20%3D%20(EXPRESSION)%0Aenter

This is only safe because of how __enter__ and __exit__ are defined for the MLIR context class:

• https://github.com/llvm/llvm-project/blob/2916352936097a35cdcaaf38a9097465adbf8cf5/mlir/lib/Bindings/Python/IRCore.cpp#L720-L728
• https://github.com/llvm/llvm-project/blob/2916352936097a35cdcaaf38a9097465adbf8cf5/mlir/lib/Bindings/Python/IRCore.cpp#L877-L893

For more introductory info, see https://mlir.llvm.org/docs/Bindings/Python/#context-management.

To make it more idiomatic, we should:

1. Increase the scope of the with statement in the tuner
2. Make the whole tuner context compatible with the context management by defining __enter__ and __exit__ functions
3. Making the TunerContext constructor create MLIR context and cleaning up tuner context test fixtures

[kuhar]

cc: @Max191