3 ms·
Honestly for those of us in machine learning, I think something like XLA will likely win over both paradigms. (Disclaimer, I work on XLA.) https://www.tensorf
by jlebar 10y ago
Honestly for those of us in machine learning, I think something like XLA will likely win over both paradigms. (Disclaimer, I work on XLA.)
https://www.tensorflow.org/versions/master/experimental/xla/ https://www.tensorflow.org/versions/master/experimental/xla/
XLA much more closely matches what you want for ML than CUDA/opencl. Which isn't a surprise; it was designed specifically for ML.
Kernel launches are expensive, so any fast CUDA system has to let you compose computations into a single kernel (e.g. multiply by 5 and then take tanh). It's possible to do this in CUDA, but it requires heoric C++ template metaprogramming. It's not uncommon to have files that take ten minutes to compile. Whereas in XLA kernel fusion is nbd, because it's a JIT.
Also, because XLA is generating GPU code after it's seen your model, it can specialize computations specifically to your model. In regular TensorFlow (and I presume other ML frameworks, although I'm not at all familiar with them), you have to compile all of your kernels upfront. This means that the framework probably doesn't have the ideal set of kernels for your model, because the framework's set of kernels needs to be generic. For example, the framework probably isn't going to have a "multiply by 5 and then take tanh" kernel -- if you're lucky, it might have a "multiply by X and then take tanh", but notice that this may be slower because X is now not a constant.
In contrast, not only can XLA specialize for your weird X==5 case, but it can also specialize all of the dimensions of your arrays. This is a really big advantage in many cases.
As just one example, it's common for kernels to do something like
int index = some computation based on threadIdx and blockIdx;
if (index < array_len) { ... }
But in XLA we know the size of the kernel, so we know the possible values for threadIdx and blockIdx, and we know the exact value of array_len. We can therefore often optimize out the if entirely.
- disposablename 10y agoWhy no AMD GPU support?
- jlebar 10y agoWhy no AMD GPU support? I think it just reflects the team's internal priorities. Patches are welcome; we want people to use this system. It wouldn't even be tremendously hard. The XLA IR --> LLVM IR backend is relatively simple, and LLVM already has support for compiling to AMD GPUs. You'd have to split out the nvidia-isms in the generated IR. I think the biggest challenge would just be one of software engineering, namely figuring out a way to specialize the GPU backend for each of the two architectures while allowing it to share code in general.
- doosra 10y agoJustin, XLA sounds interesting. Do you assume you always have CUDA sources for ML operations in XLA? I was under the impression that closed-source libraries like cuDNN were used. Is it possible to accurately evaluate the profitability of fusing two kernels in CUDA (effects of increased register pressure; shared memory)? On the other hand, the generic kernel and its launch parameters were probably hand tuned for performance.
- jlebar 10y ago> Do you assume you always have CUDA sources for ML operations in XLA? I was under the impression that closed-source libraries like cuDNN were used. Yes, XLA calls into cudnn and cublas. It's not a fundamental architectural thing, though; those are just the fastest matmul etc. kernels we currently have access to. > Is it possible to accurately evaluate the profitability of fusing two kernels in CUDA (effects of increased register pressure; shared memory)? For a human, yes, sure, just time both options. The system doesn't currently do this in an automated fashion, though. In a fashion similar to a CPU compiler's inliner, it has heuristics and makes its best guess. In general fusion is very profitable. > On the other hand, the generic kernel and its launch parameters were probably hand tuned for performance. Yes, and this is one of ways that XLA can lose to (say) vanilla TensorFlow today. But it's just a matter of tuning; the system is very young.