Skip to content

Conversation

vkovinicTT
Copy link
Member

@vkovinicTT vkovinicTT commented Oct 6, 2025

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:

  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

@vkovinicTT vkovinicTT changed the title enabled renaming op name without adding stack frame locations [Bug FiX] Enabling custom op name without adding stack frame locations Oct 8, 2025
@vkovinicTT vkovinicTT changed the title [Bug FiX] Enabling custom op name without adding stack frame locations [Bug Fix] Enabling custom op name without adding stack frame locations Oct 9, 2025
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

1 participant