add TaggedExpression #688
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -922,19 +922,16 @@ def __init__(self, knl, callables_table, kernel_rec, | |
| def combine(self, values): | ||
| return sum(values) | ||
|
|
||
| def map_constant(self, expr, *args): | ||
| return self.new_zero_poly_map() | ||
|
|
||
| map_tagged_variable = map_constant | ||
| map_variable = map_constant | ||
|
|
||
| def map_tagged_expression(self, expr, *args): | ||
| return self.rec(expr.expr, expr.tags) | ||
|
|
||
| def map_call(self, expr, tags): | ||
| from loopy.symbolic import ResolvedFunction | ||
| assert isinstance(expr.function, ResolvedFunction) | ||
| clbl = self.callables_table[expr.function.name] | ||
|
|
@@ -944,37 +941,40 @@ def map_call(self, expr): | |
| return self.new_poly_map( | ||
| {Op(dtype=self.type_inf(expr), | ||
| name="func:"+clbl.name, | ||
| tags=tags, | ||
| count_granularity=self.arithmetic_count_granularity, | ||
| kernel_name=self.knl.name): self.one} | ||
| ) + self.rec(expr.parameters, tags) | ||
| else: | ||
| return super().map_call(expr, tags) | ||
|
|
||
| def map_subscript(self, expr, *args): | ||
| if self.count_within_subscripts: | ||
| return self.rec(expr.index, *args) | ||
| else: | ||
| return self.new_zero_poly_map() | ||
|
|
||
| def map_sub_array_ref(self, expr, *args): | ||
| # generates an array view, considered free | ||
| return self.new_zero_poly_map() | ||
|
|
||
| def map_sum(self, expr, tags): | ||
| assert expr.children | ||
| return self.new_poly_map( | ||
| {Op(dtype=self.type_inf(expr), | ||
| name="add", | ||
| tags=tags, | ||
| count_granularity=self.arithmetic_count_granularity, | ||
| kernel_name=self.knl.name): | ||
| self.zero + (len(expr.children)-1)} | ||
| ) + sum(self.rec(child, tags) for child in expr.children) | ||
|
|
||
| def map_product(self, expr, tags): | ||
| from pymbolic.primitives import is_zero | ||
| assert expr.children | ||
| return sum(self.new_poly_map({Op(dtype=self.type_inf(expr), | ||
| name="mul", | ||
| tags=tags, | ||
| count_granularity=( | ||
| self.arithmetic_count_granularity), | ||
| kernel_name=self.knl.name): self.one}) | ||
|
|
@@ -983,97 +983,104 @@ def map_product(self, expr): | |
| if not is_zero(child + 1)) + \ | ||
| self.new_poly_map({Op(dtype=self.type_inf(expr), | ||
| name="mul", | ||
| tags=tags, | ||
| count_granularity=( | ||
| self.arithmetic_count_granularity), | ||
| kernel_name=self.knl.name): -self.one}) | ||
|
|
||
| def map_quotient(self, expr, tags): | ||
| return self.new_poly_map({Op(dtype=self.type_inf(expr), | ||
| name="div", | ||
| tags=tags, | ||
| count_granularity=self.arithmetic_count_granularity, | ||
| kernel_name=self.knl.name): self.one}) \ | ||
| + self.rec(expr.numerator, tags) \ | ||
| + self.rec(expr.denominator, tags) | ||
|
|
||
| map_floor_div = map_quotient | ||
| map_remainder = map_quotient | ||
|
|
||
| def map_power(self, expr, tags): | ||
| return self.new_poly_map({Op(dtype=self.type_inf(expr), | ||
| name="pow", | ||
| tags=tags, | ||
| count_granularity=self.arithmetic_count_granularity, | ||
| kernel_name=self.knl.name): self.one}) \ | ||
| + self.rec(expr.base, tags) \ | ||
| + self.rec(expr.exponent, tags) | ||
|
|
||
| def map_left_shift(self, expr, tags): | ||
| return self.new_poly_map({Op(dtype=self.type_inf(expr), | ||
| name="shift", | ||
| tags=tags, | ||
| count_granularity=self.arithmetic_count_granularity, | ||
| kernel_name=self.knl.name): self.one}) \ | ||
| + self.rec(expr.shiftee, tags) \ | ||
| + self.rec(expr.shift, tags) | ||
|
|
||
| map_right_shift = map_left_shift | ||
|
|
||
| def map_bitwise_not(self, expr, tags): | ||
| return self.new_poly_map({Op(dtype=self.type_inf(expr), | ||
| name="bw", | ||
| tags=tags, | ||
| count_granularity=self.arithmetic_count_granularity, | ||
| kernel_name=self.knl.name): self.one}) \ | ||
| + self.rec(expr.child, tags) | ||
|
|
||
| def map_bitwise_or(self, expr, tags): | ||
| return self.new_poly_map({Op(dtype=self.type_inf(expr), | ||
| name="bw", | ||
| tags=tags, | ||
| count_granularity=self.arithmetic_count_granularity, | ||
| kernel_name=self.knl.name): | ||
| self.zero + (len(expr.children)-1)}) \ | ||
| + sum(self.rec(child, tags) for child in expr.children) | ||
|
|
||
| map_bitwise_xor = map_bitwise_or | ||
| map_bitwise_and = map_bitwise_or | ||
|
|
||
| def map_if(self, expr, *args): | ||
|
There was a problem hiding this comment. In 0bbbaec, I changed all occurrences to |
||
| warn_with_kernel(self.knl, "summing_if_branches_ops", | ||
| "ExpressionOpCounter counting ops as sum of " | ||
| "if-statement branches.") | ||
| return self.rec(expr.condition, *args) + self.rec(expr.then, *args) \ | ||
| + self.rec(expr.else_, *args) | ||
|
|
||
| def map_if_positive(self, expr, *args): | ||
| warn_with_kernel(self.knl, "summing_ifpos_branches_ops", | ||
| "ExpressionOpCounter counting ops as sum of " | ||
| "if_pos-statement branches.") | ||
| return self.rec(expr.criterion, *args) + self.rec(expr.then, *args) \ | ||
| + self.rec(expr.else_, *args) | ||
|
|
||
| def map_min(self, expr, tags): | ||
| return self.new_poly_map({Op(dtype=self.type_inf(expr), | ||
| name="maxmin", | ||
| tags=tags, | ||
| count_granularity=self.arithmetic_count_granularity, | ||
| kernel_name=self.knl.name): | ||
| len(expr.children)-1}) \ | ||
| + sum(self.rec(child, tags) for child in expr.children) | ||
|
|
||
| map_max = map_min | ||
|
|
||
| def map_common_subexpression(self, expr, *args): | ||
| raise NotImplementedError("ExpressionOpCounter encountered " | ||
| "common_subexpression, " | ||
| "map_common_subexpression not implemented.") | ||
|
|
||
| def map_substitution(self, expr, *args): | ||
| raise NotImplementedError("ExpressionOpCounter encountered " | ||
| "substitution, " | ||
| "map_substitution not implemented.") | ||
|
|
||
| def map_derivative(self, expr, *args): | ||
| raise NotImplementedError("ExpressionOpCounter encountered " | ||
| "derivative, " | ||
| "map_derivative not implemented.") | ||
|
|
||
| def map_slice(self, expr, *args): | ||
| raise NotImplementedError("ExpressionOpCounter encountered slice, " | ||
| "map_slice not implemented.") | ||
|
|
||
|
|
||
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Could you type those newly-added mapper methods? I'm laying the groundwork for that in #897.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Done as part of 4ca4db4