From f23bf55f15a2dbbde83200bd002a1bcc144a2c22 Mon Sep 17 00:00:00 2001
From: Philippe Tillet <phil@openai.com>
Date: Fri, 3 Dec 2021 13:01:01 -0800
Subject: [PATCH] [RUNTIME] release the gil on launch (#383)

---
 python/src/triton.cc | 10 +++++++++-
 1 file changed, 9 insertions(+), 1 deletion(-)

diff --git a/python/src/triton.cc b/python/src/triton.cc
index 26c233287..7b6c1ce81 100644
--- a/python/src/triton.cc
+++ b/python/src/triton.cc
@@ -272,10 +272,18 @@ void init_triton_runtime(py::module &&m) {
         CU_LAUNCH_PARAM_END
     };
     uint64_t _stream = PyLong_AsLong(stream.ptr());
-    if(grid_0*grid_1*grid_2 > 0)
+    if(grid_0*grid_1*grid_2 > 0) {
+      // release the gil in case the enqueue blocks
+      // cuda will block if too many ops are enqueued
+      Py_BEGIN_ALLOW_THREADS
+
+
       drv::dispatch::cuLaunchKernel((CUfunction)kernel, grid_0, grid_1, grid_2, 
                                     _num_warps*32, 1, 1, shared_mem, (CUstream)_stream, 
                                      nullptr, config);
+
+       Py_END_ALLOW_THREADS
+   }
     return bin;
   });