From 216333231dc8999dbfaeb7eb98ca934567610a79 Mon Sep 17 00:00:00 2001
From: Lingfan Yu <ly979@nyu.edu>
Date: Sat, 5 Jan 2019 04:45:10 +0000
Subject: [PATCH 1/2] make program thread local

---
 python/dgl/runtime/ir/program.py | 21 ++++++++++++++++-----
 1 file changed, 16 insertions(+), 5 deletions(-)

diff --git a/python/dgl/runtime/ir/program.py b/python/dgl/runtime/ir/program.py
index fa17937fb050..3ce89f25d9ff 100644
--- a/python/dgl/runtime/ir/program.py
+++ b/python/dgl/runtime/ir/program.py
@@ -2,6 +2,7 @@
 from __future__ import absolute_import
 
 from contextlib import contextmanager
+import threading
 
 from .registry import IR_REGISTRY
 
@@ -44,18 +45,28 @@ def pprint(self):
         for exe in self.execs:
             self.pprint_exe(exe)
 
+class CurrentProgram(threading.local):
+
+    def __init__(self):
+        super(CurrentProgram, self).__init__()
+        self.prog = None
+
+    def get_prog(self):
+        return self.prog
+
+    def set_prog(self, program):
+        self.prog = program
+
 # current program
-CURRENT_PROG = None
+CURRENT_PROG = CurrentProgram()
 
 def get_current_prog():
     """Get the current program."""
-    global CURRENT_PROG
-    return CURRENT_PROG
+    return CURRENT_PROG.get_prog()
 
 def set_current_prog(program):
     """Set the current program."""
-    global CURRENT_PROG
-    CURRENT_PROG = program
+    CURRENT_PROG.set_prog(program)
 
 @contextmanager
 def prog():

From c3784819cfa5c10e924e320b549f4ae8bae38778 Mon Sep 17 00:00:00 2001
From: Lingfan Yu <ly979@nyu.edu>
Date: Sat, 5 Jan 2019 05:54:52 +0000
Subject: [PATCH 2/2] doc string

---
 python/dgl/runtime/ir/program.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/python/dgl/runtime/ir/program.py b/python/dgl/runtime/ir/program.py
index 3ce89f25d9ff..0682f4d4ba06 100644
--- a/python/dgl/runtime/ir/program.py
+++ b/python/dgl/runtime/ir/program.py
@@ -46,15 +46,17 @@ def pprint(self):
             self.pprint_exe(exe)
 
 class CurrentProgram(threading.local):
-
+    """Thread local storage to keep the reference of current thread's program"""
     def __init__(self):
         super(CurrentProgram, self).__init__()
         self.prog = None
 
     def get_prog(self):
+        """Get program"""
         return self.prog
 
     def set_prog(self, program):
+        """Set program"""
         self.prog = program
 
 # current program