Nvidia target via NVRTC and Nim ↦ CUDA DSL (#487)

* commit the initial Nim ⇒ CUDA DSL, NVRTC & CUDA execution helpers

* [tests/examples] add example for a BigInt modular addition

* remove `nimcuda` dependency, wrap everything we need manually

* merge `execCuda` logic for LLVM & NVRTC

IMPORTANT NOTE:

For LLVM we generate `array_t` types for the finite field elements. By
doing this we make it impossible to just copy over a Constantine
finite field element or elliptic curve point (which are also an
`array_t` type). Therefore, we have a `CUfunctionLLVM` type, which is
used to differentiate between different `execCuda` calls, based on
their "origin" (i.e. LLVM or NVRTC backends).

Based on that backend we either allow passing simple structs by their
host pointer or force a copy.

* add `quitOnFailure` for the `check` calls to avoid regression

This was added for a reason after all in 5d66b52

* [tests] turn big int `modadd` example into real test

* copy `libpaths.nim` over from nimcuda

* [cuda] add partial support for `const` in CUDA generator

* [cuda] fix minor issue in `if` statements in CUDA generator

* [cuda] add support for named blocks

* [cuda] add support for `bool`

* [cuda] remove unnecessary semicolon

* [cuda] add support for `{.volatile.}` variables

* [nvrtc] add `modadd`, `modsub`, `mtymul` implementations using inline PTX

* [cuda] support basic type conversions

By mapping them to a regular cast.

* [cuda] support `var` parameters in procs

* [cuda] make sure proc body is a block

* [cuda] support boolean / bitwise AND/OR and XOR, NOT

* [cuda] support int32 literals

* [cuda] handle prefix `not`

* [cuda] make sure to pass `array` types by pointer instead of copy

* [nvrtc] add more helpers, add TODO to investigate `slct` calls

* [nvrtc] add many more field arithmetic / bigint operations

```
proc setZero(a: var BigInt) {.device.} =
proc setOne(a: var BigInt) {.device.} =
proc add(r: var BigInt, a, b: BigInt) {.device.} =
proc sub(r: var BigInt, a, b: BigInt) {.device.} =
proc mul(r: var BigInt, a, b: BigInt) {.device.} =
proc ccopy(a: var BigInt, b: BigInt, condition: bool) {.device.} =
proc csetZero(r: var BigInt, condition: bool) {.device.} =
proc csetOne(r: var BigInt, condition: bool) {.device.} =
proc cadd(r: var BigInt, a: BigInt, condition: bool) {.device.} =
proc csub(r: var BigInt, a: BigInt, condition: bool) {.device.} =
proc doubleElement(r: var BigInt, a: BigInt) {.device.} =
proc nsqr(r: var BigInt, a: BigInt, count: int) {.device.} =
proc isZero(r: var bool, a: BigInt) {.device.} =
proc isOdd(r: var bool, a: BigInt) {.device.} =
proc neg(r: var BigInt, a: BigInt) {.device.} =
proc cneg(r: var BigInt, a: BigInt, condition: bool) {.device.} =
proc shiftRight(r: var BigInt, k: uint32) {.device.} =
proc div2(r: var BigInt) {.device.} =
```

[back]

* [tests] add test to pass by pointer and `var`

* [tests] add test case for modadd/sub/mtymul

* [tests] add basic test cases for all new NVRTC operations

* [tests] update modadd/sub, mtymul test for new `getFieldModulus`

* add BabyBear field

* [nvrtc] handle `mtymul` for fields with 1 limb

* [cuda] support nested array types, unpack generic instantiatons

* [cuda] correctly generate ptr to array & ptr to array return types

* [cuda] automatically generate `memcpy` for static array types

Static array types cannot be assigned in C/C++/CUDA. We generate a
`memcpy` call for them.

Obviously this does not work for arrays with structs that have pointer
fields or similar (at least not if you expect to copy the values)

* [nvrtc] use `const` for field modulus and other CT constants

* [cuda] extend error message for non copyable inputs

* [cuda] disable passStructByPointer also for CUDA

Once one wants to launch multiple blocks / threads per block this
approach breaks for slightly larger objects

* [cuda] add `{.nimonly.}` pragma one can use in `cuda` block

This pragma makes us not emit any code for that function. Useful to
call it in the definition of a `const`

* [cuda] better logic for detection of type names

* [cuda] support `const` by mapping it to a `__constant__`

* [cuda] allow type determination from array literal

* [nvrtc] get rid of complexity with custom uint32 constants

* [staticFor] add stepped variant of `staticFor`

* [cuda] support func, discard and command nnkCommand

* [cuda] better handle required semicolons

To fix:

```nim
const code = cuda:
  proc sum() {.device.} =
    let inputIdx = 0
    let rateIdx = 0
    if inputIdx > 0 and rateIdx > 0:
      discard
    return

echo code
```

which previously produced:

```
 extern "C" __device__ void sum(){
  long long inputIdx = 0;
  long long rateIdx = 0;
  if (((0 < inputIdx); && (0 < rateIdx);)) {
;
  };
  return ;
};
```

and now:
```
 extern "C" __device__ void sum(){
  long long inputIdx = 0;
  long long rateIdx = 0;
  if (((0 < inputIdx) && (0 < rateIdx))) {
;
  };
  return ;
};
```

* [cuda] extract type from `getType` for execution helper

* [cuda] special case `CUdeviceptr` as a type that *must not* be copied

This allows the user to e.g. allocate and copy memory to a CUDA device
before calling `execute`.

* [cuda] allow passing in shared memory size for a kernel

* [nvidia ABI] wrap cuModuleGetGlobal, cudaMemcpyKind and a couple more

* [cuda] support while loops

* [cuda] support void pointers and `nil` literals

* [cuda] refactor out module loading from execution

so that one can e.g. copy to a global symbol before execution

* [cuda] store PTX before echoeing it

* [cuda] add `copyToSymbol` helper to copy to constant symbol in CUDA code

* [cuda] generalize `volatile` annotation to support other pragmas

i.e. to write
```cuda
extern shared int foo[];
```
we will now support
```nim
var foo {.cuExtern, shared.}: array[0, int]
```
(the 0 size is the current placeholder on how to designate a `[]`
array from Nim)

* [cuda] `cudaName` pragma for custom name for a proc, eg __syncthreads

Allows to map a proc to a custom name. Useful for names that we can't
write due to Nim limitations (i.e. starting with an underscore or
names that match Nim keywords)

* [cuda] support float literals

* [cuda] map arrays of explicit length 0 to `[]` arrays in CUDA

* [cuda] explicitly support constants, mapped to `__constant__`

One can either define a Nim `const` for a variable that is already a
const at the Nim compile time or use
```nim
var foo {.constant.}: theType
```
if one wishes to copy to the symbol before kernel execution. This is
useful for global constants that are not filled in at CUDA compile
time, but before execution. For example:

```nim
  # Filled with `copyToSymbol` at runtime from host!
  var rc16 {.constant.}: array[30, array[BABYBEAR_WIDTH, BigInt]]
  var matInternalDiagM1 {.constant.}: array[BABYBEAR_WIDTH, BigInt]
  var montyInverse {.constant.}: BigInt
```

And in the host code:
```nim
  var nvrtc = initNvrtc(CudaCode)
  nvrtc.compile()
  nvrtc.getPtx()
  nvrtc.load()

  var p2bb = Poseidon2BabyBear.init()

  # copy Poseidon2 constants to CUDA kernel
  nvrtc.copyToSymbol("rc16", p2bb.rc16)
  nvrtc.copyToSymbol("matInternalDiagM1", p2bb.matInternalDiagM1)
  nvrtc.copyToSymbol("montyInverse", p2bb.montyInverse)
```

* [cuda] minor cleanup

* [cuda] add `gridDim`, `cuExtern` and `share` + device malloc/free

* force compilation with `-d:CTT_32` for the moment

Need to finalize the logic of mapping 64 bit limbs to 32 bit limb
constants

Author: Vindaar

constantine/math_compiler/codegen_nvidia.nim

@@ -71,18 +71,6 @@ export
 # Cuda Driver API
 # ------------------------------------------------------------
 
-template check*(status: CUresult, quitOnFailure = true) =
-  ## Check the status code of a CUDA operation
-  ## Exit program with error if failure
-
-  let code = status # ensure that the input expression is evaluated once only
-
-  if code != CUDA_SUCCESS:
-    writeStackTrace()
-    stderr.write(astToStr(status) & " " & $instantiationInfo() & " exited with error: " & $code & '\n')
-    if quitOnFailure:
-      quit 1 # NOTE: this hides exceptions if they are thrown!
-
 func cuModuleLoadData*(module: var CUmodule, sourceCode: openArray[char]): CUresult {.inline.}=
   cuModuleLoadData(module, sourceCode[0].unsafeAddr)
 func cuModuleGetFunction*(kernel: var CUfunction, module: CUmodule, fnName: openArray[char]): CUresult {.inline.}=
@@ -221,277 +209,6 @@ proc exec*[T](jitFn: CUfunction, r: var T, a, b: T) =
   check cuMemFree(aGPU)
   check cuMemFree(bGPU)
 
-proc getTypes(n: NimNode): seq[NimNode] =
-  case n.kind
-  of nnkIdent, nnkSym: result.add getTypeInst(n)
-  of nnkLiterals: result.add getTypeInst(n)
-  of nnkBracket, nnkTupleConstr, nnkPar:
-    for el in n:
-      result.add getTypes(el)
-  else:
-    case n.typeKind
-    of ntyPtr: result.add getTypeInst(n)
-    else:
-      error("Arguments to `execCuda` must be given as a bracket, tuple or typed expression. Instead: " & $n.treerepr)
-
-proc requiresCopy(n: NimNode): bool =
-  ## Returns `true` if the given type is not a trivial data type, which implies
-  ## it will require copying its value manually.
-  case n.typeKind
-  of ntyBool, ntyChar, ntyInt .. ntyUint64: # range includes all floats
-    result = false
-  else:
-    result = true
-
-proc allowsCopy(n: NimNode): bool =
-  ## Returns `true` if the given type is allowed to be copied. That means it is
-  ## either `requiresCopy` or a `var` symbol.
-  result = n.requiresCopy or n.symKind == nskVar
-
-proc getIdent(n: NimNode): NimNode =
-  ## Generate a `GPU` suffixed ident
-  # Note: We want a deterministic name, because we call `getIdent` for the same symbol
-  # in multiple places atm.
-  case n.kind
-  of nnkIdent, nnkSym: result = ident(n.strVal & "GPU")
-  else: result = ident("`" & n.repr & "`GPU")
-
-proc determineDevicePtrs(r, i: NimNode, iTypes: seq[NimNode]): seq[(NimNode, NimNode)] =
-  ## Returns the device pointer ident and its associated original symbol.
-  for el in r:
-    if not el.allowsCopy:
-      error("The argument for `res`: " & $el.repr & " of type: " & $el.getTypeImpl().treerepr &
-        " does not allow copying. Copying to the address of all result variables is required.")
-    result.add (getIdent(el), el)
-  for idx in 0 ..< i.len:
-    let input = i[idx]
-    let t = iTypes[idx]
-    if t.requiresCopy():
-      result.add (getIdent(input), input)
-
-proc assembleParams(r, i: NimNode, iTypes: seq[NimNode]): seq[NimNode] =
-  ## Returns all parameters. Depending on whether they require copies or
-  ## are `res` parameters, either the input parameter or the `GPU` parameter.
-  for el in r: # for `res` we always copy!
-    result.add getIdent(el)
-  for idx in 0 ..< i.len:
-    let input = i[idx]
-    let t = iTypes[idx]
-    if t.requiresCopy():
-      result.add getIdent(input)
-    else:
-      result.add input
-
-proc sizeArg(n: NimNode): NimNode =
-  ## The argument to `sizeof` must be the size of the data we copy. If the
-  ## input type is already given as a `ptr T` type, we need the size of
-  ## `T` and not `ptr`.
-  case n.typeKind
-  of ntyPtr: result = n.getTypeInst()[0]
-  else: result = n
-
-# little helper macro constructors
-template check(arg): untyped = nnkCall.newTree(ident"check", arg)
-template size(arg): untyped = nnkCall.newTree(ident"sizeof", sizeArg arg)
-template address(arg): untyped = nnkCall.newTree(ident"addr", arg)
-template csize_t(arg): untyped = nnkCall.newTree(ident"csize_t", arg)
-template pointer(arg): untyped = nnkCall.newTree(ident"pointer", arg)
-
-proc maybeAddress(n: NimNode): NimNode =
-  ## Returns the address of the given node, *IFF* the type is not a
-  ## pointer type already
-  case n.typeKind
-  of ntyPtr: result = n
-  else: result = address(n)
-
-proc genParams(pId, r, i: NimNode, iTypes: seq[NimNode]): NimNode =
-  ## Generates the parameter `params` variable
-  let ps = assembleParams(r, i, iTypes)
-  result = nnkBracket.newTree()
-  for p in ps:
-    result.add pointer(maybeAddress p)
-  result = nnkLetSection.newTree(
-    nnkIdentDefs.newTree(pId, newEmptyNode(), result)
-  )
-
-proc genVar(n: NimNode): (NimNode, NimNode) =
-  ## Generates a let `tmp` variable and returns its identifier and
-  ## the let section.
-  result[0] = genSym(nskLet, "tmp")
-  result[1] = nnkLetSection.newTree(
-    nnkIdentDefs.newTree(
-      result[0],
-      getTypeInst(n),
-      n
-    )
-  )
-
-proc genLocalVars(inputs: NimNode): (NimNode, NimNode) =
-  result[0] = newStmtList() # defines local vars
-  result[1] = nnkBracket.newTree() # returns new bracket of vars for parameters
-  for el in inputs:
-    case el.kind
-    of nnkLiterals, nnkConstDef: # define a local with the value of it
-      let (s, v) = genVar(el)
-      result[0].add v
-      result[1].add s
-    of nnkSym:
-      if el.strVal in ["true", "false"]:
-        let (s, v) = genVar(el)
-        result[0].add v
-        result[1].add s
-      else:
-        result[1].add el # keep symbol
-    else:
-      result[1].add el # keep symbol
-
-proc maybeWrap(n: NimNode): NimNode =
-  if n.kind notin {nnkBracket, nnkTupleConstr}:
-    result = nnkBracket.newTree(n)
-  else:
-    result = n
-
-proc endianCheck(): NimNode =
-  result = quote do:
-    static: doAssert cpuEndian == littleEndian, block:
-      # From https://developer.nvidia.com/cuda-downloads?target_os=Linux
-      # Supported architectures for Cuda are:
-      # x86-64, PowerPC 64 little-endian, ARM64 (aarch64)
-      # which are all little-endian at word-level.
-      #
-      # Due to limbs being also stored in little-endian, on little-endian host
-      # the CPU and GPU will have the same binary representation
-      # whether we use 32-bit or 64-bit words, so naive memcpy can be used for parameter passing.
-
-      "Most CPUs (x86-64, ARM) are little-endian, as are Nvidia GPUs, which allows naive copying of parameters.\n" &
-      "Your architecture '" & $hostCPU & "' is big-endian and GPU offloading is unsupported on it."
-
-proc execCudaImpl(jitFn, res, inputs: NimNode): NimNode =
-  # Maybe wrap individually given arguments in a `[]` bracket, e.g.
-  # `execCuda(res = foo, inputs = bar)`
-  let res = maybeWrap res
-  let inputs = maybeWrap inputs
-
-  result = newStmtList()
-  result.add endianCheck()
-
-  # get the types of the inputs
-  let rTypes = getTypes(res)
-  let iTypes = getTypes(inputs)
-
-  # determine all required `CUdeviceptr`
-  let devPtrs = determineDevicePtrs(res, inputs, iTypes)
-
-  # generate device pointers, allocate memory and copy data
-  for x in devPtrs:
-    # `var rGPU: CUdeviceptr`
-    result.add nnkVarSection.newTree(
-      nnkIdentDefs.newTree(
-        x[0],
-        ident"CUdeviceptr",
-        newEmptyNode()
-      )
-    )
-
-    # `check cuMemAlloc(rGPU, csize_t sizeof(r))`
-    result.add(
-      check nnkCall.newTree(
-        ident"cuMemAlloc",
-        x[0],
-        csize_t size(x[1])
-      )
-    )
-    # `check cuMemcpyHtoD(aGPU, a.addr, csize_t sizeof(a))`
-    result.add(
-      check nnkCall.newTree(
-        ident"cuMemcpyHtoD",
-        x[0],
-        maybeAddress x[1],
-        csize_t size(x[1])
-      )
-    )
-
-  # Generate local variables
-  let (decl, vars) = genLocalVars(inputs)
-  result.add decl
-
-  # assemble the parameters
-  let pId = ident"params"
-  let params = genParams(pId, res, vars, iTypes)
-  result.add params
-
-  # launch the kernel
-  result.add quote do:
-    check cuLaunchKernel(
-            `jitFn`,
-            1, 1, 1, # grid(x, y, z)
-            1, 1, 1, # block(x, y, z)
-            sharedMemBytes = 0,
-            CUstream(nil),
-      `pId`[0].unsafeAddr, nil)
-
-  # copy back results
-  let devPtrsRes = determineDevicePtrs(res, nnkBracket.newTree(), @[])
-  for x in devPtrsRes:
-    result.add(
-      check nnkCall.newTree(
-        ident"cuMemcpyDtoH",
-        maybeAddress x[1],
-        x[0],
-        csize_t size(x[1])
-      )
-    )
-
-  # free memory
-  for x in devPtrs:
-    result.add(
-      check nnkCall.newTree(
-        ident"cuMemFree",
-        x[0]
-      )
-    )
-  result = quote do:
-    block:
-      `result`
-
-macro execCuda*(jitFn: CUfunction,
-                res: typed,
-                inputs: typed): untyped =
-  ## Given a CUDA function, execute the kernel. Copies all non trivial data types to
-  ## to the GPU via `cuMemcpyHtoD`. Any argument given as `res` will be copied back
-  ## from the GPU after kernel execution finishes.
-  ##
-  ## IMPORTANT:
-  ## The arguments passed to the CUDA kernel will be in the order in which they are
-  ## given to the macro. This especially means `res` arguments will be passed first.
-  ##
-  ## Example:
-  ## ```nim
-  ## execCuda(fn, res = [r, s], inputs = [a, b, c]) # if all arguments have the same type
-  ## # or
-  ## execCuda(fn, res = (r, s), inputs = (a, b, c)) # if different types
-  ## ```
-  ## will pass the parameters as `[r, s, a, b, c]`.
-  ##
-  ## For more examples see the test case `tests/gpu/t_exec_literals_consts.nim`.
-  ##
-  ## We do not perform any checks on whether the given types are valid as arguments to
-  ## the CUDA target! Also, all arguments given as `res` are expected to be copied.
-  ## To return a value for a simple data type, use a `ptr X` type. However, it is allowed
-  ## to simply pass a `var` symbol as a `res` argument. We automatically copy to the
-  ## the memory location.
-  ##
-  ## We also copy all `res` data to the GPU, so that a return value can also be used
-  ## as an input.
-  ##
-  ## NOTE: This function is mainly intended for convenient execution of a single kernel
-  result = execCudaImpl(jitFn, res, inputs)
-
-macro execCuda*(jitFn: CUfunction,
-                res: typed): untyped =
-  ## Overload of the above for empty `inputs`
-  result = execCudaImpl(jitFn, res, nnkBracket.newTree())
-
 # ############################################################
 #
 #                   Compilation helper
@@ -516,6 +233,12 @@ type
 
   NvidiaAssembler* = ref NvidiaAssemblerObj
 
+  ## We define a distinct version of the `CUfunction` type to differentiate
+  ## producing a kernel via the LLVM backend from the more direct approach
+  ## using NVRTC. This is because the data passing for field elements
+  ## is more complicated on the LLVM side (requires a manual copy).
+  CUfunctionLLVM* = distinct CUfunction
+
 proc `=destroy`*(nv: NvidiaAssemblerObj) =
   ## XXX: Need to also call the finalizer for `asy` in the future!
   # NOTE: In the destructor we don't want to quit on a `check` failure.
@@ -592,7 +315,7 @@ proc initNvAsm*[Name: static Algebra](field: type EC_ShortW_Jac[Fp[Name], G1], w
   result.fd = result.cd.fd
   result.asy.definePrimitives(result.cd)
 
-proc compile*(nv: NvidiaAssembler, kernName: string): CUfunction =
+proc compile*(nv: NvidiaAssembler, kernName: string): CUfunctionLLVM =
   ## Overload of `compile` below.
   ## Call this version if you have manually used the Assembler_LLVM object
   ## to build instructions and have a kernel name you wish to compile.
@@ -617,18 +340,32 @@ proc compile*(nv: NvidiaAssembler, kernName: string): CUfunction =
   check cuModuleLoadData(nv.cuMod, ptx)
   # will be cleaned up when `NvidiaAssembler` goes out of scope
 
-  result = nv.cuMod.getCudaKernel(kernName)
+  result = CUfunctionLLVM(nv.cuMod.getCudaKernel(kernName))
 
-proc compile*(nv: NvidiaAssembler, fn: FieldFnGenerator): CUfunction =
+proc compile*(nv: NvidiaAssembler, fn: FieldFnGenerator): CUfunctionLLVM =
   ## Given a function that generates code for a finite field operation, compile
   ## that function on the given Nvidia target and return a CUDA function.
   # execute the `fn`
   let kernName = nv.asy.fn(nv.fd)
-  result = nv.compile(kernName)
+  result = CUfunctionLLVM(nv.compile(kernName))
 
-proc compile*(nv: NvidiaAssembler, fn: CurveFnGenerator): CUfunction =
+proc compile*(nv: NvidiaAssembler, fn: CurveFnGenerator): CUfunctionLLVM =
   ## Given a function that generates code for an elliptic curve operation, compile
   ## that function on the given Nvidia target and return a CUDA function.
   # execute the `fn`
   let kernName = nv.asy.fn(nv.cd)
-  result = nv.compile(kernName)
+  result = CUfunctionLLVM(nv.compile(kernName))
+
+import ./experimental/cuda_execute_dsl
+macro execCuda*(jitFn: CUfunctionLLVM,
+                res: typed,
+                inputs: typed): untyped =
+  ## See `execCuda` in `constantine/math_compiler/experimental/cuda_execute_dsl.nim`
+  ## for an explanation.
+  ##
+  ## This LLVM overload makes sure we disallow passing simple structs
+  ## via their pointer and instead always copy them (required due to our
+  ## type definitions for finite field elements and elliptic curve points
+  ## on the LLVM target).
+  execCudaImpl(jitFn, newLit 1, newLit 1, res, inputs,
+               passStructByPointer = false)