Hzfengsy commented on code in PR #18079:
URL: https://github.com/apache/tvm/pull/18079#discussion_r2157970729
##########
src/script/ir_builder/tir/ir.cc:
##########
@@ -219,12 +219,47 @@ void Writes(Array<ObjectRef> buffer_slices) {
frame->writes = writes;
}
+/*! \brief Recursively merge two annotations, the new attrs will override the
old ones */
+Map<String, Any> MergeAnnotations(const Map<String, Any>& new_attrs,
+ const Map<String, Any>& old_attrs) {
+ Map<String, Any> result = old_attrs;
+ for (const auto& [key, value] : new_attrs) {
+ auto old_value = old_attrs.Get(key);
+ // Case 1: the key is not in the old annotations, set the key to the new
value
+ if (!old_value) {
+ result.Set(key, value);
+ continue;
+ }
+
+ // Case 2: the key is in the old annotations
+ // Case 2.1: both are dicts
+ auto old_dict = old_value->try_cast<Map<String, Any>>();
+ auto new_dict = value.try_cast<Map<String, Any>>();
+ if (old_dict && new_dict) {
+ // Recursively merge the two dicts
+ auto merged_dict = MergeAnnotations(*old_dict, *new_dict);
+ result.Set(key, merged_dict);
+ continue;
+ }
+ // Case 2.2: the values are not both dicts, check if the keys are the same
+ if (!ffi::AnyEqual()(old_value.value(), value)) {
+ LOG(FATAL) << "ValueError: Try to merge two annotations with different
values for key `"
Review Comment:
Confirmed it works well. Try with testcase:
```python
from tvm.script import tir as T
@T.prim_func
def func3():
with T.block():
T.block_attr({"key1": "block1"})
T.block_attr({"key1": "block2"})
T.evaluate(0)
```
get output:
```
error: Try to merge two annotations with different values for key `key1`,
previous one is "block1", new one is "block2"
--> /gpfs/users/syfeng/tvm/t.py:7:9
|
7 | T.block_attr({"key1": "block2"})
| ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
note: run with `TVM_BACKTRACE=1` environment variable to display a backtrace.
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
