notebooks/core_01_basics.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Warp Core Tutorial: Basics\n",
    "\n",
    "Warp is a Python framework for writing high-performance code. Warp takes regular Python functions and JIT compiles them to efficient kernel code that can run on the CPU or GPU.\n",
    "\n",
    "This notebook showcases the essential features and capabilities that form the foundation of programming with Warp.\n",
    "\n",
    "A more in-depth reference of the API can be found in the [official documentation](https://nvidia.github.io/warp/).\n",
    "\n",
    "Prerequisites:\n",
    "\n",
    "- Basic Python knowledge.\n",
    "- Understanding of NumPy arrays."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install warp-lang"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import warp as wp\n",
    "\n",
    "wp.config.quiet = True\n",
    "\n",
    "# Explicitly initializing Warp is not necessary but\n",
    "# we do it here to ensure everything is good to go.\n",
    "wp.init()\n",
    "\n",
    "if not wp.get_cuda_device_count():\n",
    "    print(\n",
    "        \"Some snippets in this notebook assume the presence of \"\n",
    "        \"a CUDA-compatible device and won't run correctly without one.\"\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Data Types\n",
    "\n",
    "Warp offers a range of data types that covers the needs in common compute workflows."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Boolean\n",
    "\n",
    "The types `wp.bool` and `bool`, which are interchangeable, can be used to represent `True`/`False` values."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Scalars\n",
    "\n",
    "Signed/unsigned integer and floating-point numbers with different widths are supported.\n",
    "\n",
    "<table>\n",
    "    <tr>\n",
    "        <th></th>\n",
    "        <th>Integer</th>\n",
    "        <th>Floating-Point</th>\n",
    "    </tr>\n",
    "    <tr>\n",
    "        <td>8-bit</td>\n",
    "        <td>wp.[u]int8</td>\n",
    "        <td></td>\n",
    "    </tr>\n",
    "    <tr>\n",
    "        <td>16-bit</td>\n",
    "        <td>wp.[u]int16</td>\n",
    "        <td>wp.float16</td>\n",
    "    </tr>\n",
    "    <tr>\n",
    "        <td>32-bit</td>\n",
    "        <td>wp.[u]int32</td>\n",
    "        <td>wp.float32</td>\n",
    "    </tr>\n",
    "    <tr>\n",
    "        <td>64-bit</td>\n",
    "        <td>wp.[u]int64</td>\n",
    "        <td>wp.float64</td>\n",
    "    </tr>\n",
    "</table>\n",
    "\n",
    "Python's `int` and `float` types can also be used in place of `wp.int32` and `wp.float32`.\n",
    "\n",
    "Note that typing in Warp is strict, and no integer promotion is done under the hood, so types need to be explicitly matched for operations to succeed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Operation between 32-bit integers.\n",
    "print(\"\\nx:\")\n",
    "x = 123 + 234\n",
    "print(x)\n",
    "\n",
    "# Operation between 32-bit floating-points.\n",
    "print(\"\\ny:\")\n",
    "y = 1.2 + 2.3\n",
    "print(y)\n",
    "\n",
    "# Operation between 8-bit integers.\n",
    "print(\"\\nz:\")\n",
    "z = wp.int8(1) + wp.int8(2)\n",
    "print(z)\n",
    "\n",
    "# Invalid operation, both integer types must match.\n",
    "print(\"\\nw:\")\n",
    "try:\n",
    "    w = wp.int8(1) + wp.int16(2)\n",
    "    print(w)\n",
    "except Exception:\n",
    "    print(\"invalid operation\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Linear Algebra\n",
    "\n",
    "Vector, matrix, and quaternion types are also provided with the most common combination of scalar types and sizes being predefined.\n",
    "\n",
    "<table>\n",
    "    <tr>\n",
    "        <th></th>\n",
    "        <th colspan=\"4\">Integer</th>\n",
    "        <th colspan=\"3\">Floating-Point</th>\n",
    "    </tr>\n",
    "    <tr>\n",
    "        <th></th>\n",
    "        <th>8-bit</th>\n",
    "        <th>16-bit</th>\n",
    "        <th>32-bit</th>\n",
    "        <th>64-bit</th>\n",
    "        <th>16-bit</th>\n",
    "        <th>32-bit</th>\n",
    "        <th>64-bit</th>\n",
    "    </tr>\n",
    "    <tr>\n",
    "        <td>2D Vector</td>\n",
    "        <td>wp.vec2[u]b</td>\n",
    "        <td>wp.vec2[u]s</td>\n",
    "        <td>wp.vec2[u]i</td>\n",
    "        <td>wp.vec2[u]l</td>\n",
    "        <td>wp.vec2h</td>\n",
    "        <td>wp.vec2f</td>\n",
    "        <td>wp.vec2d</td>\n",
    "    </tr>\n",
    "    <tr>\n",
    "        <td>3D Vector</td>\n",
    "        <td>wp.vec3[u]b</td>\n",
    "        <td>wp.vec3[u]s</td>\n",
    "        <td>wp.vec3[u]i</td>\n",
    "        <td>wp.vec3[u]l</td>\n",
    "        <td>wp.vec3h</td>\n",
    "        <td>wp.vec3f</td>\n",
    "        <td>wp.vec3d</td>\n",
    "    </tr>\n",
    "    <tr>\n",
    "        <td>4D Vector</td>\n",
    "        <td>wp.vec4[u]b</td>\n",
    "        <td>wp.vec4[u]s</td>\n",
    "        <td>wp.vec4[u]i</td>\n",
    "        <td>wp.vec4[u]l</td>\n",
    "        <td>wp.vec4h</td>\n",
    "        <td>wp.vec4f</td>\n",
    "        <td>wp.vec4d</td>\n",
    "    </tr>\n",
    "    <tr>\n",
    "        <td>2x2 Matrix</td>\n",
    "        <td></td>\n",
    "        <td></td>\n",
    "        <td></td>\n",
    "        <td></td>\n",
    "        <td>wp.mat22h</td>\n",
    "        <td>wp.mat22f</td>\n",
    "        <td>wp.mat22d</td>\n",
    "    </tr>\n",
    "    <tr>\n",
    "        <td>3x3 Matrix</td>\n",
    "        <td></td>\n",
    "        <td></td>\n",
    "        <td></td>\n",
    "        <td></td>\n",
    "        <td>wp.mat33h</td>\n",
    "        <td>wp.mat33f</td>\n",
    "        <td>wp.mat33d</td>\n",
    "    </tr>\n",
    "    <tr>\n",
    "        <td>4x4 Matrix</td>\n",
    "        <td></td>\n",
    "        <td></td>\n",
    "        <td></td>\n",
    "        <td></td>\n",
    "        <td>wp.mat44h</td>\n",
    "        <td>wp.mat44f</td>\n",
    "        <td>wp.mat44d</td>\n",
    "    </tr>\n",
    "    <tr>\n",
    "        <td>Quaternion</td>\n",
    "        <td></td>\n",
    "        <td></td>\n",
    "        <td></td>\n",
    "        <td></td>\n",
    "        <td>wp.quath</td>\n",
    "        <td>wp.quatf</td>\n",
    "        <td>wp.quatd</td>\n",
    "    </tr>\n",
    "    <tr>\n",
    "        <td>Transformation</td>\n",
    "        <td></td>\n",
    "        <td></td>\n",
    "        <td></td>\n",
    "        <td></td>\n",
    "        <td>wp.transformh</td>\n",
    "        <td>wp.transformf</td>\n",
    "        <td>wp.transformd</td>\n",
    "    </tr>\n",
    "</table>\n",
    "\n",
    "The transformation types, as defined by Warp, define a translation part `pos` and a rotation `rot`, and is primarily intended to be used in the context of rigid bodies.\n",
    "\n",
    "A few aliases defaulting to 32-bit floating-points are also available as `wp.vec2`, `wp.vec3`, `wp.vec4`, `wp.mat22`, `wp.mat33`, `wp.mat44`, `wp.quat`, and `wp.transform`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Rotate and scale a position vector.\n",
    "print(\"\\nnew_pos:\")\n",
    "pos = wp.vec3(1.0, 2.0, 3.0)\n",
    "rot = wp.mat33(1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0)\n",
    "scale = 0.5\n",
    "new_pos = (pos * rot) * scale\n",
    "print(new_pos)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Custom Linear Algebra Types\n",
    "\n",
    "It is possible to create linear algebra types of other sizes using the functions `wp.vec(length, dtype)` and `wp.mat(shape, dtype)`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create a 5D vector of 32-bit floating-points.\n",
    "print(\"\\nv:\")\n",
    "vec5f = wp.vec(length=5, dtype=wp.float32)\n",
    "v = vec5f(1.0, 2.0, 3.0, 4.0, 5.0)\n",
    "print(v)\n",
    "\n",
    "# Create a 2x3 matrix of 32-bit floating-points.\n",
    "print(\"\\nm:\")\n",
    "mat23f = wp.mat(shape=(2, 3), dtype=wp.float32)\n",
    "m = mat23f(1.0, 2.0, 3.0, 4.0, 5.0, 6.0)\n",
    "print(m)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Arrays\n",
    "\n",
    "Arrays are multidimensional containers of fixed size that can store homogeneous elements of any Warp data type either on CPU or GPU memory.\n",
    "\n",
    "They are designed to seamlessly interop with arrays from other frameworks, such as [NumPy](https://numpy.org), [PyTorch](https://pytorch.org), [JAX](https://jax.readthedocs.io), and others.\n",
    "\n",
    "A gotcha due to supporting both CPU and GPU data within a unified interface, is that accessing individual elements directly from Python's runtime isn't exposed since this would otherwise encourage suboptimal performance, as explained in this [FAQ entry](https://nvidia.github.io/warp/faq.html#why-aren-t-assignments-to-warp-arrays-supported-outside-of-kernels)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Arrays can be initialized from multidimensional sequences of scalar data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create a 1D array of integers.\n",
    "print(\"\\narr_int:\")\n",
    "arr_int = wp.array((1, 2, 3), dtype=int)\n",
    "print(f\"dtype={arr_int.dtype}, shape={arr_int.shape}\")\n",
    "print(arr_int)\n",
    "\n",
    "# Create a 1D array of vectors.\n",
    "print(\"\\narr_vec:\")\n",
    "arr_vec = wp.array(((1, 2, 3), (4, 5, 6)), dtype=wp.vec3)\n",
    "print(f\"dtype={arr_vec.dtype}, shape={arr_vec.shape}\")\n",
    "print(arr_vec)\n",
    "\n",
    "# Create a 2D array of floating-points.\n",
    "print(\"\\narr_2d:\")\n",
    "arr_2d = wp.array(((1, 2, 3), (4, 5, 6)), dtype=float)\n",
    "print(f\"dtype={arr_2d.dtype}, shape={arr_2d.shape}\")\n",
    "print(arr_2d)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "A few utilities allow to initialize arrays with a given value, or to skip initialization altogether."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create an array filled with zeros.\n",
    "print(\"\\narr_zeros:\")\n",
    "arr_zeros = wp.zeros(3)\n",
    "print(f\"dtype={arr_zeros.dtype}, shape={arr_zeros.shape}\")\n",
    "print(arr_zeros)\n",
    "\n",
    "# Create an array filled with ones.\n",
    "print(\"\\narr_ones:\")\n",
    "arr_ones = wp.ones(3)\n",
    "print(f\"dtype={arr_ones.dtype}, shape={arr_ones.shape}\")\n",
    "print(arr_ones)\n",
    "\n",
    "# Create an uninitialized array.\n",
    "print(\"\\narr_empty:\")\n",
    "arr_empty = wp.empty(3)\n",
    "print(f\"dtype={arr_empty.dtype}, shape={arr_empty.shape}\")\n",
    "print(arr_empty)\n",
    "\n",
    "# Create an array filled with a custom value.\n",
    "print(\"\\narr_custom:\")\n",
    "arr_custom = wp.full(3, 123)\n",
    "print(f\"dtype={arr_custom.dtype}, shape={arr_custom.shape}\")\n",
    "print(arr_custom)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Initializing arrays from NumPy objects, or other frameworks like Torch, is also supported."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initialize an array from NumPy.\n",
    "print(\"\\narr_from_np:\")\n",
    "rng = np.random.default_rng(seed=123)\n",
    "arr_np = rng.standard_normal((4, 2)).astype(np.float16)\n",
    "arr_from_np = wp.from_numpy(arr_np)\n",
    "print(f\"dtype={arr_from_np.dtype}, {arr_from_np.shape}\")\n",
    "print(arr_from_np)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Structs\n",
    "\n",
    "When composite types are desired, it's possible to define Python classes decorated with `@wp.struct`, where each field is a class member that must be annotated with a Warp data type.\n",
    "\n",
    "Structs, like every other data types, are supported by arrays."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create a new data type made of 2 fields.\n",
    "@wp.struct\n",
    "class Obstacle:\n",
    "    pos: wp.vec3\n",
    "    radius: float\n",
    "\n",
    "\n",
    "# Create a first instance.\n",
    "print(\"\\no1:\")\n",
    "o1 = Obstacle()\n",
    "o1.pos = wp.vec3(1.0, 2.0, 3.0)\n",
    "o1.radius = 0.75\n",
    "print(o1)\n",
    "\n",
    "# Create a second instance.\n",
    "print(\"\\no2:\")\n",
    "o2 = Obstacle()\n",
    "o2.pos = wp.vec3(2.0, 3.0, 4.0)\n",
    "o2.radius = 0.5\n",
    "print(o2)\n",
    "\n",
    "# Create an array with these instances.\n",
    "print(\"\\narr_struct:\")\n",
    "arr_struct = wp.array((o1, o2), dtype=Obstacle)\n",
    "print(f\"dtype={arr_struct.dtype}, shape={arr_struct.shape}\")\n",
    "print(arr_struct)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Kernels\n",
    "\n",
    "In a typical Warp program, Python's runtime is used to allocate data and orchestrate operations, whereas the computationally intensive tasks are expected to be implemented as kernels.\n",
    "\n",
    "These kernels are functions decorated with `@wp.kernel`, however one notable difference with the usual Python functions is that they don't return values—all inputs and outputs must be defined as parameters with typed annotations, and all output parameters must be arrays.\n",
    "\n",
    "Passing data to these kernels and evaluating them on the desired device (CPU or GPU) is done with the `wp.launch()` function.\n",
    "\n",
    "Additionally, `wp.launch()` expects a `dim` argument that allows executing the same kernel many times in parallel, using threads, which is how the massively parallel architecture of modern GPUs and its associated performance boost can be leveraged.\n",
    "\n",
    "The `dim` argument expects either a single integer or a tuple with up to 4 values for multidimensional launches. To know which thread ID is currently being evaluated, we can call `wp.tid()` from within the kernel, which accordingly returns either a single value, or multiple ones."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define a kernel that performs a component-wise average of two arrays.\n",
    "@wp.kernel\n",
    "def avg_kernel(\n",
    "    a: wp.array(dtype=float),\n",
    "    b: wp.array(dtype=float),\n",
    "    out_avg: wp.array(dtype=float),\n",
    "):\n",
    "    i = wp.tid()\n",
    "    out_avg[i] = (a[i] + b[i]) * 0.5\n",
    "\n",
    "\n",
    "# Initialize the arrays to operate on and the output one storing their average.\n",
    "shape = (32,)\n",
    "rng = np.random.default_rng(seed=123)\n",
    "a = wp.array(rng.standard_normal(shape).astype(np.float32))\n",
    "b = wp.array(rng.standard_normal(shape).astype(np.float32))\n",
    "out_avg = wp.empty_like(a)\n",
    "\n",
    "# Launch the kernel.\n",
    "print(\"\\navg:\")\n",
    "wp.launch(avg_kernel, dim=shape, inputs=(a, b), outputs=(out_avg,))\n",
    "print(out_avg)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Devices\n",
    "\n",
    "We mentioned earlier that arrays can live either on CPU or GPU memory and, similarly, that kernels can be evaluated on either device, but we didn't mention how to specify that.\n",
    "\n",
    "Arrays as well as many other functions from the API come with a `device` parameter that can either be left to the default value of `None`. or to a value representing the target device. When set to `None`, the default device currently set is used, otherwise `\"cpu\"`, and `\"cuda\"` can be set to pick either CPU or GPU memory. In the case of configurations with multiple GPUs, it's also possible to specify the device index, such as `\"cuda:0\"`, `\"cuda:1\"`, and so on."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define a kernel that fills an array with range values.\n",
    "@wp.kernel\n",
    "def range_fill_kernel(\n",
    "    out: wp.array(dtype=int),\n",
    "):\n",
    "    i = wp.tid()\n",
    "    out[i] = i\n",
    "\n",
    "\n",
    "# Retrieve the current default device.\n",
    "print(\"\\ncurrent_device:\")\n",
    "current_device = wp.get_device()\n",
    "print(current_device)\n",
    "\n",
    "# Fill an array on the current default device.\n",
    "print(\"\\narr:\")\n",
    "arr = wp.zeros(3, dtype=int)\n",
    "wp.launch(range_fill_kernel, dim=arr.shape, outputs=(arr,))\n",
    "print(f\"device={arr.device}\")\n",
    "print(arr)\n",
    "\n",
    "# Fill an array on a specified device.\n",
    "print(f\"\\narr_explicit:\")\n",
    "device = \"cpu\"\n",
    "arr_explicit = wp.zeros(3, dtype=int, device=device)\n",
    "wp.launch(range_fill_kernel, dim=arr_explicit.shape, outputs=(arr_explicit,), device=device)\n",
    "print(f\"device={arr_explicit.device}\")\n",
    "print(arr_explicit)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In applications where all compute is intended to be run on a same device, it is recommended to not pass any `device` argument for individual API calls and, instead, wrap all code within a `wp.ScopedDevice()` context that sets the default device for all the API calls within that scope."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define a kernel that fills an array with a fibonacci sequence.\n",
    "@wp.kernel\n",
    "def fibonacci_fill_kernel(\n",
    "    out: wp.array(dtype=int),\n",
    "):\n",
    "    i = wp.tid()\n",
    "    sqrt_5 = wp.sqrt(5.0)\n",
    "    p = (1.0 + sqrt_5) / 2.0\n",
    "    q = 1.0 / p\n",
    "    out[i] = int((p ** float(i) + q ** float(i)) / sqrt_5 + 0.5)\n",
    "\n",
    "\n",
    "# Ensure that all nested code is set to operate on a specified device.\n",
    "device = \"cuda\"\n",
    "with wp.ScopedDevice(device):\n",
    "    print(f\"\\narr_scoped:\")\n",
    "    arr_scoped = wp.zeros(8, dtype=int)\n",
    "    wp.launch(fibonacci_fill_kernel, dim=arr_scoped.shape, outputs=(arr_scoped,))\n",
    "    print(f\"device={arr_scoped.device}\")\n",
    "    print(arr_scoped)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Transferring data between CPU and GPU memory is made easy across the API. For example, initializing an array on the GPU from an array on the CPU is handled seamlessly. More explicit functions are also exposed, such as `wp.copy()`, `wp.clone()`, or `wp.array.numpy()`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Clone a CPU array onto GPU memory.\n",
    "print(\"\\narr_clone_gpu:\")\n",
    "arr_clone_cpu = wp.array((1, 2, 3), dtype=int, device=\"cpu\")\n",
    "arr_clone_gpu = wp.clone(arr_clone_cpu, device=\"cuda\")\n",
    "print(f\"device={arr_clone_gpu.device}\")\n",
    "print(arr_clone_gpu)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Built-Ins\n",
    "\n",
    "Similarly to Python's built-in functions, Warp ships with a set of functions that aims to cover the most common operations in areas such as scalar math (e.g.: `wp.min()`, `wp.abs()`, ...), vector math (e.g.: `wp.dot()`, `wp.length()`, ...), quaternion math (e.g.: `wp.quat_from_axis_angle()`, `wp.quat_rotate()`, ...), random numbers (e.g.: `wp.noise()`, `wp.sample_unit_sphere()`, ...), and others.\n",
    "\n",
    "Some math functions like `math.cos()` and `math.sin()` are available as part of Python's standard `math` module, however only their Warp counterpart such as `wp.cos()` and `wp.sin()` can be used within Warp kernels.\n",
    "\n",
    "All of these built-ins are available from kernels but, where possible, they can also be called directly from Python's runtime.\n",
    "\n",
    "The full list of built-ins is available in the documentation: https://nvidia.github.io/warp/modules/functions.html"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define a kernel that computes the sine of each element from an array.\n",
    "@wp.kernel\n",
    "def sine_kernel(\n",
    "    values: wp.array(dtype=float),\n",
    "    out_sine: wp.array(dtype=float),\n",
    "):\n",
    "    i = wp.tid()\n",
    "    out_sine[i] = wp.sin(values[i])\n",
    "\n",
    "\n",
    "# Launch the sine kernel, once for each element.\n",
    "print(\"\\nsine (kernel):\")\n",
    "values = wp.array((1.0, 2.0, 3.0), dtype=float)\n",
    "out_sine = wp.empty_like(values)\n",
    "wp.launch(sine_kernel, dim=out_sine.shape, inputs=(values,), outputs=(out_sine,))\n",
    "print(out_sine)\n",
    "\n",
    "# Try the same `wp.sin()` built-in from Python.\n",
    "print(\"\\nsine (runtime):\")\n",
    "x = wp.sin(1.0)\n",
    "y = wp.sin(2.0)\n",
    "z = wp.sin(3.0)\n",
    "print(x, y, z)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Random Numbers\n",
    "\n",
    "Random numbers is made available from within Warp kernels using the `wp.rand_init()` built-in to initialize the state of the generator, followed by any of `wp.randf()`, `wp.randi()`, or `wp.randn()` calls."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define a kernel that generates random numbers.\n",
    "@wp.kernel\n",
    "def rand_kernel(\n",
    "    seed: int,\n",
    "    out_rand: wp.array(dtype=float),\n",
    "):\n",
    "    i = wp.tid()\n",
    "    rng = wp.rand_init(seed, i)\n",
    "    out_rand[i] = wp.randf(rng)\n",
    "\n",
    "\n",
    "# Launch the rand kernel.\n",
    "print(\"\\nrand:\")\n",
    "out_rand = wp.empty(3, dtype=float)\n",
    "wp.launch(rand_kernel, dim=out_rand.shape, inputs=(123,), outputs=(out_rand,))\n",
    "print(out_rand)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Geometric sampling is available through built-ins like `wp.sample_triangle()`, `wp.sample_unit_disk()`, `wp.sample_unit_sphere()`, `wp.sample_unit_cube()`, and others."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define a kernel that samples random points within a unit hemisphere.\n",
    "@wp.kernel\n",
    "def sample_unit_hemisphere_kernel(\n",
    "    seed: int,\n",
    "    out_sample: wp.array(dtype=wp.vec3),\n",
    "):\n",
    "    i = wp.tid()\n",
    "    rng = wp.rand_init(seed, i)\n",
    "    out_sample[i] = wp.sample_unit_hemisphere(rng)\n",
    "\n",
    "\n",
    "# Launch the rand kernel.\n",
    "print(\"\\nsample:\")\n",
    "out_sample = wp.empty(3, dtype=wp.vec3)\n",
    "wp.launch(sample_unit_hemisphere_kernel, dim=out_sample.shape, inputs=(123,), outputs=(out_sample,))\n",
    "print(out_sample)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, Perlin-based noise functions are exposed using `wp.noise()`, `wp.pnoise()`, and `wp.curlnoise()`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define a kernel that outputs a curl noise for a 2D value.\n",
    "@wp.kernel\n",
    "def noise_kernel(\n",
    "    seed: int,\n",
    "    out_noise: wp.array(dtype=wp.vec2),\n",
    "):\n",
    "    i = wp.tid()\n",
    "    rng = wp.rand_init(seed, i)\n",
    "    xy = wp.vec2(float(123 + i * 2), float(234 + i * 3))\n",
    "    out_noise[i] = wp.curlnoise(rng, xy)\n",
    "\n",
    "\n",
    "# Launch the noise kernel.\n",
    "print(\"\\nnoise:\")\n",
    "out_noise = wp.empty(3, dtype=wp.vec2)\n",
    "wp.launch(noise_kernel, dim=out_noise.shape, inputs=(12,), outputs=(out_noise,))\n",
    "print(out_noise)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## User Functions\n",
    "\n",
    "For a function to be available in kernels, it needs to be decorated with `@wp.func`. However, unlike with kernels, these functions cannot be passed to `wp.launch()` directly, instead they are meant to be called either by a kernel or by another Warp function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define a function that computes the component-wise product of a 2D vector.\n",
    "# Providing the return type hint is optional\n",
    "@wp.func\n",
    "def product(\n",
    "    v: wp.vec2,\n",
    ") -> float:\n",
    "    return v[0] * v[1]\n",
    "\n",
    "\n",
    "# Define a kernel that computes the component-wise product of 2 vectors.\n",
    "@wp.kernel\n",
    "def product_kernel(\n",
    "    v1: wp.vec2,\n",
    "    v2: wp.vec2,\n",
    "    out_product: wp.array(dtype=float),\n",
    "):\n",
    "    out_product[0] = product(v1) * product(v2)\n",
    "\n",
    "\n",
    "# Launch the product kernel once.\n",
    "print(\"\\nproduct:\")\n",
    "v1 = wp.vec2(2.0, 4.0)\n",
    "v2 = wp.vec2(3.0, 5.0)\n",
    "out_product = wp.empty(1, dtype=float)\n",
    "wp.launch(product_kernel, dim=1, inputs=(v1, v2), outputs=(out_product,))\n",
    "print(out_product)"
   ]
  }
 ],
 "metadata": {
    "accelerator": "GPU",
    "colab": {
      "gpuType": "T4",
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}