Add support for the Einsum operator

docs/dev/onnx_operators.rst

@@ -210,7 +210,14 @@ Operator Support Matrix
 |                          |           |                 | shape is not                 |
 |                          |           |                 | supported                    |
 +--------------------------+-----------+-----------------+------------------------------+
-| Einsum                   | 👷        | 👷              |                              |
+| Einsum                   | ✅        | Any             | more than 1 diagonal per     |
+|                          |           |                 | input is not supported       |
+|                          |           |                 | e.g. ``iijj->ij``            |
+|                          |           |                 |                              |
+|                          |           |                 | batch diagonal where batches |
+|                          |           |                 | are not the leading dims is  |
+|                          |           |                 | not supported                |
+|                          |           |                 | e.g. ``ii...->i...``         |
 +--------------------------+-----------+-----------------+------------------------------+
 | Elu                      | ✅        | FP8, FP16,      |                              |
 |                          |           | FP32, FP64      |                              |

src/convert_to_json.cpp

@@ -1,7 +1,7 @@
 /*
  * The MIT License (MIT)
  *
- * Copyright (c) 2015-2022 Advanced Micro Devices, Inc. All rights reserved.
+ * Copyright (c) 2015-2024 Advanced Micro Devices, Inc. All rights reserved.
  *
  * Permission is hereby granted, free of charge, to any person obtaining a copy
  * of this software and associated documentation files (the "Software"), to deal
@@ -23,66 +23,17 @@
  */
 #include <algorithm>
 #include <string>
-#include <vector>
-#include <functional>
 #include <sstream>
 #include <migraphx/errors.hpp>
 #include <migraphx/ranges.hpp>
 #include <migraphx/convert_to_json.hpp>
 #include <migraphx/stringutils.hpp>
+#include <migraphx/lexing.hpp>
 
 namespace migraphx {
 inline namespace MIGRAPHX_INLINE_NS {
 
-using token = std::pair<const char*, const char*>;
-using lexer = std::function<const char*(const char* start, const char* end)>;
-
-template <class P>
-auto lex_while(P p)
-{
-    return [=](const char* start, const char* end) {
-        return std::find_if(start, end, [&](char c) { return not p(c); });
-    };
-}
-
-template <class P>
-auto lex_if(P p)
-{
-    return [=](const char* start, const char*) {
-        if(p(*start))
-            return start + 1;
-        return start;
-    };
-}
-
-std::vector<token> tokenize(const char* start, const char* end, const std::vector<lexer>& lexers)
-{
-    std::vector<token> result;
-    while(start != end)
-    {
-        bool error = true;
-        for(const auto& l : lexers)
-        {
-            const auto* next = l(start, end);
-            if(next != start)
-            {
-                result.emplace_back(start, next);
-                start = next;
-                error = false;
-                break;
-            }
-        }
-
-        if(error)
-        {
-            MIGRAPHX_THROW("TOKENIZE: no token found!");
-        }
-    }
-
-    return result;
-}
-
-std::vector<token> json_tokenize(const std::string& s)
+std::vector<std::string_view> json_tokenize(const std::string& s)
 {
     std::vector<lexer> lexers;
 
@@ -133,7 +84,7 @@ std::string convert_to_json(const std::string& str)
 
     for(auto& token : tokens)
     {
-        std::string s(token.first, token.second);
+        std::string s(token);
         if(starts_with(s, "#") or starts_with(s, "//"))
             continue;
         if(std::isalpha(s.front()) != 0 and

src/include/migraphx/lexing.hpp

@@ -0,0 +1,98 @@
+/*
+ * The MIT License (MIT)
+ *
+ * Copyright (c) 2015-2024 Advanced Micro Devices, Inc. All rights reserved.
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a copy
+ * of this software and associated documentation files (the "Software"), to deal
+ * in the Software without restriction, including without limitation the rights
+ * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+ * copies of the Software, and to permit persons to whom the Software is
+ * furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in
+ * all copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL THE
+ * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+ * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+ * THE SOFTWARE.
+ */
+#ifndef MIGRAPHX_GUARD_MIGRAPHLIB_LEXING_HPP
+#define MIGRAPHX_GUARD_MIGRAPHLIB_LEXING_HPP
+
+#include <functional>
+#include <string>
+#include <vector>
+#include <migraphx/errors.hpp>
+
+namespace migraphx {
+inline namespace MIGRAPHX_INLINE_NS {
+
+using lexer = std::function<const char*(const char* start, const char* end)>;
+
+template <class P>
+inline auto lex_while(P p)
+{
+    return [=](const char* start, const char* end) {
+        return std::find_if(start, end, [&](char c) { return not p(c); });
+    };
+}
+
+template <class P>
+inline auto lex_if(P p)
+{
+    return [=](const char* start, const char*) {
+        if(p(*start))
+            return start + 1;
+        return start;
+    };
+}
+
+inline auto lex_equal(const std::string& s)
+{
+    return [=](const char* start, const char* end) {
+        auto n = end - start;
+        if(n < s.size())
+            return start;
+        if(std::equal(start, start + s.size(), s.data()))
+            return start + s.size();
+        return start;
+    };
+}
+
+inline std::vector<std::string_view>
+tokenize(const char* start, const char* end, const std::vector<lexer>& lexers)
+{
+    std::vector<std::string_view> result;
+    while(start != end)
+    {
+        bool error = true;
+        for(const auto& l : lexers)
+        {
+            const auto* next = l(start, end);
+            if(next != start)
+            {
+                result.emplace_back(start, next - start);
+                start = next;
+                error = false;
+                break;
+            }
+        }
+
+        if(error)
+        {
+            MIGRAPHX_THROW("TOKENIZE: no token found!");
+        }
+    }
+
+    return result;
+}
+
+} // namespace MIGRAPHX_INLINE_NS
+} // namespace migraphx
+
+#endif