[Mlir-commits] [mlir] [MLIR] Auto generate source location for python bindings (PR #112923)
Guray Ozen
llvmlistbot at llvm.org
Fri Oct 18 09:14:05 PDT 2024
https://github.com/grypp updated https://github.com/llvm/llvm-project/pull/112923
>From a5437c05a96cedf538eea98b245e208b1d61d51a Mon Sep 17 00:00:00 2001
From: Guray Ozen <gozen at nvidia.com>
Date: Fri, 18 Oct 2024 17:31:51 +0200
Subject: [PATCH 1/3] [MLIR] Auto generate source location for python bindings
This PR introduces the `get_source_location()` function, which:
Traverses the call stack to locate the frame just before the one whose filename contains 'python_packages'.
Returns a source location using the selected frame, useful for debugging or tracking the source of execution. The `python_packages` is the install directory for the python bindings, it gives the most meaningful source location.
The `get_source_location()` gets called for each OP building.
---
mlir/python/mlir/dialects/_ods_common.py | 33 +++++++++++++++++++
mlir/test/python/ir/location.py | 30 +++++++++++++++++
mlir/tools/mlir-tblgen/OpPythonBindingGen.cpp | 3 ++
3 files changed, 66 insertions(+)
diff --git a/mlir/python/mlir/dialects/_ods_common.py b/mlir/python/mlir/dialects/_ods_common.py
index d40d936cdc83d6..768a09fcc0e540 100644
--- a/mlir/python/mlir/dialects/_ods_common.py
+++ b/mlir/python/mlir/dialects/_ods_common.py
@@ -2,6 +2,7 @@
# See https://llvm.org/LICENSE.txt for license information.
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+import inspect
from typing import (
List as _List,
Optional as _Optional,
@@ -50,6 +51,38 @@ def segmented_accessor(elements, raw_segments, idx):
return elements[start:end]
+
+
+def get_source_location():
+ """
+ Returns a source location from the frame just before the one whose
+ filename includes 'python_packages'.
+ """
+ frame = inspect.currentframe()
+ outer_frames = inspect.getouterframes(frame)
+
+ # Traverse the frames in reverse order, excluding the current frame
+ selected_frame = None
+ for i in range(len(outer_frames)-1, -1, -1):
+ current_frame = outer_frames[i]
+ if 'python_packages' in current_frame.filename:
+ # Select the frame before the one containing 'python_packages'
+ selected_frame = outer_frames[i+1] if i-1 >= 0 else current_frame
+ break
+ if selected_frame is None:
+ # If no frame containing 'python_packages' is found, use the last frame
+ selected_frame = outer_frames[-1]
+
+ # Create file location using the selected frame
+ file_loc = _cext.ir.Location.file(
+ selected_frame.filename,
+ selected_frame.lineno,
+ 0
+ )
+ loc = _cext.ir.Location.name(selected_frame.function, childLoc = file_loc)
+ return loc
+
+
def equally_sized_accessor(
elements, n_simple, n_variadic, n_preceding_simple, n_preceding_variadic
):
diff --git a/mlir/test/python/ir/location.py b/mlir/test/python/ir/location.py
index f66d6c501dcf5c..6503b76f5c9c80 100644
--- a/mlir/test/python/ir/location.py
+++ b/mlir/test/python/ir/location.py
@@ -2,6 +2,7 @@
import gc
from mlir.ir import *
+from mlir.dialects import arith
def run(f):
@@ -150,3 +151,32 @@ def testLocationCapsule():
run(testLocationCapsule)
+
+
+
+# CHECK-LABEL: TEST: autoGeneratedLocation
+def autoGeneratedLocation():
+ def generator() :
+ return arith.ConstantOp(value=123, result=IntegerType.get_signless(32))
+ with Context() as ctx, Location.unknown():
+ module = Module.create()
+ with InsertionPoint(module.body):
+ a = arith.ConstantOp(value=42, result=IntegerType.get_signless(32))
+ b = arith.AddIOp(a, generator())
+ module.operation.print(enable_debug_info=True)
+
+#CHECK: module {
+#CHECK: %{{.*}} = arith.constant 42 : i32 loc(#loc4)
+#CHECK: %{{.*}} = arith.constant 123 : i32 loc(#loc5)
+#CHECK: %{{.*}} = arith.addi %{{.*}}, %{{.*}} : i32 loc(#loc6)
+#CHECK: } loc(#loc)
+
+#CHECK: #loc = loc(unknown)
+#CHECK: #loc1 = loc({{.*}}:164:0)
+#CHECK: #loc2 = loc({{.*}}:160:0)
+#CHECK: #loc3 = loc({{.*}}:165:0)
+#CHECK: #loc4 = loc("autoGeneratedLocation"(#loc1))
+#CHECK: #loc5 = loc("generator"(#loc2))
+#CHECK: #loc6 = loc("autoGeneratedLocation"(#loc3))
+
+run(autoGeneratedLocation)
\ No newline at end of file
diff --git a/mlir/tools/mlir-tblgen/OpPythonBindingGen.cpp b/mlir/tools/mlir-tblgen/OpPythonBindingGen.cpp
index 0c5c936f5addee..bcea716748bd31 100644
--- a/mlir/tools/mlir-tblgen/OpPythonBindingGen.cpp
+++ b/mlir/tools/mlir-tblgen/OpPythonBindingGen.cpp
@@ -35,6 +35,7 @@ constexpr const char *fileHeader = R"Py(
from ._ods_common import _cext as _ods_cext
from ._ods_common import (
equally_sized_accessor as _ods_equally_sized_accessor,
+ get_source_location as _get_source_location,
get_default_loc_context as _ods_get_default_loc_context,
get_op_result_or_op_results as _get_op_result_or_op_results,
get_op_result_or_value as _get_op_result_or_value,
@@ -491,6 +492,8 @@ constexpr const char *initTemplate = R"Py(
attributes = {{}
regions = None
{1}
+ if loc is None:
+ loc = _get_source_location()
super().__init__(self.build_generic({2}))
)Py";
>From adb1f823ec0f4c8db8f56180336bf9610e15bbac Mon Sep 17 00:00:00 2001
From: Guray Ozen <gozen at nvidia.com>
Date: Fri, 18 Oct 2024 17:39:51 +0200
Subject: [PATCH 2/3] fix indent
---
mlir/python/mlir/dialects/_ods_common.py | 25 +++++++++---------------
1 file changed, 9 insertions(+), 16 deletions(-)
diff --git a/mlir/python/mlir/dialects/_ods_common.py b/mlir/python/mlir/dialects/_ods_common.py
index 768a09fcc0e540..df835d63dc2d5d 100644
--- a/mlir/python/mlir/dialects/_ods_common.py
+++ b/mlir/python/mlir/dialects/_ods_common.py
@@ -51,8 +51,6 @@ def segmented_accessor(elements, raw_segments, idx):
return elements[start:end]
-
-
def get_source_location():
"""
Returns a source location from the frame just before the one whose
@@ -60,26 +58,22 @@ def get_source_location():
"""
frame = inspect.currentframe()
outer_frames = inspect.getouterframes(frame)
-
+
# Traverse the frames in reverse order, excluding the current frame
selected_frame = None
- for i in range(len(outer_frames)-1, -1, -1):
+ for i in range(len(outer_frames) - 1, -1, -1):
current_frame = outer_frames[i]
- if 'python_packages' in current_frame.filename:
+ if "python_packages" in current_frame.filename:
# Select the frame before the one containing 'python_packages'
- selected_frame = outer_frames[i+1] if i-1 >= 0 else current_frame
+ selected_frame = outer_frames[i + 1] if i - 1 >= 0 else current_frame
break
if selected_frame is None:
# If no frame containing 'python_packages' is found, use the last frame
selected_frame = outer_frames[-1]
-
+
# Create file location using the selected frame
- file_loc = _cext.ir.Location.file(
- selected_frame.filename,
- selected_frame.lineno,
- 0
- )
- loc = _cext.ir.Location.name(selected_frame.function, childLoc = file_loc)
+ file_loc = _cext.ir.Location.file(selected_frame.filename, selected_frame.lineno, 0)
+ loc = _cext.ir.Location.name(selected_frame.function, childLoc=file_loc)
return loc
@@ -171,11 +165,10 @@ def get_op_result_or_op_results(
return (
list(get_op_results_or_values(op))
if len(op.results) > 1
- else get_op_result_or_value(op)
- if len(op.results) > 0
- else op
+ else get_op_result_or_value(op) if len(op.results) > 0 else op
)
+
ResultValueTypeTuple = _cext.ir.Operation, _cext.ir.OpView, _cext.ir.Value
ResultValueT = _Union[ResultValueTypeTuple]
VariadicResultValueT = _Union[ResultValueT, _Sequence[ResultValueT]]
>From c18c3370b70cb8746c0a85cab8600bcefa3455a1 Mon Sep 17 00:00:00 2001
From: Guray Ozen <gozen at nvidia.com>
Date: Fri, 18 Oct 2024 18:13:54 +0200
Subject: [PATCH 3/3] indent
---
mlir/test/python/ir/location.py | 38 +++++++++++++++++----------------
1 file changed, 20 insertions(+), 18 deletions(-)
diff --git a/mlir/test/python/ir/location.py b/mlir/test/python/ir/location.py
index 6503b76f5c9c80..db87ef450cb6af 100644
--- a/mlir/test/python/ir/location.py
+++ b/mlir/test/python/ir/location.py
@@ -44,6 +44,7 @@ def testLocationAttr():
run(testLocationAttr)
+
# CHECK-LABEL: TEST: testFileLineCol
def testFileLineCol():
with Context() as ctx:
@@ -153,11 +154,11 @@ def testLocationCapsule():
run(testLocationCapsule)
-
# CHECK-LABEL: TEST: autoGeneratedLocation
def autoGeneratedLocation():
- def generator() :
- return arith.ConstantOp(value=123, result=IntegerType.get_signless(32))
+ def generator():
+ return arith.ConstantOp(value=123, result=IntegerType.get_signless(32))
+
with Context() as ctx, Location.unknown():
module = Module.create()
with InsertionPoint(module.body):
@@ -165,18 +166,19 @@ def generator() :
b = arith.AddIOp(a, generator())
module.operation.print(enable_debug_info=True)
-#CHECK: module {
-#CHECK: %{{.*}} = arith.constant 42 : i32 loc(#loc4)
-#CHECK: %{{.*}} = arith.constant 123 : i32 loc(#loc5)
-#CHECK: %{{.*}} = arith.addi %{{.*}}, %{{.*}} : i32 loc(#loc6)
-#CHECK: } loc(#loc)
-
-#CHECK: #loc = loc(unknown)
-#CHECK: #loc1 = loc({{.*}}:164:0)
-#CHECK: #loc2 = loc({{.*}}:160:0)
-#CHECK: #loc3 = loc({{.*}}:165:0)
-#CHECK: #loc4 = loc("autoGeneratedLocation"(#loc1))
-#CHECK: #loc5 = loc("generator"(#loc2))
-#CHECK: #loc6 = loc("autoGeneratedLocation"(#loc3))
-
-run(autoGeneratedLocation)
\ No newline at end of file
+
+# CHECK: module {
+# CHECK: %{{.*}} = arith.constant 42 : i32 loc(#loc4)
+# CHECK: %{{.*}} = arith.constant 123 : i32 loc(#loc5)
+# CHECK: %{{.*}} = arith.addi %{{.*}}, %{{.*}} : i32 loc(#loc6)
+# CHECK: } loc(#loc)
+
+# CHECK: #loc = loc(unknown)
+# CHECK: #loc1 = loc({{.*}}:164:0)
+# CHECK: #loc2 = loc({{.*}}:160:0)
+# CHECK: #loc3 = loc({{.*}}:165:0)
+# CHECK: #loc4 = loc("autoGeneratedLocation"(#loc1))
+# CHECK: #loc5 = loc("generator"(#loc2))
+# CHECK: #loc6 = loc("autoGeneratedLocation"(#loc3))
+
+run(autoGeneratedLocation)
More information about the Mlir-commits
mailing list