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,10 +35,10 @@ 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.
Can either be displayed as HTML in Jupyter notebooks or printed as normal string.
"""
from pystencils.backends.cbackend import generate_c
......@@ -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