Class: Cumo::CUDA::Compiler

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

Constant Summary collapse

VALID_KERNEL_NAME = 
/\A[a-zA-Z_][a-zA-Z_0-9]*\z/
DEFAULT_CACHE_DIR = 
File.expand_path('~/.cumo/kernel_cache')

Class Method Summary collapse

Instance Method Summary collapse

Class Method Details

.clear_modules ⇒ Object

The next compile_with_cache reads the disk cache or compiles again. A module nobody holds any more is unloaded when Ruby collects it.



24
25
26
 
# File 'lib/cumo/cuda/compiler.rb', line 24

def self.clear_modules
  @@modules.clear
end

.remove_module(mod) ⇒ Object 



28
29
30
 
# File 'lib/cumo/cuda/compiler.rb', line 28

def self.remove_module(mod)
  @@modules.delete_if { |_, m| m.equal?(mod) }
end

.valid_kernel_name?(name) ⇒ Boolean

Returns:

  • (Boolean) 


18
19
20
 
# File 'lib/cumo/cuda/compiler.rb', line 18

def self.valid_kernel_name?(name)
  VALID_KERNEL_NAME.match?(name)
end

Instance Method Details

#compile_using_nvrtc(source, options: [], arch: nil, name_expressions: nil) ⇒ Object

With name_expressions: answers the PTX and a Hash of each expression to its mangled name; without, the PTX.



34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
 
# File 'lib/cumo/cuda/compiler.rb', line 34

def compile_using_nvrtc(source, options: [], arch: nil, name_expressions: nil)
  arch ||= get_arch
  options += ["-arch=#{arch}"]
  name_expressions = nil if name_expressions && name_expressions.empty?

  Dir.mktmpdir do |root_dir|
    path = File.join(root_dir, 'kern')
    cu_path = "#{path}.cu"

    File.open(cu_path, 'w') do |cu_file|
      cu_file.write(source)
    end

    prog = NVRTCProgram.new(source, name: cu_path, name_expressions: name_expressions || [])
    begin
      ptx = prog.compile(options: options)
      return ptx if name_expressions.nil?
      return [ptx, name_expressions.to_h { |expr| [expr, prog.lowered_name(expr)] }]
    rescue CompileError => e
      if get_bool_env_variable('CUMO_DUMP_CUDA_SOURCE_ON_ERROR', false)
        e.dump($stderr)
      end
      raise e
    ensure
      prog.destroy
    end
  end
end

#compile_with_cache(source, options: [], arch: nil, cache_dir: nil, extra_source: nil, name_expressions: []) ⇒ Object

name_expressions: names such as "kernel" that the Module then answers get_function for. Their mangled names are cached with the cubin. The same source, options and device answer the same Module for the rest of the process.



67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
 
# File 'lib/cumo/cuda/compiler.rb', line 67

def compile_with_cache(source, options: [], arch: nil, cache_dir: nil, extra_source: nil, name_expressions: [])
  # NVRTC does not use extra_source. extra_source is used for cache key.
  cache_dir ||= get_cache_dir
  arch ||= get_arch
  name_expressions = name_expressions.uniq
  memo_key = [source, options, arch, cache_dir, extra_source, name_expressions, Runtime.cudaGetDevice]
  mod = @@modules[memo_key]
  return mod if mod

  options += ['-ftz=true']

  env = [arch, options, get_nvrtc_version]
  base = @@empty_file_preprocess_cache[env]
  if base.nil?
    # This is checking of NVRTC compiler internal version
    base = preprocess('', options, arch)
    @@empty_file_preprocess_cache[env] = base
  end
  key_src = "#{env} #{base} #{source} #{extra_source}"
  key_src += " #{name_expressions}" unless name_expressions.empty?

  key_src.encode!('utf-8')
  digest = Digest::MD5.hexdigest(key_src)
  name = "#{digest}_#{name_expressions.empty? ? 2 : 3}.cubin"

  unless Dir.exist?(cache_dir)
    FileUtils.mkdir_p(cache_dir)
  end

  # TODO(sonots): thread-safe?
  path = File.join(cache_dir, name)
  cubin, lowered = load_cache(path, name_expressions)
  if cubin
    mod = Module.new
    mod.load(cubin)
    mod.lowered_names = lowered
    return remember(memo_key, mod)
  end

  ptx = compile_using_nvrtc(source, options: options, arch: arch, name_expressions: name_expressions)
  ptx, lowered = name_expressions.empty? ? [ptx, {}] : ptx
  cubin = nil
  LinkState.new do |ls|
    ls.add_ptr_data(ptx, 'cumo.ptx')
    cubin = ls.complete()
  end

  save_cache(path, cubin, lowered)

  # Save .cu source file along with .cubin
  if get_bool_env_variable('CUMO_CACHE_SAVE_CUDA_SOURCE', false)
    File.open("#{path}.cu", 'w') do |f|
      f.write(source)
    end
  end

  mod = Module.new
  mod.load(cubin)
  mod.lowered_names = lowered
  return remember(memo_key, mod)
end