cuda.sl.in

import("cuda");
require("rand");

variable CURAND_DEFAULT_TYPE=Float_Type;
variable CURAND_DEFAULT_GEN_TYPE=CURAND_RNG_PSEUDO_DEFAULT;
variable CURAND_DEFAULT_GEN;
private variable TYPE_MAP=Assoc_Type[String_Type];
TYPE_MAP["double"]="SLANG_DOUBLE_TYPE";
TYPE_MAP["float"]="SLANG_FLOAT_TYPE";
TYPE_MAP["int"]="SLANG_INT_TYPE";

private define nvcompileandimport(modname){
    variable r;
    variable odir=getcwd();
    variable opath=get_import_module_path();
    ()=chdir("/tmp");
    variable cmd=strjoin([NVCC_BIN,
                      NVCC_FLAGS,
                      "--compiler-options",
                      NVCC_C_FLAGS,
                      NVCC_INC,
                      NVCC_LIB,
                      modname+".cu",
                      "-o "+modname+"-module.so"],
                     " ");
    r=system(cmd+" >/dev/null");% 2>&1");
    if (r!=0){
	fprintf(stderr,"Problems compiling dynamic CUDA code at /tmp/%s.cu",modname);
	()=chdir(odir);
	return 0;
    }
    set_import_module_path(".");
    try {
	import(modname);
	r=1;
    }
    catch ImportError: {
	fprintf(stderr,"Unable to import dynamic CUDA code at /tmp/%s.cu",modname);
	r=0;
    }
    set_import_module_path(opath);
    ()=chdir(odir);
    return r;
}

private define parse_kernel_def(kernel){
    variable name,i;
    variable args=DataType_Type[0];
    variable types=String_Type[0];
    % 
    % compact spaces
    % 
    kernel=strcompress(kernel," \n");
    %
    % first get the name
    % 
    variable tmp=strchop(kernel,'(',0);
    tmp=strchop(tmp[0],' ',0);
    name=tmp[2];
    tmp=strchop(kernel,'(',0);
    tmp=strchop(tmp[1],')',0)[0];
    tmp=strchop(tmp,',',0);
    foreach i (tmp){
	i=strtrim(i);
	if (string_match(i,"*")){
	    args=[args,SLcuda_Type];
	    types=[types,strchop(i,' ',0)[0]];
	}
	else if (string_match(i,"^int")){
	    args=[args,Integer_Type];
	    types=[types,"int"];
	}
	else {
	    args=[args,Double_Type];
	    types=[types,"double"];
	}
    }
    return (name,types,args);
}

% 
% Generate code from a kernel definition
%
define cuda_add_function(name, kernel){
    variable i,n=0;
    variable kname,data_types,input_types;
    (kname,data_types,input_types)=parse_kernel_def(kernel);
    if (input_types[-1]!=Integer_Type)
	verror("The last argument to kernel must be integer type denoting data length");
    else
        input_types=input_types[[:-2]];
    variable modname="slcuda_"+name+string(rand);
    variable fp=fopen("/tmp/"+modname+".cu","w");
    %
    % write header
    % 
    ()=fputs("#include <slang.h>\n#include <cuda.h>\n",fp);
    ()=fputs("#include \"slcuda.h\"\n",fp);
    ()=fputs("SLANG_MODULE("+modname+");\n",fp);
    %
    % write the kernel definition
    % 
    ()=fputs(kernel+"\n",fp);
    % 
    % Now, write the handler function, using the input_types list to
    % parse input args
    % 
    ()=fputs("static void handler (void)\n{\n",fp);
    for (i=length(input_types)-1;i>=0;i--){
	if (input_types[i]==SLcuda_Type){
	    ()=fprintf(fp,"SLcuda_Type *cuda_%d;\n",i);
	    ()=fprintf(fp,"%s *arg_%d;\n",data_types[i],i);
	    ()=fprintf(fp,"if (NULL==(cuda_%d=slcuda_pop_cuda())){return;}",i);
	    ()=fprintf(fp,"if (cuda_%d->type != %s){\n",i,TYPE_MAP[data_types[i]]);
	    ()=fprintf(fp,"SLang_verror(SL_USAGE_ERROR,");
	    ()=fprintf(fp,"\"Wrong type argument %d\");\nreturn; }\n",i);
	    ()=fprintf(fp,"arg_%d=(%s *)cuda_%d->dptr;\n",i,data_types[i],i);
	    n=i;
	}
	else if (input_types[i]==Integer_Type){
	    ()=fprintf(fp,"int arg_%d;\n",i);
	    ()=fprintf(fp,"if(-1==SLang_pop_int(&arg_%d)){ return; }\n",i);
	}
	else if (input_types[i]==Double_Type){
	    ()=fprintf(fp,"double arg_%d_in;\n",i);
	    ()=fprintf(fp,"if(-1==SLang_pop_double(&arg_%d_in)){ return; }\n",i);
	    if (data_types[i]=="float"){
		()=fprintf(fp,"float  arg_%d;\n",i);
		()=fprintf(fp,"arg_%d=(float)arg_%d_in;\n",i,i);
	    }
	    else {
		()=fprintf(fp,"arg_%d=(double)arg_%d_in;\n",i,i);
	    }
	}
    }
    ()=fprintf(fp,"int N=cuda_%d->nelems;\n",n);
    ()=fprintf(fp,"int bx, by, block_size = SLCUDA_BLOCK_SIZE;\n");
    ()=fprintf(fp,"slcuda_compute_dims2d( N, block_size, &bx, &by);\n");
    ()=fprintf(fp,"dim3 n_blocks(bx, by);\n");
    ()=fputs(kname+" <<< n_blocks, block_size >>> (",fp);
    for (i=0;i<length(input_types);i++){
	if (i==0) ()=fprintf(fp,"arg_0");
	else  ()=fprintf(fp,",arg_%d",i);
    }
    ()=fputs(",N",fp);
    ()=fputs(");\n}\n",fp);
    ()=fputs("int init_"+modname+"_module_ns (char *ns_name){\n",fp);
    ()=fputs("if (NULL == SLns_create_namespace (ns_name))\n",fp);
    ()=fputs("return -1;\n",fp);
    ()=fprintf(fp,"SLadd_intrinsic_function(\"%s\",",name);
    ()=fprintf(fp,"(FVOID_STAR)handler,SLANG_VOID_TYPE,0);\n");
    ()=fputs("return 0;\n}",fp);
    ()=fclose(fp);
    if (nvcompileandimport(modname)){
	()=remove("/tmp/"+modname+".cu");
	()=remove("/tmp/"+modname+"-module.so");
    }
}

define cuda_add_function_file(filename){
    variable name;
    if (_NARGS==2){
	name=();
    }
    else {
	name=path_basename_sans_extname(filename);
    }
    variable fp=fopen(filename,"r");
    variable kernel=strjoin(fgetslines(fp),"\n");
    ()=fclose(fp);
    cuda_add_function(name,kernel);
}

define curand_get_default_gen (){
    if (not __is_initialized(&CURAND_DEFAULT_GEN)){
	CURAND_DEFAULT_GEN=curand_new(CURAND_DEFAULT_GEN_TYPE,[_time,getpid]);
    }
    return CURAND_DEFAULT_GEN;
}

define curand (){
    variable arg=();
    variable gen;
    variable cuda;

    if (_NARGS==2){
	gen=();
    }
    else {
	gen=curand_get_default_gen();
    }
    if (typeof(arg)==Integer_Type){
	cuda=cuarr(arg,Integer_Type);
    }
    else {
	cuda=arg;
    }
    curand_gen(CURAND_DEFAULT, gen, cuda);
    if (typeof(arg)==Integer_Type){
	return cuda;
    }
}

define curand_uniform (){
    variable arg=();
    variable gen;
    variable cuda;

    if (_NARGS==2){
	gen=();
    }
    else {
	gen=curand_get_default_gen();
    }
    if (typeof(arg)==Integer_Type){
	cuda=cuarr(arg,CURAND_DEFAULT_TYPE);
    }
    else {
	cuda=arg;
    }
    curand_gen(CURAND_UNIFORM, gen, cuda);
    if (typeof(arg)==Integer_Type){
	return cuda;
    }
}

define curand_normal (){
    variable arg=();
    variable sigma=();
    variable mean=();
    variable gen;
    variable cuda;

    if (_NARGS==4){
	gen=();
    }
    else {
	gen=curand_get_default_gen();
    }
    if (typeof(arg)==Integer_Type){
	cuda=cuarr(arg,CURAND_DEFAULT_TYPE);
    }
    else {
	cuda=arg;
    }
    curand_gen(CURAND_NORMAL, gen, cuda, mean, sigma);
    if (typeof(arg)==Integer_Type){
	return cuda;
    }
}

define curand_lognormal(){
    variable arg=();
    variable sigma=();
    variable mean=();
    variable gen;
    variable cuda;

    if (_NARGS==4){
	gen=();
    }
    else {
	gen=curand_get_default_gen();
    }
    if (typeof(arg)==Integer_Type){
	cuda=cuarr(arg,CURAND_DEFAULT_TYPE);
    }
    else {
	cuda=arg;
    }
    curand_gen(CURAND_LOGNORMAL, gen, cuda, mean, sigma);
    if (typeof(arg)==Integer_Type){
	return cuda;
    }
}