[ONNX] Add Einsum converter

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

@@ -3517,6 +3517,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"""
 
@@ -3763,6 +3772,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/_transform.py

@@ -182,6 +182,11 @@ 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)
 
+
+# einsum
+_reg.register_strategy("einsum", strategy.einsum_strategy)
+
+
 #####################
 #  Shape functions  #
 #####################

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

@@ -1210,3 +1210,15 @@ 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()
+    strategy.add_implementation(
+        wrap_compute_einsum(topi.cuda.einsum),
+        wrap_topi_schedule(topi.cuda.schedule_einsum),
+        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/search.py

@@ -86,3 +86,19 @@ def schedule_unique(outs):
       The computation schedule for the op.
     """
     return _default_schedule(outs, False)
+
+
+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)

src/relay/op/tensor/transform.cc

@@ -33,6 +33,7 @@
 #include <tvm/tir/op.h>
 #include <tvm/topi/broadcast.h>
 #include <tvm/topi/detail/constant_utils.h>
+#include <tvm/topi/einsum.h>
 #include <tvm/topi/elemwise.h>
 #include <tvm/topi/nn.h>
 #include <tvm/topi/reduction.h>
@@ -2431,6 +2432,85 @@ RELAY_REGISTER_OP("broadcast_to_like")
     .set_attr<FTVMCompute>("FTVMCompute", BroadCastToLikeCompute)
     .set_attr<TOpPattern>("TOpPattern", kBroadcast);
 
+// relay.einsum
+TVM_REGISTER_NODE_TYPE(EinsumAttrs);
+
+bool EinsumRel(const Array<Type>& types, int num_inputs, const Attrs& attrs,
+               const TypeReporter& reporter) {
+  // Check attrs
+  const EinsumAttrs* param = attrs.as<EinsumAttrs>();
+  if (param == nullptr) {
+    reporter->GetDiagCtx().EmitFatal(Diagnostic::Error(reporter->GetSpan())
+                                     << "the call attributes are not defined");
+    return false;
+  }
+
+  // types: [data, result]
+  ICHECK_EQ(types.size(), 2) << "the arity of einsum is 2, not " << types.size();
+
+  // Check input type is a tuple.
+  const auto* tensor_tuple = types[0].as<TupleTypeNode>();
+  if (tensor_tuple == nullptr) {
+    reporter->GetDiagCtx().EmitFatal(
+        Diagnostic::Error(reporter->GetSpan())
+        << "einsum requires a tuple of tensors as the first argument, found "
+        << PrettyPrint(types[0]));
+    return false;
+  }
+
+  // Check the input tuple consists of tensors with consistent dtype.
+  const auto& first = Downcast<TensorType>(tensor_tuple->fields[0]);
+  const DataType dtype = first->dtype;
+  std::vector<Array<PrimExpr>> input_shapes;
+  for (const Type& ele : tensor_tuple->fields) {
+    if (ele.as<IncompleteTypeNode>()) {
+      return false;
+    }
+
+    const auto& e = Downcast<TensorType>(ele);
+
+    const DataType& e_dtype = e->dtype;
+    if (e_dtype != dtype) {
+      throw Error("relay.einsum requires all tensors have the same dtype");
+    }
+    input_shapes.push_back(e->shape);
+  }
+
+  // Calculate output shape
+  Array<IndexExpr> oshape = topi::NumpyEinsumShape(param->equation, input_shapes);
+
+  auto rtype = TensorType(oshape, dtype);
+  reporter->Assign(types[1], rtype);
+  return true;
+}
+
+Array<te::Tensor> EinsumCompute(const Attrs& attrs, const Array<te::Tensor>& inputs,
+                                const Type& out_type) {
+  const EinsumAttrs* param = attrs.as<EinsumAttrs>();
+  ICHECK(param != nullptr);
+  return Array<te::Tensor>{topi::einsum(param->equation, inputs)};
+}
+
+Expr MakeEinsum(Expr data, String equation) {
+  auto attrs = make_object<EinsumAttrs>();
+  attrs->equation = std::move(equation);
+  static const Op& op = Op::Get("einsum");
+  return Call(op, {data}, Attrs(attrs), {});
+}
+
+TVM_REGISTER_GLOBAL("relay.op._make.einsum").set_body_typed(MakeEinsum);
+
+RELAY_REGISTER_OP("einsum")
+    .describe(R"doc(Evaluates the Einstein summation convention
+on the operands)doc" TVM_ADD_FILELINE)
+    .set_attrs_type<EinsumAttrs>()
+    .set_num_inputs(1)
+    .add_argument("data", "Tuple of Tensors", "The input list of tensors.")
+    .set_support_level(11)
+    .add_type_rel("Einsum", EinsumRel)
+    .set_attr<FTVMCompute>("FTVMCompute", EinsumCompute)
+    .set_attr<TOpPattern>("TOpPattern", kInjective);
+
 // Adapter function to make int array.
 Array<Integer> GetIntArray(Array<IndexExpr> arr) {
   for (size_t i = 0; i < arr.size(); ++i) {

tests/python/frontend/onnx/test_forward.py

@@ -4739,11 +4739,6 @@ def verify_eyelike(indata):
     "test_dropout_default_mask",
     "test_dropout_default_mask_ratio",
     "test_dropout_default_ratio",
-    "test_einsum_batch_diagonal",
-    "test_einsum_batch_matmul",
-    "test_einsum_inner_prod",
-    "test_einsum_sum",
-    "test_einsum_transpose",
     "test_greater_equal",
     "test_greater_equal_bcast",
     "test_if_seq",