Fix sym name.

DeepLink-org · Nov 13, 2023 · d67e86a · d67e86a
1 parent 442b3db
commit d67e86a
Showing 1 changed file with 6 additions and 2 deletions.
diff --git a/dicp/dicp/vendor/AscendGraph/codegen/ascend.py b/dicp/dicp/vendor/AscendGraph/codegen/ascend.py
@@ -219,6 +219,8 @@ def process_sym_name(self, st):
             return st
         elif '+' in st:
             sp = st.split('+')
+            if len(sp) > 2:
+                sp = [sp[0], '+'.join(sp[1:])]
             assert (len(sp) == 2)
             sp = [elem.strip() for elem in sp]
             if sp[0].isdigit():
@@ -232,10 +234,10 @@ def process_sym_name(self, st):
                 return self.process_sym_name(sp[0]) + '+' + sp[1]
         elif '-' in st:
             sp = st.split('-')
+            if len(sp) > 2:
+                sp = [sp[0], '-'.join(sp[1:])]
             assert (len(sp) == 2)
             sp = [elem.strip() for elem in sp]
-            if sp[0].isdigit():
-                (sp[1], sp[0]) = (sp[0], sp[1])
             if sp[0] in self.sym_in_args:
                 arg, idx = self.sym_in_args[sp[0]]
                 return "{}.shape[{}]".format(arg, idx) + '-' + sp[1]
@@ -245,6 +247,8 @@ def process_sym_name(self, st):
                 return self.process_sym_name(sp[0]) + '-' + sp[1]
         elif '*' in st:
             sp = st.split('*')
+            if len(sp) > 2:
+                sp = [sp[0], '*'.join(sp[1:])]
             assert (len(sp) == 2)
             sp = [elem.strip() for elem in sp]
             if sp[0].isdigit():