From bd0e97eee1e2f4ab44e95368a2d645783ee49643 Mon Sep 17 00:00:00 2001 From: Lukasz Wawrzyniak Date: Wed, 14 Feb 2024 16:25:49 -0500 Subject: [PATCH 01/16] Added example of using external native types without modifying Warp --- external/README.md | 25 ++++++ external/wim/__init__.py | 35 ++++++++ external/wim/wim.cpp | 92 ++++++++++++++++++++ external/wim/wim.h | 110 +++++++++++++++++++++++ external/wim_paint.py | 155 +++++++++++++++++++++++++++++++++ external/wim_types.py | 184 +++++++++++++++++++++++++++++++++++++++ external/wim_warp.h | 68 +++++++++++++++ 7 files changed, 669 insertions(+) create mode 100644 external/README.md create mode 100644 external/wim/__init__.py create mode 100644 external/wim/wim.cpp create mode 100644 external/wim/wim.h create mode 100644 external/wim_paint.py create mode 100644 external/wim_types.py create mode 100644 external/wim_warp.h diff --git a/external/README.md b/external/README.md new file mode 100644 index 0000000000..78ef24b2eb --- /dev/null +++ b/external/README.md @@ -0,0 +1,25 @@ +# 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, running `import wim` will build and load the native library and initialize the Python bindings. + +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. + +# Running + +From the repo root: + +```python +$ python external/wim_paint.py +``` diff --git a/external/wim/__init__.py b/external/wim/__init__.py new file mode 100644 index 0000000000..1abd88ecbd --- /dev/null +++ b/external/wim/__init__.py @@ -0,0 +1,35 @@ +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 and set up Python bindings +_core = _load_lib() +_core.create_image_cpu.argtypes = [ctypes.c_int, ctypes.c_int, ctypes.POINTER(ctypes.c_void_p)] +_core.create_image_cpu.restype = ctypes.c_void_p +_core.destroy_image_cpu.argtypes = [ctypes.c_void_p] +_core.destroy_image_cpu.restype = None +_core.create_image_cuda.argtypes = [ctypes.c_int, ctypes.c_int, ctypes.c_int, ctypes.POINTER(ctypes.c_void_p)] +_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 diff --git a/external/wim/wim.cpp b/external/wim/wim.cpp new file mode 100644 index 0000000000..cc849b4609 --- /dev/null +++ b/external/wim/wim.cpp @@ -0,0 +1,92 @@ +#include "wim.h" + +#include + +#include + +#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_ret) +{ + wim::Color* data = new wim::Color[width * height]; + wim::Image* img = new wim::Image(width, height, data); + + if (data_ret) + *data_ret = data; + + return img; +} + +WIM_API void destroy_image_cpu(wim::Image* img) +{ + delete [] img->getData(); + delete img; +} + +WIM_API wim::Image* create_image_cuda(int device, int width, int height, wim::Color** data_ret) +{ + 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_h(width, height, data); + + wim::Image* img_d = nullptr; + if (!check_cuda(cudaMalloc(&img_d, sizeof(wim::Image)))) + return nullptr; + if (!check_cuda(cudaMemcpy(img_d, &img_h, sizeof(wim::Image), cudaMemcpyHostToDevice))) + return nullptr; + + if (data_ret) + *data_ret = data; + + return img_d; +} + +WIM_API void destroy_image_cuda(int device, wim::Image* img_d) +{ + if (!img_d) + return; + + if (!check_cuda(cudaSetDevice(device))) + return; + + wim::Image img_h; + if (!check_cuda(cudaMemcpy(&img_h, img_d, sizeof(wim::Image), cudaMemcpyDeviceToHost))) + return; + + cudaFree(img_h.getData()); + cudaFree(img_d); +} + +} diff --git a/external/wim/wim.h b/external/wim/wim.h new file mode 100644 index 0000000000..537cf5f456 --- /dev/null +++ b/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 diff --git a/external/wim_paint.py b/external/wim_paint.py new file mode 100644 index 0000000000..555ad57747 --- /dev/null +++ b/external/wim_paint.py @@ -0,0 +1,155 @@ +import os +import numpy as np +import warp as wp + +import wim_types +from wim_types import Coord, Color, Image + + +@wp.kernel +def print_image_info_kernel(img_handle: wp.uint64): + width = wp.img_width(img_handle) + height = wp.img_height(img_handle) + wp.printf("Image: %dx%d\n", width, height) + + +@wp.kernel +def fill_kernel(img_handle: wp.uint64, color: Color): + y, x = wp.tid() + coord = wp.Coord_(x, y) + wp.img_set_pixel(img_handle, coord, color) + + +@wp.kernel +def fill_rect_kernel(img_handle: wp.uint64, half_width: int, half_height: int, pos: Coord, color: Color): + j, i = wp.tid() + i -= half_width + j -= half_height + coord = wp.Coord_(pos.x + i, pos.y + j) + wp.img_set_pixel(img_handle, coord, color) + + +@wp.kernel +def fill_circle_kernel(img_handle: wp.uint64, radius: int, pos: Coord, color: Color): + j, i = wp.tid() + i -= radius + j -= radius + if i * i + j * j <= radius * radius: + x = pos.x + i + y = pos.y + j + coord = wp.Coord_(x, y) + wp.img_set_pixel(img_handle, coord, color) + + +@wp.kernel +def blur_kernel(img_handle: wp.uint64): + y, x = wp.tid() + + c00 = wp.img_get_pixel(img_handle, wp.Coord_(x - 1, y - 1)) + c01 = wp.img_get_pixel(img_handle, wp.Coord_(x, y - 1)) + c02 = wp.img_get_pixel(img_handle, wp.Coord_(x + 1, y - 1)) + c10 = wp.img_get_pixel(img_handle, wp.Coord_(x - 1, y)) + c11 = wp.img_get_pixel(img_handle, wp.Coord_(x, y)) + c12 = wp.img_get_pixel(img_handle, wp.Coord_(x + 1, y)) + c20 = wp.img_get_pixel(img_handle, wp.Coord_(x - 1, y + 1)) + c21 = wp.img_get_pixel(img_handle, wp.Coord_(x, y + 1)) + c22 = wp.img_get_pixel(img_handle, wp.Coord_(x + 1, y + 1)) + + c = (c00 + c02 + c20 + c22) + 2.0 * (c01 + c21 + c10 + c12) + 4.0 * c11 + c = (1.0 / 16.0) * c + + wp.img_set_pixel(img_handle, wp.Coord_(x, y), c) + + +def fill(img, color: Color): + img_handle = wp.uint64(img.ptr) + img_shape = (img.height, img.width) + wp.launch(fill_kernel, dim=img_shape, inputs=[img_handle, color]) + + +def draw_rect(img, width: int, height: int, pos: Coord, color: Color): + img_handle = wp.uint64(img.ptr) + launch_dim = (height, width) + wp.launch(fill_rect_kernel, dim=launch_dim, inputs=[img_handle, width//2, height//2, pos, color]) + + +def draw_circle(img, radius: int, pos: Coord, color: Color): + img_handle = wp.uint64(img.ptr) + launch_dim = (2 * radius, 2 * radius) + wp.launch(fill_circle_kernel, dim=launch_dim, inputs=[img_handle, radius, pos, color]) + + +def blur(img): + img_handle = wp.uint64(img.ptr) + img_shape = (img.height, img.width) + wp.launch(blur_kernel, dim=img_shape, inputs=[img_handle]) + + +def draw_picture(img): + + # background color + fill(img, Color(0.3, 0.0, 0.3)) + + # concentric circles in the corners + for iter in range(10): + g = iter / 10 + r = 20 + (10 - iter - 1) * 20 + draw_circle(img, r, Coord(0, 0), Color(0, g, 1)) + + for _ in range(500): + blur(img) + + for iter in range(10): + # rectangle crossing the vertical edges + if iter == 0: + draw_rect(img, 200, 300, Coord(img.width//2, 0), Color(1, 0, 0)) + # rectangle crossing the horizontal edges + elif iter == 9: + draw_rect(img, 200, 50, Coord(0, img.height//2), Color(1, 1, 0)) + + for _ in range(20): + blur(img) + + # center pieces + draw_rect(img, 100, 100, Coord(img.width//2, img.height//2), Color(0.5, 0.2, 0.5)) + draw_circle(img, 30, Coord(img.width//2, img.height//2), Color(0.9, 0.7, 0.9)) + + +def show(img, save_path=None): + img_shape = (img.height, img.width) + img_array = wp.array(ptr=img.data_ptr, shape=img_shape, dtype=wp.vec3f, owner=False) + + img_data = img_array.numpy() + + if save_path is not None: + import matplotlib.image as img + img.imsave(save_path, img_data) + + import matplotlib.pyplot as plt + plt.imshow(img_data) + plt.show() + + +wp.init() + +# It's a good idea to always clear the kernel when developing new native or codegen features +wp.build.clear_kernel_cache() + +# !!! DO THIS BEFORE LOADING MODULES OR LAUNCHING KERNELS +wim_types.register() + +with wp.ScopedDevice("cuda:0"): + + img = Image(800, 600) + print(img) + + # run a kernel to print image info + img_handle = wp.uint64(img.ptr) + wp.launch(print_image_info_kernel, dim=1, inputs=[img_handle]) + wp.synchronize_device() + + # make a drawing + draw_picture(img) + + # show and save the image + show(img, save_path="result.png") diff --git a/external/wim_types.py b/external/wim_types.py new file mode 100644 index 0000000000..8d1e984d53 --- /dev/null +++ b/external/wim_types.py @@ -0,0 +1,184 @@ +import ctypes +import os +import warp as wp + +# our external lib bindings (importing it will build the native lib and initialize the Python bindings) +import wim + + +class Coord: + + # define variables accessible in kernels (e.g., coord.x) + vars = { + "x": wp.codegen.Var("x", int), + "y": wp.codegen.Var("y", int), + } + + # struct that corresponds to the native Foo type + # - used when packing arguments for kernels (pass-by-value) + # - binary layout of fields must match native type + class _type_(ctypes.Structure): + + _fields_ = [ + ("x", ctypes.c_int), + ("y", ctypes.c_int), + ] + + def __init__(self, coord): + self.x = coord.x + self.y = coord.y + + def __init__(self, x=0, y=0): + self.x = x + self.y = y + + # HACK: used when packing kernel argument as `arg_type._type_(value.value)` in `pack_arg()` during `wp.launch()` + @property + def value(self): + return self + + +class Color: + + # define variables accessible in kernels + vars = { + "r": wp.codegen.Var("r", float), + "g": wp.codegen.Var("g", float), + "b": wp.codegen.Var("b", float), + } + + # struct that corresponds to the native Foo type + # - used when packing arguments for kernels (pass-by-value) + # - binary layout of fields must match native type + class _type_(ctypes.Structure): + + _fields_ = [ + ("r", ctypes.c_float), + ("g", ctypes.c_float), + ("b", ctypes.c_float), + ] + + def __init__(self, color): + self.r = color.r + self.g = color.g + self.b = color.b + + def __init__(self, r=0, g=0, b=0): + self.r = r + self.g = g + self.b = b + + # HACK: used when packing kernel argument as `arg_type._type_(value.value)` in `pack_arg()` during `wp.launch()` + @property + def value(self): + return self + + +class Image: + def __init__(self, width: int, height: int, device=None): + + # image shape + self.width = width + self.height = height + + self.device = wp.get_device(device) + + # pointer to the native wim::Image class, either on CPU or GPU + self.ptr = None + + # pointer to the image data (Color array), either on CPU or GPU + self.data_ptr = None + + if self.device.is_cpu: + data_ptr = ctypes.c_void_p() + self.ptr = wim._core.create_image_cpu(width, height, ctypes.byref(data_ptr)) + self.data_ptr = data_ptr.value + elif self.device.is_cuda: + data_ptr = ctypes.c_void_p() + self.ptr = wim._core.create_image_cuda(self.device.ordinal, width, height, ctypes.byref(data_ptr)) + self.data_ptr = data_ptr.value + else: + raise ValueError(f"Invalid device {device}") + + def __del__(self): + if self.ptr: + if self.device.is_cpu: + wim._core.destroy_image_cpu(self.ptr) + else: + wim._core.destroy_image_cuda(self.device.ordinal, self.ptr) + + +def _add_header(path): + include_directive = f"#include \"{path}\"\n" + # add this header for all native modules + wp.codegen.cpu_module_header += include_directive + wp.codegen.cuda_module_header += include_directive + + +def _register_headers(): + include_path = os.path.abspath(os.path.dirname(__file__)) + _add_header(f"{include_path}/wim_warp.h") + + +def _register_builtins(): + + # Coord constructor + wp.context.add_builtin( + "Coord_", + input_types={"x": int, "y": int}, + value_type=Coord, + missing_grad=True, + ) + + # Color addition + wp.context.add_builtin( + "add", + input_types={"a": Color, "b": Color}, + value_type=Color, + missing_grad=True, + ) + + # Color scaling + wp.context.add_builtin( + "mul", + input_types={"s": float, "c": Color}, + value_type=Color, + missing_grad=True, + ) + + # get image width + wp.context.add_builtin( + "img_width", + input_types={"handle": wp.uint64}, + value_type=int, + missing_grad=True, + ) + + # get image height + wp.context.add_builtin( + "img_height", + input_types={"handle": wp.uint64}, + value_type=int, + missing_grad=True, + ) + + # get pixel + wp.context.add_builtin( + "img_get_pixel", + input_types={"handle": wp.uint64, "coord": Coord}, + value_type=Color, + missing_grad=True, + ) + + # get pixel + wp.context.add_builtin( + "img_set_pixel", + input_types={"handle": wp.uint64, "coord": Coord, "color": Color}, + value_type=None, + missing_grad=True, + ) + + +def register(): + _register_headers() + _register_builtins() diff --git a/external/wim_warp.h b/external/wim_warp.h new file mode 100644 index 0000000000..5682f356f5 --- /dev/null +++ b/external/wim_warp.h @@ -0,0 +1,68 @@ +#pragma once + +// TODO: may need to add a mechanism for include paths +#include "wim/wim.h" + +// TODO: currently, all types and builtins need to be in the wp:: namespace +namespace wp +{ + +// import types into this namespace +using Color = ::wim::Color; +using Coord = ::wim::Coord; +using Image = ::wim::Image; + +// Coord constructor exposed as a free function +CUDA_CALLABLE inline Coord Coord_(int x, int y) +{ + return Coord(x, y); +} + +// overload operator+ for colors +CUDA_CALLABLE inline Color add(const Color& a, const Color& b) +{ + return Color(a.r + b.r, a.g + b.g, a.b + b.b); +} + +// overload operator* for scaling colors +CUDA_CALLABLE inline Color mul(float s, const Color& c) +{ + return Color(s * c.r, s * c.g, s * c.b); +} + +// +// TODO: Integer handles don't play well with polymorphism or explicit overloading for different Image subclasses. +// Would be better to pass specific types, pointers, or references. +// + +// get image pointer from handle +CUDA_CALLABLE inline Image& img_get(uint64_t handle) +{ + return *(Image*)(handle); +} + +// get image width (can't be exposed as a named var directly, because the member is private) +CUDA_CALLABLE inline int img_width(uint64_t handle) +{ + return img_get(handle).getWidth(); +} + +// get image height (can't be exposed as a named var directly, because the member is private) +CUDA_CALLABLE inline int img_height(uint64_t handle) +{ + return img_get(handle).getHeight(); +} + +// get pixel +CUDA_CALLABLE inline Color img_get_pixel(uint64_t handle, const Coord& coord) +{ + return img_get(handle).getPixel(coord); +} + +// set pixel +CUDA_CALLABLE inline void img_set_pixel(uint64_t handle, const Coord& coord, const Color& color) +{ + img_get(handle).setPixel(coord, color); +} + +} From cbb01a6a76c060210d0f2cb8cda204b8e6369bae Mon Sep 17 00:00:00 2001 From: Lukasz Wawrzyniak Date: Wed, 14 Feb 2024 18:54:41 -0500 Subject: [PATCH 02/16] Eliminated image handles, exposed image data as a Warp array, added PyTorch interop example --- external/README.md | 3 ++ external/wim/__init__.py | 10 +++-- external/wim/wim.cpp | 40 ++++++------------ external/wim_paint.py | 91 ++++++++++++++++++++++------------------ external/wim_types.py | 66 ++++++++++++++++++++++------- external/wim_warp.h | 38 +++++++++-------- 6 files changed, 143 insertions(+), 105 deletions(-) diff --git a/external/README.md b/external/README.md index 78ef24b2eb..2a2a1c9e5f 100644 --- a/external/README.md +++ b/external/README.md @@ -15,6 +15,7 @@ The file [wim_paint.py](wim_paint.py) is the main program for the example. It c * 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 @@ -23,3 +24,5 @@ From the repo root: ```python $ python external/wim_paint.py ``` + +If PyTorch is installed, the example will also demonstrate inverting the image using PyTorch. diff --git a/external/wim/__init__.py b/external/wim/__init__.py index 1abd88ecbd..7ecc7643e5 100644 --- a/external/wim/__init__.py +++ b/external/wim/__init__.py @@ -23,13 +23,17 @@ def _load_lib(): # build the lib _build_lib() -# load the lib and set up Python bindings +# load the lib _core = _load_lib() -_core.create_image_cpu.argtypes = [ctypes.c_int, ctypes.c_int, ctypes.POINTER(ctypes.c_void_p)] + +# 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 -_core.create_image_cuda.argtypes = [ctypes.c_int, ctypes.c_int, ctypes.c_int, ctypes.POINTER(ctypes.c_void_p)] + +# 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 diff --git a/external/wim/wim.cpp b/external/wim/wim.cpp index cc849b4609..d97cfe8a1d 100644 --- a/external/wim/wim.cpp +++ b/external/wim/wim.cpp @@ -31,24 +31,23 @@ bool check_cuda_result(cudaError_t code, const char* file, int line) extern "C" { -WIM_API wim::Image* create_image_cpu(int width, int height, wim::Color** data_ret) +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); - - if (data_ret) - *data_ret = data; - return img; } WIM_API void destroy_image_cpu(wim::Image* img) { - delete [] img->getData(); - delete img; + if (img) + { + delete [] img->getData(); + delete img; + } } -WIM_API wim::Image* create_image_cuda(int device, int width, int height, wim::Color** data_ret) +WIM_API wim::Image* create_image_cuda(int device, int width, int height) { if (!check_cuda(cudaSetDevice(device))) return nullptr; @@ -59,34 +58,21 @@ WIM_API wim::Image* create_image_cuda(int device, int width, int height, wim::Co if (!check_cuda(cudaMemset(data, 0, width * height * sizeof(wim::Color)))) return nullptr; - wim::Image img_h(width, height, data); - - wim::Image* img_d = nullptr; - if (!check_cuda(cudaMalloc(&img_d, sizeof(wim::Image)))) - return nullptr; - if (!check_cuda(cudaMemcpy(img_d, &img_h, sizeof(wim::Image), cudaMemcpyHostToDevice))) - return nullptr; - - if (data_ret) - *data_ret = data; - - return img_d; + wim::Image* img = new wim::Image(width, height, data); + return img; } -WIM_API void destroy_image_cuda(int device, wim::Image* img_d) +WIM_API void destroy_image_cuda(int device, wim::Image* img) { - if (!img_d) + if (!img) return; if (!check_cuda(cudaSetDevice(device))) return; - wim::Image img_h; - if (!check_cuda(cudaMemcpy(&img_h, img_d, sizeof(wim::Image), cudaMemcpyDeviceToHost))) - return; + check_cuda(cudaFree(img->getData())); - cudaFree(img_h.getData()); - cudaFree(img_d); + delete img; } } diff --git a/external/wim_paint.py b/external/wim_paint.py index 555ad57747..e596c2191f 100644 --- a/external/wim_paint.py +++ b/external/wim_paint.py @@ -7,30 +7,31 @@ @wp.kernel -def print_image_info_kernel(img_handle: wp.uint64): - width = wp.img_width(img_handle) - height = wp.img_height(img_handle) - wp.printf("Image: %dx%d\n", width, height) +def print_image_info_kernel(img: Image): + width = wp.img_width(img) + height = wp.img_height(img) + data = wp.img_data(img) # this is a Warp array which wraps the image data + wp.printf("Image: %dx%d, data array shape: (%d, %d)\n", width, height, data.shape[0], data.shape[1]) @wp.kernel -def fill_kernel(img_handle: wp.uint64, color: Color): +def fill_kernel(img: Image, color: Color): y, x = wp.tid() coord = wp.Coord_(x, y) - wp.img_set_pixel(img_handle, coord, color) + wp.img_set_pixel(img, coord, color) @wp.kernel -def fill_rect_kernel(img_handle: wp.uint64, half_width: int, half_height: int, pos: Coord, color: Color): +def fill_rect_kernel(img: Image, half_width: int, half_height: int, pos: Coord, color: Color): j, i = wp.tid() i -= half_width j -= half_height coord = wp.Coord_(pos.x + i, pos.y + j) - wp.img_set_pixel(img_handle, coord, color) + wp.img_set_pixel(img, coord, color) @wp.kernel -def fill_circle_kernel(img_handle: wp.uint64, radius: int, pos: Coord, color: Color): +def fill_circle_kernel(img: Image, radius: int, pos: Coord, color: Color): j, i = wp.tid() i -= radius j -= radius @@ -38,51 +39,47 @@ def fill_circle_kernel(img_handle: wp.uint64, radius: int, pos: Coord, color: Co x = pos.x + i y = pos.y + j coord = wp.Coord_(x, y) - wp.img_set_pixel(img_handle, coord, color) + wp.img_set_pixel(img, coord, color) @wp.kernel -def blur_kernel(img_handle: wp.uint64): +def blur_kernel(img: Image): y, x = wp.tid() - c00 = wp.img_get_pixel(img_handle, wp.Coord_(x - 1, y - 1)) - c01 = wp.img_get_pixel(img_handle, wp.Coord_(x, y - 1)) - c02 = wp.img_get_pixel(img_handle, wp.Coord_(x + 1, y - 1)) - c10 = wp.img_get_pixel(img_handle, wp.Coord_(x - 1, y)) - c11 = wp.img_get_pixel(img_handle, wp.Coord_(x, y)) - c12 = wp.img_get_pixel(img_handle, wp.Coord_(x + 1, y)) - c20 = wp.img_get_pixel(img_handle, wp.Coord_(x - 1, y + 1)) - c21 = wp.img_get_pixel(img_handle, wp.Coord_(x, y + 1)) - c22 = wp.img_get_pixel(img_handle, wp.Coord_(x + 1, y + 1)) + c00 = wp.img_get_pixel(img, wp.Coord_(x - 1, y - 1)) + c01 = wp.img_get_pixel(img, wp.Coord_(x, y - 1)) + c02 = wp.img_get_pixel(img, wp.Coord_(x + 1, y - 1)) + c10 = wp.img_get_pixel(img, wp.Coord_(x - 1, y)) + c11 = wp.img_get_pixel(img, wp.Coord_(x, y)) + c12 = wp.img_get_pixel(img, wp.Coord_(x + 1, y)) + c20 = wp.img_get_pixel(img, wp.Coord_(x - 1, y + 1)) + c21 = wp.img_get_pixel(img, wp.Coord_(x, y + 1)) + c22 = wp.img_get_pixel(img, wp.Coord_(x + 1, y + 1)) c = (c00 + c02 + c20 + c22) + 2.0 * (c01 + c21 + c10 + c12) + 4.0 * c11 c = (1.0 / 16.0) * c - wp.img_set_pixel(img_handle, wp.Coord_(x, y), c) + wp.img_set_pixel(img, wp.Coord_(x, y), c) -def fill(img, color: Color): - img_handle = wp.uint64(img.ptr) +def fill(img: Image, color: Color): img_shape = (img.height, img.width) - wp.launch(fill_kernel, dim=img_shape, inputs=[img_handle, color]) + wp.launch(fill_kernel, dim=img_shape, inputs=[img, color]) -def draw_rect(img, width: int, height: int, pos: Coord, color: Color): - img_handle = wp.uint64(img.ptr) +def draw_rect(img: Image, width: int, height: int, pos: Coord, color: Color): launch_dim = (height, width) - wp.launch(fill_rect_kernel, dim=launch_dim, inputs=[img_handle, width//2, height//2, pos, color]) + wp.launch(fill_rect_kernel, dim=launch_dim, inputs=[img, width//2, height//2, pos, color]) -def draw_circle(img, radius: int, pos: Coord, color: Color): - img_handle = wp.uint64(img.ptr) +def draw_circle(img: Image, radius: int, pos: Coord, color: Color): launch_dim = (2 * radius, 2 * radius) - wp.launch(fill_circle_kernel, dim=launch_dim, inputs=[img_handle, radius, pos, color]) + wp.launch(fill_circle_kernel, dim=launch_dim, inputs=[img, radius, pos, color]) -def blur(img): - img_handle = wp.uint64(img.ptr) +def blur(img: Image): img_shape = (img.height, img.width) - wp.launch(blur_kernel, dim=img_shape, inputs=[img_handle]) + wp.launch(blur_kernel, dim=img_shape, inputs=[img]) def draw_picture(img): @@ -115,17 +112,16 @@ def draw_picture(img): draw_circle(img, 30, Coord(img.width//2, img.height//2), Color(0.9, 0.7, 0.9)) -def show(img, save_path=None): - img_shape = (img.height, img.width) - img_array = wp.array(ptr=img.data_ptr, shape=img_shape, dtype=wp.vec3f, owner=False) +def show(img, title, save_path=None): - img_data = img_array.numpy() + img_data = img.data_array.numpy() if save_path is not None: import matplotlib.image as img img.imsave(save_path, img_data) import matplotlib.pyplot as plt + fig = plt.figure(title) plt.imshow(img_data) plt.show() @@ -141,15 +137,28 @@ def show(img, save_path=None): with wp.ScopedDevice("cuda:0"): img = Image(800, 600) - print(img) # run a kernel to print image info - img_handle = wp.uint64(img.ptr) - wp.launch(print_image_info_kernel, dim=1, inputs=[img_handle]) + wp.launch(print_image_info_kernel, dim=1, inputs=[img]) wp.synchronize_device() # make a drawing draw_picture(img) # show and save the image - show(img, save_path="result.png") + show(img, "Result", save_path="result.png") + + try: + import torch + + # wrapt the image data as a PyTorch tensor (no copy) + t = wp.to_torch(img.data_array, requires_grad=False) + + # invert the image in-place using PyTorch + torch.sub(1, t, out=t) + + # show and save the image + show(img, "Inverted using PyTorch", save_path="result_inverted.png") + + except ImportError: + print("Torch is not installed, couldn't post-process image") diff --git a/external/wim_types.py b/external/wim_types.py index 8d1e984d53..85be14ccdb 100644 --- a/external/wim_types.py +++ b/external/wim_types.py @@ -14,7 +14,7 @@ class Coord: "y": wp.codegen.Var("y", int), } - # struct that corresponds to the native Foo type + # struct that corresponds to the native Coord type # - used when packing arguments for kernels (pass-by-value) # - binary layout of fields must match native type class _type_(ctypes.Structure): @@ -47,7 +47,7 @@ class Color: "b": wp.codegen.Var("b", float), } - # struct that corresponds to the native Foo type + # struct that corresponds to the native Color type # - used when packing arguments for kernels (pass-by-value) # - binary layout of fields must match native type class _type_(ctypes.Structure): @@ -75,6 +75,23 @@ def value(self): class Image: + + # struct that corresponds to the native Image type + # - used when packing arguments for kernels (pass-by-value) + # - binary layout of fields must match native type + class _type_(ctypes.Structure): + + _fields_ = [ + ("width", ctypes.c_int), + ("height", ctypes.c_int), + ("data", ctypes.c_void_p), + ] + + def __init__(self, img): + self.width = img.width + self.height = img.height + self.data = img.data + def __init__(self, width: int, height: int, device=None): # image shape @@ -83,23 +100,20 @@ def __init__(self, width: int, height: int, device=None): self.device = wp.get_device(device) - # pointer to the native wim::Image class, either on CPU or GPU + # pointer to the native wim::Image class (on CPU) self.ptr = None - # pointer to the image data (Color array), either on CPU or GPU - self.data_ptr = None - if self.device.is_cpu: - data_ptr = ctypes.c_void_p() - self.ptr = wim._core.create_image_cpu(width, height, ctypes.byref(data_ptr)) - self.data_ptr = data_ptr.value + self.ptr = wim._core.create_image_cpu(width, height) elif self.device.is_cuda: - data_ptr = ctypes.c_void_p() - self.ptr = wim._core.create_image_cuda(self.device.ordinal, width, height, ctypes.byref(data_ptr)) - self.data_ptr = data_ptr.value + self.ptr = wim._core.create_image_cuda(self.device.ordinal, width, height) else: raise ValueError(f"Invalid device {device}") + # get pointer to the data, which could be on CPU or GPU + img_ptr = ctypes.cast(self.ptr, ctypes.POINTER(self._type_)) + self.data = img_ptr.contents.data + def __del__(self): if self.ptr: if self.device.is_cpu: @@ -107,6 +121,18 @@ def __del__(self): else: wim._core.destroy_image_cuda(self.device.ordinal, self.ptr) + # HACK: used when packing kernel argument as `arg_type._type_(value.value)` in `pack_arg()` during `wp.launch()` + @property + def value(self): + return self + + # return the data as a Warp array on the correct device + # TODO: can't currently use arrays of custom native types, so using vec3f instead + @property + def data_array(self): + shape = (self.height, self.width) + return wp.array(ptr=self.data, shape=shape, dtype=wp.vec3f, owner=False) + def _add_header(path): include_directive = f"#include \"{path}\"\n" @@ -149,7 +175,7 @@ def _register_builtins(): # get image width wp.context.add_builtin( "img_width", - input_types={"handle": wp.uint64}, + input_types={"img": Image}, value_type=int, missing_grad=True, ) @@ -157,15 +183,23 @@ def _register_builtins(): # get image height wp.context.add_builtin( "img_height", - input_types={"handle": wp.uint64}, + input_types={"img": Image}, value_type=int, missing_grad=True, ) + # get image data as a Warp array + wp.context.add_builtin( + "img_data", + input_types={"img": Image}, + value_type=wp.array2d(dtype=wp.vec3f), + missing_grad=True, + ) + # get pixel wp.context.add_builtin( "img_get_pixel", - input_types={"handle": wp.uint64, "coord": Coord}, + input_types={"img": Image, "coord": Coord}, value_type=Color, missing_grad=True, ) @@ -173,7 +207,7 @@ def _register_builtins(): # get pixel wp.context.add_builtin( "img_set_pixel", - input_types={"handle": wp.uint64, "coord": Coord, "color": Color}, + input_types={"img": Image, "coord": Coord, "color": Color}, value_type=None, missing_grad=True, ) diff --git a/external/wim_warp.h b/external/wim_warp.h index 5682f356f5..6e1d17df91 100644 --- a/external/wim_warp.h +++ b/external/wim_warp.h @@ -3,6 +3,10 @@ // TODO: may need to add a mechanism for include paths #include "wim/wim.h" +// include some Warp types so we can expose the image data as a Warp array +#include "../warp/native/array.h" +#include "../warp/native/vec.h" + // TODO: currently, all types and builtins need to be in the wp:: namespace namespace wp { @@ -30,39 +34,37 @@ CUDA_CALLABLE inline Color mul(float s, const Color& c) return Color(s * c.r, s * c.g, s * c.b); } -// -// TODO: Integer handles don't play well with polymorphism or explicit overloading for different Image subclasses. -// Would be better to pass specific types, pointers, or references. -// - -// get image pointer from handle -CUDA_CALLABLE inline Image& img_get(uint64_t handle) +// get image width (can't be exposed as a named var directly, because the member is private) +CUDA_CALLABLE inline int img_width(const Image& img) { - return *(Image*)(handle); + return img.getWidth(); } -// get image width (can't be exposed as a named var directly, because the member is private) -CUDA_CALLABLE inline int img_width(uint64_t handle) +// get image height (can't be exposed as a named var directly, because the member is private) +CUDA_CALLABLE inline int img_height(const Image& img) { - return img_get(handle).getWidth(); + return img.getHeight(); } -// get image height (can't be exposed as a named var directly, because the member is private) -CUDA_CALLABLE inline int img_height(uint64_t handle) +// get image data as a Warp array +CUDA_CALLABLE inline array_t img_data(Image& img) { - return img_get(handle).getHeight(); + Color* data = img.getData(); + + // TODO: can't currently use array of custom native types, so use vec3f + return array_t((vec3f*)data, img.getWidth(), img.getHeight()); } // get pixel -CUDA_CALLABLE inline Color img_get_pixel(uint64_t handle, const Coord& coord) +CUDA_CALLABLE inline Color img_get_pixel(const Image& img, const Coord& coord) { - return img_get(handle).getPixel(coord); + return img.getPixel(coord); } // set pixel -CUDA_CALLABLE inline void img_set_pixel(uint64_t handle, const Coord& coord, const Color& color) +CUDA_CALLABLE inline void img_set_pixel(Image& img, const Coord& coord, const Color& color) { - img_get(handle).setPixel(coord, color); + img.setPixel(coord, color); } } From 7f69a71c5af3d7c2c7942278f82f1e607baf5c39 Mon Sep 17 00:00:00 2001 From: Lukasz Wawrzyniak Date: Wed, 14 Feb 2024 19:01:01 -0500 Subject: [PATCH 03/16] Fixed typo --- external/wim_paint.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/external/wim_paint.py b/external/wim_paint.py index e596c2191f..7c1882cf8b 100644 --- a/external/wim_paint.py +++ b/external/wim_paint.py @@ -151,7 +151,7 @@ def show(img, title, save_path=None): try: import torch - # wrapt the image data as a PyTorch tensor (no copy) + # wrap the image data as a PyTorch tensor (no copy) t = wp.to_torch(img.data_array, requires_grad=False) # invert the image in-place using PyTorch From a5906850d164440debe4928012eeaa8beb71eb19 Mon Sep 17 00:00:00 2001 From: Lukasz Wawrzyniak Date: Wed, 14 Feb 2024 19:07:28 -0500 Subject: [PATCH 04/16] Removed unused imports --- external/wim_paint.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/external/wim_paint.py b/external/wim_paint.py index 7c1882cf8b..43fabd1f13 100644 --- a/external/wim_paint.py +++ b/external/wim_paint.py @@ -1,5 +1,3 @@ -import os -import numpy as np import warp as wp import wim_types From e2525574e2032ae2d62065fa59e3a18460f4735b Mon Sep 17 00:00:00 2001 From: Lukasz Wawrzyniak Date: Wed, 14 Feb 2024 19:08:56 -0500 Subject: [PATCH 05/16] Fixed a comment --- external/wim_paint.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/external/wim_paint.py b/external/wim_paint.py index 43fabd1f13..ad1a424804 100644 --- a/external/wim_paint.py +++ b/external/wim_paint.py @@ -126,7 +126,7 @@ def show(img, title, save_path=None): wp.init() -# It's a good idea to always clear the kernel when developing new native or codegen features +# It's a good idea to always clear the kernel cache when developing new native or codegen features wp.build.clear_kernel_cache() # !!! DO THIS BEFORE LOADING MODULES OR LAUNCHING KERNELS From c962b58a4e1fb2b253bedaad1ed15a272081a083 Mon Sep 17 00:00:00 2001 From: Lukasz Wawrzyniak Date: Wed, 14 Feb 2024 19:18:43 -0500 Subject: [PATCH 06/16] Updated README --- external/README.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/external/README.md b/external/README.md index 2a2a1c9e5f..067e2c3337 100644 --- a/external/README.md +++ b/external/README.md @@ -2,7 +2,7 @@ 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, running `import wim` will build and load the native library and initialize the Python bindings. +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. @@ -21,7 +21,7 @@ The file [wim_paint.py](wim_paint.py) is the main program for the example. It c From the repo root: -```python +```bash $ python external/wim_paint.py ``` From 02ba2b11c428ab0d1498f42fe8e285b9a2438c4f Mon Sep 17 00:00:00 2001 From: Lukasz Wawrzyniak Date: Wed, 14 Feb 2024 19:37:37 -0500 Subject: [PATCH 07/16] Added limitations to README --- external/README.md | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/external/README.md b/external/README.md index 067e2c3337..9cd00ef701 100644 --- a/external/README.md +++ b/external/README.md @@ -26,3 +26,13 @@ $ 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. From b793cedfb172d6936b4ad11d38ecbc76e23ce3da Mon Sep 17 00:00:00 2001 From: Lukasz Wawrzyniak Date: Wed, 14 Feb 2024 20:30:47 -0500 Subject: [PATCH 08/16] Tweaked example output --- external/wim_paint.py | 42 ++++++++++++++++++++++++++++++------------ 1 file changed, 30 insertions(+), 12 deletions(-) diff --git a/external/wim_paint.py b/external/wim_paint.py index ad1a424804..48e8a33ce3 100644 --- a/external/wim_paint.py +++ b/external/wim_paint.py @@ -9,7 +9,17 @@ def print_image_info_kernel(img: Image): width = wp.img_width(img) height = wp.img_height(img) data = wp.img_data(img) # this is a Warp array which wraps the image data - wp.printf("Image: %dx%d, data array shape: (%d, %d)\n", width, height, data.shape[0], data.shape[1]) + + wp.printf("Dimensions: %dx%d, data array shape: (%d, %d)\n", width, height, data.shape[0], data.shape[1]) + + if width > 0 and height > 0: + # demonstrate accessing elements as Colors using new builtins + color = wp.img_get_pixel(img, wp.Coord_(0, 0)) + wp.printf("First pixel color: (%f, %f, %f)\n", color.r, color.g, color.b) + + # demonstrate accessing elements through Warp array + value = data[0, 0] + wp.printf("First data value: (%f, %f, %f)\n", value[0], value[1], value[2]) @wp.kernel @@ -80,9 +90,12 @@ def blur(img: Image): wp.launch(blur_kernel, dim=img_shape, inputs=[img]) -def draw_picture(img): +def create_example_image(): + + # create image + img = Image(800, 600) - # background color + # fill with background color fill(img, Color(0.3, 0.0, 0.3)) # concentric circles in the corners @@ -109,8 +122,10 @@ def draw_picture(img): draw_rect(img, 100, 100, Coord(img.width//2, img.height//2), Color(0.5, 0.2, 0.5)) draw_circle(img, 30, Coord(img.width//2, img.height//2), Color(0.9, 0.7, 0.9)) + return img -def show(img, title, save_path=None): + +def show_image(img, title, save_path=None): img_data = img.data_array.numpy() @@ -134,18 +149,17 @@ def show(img, title, save_path=None): with wp.ScopedDevice("cuda:0"): - img = Image(800, 600) + # create an image + img = create_example_image() - # run a kernel to print image info + # run a kernel to print some image info + print("===== Image info:") wp.launch(print_image_info_kernel, dim=1, inputs=[img]) - wp.synchronize_device() - - # make a drawing - draw_picture(img) # show and save the image - show(img, "Result", save_path="result.png") + show_image(img, "Result", save_path="result.png") + # run some post-processing using PyTorch if it's installed to demonstrate interop try: import torch @@ -155,8 +169,12 @@ def show(img, title, save_path=None): # invert the image in-place using PyTorch torch.sub(1, t, out=t) + # run a kernel to print some image info + print("===== Inverted image info:") + wp.launch(print_image_info_kernel, dim=1, inputs=[img]) + # show and save the image - show(img, "Inverted using PyTorch", save_path="result_inverted.png") + show_image(img, "Inverted using PyTorch", save_path="result_inverted.png") except ImportError: print("Torch is not installed, couldn't post-process image") From f757ea8ab8039a6d16e481614499c1520a3617c5 Mon Sep 17 00:00:00 2001 From: Lukasz Wawrzyniak Date: Wed, 14 Feb 2024 20:36:50 -0500 Subject: [PATCH 09/16] Properly print data array shape --- external/wim_paint.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/external/wim_paint.py b/external/wim_paint.py index 48e8a33ce3..ba97003ff5 100644 --- a/external/wim_paint.py +++ b/external/wim_paint.py @@ -10,7 +10,7 @@ def print_image_info_kernel(img: Image): height = wp.img_height(img) data = wp.img_data(img) # this is a Warp array which wraps the image data - wp.printf("Dimensions: %dx%d, data array shape: (%d, %d)\n", width, height, data.shape[0], data.shape[1]) + wp.printf("Dimensions: %dx%d, data array shape: (%d, %d)\n", width, height, data.shape[1], data.shape[0]) if width > 0 and height > 0: # demonstrate accessing elements as Colors using new builtins From 7d8ab29d48a56354acb9c5304f68a8254f03a3b3 Mon Sep 17 00:00:00 2001 From: Lukasz Wawrzyniak Date: Wed, 14 Feb 2024 20:44:05 -0500 Subject: [PATCH 10/16] Fixed comment --- external/wim_types.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/external/wim_types.py b/external/wim_types.py index 85be14ccdb..805df8ca64 100644 --- a/external/wim_types.py +++ b/external/wim_types.py @@ -204,7 +204,7 @@ def _register_builtins(): missing_grad=True, ) - # get pixel + # set pixel wp.context.add_builtin( "img_set_pixel", input_types={"img": Image, "coord": Coord, "color": Color}, From c95f54820bdd06f561b22ac9b319b7022a08d00c Mon Sep 17 00:00:00 2001 From: Lukasz Wawrzyniak Date: Wed, 14 Feb 2024 20:57:02 -0500 Subject: [PATCH 11/16] Fixed Warp data array dimensions in kernels --- external/wim_paint.py | 8 ++++++-- external/wim_warp.h | 2 +- 2 files changed, 7 insertions(+), 3 deletions(-) diff --git a/external/wim_paint.py b/external/wim_paint.py index ba97003ff5..177dcde582 100644 --- a/external/wim_paint.py +++ b/external/wim_paint.py @@ -13,12 +13,16 @@ def print_image_info_kernel(img: Image): wp.printf("Dimensions: %dx%d, data array shape: (%d, %d)\n", width, height, data.shape[1], data.shape[0]) if width > 0 and height > 0: + # middle pixel coordinates + x = width // 2 + y = height // 2 + # demonstrate accessing elements as Colors using new builtins - color = wp.img_get_pixel(img, wp.Coord_(0, 0)) + color = wp.img_get_pixel(img, wp.Coord_(x, y)) wp.printf("First pixel color: (%f, %f, %f)\n", color.r, color.g, color.b) # demonstrate accessing elements through Warp array - value = data[0, 0] + value = data[y, x] wp.printf("First data value: (%f, %f, %f)\n", value[0], value[1], value[2]) diff --git a/external/wim_warp.h b/external/wim_warp.h index 6e1d17df91..5fd606b869 100644 --- a/external/wim_warp.h +++ b/external/wim_warp.h @@ -52,7 +52,7 @@ CUDA_CALLABLE inline array_t img_data(Image& img) Color* data = img.getData(); // TODO: can't currently use array of custom native types, so use vec3f - return array_t((vec3f*)data, img.getWidth(), img.getHeight()); + return array_t((vec3f*)data, img.getHeight(), img.getWidth()); } // get pixel From 718bdbb00f4165ccf37544c7b14108a5f8251e90 Mon Sep 17 00:00:00 2001 From: Lukasz Wawrzyniak Date: Wed, 14 Feb 2024 21:04:49 -0500 Subject: [PATCH 12/16] Updated comments and printouts --- external/wim_paint.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/external/wim_paint.py b/external/wim_paint.py index 177dcde582..ec5aacf605 100644 --- a/external/wim_paint.py +++ b/external/wim_paint.py @@ -4,6 +4,7 @@ from wim_types import Coord, Color, Image +# print some info about an image @wp.kernel def print_image_info_kernel(img: Image): width = wp.img_width(img) @@ -19,13 +20,14 @@ def print_image_info_kernel(img: Image): # demonstrate accessing elements as Colors using new builtins color = wp.img_get_pixel(img, wp.Coord_(x, y)) - wp.printf("First pixel color: (%f, %f, %f)\n", color.r, color.g, color.b) + wp.printf("Middle pixel color: (%g, %g, %g)\n", color.r, color.g, color.b) # demonstrate accessing elements through Warp array value = data[y, x] - wp.printf("First data value: (%f, %f, %f)\n", value[0], value[1], value[2]) + wp.printf("Middle array value: (%g, %g, %g)\n", value[0], value[1], value[2]) +# fill image with a constant color @wp.kernel def fill_kernel(img: Image, color: Color): y, x = wp.tid() @@ -33,6 +35,7 @@ def fill_kernel(img: Image, color: Color): wp.img_set_pixel(img, coord, color) +# fill a rectangle centered at `pos` @wp.kernel def fill_rect_kernel(img: Image, half_width: int, half_height: int, pos: Coord, color: Color): j, i = wp.tid() @@ -42,6 +45,7 @@ def fill_rect_kernel(img: Image, half_width: int, half_height: int, pos: Coord, wp.img_set_pixel(img, coord, color) +# fill a circle centered at `pos` @wp.kernel def fill_circle_kernel(img: Image, radius: int, pos: Coord, color: Color): j, i = wp.tid() @@ -54,6 +58,7 @@ def fill_circle_kernel(img: Image, radius: int, pos: Coord, color: Color): wp.img_set_pixel(img, coord, color) +# blur the image using a simple weighted sum over the neighbours @wp.kernel def blur_kernel(img: Image): y, x = wp.tid() From 6f90f0797a693c02ee901612a71a6a42b14d3387 Mon Sep 17 00:00:00 2001 From: Lukasz Wawrzyniak Date: Wed, 14 Feb 2024 21:14:30 -0500 Subject: [PATCH 13/16] Tweaked comments and naming --- external/wim_paint.py | 21 +++++++++++++-------- 1 file changed, 13 insertions(+), 8 deletions(-) diff --git a/external/wim_paint.py b/external/wim_paint.py index ec5aacf605..0407cd1932 100644 --- a/external/wim_paint.py +++ b/external/wim_paint.py @@ -79,26 +79,31 @@ def blur_kernel(img: Image): wp.img_set_pixel(img, wp.Coord_(x, y), c) +# fill image with constant color def fill(img: Image, color: Color): - img_shape = (img.height, img.width) - wp.launch(fill_kernel, dim=img_shape, inputs=[img, color]) + domain = (img.height, img.width) + wp.launch(fill_kernel, dim=domain, inputs=[img, color]) +# draw a width x height rectangle centered at pos def draw_rect(img: Image, width: int, height: int, pos: Coord, color: Color): - launch_dim = (height, width) - wp.launch(fill_rect_kernel, dim=launch_dim, inputs=[img, width//2, height//2, pos, color]) + domain = (height, width) + wp.launch(fill_rect_kernel, dim=domain, inputs=[img, width//2, height//2, pos, color]) +# draw a circle with the given radius centered at pos def draw_circle(img: Image, radius: int, pos: Coord, color: Color): - launch_dim = (2 * radius, 2 * radius) - wp.launch(fill_circle_kernel, dim=launch_dim, inputs=[img, radius, pos, color]) + domain = (2 * radius, 2 * radius) + wp.launch(fill_circle_kernel, dim=domain, inputs=[img, radius, pos, color]) +# blur the image def blur(img: Image): - img_shape = (img.height, img.width) - wp.launch(blur_kernel, dim=img_shape, inputs=[img]) + domain = (img.height, img.width) + wp.launch(blur_kernel, dim=domain, inputs=[img]) +# make awesome art def create_example_image(): # create image From c10dc1663759963e88a900e2843f5d426f152650 Mon Sep 17 00:00:00 2001 From: Lukasz Wawrzyniak Date: Wed, 14 Feb 2024 21:29:36 -0500 Subject: [PATCH 14/16] More comments --- external/wim_paint.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/external/wim_paint.py b/external/wim_paint.py index 0407cd1932..3c1dec0991 100644 --- a/external/wim_paint.py +++ b/external/wim_paint.py @@ -139,8 +139,10 @@ def create_example_image(): return img +# show the image and optionally save it if save_path is specified def show_image(img, title, save_path=None): + # get the image data as a numpy array img_data = img.data_array.numpy() if save_path is not None: From 4bd50dd2741d84c4616bd02b950507baefad7092 Mon Sep 17 00:00:00 2001 From: Lukasz Wawrzyniak Date: Wed, 14 Feb 2024 21:40:38 -0500 Subject: [PATCH 15/16] Fixed printing data array shape --- external/wim_paint.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/external/wim_paint.py b/external/wim_paint.py index 3c1dec0991..27392b8e30 100644 --- a/external/wim_paint.py +++ b/external/wim_paint.py @@ -11,7 +11,7 @@ def print_image_info_kernel(img: Image): height = wp.img_height(img) data = wp.img_data(img) # this is a Warp array which wraps the image data - wp.printf("Dimensions: %dx%d, data array shape: (%d, %d)\n", width, height, data.shape[1], data.shape[0]) + wp.printf("Dimensions: %dx%d, data array shape: (%d, %d)\n", width, height, data.shape[0], data.shape[1]) if width > 0 and height > 0: # middle pixel coordinates From 4b756aaa6fcd720fb05bea82ddd35f4a6759b6da Mon Sep 17 00:00:00 2001 From: Lukasz Wawrzyniak Date: Wed, 14 Feb 2024 21:56:30 -0500 Subject: [PATCH 16/16] Tweaked comments --- external/wim_paint.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/external/wim_paint.py b/external/wim_paint.py index 27392b8e30..5a5057ca8a 100644 --- a/external/wim_paint.py +++ b/external/wim_paint.py @@ -165,10 +165,10 @@ def show_image(img, title, save_path=None): with wp.ScopedDevice("cuda:0"): - # create an image + # create an image on the current device img = create_example_image() - # run a kernel to print some image info + # print image info print("===== Image info:") wp.launch(print_image_info_kernel, dim=1, inputs=[img]) @@ -185,7 +185,7 @@ def show_image(img, title, save_path=None): # invert the image in-place using PyTorch torch.sub(1, t, out=t) - # run a kernel to print some image info + # print image info print("===== Inverted image info:") wp.launch(print_image_info_kernel, dim=1, inputs=[img])