[microNPU] Determine block configs using the cascader (apache#10695)

The cascader needs to be able to choose the block config for operations in order to accurately model their performance.
 The cascader must attach the chosen block config to the te.Schedule. This is done using a pragma. The chosen block config
is also added to the TIR spec. If the cascader hasn't set a block config, it defaults to the existing block config
selection behaviour.

Co-authored-by: Matthew Barrett <[email protected]>

Co-authored-by: Matthew Barrett <[email protected]>

Author: 2 people

python/tvm/contrib/ethosu/cascader/scheduler.py

@@ -24,6 +24,7 @@
 from tvm import tir
 from .cascader_options import CascaderOptions
 from .graph import CascaderGraph, Part, Tensor, TESubgraph
+from .parts import EthosuPart
 from .tensor_config import MemoryRegion
 from .proposal import Proposal
 from .proposal_generator import generate_proposals
@@ -125,6 +126,23 @@ def apply_proposal(proposal: Proposal, sch: te.Schedule) -> None:
 
     """
     for plan in proposal.plans:
+        for part in plan.part_group:
+            if isinstance(part, EthosuPart):
+                tensor_config = plan.tensor_configs[part.output_tensor]
+                stripe_config = tensor_config.stripe_configs[0]
+                block_config = part.get_block_config(stripe_config)
+                iv = part.subgraph.output_tensor.op.axis[0]
+                block_shape = block_config.output_shape
+                if len(block_shape) == 4:
+                    height, width, depth = block_shape[1:]
+                else:
+                    height = block_shape[1]
+                    width = block_shape[3]
+                    depth = block_shape[2] * block_shape[4]
+                sch[part.subgraph.output_tensor].pragma(iv, "block_config_height", height)
+                sch[part.subgraph.output_tensor].pragma(iv, "block_config_width", width)
+                sch[part.subgraph.output_tensor].pragma(iv, "block_config_depth", depth)
+
         output_tensor_config = plan.output_config
         output_tensor = output_tensor_config.tensor
         output_part = output_tensor.producers[0]

python/tvm/relay/backend/contrib/ethosu/tir/binary_elementwise.py

@@ -86,7 +86,9 @@ def get_binary_elementwise_params(
     # Get feature map info
     serial_ifm, _ = get_ifm_params(input_pointer, producers)
     serial_ifm2, _ = get_ifm_params(input_pointer1, producers)
-    serial_ofm, replace_pointer, is_allocator = get_ofm_params(output_pointer, consumers, producers)
+    serial_ofm, serial_block_config, replace_pointer, is_allocator = get_ofm_params(
+        output_pointer, consumers, producers
+    )
     # Get activation info
     serial_activation = SerialActivation(
         op=attrs["activation"], clip_min=attrs["clip_min"], clip_max=attrs["clip_max"]
@@ -100,6 +102,7 @@ def get_binary_elementwise_params(
             reversed_operands=reversed_operands,
             activation=serial_activation,
             rounding_mode=attrs["rounding_mode"],
+            block_config=serial_block_config,
         ),
         output_pointer,
         replace_pointer,

python/tvm/relay/backend/contrib/ethosu/tir/convolution.py

@@ -63,7 +63,9 @@ def get_conv2d_params(stmt, producers, consumers):
     output_pointer = stores[0].buffer.data
     # Get feature map info
     serial_ifm, serial_padding = get_ifm_params(input_pointer, producers)
-    serial_ofm, replace_pointer, is_allocator = get_ofm_params(output_pointer, consumers, producers)
+    serial_ofm, serial_block_config, replace_pointer, is_allocator = get_ofm_params(
+        output_pointer, consumers, producers
+    )
     # Get kernel info
     serial_kernel = SerialKernel(
         width=int(rw.extent),
@@ -103,6 +105,7 @@ def get_conv2d_params(stmt, producers, consumers):
             activation=serial_activation,
             rounding_mode=attrs["rounding_mode"],
             upscale=attrs["upscale"],
+            block_config=serial_block_config,
         ),
         output_pointer,
         replace_pointer,

python/tvm/relay/backend/contrib/ethosu/tir/depthwise.py

@@ -72,7 +72,9 @@ def get_depthwise_conv2d_params(
     output_pointer = stores[0].buffer.data
     # Get feature map info
     serial_ifm, serial_padding = get_ifm_params(input_pointer, producers)
-    serial_ofm, replace_pointer, is_allocator = get_ofm_params(output_pointer, consumers, producers)
+    serial_ofm, serial_block_config, replace_pointer, is_allocator = get_ofm_params(
+        output_pointer, consumers, producers
+    )
     # Get kernel info
     serial_kernel = SerialKernel(
         width=int(rw.extent),
@@ -113,6 +115,7 @@ def get_depthwise_conv2d_params(
             activation=serial_activation,
             rounding_mode=attrs["rounding_mode"],
             upscale="NONE",
+            block_config=serial_block_config,
         ),
         output_pointer,
         replace_pointer,

python/tvm/relay/backend/contrib/ethosu/tir/dma.py

@@ -18,7 +18,7 @@
 """Extract parameters from the DMA operators in TIR."""
 import tvm
 from .utils import get_outer_loops, get_base_address, get_strides, get_op_attrs
-from .spec import SerialFeatureMap, SerialPadding
+from .spec import SerialBlockConfig, SerialFeatureMap, SerialPadding
 
 
 def get_pad_params(stmt):
@@ -253,6 +253,14 @@ def get_write_params(stmt):
 
     base_address = [get_base_address(index) for index in inner.indices]
     data_type = inner.buffer.data.type_annotation.element_type.dtype
+    if "block_config_height" in attrs:
+        block_config = SerialBlockConfig(
+            height=int(attrs["block_config_height"]),
+            width=int(attrs["block_config_width"]),
+            depth=int(attrs["block_config_depth"]),
+        )
+    else:
+        block_config = SerialBlockConfig(0, 0, 0)
     return (
         SerialFeatureMap(
             data_type=data_type,
@@ -273,6 +281,7 @@ def get_write_params(stmt):
             stride_w=strides[1],
             stride_c=strides[2],
         ),
+        block_config,
         input_pointer,
         output_pointer,
     )
@@ -327,6 +336,8 @@ def get_ofm_params(pointer, consumers, producers):
     -------
     serial_ifm : SerialFeatureMap
         The serializable OFM.
+    serial_block_config : SerialBlockConfig
+        The serializable block config.
     output_pointer : tvm.tir.Var
         The pointer that the OFM DMA pipeline produces.
     is_allocator : bool
@@ -336,11 +347,11 @@ def get_ofm_params(pointer, consumers, producers):
     convert_to_nhcwb16 = consumers[pointer]
     out_channels, _, output_pointer = get_convert_to_nhcwb16_params(convert_to_nhcwb16)
     write = consumers[output_pointer]
-    serial_ofm, _, output_pointer = get_write_params(write)
+    serial_ofm, serial_block_config, _, output_pointer = get_write_params(write)
     is_allocator = True
     if output_pointer not in producers:
         is_allocator = False
     elif producers[output_pointer] != write:
         is_allocator = False
     serial_ofm.channels = out_channels
-    return serial_ofm, output_pointer, is_allocator
+    return serial_ofm, serial_block_config, output_pointer, is_allocator

python/tvm/relay/backend/contrib/ethosu/tir/identity.py

@@ -18,7 +18,14 @@
 """Extract information from the identity operator in TIR."""
 from typing import Dict, Tuple
 import tvm
-from .spec import SerialKernel, SerialActivation, SerialPooling, SerialPadding, SerialFeatureMap
+from .spec import (
+    SerialBlockConfig,
+    SerialKernel,
+    SerialActivation,
+    SerialPooling,
+    SerialPadding,
+    SerialFeatureMap,
+)
 from .utils import get_op_attrs, get_base_address, get_strides, get_loads
 
 
@@ -164,6 +171,7 @@ def get_identity_params(
             activation=serial_activation,
             upscale="NONE",
             rounding_mode="TFL",
+            block_config=SerialBlockConfig(0, 0, 0),
         ),
         output_pointer,
         replace_pointer,

python/tvm/relay/backend/contrib/ethosu/tir/pooling.py

@@ -65,7 +65,9 @@ def get_pooling_params(
     output_pointer = stores[0].buffer.data
     # Get feature map info
     serial_ifm, serial_padding = get_ifm_params(input_pointer, producers)
-    serial_ofm, replace_pointer, is_allocator = get_ofm_params(output_pointer, consumers, producers)
+    serial_ofm, serial_block_config, replace_pointer, is_allocator = get_ofm_params(
+        output_pointer, consumers, producers
+    )
     # Get kernel info
     serial_kernel = SerialKernel(
         width=int(rw.extent),
@@ -90,6 +92,7 @@ def get_pooling_params(
             activation=serial_activation,
             rounding_mode=attrs["rounding_mode"],
             upscale=attrs["upscale"],
+            block_config=serial_block_config,
         ),
         output_pointer,
         replace_pointer,

python/tvm/relay/backend/contrib/ethosu/tir/scheduler.py

@@ -209,6 +209,13 @@ def _add_pragmas(stage, ax):
             for attr, val in stage.op.attrs.items():
                 if attr not in ("op", "lut") and not isinstance(val, Propagator):
                     stage.pragma(ax, str(attr), val)
+        if stage.op.axis[0] in stage.iter_var_attrs:
+            attrs = stage.iter_var_attrs[stage.op.axis[0]]
+            if "block_config_height" in attrs.pragma_keys:
+                pragmas = dict(zip([k.value for k in attrs.pragma_keys], attrs.pragma_values))
+                stage.pragma(ax, "block_config_height", pragmas["block_config_height"])
+                stage.pragma(ax, "block_config_width", pragmas["block_config_width"])
+                stage.pragma(ax, "block_config_depth", pragmas["block_config_depth"])
 
     for stage in sch.stages:
         if (

python/tvm/relay/backend/contrib/ethosu/tir/spec.py

@@ -174,6 +174,16 @@ def __init__(self, op: str, clip_min: int, clip_max: int):
         self.clip_max = clip_max
 
 
+class SerialBlockConfig(SerializableFormat):
+    """Specialization class to retrieve arguments of a BlockConfig
+    (similar to NpuBlockConfig of Vela) on a predefined ordering"""
+
+    def __init__(self, height: int, width: int, depth: int):
+        self.height = height
+        self.width = width
+        self.depth = depth
+
+
 class Serial2DConvolution(SerializableFormat):
     """Specialization class to retrieve arguments of
     a ethosu.conv2d tir extern call on a predefined ordering"""
@@ -190,6 +200,7 @@ def __init__(
         activation: SerialActivation,
         rounding_mode: str,
         upscale: str,
+        block_config: SerialBlockConfig,
     ):
         self.ifm = ifm
         self.ofm = ofm
@@ -201,6 +212,7 @@ def __init__(
         self.activation = activation
         self.rounding_mode = rounding_mode
         self.upscale = upscale
+        self.block_config = block_config
 
 
 class Serial2DDepthwise(SerializableFormat):
@@ -219,6 +231,7 @@ def __init__(
         activation: SerialActivation,
         rounding_mode: str,
         upscale: str,
+        block_config: SerialBlockConfig,
     ):
         self.ifm = ifm
         self.ofm = ofm
@@ -230,6 +243,7 @@ def __init__(
         self.activation = activation
         self.rounding_mode = rounding_mode
         self.upscale = upscale
+        self.block_config = block_config
 
 
 class SerialCopy(SerializableFormat):
@@ -261,6 +275,7 @@ def __init__(
         activation: SerialActivation,
         rounding_mode: str,
         upscale: str,
+        block_config: SerialBlockConfig,
     ):
         self.ifm = ifm
         self.ofm = ofm
@@ -270,6 +285,7 @@ def __init__(
         self.activation = activation
         self.rounding_mode = rounding_mode
         self.upscale = upscale
+        self.block_config = block_config
 
 
 class SerialBinaryElementwise(SerializableFormat):
@@ -285,6 +301,7 @@ def __init__(
         reversed_operands: bool,
         activation: SerialActivation,
         rounding_mode: str,
+        block_config: SerialBlockConfig,
     ):
         self.ifm = ifm
         self.ifm2 = ifm2
@@ -293,6 +310,7 @@ def __init__(
         self.reversed_operands = reversed_operands
         self.activation = activation
         self.rounding_mode = rounding_mode
+        self.block_config = block_config
 
 
 class SerialUnaryElementwise(SerializableFormat):
@@ -306,9 +324,11 @@ def __init__(
         operator_type: str,
         activation: SerialActivation,
         rounding_mode: str,
+        block_config: SerialBlockConfig,
     ):
         self.ifm = ifm
         self.ofm = ofm
         self.operator_type = operator_type
         self.activation = activation
         self.rounding_mode = rounding_mode
+        self.block_config = block_config

python/tvm/relay/backend/contrib/ethosu/tir/unary_elementwise.py

@@ -61,7 +61,9 @@ def get_unary_elementwise_params(stmt, producers, consumers):
     output_pointer = inner.buffer.data
     # Get feature map info
     serial_ifm, _ = get_ifm_params(input_pointer, producers)
-    serial_ofm, replace_pointer, is_allocator = get_ofm_params(output_pointer, consumers, producers)
+    serial_ofm, serial_block_config, replace_pointer, is_allocator = get_ofm_params(
+        output_pointer, consumers, producers
+    )
     # Get activation info
     serial_activation = SerialActivation(
         op=attrs["activation"], clip_min=attrs["clip_min"], clip_max=attrs["clip_max"]
@@ -73,6 +75,7 @@ def get_unary_elementwise_params(stmt, producers, consumers):
             operator_type=attrs["operator_type"],
             activation=serial_activation,
             rounding_mode=attrs["rounding_mode"],
+            block_config=serial_block_config,
         ),
         output_pointer,
         replace_pointer,