Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add support adding context instructions to basic block graphs. #289

Open
wants to merge 5 commits into
base: main
Choose a base branch
from

Conversation

virajbshah
Copy link
Collaborator

  • Update the graph builder and its Python bindings to add context
    instructions to basic block graphs and store context node mask to later
    be used by models.
  • Add tests for the new graph builder functionality.

 * Add fields to the `proto` specification to store context.
 * Add members to the Gematria `BasicBlock` data structure to store
   context and update methods on it and its Python binding accordingly.
 * Bonus: Remove dangling TODO.
Comment on lines +242 to +244
AddEdge(EdgeType::kStructuralDependency, previous_instruction_node,
instruction_node);
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think it would make sense to add a new edge type that connects the preceding context to the main block, and another one to connect the main block to the following context.

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I have an upcoming PR that adds EdgeType::kTakenBranch, which will always connect any preceding context with the main block and the main block with any following context when the LBR methodology is used to collect context (which also supports tracing multiple context blocks in either direction). Though, this does not necessarily have to be the case for any other data/context collection method.

Copy link
Collaborator

@ondrasej ondrasej Feb 17, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We could give the new edge type(s) a more generic name, e.g. kFromPrecedingContext and kToFollowingContext so that it doesn't suggest any specific type of context. But what I'd prefer to do is to give the model a signal where the boundary is.

SGTM for postponing this until a later PR.

ElementsAre(true, false, false, true, false, true, false, true,
false, false, true, false));
EXPECT_THAT(builder_->context_node_mask(),
ElementsAre(true, _, _, false, _, true, _, true, _, _, false, _));
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'd prefer to check all the values (to check that the context and non-context parts are contiguous).

 * Change `back_context` -> `preceding_context` and
  `front_context` -> `following_context`.
 * Add test for updated `BasicBlock::ToString`.
 * Keep previous proto field numbering.
 * Update the graph builder and its Python bindings to add context
   instructions to basic block graphs and store context node mask to later
   be used by models.
 * Add tests for the new graph builder functionality.
 * Change `back_context` -> `preceding_context` and
   `front_context` -> `following_context`.
 * Set `context_node_mask_` appropriately for non-instruction nodes as
   well, and update tests to reflect this.
@virajbshah virajbshah force-pushed the block-context-graphs branch from fcc10fd to 29fb29d Compare February 8, 2025 19:45
@virajbshah
Copy link
Collaborator Author

I have made some newer changes to the way context_node_mask_ is populated. Previously, only nodes having NodeType::kInstruction were ensured to have their corresponding entries in context_node_mask_ set correctly, as only these values mattered for usage in downstream models (and not checking the other values in BasicBlockGraphBuilderTest::MultipleBasicBlocksWithContext was to reflect this).

After this change, this behavior is better defined with context_node_mask_ set to true for non-instruction nodes connected exclusively to context instruction nodes.

Please have a look.

@virajbshah virajbshah requested a review from ondrasej February 8, 2025 19:45
@ondrasej
Copy link
Collaborator

I have made some newer changes to the way context_node_mask_ is populated. Previously, only nodes having NodeType::kInstruction were ensured to have their corresponding entries in context_node_mask_ set correctly, as only these values mattered for usage in downstream models (and not checking the other values in BasicBlockGraphBuilderTest::MultipleBasicBlocksWithContext was to reflect this).

After this change, this behavior is better defined with context_node_mask_ set to true for non-instruction nodes connected exclusively to context instruction nodes.

Please have a look.

I've added one technical comment, but otherwise this makes sense.

@@ -291,49 +302,53 @@ bool BasicBlockGraphBuilder::AddInputOperand(
assert(instruction_node >= 0);
assert(instruction_node < num_nodes());

bool is_context = context_node_mask_[instruction_node];
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'd rather pass this as an argument:

  • it's probably faster (it will go through registers, with at least one fewer memory access),
  • it reduces the risk that somewhere, sometimes we forget to set the the context mask for the instruction before adding its operands.

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.

2 participants