SHARK

High Performance Machine Learning and Data Analytics for CPUs, GPUs, Accelerators and Heterogeneous Clusters

Communication Channels

• Nod.ai SHARK Discord server: Real time discussions with the nod.ai team and other users
• GitHub issues: Feature requests, bugs etc

Installation

Installation (Linux and macOS)

pip install SHARK and related packages on Linux Python 3.7, 3.8, 3.9, 3.10 and macOS Python 3.10

pip install nodai-shark -f https://github.com/nod-ai/SHARK/releases -f https://github.com/llvm/torch-mlir/releases -f https://github.com/nod-ai/shark-runtime/releases --extra-index-url https://download.pytorch.org/whl/nightly/cpu

Download and run Resnet50 sample

curl -O https://raw.githubusercontent.com/nod-ai/SHARK/main/shark/examples/shark_inference/resnet50_script.py
#Install deps for test script
pip install --pre torch torchvision torchaudio tqdm pillow --extra-index-url https://download.pytorch.org/whl/nightly/cpu
python ./resnet50_script.py --device="cpu"  #use cuda or vulkan or metal 

Download and run BERT (MiniLM) sample

curl -O https://raw.githubusercontent.com/nod-ai/SHARK/main/shark/examples/shark_inference/minilm_jit.py
#Install deps for test script
pip install transformers torch --extra-index-url https://download.pytorch.org/whl/nightly/cpu
python ./minilm_jit.py --device="cpu"  #use cuda or vulkan or metal 

Source Installation

Check out the code

git clone https://github.com/nod-ai/SHARK.git 

Setup your Python VirtualEnvironment and Dependencies

# Setup venv and install necessary packages (torch-mlir, nodLabs/Shark, ...).
./setup_venv.sh
# Please activate the venv after installation.

Run a demo script

python -m  shark.examples.shark_inference.resnet50_script --device="cpu" # Use gpu | vulkan

Run all tests on CPU/GPU/VULKAN/Metal

pytest

# If on Linux for quicker results:
pytest --workers auto

API Reference

### Shark Inference API

from shark_runner import SharkInference

shark_module = SharkInference(
        module = model class.
        (input,)  = inputs to model (must be a torch-tensor)
        dynamic (boolean) = Pass the input shapes as static or dynamic.
        device = `cpu`, `gpu` or `vulkan` is supported.
        tracing_required = (boolean) = Jit trace the module with the given input, useful in the case where jit.script doesn't work. )
shark_module.set_frontend("pytorch") # Use tensorflow, mhlo, linalg, tosa
shark_module.compile()

result = shark_module.forward(inputs)

Example demonstrating running MHLO IR.

from shark.shark_inference import SharkInference
import numpy as np

mhlo_ir = r"""builtin.module  {
      func.func @forward(%arg0: tensor<1x4xf32>, %arg1: tensor<4x1xf32>) -> tensor<4x4xf32> {
        %0 = chlo.broadcast_add %arg0, %arg1 : (tensor<1x4xf32>, tensor<4x1xf32>) -> tensor<4x4xf32>
        %1 = "mhlo.abs"(%0) : (tensor<4x4xf32>) -> tensor<4x4xf32>
        return %1 : tensor<4x4xf32>
      }
}"""

arg0 = np.ones((1, 4)).astype(np.float32)
arg1 = np.ones((4, 1)).astype(np.float32)

shark_module = SharkInference(mhlo_ir, (arg0, arg1))
shark_module.set_frontend("mhlo")
shark_module.compile()
print(shark_module.forward((arg0, arg1)))

Supported and Validated Models

PyTorch Models

Huggingface PyTorch Models

Hugging Face Models	Torch-MLIR lowerable	SHARK-CPU	SHARK-CUDA	SHARK-METAL
BERT	✔️ (JIT)	✔️
Albert	✔️ (JIT)	✔️
BigBird	✔️ (AOT)
DistilBERT	✔️ (JIT)	✔️
GPT2	❌ (AOT)

Torchvision Models

TORCHVISION Models	Torch-MLIR lowerable	SHARK-CPU	SHARK-CUDA	SHARK-METAL
AlexNet	✔️ (Script)	✔️	✔️
DenseNet121	✔️ (Script)
MNasNet1_0	✔️ (Script)
MobileNetV2	✔️ (Script)
MobileNetV3	✔️ (Script)
Unet	❌ (Script)
Resnet18	✔️ (Script)	✔️	✔️
Resnet50	✔️ (Script)	✔️	✔️
Resnet101	✔️ (Script)	✔️	✔️
Resnext50_32x4d	✔️ (Script)
ShuffleNet_v2	❌ (Script)
SqueezeNet	✔️ (Script)	✔️	✔️
EfficientNet	✔️ (Script)
Regnet	✔️ (Script)
Resnest	❌ (Script)
Vision Transformer	✔️ (Script)
VGG 16	✔️ (Script)	✔️	✔️
Wide Resnet	✔️ (Script)	✔️	✔️
RAFT	❌ (JIT)

For more information refer to MODEL TRACKING SHEET

PyTorch Training Models

Models	Torch-MLIR lowerable	SHARK-CPU	SHARK-CUDA	SHARK-METAL
BERT	❌	❌
FullyConnected	✔️	✔️

JAX Models

JAX Models

Models	Torch-MLIR lowerable	SHARK-CPU	SHARK-CUDA	SHARK-METAL
BERT	❌	❌
FullyConnected	✔️	✔️

TFLite Models

TFLite Models

Models	Torch-MLIR lowerable	SHARK-CPU	SHARK-CUDA	SHARK-METAL
BERT	❌	❌
FullyConnected	✔️	✔️

TF Models

Tensorflow Models

Models	Torch-MLIR lowerable	SHARK-CPU	SHARK-CUDA	SHARK-METAL
BERT	❌	❌
FullyConnected	✔️	✔️

Related Projects

Related Project Channels

• Upstream IREE issues: Feature requests, bugs, and other work tracking
• Upstream IREE Discord server: Daily development discussions with the core team and collaborators
• iree-discuss email list: Announcements, general and low-priority discussion
• MLIR topic within LLVM Discourse: IREE is enabled by and heavily relies on MLIR. IREE sometimes is referred to in certain MLIR discussions. Useful if you are also interested in MLIR evolution.

License

nod.ai SHARK is licensed under the terms of the Apache 2.0 License with LLVM Exceptions. See LICENSE for more information.