Fix MatmulOp IterDomain mapping by jacobhinkle · Pull Request #2246 · NVIDIA/Fuser

jacobhinkle · 2024-05-14T20:31:05Z

This PR does the following:

Add MatmulOp to ir_utils::isTvOp so that its IterDomains will be automatically propagated by IdModel.
Updates the tests to check that all non-Broadcast axes are properly mapped by IdModel through the MatmulOp.
Changes the output of MatmulOp to have an IterType::Reduction axis in the last position of its root domain to represent the K dimension. This change was motivated by needing a way to have both operand K dimensions exact mapped together, as they would be if the op were translated to a mul+sum+cast.
Updates the matmul op to translate trivial cases where K=1 to simple multiply+cast patterns.

Fixes #1707. In fact, that test was actually fixed by #2175 but the test validation was failing because isTvOp was not picking up the matmul as a reduction.

This also adds testing of exact mapping to the node tests (WIP).

tests/cpp/test_matmul_aten_evaluation.cpp

Tests pass!

jacobhinkle · 2024-05-15T12:32:48Z

csrc/ops/composite.cpp

    out_domain[idx] = ops::newOutputIterDomain(input_ids);
  }

+  out_domain[ndims_out - 1] = IterDomainBuilder(mapping_a.back())


I placed the rK dimension last in the output, just because that makes it a bit easier to handle the many cases we can encounter for matmul. Note that this does not need to match the position of the K dimension in a translated MmaOp.

Should we use both, mapping_a and mapping_b here?

Good question. I was assuming the K dimension would be IterType::Iteration in both operands. They will be exact mapped, so they should be equivalent to one another so I just took the ID from A.

What happens if one is Symbolic and the other Concrete tensor?

Good point. I updated ops::newOutputIterDomain to have a std::optional<IterType> force_iter_type argument that I am now using to create the reduction domain. That way, if the A ID is Symbolic but the B ID is not, then we'll use the B extent.

jacobhinkle · 2024-05-15T12:33:43Z

csrc/ops/composite.cpp

      " and ",
      tv_b->dtype());

+  // Check for K=1 i.e. reduction of broadcast. In these cases we don't need a


I considered placing this translation to mul in another PR, but it is relatively compact and it is difficult to work around if we want to exhaustively check ID mappings with all different combinations of inputs.

Thanks for adding these cases!

jacobhinkle · 2024-05-15T12:34:40Z

csrc/ops/utils.cpp


+  // Input A to matmul: {*, M, K}
+  // Input B to matmul: {*, K, N}
+  auto kpos = input_role == MatmulRole::INPUT_A ? inp_size - 1 : inp_size - 2;


Changes in this file are to accomodate the new Reduction output domain by mapping it to the k position in each operand.

jacobhinkle · 2024-05-15T12:38:29Z

tests/cpp/test_matmul_aten_evaluation.cpp

    ReductionAxisIsOne,
    ATenNodesParametrizedTest,
-    testing::Values(std::make_tuple(Sizes({m, 1}), Sizes({1, n}))));
+    testing::Combine(


Off-topic: It's not needed here, but for new tests I wonder if we could combine these kinds of parametrizations? For example since we use k=32, we could parametrize the test with a boolean flag k_is_one and if the flag is true we would replace any 32s in the shapes with 1 before running the test. That would reduce the need to repeat all the inputs for each case, but it makes it harder to filter to only the concrete K tests since the tests are simply numbered. It would be great if the parameters were reflected in the name like they are in pytest...

We could add another flag, and replace the value of K within tests. I am not sure if there is a way to replace the values before they are passed to a test: that is, the test directly views K=1. I'll look into it.

jacobhinkle · 2024-05-15T12:40:07Z

csrc/ops/composite.cpp

-  // > 1D.
-  auto ndims_out = std::max(ndims_a, ndims_b);
+  // > 1D, but with 1 additional IterType::Reduction axis rK.
+  auto ndims_out = std::max(ndims_a, ndims_b) + 1;


Changes in this file are for adding a new Reduction output domain to facilitate easier mapping of K axes.

jacobhinkle · 2024-05-15T12:44:49Z

!build --diff

csrc/ops/composite.cpp

tests/cpp/test_matmul_aten_evaluation.cpp

Priya2698 · 2024-05-15T19:16:08Z

csrc/ops/composite.cpp

      tv_b->dtype());

+  // Check for K=1 i.e. reduction of broadcast. In these cases we don't need a
+  // matmul so we translate it to a multiplication+cast


Can we have a comment at the matmul definition or in the header enumerating the different cases and how they are handled for an easier summary.

Priya2698

LGTM overall.

My main question is around using iterdomains from both A and B for mapping. While they will be exact mapped, what happens we may have one Symbolic and one Concrete tensor (although rare, and we do not test this). The newOutputIterDomain call will resolve the extents using one of the tensors if we use both. Wdyt?

jacobhinkle · 2024-05-15T20:34:08Z

I agree about symbolic domains and I think we might need to handle this at concretization. Namely if K is a broadcast dim then we should translate to a mul op before segmentation. I'll work on an example for a case like that.

This lets us create a new Reduction IterDomain from the inputs. Typically this will just use the mapping_a.back() extent, but if that ID is Symbolic but the B ID is not, then it will use that extent instead.

Add MatmulOp to ir_utils::isTvOp

a697b55

This also adds testing of exact mapping to the node tests (WIP).

jacobhinkle commented May 14, 2024

View reviewed changes

tests/cpp/test_matmul_aten_evaluation.cpp Outdated Show resolved Hide resolved

jacobhinkle added 3 commits May 15, 2024 00:34

Add IterType::Reduction domain for K dim in output

04522b8

Translate bcast K as simple product

7b04286

Finish testing all combinations of mappings

dcaf349

Tests pass!

jacobhinkle commented May 15, 2024

View reviewed changes

jacobhinkle changed the title ~~[WIP] Add MatmulOp to ir_utils::isTvOp~~ Fix MatmulOp IterDomain mapping May 15, 2024

jacobhinkle requested a review from Priya2698 May 15, 2024 12:43

jacobhinkle marked this pull request as ready for review May 15, 2024 12:43

jacobhinkle mentioned this pull request May 15, 2024

Translate MatmulOp and LinearOp #2236

Merged

Priya2698 reviewed May 15, 2024

View reviewed changes

csrc/ops/composite.cpp Outdated Show resolved Hide resolved

Priya2698 reviewed May 15, 2024

View reviewed changes

tests/cpp/test_matmul_aten_evaluation.cpp Show resolved Hide resolved

Remove redundant bcast K check

e8f77e0

Priya2698 reviewed May 15, 2024

View reviewed changes

Priya2698 approved these changes May 15, 2024

View reviewed changes

Add force_iter_type argument to newOutputIterDomain

7b30790

This lets us create a new Reduction IterDomain from the inputs. Typically this will just use the mapping_a.back() extent, but if that ID is Symbolic but the B ID is not, then it will use that extent instead.

jacobhinkle mentioned this pull request May 16, 2024

LinearOp IR node #2240

Merged

jacobhinkle merged commit a6ce3e1 into main May 16, 2024

jacobhinkle deleted the matmul_op_id_mapping branch May 16, 2024 13:11

Conversation

jacobhinkle commented May 14, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

jacobhinkle commented May 15, 2024

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Priya2698 left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

jacobhinkle commented May 15, 2024

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

jacobhinkle commented May 14, 2024 •

edited

Loading

Priya2698 left a comment •

edited

Loading