wehub/apache--tvm
1
0
Fork 0
Code Issues Pull Requests Actions Deployments Models Agent Notes Packages Projects Releases Wiki Activity

chore: import upstream snapshot with attribution
Lint / lint (push) Has been cancelled
Details
CI / MacOS (push) Has been cancelled
Details
CI / Windows (push) Has been cancelled
Details

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 13:36:25 +08:00
commit 26446540fa
3151 changed files with 974126 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/tvm/relax/backend/contrib/__init__.py
+17
View File
@@ -0,0 +1,17 @@
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements. See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership. The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License. You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
"""Relax backends contrib"""
python/tvm/relax/backend/contrib/example_npu/README.md
+243
View File
@@ -0,0 +1,243 @@
<!--- Licensed to the Apache Software Foundation (ASF) under one -->
<!--- or more contributor license agreements. See the NOTICE file -->
<!--- distributed with this work for additional information -->
<!--- regarding copyright ownership. The ASF licenses this file -->
<!--- to you under the Apache License, Version 2.0 (the -->
<!--- "License"); you may not use this file except in compliance -->
<!--- with the License. You may obtain a copy of the License at -->
<!--- http://www.apache.org/licenses/LICENSE-2.0 -->
<!--- Unless required by applicable law or agreed to in writing, -->
<!--- software distributed under the License is distributed on an -->
<!--- "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -->
<!--- KIND, either express or implied. See the License for the -->
<!--- specific language governing permissions and limitations -->
<!--- under the License. -->
# Example NPU Backend
A hands-on example showing how to build a Neural Processing Unit (NPU) backend for TVM's Relax framework using Bring Your Own Codegen (BYOC).
## Context
NPUs are purpose-built accelerators designed around a fixed set of operations common in neural network inference, such as matrix multiplication, convolution, and activation functions. This example shows the architectural patterns you will encounter when building real NPU backends, making it easier to adapt to specific hardware like:
- Mobile NPUs (AMD XDNA, Google Edge TPU, Samsung NPU)
- Dedicated AI chips (Intel Movidius, Qualcomm Hexagon, MediaTek APU)
- Cloud AI accelerators (AWS Inferentia, Google TPU, Microsoft Azure Maia)
- Custom ASIC designs and embedded AI processors
## What This Is
This is an educational template that demonstrates real NPU concepts without requiring actual NPU hardware. It shows developers how to:
- **Pattern-based partitioning**: Identify and group operations that should run on specialized hardware
- **Memory hierarchy management**: Handle different memory tiers (L0/L1/L2/L3) common in NPUs
- **Automatic tiling**: Break large tensors into smaller chunks that fit in on-chip memory
- **Quantization support**: Handle different data precisions efficiently
- **BYOC integration**: Connect custom backends to TVM's compilation pipeline
## Building TVM with Example NPU Support
Add the following flags when configuring TVM with CMake:
```bash
cmake -DUSE_EXAMPLE_NPU_CODEGEN=ON -DUSE_EXAMPLE_NPU_RUNTIME=ON ..
```
Or set them in your `config.cmake`:
```cmake
set(USE_EXAMPLE_NPU_CODEGEN ON)
set(USE_EXAMPLE_NPU_RUNTIME ON)
```
## Quick Start
```python
import tvm
from tvm import relax
from tvm.relax.backend.pattern_registry import get_patterns_with_prefix
from tvm.relax.transform import FuseOpsByPattern, RunCodegen
# Import to register patterns
import tvm.relax.backend.contrib.example_npu
# Get available patterns
patterns = get_patterns_with_prefix("example_npu")
print(f"Available patterns: {[p.name for p in patterns]}")
# Your model gets automatically partitioned
# Operations matching patterns get fused into "Composite" functions
# Those get lowered to the example NPU backend
```
The snippet above shows how to discover registered patterns. A minimal runnable example that demonstrates the BYOC flow (partition -> merge -> codegen) looks like this:
```python
import tvm
from tvm import relax
from tvm.script import relax as R
from tvm.relax.backend.pattern_registry import get_patterns_with_prefix
from tvm.relax.transform import FuseOpsByPattern, MergeCompositeFunctions, RunCodegen
import tvm.relax.backend.contrib.example_npu # registers patterns
@tvm.script.ir_module
class MatmulReLU:
@R.function
def main(
x: R.Tensor((2, 4), "float32"),
w: R.Tensor((4, 8), "float32"),
) -> R.Tensor((2, 8), "float32"):
with R.dataflow():
y = relax.op.matmul(x, w)
z = relax.op.nn.relu(y)
R.output(z)
return z
mod = MatmulReLU
patterns = get_patterns_with_prefix("example_npu")
# Apply partitioning and codegen annotation
mod = FuseOpsByPattern(patterns, bind_constants=False, annotate_codegen=True)(mod)
mod = MergeCompositeFunctions()(mod)
mod = RunCodegen()(mod)
print(mod)
```
A compact visualization of the BYOC flow:
```
Model source (Relax)
│
▼
Pattern-based partition (FuseOpsByPattern)
│
▼
Composite functions (MergeCompositeFunctions)
│
▼
Lower/Codegen for example NPU (RunCodegen / relax.ext.example_npu)
│
▼
Runtime dispatch to NPU runtime (runtime.ExampleNPUJSONRuntimeCreate)
```
## Supported Operations
The backend recognizes these common neural network patterns:
### Core Operations
- `example_npu.dense` - Dense/fully connected layers
- `example_npu.matmul` - Matrix multiplication operations
- `example_npu.conv1d` - 1D convolution for sequence processing
- `example_npu.conv2d` - 2D convolution for image processing
- `example_npu.depthwise_conv2d` - Depthwise separable convolutions
- `example_npu.max_pool2d` - 2D max pooling
- `example_npu.avg_pool2d` - 2D average pooling
- `example_npu.batch_norm` - Batch normalization
- `example_npu.softmax` - Softmax
- `example_npu.add` - Element-wise addition
- `example_npu.multiply` - Element-wise multiplication
- `example_npu.subtract` - Element-wise subtraction
- `example_npu.divide` - Element-wise division
- `example_npu.relu` - ReLU activation
- `example_npu.gelu` - Gaussian Error Linear Unit
- `example_npu.quantize` - Quantization
- `example_npu.dequantize` - Dequantization
### Build-dependent Operations
These patterns are registered only when the corresponding Relax op is present
in the TVM build:
- `example_npu.relu6` - ReLU6 activation (`relax.nn.relu6`)
- `example_npu.sigmoid` - Sigmoid activation (`relax.nn.sigmoid`)
- `example_npu.tanh` - Hyperbolic tangent (`relax.nn.tanh`)
### Fused Patterns
- `example_npu.conv2d_relu_fused` - Optimized Conv2D+ReLU fusion
## Files
### Backend Implementation
- `patterns.py` - Defines which operations get fused together, along with pattern metadata and architectural annotations used by the partitioner. Includes operator availability checking and NPU-specific constraints.
- `__init__.py` - Registers the backend and its BYOC entry points with TVM so the compiler can discover and use the example NPU.
### Runtime Implementation
- `src/runtime/extra/contrib/example_npu/example_npu_runtime.cc` - C++ runtime implementation that handles JSON-based graph execution for the NPU backend.
### Tests and Examples
- `tests/python/contrib/test_example_npu.py` - Comprehensive test suite containing example IRModules (e.g. `MatmulReLU`, `Conv2dReLU`) and demonstrating the complete BYOC flow from pattern registration to runtime execution.
## Status / Build
- The example backend is an educational, CPU-backed emulation. It does not require real NPU hardware.
- Tests are skipped automatically when the example codegen/runtime are not built into TVM. The test checks for the presence of these global functions before running:
```python
import tvm
has_codegen = tvm.get_global_func("relax.ext.example_npu", True)
has_runtime = tvm.get_global_func("runtime.ExampleNPUJSONRuntimeCreate", True)
has_example_npu = has_codegen and has_runtime
```
If `has_example_npu` is False, tests are skipped. This ensures compatibility across different TVM build configurations.
## Testing
Run the tests to see it in action:
```bash
pytest tests/python/contrib/test_example_npu.py -v
```
Tests are skipped if the backend isn't built — see the test file for the exact runtime/codegen checks.
The test suite includes:
- Pattern registration verification (checks that core patterns are available)
- Graph partitioning validation (ensures operations get grouped correctly)
- End-to-end execution testing (verifies runtime integration)
- Build-dependent pattern verification (confirms build-dependent ops register when present)
### Example output
When you run the quick-start snippet or the test, you should see output similar to the following (truncated for brevity):
```
Available patterns: ['example_npu.dense', 'example_npu.matmul', 'example_npu.conv1d', 'example_npu.conv2d', 'example_npu.depthwise_conv2d', 'example_npu.max_pool2d', 'example_npu.avg_pool2d', 'example_npu.batch_norm', 'example_npu.relu', 'example_npu.add', 'example_npu.multiply', 'example_npu.conv2d_relu_fused']
Relax IRModule
def @main(...) -> ...
%0 = call_extern("relax.ext.example_npu", ...)
# composite functions
def @composite_0(...) /* Composite */ = ...
```
This shows the registered patterns and that matched subgraphs were turned into composite functions and lowered to the example NPU codegen/runtime.
## Key Features Demonstrated
### NPU Architectural Concepts
- **Multi-tier memory hierarchy**: SRAM (256KB), CMX (512KB), and DRAM management
- **Tiling constraints**: 32x32 tiles with 16-element vectors for optimal NPU utilization
- **Quantization support**: INT8/INT16 for inference acceleration, mixed precision handling
- **Specialized execution units**: Matrix engines (16x16), vector units (64-wide), pooling units
- **Power management**: Support for different power modes (high_performance, balanced, low_power)
### Pattern Matching Features
- **Memory constraint hooks**: Placeholder checks where a real backend would reject tensors that exceed on-chip memory; the example accepts all
- **Fusion opportunities**: Identifies conv+activation and other beneficial fusions
- **Layout preferences**: NHWC channel-last layouts preferred by NPUs
### Error Handling
- **Robust exception handling**: Catches specific exception types instead of generic exceptions
- **Comprehensive testing**: Validates both successful cases and error conditions
## Learn More
This backend serves as both a working example and educational resource for understanding NPU integration patterns. The implementation demonstrates vendor-neutral concepts that apply across different NPU architectures, making it a valuable starting point for real NPU backend development.
python/tvm/relax/backend/contrib/example_npu/__init__.py
+31
View File
@@ -0,0 +1,31 @@
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements. See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership. The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License. You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
"""
Example NPU Backend for BYOC Integration
This module provides an educational example of how to implement
a custom NPU backend in TVM using the Bring Your Own Codegen (BYOC)
framework. It demonstrates key NPU architectural concepts including
memory hierarchy, tiling, quantization, and operation fusion.
The patterns module registers all supported NPU operations and their
constraints, making them available for graph partitioning.
"""
from . import patterns
__all__ = ["patterns"]
python/tvm/relax/backend/contrib/example_npu/patterns.py
+543
View File
@@ -0,0 +1,543 @@
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements. See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership. The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License. You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
"""
Example NPU Pattern Table with Architectural Concepts
This module demonstrates NPU-specific architectural patterns that are common
across different NPU vendors, including memory hierarchy, quantization,
tiling, and fusion strategies.
"""
from typing import ClassVar
from tvm.ir import Op
from tvm.relax.dpl.pattern import is_op, wildcard
from tvm.relax.transform import PatternCheckContext
from ...pattern_registry import register_patterns
# NPU-specific configuration constants (vendor-neutral)
class NPUConfig:
"""NPU architectural parameters common across vendors"""
# Memory hierarchy sizes (in KB) - typical NPU values
SRAM_SIZE_KB = 256 # On-chip SRAM/scratchpad
CMX_SIZE_KB = 512 # Compute memory (near compute units)
# Tiling constraints
TILE_HEIGHT = 32
TILE_WIDTH = 32
VECTOR_SIZE = 16
# Supported data types for NPU acceleration
SUPPORTED_DTYPES: ClassVar[list[str]] = ["int8", "int16", "float16", "float32"]
QUANTIZED_DTYPES: ClassVar[list[str]] = ["int8", "int16"]
# NPU execution units
MATRIX_ENGINE_SIZE = 16 # MxN matrix engine
VECTOR_ENGINE_WIDTH = 64 # Vector processing width
# Power modes
POWER_MODES: ClassVar[list[str]] = ["high_performance", "balanced", "low_power"]
def _check_npu_memory_constraints(
context: PatternCheckContext, # pylint: disable=unused-argument
) -> bool:
"""
Placeholder for NPU memory hierarchy constraint checking.
A real implementation would inspect the annotated expression's
TensorType to verify the tensor fits within the NPU's
on-chip SRAM (L1) or compute memory (L2/CMX). Tensors that
exceed on-chip capacity require tiling before offload.
"""
return True
def _check_npu_quantization(
context: PatternCheckContext, # pylint: disable=unused-argument
) -> bool:
"""
Placeholder for NPU quantization requirement checking.
A real implementation would verify the op's dtype falls within
the set supported by the NPU (e.g. int8, int16, float16, float32)
and reject ops with unsupported dtypes so they fall back to CPU.
"""
return True
def conv2d_relu_fused_pattern():
"""
NPU-optimized Conv2D+ReLU fusion pattern.
This is a key NPU optimization - fusing convolution with activation
avoids memory traffic between operations.
"""
def _make_conv2d_relu_pattern():
input_tensor = wildcard()
weight = wildcard()
conv = is_op("relax.nn.conv2d")(input_tensor, weight)
relu = is_op("relax.nn.relu")(conv)
annotations = {
"input": input_tensor,
"weight": weight,
"conv": conv,
"root": relu,
}
return relu, annotations
def _check_conv2d_relu(context: PatternCheckContext) -> bool:
"""Check if Conv2D+ReLU fusion is beneficial for NPU"""
if not _check_npu_memory_constraints(context):
return False
if not _check_npu_quantization(context):
return False
return True
return ("example_npu.conv2d_relu_fused", *_make_conv2d_relu_pattern(), _check_conv2d_relu)
def matmul_relu_fused_pattern():
"""
NPU-optimized MatMul+ReLU fusion pattern.
Fusing the matrix engine output with the activation unit avoids a
write/read round-trip through L1 SRAM, mirroring the conv2d+relu
fusion below.
"""
def _make_matmul_relu_pattern():
input_tensor = wildcard()
weight = wildcard()
matmul = is_op("relax.matmul")(input_tensor, weight)
relu = is_op("relax.nn.relu")(matmul)
annotations = {
"input": input_tensor,
"weight": weight,
"matmul": matmul,
"root": relu,
}
return relu, annotations
def _check_matmul_relu(context: PatternCheckContext) -> bool:
"""Check if MatMul+ReLU fusion is beneficial for NPU"""
if not _check_npu_memory_constraints(context):
return False
if not _check_npu_quantization(context):
return False
return True
return ("example_npu.matmul_relu_fused", *_make_matmul_relu_pattern(), _check_matmul_relu)
def matmul_patterns():
"""
NPU-optimized matrix multiplication patterns.
NPUs typically have dedicated matrix engines (systolic arrays,
tensor cores) that require specific layouts and sizes.
"""
def _make_matmul_pattern():
input_tensor = wildcard()
weight = wildcard()
output = is_op("relax.matmul")(input_tensor, weight)
annotations = {
"input": input_tensor,
"weight": weight,
"root": output,
}
return output, annotations
def _check_matmul(context: PatternCheckContext) -> bool:
"""Check if matmul can use NPU matrix engine"""
return _check_npu_memory_constraints(context) and _check_npu_quantization(context)
def _matmul_pattern(pattern_name):
return (pattern_name, *_make_matmul_pattern(), _check_matmul)
# Register both common names used for matrix multiplication in patterns/tests
return [
_matmul_pattern("example_npu.dense"),
_matmul_pattern("example_npu.matmul"),
]
def conv1d_patterns():
"""
1D Convolution patterns optimized for NPU execution.
NPUs handle 1D convolution by mapping to 2D operations
or using specialized 1D processing units.
"""
def _make_conv1d_pattern():
input_tensor = wildcard()
weight = wildcard()
output = is_op("relax.nn.conv1d")(input_tensor, weight)
annotations = {
"input": input_tensor,
"weight": weight,
"root": output,
}
return output, annotations
def _check_conv1d(context: PatternCheckContext) -> bool:
"""Check if conv1d can use NPU vector engine"""
return _check_npu_memory_constraints(context) and _check_npu_quantization(context)
def _conv1d_pattern(pattern_name):
return (pattern_name, *_make_conv1d_pattern(), _check_conv1d)
return [_conv1d_pattern("example_npu.conv1d")]
def conv2d_patterns():
"""
2D Convolution patterns with NPU tiling and memory management.
2D convolution is the most important NPU operation, with
dedicated hardware for efficient processing.
"""
def _make_conv2d_pattern():
input_tensor = wildcard()
weight = wildcard()
output = is_op("relax.nn.conv2d")(input_tensor, weight)
annotations = {
"input": input_tensor,
"weight": weight,
"root": output,
}
return output, annotations
def _check_conv2d(context: PatternCheckContext) -> bool:
"""Check conv2d NPU constraints"""
return _check_npu_memory_constraints(context) and _check_npu_quantization(context)
def _conv2d_pattern(pattern_name):
return (pattern_name, *_make_conv2d_pattern(), _check_conv2d)
return [_conv2d_pattern("example_npu.conv2d")]
def depthwise_conv2d_patterns():
"""
Depthwise convolution - critical for mobile NPUs.
Many NPUs have specialized units for depthwise operations
used in MobileNet-style architectures.
"""
def _make_depthwise_pattern():
input_tensor = wildcard()
weight = wildcard()
output = is_op("relax.nn.conv2d")(input_tensor, weight)
annotations = {
"input": input_tensor,
"weight": weight,
"root": output,
}
return output, annotations
def _check_depthwise(context: PatternCheckContext) -> bool:
"""Check if this is a depthwise conv that NPU can accelerate"""
conv_call = context.annotated_expr["root"]
# groups > 1 distinguishes depthwise/grouped conv from standard conv2d.
# True depthwise has groups == in_channels; we accept any grouped variant
# here since the NPU's depthwise unit handles all grouped convolutions.
if conv_call.attrs.groups <= 1:
return False
return _check_npu_memory_constraints(context) and _check_npu_quantization(context)
return [("example_npu.depthwise_conv2d", *_make_depthwise_pattern(), _check_depthwise)]
def pooling_patterns():
"""
Pooling operations with NPU memory streaming.
NPUs often process pooling with the convolution engine
or dedicated pooling units.
"""
def _make_maxpool2d_pattern():
input_tensor = wildcard()
output = is_op("relax.nn.max_pool2d")(input_tensor)
annotations = {
"input": input_tensor,
"root": output,
}
return output, annotations
def _make_avgpool2d_pattern():
input_tensor = wildcard()
output = is_op("relax.nn.avg_pool2d")(input_tensor)
annotations = {
"input": input_tensor,
"root": output,
}
return output, annotations
def _check_pooling(context: PatternCheckContext) -> bool:
"""Check pooling NPU constraints"""
return _check_npu_memory_constraints(context)
return [
("example_npu.max_pool2d", *_make_maxpool2d_pattern(), _check_pooling),
("example_npu.avg_pool2d", *_make_avgpool2d_pattern(), _check_pooling),
]
def batch_norm_patterns():
"""
Batch normalization - often fused with conv on NPUs.
NPUs typically fuse BN into convolution to avoid
separate memory passes.
"""
def _make_batch_norm_pattern():
input_tensor = wildcard()
gamma = wildcard()
beta = wildcard()
moving_mean = wildcard()
moving_var = wildcard()
output = is_op("relax.nn.batch_norm")(input_tensor, gamma, beta, moving_mean, moving_var)
annotations = {
"input": input_tensor,
"root": output,
}
return output, annotations
def _check_batch_norm(context: PatternCheckContext) -> bool:
"""Check if batch norm should be offloaded or fused"""
return _check_npu_quantization(context)
return [("example_npu.batch_norm", *_make_batch_norm_pattern(), _check_batch_norm)]
def softmax_patterns():
"""
Softmax - used in classification heads and attention mechanisms.
NPUs typically implement softmax via dedicated hardware or
a combination of exp, sum, and divide operations.
"""
def _make_softmax_pattern():
input_tensor = wildcard()
output = is_op("relax.nn.softmax")(input_tensor)
annotations = {
"input": input_tensor,
"root": output,
}
return output, annotations
def _check_softmax(context: PatternCheckContext) -> bool:
"""Check if softmax can use NPU activation unit"""
return _check_npu_memory_constraints(context) and _check_npu_quantization(context)
patterns = []
try:
Op.get("relax.nn.softmax")
patterns.append(("example_npu.softmax", *_make_softmax_pattern(), _check_softmax))
except (KeyError, AttributeError):
pass
return patterns
def activation_patterns():
"""
NPU activation functions with specialized hardware.
NPUs have dedicated activation units that can handle
various functions efficiently.
"""
def _make_activation_pattern(op_name: str):
def _pattern():
input_tensor = wildcard()
output = is_op(op_name)(input_tensor)
annotations = {
"input": input_tensor,
"root": output,
}
return output, annotations
return _pattern
def _check_activation(context: PatternCheckContext) -> bool:
"""Check if activation can use NPU activation unit"""
return _check_npu_quantization(context)
activations = [
("example_npu.relu", "relax.nn.relu"),
("example_npu.relu6", "relax.nn.relu6"),
("example_npu.sigmoid", "relax.nn.sigmoid"),
("example_npu.tanh", "relax.nn.tanh"),
("example_npu.gelu", "relax.nn.gelu"),
]
patterns = []
for pattern_name, op_name in activations:
try:
Op.get(op_name)
except (KeyError, AttributeError):
continue
pattern_fn = _make_activation_pattern(op_name)
patterns.append((pattern_name, *pattern_fn(), _check_activation))
return patterns
def elementwise_patterns():
"""
Element-wise operations that NPUs can vectorize.
NPUs process element-wise ops using vector units
with SIMD capabilities.
"""
def _make_elementwise_pattern(op_name: str):
def _pattern():
input1 = wildcard()
input2 = wildcard()
output = is_op(op_name)(input1, input2)
annotations = {
"input1": input1,
"input2": input2,
"root": output,
}
return output, annotations
return _pattern
def _check_elementwise(context: PatternCheckContext) -> bool:
"""Check if elementwise op can use NPU vector unit"""
return _check_npu_memory_constraints(context) and _check_npu_quantization(context)
ops = ["relax.add", "relax.multiply", "relax.subtract", "relax.divide"]
patterns = []
for op in ops:
try:
Op.get(op)
except (KeyError, AttributeError):
continue
op_short = op.split(".")[-1]
pattern_fn = _make_elementwise_pattern(op)
patterns.append((f"example_npu.{op_short}", *pattern_fn(), _check_elementwise))
return patterns
def quantization_patterns():
"""
Quantization/dequantization patterns for NPU.
NPUs need explicit quantization boundaries to switch
between precision levels.
"""
def _make_quantize_pattern():
input_tensor = wildcard()
output = is_op("relax.quantize")(input_tensor)
annotations = {
"input": input_tensor,
"root": output,
}
return output, annotations
def _make_dequantize_pattern():
input_tensor = wildcard()
output = is_op("relax.dequantize")(input_tensor)
annotations = {
"input": input_tensor,
"root": output,
}
return output, annotations
def _check_quantization(
context: PatternCheckContext, # pylint: disable=unused-argument
) -> bool:
"""Check quantization operations"""
return True
patterns = []
try:
Op.get("relax.quantize")
patterns.append(("example_npu.quantize", *_make_quantize_pattern(), _check_quantization))
except (KeyError, AttributeError):
pass
try:
Op.get("relax.dequantize")
patterns.append(
("example_npu.dequantize", *_make_dequantize_pattern(), _check_quantization)
)
except (KeyError, AttributeError):
pass
return patterns
# Register all NPU patterns with architectural awareness
# register_patterns priority: patterns that appear LATER in the list win.
# So we place general / standalone patterns first, and fused (more
# specific) patterns last so they take precedence over their constituents.
register_patterns(
[
*quantization_patterns(),
*elementwise_patterns(),
*activation_patterns(),
*softmax_patterns(),
*batch_norm_patterns(),
*pooling_patterns(),
*matmul_patterns(),
*conv1d_patterns(),
# Plain conv2d is more general than depthwise (groups>1); list
# plain first so depthwise wins on grouped convs.
*conv2d_patterns(),
*depthwise_conv2d_patterns(),
# Fused patterns last (highest priority).
matmul_relu_fused_pattern(),
conv2d_relu_fused_pattern(),
]
)
python/tvm/relax/backend/contrib/nnapi.py
+321
View File
@@ -0,0 +1,321 @@
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements. See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership. The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License. You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
"""Pattern table for NNAPI backend"""
from collections.abc import Mapping
from tvm.ir import IRModule
from tvm.relax.dpl.pattern import (
DFPattern,
is_op,
wildcard,
)
from tvm.relax.transform import FuseOpsByPattern, MergeCompositeFunctions
from ..pattern_registry import get_patterns_with_prefix, register_patterns
def elementwise_binary_patterns() -> list[tuple[str, DFPattern, Mapping[str, DFPattern]]]:
"""
Returns a list of tuples representing elementwise binary operation patterns mapped
between NNAPI and Relax frameworks.
"""
def _elementwise_binary_pattern(
pattern_name: str,
op_name: str,
) -> tuple[str, DFPattern, Mapping[str, DFPattern]]:
input0 = wildcard()
input1 = wildcard()
pattern = is_op(op_name)(input0, input1)
return (pattern_name, pattern, {})
return [
_elementwise_binary_pattern("nnapi.add", "relax.add"),
_elementwise_binary_pattern("nnapi.mul", "relax.multiply"),
_elementwise_binary_pattern("nnapi.div", "relax.divide"),
_elementwise_binary_pattern("nnapi.sub", "relax.subtract"),
_elementwise_binary_pattern("nnapi.pow", "relax.power"),
_elementwise_binary_pattern("nnapi.equal", "relax.equal"),
_elementwise_binary_pattern("nnapi.greater", "relax.greater"),
_elementwise_binary_pattern("nnapi.greater_equal", "relax.greater_equal"),
_elementwise_binary_pattern("nnapi.less", "relax.less"),
_elementwise_binary_pattern("nnapi.less_equal", "relax.less_equal"),
_elementwise_binary_pattern("nnapi.not_equal", "relax.not_equal"),
_elementwise_binary_pattern("nnapi.maximum", "relax.maximum"),
_elementwise_binary_pattern("nnapi.minimum", "relax.minimum"),
]
def unary_patterns() -> list[tuple[str, DFPattern, Mapping[str, DFPattern]]]:
"""
Returns a list of tuples representing unary operation patterns mapped
between NNAPI and Relax frameworks.
"""
def _unary_pattern(
pattern_name: str, op_name: str
) -> tuple[str, DFPattern, Mapping[str, DFPattern]]:
input0 = wildcard()
pattern = is_op(op_name)(input0)
return (pattern_name, pattern, {})
return [
_unary_pattern("nnapi.floor", "relax.floor"),
_unary_pattern("nnapi.relu", "relax.nn.relu"),
_unary_pattern("nnapi.logistic", "relax.sigmoid"),
_unary_pattern("nnapi.softmax", "relax.nn.softmax"),
_unary_pattern("nnapi.tanh", "relax.tanh"),
_unary_pattern("nnapi.abs", "relax.abs"),
_unary_pattern("nnapi.exp", "relax.exp"),
_unary_pattern("nnapi.log", "relax.log"),
_unary_pattern("nnapi.neg", "relax.negative"),
_unary_pattern("nnapi.cast", "relax.astype"),
_unary_pattern("nnapi.sqrt", "relax.sqrt"),
_unary_pattern("nnapi.rsqrt", "relax.rsqrt"),
]
def matmul_pattern() -> tuple[str, DFPattern, Mapping[str, DFPattern]]:
"""
Returns a tuple representing matmul operation patterns mapped
between NNAPI and Relax frameworks.
"""
input0 = wildcard()
input1 = wildcard()
pattern = is_op("relax.matmul")(input0, input1)
return ("nnapi.batch_matmul", pattern, {})
def permute_dims_pattern() -> tuple[str, DFPattern, Mapping[str, DFPattern]]:
"""
Returns a tuple representing permute operation patterns mapped
between NNAPI and Relax frameworks.
"""
input0 = wildcard()
pattern = is_op("relax.permute_dims")(input0)
return ("nnapi.transpose", pattern, {})
def astype_pattern() -> tuple[str, DFPattern, Mapping[str, DFPattern]]:
"""
Returns a tuple representing astype operation patterns mapped
between NNAPI and Relax frameworks.
"""
input0 = wildcard().has_dtype("float16") | wildcard().has_dtype("float32")
pattern = is_op("relax.astype")(input0).has_dtype("float16") | is_op("relax.astype")(
input0
).has_dtype("float32")
return ("nnapi.cast", pattern, {})
def mean_pattern() -> tuple[str, DFPattern, Mapping[str, DFPattern]]:
"""
Returns a tuple representing mean operation patterns mapped
between NNAPI and Relax frameworks.
"""
input0 = wildcard()
pattern = is_op("relax.mean")(input0)
return ("nnapi.mean", pattern, {})
def conv2d_pattern() -> tuple[str, DFPattern, Mapping[str, DFPattern]]:
"""
Returns a tuple representing conv2d operation patterns mapped
between NNAPI and Relax frameworks.
"""
input0 = wildcard()
input1 = wildcard()
input2 = wildcard()
conv = is_op("relax.nn.conv2d")(input0, input1)
pattern = is_op("relax.add")(conv, input2)
return ("nnapi.conv2d", pattern, {})
def max_pool2d_pattern() -> tuple[str, DFPattern, Mapping[str, DFPattern]]:
"""
Returns a tuple representing max_pool2d operation patterns mapped
between NNAPI and Relax frameworks.
"""
input0 = wildcard()
pattern = is_op("relax.nn.max_pool2d")(input0)
return ("nnapi.max_pool_2d", pattern, {})
register_patterns(
[
*elementwise_binary_patterns(),
*unary_patterns(),
matmul_pattern(),
permute_dims_pattern(),
astype_pattern(),
mean_pattern(),
conv2d_pattern(),
max_pool2d_pattern(),
]
)
def min_feature_level(pattern_name: str) -> int:
"""
Returns the minimum feature level required to support a given NNAPI operation pattern.
Args:
pattern_name (str): The name of the NNAPI operation pattern
(e.g., "nnapi.add", "nnapi.conv2d").
Returns:
int: The minimum feature level for the specified pattern, or 1 if the pattern is not found.
"""
levels = {
"nnapi.add": 1,
"nnapi.average_pool_2d": 1,
"nnapi.concatenation": 1,
"nnapi.conv2d": 1,
"nnapi.depthwise_conv_2d": 1,
"nnapi.depth_to_space": 1,
"nnapi.dequantize": 1,
"nnapi.embedding_lookup": 1,
"nnapi.floor": 1,
"nnapi.fully_connected": 1,
"nnapi.hashtable_lookup": 1,
"nnapi.l2_normalization": 1,
"nnapi.l2_pool_2d": 1,
"nnapi.local_response_normalization": 1,
"nnapi.logistic": 1,
"nnapi.lsh_projection": 1,
"nnapi.lstm": 1,
"nnapi.max_pool_2d": 1,
"nnapi.mul": 1,
"nnapi.relu": 1,
"nnapi.relu1": 1,
"nnapi.relu6": 1,
"nnapi.reshape": 1,
"nnapi.resize_bilinear": 1,
"nnapi.rnn": 1,
"nnapi.softmax": 1,
"nnapi.space_to_depth": 1,
"nnapi.svdf": 1,
"nnapi.tanh": 1,
"nnapi.batch_to_space_nd": 2,
"nnapi.div": 2,
"nnapi.mean": 2,
"nnapi.pad": 2,
"nnapi.space_to_batch_nd": 2,
"nnapi.squeeze": 2,
"nnapi.strided_slice": 2,
"nnapi.sub": 2,
"nnapi.transpose": 2,
"nnapi.abs": 3,
"nnapi.argmax": 3,
"nnapi.argmin": 3,
"nnapi.axis_aligned_bbox_transform": 3,
"nnapi.bidirectional_sequence_lstm": 3,
"nnapi.bidirectional_sequence_rnn": 3,
"nnapi.box_with_nms_limit": 3,
"nnapi.cast": 3,
"nnapi.channel_shuffle": 3,
"nnapi.detection_postprocessing": 3,
"nnapi.equal": 3,
"nnapi.exp": 3,
"nnapi.expand_dims": 3,
"nnapi.gather": 3,
"nnapi.generate_proposals": 3,
"nnapi.greater": 3,
"nnapi.greater_equal": 3,
"nnapi.grouped_conv_2d": 3,
"nnapi.heatmap_max_keypoint": 3,
"nnapi.instance_normalization": 3,
"nnapi.less": 3,
"nnapi.less_equal": 3,
"nnapi.log": 3,
"nnapi.logical_and": 3,
"nnapi.logical_not": 3,
"nnapi.logical_or": 3,
"nnapi.log_softmax": 3,
"nnapi.maximum": 3,
"nnapi.minimum": 3,
"nnapi.neg": 3,
"nnapi.not_equal": 3,
"nnapi.pad_v2": 3,
"nnapi.pow": 3,
"nnapi.prelu": 3,
"nnapi.quantize": 3,
"nnapi.quantized_16bit_lstm": 3,
"nnapi.random_multinomial": 3,
"nnapi.reduce_all": 3,
"nnapi.reduce_any": 3,
"nnapi.reduce_max": 3,
"nnapi.reduce_min": 3,
"nnapi.reduce_prod": 3,
"nnapi.reduce_sum": 3,
"nnapi.roi_align": 3,
"nnapi.roi_pooling": 3,
"nnapi.rsqrt": 3,
"nnapi.select": 3,
"nnapi.sin": 3,
"nnapi.slice": 3,
"nnapi.split": 3,
"nnapi.sqrt": 3,
"nnapi.tile": 3,
"nnapi.topk_v2": 3,
"nnapi.transpose_conv_2d": 3,
"nnapi.unidirectional_sequence_lstm": 3,
"nnapi.unidirectional_sequence_rnn": 3,
"nnapi.resize_nearest_neighbor": 3,
"nnapi.quantized_lstm": 4,
"nnapi.if": 4,
"nnapi.while": 4,
"nnapi.elu": 4,
"nnapi.hard_swish": 4,
"nnapi.fill": 4,
"nnapi.rank": 4,
"nnapi.batch_matmul": 6,
"nnapi.pack": 6,
"nnapi.mirror_pad": 7,
"nnapi.reverse": 7,
}
return levels[pattern_name]
def partition_for_nnapi(mod: IRModule, feature_level: int | None = None) -> IRModule:
"""Partition the graph greedily offloading supported operators to NNAPI.
Parameters
----------
mod : tvm.ir.IRModule
The module to run passes on.
feature_level : Optional[int]
The maximum NNAPI feature level.
Returns
-------
mod : tvm.ir.IRModule
Annotated and partitioned module.
"""
patterns = get_patterns_with_prefix("nnapi")
if feature_level is not None:
patterns = [pat for pat in patterns if feature_level >= min_feature_level(pat.name)]
mod = FuseOpsByPattern(patterns, bind_constants=False, annotate_codegen=False)(mod)
mod = MergeCompositeFunctions()(mod)
return mod
python/tvm/relax/backend/contrib/tensorrt.py
+140
View File
@@ -0,0 +1,140 @@
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements. See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership. The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License. You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
"""Pattern table and partitioning for the TensorRT BYOC backend.
The composite name of each pattern is "tensorrt.<op>", matching the runtime
converter registered under the same name (the converters are keyed by
"tensorrt." + op_name). ``partition_for_tensorrt`` carves the matched subgraphs
out of the module and annotates them for the ``tensorrt`` codegen.
"""
from collections.abc import Mapping
from tvm.ir import IRModule
from tvm.relax.dpl.pattern import DFPattern, is_op, wildcard
from tvm.relax.transform import FuseOpsByPattern, MergeCompositeFunctions
from ..pattern_registry import get_patterns_with_prefix, register_patterns
Pattern = tuple[str, DFPattern, Mapping[str, DFPattern]]
def _op_pattern(composite_name: str, op_name: str, num_args: int) -> Pattern:
"""A pattern matching a single op called with ``num_args`` wildcard arguments."""
args = [wildcard() for _ in range(num_args)]
return (composite_name, is_op(op_name)(*args), {})
def _tensorrt_patterns() -> list[Pattern]:
patterns: list[Pattern] = []
# Activations and unary elementwise ops (single tensor argument).
for composite, op in [
("tensorrt.nn.relu", "relax.nn.relu"),
("tensorrt.sigmoid", "relax.sigmoid"),
("tensorrt.tanh", "relax.tanh"),
("tensorrt.exp", "relax.exp"),
("tensorrt.log", "relax.log"),
("tensorrt.sqrt", "relax.sqrt"),
("tensorrt.abs", "relax.abs"),
("tensorrt.negative", "relax.negative"),
("tensorrt.sin", "relax.sin"),
("tensorrt.cos", "relax.cos"),
("tensorrt.atan", "relax.atan"),
("tensorrt.ceil", "relax.ceil"),
("tensorrt.floor", "relax.floor"),
("tensorrt.erf", "relax.erf"),
("tensorrt.nn.softmax", "relax.nn.softmax"),
("tensorrt.nn.batch_flatten", "relax.nn.batch_flatten"),
("tensorrt.expand_dims", "relax.expand_dims"),
("tensorrt.squeeze", "relax.squeeze"),
("tensorrt.transpose", "relax.permute_dims"),
("tensorrt.layout_transform", "relax.layout_transform"),
("tensorrt.nn.max_pool2d", "relax.nn.max_pool2d"),
("tensorrt.nn.avg_pool2d", "relax.nn.avg_pool2d"),
("tensorrt.nn.max_pool3d", "relax.nn.max_pool3d"),
("tensorrt.nn.avg_pool3d", "relax.nn.avg_pool3d"),
("tensorrt.nn.adaptive_avg_pool2d", "relax.nn.adaptive_avg_pool2d"),
("tensorrt.sum", "relax.sum"),
("tensorrt.prod", "relax.prod"),
("tensorrt.max", "relax.max"),
("tensorrt.min", "relax.min"),
("tensorrt.mean", "relax.mean"),
("tensorrt.concatenate", "relax.concat"),
("tensorrt.split", "relax.split"),
]:
patterns.append(_op_pattern(composite, op, 1))
# Binary elementwise ops (two tensor arguments).
for composite, op in [
("tensorrt.add", "relax.add"),
("tensorrt.subtract", "relax.subtract"),
("tensorrt.multiply", "relax.multiply"),
("tensorrt.divide", "relax.divide"),
("tensorrt.power", "relax.power"),
("tensorrt.maximum", "relax.maximum"),
("tensorrt.minimum", "relax.minimum"),
]:
patterns.append(_op_pattern(composite, op, 2))
# Convolutions and matmul (data + weight).
for composite, op in [
("tensorrt.nn.conv1d", "relax.nn.conv1d"),
("tensorrt.nn.conv2d", "relax.nn.conv2d"),
("tensorrt.nn.conv3d", "relax.nn.conv3d"),
("tensorrt.nn.conv2d_transpose", "relax.nn.conv2d_transpose"),
("tensorrt.nn.conv3d_transpose", "relax.nn.conv3d_transpose"),
("tensorrt.nn.batch_matmul", "relax.matmul"),
("tensorrt.reshape", "relax.reshape"),
]:
patterns.append(_op_pattern(composite, op, 2))
# layer_norm (data, gamma, beta) and clip (data, min, max).
patterns.append(_op_pattern("tensorrt.nn.layer_norm", "relax.nn.layer_norm", 3))
patterns.append(_op_pattern("tensorrt.clip", "relax.clip", 3))
# strided_slice is called either with or without the optional strides argument.
patterns.append(_op_pattern("tensorrt.strided_slice", "relax.strided_slice", 5))
patterns.append(_op_pattern("tensorrt.strided_slice", "relax.strided_slice", 4))
return patterns
register_patterns(_tensorrt_patterns())
def partition_for_tensorrt(mod: IRModule) -> IRModule:
"""Partition the module, offloading TensorRT-supported subgraphs.
Parameters
----------
mod : tvm.ir.IRModule
The module to partition. Bind model parameters (e.g. via
``relax.transform.BindParams``) before calling this so that weights are
available to TensorRT as constants.
Returns
-------
mod : tvm.ir.IRModule
The module with TensorRT-supported subgraphs grouped into composite
functions annotated for the ``tensorrt`` codegen.
"""
patterns = get_patterns_with_prefix("tensorrt")
mod = FuseOpsByPattern(patterns, bind_constants=True, annotate_codegen=False)(mod)
mod = MergeCompositeFunctions()(mod)
return mod