1-1-array_programming.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[32m\u001b[1m  Activating\u001b[22m\u001b[39m project at `~/cscs2023`\n"
     ]
    }
   ],
   "source": [
    "using Pkg\n",
    "Pkg.activate(@__DIR__)\n",
    "Pkg.instantiate()\n",
    "\n",
    "using CUDA"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Array programming\n",
    "\n",
    "In this notebook, I'll explain how to use the `CuArray` type to program the GPU. This is a convenient programming model that does not require detailed knowledge of the GPU, but there's still some noteworthy tips and tricks that can significantly impact performance."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "It all starts with the `CuArray` type, which serves a dual purpose:\n",
    "\n",
    "- a managed container for GPU memory\n",
    "- a way to dispatch to operations that execute on the GPU"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "2×2 CuArray{Float64, 2, CUDA.Mem.DeviceBuffer}:\n",
       " 1.0  2.0\n",
       " 3.0  4.0"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "A = CuArray([1. 2.; 3. 4.])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "A common shorthand way to create a `CuArray` is to call the `cu` function. It behaves like an recursive, but opiniated constructor:\n",
    "- it descends into structures, e.g., `cu(Adjoint([1, 2])) == Adjoint(CuArray([1, 2]))`\n",
    "- it convers slow `Float64` into much faster `Float32`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "2×2 adjoint(::CuArray{Float32, 2, CUDA.Mem.DeviceBuffer}) with eltype Float32:\n",
       " 1.0  3.0\n",
       " 2.0  4.0"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "cu([1. 2.; 3. 4.]')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "2×2 CuArray{Float64, 2, CUDA.Mem.DeviceBuffer}:\n",
       " 1.0  3.0\n",
       " 2.0  4.0"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# compare to\n",
    "CuArray([1. 2.; 3. 4.]')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Memory management will be discussed in detail in a later notebook, but for now it's enough to remember that a CuArray is **a CPU object representing memory on the GPU**. It will be automatically freed when all references have been removed, and the garbage collector runs."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The goal of `CuArray` is to make it easy to program GPUs using array operations:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "2×2 CuArray{Float64, 2, CUDA.Mem.DeviceBuffer}:\n",
       "  7.0  10.0\n",
       " 15.0  22.0"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# this will automatically use CUBLAS\n",
    "A * A"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "2×2 CuArray{Float64, 2, CUDA.Mem.DeviceBuffer}:\n",
       " 1.0   4.0\n",
       " 9.0  16.0"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# whereas this operation will use a native broadcast kernel\n",
    "A .* A"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This works by specializing certain methods with a GPU-specialized implementation, either for:\n",
    "- compatibility: not all CPU implementations work on the GPU\n",
    "- performance: GPUs have a different programming model so might require optimized implementations\n",
    "\n",
    "This generally works pretty well, the goal is to get as close to the CPU `Array` type's functionality as possible, and entire applications have been built on top of CuArray's array functionality."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Higher-order functionality\n",
    "\n",
    "The broadcast expression `A .* A` may look like a simple, special-purpose element-wise multiplication, but is syntactical sugar for a much more generic operation:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "2×2 CuArray{Float64, 2, CUDA.Mem.DeviceBuffer}:\n",
       " 1.0   4.0\n",
       " 9.0  16.0"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "elwise_op(a, b) = a * b\n",
    "broadcast(elwise_op, A, A)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This is an example of a higher-order operation, i.e., an operation that takes a function as an argument. This is a very powerful concept, because it makes it possible to *compose* the library definition of an operation, here `broadcast`, with user-provided code. This is possible in Julia because we have a JIT compiler, and in many cases **makes it possible to write custom GPU code without ever having to write a kernel**."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Passing user code to a function in Julia can also be done using the `do-block` syntax:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "2×2 CuArray{Float64, 2, CUDA.Mem.DeviceBuffer}:\n",
       " 1.0   4.0\n",
       " 9.0  16.0"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "broadcast(A, A) do a, b\n",
    "    a * b\n",
    "end\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "But the most convenient syntax is of course the dot syntax, where e.g. `f.(A .+ B)` is equivalent to broadcasting a function that computes `f(a + b)` for each element `a` in `A` and `b` in `B`. It's important to note that this dot expression performing multiple operations resulted in only a single broadcast invocation, i.e., we have **syntactical dot fusion**."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       ":($(Expr(:thunk, CodeInfo(\n",
       "   \u001b[33m @ none within `top-level scope`\u001b[39m\n",
       "\u001b[90m1 ─\u001b[39m %1 = -\n",
       "\u001b[90m│  \u001b[39m %2 = f\n",
       "\u001b[90m│  \u001b[39m %3 = Base.broadcasted(+, A, B)\n",
       "\u001b[90m│  \u001b[39m %4 = Base.broadcasted(%2, %3)\n",
       "\u001b[90m│  \u001b[39m %5 = Base.broadcasted(%1, %4, 2)\n",
       "\u001b[90m│  \u001b[39m %6 = Base.materialize(%5)\n",
       "\u001b[90m└──\u001b[39m      return %6\n",
       "))))"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "Meta.@lower f.(A .+ B) .- 2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Julia defines a variety of these higher-order operations, many of which are implemented by the GPU back-ends. For example, there's also `map`, similar to `broadcast` but without the, well, broadcasting property that allows for mismatching sizes:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "2×2 CuArray{Float64, 2, CUDA.Mem.DeviceBuffer}:\n",
       " 2.0  4.0\n",
       " 6.0  8.0"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "map(A) do a\n",
    "    a * 2\n",
    "end\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Reductions\n",
    "\n",
    "Another important higher-order operation is `mapreduce`, which can be used to map & reduce any part of an N-dimensional array. For example, one of the simplest invocations:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "3"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "A = cu([1, 2])\n",
    "reduce(+, A)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This reduces to a scalar, which requires synchronizing the GPU. Instead, you can also synchronize to a one-element array, which will then only synchronize the GPU when fetching the contents of that array. This is done by specifying the `dims` keyword, which specifies which dimensions to reduce over:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "1-element CuArray{Int64, 1, CUDA.Mem.DeviceBuffer}:\n",
       " 3"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "reduce(+, A; dims=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The `dims` keyword also makes it possible to perform multiple reductions at once, e.g., in the case the latest dimension of an array represents the batch:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "1×1×3 CuArray{Float32, 3, CUDA.Mem.DeviceBuffer}:\n",
       "[:, :, 1] =\n",
       " 49.22279\n",
       "\n",
       "[:, :, 2] =\n",
       " 51.255432\n",
       "\n",
       "[:, :, 3] =\n",
       " 46.279785"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "A = CUDA.rand(10, 10, 3)\n",
    "reduce(+, A; dims=[1,2])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The `reduce` operation is part of a family of operations that all build on the `mapreduce` operation, with specializations like `sum`, `prod`, `any`, `all`, etc.\n",
    "\n",
    "A version of `reduce` that maintains the intermediate results is `accumulate`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "5-element CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}:\n",
       " 1.0\n",
       " 2.0\n",
       " 3.0\n",
       " 4.0\n",
       " 5.0"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "A = cu(ones(5))\n",
    "accumulate(+, A)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Module-level functionality\n",
    "\n",
    "Some Julia APIs do not take an array argument, and as such cannot be specialized for GPU execution. Examples include: `rand`, `fill`, `zeros`, etc. For these functions, CUDA provides unexported replacements, e.g., `CUDA.rand`, `CUDA.fill`, etc."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "1-element CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}:\n",
       " 0.0"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "CUDA.zeros(1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "2×2 CuArray{Float64, 2, CUDA.Mem.DeviceBuffer}:\n",
       " 0.191234  0.247368\n",
       " 0.33033   0.617222"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "CUDA.rand(Float64, 2, 2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Common issues\n",
    "\n",
    "Before trying this out, let's take a look at some issues that are common when using arrays to program the GPU."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Scalar iteration\n",
    "\n",
    "A key performance issue comes from the fact that a `CuArray` instance is a CPU object representing a chunk of memory on the GPU. That means we invoke the GPU for every CPU operation invoked on a CuArray. That is OK for array operations, where the GPU will have to do a bunch of work, but is very bad when you have CPU code performing a bunch of small scalar operations:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[33m\u001b[1m┌ \u001b[22m\u001b[39m\u001b[33m\u001b[1mWarning: \u001b[22m\u001b[39mPerforming scalar indexing on task Task (runnable) @0x000015554c14ed60.\n",
      "\u001b[33m\u001b[1m│ \u001b[22m\u001b[39mInvocation of getindex resulted in scalar indexing of a GPU array.\n",
      "\u001b[33m\u001b[1m│ \u001b[22m\u001b[39mThis is typically caused by calling an iterating implementation of a method.\n",
      "\u001b[33m\u001b[1m│ \u001b[22m\u001b[39mSuch implementations *do not* execute on the GPU, but very slowly on the CPU,\n",
      "\u001b[33m\u001b[1m│ \u001b[22m\u001b[39mand therefore are only permitted from the REPL for prototyping purposes.\n",
      "\u001b[33m\u001b[1m│ \u001b[22m\u001b[39mIf you did intend to index this array, annotate the caller with @allowscalar.\n",
      "\u001b[33m\u001b[1m└ \u001b[22m\u001b[39m\u001b[90m@ GPUArraysCore /scratch/snx3000/julia/tbesard/daint-gpu/packages/GPUArraysCore/uOYfN/src/GPUArraysCore.jl:106\u001b[39m\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "55"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "A = CuArray(1:10)\n",
    "A_sum = zero(eltype(A))\n",
    "for I in eachindex(A)\n",
    "    A_sum += A[I]\n",
    "end\n",
    "A_sum"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Because of this kind of programming pattern, iterating the array and fetching one scalar at a time (hence 'scalar iteration'), being so slow CUDA.jl warns about it. With the above snippet, the situation is actually even worse: Not only does every iteration require a GPU operation to fetch an element, the `getindex` call is also the only array operation meaning that the actual summation won't even run on the GPU!\n",
    "\n",
    "The solution here is to use the `sum` function that performs the entire operation as a single step.\n",
    "\n",
    "To disallow scalar iteration, use the `allowscalar` function:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "ename": "LoadError",
     "evalue": "Scalar indexing is disallowed.\nInvocation of getindex resulted in scalar indexing of a GPU array.\nThis is typically caused by calling an iterating implementation of a method.\nSuch implementations *do not* execute on the GPU, but very slowly on the CPU,\nand therefore are only permitted from the REPL for prototyping purposes.\nIf you did intend to index this array, annotate the caller with @allowscalar.",
     "output_type": "error",
     "traceback": [
      "Scalar indexing is disallowed.\nInvocation of getindex resulted in scalar indexing of a GPU array.\nThis is typically caused by calling an iterating implementation of a method.\nSuch implementations *do not* execute on the GPU, but very slowly on the CPU,\nand therefore are only permitted from the REPL for prototyping purposes.\nIf you did intend to index this array, annotate the caller with @allowscalar.",
      "",
      "Stacktrace:",
      " [1] error(s::String)",
      "   @ Base ./error.jl:35",
      " [2] assertscalar(op::String)",
      "   @ GPUArraysCore /scratch/snx3000/julia/tbesard/daint-gpu/packages/GPUArraysCore/uOYfN/src/GPUArraysCore.jl:103",
      " [3] getindex(A::CuArray{Int64, 1, CUDA.Mem.DeviceBuffer}, I::Int64)",
      "   @ GPUArrays /scratch/snx3000/julia/tbesard/daint-gpu/packages/GPUArrays/dAUOE/src/host/indexing.jl:48",
      " [4] top-level scope",
      "   @ In[20]:2"
     ]
    }
   ],
   "source": [
    "CUDA.allowscalar(false)\n",
    "A[1]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You should generally always enable this option! It's not by default in interactive sessions because it simplifies porting CPU code, and it's easy to trigger scalar iteration from non performance-sensitive paths (e.g. display methods):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "1×10 adjoint(::CuArray{Int64, 1, CUDA.Mem.DeviceBuffer}) with eltype Int64:\n",
       " 1  2  3  4  5  6  7  8  9  10"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "A'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "ename": "ErrorException",
     "evalue": "Scalar indexing is disallowed.\nInvocation of getindex resulted in scalar indexing of a GPU array.\nThis is typically caused by calling an iterating implementation of a method.\nSuch implementations *do not* execute on the GPU, but very slowly on the CPU,\nand therefore are only permitted from the REPL for prototyping purposes.\nIf you did intend to index this array, annotate the caller with @allowscalar.",
     "output_type": "error",
     "traceback": [
      "Scalar indexing is disallowed.\nInvocation of getindex resulted in scalar indexing of a GPU array.\nThis is typically caused by calling an iterating implementation of a method.\nSuch implementations *do not* execute on the GPU, but very slowly on the CPU,\nand therefore are only permitted from the REPL for prototyping purposes.\nIf you did intend to index this array, annotate the caller with @allowscalar.",
      "",
      "Stacktrace:",
      "  [1] error(s::String)",
      "    @ Base ./error.jl:35",
      "  [2] assertscalar(op::String)",
      "    @ GPUArraysCore /scratch/snx3000/julia/tbesard/daint-gpu/packages/GPUArraysCore/uOYfN/src/GPUArraysCore.jl:103",
      "  [3] getindex",
      "    @ /scratch/snx3000/julia/tbesard/daint-gpu/packages/GPUArrays/dAUOE/src/host/indexing.jl:48 [inlined]",
      "  [4] isassigned(A::CuArray{Int64, 1, CUDA.Mem.DeviceBuffer}, i::Int64)",
      "    @ Base ./multidimensional.jl:1577",
      "  [5] isassigned",
      "    @ LinearAlgebra /apps/daint/UES/omlins/course2023/software/Julia/1.10.0-rc1-CrayGNU-21.09-cuda/share/julia/stdlib/v1.10/LinearAlgebra/src/adjtrans.jl:326 [inlined]",
      "  [6] isassigned",
      "    @ Base ./multidimensional.jl:1572 [inlined]",
      "  [7] isassigned(::SubArray{Int64, 2, LinearAlgebra.Adjoint{Int64, CuArray{Int64, 1, CUDA.Mem.DeviceBuffer}}, Tuple{Base.Slice{Base.OneTo{Int64}}, Base.Slice{Base.OneTo{Int64}}}, true}, ::Int64, ::Int64)",
      "    @ Base ./subarray.jl:360",
      "  [8] alignment(io::IOContext{IOBuffer}, X::AbstractVecOrMat, rows::Vector{Int64}, cols::Vector{Int64}, cols_if_complete::Int64, cols_otherwise::Int64, sep::Int64, ncols::Int64)",
      "    @ Base ./arrayshow.jl:68",
      "  [9] _print_matrix(io::IOContext{IOBuffer}, X::AbstractVecOrMat, pre::String, sep::String, post::String, hdots::String, vdots::String, ddots::String, hmod::Int64, vmod::Int64, rowsA::UnitRange{Int64}, colsA::UnitRange{Int64})",
      "    @ Base ./arrayshow.jl:207",
      " [10] print_matrix(io::IOContext{IOBuffer}, X::SubArray{Int64, 2, LinearAlgebra.Adjoint{Int64, CuArray{Int64, 1, CUDA.Mem.DeviceBuffer}}, Tuple{Base.Slice{Base.OneTo{Int64}}, Base.Slice{Base.OneTo{Int64}}}, true}, pre::String, sep::String, post::String, hdots::String, vdots::String, ddots::String, hmod::Int64, vmod::Int64)",
      "    @ Base ./arrayshow.jl:171",
      " [11] print_matrix(io::IO, X::AbstractVecOrMat, pre::AbstractString, sep::AbstractString, post::AbstractString, hdots::AbstractString, vdots::AbstractString, ddots::AbstractString, hmod::Integer, vmod::Integer)",
      "    @ Base ./arrayshow.jl:171 [inlined]",
      " [12] print_array",
      "    @ ./arrayshow.jl:358 [inlined]",
      " [13] show(io::IOContext{IOBuffer}, ::MIME{Symbol(\"text/plain\")}, X::SubArray{Int64, 2, LinearAlgebra.Adjoint{Int64, CuArray{Int64, 1, CUDA.Mem.DeviceBuffer}}, Tuple{Base.Slice{Base.OneTo{Int64}}, Base.Slice{Base.OneTo{Int64}}}, true})",
      "    @ Base ./arrayshow.jl:399",
      " [14] limitstringmime(mime::MIME{Symbol(\"text/plain\")}, x::SubArray{Int64, 2, LinearAlgebra.Adjoint{Int64, CuArray{Int64, 1, CUDA.Mem.DeviceBuffer}}, Tuple{Base.Slice{Base.OneTo{Int64}}, Base.Slice{Base.OneTo{Int64}}}, true}, forcetext::Bool)",
      "    @ IJulia /apps/daint/UES/omlins/course2023/software/jupyterlab/3.4.5-CrayGNU-21.09-batchspawner-cuda/share/IJulia/packages/IJulia/Vo51o/src/inline.jl:43",
      " [15] limitstringmime(mime::MIME, x::Any, forcetext::Any)",
      "    @ IJulia /apps/daint/UES/omlins/course2023/software/jupyterlab/3.4.5-CrayGNU-21.09-batchspawner-cuda/share/IJulia/packages/IJulia/Vo51o/src/inline.jl:38 [inlined]",
      " [16] display_mimestring(m::MIME{Symbol(\"text/plain\")}, x::SubArray{Int64, 2, LinearAlgebra.Adjoint{Int64, CuArray{Int64, 1, CUDA.Mem.DeviceBuffer}}, Tuple{Base.Slice{Base.OneTo{Int64}}, Base.Slice{Base.OneTo{Int64}}}, true})",
      "    @ IJulia /apps/daint/UES/omlins/course2023/software/jupyterlab/3.4.5-CrayGNU-21.09-batchspawner-cuda/share/IJulia/packages/IJulia/Vo51o/src/display.jl:71",
      " [17] display_dict(x::SubArray{Int64, 2, LinearAlgebra.Adjoint{Int64, CuArray{Int64, 1, CUDA.Mem.DeviceBuffer}}, Tuple{Base.Slice{Base.OneTo{Int64}}, Base.Slice{Base.OneTo{Int64}}}, true})",
      "    @ IJulia /apps/daint/UES/omlins/course2023/software/jupyterlab/3.4.5-CrayGNU-21.09-batchspawner-cuda/share/IJulia/packages/IJulia/Vo51o/src/display.jl:102",
      " [18] #invokelatest#2",
      "    @ ./essentials.jl:887 [inlined]",
      " [19] invokelatest",
      "    @ ./essentials.jl:884 [inlined]",
      " [20] execute_request(socket::ZMQ.Socket, msg::IJulia.Msg)",
      "    @ IJulia /apps/daint/UES/omlins/course2023/software/jupyterlab/3.4.5-CrayGNU-21.09-batchspawner-cuda/share/IJulia/packages/IJulia/Vo51o/src/execute_request.jl:112",
      " [21] #invokelatest#2",
      "    @ ./essentials.jl:887 [inlined]",
      " [22] invokelatest",
      "    @ ./essentials.jl:884 [inlined]",
      " [23] eventloop(socket::ZMQ.Socket)",
      "    @ IJulia /apps/daint/UES/omlins/course2023/software/jupyterlab/3.4.5-CrayGNU-21.09-batchspawner-cuda/share/IJulia/packages/IJulia/Vo51o/src/eventloop.jl:8",
      " [24] (::IJulia.var\"#15#18\")()",
      "    @ IJulia /apps/daint/UES/omlins/course2023/software/jupyterlab/3.4.5-CrayGNU-21.09-batchspawner-cuda/share/IJulia/packages/IJulia/Vo51o/src/eventloop.jl:38"
     ]
    }
   ],
   "source": [
    "view(A', :, :)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Because of how Julia's type system works, it's easy to trigger non GPU-specialized methods when using array wrappers. Still, for non-interactive code it's recommended to always disable scalar iteration."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Sometimes, however, scalar iteration is perfectly fine. For example, when fetching the result of a reduction:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "505.774f0"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "A = CUDA.rand(1024)\n",
    "R = sum(A; dims=1)\n",
    "CUDA.@allowscalar R[]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The above can be useful when the result of a reduction isn't immediately used, because reducing to an array can be executed asynchronously, while reducing to a scalar cannot."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "It's also possible to avoid the issue of scalar iteration altogether by using unified memory, but more on that in a later notebook."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Calling into C libraries\n",
    "\n",
    "Another common issue arises when calling CPU-specific code, e.g. in some C library, using a GPU array. This generally does not work, because GPU pointers are not dereferencable on the CPU. To prevent this from crashing, we introduce a GPU-specific pointer type and disallow conversions:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "ename": "LoadError",
     "evalue": "ArgumentError: cannot take the CPU address of a CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}",
     "output_type": "error",
     "traceback": [
      "ArgumentError: cannot take the CPU address of a CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}",
      "",
      "Stacktrace:",
      " [1] unsafe_convert(::Type{Ptr{Float32}}, x::CuArray{Float32, 1, CUDA.Mem.DeviceBuffer})",
      "   @ CUDA /scratch/snx3000/julia/tbesard/daint-gpu/packages/CUDA/hfp7r/src/array.jl:478",
      " [2] top-level scope",
      "   @ ./In[24]:1"
     ]
    }
   ],
   "source": [
    "ccall(:whatever, Nothing, (Ptr{Float32},), CUDA.rand(1))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In that case, either you need to use different (supported) array operations, or fix the implementation in CUDA.jl. Such a fix can mean using functions from a CUDA library, using existing operations, or writing your own kernel."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Once again, if you need to support this, you can also consider using unified memory. More on that later."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Exercise: Matrix RMSE\n",
    "\n",
    "As a simple exercise, try to implement a function that computes the RMSE of two matrices on the GPU using array operations:\n",
    "\n",
    "$$\n",
    "    RMSE = \\sqrt{\\frac{1}{N} \\sum_{i=1}^N (A_i - B_i)^2}\n",
    "$$\n",
    "\n",
    "Benchmark the implementation against the CPU version."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's start with a CPU implementation:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.4080339174096138"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "rmse(A, B) = sqrt(sum((A-B).^2) / length(A))\n",
    "\n",
    "A = rand(1024, 1024)\n",
    "B = rand(1024, 1024)\n",
    "rmse(A, B)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To 'port' this to the GPU, just change the type of the input arrays to `CuArray` and the computation of C just works:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.4080339174096138"
      ]
     },
     "execution_count": 26,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dA = CuArray(A)\n",
    "dB = CuArray(B)\n",
    "rmse(dA, dB)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Results look identical, so let's try benchmarking."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[36m\u001b[1m[ \u001b[22m\u001b[39m\u001b[36m\u001b[1mInfo: \u001b[22m\u001b[39mPrecompiling BenchmarkTools [6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf]\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "BenchmarkTools.Trial: 1219 samples with 1 evaluation.\n",
       " Range \u001b[90m(\u001b[39m\u001b[36m\u001b[1mmin\u001b[22m\u001b[39m … \u001b[35mmax\u001b[39m\u001b[90m):  \u001b[39m\u001b[36m\u001b[1m3.052 ms\u001b[22m\u001b[39m … \u001b[35m8.698 ms\u001b[39m  \u001b[90m┊\u001b[39m GC \u001b[90m(\u001b[39mmin … max\u001b[90m): \u001b[39m0.00% … 5.26%\n",
       " Time  \u001b[90m(\u001b[39m\u001b[34m\u001b[1mmedian\u001b[22m\u001b[39m\u001b[90m):     \u001b[39m\u001b[34m\u001b[1m3.112 ms             \u001b[22m\u001b[39m\u001b[90m┊\u001b[39m GC \u001b[90m(\u001b[39mmedian\u001b[90m):    \u001b[39m0.00%\n",
       " Time  \u001b[90m(\u001b[39m\u001b[32m\u001b[1mmean\u001b[22m\u001b[39m ± \u001b[32mσ\u001b[39m\u001b[90m):   \u001b[39m\u001b[32m\u001b[1m4.084 ms\u001b[22m\u001b[39m ± \u001b[32m1.722 ms\u001b[39m  \u001b[90m┊\u001b[39m GC \u001b[90m(\u001b[39mmean ± σ\u001b[90m):  \u001b[39m5.82% ± 9.42%\n",
       "\n",
       "  \u001b[39m█\u001b[34m \u001b[39m\u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[32m \u001b[39m\u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \n",
       "  \u001b[39m█\u001b[34m▂\u001b[39m\u001b[39m▁\u001b[39m▁\u001b[39m▁\u001b[39m▁\u001b[39m▁\u001b[39m▁\u001b[39m▁\u001b[39m▁\u001b[39m▁\u001b[32m▁\u001b[39m\u001b[39m▁\u001b[39m▁\u001b[39m▁\u001b[39m▁\u001b[39m▃\u001b[39m▃\u001b[39m▂\u001b[39m▁\u001b[39m▁\u001b[39m▁\u001b[39m▁\u001b[39m▁\u001b[39m▁\u001b[39m▁\u001b[39m▂\u001b[39m▁\u001b[39m▁\u001b[39m▁\u001b[39m▂\u001b[39m▄\u001b[39m▂\u001b[39m▁\u001b[39m▁\u001b[39m▁\u001b[39m▁\u001b[39m▁\u001b[39m▁\u001b[39m▁\u001b[39m▁\u001b[39m▁\u001b[39m▁\u001b[39m▁\u001b[39m▁\u001b[39m▁\u001b[39m▁\u001b[39m▁\u001b[39m▁\u001b[39m▁\u001b[39m▂\u001b[39m▂\u001b[39m▁\u001b[39m▁\u001b[39m▁\u001b[39m▁\u001b[39m▄\u001b[39m \u001b[39m▂\n",
       "  3.05 ms\u001b[90m        Histogram: frequency by time\u001b[39m        8.4 ms \u001b[0m\u001b[1m<\u001b[22m\n",
       "\n",
       " Memory estimate\u001b[90m: \u001b[39m\u001b[33m16.00 MiB\u001b[39m, allocs estimate\u001b[90m: \u001b[39m\u001b[33m4\u001b[39m."
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "using BenchmarkTools\n",
    "@benchmark rmse($A, $B)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "BenchmarkTools.Trial: 10000 samples with 1 evaluation.\n",
       " Range \u001b[90m(\u001b[39m\u001b[36m\u001b[1mmin\u001b[22m\u001b[39m … \u001b[35mmax\u001b[39m\u001b[90m):  \u001b[39m\u001b[36m\u001b[1m146.101 μs\u001b[22m\u001b[39m … \u001b[35m  9.625 ms\u001b[39m  \u001b[90m┊\u001b[39m GC \u001b[90m(\u001b[39mmin … max\u001b[90m): \u001b[39m0.00% … 19.85%\n",
       " Time  \u001b[90m(\u001b[39m\u001b[34m\u001b[1mmedian\u001b[22m\u001b[39m\u001b[90m):     \u001b[39m\u001b[34m\u001b[1m149.407 μs               \u001b[22m\u001b[39m\u001b[90m┊\u001b[39m GC \u001b[90m(\u001b[39mmedian\u001b[90m):    \u001b[39m0.00%\n",
       " Time  \u001b[90m(\u001b[39m\u001b[32m\u001b[1mmean\u001b[22m\u001b[39m ± \u001b[32mσ\u001b[39m\u001b[90m):   \u001b[39m\u001b[32m\u001b[1m158.961 μs\u001b[22m\u001b[39m ± \u001b[32m290.588 μs\u001b[39m  \u001b[90m┊\u001b[39m GC \u001b[90m(\u001b[39mmean ± σ\u001b[90m):  \u001b[39m1.10% ±  0.59%\n",
       "\n",
       "  \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m▂\u001b[39m▄\u001b[39m▆\u001b[39m▇\u001b[39m█\u001b[34m▇\u001b[39m\u001b[39m▅\u001b[39m▄\u001b[39m▂\u001b[39m▁\u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[32m \u001b[39m\u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \u001b[39m \n",
       "  \u001b[39m▂\u001b[39m▂\u001b[39m▂\u001b[39m▂\u001b[39m▃\u001b[39m▃\u001b[39m▅\u001b[39m▇\u001b[39m█\u001b[39m█\u001b[39m█\u001b[39m█\u001b[39m█\u001b[34m█\u001b[39m\u001b[39m█\u001b[39m█\u001b[39m█\u001b[39m█\u001b[39m▇\u001b[39m▅\u001b[39m▄\u001b[39m▄\u001b[39m▄\u001b[39m▃\u001b[39m▃\u001b[39m▂\u001b[39m▂\u001b[39m▂\u001b[39m▂\u001b[39m▂\u001b[39m▂\u001b[39m▂\u001b[39m▂\u001b[39m▂\u001b[39m▂\u001b[39m▂\u001b[39m▂\u001b[39m▂\u001b[39m▂\u001b[39m▂\u001b[39m▂\u001b[39m▂\u001b[39m▂\u001b[39m▂\u001b[39m▂\u001b[39m▁\u001b[39m▂\u001b[39m▁\u001b[39m▂\u001b[39m▂\u001b[39m▂\u001b[32m▂\u001b[39m\u001b[39m▂\u001b[39m▂\u001b[39m▂\u001b[39m▂\u001b[39m▂\u001b[39m▂\u001b[39m▂\u001b[39m▂\u001b[39m▂\u001b[39m \u001b[39m▃\n",
       "  146 μs\u001b[90m           Histogram: frequency by time\u001b[39m          162 μs \u001b[0m\u001b[1m<\u001b[22m\n",
       "\n",
       " Memory estimate\u001b[90m: \u001b[39m\u001b[33m6.17 KiB\u001b[39m, allocs estimate\u001b[90m: \u001b[39m\u001b[33m120\u001b[39m."
      ]
     },
     "execution_count": 28,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "@benchmark rmse($dA, $dB)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Impressive speed-up! Of course, note that we're only measuring the actual computation, and not the time to transfer the data."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Although we're simply using `BenchmarkTools.@benchmark` here, measuring execution time of code running on the GPU is generally a little tricky. So let's delve into how to do so in the next notebook."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Julia 1.10.0-rc1",
   "language": "julia",
   "name": "julia-1.10"
  },
  "language_info": {
   "file_extension": ".jl",
   "mimetype": "application/julia",
   "name": "julia",
   "version": "1.10.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}