Skip to content

Commit

Permalink
added a function pointer field to CpuStorageObj for custom destructio…
Browse files Browse the repository at this point in the history
…n logic
  • Loading branch information
darkestpigeon committed Aug 20, 2024
1 parent caf0d8d commit 65356f1
Showing 1 changed file with 38 additions and 6 deletions.
44 changes: 38 additions & 6 deletions src/arraymancer/laser/tensor/datatypes.nim
Original file line number Diff line number Diff line change
Expand Up @@ -40,9 +40,10 @@ type
CpuStorageObj[T] {.shallow.} = object
# Workaround supportsCopyMem in type section - https://github.com/nim-lang/Nim/issues/13193
when T is KnownSupportsCopyMem:
raw_buffer*: ptr UncheckedArray[T] # 8 bytes
memalloc*: pointer # 8 bytes
isMemOwner*: bool # 1 byte
raw_buffer*: ptr UncheckedArray[T] # 8 bytes
memalloc*: pointer # 8 bytes
isMemOwner*: bool # 1 byte
deleter*: proc(memalloc: pointer) {.noconv, gcsafe, raises: [].} # 8 bytes
else: # Tensors of strings, other ref types or non-trivial destructors
raw_buffer*: seq[T] # 8 bytes (16 for seq v2 backed by destructors?)

Expand Down Expand Up @@ -74,14 +75,21 @@ when not defined(gcDestructors):
proc finalizer[T](storage: CpuStorage[T]) =
static: assert T is KnownSupportsCopyMem, "Tensors of seq, strings, ref types and types with non-trivial destructors cannot be finalized by this proc"
if storage.isMemOwner and not storage.memalloc.isNil:
storage.memalloc.deallocShared()
if storage.deleter.isNil:
storage.memalloc.deallocShared()
else:
storage.deleter(storage.memalloc)
storage.deleter = nil
storage.memalloc = nil
else:
when (NimMajor, NimMinor, NimPatch) >= (2, 0, 0):
proc `=destroy`[T](storage: CpuStorageObj[T]) =
when T is KnownSupportsCopyMem:
if storage.isMemOwner and not storage.memalloc.isNil:
storage.memalloc.deallocShared()
if storage.deleter.isNil:
storage.memalloc.deallocShared()
else:
storage.deleter(storage.memalloc)
else:
# The following cast removes the following warning:
# =destroy(storage.raw_buffer) can raise an unlisted exception: Exception
Expand All @@ -91,7 +99,11 @@ else:
proc `=destroy`[T](storage: var CpuStorageObj[T]) =
when T is KnownSupportsCopyMem:
if storage.isMemOwner and not storage.memalloc.isNil:
storage.memalloc.deallocShared()
if storage.deleter.isNil:
storage.memalloc.deallocShared()
else:
storage.deleter(storage.memalloc)
storage.deleter = nil
storage.memalloc = nil
else:
`=destroy`(storage.raw_buffer)
Expand All @@ -113,6 +125,7 @@ proc allocCpuStorage*[T](storage: var CpuStorage[T], size: int) =
storage.memalloc = allocShared(sizeof(T) * size + LASER_MEM_ALIGN - 1)
storage.isMemOwner = true
storage.raw_buffer = align_raw_data(T, storage.memalloc)
storage.deleter = nil
else: # Always 0-initialize Tensors of seq, strings, ref types and types with non-trivial destructors
new(storage)
storage.raw_buffer.newSeq(size)
Expand All @@ -133,6 +146,25 @@ proc cpuStorageFromBuffer*[T: KnownSupportsCopyMem](
storage.memalloc = rawBuffer
storage.isMemOwner = false
storage.raw_buffer = cast[ptr UncheckedArray[T]](storage.memalloc)
storage.deleter = nil

proc cpuStorageFromBuffer*[T: KnownSupportsCopyMem](
storage: var CpuStorage[T],
rawBuffer: pointer,
memalloc: pointer,
deleter: proc(memalloc: pointer) {.noconv, gcsafe, raises: [].}) =
## Create a `CpuStorage`, which stores data from a given raw pointer, which
## it will own. The destructor/finalizer will do a `deleter(memalloc)` call.
##
## The input buffer must be a raw `pointer`.
when not defined(gcDestructors):
new(storage, finalizer[T])
else:
new(storage)
storage.memalloc = memalloc
storage.isMemOwner = true
storage.raw_buffer = cast[ptr UncheckedArray[T]](rawBuffer)
storage.deleter = deleter

func is_C_contiguous*(t: Tensor): bool =
## Check if the tensor follows C convention / is row major
Expand Down

0 comments on commit 65356f1

Please sign in to comment.