fix omp target

jcosborn · jcosborn · commit b1aac2f9e682 · 2026-02-20T23:15:47.000-06:00
diff --git a/src/backend/coalesced.nim b/src/backend/coalesced.nim
@@ -109,10 +109,11 @@ proc newCoalesced*[T](V,M:static[int], p:ptr T, n:int):auto {.noinit.} =
 template `[]`*(x:Coalesced, ix:int):untyped = CoalescedObj[x.V,x.M,x.T](o:x, i:ix)
 template len*(x:Coalesced):untyped = x.n
 
-template fromCoalesced*(x:CoalescedObj):untyped =
+template fromCoalesced*(x:CoalescedObj): auto =
   const N = getSize(x.T) div (x.M*sizeof(RegisterWord))
   type A = ptr Uncheckedarray[MemoryWord(x.M)]
-  var r {.noinit.}: x.T
+  #var r {.noinit.}: x.T
+  var r: x.T
   let offset = (x.i div x.V)*N*x.V + x.i mod x.V
   unrollfor:
     for j in 0..N-1: cast[A](r.addr)[j] = cast[A](x.o.p)[offset + j*x.V]
@@ -145,8 +146,8 @@ template `:=`*[V,M:static[int],X,Y](x:CoalescedObj[V,M,X], y:Y) =
 
 proc `*`*[VX,MX,VY,MY:static[int],X,Y](x:CoalescedObj[VX,MX,X], y:CoalescedObj[VY,MY,Y]):auto {.noinit,inline.} =
   let
-    tx {.noinit.} = fromCoalesced(x)
-    ty {.noinit.} = fromCoalesced(y)
+    tx = fromCoalesced(x)
+    ty = fromCoalesced(y)
   mixin `*`
   tx * ty
 
diff --git a/src/backend/linalg.nim b/src/backend/linalg.nim
@@ -4,10 +4,10 @@ type SomeNumber2* = SomeNumber
 template `:=`*(x: var SomeNumber, y: SomeNumber2) =
   type tx = type(x)
   x = (tx)(y)
-template `+=`*(x: var SomeNumber, y: SomeNumber2) =
-  bind `+=`    # So the following += doesn't call this template again.
-  type tx = type(x)
-  x += (tx)(y)
+#template `+=`*(x: var SomeNumber, y: SomeNumber2) =
+#  bind `+=`    # So the following += doesn't call this template again.
+#  type tx = type(x)
+#  x += (tx)(y)
 
 type
   Complex*[T] = object
@@ -63,7 +63,8 @@ template `*`*[T](x,y: Complex[T]): untyped =
     xx = x
     yy = y
   type tx = type(xx)
-  var r {.noInit.}: tx #Complex[x.T] #Complex[type(x.re*y.re)]
+  #var r {.noInit.}: tx #Complex[x.T] #Complex[type(x.re*y.re)]
+  var r: tx #Complex[x.T] #Complex[type(x.re*y.re)]
   r.re := xx.re*yy.re - xx.im*yy.im
   r.im := xx.re*yy.im + xx.im*yy.re
   r
@@ -110,7 +111,8 @@ template `+`*[N:static[int],T](x,y: Colmat[N,T]): untyped =
 template `*`*[N:static[int],T](x,y: Colmat[N,T]): untyped =
   let xx = x
   let yy = y
-  var r {.noInit.}: Colmat[N,type(x.d[0][0].re)]
+  #var r {.noInit.}: Colmat[N,type(x.d[0][0].re)]
+  var r: Colmat[N,type(x.d[0][0].re)]
   forstatic i, 0, N-1:
     forstatic j, 0, N-1:
       r.d[i][j] := xx.d[i][0] * yy.d[0][j]
diff --git a/src/backend/openmp.nim b/src/backend/openmp.nim
@@ -3,6 +3,9 @@ import base/metaUtils
 import base/omp
 
 {.pragma: omp, header:"omp.h".}
+{.passC:"-fcf-protection=none -no-pie -fno-stack-protector" .}
+{.passL:"-fcf-protection=none -no-pie -fno-stack-protector" .}
+
 template mkMemoryPragma*:untyped =
   {.pragma: restrict, codegenDecl: "$# __restrict__ $#".}
   {.pragma: aligned, codegenDecl: "$# $# __attribute__((aligned))".}
@@ -34,7 +37,7 @@ proc alignatImpl(n:NimNode, byte:int): NimNode =
     for c in n:
       result.add c.alignatImpl byte
 macro alignat*(byte:static[int], n:untyped): untyped =
-  if byte notin {1,2,4,8,16,32,64,128,256}:
+  if byte notin [1,2,4,8,16,32,64,128,256]:
     error("macro alignat: unsupported alignment: " & $byte, n)
   #echo "alignatImpl ", byte
   #echo n.treerepr
@@ -61,11 +64,11 @@ proc omp_get_initial_device*: cint {.omp.}
 proc omp_get_num_teams*: cint {.omp.}
 proc omp_get_team_num*: cint {.omp.}
 
-template omp_target_alloc*(size: int): pointer =
+template omp_target_alloc*(size: SomeNumber): pointer =
   omp_target_alloc(csize_t size, omp_get_default_device())
 template omp_target_memcpy_tocpu*(dst: pointer, src: pointer; length: csize_t): cint =
   omp_target_memcpy(dst, src, length, 0, 0, omp_get_initial_device(), omp_get_default_device())
-template omp_target_memcpy_togpu*(dst: pointer, src: pointer; length: int): cint =
+template omp_target_memcpy_togpu*(dst: pointer, src: pointer; length: csize_t): cint =
   omp_target_memcpy(dst, src, csize_t length, 0, 0, omp_get_default_device(), omp_get_initial_device())
 template omp_target_free*(device_ptr: pointer) =
   omp_target_free(device_ptr, omp_get_default_device())
@@ -130,7 +133,8 @@ proc prepareVars(n:NimNode):seq[NimNode] =
         nnkWhileStmt, nnkForStmt} + RoutineNodes:
       # New lexical scope
       newscope = true
-      ignoreStack.add newPar()
+      #ignoreStack.add newPar()
+      ignoreStack.add newNimNode(nnkTupleConstr)
     for i in 0..<n.len:
       #echo "### ",n[i].lisprepr
       case n[i].kind
@@ -176,7 +180,8 @@ proc prepareVars(n:NimNode):seq[NimNode] =
             let np = gensym(nsklet, "gpu_ptr_" & $n[i])
             ignoreStack[0].add nv
             ignoreStack[0].add np
-            openvars.add newpar(n[i], nv, np)
+            #openvars.add newpar(n[i], nv, np)
+            openvars.add newNimNode(nnkTupleConstr).add(n[i], nv, np)
             n[i] = newcall("gpuVarPtr",nv,np)
       else:
         discard
@@ -199,6 +204,7 @@ proc genCpuPrepare(n:seq[NimNode]):NimNode =
       var v{.noinit.}:OffloadDummy[typeof(x)]
   result = newstmtlist()
   for c in n:
+    #echo c.treerepr
     result.add getast r(c[0],c[1],c[2])
 proc genGpuPrepare(n:seq[NimNode]):NimNode =
   template r(x,v,p:untyped):untyped =
@@ -214,31 +220,55 @@ proc genCpuFinalize(n:seq[NimNode]):NimNode =
   result = newstmtlist()
   for c in n:
     result.add getast r(c[0],c[1],c[2])
-proc declarePtrString(n:seq[NimNode]):NimNode =
-  template res(ptrlist:untyped):untyped =
-    const s = ptrlist
-    when s.len == 0: "" else: "is_device_ptr(" & s[0..^2] & ")"
-  template varname(x, xp:untyped):untyped =
-    mixin offloadPtr
-    when compiles(offloadPtr(x)): xp&"," else: ""
-  var ps = newlit""
+#proc declarePtrString(n:seq[NimNode]):NimNode =
+#  template res(ptrlist:untyped):untyped =
+#    const s = ptrlist
+#    when s.len == 0: "" else: "is_device_ptr(" & s[0..^2] & ")"
+#  template varname(x, xp:untyped):untyped =
+#    mixin offloadPtr
+#    when compiles(offloadPtr(x)): xp&"," else: ""
+#  var ps = newlit""
+#  for c in n:
+#    ps = infix(getast varname(c[0], $c[2]), "&", ps)
+#  result = getast res(ps)
+proc declarePtrTuple(n:seq[NimNode]):NimNode =
+  mixin offloadPtr
+  var ps = newNimNode(nnkTupleConstr)
+  ps.add newLit"is_device_ptr("
   for c in n:
-    ps = infix(getast varname(c[0], $c[2]), "&", ps)
-  result = getast res(ps)
+    when compiles(offloadPtr(c[0])):
+      ps.add c[2]
+  if ps.len == 1:
+    result = newNimNode(nnkTupleConstr)
+  else:
+    ps.add newLit")"
+    result = ps
+  #echo result.treerepr
 
-macro isDevicePtr(x: typed): untyped =
-  let n = $x
-  result = newLit(" is_device_ptr("&n&")")
+#macro isDevicePtr(x: typed): untyped =
+#  let n = $x
+#  result = newLit(" is_device_ptr("&n&")")
 
-macro useDevicePtr(x: typed): untyped =
-  let n = $x
-  let p = newLit("#pragma omp target data use_device_ptr("&n&")")
-  result = quote do:
-    {.emit: `p`.}
+#macro useDevicePtr(x: auto): auto =
+  #echo x.treerepr
+  #let n = x.strVal
+  #echo "useDevicePtr: ", n
+  #let p = newLit("#pragma omp target data use_device_ptr("&n&")")
+#  result = quote do:
+#    {.emit: ["#pragma omp target data use_device_ptr(",`x`,")"].}
 
-macro mapto(x: typed): untyped =
-  let n = $x
-  result = newLit(" map(to:"&n&")")
+#macro getrepr(x: untyped): auto =
+#  echo x.treerepr
+#  result = x
+
+template useDevicePtr(x: auto) =
+  #getrepr:
+  {.emit: ["#pragma omp target data use_device_ptr(",x,")"].}
+
+#macro mapto(x: typed): untyped =
+#  let n = $x
+#  result = newLit(" map(to:"&n&")")
+#macro mapto(x: typed): untyped =
 
 macro onGpu*(body: untyped): untyped =
   # the architecture for cpugpuarray requires us replace body before it gets expanded, so we require untyped.
@@ -248,8 +278,8 @@ macro onGpu*(body: untyped): untyped =
     {.push stacktrace: off.}
     proc gpuProc {.gensym.} =
       cpuPrepare  # a let section declare and save device pointers
-      const isDevicePtrList = devicePtrDeclare  # is_device_ptr(ptrList) in string
-      ompBlock("target teams " & isDevicePtrList):
+      #const isDevicePtrList = devicePtrDeclare  # is_device_ptr(ptrList) in string
+      ompBlock2("target teams ", devicePtrDeclare):
         openmpDefs:
           gpuPrepare
           body
@@ -260,7 +290,7 @@ macro onGpu*(body: untyped): untyped =
     cpuPrepare = genCpuPrepare v
     gpuPrepare = genGpuPrepare v
     cpuFinalize = genCpuFinalize v
-    isDevicePtrs = declarePtrString v
+    isDevicePtrs = declarePtrTuple v
   result = getast(target(cpuPrepare, gpuPrepare, cpuFinalize, isDevicePtrs, body))
   #echo result.repr
 
@@ -306,10 +336,10 @@ proc identStr(n:NimNode):string =
   result = n.repr
   for i in 0..<result.len:
     if result[i] in {'.','[',']',':'}: result[i] = '_'
-proc isIndex(n,i:NimNode):bool =
-  result = n.eqident i
-  if n.kind == nnkHiddenStdConv:
-    result = n[1].eqident i
+#proc isIndex(n,i:NimNode):bool =
+#  result = n.eqident i
+#  if n.kind == nnkHiddenStdConv:
+#    result = n[1].eqident i
 macro simdForImpl(n:typed):untyped =
   proc getIndexedPtrs(n,i:NimNode):(NimNode,seq[NimNode]) =
     #echo "### getIndexedPtrs: ", i.repr
@@ -323,7 +353,8 @@ macro simdForImpl(n:typed):untyped =
               break
           if m < 0:
             let v = gensym(nskVar, n.cleanAst.identStr)
-            ptrs.add newPar(v, n)
+            #ptrs.add newPar(v, n)
+            ptrs.add newNimNode(nnkTupleConstr).add(v, n)
             return v
           else:
             return ptrs[m][0]
@@ -437,20 +468,19 @@ when isMainModule:
     useDevicePtr(y)
     discard omp_target_memcpy_togpu(y, addr x, sizeof(float32))
     #ompBlock("target teams"&isDevicePtr(x)):
-    ompBlock("target teams"&mapto(x)):
+    #ompBlock("target teams"&mapto(x)):
+    ompBlock2("target teams", " map(to:", x, ")"):
       {.emit:"#pragma omp parallel".}
       {.emit:"for(int ii=0; ii<1; ii++)".}
       block:
         x = 1.0
 
-    macro dump(n:typed):typed =
-      echo n.repr
-      n
-    #[
-    dump:
-      onGpu:
-        let i = getThreadNum()
-        if i < n:
-          c[i] = a[i] + b[i]
-    ]#
+    #macro dump(n:auto):auto =
+    #  echo n.repr
+    #  n
+    onGpu:
+      let i = getThreadNum()
+      if i < n:
+        c[i] = a[i] + b[i]
+
   test()
diff --git a/src/backend/vectorized.nim b/src/backend/vectorized.nim
@@ -29,7 +29,7 @@ import macros
 mkMemoryPragma()
 
 const CPUVLEN* {.intdefine.} = 0 ## CPU SIMD vector length in bits.  Zero lets compiler auto-vectorize.
-const SupportedCPUVLENs = {128,256,512}
+const SupportedCPUVLENs = [128,256,512]
 const oneByte = 8
 macro defsimd:auto =
   var s,d:NimNode
diff --git a/src/base/omp.nim b/src/base/omp.nim
@@ -1,4 +1,4 @@
-import os
+import os, macros
 
 when defined(noOpenmp):
   static: echo "OpenMP disabled"
@@ -42,6 +42,38 @@ else:
       body
     #{. emit:"} /* End ompBlock " & p & " */".}
 
+  macro ompPragma2*(p: varargs[untyped]): auto =
+    var b = newNimNode(nnkBracket)
+    b.add newLit "_Pragma(\"omp "
+    for i in 0..<p.len:
+      if p[i].kind == nnkTupleConstr:
+        for j in 0..<p[i].len:
+          b.add p[i][j]
+      else:
+        b.add p[i]
+    b.add newLit "\")"
+    #echo p.treerepr
+    result = quote do:
+      {. emit:`b` .}
+  #template ompBlock2*(p: varargs[untyped]) =
+  #  ompPragma2(p[0..^2])
+  #  block:
+  #    p[^1]
+  macro ompBlock2*(p: varargs[untyped]): auto =
+    {. push stackTrace:off, lineTrace:off, line_dir:off .}
+    #echo p.treerepr
+    let body = p[^1]
+    var p2 = newNimNode(nnkCall).add bindSym"ompPragma2"
+    for i in 0..(p.len-2):
+      p2.add p[i]
+    #echo body.treerepr
+    result = quote do:
+      `p2`
+      block:
+        `body`
+    #echo result.treerepr
+    {. pop .}
+
 template ompBarrier* = ompPragma("barrier")
 template ompFlush* = ompPragma("flush")
 template ompFlushAcquire* = ompPragma("flush acquire")
