From 130619f9dc008ea03963033d5719e283c32514ea Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jan=20H=C3=B6nig?= <jan.hoenig@fau.de>
Date: Mon, 30 May 2022 20:52:49 +0200
Subject: [PATCH] WIP
---
pystencils/display_utils.py | 43 ++++++++++++++++++++++++++
pystencils_tests/test_show_assembly.py | 8 +++++
2 files changed, 51 insertions(+)
create mode 100644 pystencils_tests/test_show_assembly.py
diff --git a/pystencils/display_utils.py b/pystencils/display_utils.py
index f6c32ac88..e83067e90 100644
--- a/pystencils/display_utils.py
+++ b/pystencils/display_utils.py
@@ -70,6 +70,32 @@ def get_code_obj(ast: Union[KernelFunction, KernelWrapper], custom_backend=None)
return CodeDisplay(ast)
+def get_assembly_obj(ast: Union[KernelFunction, KernelWrapper]):
+ """Returns an object to display generated assembly code (C/C++)
+
+ Can either be displayed as HTML in Jupyter notebooks or printed as normal string.
+ """
+ from pystencils.backends.cbackend import generate_c
+
+ if isinstance(ast, KernelWrapper):
+ ast = ast.ast
+ code = generate_c(ast, dialect=ast.backend)
+
+ class AssemlbyDisplay:
+ def __init__(self, assembly):
+ self.assembly = assembly
+
+ def _repr_html(self):
+ return ''
+
+ def __str__(self):
+ return ''
+
+ def __repr__(self):
+ return ''
+ return AssemlbyDisplay(assembly)
+
+
def get_code_str(ast, custom_backend=None):
return str(get_code_obj(ast, custom_backend))
@@ -102,3 +128,20 @@ def show_code(ast: Union[KernelFunction, KernelWrapper], custom_backend=None):
console.print(syntax)
except ImportError:
print(code)
+
+
+def show_assembly(ast: Union[KernelFunction, KernelWrapper]):
+ code = get_assembly_obj(ast)
+
+ if _isnotebook():
+ from IPython.display import display
+ display(code)
+ else:
+ try:
+ import rich.syntax
+ import rich.console
+ syntax = rich.syntax.Syntax(str(code), "c++", theme="monokai", line_numbers=True)
+ console = rich.console.Console()
+ console.print(syntax)
+ except ImportError:
+ print(code)
diff --git a/pystencils_tests/test_show_assembly.py b/pystencils_tests/test_show_assembly.py
new file mode 100644
index 000000000..74b2b06f1
--- /dev/null
+++ b/pystencils_tests/test_show_assembly.py
@@ -0,0 +1,8 @@
+import pytest
+import sympy as sp
+
+import pystencils as ps
+
+
+def test_simple():
+ pass
--
GitLab