Relax non-negative-integer Pow to ipow (fix symbolic-exponent codegen)

dace/codegen/codegen.py

@@ -209,6 +209,15 @@ def generate_code(sdfg: SDFG, validate=True) -> List[CodeObject]:
     infer_types.infer_connector_types(sdfg)
     infer_types.set_default_schedule_and_storage_types(sdfg, None)
 
+    # Lower ``base ** exp`` to ``ipow`` (exact integer power) wherever the exponent is a
+    # provable non-negative integer -- array sizes, subscripts, loop bounds and interstate
+    # edges. Runs here, right before codegen, rather than in ``simplify``: a size is "just a
+    # more complicated expression" that should not perturb the SDFG, and keeping powers as
+    # ``Pow`` through simplification lets SymPy's power laws fold them (``R**i * R**(K-i-1) ->
+    # R**(K-1)``) before the opaque ``ipow`` freezes the form.
+    from dace.transformation.passes.relax_integer_powers import RelaxIntegerPowers
+    RelaxIntegerPowers().apply_pass(sdfg, {})
+
     frame = framecode.DaCeCodeGenerator(sdfg)
 
     # Test for undefined symbols in SDFG arguments

dace/codegen/cppunparse.py

@@ -1000,6 +1000,18 @@ def _Call(self, t: ast.Call):
                 op = self.callbools[t.func.id]()
                 self.dispatch(ast.BoolOp(op=op, values=t.args))
                 return
+            # ``ipow`` (integer power, RelaxIntegerPowers) has no bare C++ symbol -- it
+            # lives in ``dace::math``. The symbolic printer qualifies it for descriptor
+            # expressions; qualify it here too so it is correct in loop bounds / interstate
+            # edges, which unparse through this AST path rather than the symbolic one.
+            elif t.func.id == 'ipow':
+                self.write('dace::math::ipow')
+                self.write('(')
+                for i, e in enumerate(t.args):
+                    self.write(', ' if i else '')
+                    self.dispatch(e)
+                self.write(')')
+                return
 
         self.dispatch(t.func)
         self.write("(")

dace/frontend/python/parser.py

@@ -131,7 +131,8 @@ def infer_symbols_from_datadescriptor(sdfg: SDFG,
                     sym_dim = sym_dim.subs(repldict)
 
                 if symbolic.issymbolic(sym_dim - real_dim):
-                    equations.append(sym_dim - real_dim)
+                    # ipow is semantically Pow; restore it so the solver can invert the shape.
+                    equations.append((sym_dim - real_dim).replace(symbolic.ipow, lambda b, e: b**e))
 
     if len(symbols) == 0:
         return {}

dace/runtime/include/dace/math.h

@@ -546,8 +546,21 @@ namespace dace
             return result;
         }
 
-        template<typename T>
-        DACE_HDFI T ipow(const T& a, const unsigned int& b) {
+        // Scalar types seed at ``T(1)`` so ``ipow(a, 0) == 1``.
+        template<typename T,
+                 typename std::enable_if<std::is_constructible<T, int>::value>::type* = nullptr>
+        DACE_HDFI T ipow(const T a, const unsigned int b) {
+            T result = T(1);
+            for (unsigned int i = 0; i < b; ++i)
+                result *= a;
+            return result;
+        }
+
+        // Vector types have no scalar constructor, so seed at ``a``. Only ever reached with a
+        // compile-time exponent >= 1 (the constant-power path emits a literal 1 for exponent 0).
+        template<typename T,
+                 typename std::enable_if<!std::is_constructible<T, int>::value>::type* = nullptr>
+        DACE_HDFI T ipow(const T a, const unsigned int b) {
             T result = a;
             for (unsigned int i = 1; i < b; ++i)
                 result *= a;

dace/symbolic.py

@@ -911,8 +911,8 @@ def swalk(expr, enter_functions=False):
 
 
 _builtin_userfunctions = {
-    'int_floor', 'int_ceil', 'abs', 'Abs', 'min', 'Min', 'max', 'Max', 'not', 'Not', 'Eq', 'NotEq', 'Ne', 'AND', 'OR',
-    'pow', 'round'
+    'int_floor', 'int_ceil', 'ipow', 'abs', 'Abs', 'min', 'Min', 'max', 'Max', 'not', 'Not', 'Eq', 'NotEq', 'Ne', 'AND',
+    'OR', 'pow', 'round'
 }
 
 
@@ -1093,6 +1093,28 @@ def _eval_is_integer(self):
         return True
 
 
+class ipow(sympy.Function):
+    """Integer power ``base ** exp`` with a non-negative integer exponent, lowered
+    to ``dace::math::ipow`` (repeated multiply, exact integer) -- valid where
+    ``dace::math::pow`` (libm ``double``) is not: an array size, subscript or loop
+    bound.  ``RelaxIntegerPowers`` mints these from ``Pow``."""
+
+    @classmethod
+    def eval(cls, base, exp):
+        if base.is_Number and exp.is_Number and exp.is_integer and exp.is_nonnegative:
+            return base**exp
+
+    def _eval_is_integer(self):
+        base, exp = self.args
+        if base.is_integer and exp.is_integer:
+            return True
+
+    def _eval_is_positive(self):
+        base, _exp = self.args
+        if base.is_positive:
+            return True
+
+
 class OR(sympy.Function):
 
     @classmethod
@@ -1830,6 +1852,7 @@ def _unary_minus(a):
         'Le': sympy.Le,
         'int_floor': int_floor,
         'int_ceil': int_ceil,
+        'ipow': ipow,
         'IfExpr': IfExpr,
         'Mod': sympy.Mod,
         'Attr': Attr,
@@ -2203,6 +2226,7 @@ def _rewrite_typed_complex(m: re.Match) -> str:
     'int_floor': int_floor,
     '__int_floor': __int_floor,
     'int_ceil': int_ceil,
+    'ipow': ipow,
     'IfExpr': IfExpr,
     'Mod': sympy.Mod,
     'Attr': Attr,
@@ -2348,6 +2372,8 @@ def _print_Function(self, expr):
         if base == 'int_floor' and as_operator:
             op = '/' if self.cpp_mode else '//'
             return '((%s) %s (%s))' % (self._print(expr.args[0]), op, self._print(expr.args[1]))
+        if str(expr.func) == 'ipow' and self.cpp_mode:
+            return 'dace::math::ipow(%s, %s)' % (self._print(expr.args[0]), self._print(expr.args[1]))
         if str(expr.func) == 'IfExpr':
             cond, tval, fval = (self._print(a) for a in expr.args)
             if self.cpp_mode: