Added example of using external native types without modifying Warp

external/README.md

@@ -0,0 +1,38 @@
+# Overview
+
+An example of using external native types in Warp.
+
+The [wim](wim) subdirectory contains an independent library with [header-only types](wim/wim.h).  It also defines a [C-style public interface](wim/wim.cpp) used for the [Python bindings](wim/__init__.py).  Normally, there would be more code there, but this is a minimal viable example.  For simplicity, importing the `wim` Python module will build and load the native library and initialize the Python bindings.  This will happen automatically when running the example, so no need to build it separately.
+
+The file [wim_warp.h](wim_warp.h) is a header that will be included by Warp when building kernels.  It imports the types into the `wp` namespace, which is currently necessary, but may change in the future.  It also defines some useful functions that will be exposed to Warp code generation.
+
+The file [wim_types.py](wim_types.py) defines the Python versions of the custom types and registers the native utility functions as buitin functions that are available in kernels.
+
+The file [wim_paint.py](wim_paint.py) is the main program for the example.  It creates an image and draws shapes using Warp kernels.
+
+# Prerequisites
+
+* Linux is required
+* CUDA Toolkit installed in `/usr/local/cuda` (Note that `cuda_path` can be modified in [wim/__init__.py](wim/__init__.py)).  This is needed for building the "external" `wim` library.
+* `pip install matplotlib` for showing and saving the generated image.
+* `pip install torch` for an optional interop example!
+
+# Running
+
+From the repo root:
+
+```bash
+$ python external/wim_paint.py
+```
+
+If PyTorch is installed, the example will also demonstrate inverting the image using PyTorch.
+
+# Limitations and Future Work
+
+* The fact that Warp builds kernels with its own custom CRT might be a stumbling block for users who want to include their own external headers.  For example, including standard library headers fails during kernel compilation.
+
+* Currently, Warp doesn't have "proper" support for custom native types.  Our codegen assumes that all types and builtins are in the `wp::` namespace.  I was able to hack around that, but it's not clean (or clear to external users).
+
+* Warp supports accessing public struct/class members using `Type.vars`, but there's no way to expose getters and setters (or other methods) from native classes.  Free functions/builtins can be used to get around it, but some users might prefer OO syntax.
+
+* Using custom native types in Warp arrays is not fully supported yet.  In this example, I substituted the built-in `vec3f` for `Color` to expose the image data as a Warp array, but this kind of substitution would not work for all cases.  

external/wim/__init__.py

@@ -0,0 +1,39 @@
+import ctypes
+import os
+import subprocess
+
+_lib_dir = os.path.abspath(os.path.dirname(__file__))
+_lib_path = os.path.join(_lib_dir, "wim.so")
+
+
+def _build_lib(cuda_path="/usr/local/cuda"):
+    build_cmd = [os.path.join(cuda_path, "bin", "nvcc"),
+                "-shared",
+                "-Xcompiler", "-fPIC",
+                os.path.join(_lib_dir, "wim.cpp"),
+                "-o", _lib_path]
+    subprocess.run(build_cmd, check=True)
+
+
+def _load_lib():
+    lib_dir = os.path.abspath(os.path.dirname(__file__))
+    return ctypes.CDLL(os.path.join(lib_dir, _lib_path))
+
+
+# build the lib
+_build_lib()
+
+# load the lib
+_core = _load_lib()
+
+# bindings for CPU images
+_core.create_image_cpu.argtypes = [ctypes.c_int, ctypes.c_int]
+_core.create_image_cpu.restype = ctypes.c_void_p
+_core.destroy_image_cpu.argtypes = [ctypes.c_void_p]
+_core.destroy_image_cpu.restype = None
+
+# bindings for GPU images
+_core.create_image_cuda.argtypes = [ctypes.c_int, ctypes.c_int, ctypes.c_int]
+_core.create_image_cuda.restype = ctypes.c_void_p
+_core.destroy_image_cuda.argtypes = [ctypes.c_int, ctypes.c_void_p]
+_core.destroy_image_cuda.restype = None

external/wim/wim.cpp

@@ -0,0 +1,78 @@
+#include "wim.h"
+
+#include <cuda_runtime.h>
+
+#include <unordered_map>
+
+#if defined(_WIN32)
+    #define WIM_API __declspec(dllexport)
+#else
+    #define WIM_API __attribute__ ((visibility ("default")))
+#endif
+
+#define check_cuda(code) (wim::check_cuda_result(code, __FILE__, __LINE__))
+
+// internal stuff
+namespace wim
+{
+
+bool check_cuda_result(cudaError_t code, const char* file, int line)
+{
+    if (code == cudaSuccess)
+        return true;
+
+    fprintf(stderr, "WIM CUDA error %u: %s (%s:%d)\n", unsigned(code), cudaGetErrorString(code), file, line);
+    return false;
+}
+
+}
+
+// API for Python bindings
+extern "C"
+{
+
+WIM_API wim::Image* create_image_cpu(int width, int height)
+{
+    wim::Color* data = new wim::Color[width * height];
+    wim::Image* img = new wim::Image(width, height, data);
+    return img;
+}
+
+WIM_API void destroy_image_cpu(wim::Image* img)
+{
+    if (img)
+    {
+        delete [] img->getData();
+        delete img;
+    }
+}
+
+WIM_API wim::Image* create_image_cuda(int device, int width, int height)
+{
+    if (!check_cuda(cudaSetDevice(device)))
+        return nullptr;
+
+    wim::Color* data = nullptr;
+    if (!check_cuda(cudaMalloc(&data, width * height * sizeof(wim::Color))))
+        return nullptr;
+    if (!check_cuda(cudaMemset(data, 0, width * height * sizeof(wim::Color))))
+        return nullptr;
+
+    wim::Image* img = new wim::Image(width, height, data);
+    return img;
+}
+
+WIM_API void destroy_image_cuda(int device, wim::Image* img)
+{
+    if (!img)
+        return;
+
+    if (!check_cuda(cudaSetDevice(device)))
+        return;
+
+    check_cuda(cudaFree(img->getData()));
+
+    delete img;
+}
+
+}

external/wim/wim.h

@@ -0,0 +1,110 @@
+#pragma once
+
+#if !defined(__CUDACC__)
+    #define CUDA_CALLABLE
+    #define CUDA_CALLABLE_DEVICE
+#else
+    #define CUDA_CALLABLE __host__ __device__ 
+    #define CUDA_CALLABLE_DEVICE __device__
+#endif
+
+// our amazing Wraparound IMage lib
+namespace wim
+{
+
+struct Color
+{
+    float r, g, b;
+
+    CUDA_CALLABLE Color(float r = 0.0f, float g = 0.0f, float b = 0.0f)
+        : r(r), g(g), b(b)    
+    {
+    }
+};
+
+struct Coord
+{
+    int x, y;
+
+    CUDA_CALLABLE Coord(int x = 0, int y = 0)
+        : x(x), y(y)
+    {
+    }
+};
+
+class Image
+{
+    int mWidth;
+    int mHeight;
+    Color* mData;
+
+public:
+    CUDA_CALLABLE Image()
+        : mWidth(0), mHeight(0), mData(nullptr)
+    {
+    }
+
+    CUDA_CALLABLE Image(int width, int height, Color* data)
+        : mWidth(width), mHeight(height), mData(data)
+    {
+    }
+
+    CUDA_CALLABLE int getWidth() const
+    {
+        return mWidth;
+    }
+
+    CUDA_CALLABLE int getHeight() const
+    {
+        return mHeight;
+    }
+
+    CUDA_CALLABLE const Color* getData() const
+    {
+        return mData;
+    }
+
+    CUDA_CALLABLE Color* getData()
+    {
+        return mData;
+    }
+
+    CUDA_CALLABLE Coord wrapCoord(const Coord& coord) const
+    {
+        int x = coord.x;
+        int y = coord.y;
+
+        while (x < 0)
+            x += mWidth;
+        while (x >= mWidth)
+            x -= mWidth;
+
+        while (y < 0)
+            y += mHeight;
+        while (y >= mHeight)
+            y -= mHeight;
+
+        return Coord(x, y);
+    }
+
+    CUDA_CALLABLE Color getPixel(const Coord& coord) const
+    {
+        if (mData)
+        {
+            Coord wc = wrapCoord(coord);
+            return mData[wc.y * mWidth + wc.x];
+        }
+        return Color(1.0f, 0.0f, 1.0f);
+    }
+
+    CUDA_CALLABLE void setPixel(const Coord& coord, const Color& color)
+    {
+        if (mData)
+        {
+            Coord wc = wrapCoord(coord);
+            mData[wc.y * mWidth + wc.x] = color;
+        }
+    }
+};
+
+} // end of namespace wim