[ONNX] Add Einsum converter (#8985)

* einsum

* address review

* move files around

* use generic topi op

* TODO comment

* jostle ci

* jostle ci

Author: anwang2009

include/tvm/relay/attrs/transform.h

@@ -475,7 +475,7 @@ struct ScanopAttrs : public tvm::AttrsNode<ScanopAttrs> {
         .describe("The first element is not included")
         .set_default(Bool(false));
   }
-};
+};  // struct ScanopAttrs
 
 /*! \brief Attributes used in unique operator */
 struct UniqueAttrs : public tvm::AttrsNode<UniqueAttrs> {
@@ -489,6 +489,15 @@ struct UniqueAttrs : public tvm::AttrsNode<UniqueAttrs> {
   }
 };  // struct UniqueAttrs
 
+/*! \brief Attributes used in einsum operator */
+struct EinsumAttrs : public tvm::AttrsNode<EinsumAttrs> {
+  String equation;
+
+  TVM_DECLARE_ATTRS(EinsumAttrs, "relay.attrs.EinsumAttrs") {
+    TVM_ATTR_FIELD(equation).describe("The einsum expression string");
+  }
+};  // struct EinsumAttrs
+
 }  // namespace relay
 }  // namespace tvm
 #endif  // TVM_RELAY_ATTRS_TRANSFORM_H_

python/tvm/relay/frontend/onnx.py

@@ -3501,6 +3501,15 @@ def _impl_v11(cls, inputs, attr, params):
         return _expr.TupleWrapper(_expr.Tuple([unique_vals, indices, inverse_indices, counts]), 4)
 
 
+class Einsum(OnnxOpConverter):
+    """Operator converter for Einsum"""
+
+    @classmethod
+    def _impl_v12(cls, inputs, attr, params):
+        equation = attr["equation"].decode("utf-8")
+        return _op.einsum(inputs, equation)
+
+
 class RandomUniform(OnnxOpConverter):
     """Operator converter for random_uniform"""
 
@@ -3864,6 +3873,7 @@ def _get_convert_map(opset):
         "Range": Range.get_converter(opset),
         "CumSum": CumSum.get_converter(opset),
         "Unique": Unique.get_converter(opset),
+        "Einsum": Einsum.get_converter(opset),
         # defs/control_flow
         "Loop": Loop.get_converter(opset),
         "If": If.get_converter(opset),

python/tvm/relay/op/__init__.py

@@ -54,6 +54,7 @@
 from . import _transform
 from . import _reduce
 from . import _algorithm
+from . import _math
 
 
 def _register_op_make():

python/tvm/relay/op/_math.py

@@ -0,0 +1,22 @@
+# 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.
+"""Backend compiler related feature registration"""
+from . import op as _reg
+from . import strategy
+
+# einsum
+_reg.register_strategy("einsum", strategy.einsum_strategy)

python/tvm/relay/op/_transform.py

@@ -182,6 +182,7 @@ def compute_unique(attrs, inputs, output_type):
 _reg.register_strategy("invert_permutation", strategy.invert_permutation_strategy)
 _reg.register_shape_func("invert_permutation", False, elemwise_shape_func)
 
+
 #####################
 #  Shape functions  #
 #####################

python/tvm/relay/op/strategy/cuda.py

@@ -1215,3 +1215,16 @@ def invert_permutation_strategy_cuda(attrs, inputs, out_type, target):
         name="invert_permutation.cuda",
     )
     return strategy
+
+
+@einsum_strategy.register(["cuda", "gpu"])
+def einsum_strategy_cuda(attrs, inputs, out_type, target):
+    """einsum cuda strategy"""
+    strategy = _op.OpStrategy()
+    # TODO: Add cuda-specific op implementation for einsum
+    strategy.add_implementation(
+        wrap_compute_einsum(topi.einsum),
+        wrap_topi_schedule(topi.generic.schedule_extern),
+        name="einsum.cuda",
+    )
+    return strategy

python/tvm/relay/op/strategy/generic.py

@@ -1669,3 +1669,24 @@ def invert_permutation_strategy(attrs, inputs, out_type, target):
         name="invert_permutation.generic",
     )
     return strategy
+
+
+def wrap_compute_einsum(topi_compute):
+    """Wrap einsum topi compute"""
+
+    def _compute_einsum(attrs, inputs, _):
+        return [topi_compute(attrs.equation, *inputs)]
+
+    return _compute_einsum
+
+
+@override_native_generic_func("einsum_strategy")
+def einsum_strategy(attrs, inputs, out_type, target):
+    """einsum generic strategy"""
+    strategy = _op.OpStrategy()
+    strategy.add_implementation(
+        wrap_compute_einsum(topi.einsum),
+        wrap_topi_schedule(topi.generic.schedule_einsum),
+        name="einsum.generic",
+    )
+    return strategy

python/tvm/relay/op/tensor.py

@@ -1104,6 +1104,29 @@ def concatenate(data, axis):
     return _make.concatenate(Tuple(data), axis)
 
 
+def einsum(data, equation):
+    """Evaluates the Einstein summation convention on data
+
+    Parameters
+    ----------
+    data : Union(List[relay.Expr], Tuple[relay.Expr])
+        A list of tensors.
+    equation : str
+        The einsum expression string.
+
+    Returns
+    -------
+    result : relay.Expr
+        The output tensor from the einsum op.
+    """
+    data = list(data)
+    if not data:
+        raise ValueError("relay.einsum requires data to be non-empty.")
+    if not isinstance(equation, str):
+        raise ValueError("einsum `equation` must be a str")
+    return _make.einsum(Tuple(data), equation)
+
+
 def stack(data, axis):
     """Join a sequence of arrays along a new axis.
 

python/tvm/topi/generic/__init__.py

@@ -39,3 +39,4 @@
 from .sort import *
 from .search import *
 from .image import *
+from .math import *

python/tvm/topi/generic/math.py

@@ -0,0 +1,34 @@
+# 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.
+"""Generic math operators"""
+from .default import default_schedule as _default_schedule
+
+
+def schedule_einsum(outs):
+    """Schedule for einsum operator.
+
+    Parameters
+    ----------
+    outs: Array of Tensor
+      The computation graph description of einsum.
+
+    Returns
+    -------
+    s: Schedule
+      The computation schedule for the op.
+    """
+    return _default_schedule(outs, False)