WebGPU GPU compiler backend improvements (and small CUDA)

constantine/math_compiler/experimental/backends/backends.nim

@@ -14,27 +14,11 @@ when defined(cuda):
 else:
   const Backend* = bkWGSL
 
-proc gpuTypeToString*(t: GpuTypeKind): string =
-  case Backend
-  of bkCuda: cuda.gpuTypeToString(t)
-  of bkWGSL: wgsl.gpuTypeToString(t)
-
-proc gpuTypeToString*(t: GpuType, ident = newGpuIdent(), allowArrayToPtr = false,
-                      allowEmptyIdent = false,
-                     ): string =
-  case Backend
-  of bkCuda: cuda.gpuTypeToString(t, ident.ident(), allowArrayToPtr, allowEmptyIdent)
-  of bkWGSL: wgsl.gpuTypeToString(t, ident, allowArrayToPtr, allowEmptyIdent)
-
-proc genFunctionType*(typ: GpuType, fn: string, fnArgs: string): string =
-  case Backend
-  of bkCuda: cuda.genFunctionType(typ, fn, fnArgs)
-  of bkWGSL: wgsl.genFunctionType(typ, fn, fnArgs)
-
 proc codegen*(ctx: var GpuContext, ast: GpuAst, kernel: string = ""): string =
   case Backend
   of bkCuda:
-    result = ctx.genCuda(ast)
+    cuda.preprocess(ctx, ast, kernel)
+    result = cuda.codegen(ctx)
   of bkWGSL:
-    ctx.storagePass(ast, kernel)
+    wgsl.preprocess(ctx, ast, kernel)
     result = wgsl.codegen(ctx)

constantine/math_compiler/experimental/backends/common_utils.nim

@@ -6,9 +6,36 @@
 #   * Apache v2 license (license terms in the root directory or at http://www.apache.org/licenses/LICENSE-2.0).
 # at your option. This file may not be copied, modified, or distributed except according to those terms.
 
+import std / tables
 import ../gpu_types
-# import ./backends
 
 proc address*(a: string): string = "&" & a
-
 proc size*(a: string): string = "sizeof(" & a & ")"
+
+proc isGlobal*(fn: GpuAst): bool =
+  doAssert fn.kind == gpuProc, "Not a function, but: " & $fn.kind
+  result = attGlobal in fn.pAttributes
+
+proc farmTopLevel*(ctx: var GpuContext, ast: GpuAst, kernel: string, varBlock: var GpuAst) =
+  ## Farms the top level of the code for functions, variable and type definition.
+  ## All functions are added to the `allFnTab`, while only global ones (or even only
+  ## `kernel` if any) is added to the `fnTab` as the starting point for the remaining
+  ## logic.
+  ## Variables are collected in `varBlock`.
+  case ast.kind
+  of gpuProc:
+    ctx.allFnTab[ast.pName] = ast
+    if kernel.len > 0 and ast.pName.ident() == kernel and ast.isGlobal():
+      ctx.fnTab[ast.pName] = ast.clone() # store global function extra
+    elif kernel.len == 0 and ast.isGlobal():
+      ctx.fnTab[ast.pName] = ast.clone() # store global function extra
+  of gpuBlock:
+    # could be a type definition or global variable
+    for ch in ast:
+      ctx.farmTopLevel(ch, kernel, varBlock)
+  of gpuVar, gpuConstexpr:
+    varBlock.statements.add ast
+  of gpuTypeDef, gpuAlias:
+    raiseAssert "Unexpected type def / alias def found. These should be in `ctx.types` now: " & $ast
+  else:
+    discard