We prototyped our algorithms in Python with PyTorch and optimized our implementation with custom CUDA kernels. In particular, we seek to strike a balance by executing high-level functions in PyTorch or TensorFlow and handling compute-intensive kernels with CUDA. The code for this paper is available at: https://gitlab.inria.fr/landscapes/fastflow. For ease of use, this includes our standalone unoptimized PyTorch implementation and a TensorFlow port, as well as our CUDA optimizations. The repository also includes examples of how to use our code to edit terrain with Houdini [Sid23], as demonstrated in the accompanying video.
7.1. Implementation
In terms of optimization, we reduce memory allocations and copying and prevent CPU-GPU communication in the control flow. We found that allowing the GPU to run for the maximum number of iterations (e.g., for flow routing) was more efficient than checking for early convergence. The only exception is in depression routing, where we halt the algorithm when no depressions remain.