1from numba.core import errors, ir 2from numba.core.rewrites import register_rewrite, Rewrite 3 4 5@register_rewrite('before-inference') 6class DetectStaticBinops(Rewrite): 7 """ 8 Detect constant arguments to select binops. 9 """ 10 11 # Those operators can benefit from a constant-inferred argument 12 rhs_operators = {'**'} 13 14 def match(self, func_ir, block, typemap, calltypes): 15 self.static_lhs = {} 16 self.static_rhs = {} 17 self.block = block 18 # Find binop expressions with a constant lhs or rhs 19 for expr in block.find_exprs(op='binop'): 20 try: 21 if (expr.fn in self.rhs_operators 22 and expr.static_rhs is ir.UNDEFINED): 23 self.static_rhs[expr] = func_ir.infer_constant(expr.rhs) 24 except errors.ConstantInferenceError: 25 continue 26 27 return len(self.static_lhs) > 0 or len(self.static_rhs) > 0 28 29 def apply(self): 30 """ 31 Store constant arguments that were detected in match(). 32 """ 33 for expr, rhs in self.static_rhs.items(): 34 expr.static_rhs = rhs 35 return self.block 36