From 913e36031f4fdc94135741737f7b8c3e90c3698c Mon Sep 17 00:00:00 2001 From: Andreas Kloeckner Date: Tue, 2 Jul 2024 11:46:55 -0500 Subject: [PATCH] Support *args in CombineMapper.map_sub_array_ref --- loopy/symbolic.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/loopy/symbolic.py b/loopy/symbolic.py index d912c64f..2ea506ef 100644 --- a/loopy/symbolic.py +++ b/loopy/symbolic.py @@ -280,10 +280,12 @@ class CombineMapper(CombineMapperBase): 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 -- GitLab