[RFC] UMA: Universal Modular Accelerator Interface

manupa-arm · February 24, 2022, 11:26am

Apologies for not getting back to this in time. Thanks for the proposal! and it broadly looks like wrapping the Target Hooks RFC (by @Mousius ) : tvm-rfcs/rfcs/0010-target-registered-compiler-flow-customisation.md at main · apache/tvm-rfcs · GitHub, and exposing a nice/structured interface to python. It is nice to see progress on this .

I would like to suggest potential text changes for the formal RFC to those of us who are familiar with the existing flow (specially around naming).

Maybe it is worth mentioning these are current implemented as partition_for_<(backend\target)> ?

github.com/apache/tvm

python/tvm/relay/op/contrib/dnnl.py

f1ff61a7b


      
                          "nn.conv2d",
                          "nn.conv3d",
                          "nn.conv2d_transpose",
                          "nn.conv3d_transpose",
                      ]:
                          dnnl_patterns.append(make_dnnl_pattern(conv_name, with_bias, elt))
                      dnnl_patterns.append(make_dnnl_pattern("nn.dense", with_bias, elt))
              return dnnl_patterns
          
          
          def partition_for_dnnl(mod, params=None):
              """Partition the graph greedily offloading supported operators to DNNL.
          
              Parameters
              ----------
              mod : Module
                  The module to run passes on.
              params : Optional[Dict[str, NDArray]]
                  Constant input parameters.
              Returns
              -------

github.com/apache/tvm

python/tvm/relay/op/contrib/tensorrt.py

f1ff61a7b


      
              pass_ctx = tvm.transform.PassContext.current()
              if "relay.ext.tensorrt.options" in pass_ctx.config:
                  return pass_ctx.config["relay.ext.tensorrt.options"].remove_no_mac_subgraphs
              logger.warning(
                  "PassContext has no relay.ext.tensorrt.options config, using default value "
                  "remove_no_mac_subgraphs=False."
              )
              return False
          
          
          def partition_for_tensorrt(
              mod,
              params=None,
              version=None,
              use_implicit_batch=True,
              remove_no_mac_subgraphs=False,
              max_workspace_size=1 << 30,
          ):
              """Partition the graph greedily offloading supported operators to TensorRT.
          
              Parameters

github.com/apache/tvm

python/tvm/relay/op/contrib/ethosu.py

f1ff61a7b


      
                  (
                      SqueezeParams.composite_name,
                      squeeze_pattern(),
                      lambda pat: SqueezeParams(pat).is_valid(),
                  ),
              ]
          
          
          # pylint: disable=unused-argument
          @requires_vela
          def partition_for_ethosu(
              mod: tvm.ir.IRModule, params: Optional[Dict[str, tvm.runtime.NDArray]] = None, **opts
          ):
              """This helper function partition the relay graph as produced by the
              relay frontend for a given model into external functions
              to be presented to the codegen.
          
              Parameters
              ----------
              mod : tvm.ir.IRModule
                  The IRModule that gets generated from a relay frontend

I am a bit curious, why this interface is specifically positioned as an “accelerator” (as in UMA) partitioner though ? i.e. Would it not be used to support optimized library support as we currently have today with BYOC ?

MJKlaiber:

class MyCustomAcceleratorPartitioner(UMAPartitioner):
    @property
    def target_name(self):
        return "my_custom_accelerator"

    def _register_patterns(self):
        self._register_pattern("conv1d_relu", conv1d_relu_pattern())
    
    def _register_relay_passes(self):
        self._register_relay_pass(1, ConfigGenerator())
        self._register_relay_pass(2, BufferScopeAnnotator())

Since the proposal suggests to use the properly registered targets, any reason should we stick to target_name (str) as opposed to the actual TargetKind ?

Following up on the above question, what are your thoughts on moving the UMAPartitioner inside relay.build(…) ?

Also this seemed to be proposed on using S-TIR (as opposed to “legacy” TE->TIR pipeline), would you be able to share the motivation to the partitioning of tir_schedules and tir_passes ? (Im asking mainly because they will all be S-TIR → S-TIR IRModule passes).

Following from the above question, is there an ambition to handover S-TIR back to the core compiler ?

Following up on Mark’s comments,

Mark, we are quite looking forward for the RFC for this, especially related to reference-level explanation to see where this work is headed – which I believe might be better to know in this mutual interest of structuring BYOC targets.

However, I think we all share the ambition to replace kCompiler strings to be targets if can get more support from the community.