diff --git a/pystencils/__init__.py b/pystencils/__init__.py
index ed08ff1327c68a1a116edb21547a86f8b27a8c81..18679382f0014ba12844c72fa965a3ad4fe6e196 100644
--- a/pystencils/__init__.py
+++ b/pystencils/__init__.py
@@ -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_obj, 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_obj', 'get_code_str',
'AssignmentCollection',
'Assignment',
'assignment_from_stencil',
diff --git a/pystencils/display_utils.py b/pystencils/display_utils.py
index 610c404b709885e7445eba3be412f6811fca0241..0bdad7ac16ba6c600cd66e5c82721c8ac46bb0dd 100644
--- a/pystencils/display_utils.py
+++ b/pystencils/display_utils.py
@@ -1,4 +1,4 @@
-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_obj(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_obj(ast, custom_backend))
+
+
+def show_code(ast: Union[KernelFunction, KernelWrapper], custom_backend=None):
+ code = get_code_obj(ast, custom_backend)
+
+ try:
+ from IPython.display import display
+ display(code)
+ except Exception:
+ print(code)
diff --git a/pystencils/kernel_wrapper.py b/pystencils/kernel_wrapper.py
index 3494b52a9fd060bbccbfe493f165dcc7d63c8c06..b76ec4e3d79c1b345c9f0a8b2bd8a97b59866b8d 100644
--- a/pystencils/kernel_wrapper.py
+++ b/pystencils/kernel_wrapper.py
@@ -19,4 +19,4 @@ class KernelWrapper:
@property
def code(self):
- return str(pystencils.show_code(self.ast))
+ return pystencils.get_code_str(self.ast)
diff --git a/pystencils/test_type_interference.py b/pystencils/test_type_interference.py
index 0daa6f9d2a36948184d537809457b4aaf8001d29..953b87742304b2d629a1bd564fc23e0982d4f6d9 100644
--- a/pystencils/test_type_interference.py
+++ b/pystencils/test_type_interference.py
@@ -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
diff --git a/pystencils_tests/test_address_of.py b/pystencils_tests/test_address_of.py
index b4c46c1c6045520e30826fe84ac91ae76339ad1b..717f6e43720a4ddddfca8c7a3fce3ff9e15982ff 100644
--- a/pystencils_tests/test_address_of.py
+++ b/pystencils_tests/test_address_of.py
@@ -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)
diff --git a/pystencils_tests/test_complex_numbers.py b/pystencils_tests/test_complex_numbers.py
index 9f9164640f50e417348ee635af620dd891249d77..2fb558cbc883ca89f81667ea3d03ffd46eafc3bf 100644
--- a/pystencils_tests/test_complex_numbers.py
+++ b/pystencils_tests/test_complex_numbers.py
@@ -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
diff --git a/pystencils_tests/test_conditional_field_access.py b/pystencils_tests/test_conditional_field_access.py
index f68b34679ae5de0996a09a30921c85ec49bece58..a4bd53228476ea49f977e08f71acfd1d596231fe 100644
--- a/pystencils_tests/test_conditional_field_access.py
+++ b/pystencils_tests/test_conditional_field_access.py
@@ -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)
diff --git a/pystencils_tests/test_cuda_known_functions.py b/pystencils_tests/test_cuda_known_functions.py
index fe33b98df27cc6a6c2ec77eda49e704b9afbec5e..c610d9e8517f25352a1853448f35889da16f1c23 100644
--- a/pystencils_tests/test_cuda_known_functions.py
+++ b/pystencils_tests/test_cuda_known_functions.py
@@ -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)
diff --git a/pystencils_tests/test_custom_backends.py b/pystencils_tests/test_custom_backends.py
index 1ff775d41355b79c1714a935125f27c261045011..6b0fc6fd1307426bd2b2a88549570f8e2fca2a07 100644
--- a/pystencils_tests/test_custom_backends.py
+++ b/pystencils_tests/test_custom_backends.py
@@ -1,9 +1,9 @@
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())
diff --git a/pystencils_tests/test_fast_approximation.py b/pystencils_tests/test_fast_approximation.py
index 76bef174eb7451234dcf347818fdb2d456ce7005..6c9539f64ee2a911123a8fd9f10845a3c67dcbea 100644
--- a/pystencils_tests/test_fast_approximation.py
+++ b/pystencils_tests/test_fast_approximation.py
@@ -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
diff --git a/pystencils_tests/test_interpolation.py b/pystencils_tests/test_interpolation.py
index dd4ee06f4ba9197111a04af572df6cdaf2695ab5..4ead7a0fc2b489eb874abc7e9fde8cb043d5d770 100644
--- a/pystencils_tests/test_interpolation.py
+++ b/pystencils_tests/test_interpolation.py
@@ -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)
diff --git a/pystencils_tests/test_jacobi_cbackend.py b/pystencils_tests/test_jacobi_cbackend.py
index 9206aeedaf4bb035444f9280cf58968c9ca385c9..6d86ecb05f920f7a6c54b018ea5dd3508172b727 100644
--- a/pystencils_tests/test_jacobi_cbackend.py
+++ b/pystencils_tests/test_jacobi_cbackend.py
@@ -1,6 +1,6 @@
import numpy as np
-from pystencils import show_code
+from pystencils import get_code_obj
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_obj(ast_node)
assert 'for' in str(code_display)
assert 'for' in code_display._repr_html_()
diff --git a/pystencils_tests/test_jacobi_llvm.py b/pystencils_tests/test_jacobi_llvm.py
index 13b22aa4b0235bb98b4789e3b9266bf672985111..95bd52d2edd706fc3e8b6c595c7ad9855ae5e4f6 100644
--- a/pystencils_tests/test_jacobi_llvm.py
+++ b/pystencils_tests/test_jacobi_llvm.py
@@ -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):
diff --git a/pystencils_tests/test_opencl.py b/pystencils_tests/test_opencl.py
index ea5b1ecbc83f21672731b2ecbcb54c78faa1603b..42c53336749c11c02bff612522c53905c4e2d4c5 100644
--- a/pystencils_tests/test_opencl.py
+++ b/pystencils_tests/test_opencl.py
@@ -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
diff --git a/pystencils_tests/test_source_code_comment.py b/pystencils_tests/test_source_code_comment.py
index 6cc66feb6f7f92174029e274feee22dfc4455b8d..16e6e5a9647843db4483b951aa631b591240dac5 100644
--- a/pystencils_tests/test_source_code_comment.py
+++ b/pystencils_tests/test_source_code_comment.py
@@ -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)
diff --git a/pystencils_tests/test_sum_prod.py b/pystencils_tests/test_sum_prod.py
index 4fa5c0618612b013edda4d164dd035dafdd2438a..4b4cd7131ad4b860c754865cd35f3960399e9b5a 100644
--- a/pystencils_tests/test_sum_prod.py
+++ b/pystencils_tests/test_sum_prod.py
@@ -30,7 +30,7 @@ def test_sum():
})
ast = pystencils.create_kernel(assignments)
- code = str(pystencils.show_code(ast))
+ code = str(pystencils.get_code_obj(ast))
kernel = ast.compile()
print(code)
@@ -58,11 +58,11 @@ def test_sum_use_float():
})
ast = pystencils.create_kernel(assignments, data_type=create_type('float32'))
- code = str(pystencils.show_code(ast))
+ code = str(pystencils.get_code_obj(ast))
kernel = ast.compile()
print(code)
- print(pystencils.show_code(ast))
+ print(pystencils.get_code_obj(ast))
assert 'float sum' in code
array = np.zeros((10,), np.float32)
@@ -89,7 +89,7 @@ def test_product():
})
ast = pystencils.create_kernel(assignments)
- code = str(pystencils.show_code(ast))
+ code = pystencils.get_code_str(ast)
kernel = ast.compile()
print(code)
@@ -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)
diff --git a/pystencils_tests/test_sympy_optimizations.py b/pystencils_tests/test_sympy_optimizations.py
index 745b936ea7d0e3ce5da6c684dee42c56c09dd835..0899d5b10cd87ff1e6acc630c8c88ea865005be5 100644
--- a/pystencils_tests/test_sympy_optimizations.py
+++ b/pystencils_tests/test_sympy_optimizations.py
@@ -18,7 +18,7 @@ def test_sympy_optimizations():
assignments = optimize_assignments(assignments, optims_pystencils_cpu)
ast = pystencils.create_kernel(assignments, target=target)
- code = str(pystencils.show_code(ast))
+ code = pystencils.get_code_str(ast)
assert 'expm1(' in code
@@ -35,7 +35,7 @@ def test_evaluate_constant_terms():
assignments = optimize_assignments(assignments, optims_pystencils_cpu)
ast = pystencils.create_kernel(assignments, target=target)
- code = str(pystencils.show_code(ast))
+ code = pystencils.get_code_str(ast)
assert 'cos(' not in code
print(code)
@@ -55,6 +55,6 @@ def test_do_not_evaluate_constant_terms():
optimize_assignments(assignments, optimizations)
ast = pystencils.create_kernel(assignments, target=target)
- code = str(pystencils.show_code(ast))
+ code = pystencils.get_code_str(ast)
assert 'cos(' in code
print(code)