def local_to_gpu(node): """ op(host_from_gpu()) -> host_from_gpu(op) gpu_from_host(op) -> op(gpu_from_host) """ if isinstance(node.op, op): #op(host_from_gpu()) -> host_from_gpu(op) #If any of the input that go on the GPU are on the GPU, #move the op to the gpu. if any(node.inputs[idx].owner and isinstance(node.inputs[idx].owner.op, cuda.HostFromGpu) for idx in to_gpu): new_inp = list(node.inputs) for idx in to_gpu: new_inp[idx] = cuda.gpu_from_host(new_inp[idx]) return [cuda.host_from_gpu(op()(*new_inp))] if node.op == cuda.gpu_from_host: #gpu_from_host(op) -> op(gpu_from_host) host_input = node.inputs[0] if host_input.owner and isinstance(host_input.owner.op, op): op_node = host_input.owner new_inp = list(op_node.inputs) for idx in to_gpu: new_inp[idx] = cuda.gpu_from_host(new_inp[idx]) return [op()(*new_inp)] return False