[WIP] QNN GPU Procyon recipes

Author: qti-mattsinc

Salesforce-blip-image-captioning-base/QNN/config_gpu_fp32.json

@@ -0,0 +1,54 @@
+{
+    "input_model": {
+        "type": "HfModel",
+        "model_path": "Salesforce/blip-image-captioning-base",
+        "task": "image-text-to-text",
+        "load_kwargs": {
+            "output_attentions": true
+        },
+        "io_config": {
+            "input_names": [
+                "pixel_values",
+                "input_ids",
+                "attention_mask"
+            ],
+            "input_shapes": [
+                [ 1, 3, 384, 384 ],
+                [ 1, 31 ],
+                [ 1, 31 ]
+            ],
+            "input_types": [
+                "float32",
+                "int64",
+                "int64"
+            ],
+            "output_names": [ "logits", "encoder_hidden_states" ]
+        }
+    },
+    "systems": {
+        "target_system": {
+            "type": "LocalSystem",
+            "accelerators": [ { "device": "gpu", "execution_providers": [ "QNNExecutionProvider" ] } ]
+        }
+    },
+    "passes": {
+        "cs": {
+            "type": "CaptureSplitInfo",
+            "block_to_split": [ "vision_model", "text_decoder" ],
+            "num_splits": 2
+        },
+        "conversion": { "device": "cpu", "type": "OnnxConversion", "target_opset": 20 },
+        "add_metadata": { "type": "AddOliveMetadata", "graph_name": "main_graph" },
+        "surgery": {
+            "type": "GraphSurgeries",
+            "surgeries": [ { "surgeon": "MatMulAddToGemm" } ]
+        },
+        "sp": { "type": "SplitModel" }
+    },
+    "host": "target_system",
+    "target": "target_system",
+    "cache_dir": "cache",
+    "output_dir": "output",
+    "evaluate_input_model": false,
+    "clean_cache": true
+}

caidas-swin2SR-realworld-sr-x4-64-bsrgan-psnr/QNN/config_gpu_fp32.json

@@ -0,0 +1,28 @@
+{
+    "input_model": {
+        "type": "HfModel",
+        "model_path": "caidas/swin2SR-realworld-sr-x4-64-bsrgan-psnr",
+        "task": "image-to-image",
+        "io_config": {
+            "input_names": [ "pixel_values" ],
+            "input_shapes": [ [ 1, 3, 232, 232 ] ],
+            "output_names": [ "output" ]
+        }
+    },
+    "systems": {
+        "target_system": {
+            "type": "LocalSystem",
+            "accelerators": [ { "device": "gpu", "execution_providers": [ "QNNExecutionProvider" ] } ]
+        }
+    },
+    "passes": {
+        "conversion": { "device": "cpu", "type": "OnnxConversion", "target_opset": 20 },
+        "add_metadata": { "type": "AddOliveMetadata", "graph_name": "main_graph" }
+    },
+    "host": "target_system",
+    "target": "target_system",
+    "cache_dir": "cache",
+    "output_dir": "output",
+    "evaluate_input_model": false,
+    "clean_cache": true
+}

facebook-convnext-tiny-224/QNN/config_gpu_fp32.json

@@ -0,0 +1,120 @@
+{
+    "input_model": {
+        "type": "HfModel",
+        "model_path": "facebook/convnext-tiny-224",
+        "task": "image-classification",
+        "io_config": {
+            "input_names": [ "pixel_values" ],
+            "input_shapes": [ [ 1, 3, 224, 224 ] ],
+            "output_names": [ "logits" ]
+        }
+    },
+    "systems": {
+        "target_system": {
+            "type": "LocalSystem",
+            "accelerators": [ { "device": "gpu", "execution_providers": [ "QNNExecutionProvider" ] } ]
+        }
+    }, "data_configs": [
+        {
+            "name": "evaluate_data_config",
+            "type": "HuggingfaceContainer",
+            "user_script": "imagenet.py",
+            "load_dataset_config": {
+                "data_name": "imagenet-1k",
+                "split": "validation",
+                "streaming": true,
+                "trust_remote_code": true
+            },
+            "pre_process_data_config": { "type": "dataset_pre_process", "size": 256, "cache_key": "imagenet" },
+            "post_process_data_config": { "type": "imagenet_post_fun" }
+        }
+    ],
+    "evaluators": {
+        "common_evaluator": {
+            "metrics": [
+                {
+                    "name": "accuracy_qnn",
+                    "type": "accuracy",
+                    "data_config": "evaluate_data_config",
+                    "sub_types": [
+                        {
+                            "name": "accuracy_score",
+                            "priority": 1,
+                            "metric_config": { "task": "multiclass", "num_classes": 1000 }
+                        }
+                    ],
+                    "user_config": {
+                        "inference_settings": {
+                            "onnx": {
+                                "session_options": {
+                                    "extra_session_config": { "session.disable_cpu_ep_fallback": "0" }
+                                },
+                                "execution_provider": "QNNExecutionProvider",
+                                "provider_options": [
+                                    {
+                                        "backend_path": "QnnGpu.dll"
+                                    }
+                                ]
+                            }
+                        }
+                    }
+                },
+                {
+                    "name": "accuracy_cpu",
+                    "type": "accuracy",
+                    "data_config": "evaluate_data_config",
+                    "sub_types": [
+                        {
+                            "name": "accuracy_score",
+                            "priority": 2,
+                            "metric_config": { "task": "multiclass", "num_classes": 1000 }
+                        }
+                    ],
+                    "user_config": {
+                        "inference_settings": { "onnx": { "execution_provider": "CPUExecutionProvider" } }
+                    }
+                },
+                {
+                    "name": "latency_qnn",
+                    "type": "latency",
+                    "data_config": "evaluate_data_config",
+                    "sub_types": [ { "name": "avg", "priority": 3 } ],
+                    "user_config": {
+                        "inference_settings": {
+                            "onnx": {
+                                "session_options": {
+                                    "extra_session_config": { "session.disable_cpu_ep_fallback": "0" }
+                                },
+                                "execution_provider": "QNNExecutionProvider",
+                                "provider_options": [
+                                    {
+                                        "backend_path": "QnnGpu.dll"                                    }
+                                ]
+                            }
+                        }
+                    }
+                },
+                {
+                    "name": "latency_cpu",
+                    "type": "latency",
+                    "data_config": "evaluate_data_config",
+                    "sub_types": [ { "name": "avg", "priority": 4 } ],
+                    "user_config": {
+                        "inference_settings": { "onnx": { "execution_provider": "CPUExecutionProvider" } }
+                    }
+                }
+            ]
+        }
+    },
+    "passes": {
+        "conversion": { "device": "cpu", "type": "OnnxConversion", "target_opset": 20 },
+        "add_metadata": { "type": "AddOliveMetadata", "graph_name": "main_graph" }
+    },
+    "host": "target_system",
+    "target": "target_system",
+    "evaluator": "common_evaluator",
+    "cache_dir": "cache",
+    "output_dir": "output",
+    "evaluate_input_model": false,
+    "clean_cache": true
+}

sam2.1-hiera-small/QNN/README.md

@@ -33,8 +33,13 @@ This command will return the path to the Python executable. Set the parent direc
 ### Generate ONNX Model
 Activate the **Quantization Python Environment** and run command to generate encoder and decoder models:
 
+#### NPU
 ```bash
-python generate_model.py
+python generate_model.py --device npu
+```
+#### GPU
+```bash
+python generate_model.py --device gpu
 ```
 
 ### Run the Quantization + Compilation Config

sam2.1-hiera-small/QNN/generate_model.py

@@ -3,6 +3,7 @@
 # Licensed under the MIT License.
 # --------------------------------------------------------------------------
 
+import argparse
 import os
 
 import onnx
@@ -99,7 +100,7 @@ def forward(self, x: torch.Tensor):
 
 
 class SplitHeadSAMEncoderAttention(nn.Module):
-    """SAM Attention block with the following modifications necessary to run on QNN.
+    """SAM Attention block with the following modifications necessary to run on QNN NPU.
 
     * Heads are split into separate ops, rather than all heads running in a single op.
     * QKV is unpacked from 1 tensor into 3 tensors.
@@ -176,6 +177,53 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
 
         return self.proj(x)
 
+class GpuSAMEncoderAttention(nn.Module):
+    """SAM Attention block with the following modifications necessary to run on QNN GPU."""
+
+    def __init__(self, attention_block) -> None:
+        super().__init__()
+        self.qkv = attention_block.qkv
+        self.proj = attention_block.proj
+        self.num_heads = attention_block.num_heads
+        self.q_pool = attention_block.q_pool
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        B, H, W, _ = x.shape
+
+        qkv = self.qkv(x)
+        if self.num_heads == 1:
+            # qkv with shape (B, H * W, 3, C)
+            qkv = qkv.reshape(B, H * W, 3, -1)
+            # q, k, v with shape (B, H * W, C)
+            q, k, v = torch.unbind(qkv, 2)
+        else:
+            # qkv with shape (B, H * W, 3, nHead, C)
+            qkv = qkv.reshape(B, H * W, 3, self.num_heads, -1)
+
+            # q, k, v with shape (B, H * W, nheads, C)
+            q, k, v = torch.unbind(qkv, 2)
+
+            k = k.reshape(B, H * W, self.num_heads, -1).permute(0, 2, 1, 3).reshape(B * self.num_heads, H * W, -1)
+            v = v.reshape(B, H * W, self.num_heads, -1).permute(0, 2, 1, 3).reshape(B * self.num_heads, H * W, -1)
+
+        # Q pooling (for downsample at stage changes)
+        if self.q_pool:
+            q = do_pool(q.reshape(B, H, W, -1), self.q_pool)
+            H, W = q.shape[1:3]  # downsampled shape
+
+        q = q.reshape(B, H * W, self.num_heads, -1).permute(0, 2, 1, 3).reshape(B * self.num_heads, H * W, -1)
+
+        x = F.scaled_dot_product_attention(q, k, v)
+
+        # Transpose back
+        x = x.reshape(B, self.num_heads, H * W, -1)
+        x = x.transpose(1, 2)
+        x = x.reshape(B, H, W, -1)
+
+        x = self.proj(x)
+
+        return x
+
 
 class Conv2DInplaceLinearSAMTransformerMLPBlock(nn.Module):
     """SAM MLPBlock that uses 1x1 Conv2D in place of linear layers."""
@@ -240,11 +288,14 @@ def window_unpartition(x, window_size, pad_hw, hw):
 
 
 class ModMultiScaleBlock(nn.Module):
-    def __init__(self, block):
+    def __init__(self, block, device):
         super().__init__()
         self.model = block
         self.model.mlp = Conv2DInplaceLinearSAMTransformerMLPBlock(self.model.mlp)
-        self.model.attn = SplitHeadSAMEncoderAttention(self.model.attn)
+        if device == "npu":
+            self.model.attn = SplitHeadSAMEncoderAttention(self.model.attn)
+        else:
+            self.model.attn = GpuSAMEncoderAttention(self.model.attn)
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         shortcut = x  # B, h, w, C
@@ -286,6 +337,7 @@ class SAM2Encoder(nn.Module):
     def __init__(
         self,
         sam2,
+        device,
     ) -> None:
         super().__init__()
         self.model = sam2
@@ -295,7 +347,7 @@ def __init__(
             (256, 64, 64),
         ]
         for i, block in enumerate(self.model.image_encoder.trunk.blocks):
-            self.model.image_encoder.trunk.blocks[i] = ModMultiScaleBlock(block)
+            self.model.image_encoder.trunk.blocks[i] = ModMultiScaleBlock(block, device)
 
     def forward(self, input: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
         """Run SAM2 input encoder and returns image_embeddings, high_res_features1, high_res_features2.
@@ -352,7 +404,7 @@ def _embed_masks(
         ) * self.prompt_encoder.no_mask_embed.weight.reshape(1, -1, 1, 1)
         return mask_embedding
 
-    
+
     def _embed_points(
         self,
         points: torch.Tensor,
@@ -364,7 +416,7 @@ def _embed_points(
         padding_label = -torch.ones((labels.shape[0], 1), device=labels.device)
         points = torch.cat([points, padding_point], dim=1)
         labels = torch.cat([labels, padding_label], dim=1)
-        
+
         point_embedding = self.prompt_encoder.pe_layer.forward_with_coords(
             points, self.prompt_encoder.input_image_size
         )
@@ -395,7 +447,7 @@ def _embed_points(
             point_embedding,
         )
         return point_embedding
-    
+
     def forward(
         self,
         image_embeddings: torch.Tensor,  # [1,256,64,64]
@@ -432,7 +484,7 @@ def forward(
 
         sparse_embedding = self._embed_points(point_coords, point_labels)
         dense_embedding = self._embed_masks(mask_input, has_mask_input)
-        
+
         low_res_masks, iou_predictions, _, _ = self.mask_decoder.predict_masks(
             image_embeddings=image_embeddings,
             image_pe=self.prompt_encoder.get_dense_pe(),
@@ -442,7 +494,7 @@ def forward(
             high_res_features=[high_res_features1, high_res_features2],
         )
         masks = F.interpolate(low_res_masks, size=(1024, 1024), mode="bilinear", align_corners=False)
-        return masks, iou_predictions, low_res_masks 
+        return masks, iou_predictions, low_res_masks
 
 
 model_weights_url = "https://dl.fbaipublicfiles.com/segment_anything_2/092824/sam2.1_hiera_small.pt"
@@ -452,16 +504,20 @@ def forward(
 
 
 def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--device", choices=["npu", "gpu"], default="npu")
+    args = parser.parse_args()
+
     download_file(model_weights_url, checkpoint)
     download_file(model_config_url, model_cfg)
 
     GlobalHydra.instance().clear()
     initialize(config_path="./", job_name="sam2_inference", version_base=None)
 
     sam2_model = build_sam2(model_cfg, checkpoint, device="cpu")
-    encoder = SAM2Encoder(sam2_model)
+    encoder = SAM2Encoder(sam2_model, args.device)
     decoder = SAM2Decoder(sam2_model)
-    
+
     en_inputs = {"input": torch.rand((1, 3, 1024, 1024))}
 
     with torch.no_grad():