gorgonia.org/gorgonia@v0.9.17/cuda/extension.go (about)

     1  package cuda
     2  
     3  import (
     4  	"fmt"
     5  
     6  	"github.com/pkg/errors"
     7  	"gorgonia.org/cu"
     8  )
     9  
    10  // this file relates to code that allows you to extend Engine
    11  
    12  // LoadCUDAFunc loads a string representing a CUDA PTX file into the engine, giving it the universe of computing functions.
    13  func (e *Engine) LoadCUDAFunc(moduleName, data string, funcs []string) (err error) {
    14  	fns := e.f
    15  	if fns == nil {
    16  		fns = make(map[string]cu.Function)
    17  	}
    18  	if err = cu.SetCurrentContext(e.c.Context.CUDAContext()); err != nil {
    19  		return errors.Wrapf(err, "Unable to set current context when loading module %q at device %v", moduleName, e.d)
    20  	}
    21  
    22  	var mod cu.Module
    23  	if mod, err = cu.LoadData(data); err != nil {
    24  		return errors.Wrapf(err, "Failed to load module %q data for Device %v context %x", moduleName, e.d, e.c)
    25  	}
    26  
    27  	for _, name := range funcs {
    28  		var fn cu.Function
    29  		if fn, err = mod.Function(name); err != nil {
    30  			return errors.Wrapf(err, "Unable to get function %q in Device %v context %x", name, e.d, e.c)
    31  		}
    32  		fqn := fmt.Sprintf("%v.%v", moduleName, name)
    33  		fns[fqn] = fn
    34  	}
    35  	if e.m == nil {
    36  		e.m = make(map[string]cu.Module)
    37  	}
    38  	e.m[moduleName] = mod
    39  	e.f = fns
    40  	return nil
    41  }