Class: Cumo::CUDA::Stream

Inherits:
Object
  • Object
show all
Defined in:
lib/cumo/cuda/stream.rb

Overview

A CUDA stream. Cumo launches every kernel and copy on the current stream of the thread, which is the null stream until a Stream is used.

s = Cumo::CUDA::Stream.new
s.with { c = a.gemm(b) }   # queued on s, and waited for on the way out

Fibers share a thread's current stream.

Constant Summary collapse

LIVE = 
ObjectSpace::WeakMap.new

Instance Attribute Summary collapse

Class Method Summary collapse

Instance Method Summary collapse

Constructor Details

#initialize(non_blocking: false) ⇒ Stream

Returns a new instance of Stream.



38
39
40
41
42
43
 
# File 'lib/cumo/cuda/stream.rb', line 38

def initialize(non_blocking: false)
  flags = non_blocking ? Runtime::CUDA_STREAM_NON_BLOCKING : Runtime::CUDA_STREAM_DEFAULT
  @ptr = Runtime.cudaStreamCreateWithFlags(flags)
  LIVE[@ptr] = self
  ObjectSpace.define_finalizer(self, self.class.finalizer(@ptr))
end

Instance Attribute Details

#ptr ⇒ Object (readonly)

Returns the value of attribute ptr.



14
15
16
 
# File 'lib/cumo/cuda/stream.rb', line 14

def ptr
  @ptr
end

Class Method Details

.current ⇒ Object

The Stream whose handle is the thread's current stream.



21
22
23
24
25
 
# File 'lib/cumo/cuda/stream.rb', line 21

def self.current
  ptr = Runtime.current_stream
  return null if ptr == 0
  LIVE[ptr] ||= allocate.tap { |s| s.instance_variable_set(:@ptr, ptr) }
end

.finalizer(ptr) ⇒ Object

A handle may be given out again after a destroy, so a stream that is still owned by a live object is not the one this finalizer was for.



29
30
31
32
33
34
35
36
 
# File 'lib/cumo/cuda/stream.rb', line 29

def self.finalizer(ptr)
  proc do
    unless LIVE[ptr]
      MemoryPool.free_all_blocks(ptr) rescue nil
      Runtime.cudaStreamDestroy(ptr) rescue nil
    end
  end
end

.null ⇒ Object 



16
17
18
 
# File 'lib/cumo/cuda/stream.rb', line 16

def self.null
  @null ||= allocate.tap { |s| s.instance_variable_set(:@ptr, 0) }
end

Instance Method Details

#destroy ⇒ Object

Waits for the stream, hands its pooled memory back, and destroys it. A stream that is current in a thread is refused, and stays usable.



104
105
106
107
108
109
110
111
112
 
# File 'lib/cumo/cuda/stream.rb', line 104

def destroy
  return if @ptr.nil? || null?
  Runtime.cudaStreamSynchronize(@ptr)
  Runtime.cudaStreamDestroy(@ptr)
  MemoryPool.free_all_blocks(@ptr)
  ObjectSpace.undefine_finalizer(self)
  LIVE.delete(@ptr) if LIVE.respond_to?(:delete)
  @ptr = nil
end

#done? ⇒ Boolean

Returns:

  • (Boolean) 


87
88
89
 
# File 'lib/cumo/cuda/stream.rb', line 87

def done?
  Runtime.cudaStreamQuery(handle)
end

#handle ⇒ Object

The handle of a stream that has not been destroyed.

Raises:

  • (ArgumentError) 


50
51
52
53
 
# File 'lib/cumo/cuda/stream.rb', line 50

def handle
  raise ArgumentError, "the stream is destroyed" if @ptr.nil?
  @ptr
end

#null? ⇒ Boolean

Returns:

  • (Boolean) 


45
46
47
 
# File 'lib/cumo/cuda/stream.rb', line 45

def null?
  @ptr == 0
end

#record(event = nil) ⇒ Object

Records an event after everything queued so far, and answers it.



92
93
94
 
# File 'lib/cumo/cuda/stream.rb', line 92

def record(event = nil)
  (event || Event.new).record(self)
end

#synchronize ⇒ Object 



82
83
84
85
 
# File 'lib/cumo/cuda/stream.rb', line 82

def synchronize
  Runtime.cudaStreamSynchronize(handle)
  self
end

#use ⇒ Object

Makes this the current stream of the thread until another is used.



56
57
58
59
 
# File 'lib/cumo/cuda/stream.rb', line 56

def use
  Runtime.current_stream = handle
  self
end

#wait_event(event) ⇒ Object

Makes everything queued after this call wait for the event.



97
98
99
100
 
# File 'lib/cumo/cuda/stream.rb', line 97

def wait_event(event)
  Runtime.cudaStreamWaitEvent(handle, event.handle)
  self
end

#with ⇒ Object

Runs the block with this as the current stream, waits for everything the block queued, and puts the previous stream back. The block's work is ordered after what the previous stream had queued, and nothing the block started is left running once it returns.



65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
 
# File 'lib/cumo/cuda/stream.rb', line 65

def with
  before = Stream.current
  if before.handle != handle
    wait_event(before.record(Event.new(timing: false)))
  end
  Runtime.current_stream = handle
  begin
    yield self
  ensure
    begin
      synchronize unless null?
    ensure
      Runtime.current_stream = before.handle
    end
  end
end