Class: Cumo::CUDA::Compiler
- Inherits:
-
Object
- Object
- Cumo::CUDA::Compiler
- 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.('~/.cumo/kernel_cache')
Class Method Summary collapse
Instance Method Summary collapse
- #compile_using_nvrtc(source, options: [], arch: nil) ⇒ Object
- #compile_with_cache(source, options: [], arch: nil, cache_dir: nil, extra_source: nil) ⇒ Object
Class Method Details
.valid_kernel_name?(name) ⇒ Boolean
14 15 16 |
# File 'lib/cumo/cuda/compiler.rb', line 14 def self.valid_kernel_name?(name) VALID_KERNEL_NAME.match?(name) end |
Instance Method Details
#compile_using_nvrtc(source, options: [], arch: nil) ⇒ Object
18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 |
# File 'lib/cumo/cuda/compiler.rb', line 18 def compile_using_nvrtc(source, options: [], arch: nil) arch ||= get_arch += ["-arch=#{arch}"] 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) begin ptx = prog.compile(options: ) 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 return ptx end end |
#compile_with_cache(source, options: [], arch: nil, cache_dir: nil, extra_source: nil) ⇒ Object
45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 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 |
# File 'lib/cumo/cuda/compiler.rb', line 45 def compile_with_cache(source, options: [], arch: nil, cache_dir: nil, extra_source: nil) # NVRTC does not use extra_source. extra_source is used for cache key. cache_dir ||= get_cache_dir arch ||= get_arch += ['-ftz=true'] env = [arch, , get_nvrtc_version] base = @@empty_file_preprocess_cache[env] if base.nil? # This is checking of NVRTC compiler internal version base = preprocess('', , arch) @@empty_file_preprocess_cache[env] = base end key_src = "#{env} #{base} #{source} #{extra_source}" key_src.encode!('utf-8') digest = Digest::MD5.hexdigest(key_src) name = "#{digest}_2.cubin" unless Dir.exist?(cache_dir) FileUtils.mkdir_p(cache_dir) end # TODO(sonots): thread-safe? path = File.join(cache_dir, name) cubin = load_cache(path) if cubin mod = Module.new mod.load(cubin) return mod end ptx = compile_using_nvrtc(source, options: , arch: arch) cubin = nil cubin_hash = nil LinkState.new do |ls| ls.add_ptr_data(ptx, 'cumo.ptx') cubin = ls.complete() cubin_hash = Digest::MD5.hexdigest(cubin) end save_cache(path, cubin_hash, cubin) # 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) return mod end |