aithree 0.1.0

Enables Algorithmic Selection and Customization in Deep Neural Networks

The ai3 (Algorithmic Innovations for Accelerated Implementations of Artificial Intelligence) framework provides easy-to-use fine-grain algorithmic control over an existing DNN. ai3 contains built-in high performance implementations of common deep learning operations and methods by which users can implement their own algorithms in C++. ai3 incurs no additional performance overhead, meaning that performance depends solely on the algorithms chosen by the user.

Documentation Source Code

Installation

Default Implementations: pip install aithree

Custom Implementations:

1. Download the source code

2. Create an implementation with the operations defined in custom

3. If needed, configure the build process with custom.cmake

4. pip install <path to source code>

The framework currently features two methods for algorithmic swapping. swap_backend which swaps every module type of a DNN returning an object completely managed by ai3 and swap_conv2d which swaps convolution operations out of the existing DNN.

swap_conv2d

Swaps, in-place, conv2d operations out of the existing DNN for an implementation of the user specified algorithm. After swapping, the same DNN can still be trained and compiled. If no AlgorithmicSelector is given then the default algorithm decided by the framework are used.

Example:

Swaps the first conv2d operation for an implementation of direct convolution and the second conv2d operation for an implementation of SMM convolution

>>> input_data = torch.randn(10, 3, 224, 224)
>>> orig = ConvNet()
>>> orig_out = orig(input_data)
>>> ai3.swap_conv2d(orig, ['direct', 'smm'])
>>> sc_out = orig(input_data)
>>> torch.allclose(orig_out, sc_out, atol=1e-6)
True

swap_backend

Swaps every module in an exsiting DNN for an implementation of the user specified algorithm returning a Model completly managed by the framework.

Algorithmic selection is performed by passing a mapping from strings containing names of the operations to swap to a AlgorithmicSelector. If no AlgorithmicSelector is passed for a given operation then the default algorithm decided by the framework are used.

Example: Swaps the first conv2d operation for an implementation of direct convolution and the second conv2d operation for an implementation of SMM convolution

>>> def auto_selector(orig: torch.nn.Conv2d, input_shape) -> str:
...     out_channels = orig.weight.shape[0]
...     if (out_channels < 50 and
...         input_shape[1] < 50 and
...         input_shape[2] > 150 and
...         input_shape[3] > 150):
...         return 'direct'
...     return 'smm'
...
>>> input_data = torch.randn(1, 3, 224, 224)
>>> vgg16 = torchvision.models.vgg16(weights=torchvision.models.VGG16_Weights.DEFAULT)
>>> vgg16 = vgg16.eval()
>>> with torch.inference_mode():
...     torch_out = vgg16(input_data)
...     model: ai3.Model = ai3.swap_backend(vgg16, {"conv2d": auto_selector,
...                                                 "maxpool2d": "default"},
...                                         sample_input_shape=(1, 3, 224, 224))
...     sb_out = model(input_data)
...     torch.allclose(torch_out, sb_out, atol=1e-4)
True

Supported Operations, their Algorithms, and Acceleration Platform Compatibility

2D Convolution

The guess algorithm uses the algorithm returned by cudnnGetConvolutionForwardAlgorithm_v7.

Algorithm	direct	smm	gemm	implicit precomp gemm	implicit gemm	winograd	guess	some
none	✓	✓	✗	✗	✗	✗	✗	✓
sycl	✓	✓	✗	✗	✗	✗	✗	✓
cudnn	✗	✗	✓	✓	✓	✓	✓	✓
cublas	✗	✗	✗	✗	✗	✗	✗	✗
mps	✗	✗	✗	✗	✗	✗	✗	✓
metal	✗	✗	✗	✗	✗	✗	✗	✓

Linear

Algorithm	gemm
none	✓
sycl	✗
cudnn	✗
cublas	✓
mps	✗
metal	✗

2D MaxPool

Algorithm	direct
none	✓
sycl	✗
cudnn	✗
cublas	✗
mps	✗
metal	✗

2D AvgPool

Algorithm	direct
none	✓
sycl	✗
cudnn	✗
cublas	✗
mps	✗
metal	✗

2D AdaptiveAvgPool

Algorithm	direct
none	✓
sycl	✗
cudnn	✗
cublas	✗
mps	✗
metal	✗

ReLU

Algorithm	direct
none	✓
sycl	✗
cudnn	✗
cublas	✗
mps	✗
metal	✗

Flatten

Algorithm	direct
none	✓
sycl	✗
cudnn	✗
cublas	✗
mps	✗
metal	✗