Reland upstream commit f9688ab

[whitneywhtsang]

Please do not squash and merge this PR.

[anmyachev]

@whitneywhtsang this is the only change needed to fix lit tests. In the new code, the swap occurs conditionally (if (opIdx == 1)), which apparently did not work for dpas so I returned unconditional swap.

[whitneywhtsang]

Thanks for the quick update, I wonder if the logic can be added in intel specific files instead. Do you have any suggestions @chengjunlu ?

[chengjunlu]

I think we can propose the changes to make the getOrderForDotOperand as the interface to the MmaTraits.

I found the AMD engineer refactor the code at this PR ff02a46. But their comments about the order is not general and doesn't make sense to Intel GPU. The order should be overridable by the parent layout of the DotOp layout.

[anmyachev]

I think we can propose the changes to make the getOrderForDotOperand as the interface to the MmaTraits.

@chengjunlu with this approach, BlockedEncodingAttr type needs to be handled separately, since it does not inherit MmaTraits interface. Tests in CI are now falling because of this (should be fixed in the last commit).

What can be done about this?

[chengjunlu]

There is not a easy way to unified the BlockedEncodingAttr and the MmaTraits unless let it to inherited the MmaTraits as well. But I am not sure whether it worth to do so for now.

For the simplicity of the changes, we can handle it separately and check the feed back of the public Triton.

[whitneywhtsang]

FYI @anmyachev, squashed the last few commits, so is easier to isolate changes needed to review.

[victor-eds]

Changes look sensible to me. Just a question: why does our encoding have a different order compared to the rest? What would be the cost of modifying it so the order matches other dot encodings?

[victor-eds]

As I said, changes LGTM. I am just revoking approval till we get an assessment of costs: changing our order vs upstreaming this change (we may just be asked to change our order instead).

[anmyachev]

Just a question: why does our encoding have a different order compared to the rest? What would be the cost of modifying it so the order matches other dot encodings?

Good question. There is a chance that this issue was discussed earlier during initial implementation. Let's ask more experienced Triton developers than me :) @chengjunlu @whitneywhtsang do you have an answer to this question?

I am just revoking approval till we get an assessment of costs: changing our order vs upstreaming this change (we may just be asked to change our order instead).

If there is no answer to this question, then I can research the implementation history myself and look for answers to the question of why it was done this way and not differently (any pointers and links to code will speed up the process). However, this is not fast, wouldn't it be better to merge this pull request to simplify merging subsequent commits? @whitneywhtsang

P.S. after rebase the tests don't pass, I'll have a look (fixed, my changes after rebase were lost, i returned them)

[chengjunlu]

Changes look sensible to me. Just a question: why does our encoding have a different order compared to the rest? What would be the cost of modifying it so the order matches other dot encodings?

Actually the all the DotOp layout has the same order before this PR ff02a46. (The linear ID of the register dim to coordinate are computed as in row major.)

The code before AMD's change

 else if (auto dotLayout = dyn_cast<DotOperandEncodingAttr>(layout)) {
    auto rank = getWarpsPerCTA(dotLayout.getParent()).size();
    SmallVector<unsigned> order(rank);
    for (auto i = 0; i < rank; ++i)
      order[i] = rank - 1 - i;
    return order;
  }

I think the AMD's engineer has different interpretation based on their comments about the new code:

// The 'order' field typically represents a descending sorted array of
  // dimensions based on contiguity. For instance, in axisInfo utilities that
  // retrieve tensor contiguity, it's assumed that the dimension with the
  // highest contiguity corresponds to order[0].
  //
  // The relation between contiguity and order is only relevant if the layout
  // interfaces with HBM, as is the case when we load tensor from HBM to
  // registers in the dot layout to bypass LDS. When bypassing LDS, we make the
  // following assumptions about tensor layouts:
  // - Tensor A (opIdx == 0) is considered to be row-major.
  // - Tensor B (opIdx == 1) is considered to be column-major.
  //
  // Based on these assumptions, we define the following orders:
  // - For opIdx == 0, we assume an order of [1, 0].
  // - For opIdx == 1, we assume an order of [0, 1].

For Intel GPU, the layout is only used to describe the layout of the value in register and the matrix A and B are both row-major in register.

[chengjunlu]

I didn't make a clear review at the first time. If the original code changes is only made for AMD, then we can keep all those DotOp register layout order unchanged.

[chengjunlu]

On the other hand, it is good for the third party extension can override the getOrder for DotOp layout with MmaEncodingTrait. I am neutral to this change.

[anmyachev]

I suggest speeding up the merge of this pull request, leaving only the workaround for now. We can try to upstream the interface function getOrderForDotOperand separately.

@victor-eds @whitneywhtsang @chengjunlu please approve if this makes sense.

[whitneywhtsang]

I suggest speeding up the merge of this pull request, leaving only the workaround for now. We can try to upstream the interface function getOrderForDotOperand separately.

@victor-eds @whitneywhtsang @chengjunlu please approve if this makes sense.

I think it makes sense, let's remove the last two commits, and add a FIXME comment in f9ccfeb.

[anmyachev]

I suggest speeding up the merge of this pull request, leaving only the workaround for now. We can try to upstream the interface function getOrderForDotOperand separately.
@victor-eds @whitneywhtsang @chengjunlu please approve if this makes sense.

I think it makes sense, let's remove the last two commits, and add a FIXME comment in f9ccfeb.

Done.

[victor-eds]

I'm fine with this merge 👍