[WIP] WebGL Backend (apache#672)

Basic WebGL Backend

CMakeLists.txt

@@ -22,6 +22,7 @@ endif()
 
 tvm_option(USE_CUDA "Build with CUDA" OFF)
 tvm_option(USE_OPENCL "Build with OpenCL" OFF)
+tvm_option(USE_OPENGL "Build with OpenGL" OFF)
 tvm_option(USE_METAL "Build with Metal" OFF)
 tvm_option(USE_RPC "Build with RPC" ON)
 tvm_option(USE_GRAPH_RUNTIME "Build with tiny graph runtime" ON)
@@ -61,8 +62,8 @@ if(MSVC)
 else(MSVC)
   include(CheckCXXCompilerFlag)
   check_cxx_compiler_flag("-std=c++11"    SUPPORT_CXX11)
-  set(CMAKE_C_FLAGS "-O3 -Wall -std=c++11 -fPIC")
-  set(CMAKE_CXX_FLAGS ${CMAKE_C_FLAGS})
+  set(CMAKE_C_FLAGS "-O3 -Wall -fPIC")
+  set(CMAKE_CXX_FLAGS "${CMAKE_C_FLAGS} -std=c++11")
 endif(MSVC)
 
 # add source group
@@ -87,6 +88,7 @@ file(GLOB RUNTIME_SRCS src/runtime/*.cc)
 file(GLOB COMPILER_LLVM_SRCS src/codegen/llvm/*.cc)
 file(GLOB RUNTIME_CUDA_SRCS src/runtime/cuda/*.cc)
 file(GLOB RUNTIME_OPENCL_SRCS src/runtime/opencl/*.cc)
+file(GLOB RUNTIME_OPENGL_SRCS src/runtime/opengl/*.cc)
 file(GLOB RUNTIME_METAL_SRCS src/runtime/metal/*.mm)
 file(GLOB RUNTIME_RPC_SRCS src/runtime/rpc/*.cc)
 file(GLOB RUNTIME_GRAPH_SRCS src/runtime/graph/*.cc)
@@ -135,6 +137,18 @@ else(USE_OPENCL)
   add_definitions(-DTVM_OPENCL_RUNTIME=0)
 endif(USE_OPENCL)
 
+if(USE_OPENGL)
+  find_package(OpenGL QUIET REQUIRED)
+  find_package(glfw3 QUIET REQUIRED)
+  message(STATUS "Build with OpenGL support")
+  include_directories(${OPENGL_INCLUDE_DIRS})
+  list(APPEND TVM_RUNTIME_LINKER_LIBS ${OpenGL_LIBRARIES} glfw)
+  list(APPEND RUNTIME_SRCS ${RUNTIME_OPENGL_SRCS})
+  add_definitions(-DTVM_OPENGL_RUNTIME=1)
+else(USE_OPENGL)
+  add_definitions(-DTVM_OPENGL_RUNTIME=0)
+endif(USE_OPENGL)
+
 if(USE_METAL)
   find_package(OpenCL QUIET REQUIRED)
   message(STATUS "Build with Metal support")

Jenkinsfile

@@ -88,6 +88,7 @@ stage('Build') {
            echo USE_CUDNN=1 >> config.mk
            echo USE_CUDA=1 >> config.mk
            echo USE_OPENCL=1 >> config.mk
+           echo USE_OPENGL=1 >> config.mk
            echo LLVM_CONFIG=llvm-config-4.0 >> config.mk
            echo USE_RPC=1 >> config.mk
            echo USE_GRAPH_RUNTIME=1 >> config.mk
@@ -120,6 +121,7 @@ stage('Build') {
            echo USE_CUDA=0 >> config.mk
            echo USE_OPENCL=0 >> config.mk
            echo USE_RPC=0 >> config.mk
+           echo USE_OPENGL=1 >> config.mk
            echo LLVM_CONFIG=llvm-config-4.0 >> config.mk
            """
         make('cpu', '-j2')

Makefile

@@ -32,8 +32,8 @@ OBJCFLAGS = -fno-objc-arc
 EMCC_FLAGS= -std=c++11 -DDMLC_LOG_STACK_TRACE=0\
 	-Oz -s RESERVED_FUNCTION_POINTERS=2 -s MAIN_MODULE=1 -s NO_EXIT_RUNTIME=1\
 	-s EXTRA_EXPORTED_RUNTIME_METHODS="['cwrap','getValue','setValue','addFunction']"\
+	-s USE_GLFW=3 -s USE_WEBGL2=1 -lglfw\
 	$(INCLUDE_FLAGS)
-
 # llvm configuration
 ifdef LLVM_CONFIG
 	LLVM_VERSION=$(shell $(LLVM_CONFIG) --version| cut -b 1,3)
@@ -54,6 +54,7 @@ METAL_SRC = $(wildcard src/runtime/metal/*.mm)
 CUDA_SRC = $(wildcard src/runtime/cuda/*.cc)
 ROCM_SRC = $(wildcard src/runtime/rocm/*.cc)
 OPENCL_SRC = $(wildcard src/runtime/opencl/*.cc)
+OPENGL_SRC = $(wildcard src/runtime/opengl/*.cc)
 RPC_SRC = $(wildcard src/runtime/rpc/*.cc)
 GRAPH_SRC = $(wildcard src/runtime/graph/*.cc)
 RUNTIME_SRC = $(wildcard src/runtime/*.cc)
@@ -65,6 +66,7 @@ METAL_OBJ = $(patsubst src/%.mm, build/%.o, $(METAL_SRC))
 CUDA_OBJ = $(patsubst src/%.cc, build/%.o, $(CUDA_SRC))
 ROCM_OBJ = $(patsubst src/%.cc, build/%.o, $(ROCM_SRC))
 OPENCL_OBJ = $(patsubst src/%.cc, build/%.o, $(OPENCL_SRC))
+OPENGL_OBJ = $(patsubst src/%.cc, build/%.o, $(OPENGL_SRC))
 RPC_OBJ = $(patsubst src/%.cc, build/%.o, $(RPC_SRC))
 GRAPH_OBJ = $(patsubst src/%.cc, build/%.o, $(GRAPH_SRC))
 CC_OBJ = $(patsubst src/%.cc, build/%.o, $(CC_SRC)) $(LLVM_OBJ)
@@ -119,6 +121,19 @@ else
 	CFLAGS += -DTVM_OPENCL_RUNTIME=0
 endif
 
+ifeq ($(USE_OPENGL), 1)
+	CFLAGS += -DTVM_OPENGL_RUNTIME=1
+	EMCC_FLAGS += -DTVM_OPENGL_RUNTIME=1
+	ifeq ($(UNAME_S), Darwin)
+		FRAMEWORKS += -framework OpenGL
+	else
+		LDFLAGS += -lGL -lglfw
+	endif
+	RUNTIME_DEP += $(OPENGL_OBJ)
+else
+	CFLAGS += -DTVM_OPENGL_RUNTIME=0
+endif
+
 ifeq ($(USE_METAL), 1)
 	CFLAGS += -DTVM_METAL_RUNTIME=1
 	LDFLAGS += -lobjc

include/tvm/runtime/c_runtime_api.h

@@ -55,9 +55,11 @@ typedef int64_t tvm_index_t;
 
 /*! \brief Extension device types in TVM */
 typedef enum {
+  kOpenGL = 11,
+
   // Extension DRAM type, used for quickly test extension device
   // The device api can differ depending on the xpu driver registered.
-  kExtDev = 12
+  kExtDev = 12,
   // AddExtraTVMType which is not in DLPack here
 } TVMDeviceExtType;
 

include/tvm/runtime/device_api.h

@@ -55,11 +55,16 @@ class DeviceAPI {
   /*!
    * \brief Allocate a data space on device.
    * \param ctx The device context to perform operation.
-   * \param size The size of the memory
+   * \param nbytes The number of bytes in memory.
    * \param alignment The alignment of the memory.
-   * \return The allocated device pointer
+   * \param type_hint The type of elements. Only needed by certain backends such
+   * as OpenGL, as nbytes & alignment are sufficient for most backends.
+   * \return The allocated device pointer.
    */
-  virtual void* AllocDataSpace(TVMContext ctx, size_t size, size_t alignment) = 0;
+  virtual void* AllocDataSpace(TVMContext ctx,
+                               size_t nbytes,
+                               size_t alignment,
+                               TVMType type_hint) = 0;
   /*!
    * \brief Free a data space on device.
    * \param ctx The device context to perform operation.

include/tvm/schedule.h

@@ -213,6 +213,11 @@ class Stage : public NodeRef {
    * \return reference to self.
    */
   Stage& double_buffer();   // NOLINT(*)
+  /*!
+   * \brief Schedule for OpenGL fragment shader.
+   * \return reference to self.
+   */
+  Stage& opengl(); // NOLINT(*)
   /*!
    * \brief whether the stage has been scheduled.
    * \return whether the stage has been scheduled.

python/tvm/__init__.py

@@ -17,7 +17,7 @@
 from . import target
 
 from . import ndarray as nd
-from .ndarray import context, cpu, gpu, opencl, cl, metal, mtl, vpi, rocm, ext_dev
+from .ndarray import context, cpu, gpu, opencl, cl, metal, mtl, vpi, rocm, opengl, ext_dev
 
 from ._ffi.runtime_ctypes import TypeCode
 from ._ffi.function import Function

python/tvm/_ffi/runtime_ctypes.py

@@ -97,6 +97,7 @@ class TVMContext(ctypes.Structure):
         8 : 'metal',
         9 : 'vpi',
         10: 'rocm',
+        11: 'opengl',
         12: 'ext_dev',
     }
     STR2MASK = {
@@ -111,6 +112,7 @@ class TVMContext(ctypes.Structure):
         'metal': 8,
         'vpi': 9,
         'rocm': 10,
+        'opengl': 11,
         'ext_dev': 12,
     }
     def __init__(self, device_type, device_id):

python/tvm/contrib/rpc.py

@@ -285,6 +285,10 @@ def metal(self, dev_id=0):
         """Construct remote Metal device."""
         return self.context(8, dev_id)
 
+    def opengl(self, dev_id=0):
+        """Construct remote OpenGL device."""
+        return self.context(11, dev_id)
+
     def ext_dev(self, dev_id=0):
         """Construct remote extension device."""
         return self.context(12, dev_id)

python/tvm/ndarray.py

@@ -120,6 +120,21 @@ def vpi(dev_id=0):
     """
     return TVMContext(9, dev_id)
 
+def opengl(dev_id=0):
+    """Construct a OpenGL device
+
+    Parameters
+    ----------
+    dev_id : int, optional
+        The integer device id
+
+    Returns
+    -------
+    ctx : TVMContext
+        The created context
+    """
+    return TVMContext(11, dev_id)
+
 def ext_dev(dev_id=0):
     """Construct a extension device
 

python/tvm/schedule.py

@@ -611,4 +611,11 @@ def double_buffer(self):
         """
         _api_internal._StageDoubleBuffer(self)
 
+    def opengl(self):
+        """The special OpenGL schedule
+
+        Maps each output element to a pixel.
+        """
+        _api_internal._StageOpenGL(self)
+
 _init_api("tvm.schedule")

python/tvm/target.py

@@ -67,7 +67,7 @@ class Target(object):
 
     Parameters
     ----------
-    target_name : {"llvm", "cuda", "opencl", "metal", "rocm", "stackvm", "ext_dev"}
+    target_name : {"llvm", "cuda", "opencl", "metal", "rocm", "stackvm", "opengl", "ext_dev"}
         The major target name.
 
     options : list of str, optional
@@ -119,6 +119,8 @@ def __init__(self,
         elif target_name in ("metal",):
             self.keys += ("gpu",)
             self.max_num_threads = 256
+        elif target_name in ("opengl",):
+            self.keys += ("opengl",)
         elif target_name in ("stackvm", "ext_dev"):
             # Do not now class for stacvm or ext_dev
             pass

src/api/api_lang.cc

@@ -399,6 +399,11 @@ TVM_REGISTER_API("_StageDoubleBuffer")
     args[0].operator Stage().double_buffer();
   });
 
+TVM_REGISTER_API("_StageOpenGL")
+  .set_body([](TVMArgs args, TVMRetValue *ret) {
+    args[0].operator Stage().opengl();
+  });
+
 TVM_REGISTER_API("_ScheduleNormalize")
 .set_body([](TVMArgs args, TVMRetValue* ret) {
     *ret = args[0].operator Schedule()

src/codegen/build_opengl.cc

@@ -0,0 +1,35 @@
+/*!
+ *  Copyright (c) 2017 by Contributors
+ *  Build opengl modules from source.
+ * \file build_opengl.cc
+ */
+#include <tvm/base.h>
+#include "./codegen_opengl.h"
+#include "./build_common.h"
+
+namespace tvm {
+namespace codegen {
+
+runtime::Module BuildOpenGL(Array<LoweredFunc> funcs) {
+  bool output_ssa = false;
+  CodeGenOpenGL cg;
+  cg.Init(output_ssa);
+  for (LoweredFunc f : funcs) {
+    cg.AddFunction(f);
+  }
+  auto shaders = cg.Finish();
+#if TVM_OPENGL_RUNTIME
+  return OpenGLModuleCreate(shaders, "gl", ExtractFuncInfo(funcs));
+#else
+  LOG(WARNING) << "OpenGL runtime not enabled, return a source module...";
+  auto data = ToJSON(shaders);
+  return DeviceSourceModuleCreate(data, "gl", ExtractFuncInfo(funcs), "opengl");
+#endif  // TVM_OPENGL_RUNTIME
+}
+
+TVM_REGISTER_API("codegen.build_opengl")
+.set_body([](TVMArgs args, TVMRetValue* rv) {
+  *rv = BuildOpenGL(args[0]);
+});
+}  // namespace codegen
+}  // namespace tvm

src/codegen/codegen_c.h

@@ -150,7 +150,7 @@ class CodeGenC :
   std::string GetStructRef(
       Type t, const Expr& buffer, const Expr& index, int kind);
   // print reference to a buffer as type t in index.
-  std::string GetBufferRef(
+  virtual std::string GetBufferRef(
       Type t, const Variable* buffer, Expr index);
   /*!
    * \brief If buffer is allocated as type t.