Skip to content

Commit d0b43cc

Browse files
authored
[Bug Fix] Enabling custom op name without adding stack frame locations (#9676)
This PR fixes two related issues in `PopulateXlaOpMetadata` that prevented proper usage of `CustomOpNameMetaData`: ## Fix 1: Support custom op names without stack frame locations This will enable user to change the name of the op using `CustomOpNameMetaData` without necessarily having to add stack frame locations. Example of usage (if we pass 0 for `max_stack_depth` field it will just prepend custom prefix to the current op name): ``` torch_xla._XLAC._set_xla_custom_op_name_prefix(tensor, your_custom_name_prefix, 0) ``` Additionaly, this will guard the `AddStackFrameLocations()` function from passing invalid number for the `max_stack_depth`. If we were to pass number that is <= 0, we would get a **segmentation fault** due to improper iterator dereferencing (which could've happen before this change). The `AddStackFrameLocations()` function in `stack_frame_index_builder.cpp` uses reverse iterators and assumes at least one iteration occurs: ```cpp auto frame_it = frame_info.rbegin(); for (; frame_it != frame_info.rend() && depth < max_stack_depth; ++frame_it) { // Loop never executes when max_stack_depth == 0 } --frame_it; // ← Segfault: iterator is still at rbegin(), decrement goes past-the-end metadata_to_populate.set_source_file(frame_it->file); // ← Dereference invalid iterator ``` ## Fix 2: Prevent scope from overwriting custom metadata **Problem:** Even when users set custom metadata via `_set_xla_custom_op_name_prefix()`, the `nmeta.scope` field was unconditionally overwriting the custom `op_name_prefix`. This affected operations like `add` and `mul` which have scope set (e.g., `aten::add.3`), resulting in loss of user-provided semantic location information. **Changes:** - Modified the condition from `if (!nmeta.scope.empty())` to `else if (!nmeta.scope.empty())` - This ensures custom metadata takes precedence: custom metadata is used if available, otherwise scope is used as fallback **Precedence hierarchy** (now correctly implemented): 1. Custom user metadata (via `SetUserMetadata` APIs) - highest priority 2. Scope-based naming (auto-generated by torch-xla) - fallback 3. Bare op_type - default
1 parent d0bcacd commit d0b43cc

File tree

1 file changed

+8
-4
lines changed

1 file changed

+8
-4
lines changed

torch_xla/csrc/lowering_context.cpp

Lines changed: 8 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -78,15 +78,19 @@ class HloMetadataSetter {
7878
max_stack_depth = custom_opname_meta->max_stack_depth;
7979
}
8080

81-
if (!nmeta.scope.empty()) {
81+
else if (!nmeta.scope.empty()) {
8282
op_name_prefix =
8383
absl::StrCat(absl::StrReplaceAll(nmeta.scope, {{":", "_"}}), "/");
8484
}
8585
metadata.set_op_name(absl::StrCat(op_name_prefix, op_type));
8686

87-
// Sets file, line and stack_frame_id in metadata
88-
lowering_context.stack_frame_index_builder()->AddStackFrameLocations(
89-
nmeta.frame_info, static_cast<int>(max_stack_depth), metadata);
87+
// NOTE: if max_stack_depth is 0, we are just renaming the op, so we don't
88+
// need to add stack frame locations
89+
if (max_stack_depth > 0) {
90+
// Sets file, line and stack_frame_id in metadata
91+
lowering_context.stack_frame_index_builder()->AddStackFrameLocations(
92+
nmeta.frame_info, static_cast<int>(max_stack_depth), metadata);
93+
}
9094

9195
lowering_context.builder()->SetOpMetadata(std::move(metadata));
9296
}

0 commit comments

Comments
 (0)