Choose a GPU in main.py

main.py

@@ -118,9 +118,12 @@ def main(args, prof):
     frames = [('%s%05d%s' % (prefix_plus, i, args.suffix), t)
               for i, t in profile.enumerate_times(gprof)]
 
-    # We don't initialize a CUDA context until here. This keeps other
-    # functions like --help and --print snappy.
-    import pycuda.autoinit
+    import pycuda.driver as cuda
+    cuda.init()
+    dev = cuda.Device(args.device or 0)
+    cuctx = dev.make_context()
+
+    try:
       rmgr = render.RenderManager()
       rdr = render.Renderer(gnm, gprof)
 
@@ -154,6 +157,9 @@ def main(args, prof):
           pyglet_preview(args, gprof, render_iter())
       else:
           for i in render_iter(): pass
+    finally:
+      cuda.Context.pop()
+
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser(description='Render fractal flames.')
@@ -183,6 +189,8 @@ if __name__ == "__main__":
         help='Use half-loops when converting nodes to animations')
     parser.add_argument('--print', action='store_true',
         help="Print the blended animation and exit.")
+    parser.add_argument('--device', metavar='NUM', type=int,
+                        help="GPU device number to use (from nvidia-smi).")
     profile.add_args(parser)
 
     args = parser.parse_args()