Fix conflicting multithreading and fork management

[binarman]

This PR disables multithreading in MLIR context after compilation ends. This is done to safely finalize thread pool implemented in MLIRContext. Not properly finalized thread pool can hang or crash in child process after fork.

[binarman]

Details about this issue

In compilation Triton creates native MLIRContext object managed by python:

triton/python/triton/compiler/compiler.py

Line 271 in f04df24

context = ir.context()

MLIRContext object holds pool of worker threads, which are allocated during compilation.

This object is not removed instantly after compilation ends. I've tried to remove it by python del statement, but seems that something referencing it and it is not trivial to delete on request.

If program forks right after compiling one kernel, MLIRContext has a chance to be duplicated along with thread pool inside. This invalidates pool in child process and subsequent attempt to delete MLIRContext hangs or crashes compilation.

[zhanglx13]

@binarman Then why this is not an issue on nv backend?

[binarman]

@zhanglx13
I've checked test from this PR and it hangs on NV backend as well if do not finalize thread pool.

My conclusion is that we were just lucky this issue did not showed before.

P.s. pay attention that I've disabled automatic GC in test, otherwise this test is flickering

[binarman]

@antiagainst please, take a look

[antiagainst]

Awesome findings! Just some nits around comments.

[alefimov-amd]

@antiagainst added more comments, PTAL

[ThomasRaoux]

how would you reproduce the problem?

[binarman]

@ThomasRaoux

hi!
It is in test_compile_in_forked_subproc_with_forced_gc test.

[antiagainst]

Drop the change in python/triton/compiler/compiler.py and run the test I believe.

[binarman]

@ThomasRaoux kind reminder