Support *args in CombineMapper.map_sub_array_ref

inducer · Jul 2, 2024 · 913e360 · 913e360
1 parent 8116d69
commit 913e360
Showing 1 changed file with 5 additions and 3 deletions.
diff --git a/loopy/symbolic.py b/loopy/symbolic.py
@@ -280,10 +280,12 @@ def map_reduction(self, expr, *args, **kwargs):
     def map_type_cast(self, expr, *args, **kwargs):
         return self.rec(expr.child, *args, **kwargs)
 
-    def map_sub_array_ref(self, expr):
+    def map_sub_array_ref(self, expr, *args, **kwargs):
         return self.combine((
-            self.rec(expr.subscript),
-            self.combine(tuple(self.rec(idx) for idx in expr.swept_inames))))
+            self.rec(expr.subscript, *args, **kwargs),
+            self.combine(tuple(
+                         self.rec(idx, *args, **kwargs)
+                         for idx in expr.swept_inames))))
 
     map_linear_subscript = CombineMapperBase.map_subscript