diff options
-rw-r--r-- | Tools/CodeGen/CaKernel.m | 39 |
1 files changed, 30 insertions, 9 deletions
diff --git a/Tools/CodeGen/CaKernel.m b/Tools/CodeGen/CaKernel.m index 870989e..ba5bca2 100644 --- a/Tools/CodeGen/CaKernel.m +++ b/Tools/CodeGen/CaKernel.m @@ -202,16 +202,37 @@ DefFn[ If[lookup[calc,ExecuteOn,Automatic] === Host, {calc}, - {mapReplaceAdd[ + + Module[ + {deviceCalc, hostCalc, newCalcs}, + + deviceCalc = + mapReplaceAdd[ + mapReplaceAdd[ + mapReplace[ + calc, + Name,"DEVICE_"<>lookup[calc,Name]], + UseCaKernel, True], + ExecuteOn, Device]; + + hostCalc = mapReplaceAdd[ - AddConditionSuffix[calc, "Accelerator::device_process"], - UseCaKernel, True], - ExecuteOn, Device], - mapReplaceAdd[ - mapReplaceAdd[ - AddConditionSuffix[mapReplace[calc,Name,"HOST__"<>lookup[calc,Name]], "Accelerator::host_process"], - UseCaKernel,False], - ExecuteOn, Host]}]]; + mapReplaceAdd[ + mapReplace[ + calc, + Name,"HOST_"<>lookup[calc,Name]], + UseCaKernel,False], + ExecuteOn, Host]; + + newCalcs = Map[InNewScheduleGroup[lookup[calc,Name],#] &, {hostCalc, deviceCalc}]; + + (* Apply the conditional to the routines, not to the group *) + newCalcs = MapThread[AddConditionSuffix, + {newCalcs, + {"Accelerator::host_process", + "Accelerator::device_process"}}]; + + newCalcs]]]; DefFn[ WithHostCalculations[calcs_List] := |