DirectML Acceleration for TensorFlow and PyTorch in WSL2

DirectML brings high-performance, hardware-accelerated machine learning workloads to TensorFlow and PyTorch inside Windows Subsystem for Linux 2 (WSL2) by bridging Linux-based frameworks with the host Windows DirectX 12 stack. By leveraging Microsoft's GPU virtualization architecture, DirectML enables developers to utilize diverse GPU hardware—including AMD, Intel, Nvidia, and Qualcomm—within a native Linux environment without requiring proprietary vendor-specific Linux drivers.

The Virtualization Bridge: dxgkrnl and /dev/dxg

WSL2 runs a complete Linux kernel inside a lightweight virtual machine. To expose the host system’s GPU to this guest environment, Microsoft introduced a virtualized GPU driver known as dxgkrnl into the WSL2 Linux kernel.

When WSL2 initializes, dxgkrnl exposes a virtualization device node at /dev/dxg. This virtual device establishes a low-latency communication channel between the Linux user space and the Windows Direct3D graphics kernel on the host OS. Rather than passing through the entire physical GPU PCIe device, WSL2 shares the host GPU resources, allowing both the Windows host and Linux guest to share the hardware concurrently.

The Direct3D 12 User-Mode Driver

Above the kernel layer, WSL2 includes a specialized Linux user-mode driver: libd3d12.so. This library provides a functional Direct3D 12 API surface directly inside Linux.

When an application inside WSL2 issues graphics or compute commands, libd3d12.so packages these calls and forwards them across the /dev/dxg boundary to the Windows host D3D12 runtime. The host runtime then hands the commands off to the installed Windows GPU driver, which executes the workload directly on the physical hardware.

DirectML as the Execution Engine

DirectML is a low-level, hardware-agnostic API built directly on top of Direct3D 12. It provides high-performance implementations of mathematical operations standard in neural networks, such as matrix multiplications, 2D/3D convolutions, activations, and pooling layers.

Because DirectML abstracts the hardware behind the DirectX 12 compute model, it compiles these operations into optimized GPU compute shaders at runtime. When running inside WSL2, the DirectML library translates framework-level tensor operations into D3D12 compute dispatches, which stream through libd3d12.so directly to the host GPU.

Framework Integration: PyTorch and TensorFlow

Machine learning frameworks communicate with DirectML using framework-specific plugins and packages:

  • PyTorch: Acceleration is enabled via the torch-directml extension package. It registers a custom backend device (accessible via torch_directml.device()). When tensors and model parameters are moved to this device, PyTorch routes the forward and backward computation graphs through DirectML operator kernels instead of the default CPU or CUDA backends.
  • TensorFlow: Acceleration is delivered via the tensorflow-directml distribution. DirectML acts as a registered device vendor within the TensorFlow runtime, intercepting execution graphs and mapping supported TensorFlow operations directly to DirectML compute kernels.

Through this multi-layered pipeline—traversing the framework plugin, DirectML, libd3d12.so, the /dev/dxg device interface, and the host Windows display driver—developers achieve cross-vendor, near-native GPU acceleration inside their WSL2 development workflows.