From 1e7a179a846a240f0c73c23f89df489dd7edd1ed Mon Sep 17 00:00:00 2001
From: Stephan Seitz <stephan.seitz@fau.de>
Date: Sat, 18 Jan 2020 15:53:32 +0100
Subject: [PATCH] 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(...))`
---
pystencils/__init__.py | 4 ++--
pystencils/display_utils.py | 20 ++++++++++++++++---
pystencils/kernel_wrapper.py | 2 +-
pystencils/test_type_interference.py | 3 +--
pystencils_tests/test_address_of.py | 11 ++++------
pystencils_tests/test_complex_numbers.py | 4 ++--
.../test_conditional_field_access.py | 2 +-
pystencils_tests/test_cuda_known_functions.py | 7 +++----
pystencils_tests/test_custom_backends.py | 6 +++---
pystencils_tests/test_fast_approximation.py | 6 +++---
pystencils_tests/test_interpolation.py | 18 ++++++++---------
pystencils_tests/test_jacobi_cbackend.py | 4 ++--
pystencils_tests/test_jacobi_llvm.py | 2 +-
pystencils_tests/test_opencl.py | 12 +++++------
pystencils_tests/test_source_code_comment.py | 3 ++-
pystencils_tests/test_sum_prod.py | 4 +---
16 files changed, 57 insertions(+), 51 deletions(-)
diff --git a/pystencils/__init__.py b/pystencils/__init__.py
index ed08ff132..b9002ddbb 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, 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',
diff --git a/pystencils/display_utils.py b/pystencils/display_utils.py
index 610c404b7..6d817e301 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(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)
diff --git a/pystencils/kernel_wrapper.py b/pystencils/kernel_wrapper.py
index 0e327711e..94c2a7c24 100644
--- a/pystencils/kernel_wrapper.py
+++ b/pystencils/kernel_wrapper.py
@@ -16,4 +16,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 0daa6f9d2..953b87742 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 b4c46c1c6..717f6e437 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 9f9164640..2fb558cbc 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 f68b34679..a4bd53228 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 fe33b98df..c610d9e85 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 1ff775d41..6b0fc6fd1 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 76bef174e..6c9539f64 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 dd4ee06f4..4ead7a0fc 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 9206aeeda..9a2fb18f6 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
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_()
diff --git a/pystencils_tests/test_jacobi_llvm.py b/pystencils_tests/test_jacobi_llvm.py
index 13b22aa4b..95bd52d2e 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 ea5b1ecbc..42c533367 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 6cc66feb6..16e6e5a96 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 4fa5c0618..356966918 100644
--- a/pystencils_tests/test_sum_prod.py
+++ b/pystencils_tests/test_sum_prod.py
@@ -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)
--
GitLab