def test_shape_unit_plus_symbolic(left_shape, right_shape):
symbol_table = {
'A': ast.SymbolNode(ast.NodeSymbol.ARRAY, left_shape, None, None),
'B': ast.SymbolNode(ast.NodeSymbol.ARRAY, right_shape, None, None),
'n': ast.SymbolNode(ast.NodeSymbol.ARRAY, (), None, None),
}
tree = ast.Node((ast.NodeSymbol.PLUS,), None, (), (
ast.Node((ast.NodeSymbol.ARRAY,), None, ('A',), ()),
ast.Node((ast.NodeSymbol.ARRAY,), None, ('B',), ()),))
expected_tree = ast.Node((ast.NodeSymbol.PLUS,), (3, 4, 5), (), (
ast.Node((ast.NodeSymbol.ARRAY,), left_shape, ('A',), ()),
ast.Node((ast.NodeSymbol.ARRAY,), right_shape, ('B',), ()),))
context = ast.create_context(ast=tree, symbol_table=symbol_table)
expected_context = ast.create_context(ast=expected_tree, symbol_table={
**symbol_table,
'_a3': ast.SymbolNode(ast.NodeSymbol.ARRAY, (), None, (5,)),
})
context_copy = copy.deepcopy(context)
new_context = shape.calculate_shapes(context)
exclude_condition_node = ast.select_node(new_context, (1,))
print(new_context.symbol_table)
testing.assert_context_equal(context, context_copy)
testing.assert_context_equal(expected_context, exclude_condition_node)
def test_shape_unit(symbol_table, tree, shape_symbol_table, shape_tree):
context = ast.create_context(ast=tree, symbol_table=symbol_table)
expected_context = ast.create_context(ast=shape_tree, symbol_table=shape_symbol_table)
context_copy = copy.deepcopy(context)
new_context = shape.calculate_shapes(context)
testing.assert_context_equal(context, context_copy)
testing.assert_context_equal(new_context, expected_context)
def test_ast_replace_node_shape(node, replacement_shape, selection,
result_node):
context = ast.create_context(ast=node)
result_context = ast.create_context(ast=result_node)
testing.assert_context_equal(
ast.replace_node_shape(context, replacement_shape, selection),
result_context)
def test_ast_select_node(selection, result_node):
node = ast.Node((ast.NodeSymbol.PLUS, ), None, (), (
ast.Node((ast.NodeSymbol.ARRAY, ), None, (0, ), ()),
ast.Node((ast.NodeSymbol.TRANSPOSE, ), None, (1, ), (ast.Node(
(ast.NodeSymbol.ARRAY, ), None, (2, ), ()), )),
))
context = ast.create_context(ast=node)
result_context = ast.create_context(ast=result_node)
testing.assert_context_equal(ast.select_node(context, selection),
result_context)
def test_shape_unit_reduce_plus_minus_multiply_divide_no_symbol(operation):
symbol_table = {
'A': ast.SymbolNode(ast.NodeSymbol.ARRAY, (1, 2, 3), None, None),
}
tree = ast.Node((ast.NodeSymbol.REDUCE, operation), None, (), (
ast.Node((ast.NodeSymbol.ARRAY,), None, ('A',), ()),))
expected_tree = ast.Node((ast.NodeSymbol.REDUCE, operation), (2, 3), (), (
ast.Node((ast.NodeSymbol.ARRAY,), (1, 2, 3), ('A',), ()),))
context = ast.create_context(ast=tree, symbol_table=symbol_table)
expected_context = ast.create_context(ast=expected_tree, symbol_table=symbol_table)
context_copy = copy.deepcopy(context)
new_context = shape.calculate_shapes(context)
testing.assert_context_equal(context, context_copy)
testing.assert_context_equal(expected_context, new_context)
def test_array_single_array():
expression = LazyArray(name='A', shape=(2, 3))
node = ast.Node((ast.NodeSymbol.ARRAY,), None, ('A',), ())
symbol_table = {'A': ast.SymbolNode(ast.NodeSymbol.ARRAY, (2, 3), None, None)}
context = ast.create_context(ast=node, symbol_table=symbol_table)
testing.assert_context_equal(context, expression.context)
def test_dimension_operation():
symbol_table = {'_a1': ast.SymbolNode(ast.NodeSymbol.ARRAY, (3, 4), None, None)}
node = ast.Node((ast.NodeSymbol.TRANSPOSE,), (4, 3), (), (
ast.Node((ast.NodeSymbol.ARRAY,), (3, 4), ('_a1',), ()),))
context = ast.create_context(ast=node, symbol_table=symbol_table)
assert shape.dimension(context) == 2
def test_array_transpose_default():
expression = LazyArray(name='A', shape=(2, 3)).transpose()
node = ast.Node((ast.NodeSymbol.TRANSPOSE,), None, (), (
ast.Node((ast.NodeSymbol.ARRAY,), None, ('A',), ()),))
symbol_table = {'A': ast.SymbolNode(ast.NodeSymbol.ARRAY, (2, 3), None, None)}
context = ast.create_context(ast=node, symbol_table=symbol_table)
testing.assert_context_equal(context, expression.context)
def test_ast_select_node_shape(selection, result_shape):
node = ast.Node((ast.NodeSymbol.PLUS, ), (1, 2, 3), (), (
ast.Node((ast.NodeSymbol.ARRAY, ), (4, 5, 6), (0, ), ()),
ast.Node((ast.NodeSymbol.TRANSPOSE, ), (7, 8, 9), (1, ), (ast.Node(
(ast.NodeSymbol.ARRAY, ), (10, 11, 12), (2, ), ()), )),
))
context = ast.create_context(ast=node)
assert ast.select_node_shape(context, selection) == result_shape
def test_symbol_table_unique_index():
context = ast.create_context()
context_copy = copy.deepcopy(context)
symbol_index_name_1 = ast.generate_unique_index_name(context)
new_context = ast.add_symbol(context, symbol_index_name_1,
ast.NodeSymbol.ARRAY, (3, 4), None, None)
symbol_index_name_2 = ast.generate_unique_index_name(new_context)
assert context == context_copy
assert symbol_index_name_1 != symbol_index_name_2
def test_add_symbol_table_idempodent():
context = ast.create_context()
context_copy = copy.deepcopy(context)
new_context = ast.add_symbol(context, 'A', ast.NodeSymbol.ARRAY, (3, 4),
None, None)
assert context == context_copy
assert new_context == ast.Context(
ast=None,
symbol_table={'A': (ast.NodeSymbol.ARRAY, (3, 4), None, None)})
def test_not_matches_rule_simple():
tree = ast.Node((ast.NodeSymbol.ARRAY, ), (2, 3), ('A', ), ())
symbol_table = {
'A': ast.SymbolNode(ast.NodeSymbol.ARRAY, (2, 3), None,
(1, 2, 3, 4, 5, 6))
}
context = ast.create_context(ast=tree, symbol_table=symbol_table)
rule = ((ast.NodeSymbol.PSI, ), )
assert not dnf.matches_rule(rule, context)
def test_shape_scalar_plus_minus_multiply_divide_no_symbol(operation):
symbol_table = {
'A': ast.SymbolNode(ast.NodeSymbol.ARRAY, (3, 4, 5), None, None),
'B': ast.SymbolNode(ast.NodeSymbol.ARRAY, (), None, (0,))
}
tree = ast.Node((operation,), None, (), (
ast.Node((ast.NodeSymbol.ARRAY,), None, ('A',), ()),
ast.Node((ast.NodeSymbol.ARRAY,), None, ('B',), ()),))
expected_tree = ast.Node((operation,), (3, 4, 5), (), (
ast.Node((ast.NodeSymbol.ARRAY,), (3, 4, 5), ('A',), ()),
ast.Node((ast.NodeSymbol.ARRAY,), (), ('B',), ()),))
context = ast.create_context(ast=tree, symbol_table=symbol_table)
expected_context = ast.create_context(ast=expected_tree, symbol_table=symbol_table)
context_copy = copy.deepcopy(context)
new_context = shape.calculate_shapes(context)
testing.assert_context_equal(context, context_copy)
testing.assert_context_equal(expected_context, new_context)
def test_array_reduce(symbol):
expression = LazyArray(name='A', shape=(2, 3)).reduce(symbol)
expected_tree = ast.Node((ast.NodeSymbol.REDUCE, LazyArray.OPPERATION_MAP[symbol]), None, (), (
ast.Node((ast.NodeSymbol.ARRAY,), None, ('A',), ()),))
expected_symbol_table = {
'A': ast.SymbolNode(ast.NodeSymbol.ARRAY, (2, 3), None, None),
}
expected_context = ast.create_context(ast=expected_tree, symbol_table=expected_symbol_table)
testing.assert_context_equal(expected_context, expression.context)
def test_matches_rule_nested():
tree = ast.Node((ast.NodeSymbol.TRANSPOSE, ), None, (), (ast.Node(
(ast.NodeSymbol.ARRAY, ), (2, 3), ('A', ), ()), ))
symbol_table = {
'A': ast.SymbolNode(ast.NodeSymbol.ARRAY, (2, 3), None,
(1, 2, 3, 4, 5, 6))
}
context = ast.create_context(ast=tree, symbol_table=symbol_table)
rule = ((ast.NodeSymbol.TRANSPOSE, ), (((ast.NodeSymbol.ARRAY, ), ), ))
assert dnf.matches_rule(rule, context)
def test_array_index_tuple():
expression = LazyArray(name='A', shape=(2, 3))[1, 0]
node = ast.Node((ast.NodeSymbol.PSI,), None, (), (
ast.Node((ast.NodeSymbol.ARRAY,), None, ('_a1',), ()),
ast.Node((ast.NodeSymbol.ARRAY,), None, ('A',), ()),))
symbol_table = {
'A': ast.SymbolNode(ast.NodeSymbol.ARRAY, (2, 3), None, None),
'_a1': ast.SymbolNode(ast.NodeSymbol.ARRAY, (2,), None, (1, 0)),
}
context = ast.create_context(ast=node, symbol_table=symbol_table)
testing.assert_context_equal(context, expression.context)
def test_array_addition():
expression = LazyArray(name='A', shape=(2, 3)) + LazyArray(name='B', shape=(2, 3))
tree = ast.Node((ast.NodeSymbol.PLUS,), None, (), (
ast.Node((ast.NodeSymbol.ARRAY,), None, ('A',), ()),
ast.Node((ast.NodeSymbol.ARRAY,), None, ('B',), ())))
symbol_table = {
'A': ast.SymbolNode(ast.NodeSymbol.ARRAY, (2, 3), None, None),
'B': ast.SymbolNode(ast.NodeSymbol.ARRAY, (2, 3), None, None)
}
context = ast.create_context(ast=tree, symbol_table=symbol_table)
testing.assert_context_equal(context, expression.context)
def test_array_transpose_with_vector():
expression = LazyArray(name='A', shape=(2, 3)).transpose([1, 0])
node = ast.Node((ast.NodeSymbol.TRANSPOSEV,), None, (), (
ast.Node((ast.NodeSymbol.ARRAY,), None, ('_a1',), ()),
ast.Node((ast.NodeSymbol.ARRAY,), None, ('A',), ()),))
symbol_table = {
'A': ast.SymbolNode(ast.NodeSymbol.ARRAY, (2, 3), None, None),
'_a1': ast.SymbolNode(ast.NodeSymbol.ARRAY, (2,), None, (1, 0)),
}
context = ast.create_context(ast=node, symbol_table=symbol_table)
testing.assert_context_equal(context, expression.context)
def test_array_index_stride_reverse():
expression = LazyArray(name='A', shape=(2, 3))[1:2:-1]
tree = ast.Node(ast.NodeSymbol.PSI, None, (), (
ast.Node(ast.NodeSymbol.ARRAY, None, ('_a2',), ()),
ast.Node(ast.NodeSymbol.ARRAY, None, ('A',), ())))
symbol_table = {
'A': SymbolNode(ast.NodeSymbol.ARRAY, (2, 3), None, None),
'n': SymbolNode(ast.NodeSymbol.ARRAY, (), None, None),
'_a2': SymbolNode(ast.NodeSymbol.ARRAY, (1,), None, (Node(ast.NodeSymbol.ARRAY, (), 'n'),)),
}
context = ast.create_context(ast=tree, symbol_table=symbol_table)
testing.assert_context_equal(context, expression.context)
def test_preorder_replacement():
counter = itertools.count()
def replacement_function(context):
if context.ast.shape is None:
return ast.replace_node_shape(context, (next(counter), ), ())
return None
tree = ast.Node((ast.NodeSymbol.CAT, ), None, (), (
ast.Node((ast.NodeSymbol.TRANSPOSE, ), None, (), (ast.Node(
(ast.NodeSymbol.ARRAY, ), None, (), ()), )),
ast.Node((ast.NodeSymbol.PLUS, ), None, (), (
ast.Node((ast.NodeSymbol.SHAPE, ), None, (), (ast.Node(
(ast.NodeSymbol.ARRAY, ), None, (), ()), )),
ast.Node((ast.NodeSymbol.RAV, ), None, (), (ast.Node(
(ast.NodeSymbol.ARRAY, ), None, (), ()), )),
)),
))
context = ast.create_context(ast=tree)
context_copy = copy.deepcopy(context)
expected_tree = ast.Node((ast.NodeSymbol.CAT, ), (0, ), (), (
ast.Node((ast.NodeSymbol.TRANSPOSE, ), (1, ), (), (ast.Node(
(ast.NodeSymbol.ARRAY, ), (2, ), (), ()), )),
ast.Node((ast.NodeSymbol.PLUS, ), (3, ), (), (
ast.Node((ast.NodeSymbol.SHAPE, ), (4, ), (), (ast.Node(
(ast.NodeSymbol.ARRAY, ), (5, ), (), ()), )),
ast.Node((ast.NodeSymbol.RAV, ), (6, ), (), (ast.Node(
(ast.NodeSymbol.ARRAY, ), (7, ), (), ()), )),
)),
))
expected_context = ast.create_context(ast=expected_tree)
new_context = ast.node_traversal(context,
replacement_function,
traversal='preorder')
testing.assert_context_equal(context, context_copy)
testing.assert_context_equal(new_context, expected_context)
def test_array_inner_product(left_symbol, right_symbol):
expression = LazyArray(name='A', shape=(2, 3)).inner(left_symbol, right_symbol, LazyArray(name='B', shape=(3, 4)))
expected_tree = ast.Node((ast.NodeSymbol.DOT, LazyArray.OPPERATION_MAP[left_symbol], LazyArray.OPPERATION_MAP[right_symbol]), None, (), (
ast.Node((ast.NodeSymbol.ARRAY,), None, ('A',), ()),
ast.Node((ast.NodeSymbol.ARRAY,), None, ('B',), ())))
expected_symbol_table = {
'A': ast.SymbolNode(ast.NodeSymbol.ARRAY, (2, 3), None, None),
'B': ast.SymbolNode(ast.NodeSymbol.ARRAY, (3, 4), None, None),
}
expected_context = ast.create_context(ast=expected_tree, symbol_table=expected_symbol_table)
testing.assert_context_equal(expected_context, expression.context)
def test_array_single_array_binary_operation_cast(function, side, operation):
expression = function()
if side == 'right':
tree = ast.Node((operation,), None, (), (
ast.Node((ast.NodeSymbol.ARRAY,), None, ('A',), ()),
ast.Node((ast.NodeSymbol.ARRAY,), None, ('_a1',), ())))
else:
tree = ast.Node((operation,), None, (), (
ast.Node((ast.NodeSymbol.ARRAY,), None, ('_a1',), ()),
ast.Node((ast.NodeSymbol.ARRAY,), None, ('A',), ())))
symbol_table = {
'A': ast.SymbolNode(ast.NodeSymbol.ARRAY, (2, 3), None, None),
'_a1': ast.SymbolNode(ast.NodeSymbol.ARRAY, (), None, (1,))
}
context = ast.create_context(ast=tree, symbol_table=symbol_table)
testing.assert_context_equal(context, expression.context)
def test_graphviz_visualization():
symbol_table = {
'_a1': ast.SymbolNode(ast.NodeSymbol.ARRAY, (1, ), None, (0, )),
'A': ast.SymbolNode(ast.NodeSymbol.ARRAY, (3, 4), None, None),
'B': ast.SymbolNode(ast.NodeSymbol.ARRAY, (3, 4), None, None)
}
tree = ast.Node(
(ast.NodeSymbol.PSI, ), None, (),
(ast.Node((ast.NodeSymbol.ARRAY, ), None, ('_a1', ), ()),
ast.Node((ast.NodeSymbol.TRANSPOSE, ), None, (), (ast.Node(
(ast.NodeSymbol.PLUS, ), None, (), (
ast.Node((ast.NodeSymbol.ARRAY, ), None, ('A', ), ()),
ast.Node((ast.NodeSymbol.ARRAY, ), None, ('B', ), ()),
)), ))))
context = ast.create_context(ast=tree, symbol_table=symbol_table)
visualize_ast(context)
def test_metric_flops():
tree = ast.Node((ast.NodeSymbol.PSI,), None, (), (
ast.Node((ast.NodeSymbol.ARRAY,), None, ('_a1',), ()),
ast.Node((ast.NodeSymbol.TRANSPOSE,), None, (), (
ast.Node((ast.NodeSymbol.PLUS,), None, (), (
ast.Node((ast.NodeSymbol.ARRAY,), None, ('A',), ()),
ast.Node((ast.NodeSymbol.ARRAY,), None, ('B',), ()))),))))
symbol_table = {
'_a1': ast.SymbolNode(ast.NodeSymbol.ARRAY, (1,), None, (0,)),
'A': ast.SymbolNode(ast.NodeSymbol.ARRAY, (10, 100), None, None),
'B': ast.SymbolNode(ast.NodeSymbol.ARRAY, (10, 100), None, None)
}
context = ast.create_context(ast=tree, symbol_table=symbol_table)
context = shape.calculate_shapes(context)
assert analysis.metric_flops(context) == 1000
context = dnf.reduce_to_dnf(context)
assert analysis.metric_flops(context) == 10
def test_is_not_vector_2d(): # 2D array
symbol_table = {'A': ast.SymbolNode(ast.NodeSymbol.ARRAY, (5, 1), None, (1, 2, 3, 4, 5))}
node = ast.Node((ast.NodeSymbol.ARRAY,), None, ('A',), ())
context = ast.create_context(ast=node, symbol_table=symbol_table)
assert not shape.is_vector(context)
def test_is_not_vector_1d(): # scalar
symbol_table = {'asdf': ast.SymbolNode(ast.NodeSymbol.ARRAY, (), None, (3,))}
node = ast.Node((ast.NodeSymbol.ARRAY,), None, ('asdf',), ())
context = ast.create_context(ast=node, symbol_table=symbol_table)
assert not shape.is_vector(context)
def test_ast_num_node_children(node, result):
context = ast.create_context(ast=node)
assert ast.num_node_children(context) == result
def test_join_symbol_tables_simple():
left_tree = ast.Node((ast.NodeSymbol.PLUS, ), None, (), (
ast.Node((ast.NodeSymbol.ARRAY, ), None, ('A', ), ()),
ast.Node((ast.NodeSymbol.ARRAY, ), None, ('B', ), ()),
))
left_symbol_table = {
'A':
ast.SymbolNode(ast.NodeSymbol.ARRAY, (3, 4), None, None),
'_a1':
ast.SymbolNode(ast.NodeSymbol.ARRAY,
(1, ast.Node((ast.NodeSymbol.ARRAY, ), (), ('m', ),
())), None, None),
'B':
ast.SymbolNode(ast.NodeSymbol.ARRAY, (2, 4), None, None)
}
right_tree = ast.Node((ast.NodeSymbol.MINUS, ), None, (), (
ast.Node((ast.NodeSymbol.ARRAY, ), None, ('A', ), ()),
ast.Node((ast.NodeSymbol.ARRAY, ), None, ('_a3', ), ()),
))
right_symbol_table = {
'A':
ast.SymbolNode(ast.NodeSymbol.ARRAY, (3, 4), None, None),
'_a3':
ast.SymbolNode(ast.NodeSymbol.ARRAY, (ast.Node(
(ast.NodeSymbol.ARRAY, ), (), ('_a10', ),
()), ast.Node((ast.NodeSymbol.ARRAY, ), (), ('m', ), ())), None,
None),
'_a10':
ast.SymbolNode(ast.NodeSymbol.ARRAY, (), None, (1, )),
'm':
ast.SymbolNode(ast.NodeSymbol.ARRAY, (), None, None),
'B':
ast.SymbolNode(ast.NodeSymbol.ARRAY, (2, 4), None, None),
'n':
ast.SymbolNode(ast.NodeSymbol.ARRAY, (), None, None),
}
left_context = ast.create_context(ast=left_tree,
symbol_table=left_symbol_table)
right_context = ast.create_context(ast=right_tree,
symbol_table=right_symbol_table)
new_symbol_table, new_left_context, new_right_context = ast.join_symbol_tables(
left_context, right_context)
assert new_symbol_table == {
'A':
ast.SymbolNode(ast.NodeSymbol.ARRAY, (3, 4), None, None),
'B':
ast.SymbolNode(ast.NodeSymbol.ARRAY, (2, 4), None, None),
'm':
ast.SymbolNode(ast.NodeSymbol.ARRAY, (), None, None),
'_a0':
ast.SymbolNode(ast.NodeSymbol.ARRAY, (), None, (1, )),
'_a1':
ast.SymbolNode(ast.NodeSymbol.ARRAY, (ast.Node(
(ast.NodeSymbol.ARRAY, ), (), ('_a0', ),
()), ast.Node((ast.NodeSymbol.ARRAY, ), (), ('m', ), ())), None,
None),
}
expected_left_context = ast.create_context(
ast=ast.Node((ast.NodeSymbol.PLUS, ), None, (),
(
ast.Node((ast.NodeSymbol.ARRAY, ), None, ('A', ), ()),
ast.Node((ast.NodeSymbol.ARRAY, ), None, ('B', ), ()),
)),
symbol_table={
'A': ast.SymbolNode(ast.NodeSymbol.ARRAY, (3, 4), None, None),
'B': ast.SymbolNode(ast.NodeSymbol.ARRAY, (2, 4), None, None)
})
expected_right_context = ast.create_context(
ast=ast.Node(
(ast.NodeSymbol.MINUS, ), None, (), (
ast.Node((ast.NodeSymbol.ARRAY, ), None, ('A', ), ()),
ast.Node((ast.NodeSymbol.ARRAY, ), None, ('_a1', ), ()),
)),
symbol_table={
'A':
ast.SymbolNode(ast.NodeSymbol.ARRAY, (3, 4), None, None),
'm':
ast.SymbolNode(ast.NodeSymbol.ARRAY, (), None, None),
'_a0':
ast.SymbolNode(ast.NodeSymbol.ARRAY, (), None, (1, )),
'_a1':
ast.SymbolNode(ast.NodeSymbol.ARRAY, (ast.Node(
(ast.NodeSymbol.ARRAY, ), (), ('_a0', ),
()), ast.Node((ast.NodeSymbol.ARRAY, ), (), ('m', ), ())),
None, None)
})
testing.assert_context_equal(new_left_context, expected_left_context)
testing.assert_context_equal(new_right_context, expected_right_context)
def test_ast_nodes(node, result):
context = ast.create_context(ast=node)
assert result == (ast.is_array(context), ast.is_operation(context),
ast.is_unary_operation(context),
ast.is_binary_operation(context))
def test_dimension_array():
symbol_table = {'_a1': ast.SymbolNode(ast.NodeSymbol.ARRAY, (), None, (3,))}
node = ast.Node((ast.NodeSymbol.ARRAY,), None, ('_a1',), ())
context = ast.create_context(ast=node, symbol_table=symbol_table)
assert shape.dimension(context) == 0
