def _get_pair_first_container(func_ir, rhs):
assert isinstance(rhs, ir.Expr) and rhs.op == 'pair_first'
iternext = get_definition(func_ir, rhs.value)
require(isinstance(iternext, ir.Expr) and iternext.op == 'iternext')
getiter = get_definition(func_ir, iternext.value)
require(isinstance(iternext, ir.Expr) and getiter.op == 'getiter')
return getiter.value
def _get_const_binary_expr(stencil_ir, func_ir, index_def):
"""evaluate constant binary expr if possible
otherwise, raise GuardException
"""
require(isinstance(index_def, ir.Expr) and index_def.op == 'binop')
arg1 = _get_const_index_expr_inner(stencil_ir, func_ir, index_def.lhs)
arg2 = _get_const_index_expr_inner(stencil_ir, func_ir, index_def.rhs)
op = OPERATORS_TO_BUILTINS[index_def.fn]
return eval("{}{}{}".format(arg1, op, arg2))
def _get_const_unary_expr(stencil_ir, func_ir, index_def):
"""evaluate constant unary expr if possible
otherwise, raise GuardException
"""
require(isinstance(index_def, ir.Expr) and index_def.op == 'unary')
inner_var = index_def.value
# return -c as constant
const_val = _get_const_index_expr_inner(stencil_ir, func_ir, inner_var)
op = OPERATORS_TO_BUILTINS[index_def.fn]
return eval("{}{}".format(op, const_val))
def _get_const_index_expr_inner(stencil_ir, func_ir, index_var):
"""inner constant inference function that calls constant, unary and binary
cases.
"""
require(isinstance(index_var, ir.Var))
# case where the index is a const itself in outer function
var_const = guard(_get_const_two_irs, stencil_ir, func_ir, index_var)
if var_const is not None:
return var_const
# get index definition
index_def = ir_utils.get_definition(stencil_ir, index_var)
# match inner_var = unary(index_var)
var_const = guard(_get_const_unary_expr, stencil_ir, func_ir, index_def)
if var_const is not None:
return var_const
# match inner_var = arg1 + arg2
var_const = guard(_get_const_binary_expr, stencil_ir, func_ir, index_def)
if var_const is not None:
return var_const
raise GuardException
def is_const_slice(typemap, func_ir, var, accept_stride=False):
""" return True if var can be determined to be a constant size slice """
require(typemap[var.name] == types.slice2_type
or (accept_stride and typemap[var.name] == types.slice3_type))
call_expr = get_definition(func_ir, var)
require(isinstance(call_expr, ir.Expr) and call_expr.op == 'call')
assert (len(call_expr.args) == 2
or (accept_stride and len(call_expr.args) == 3))
assert find_callname(func_ir, call_expr) == ('slice', 'builtins')
arg0_def = get_definition(func_ir, call_expr.args[0])
require(isinstance(arg0_def, ir.Const) and arg0_def.value is None)
size_const = find_const(func_ir, call_expr.args[1])
require(isinstance(size_const, int))
return True
def find_build_tuple(func_ir, var):
"""Check if a variable is constructed via build_tuple
and return the sequence or raise GuardException otherwise.
"""
# variable or variable name
require(isinstance(var, (ir.Var, str)))
var_def = get_definition(func_ir, var)
require(isinstance(var_def, ir.Expr))
require(var_def.op == 'build_tuple')
return var_def.items
def find_str_const(func_ir, var):
"""Check if a variable can be inferred as a string constant, and return
the constant value, or raise GuardException otherwise.
"""
require(isinstance(var, ir.Var))
var_def = get_definition(func_ir, var)
if isinstance(var_def, ir.Const):
val = var_def.value
require(isinstance(val, str))
return val
# only add supported (s1+s2), TODO: extend to other expressions
require(isinstance(var_def, ir.Expr) and var_def.op == 'binop'
and var_def.fn == operator.add)
arg1 = find_str_const(func_ir, var_def.lhs)
arg2 = find_str_const(func_ir, var_def.rhs)
return arg1 + arg2
def get_slice_step(typemap, func_ir, var):
require(typemap[var.name] == types.slice3_type)
call_expr = get_definition(func_ir, var)
require(isinstance(call_expr, ir.Expr) and call_expr.op == 'call')
assert len(call_expr.args) == 3
return call_expr.args[2]
