Commit 1e7a179a authored by Stephan Seitz's avatar Stephan Seitz
Browse files

Fix #6: let show_code display/print code

- `get_code` will keep the previous behavior of `show_code`
- `get_code_str` will be equivalent to previous `str(show_code(...))`
parent f16225d5
Pipeline #21146 failed with stage
in 7 minutes and 44 seconds
......@@ -5,7 +5,7 @@ from . import stencil as stencil
from .assignment import Assignment, assignment_from_stencil
from .data_types import TypedSymbol
from .datahandling import create_data_handling
from .display_utils import show_code, to_dot
from .display_utils import show_code, get_code, get_code_str, to_dot
from .field import Field, FieldType, fields
from .kernel_decorator import kernel
from .kernelcreation import create_indexed_kernel, create_kernel, create_staggered_kernel
......@@ -25,7 +25,7 @@ __all__ = ['Field', 'FieldType', 'fields',
'TypedSymbol',
'make_slice',
'create_kernel', 'create_indexed_kernel', 'create_staggered_kernel',
'show_code', 'to_dot',
'show_code', 'to_dot', 'get_code', 'get_code_str',
'AssignmentCollection',
'Assignment',
'assignment_from_stencil',
......
from typing import Any, Dict, Optional
from typing import Any, Dict, Optional, Union
import sympy as sp
......@@ -35,7 +35,7 @@ def highlight_cpp(code: str):
return HTML(highlight(code, CppLexer(), HtmlFormatter()))
def show_code(ast: KernelFunction, custom_backend=None):
def get_code(ast: Union[KernelFunction, KernelWrapper], custom_backend=None):
"""Returns an object to display generated code (C/C++ or CUDA)
Can either be displayed as HTML in Jupyter notebooks or printed as normal string.
......@@ -65,3 +65,17 @@ def show_code(ast: KernelFunction, custom_backend=None):
def __repr__(self):
return generate_c(self.ast, dialect=dialect, custom_backend=custom_backend)
return CodeDisplay(ast)
def get_code_str(ast, custom_backend=None):
return str(get_code(ast, custom_backend))
def show_code(ast: Union[KernelFunction, KernelWrapper], custom_backend=None):
code = get_code(ast, custom_backend)
try:
from IPython.display import display
display(code)
except Exception:
print(code)
......@@ -16,4 +16,4 @@ class KernelWrapper:
@property
def code(self):
return str(pystencils.show_code(self.ast))
return pystencils.get_code_str(self.ast)
......@@ -18,8 +18,7 @@ def test_type_interference():
ast = pystencils.create_kernel(assignments)
code = str(pystencils.show_code(ast))
print(code)
code = str(pystencils.get_code_str(ast))
assert 'double a' in code
assert 'uint16_t b' in code
assert 'uint16_t f' in code
......
......@@ -17,16 +17,14 @@ def test_address_of():
}, {})
ast = pystencils.create_kernel(assignments)
code = pystencils.show_code(ast)
print(code)
pystencils.show_code(ast)
assignments = pystencils.AssignmentCollection({
y[0, 0]: cast_func(address_of(x[0, 0]), create_type('int64'))
}, {})
ast = pystencils.create_kernel(assignments)
code = pystencils.show_code(ast)
print(code)
pystencils.show_code(ast)
def test_address_of_with_cse():
......@@ -39,9 +37,8 @@ def test_address_of_with_cse():
}, {})
ast = pystencils.create_kernel(assignments)
code = pystencils.show_code(ast)
pystencils.show_code(ast)
assignments_cse = sympy_cse(assignments)
ast = pystencils.create_kernel(assignments_cse)
code = pystencils.show_code(ast)
print(code)
pystencils.show_code(ast)
......@@ -52,7 +52,7 @@ def test_complex_numbers(assignment, scalar_dtypes, target):
ast = pystencils.create_kernel(assignment,
target=target,
data_type=scalar_dtypes)
code = str(pystencils.show_code(ast))
code = pystencils.get_code_str(ast)
print(code)
assert "Not supported" not in code
......@@ -95,7 +95,7 @@ def test_complex_numbers_64(assignment, target):
ast = pystencils.create_kernel(assignment,
target=target,
data_type='double')
code = str(pystencils.show_code(ast))
code = pystencils.get_code_str(ast)
print(code)
assert "Not supported" not in code
......
......@@ -63,7 +63,7 @@ def test_boundary_check(with_cse):
print(assignments)
kernel_checked = ps.create_kernel(assignments, ghost_layers=0).compile()
print(ps.show_code(kernel_checked))
ps.show_code(kernel_checked)
# No SEGFAULT, please!!
kernel_checked(f=f_arr, g=g_arr)
......@@ -18,7 +18,7 @@ def test_cuda_known_functions():
})
ast = pystencils.create_kernel(assignments, 'gpu')
print(pystencils.show_code(ast))
pystencils.show_code(ast)
kernel = ast.compile()
assert(kernel is not None)
......@@ -32,7 +32,7 @@ def test_cuda_but_not_c():
})
ast = pystencils.create_kernel(assignments, 'cpu')
print(pystencils.show_code(ast))
pystencils.show_code(ast)
def test_cuda_unknown():
......@@ -43,5 +43,4 @@ def test_cuda_unknown():
})
ast = pystencils.create_kernel(assignments, 'gpu')
code = str(pystencils.show_code(ast))
print(code)
pystencils.show_code(ast)
from subprocess import CalledProcessError
import pycuda.driver
import pytest
import sympy
import pycuda.driver
import pystencils
import pystencils.cpu.cpujit
import pystencils.gpucuda.cudajit
......@@ -32,11 +32,11 @@ def test_custom_backends():
z[0, 0], x[0, 0] * sympy.log(x[0, 0] * y[0, 0]))], [])
ast = pystencils.create_kernel(normal_assignments, target='cpu')
print(pystencils.show_code(ast, ScreamingBackend()))
pystencils.show_code(ast, ScreamingBackend())
with pytest.raises(CalledProcessError):
pystencils.cpu.cpujit.make_python_function(ast, custom_backend=ScreamingBackend())
ast = pystencils.create_kernel(normal_assignments, target='gpu')
print(pystencils.show_code(ast, ScreamingGpuBackend()))
pystencils.show_code(ast, ScreamingGpuBackend())
with pytest.raises(pycuda.driver.CompileError):
pystencils.gpucuda.cudajit.make_python_function(ast, custom_backend=ScreamingGpuBackend())
......@@ -12,7 +12,7 @@ def test_fast_sqrt():
assert len(insert_fast_sqrts(expr).atoms(fast_sqrt)) == 1
assert len(insert_fast_sqrts([expr])[0].atoms(fast_sqrt)) == 1
ast = ps.create_kernel(ps.Assignment(g[0, 0], insert_fast_sqrts(expr)), target='gpu')
code_str = str(ps.show_code(ast))
code_str = ps.get_code_str(ast)
assert '__fsqrt_rn' in code_str
expr = ps.Assignment(sp.Symbol("tmp"), 3 / sp.sqrt(f[0, 0] + f[1, 0]))
......@@ -21,7 +21,7 @@ def test_fast_sqrt():
ac = ps.AssignmentCollection([expr], [])
assert len(insert_fast_sqrts(ac).main_assignments[0].atoms(fast_inv_sqrt)) == 1
ast = ps.create_kernel(insert_fast_sqrts(ac), target='gpu')
code_str = str(ps.show_code(ast))
code_str = ps.get_code_str(ast)
assert '__frsqrt_rn' in code_str
......@@ -34,5 +34,5 @@ def test_fast_divisions():
assert len(insert_fast_divisions(expr).atoms(fast_division)) == 1
ast = ps.create_kernel(ps.Assignment(g[0, 0], insert_fast_divisions(expr)), target='gpu')
code_str = str(ps.show_code(ast))
code_str = ps.get_code_str(ast)
assert '__fdividef' in code_str
......@@ -11,11 +11,11 @@ import itertools
from os.path import dirname, join
import numpy as np
import pycuda.autoinit # NOQA
import pycuda.gpuarray as gpuarray
import pytest
import sympy
import pycuda.autoinit # NOQA
import pycuda.gpuarray as gpuarray
import pystencils
from pystencils.interpolation_astnodes import LinearInterpolator
from pystencils.spatial_coordinates import x_, y_
......@@ -51,7 +51,7 @@ def test_interpolation():
print(assignments)
ast = pystencils.create_kernel(assignments)
print(ast)
print(pystencils.show_code(ast))
pystencils.show_code(ast)
kernel = ast.compile()
pyconrad.imshow(lenna)
......@@ -71,7 +71,7 @@ def test_scale_interpolation():
print(assignments)
ast = pystencils.create_kernel(assignments)
print(ast)
print(pystencils.show_code(ast))
pystencils.show_code(ast)
kernel = ast.compile()
out = np.zeros_like(lenna)
......@@ -102,7 +102,7 @@ def test_rotate_interpolation(address_mode):
print(assignments)
ast = pystencils.create_kernel(assignments)
print(ast)
print(pystencils.show_code(ast))
pystencils.show_code(ast)
kernel = ast.compile()
out = np.zeros_like(lenna)
......@@ -135,7 +135,7 @@ def test_rotate_interpolation_gpu(address_mode):
print(assignments)
ast = pystencils.create_kernel(assignments, target='gpu', use_textures_for_interpolation=use_textures)
print(ast)
print(pystencils.show_code(ast))
pystencils.show_code(ast)
kernel = ast.compile()
out = gpuarray.zeros_like(lenna_gpu)
......@@ -182,7 +182,7 @@ def test_shift_interpolation_gpu(address_mode):
# print(assignments)
ast = pystencils.create_kernel(assignments, target='gpu', use_textures_for_interpolation=use_textures)
# print(ast)
print(pystencils.show_code(ast))
pystencils.show_code(ast)
kernel = ast.compile()
out = gpuarray.zeros_like(lenna_gpu)
......@@ -209,7 +209,7 @@ def test_rotate_interpolation_size_change(address_mode):
print(assignments)
ast = pystencils.create_kernel(assignments)
print(ast)
print(pystencils.show_code(ast))
pystencils.show_code(ast)
kernel = ast.compile()
out = np.zeros((100, 150), np.float64)
......@@ -228,7 +228,7 @@ def test_field_interpolated(address_mode, target):
print(assignments)
ast = pystencils.create_kernel(assignments)
print(ast)
print(pystencils.show_code(ast))
pystencils.show_code(ast)
kernel = ast.compile()
out = np.zeros_like(lenna)
......
import numpy as np
from pystencils import show_code
from pystencils import get_code
from pystencils.astnodes import Block, KernelFunction, SympyAssignment
from pystencils.cpu import make_python_function
from pystencils.field import Field
......@@ -36,7 +36,7 @@ def test_jacobi_fixed_field_size():
error = np.sum(np.abs(dst_field_py - dst_field_c))
np.testing.assert_allclose(error, 0.0, atol=1e-13)
code_display = show_code(ast_node)
code_display = get_code(ast_node)
assert 'for' in str(code_display)
assert 'for' in code_display._repr_html_()
......
......@@ -52,7 +52,7 @@ def test_jacobi_fixed_field_size_gpu():
jacobi = Assignment(d[0, 0], (f[1, 0] + f[-1, 0] + f[0, 1] + f[0, -1]) / 4)
ast = create_kernel([jacobi], target='gpu')
print(show_code(ast))
show_code(ast)
for x in range(1, size[0] - 1):
for y in range(1, size[1] - 1):
......
......@@ -27,10 +27,9 @@ def test_print_opencl():
print(ast)
code = pystencils.show_code(ast, custom_backend=CudaBackend())
print(code)
pystencils.show_code(ast, custom_backend=CudaBackend())
opencl_code = pystencils.show_code(ast, custom_backend=OpenClBackend())
opencl_code = pystencils.get_code_str(ast, custom_backend=OpenClBackend())
print(opencl_code)
assert "__global double * RESTRICT const _data_x" in str(opencl_code)
......@@ -108,10 +107,9 @@ def test_opencl_jit():
print(ast)
code = pystencils.show_code(ast, custom_backend=CudaBackend())
print(code)
opencl_code = pystencils.show_code(ast, custom_backend=OpenClBackend())
print(opencl_code)
pystencils.show_code(ast, custom_backend=CudaBackend())
pystencils.show_code(ast, custom_backend=OpenClBackend())
cuda_kernel = ast.compile()
assert cuda_kernel is not None
......
......@@ -27,4 +27,5 @@ def test_source_code_comment():
print(ast)
compiled = ast.compile()
assert compiled is not None
print(pystencils.show_code(ast))
pystencils.show_code(ast)
......@@ -120,11 +120,9 @@ def test_prod_var_limit():
})
ast = pystencils.create_kernel(assignments)
code = str(pystencils.show_code(ast))
pystencils.show_code(ast)
kernel = ast.compile()
print(code)
array = np.zeros((10,), np.int64)
kernel(x=array, limit=100)
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment