Stream (cudajit.Cuda.Stream)

type t

Stores a stream pointer and manages lifetimes of kernel launch arguments. See CUstream.

val sexp_of_t : t -> Sexplib0.Sexp.t

val mem_alloc : t -> size_in_bytes:int -> Deviceptr.t

The pointer is finalized using cuMemFreeAsync.

val mem_free : t -> Deviceptr.t -> unit

val memcpy_H_to_D_unsafe : 
  dst:Deviceptr.t ->
  src:unit Ctypes.ptr ->
  size_in_bytes:int ->
  t ->
  unit

See cuMemcpyHtoDAsync.

val memcpy_H_to_D : 
  ?host_offset:int ->
  ?length:int ->
  ?dst_offset:int ->
  dst:Deviceptr.t ->
  src:('a, 'b, 'c) Stdlib.Bigarray.Genarray.t ->
  t ->
  unit

Copies the bigarray (or its interval) into the device memory asynchronously. host_offset and length are in numbers of elements. dst_offset is a device-side byte offset into dst (default 0), letting the copy target a sub-region of a larger allocation; when length is not given the copied size is reduced by dst_offset so it does not write past the end of dst. See memcpy_H_to_D_unsafe.

type kernel_param =

| Tensor of Deviceptr.t
| Tensor_at of Deviceptr.region
(*
Tensor_at { base; offset_bytes } passes base + offset_bytes as the CUdeviceptr kernel argument (tinygrad's buf.value + off pattern). The kernel receives an address displaced by offset_bytes bytes and dereferences from index 0 of that address; the kernel signature is unchanged. check_freed and lifetime bookkeeping operate on base.
*)
| Int of int
(*
Passed as C int.
*)
| Size_t of Unsigned.size_t
| Single of float
(*
Passed as C float.
*)
| Double of float
(*
Passed as C double.
*)

Parameters to pass to a kernel.

val sexp_of_kernel_param : kernel_param -> Sexplib0.Sexp.t

val no_stream : t

The NULL stream which is the main synchronization stream of a device. Manages lifetimes of the corresponding kernel launch parameters.

val launch_kernel : 
  Module.func ->
  grid_dim_x:int ->
  ?grid_dim_y:int ->
  ?grid_dim_z:int ->
  block_dim_x:int ->
  ?block_dim_y:int ->
  ?block_dim_z:int ->
  shared_mem_bytes:int ->
  t ->
  kernel_param list ->
  unit

See cuLaunchKernel.

val memcpy_D_to_H_unsafe : 
  dst:unit Ctypes.ptr ->
  src:Deviceptr.t ->
  size_in_bytes:int ->
  t ->
  unit

See cuMemcpyDtoHAsync.

val memcpy_D_to_H : 
  ?host_offset:int ->
  ?length:int ->
  ?src_offset:int ->
  dst:('a, 'b, 'c) Stdlib.Bigarray.Genarray.t ->
  src:Deviceptr.t ->
  t ->
  unit

Copies from the device memory into the bigarray (or its interval) asynchronously. host_offset and length are in numbers of elements: length is the count of elements copied (not an end index), filling dst.(host_offset .. host_offset + length) and mirroring how memcpy_H_to_D reads src.(host_offset .. host_offset + length). src_offset is a device-side byte offset into src (default 0), letting the copy read from a sub-region of a larger allocation; when length is not given the copied size is reduced by src_offset so it does not read past the end of src. See memcpy_D_to_H_unsafe and cuMemcpyDtoHAsync.

val memcpy_D_to_D : 
  ?kind:('a, 'b) Stdlib.Bigarray.kind ->
  ?length:int ->
  ?size_in_bytes:int ->
  ?dst_offset:int ->
  ?src_offset:int ->
  dst:Deviceptr.t ->
  src:Deviceptr.t ->
  t ->
  unit

Copies between two memory positions on the same device asynchronously. The size to copy can optionally be provided in numbers of elements via kind and length. Provide either both kind and length, or just size_in_bytes. dst_offset and src_offset are device-side byte offsets into dst and src respectively (default 0), letting either end target a sub-region of a larger allocation. See cuMemcpyDtoDAsync.

val memcpy_peer : 
  ?kind:('a, 'b) Stdlib.Bigarray.kind ->
  ?length:int ->
  ?size_in_bytes:int ->
  ?dst_offset:int ->
  ?src_offset:int ->
  dst:Deviceptr.t ->
  dst_ctx:Context.t ->
  src:Deviceptr.t ->
  src_ctx:Context.t ->
  t ->
  unit

Copies between memory positions on two different devices asynchronously. The size to copy can optionally be provided in numbers of elements via kind and length. Provide either both kind and length, or just size_in_bytes. dst_offset and src_offset are device-side byte offsets into dst and src respectively (default 0), letting either end target a sub-region of a larger allocation. See cuMemcpyPeerAsync.

type attach_mem =

| GLOBAL
(*
Memory can be accessed by any stream on any device.
*)
| HOST
(*
Memory cannot be accessed from devices.
*)
| SINGLE_stream
(*
Memory can only be accessed by a single stream.
*)

See CUmemAttach_flags.

val sexp_of_attach_mem : attach_mem -> Sexplib0.Sexp.t

val attach_mem : t -> Deviceptr.t -> int -> attach_mem -> unit

See cuStreamAttachMemAsync.

val create : ?non_blocking:bool -> ?lower_priority:int -> unit -> t

Lower lower_priority numbers represent higher priorities, the default is 0. See cuStreamCreateWithPriority.

The stream value is finalized using cuStreamDestroy. This is meant to be safe without needing to set the proper context.

val get_context : t -> Context.t

See cuStreamGetCtx.

val get_id : t -> Unsigned.uint64

See cuStreamGetId.

val is_ready : t -> bool

Returns false when the querying status is CUDA_ERROR_NOT_READY, and true if it is CUDA_SUCCESS. See cuStreamQuery.

val synchronize : t -> unit

Waits until a stream's tasks are completed. See cuStreamSynchronize.

val memset_d8 : 
  ?offset:int ->
  Deviceptr.t ->
  Unsigned.uchar ->
  length:int ->
  t ->
  unit

offset is a device-side byte offset into Deviceptr.t (default 0), letting the memset target a sub-region of a larger allocation. See cuMemsetD8Async.

val memset_d16 : 
  ?offset:int ->
  Deviceptr.t ->
  Unsigned.ushort ->
  length:int ->
  t ->
  unit

length is in number of elements. offset is a device-side byte offset into Deviceptr.t (default 0). See cuMemsetD16Async.

val memset_d32 : 
  ?offset:int ->
  Deviceptr.t ->
  Unsigned.uint32 ->
  length:int ->
  t ->
  unit

length is in number of elements. offset is a device-side byte offset into Deviceptr.t (default 0). See cuMemsetD32Async.

val total_unreleased_unfinished_delimited_events : t -> int * int * int

Debug information about delimited events carried by the stream: total, unreleased (i.e. not destroyed), unfinished.

val get_total_live_streams : unit -> int

The total non-garbage-collected streams across all devices.