Skip to content

Evaluate multiple ops together#1921

Merged
Priya2698 merged 9 commits intomainfrom
pm/evaluate_segments
Mar 15, 2024
Merged

Evaluate multiple ops together#1921
Priya2698 merged 9 commits intomainfrom
pm/evaluate_segments

Conversation

@Priya2698
Copy link
Collaborator

This PR is intended to add support for evaluating multiple ops together and skip evaluating intermediate inputs.
Current use cases:

  1. CatOp: Currently, CatOp::evaluate uses the unpadded inputs, however the padded inputs are still evaluated.
  2. Matmul Fallback path: We need the expression evaluator to evaluate patterns such as MmaOp + Cast, MmaOp + Bias + Cast in a single call instead of one op at a time.

@Priya2698
Copy link
Collaborator Author

!build

@Priya2698 Priya2698 marked this pull request as ready for review March 13, 2024 18:36
Copy link
Collaborator

@wujingyue wujingyue left a comment

Choose a reason for hiding this comment

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

LGTM with comments

@@ -59,6 +59,11 @@ class ExpressionEvaluator {
//! Try to evaluate a value using const evaluator ref
Copy link
Collaborator

Choose a reason for hiding this comment

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

Suggested change
//! Try to evaluate a value using const evaluator ref
//! Try to evaluate a value without using known_values_.

const ExpressionEvaluator& ee,
const std::vector<PolymorphicValue>& inputs) const;

virtual std::vector<PolymorphicValue> evaluate(
Copy link
Collaborator

Choose a reason for hiding this comment

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

There should be a code comment here about why this version of evaluate (and thus the added complexity) is needed and how it's different from the version above.

Copy link
Collaborator Author

Choose a reason for hiding this comment

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

Added. Should be more clear now.

Copy link
Collaborator

Choose a reason for hiding this comment

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

Can you fix the format of the comment? It shows some unnecessary indentation.

@Priya2698
Copy link
Collaborator Author

!build

Copy link
Collaborator

@wujingyue wujingyue left a comment

Choose a reason for hiding this comment

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

🚢

const ExpressionEvaluator& ee,
const std::vector<PolymorphicValue>& inputs) const;

virtual std::vector<PolymorphicValue> evaluate(
Copy link
Collaborator

Choose a reason for hiding this comment

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

Can you fix the format of the comment? It shows some unnecessary indentation.

@Priya2698
Copy link
Collaborator Author

!build

@Priya2698 Priya2698 merged commit c681f1f into main Mar 15, 2024
@Priya2698 Priya2698 deleted the pm/evaluate_segments branch March 15, 2024 05:54
@Priya2698 Priya2698 mentioned this pull request Mar 25, 2024
Priya2698 added a commit that referenced this pull request Apr 2, 2024
Adds ATen evaluation for Matmul and Matmul + Bias. Based on PR #1921,
when evaluating a `castOp`, we _look back_ to see if there is a
preceding MmaOp and evaluate them together.

Issue #1775.
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.

3 participants