This is an archive of the discontinued LLVM Phabricator instance.

[mlir] Broadcast scalars when vectorising tensor.extract
ClosedPublic

Authored by awarzynski on Dec 30 2022, 7:36 AM.

Details

Summary

When vectorizing tensor.extract embedded within linalg.generic, the
default option is to rewrite it as vector.gather. When doing so, we need
to make sure that the corresponding indices are vectorized accordingly.
However, the Linalg vectorizer will not vectorize constants like in the
following example. This is fixed by simply broadcasting %c0 and %c1.

func.func @example(%arg0: tensor<3x3xf32>, %arg2: tensor<1x1x3xf32>) -> tensor<1x1x3xf32> {
  %c0 = arith.constant 1 : index
  %c1 = arith.constant 2 : index
  %1 = linalg.generic {
    (...)
  } outs(...) {
  ^bb0(...):
    %2 = tensor.extract %arg0[%c0, %c1] : tensor<3x3xf32>
    linalg.yield %2 : f32
  } -> tensor<1x1x3xf32>
  return %1 : tensor<1x1x3xf32>
}

This patch makes sure that in this case the vectorizer broadcasts %c0 and %c1.
Additional tests are added to check other scenarios as well.

Diff Detail

Event Timeline

awarzynski created this revision.Dec 30 2022, 7:36 AM
Herald added a project: Restricted Project. · View Herald Transcript
awarzynski requested review of this revision.Dec 30 2022, 7:36 AM

This is just a small follow-up for https://reviews.llvm.org/D137660. My main goal was to make sure that the scenario tested in "vectorize_nd_tensor_extract_constant_idx" behaves sanely rather than throwing an error:

within split at test.mlir:1 offset :11:10: error: 'arith.addi' op requires the same type for all operands and results
    %7 = tensor.extract %arg0[%c0, %c0, %3] : tensor<3x3x3xf32>
         ^
within split at test.mlir:1 offset :11:10: note: see current operation: %9 = "arith.addi"(%4, %0) : (vector<1x1x3xindex>, index) -> index

Although the change is (hopefully) straightforward, I'm not 100% sure whether this is the right approach. @dcaballe did point out previously, that I should avoid broadcasting scalars and instead do the address calculation with scalars. That's on my TODO list.

awarzynski edited the summary of this revision. (Show Details)
awarzynski edited the summary of this revision. (Show Details)Jan 9 2023, 10:01 AM

Refine the comment inside Vectorizer.cpp

I've realised that the Linalg Vectorizer does not vectorize scalar constant "by design". Thanks @dcaballe for the pointer!

dcaballe added inline comments.Jan 9 2023, 12:09 PM
mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
636

Could we use one of the existing utilities that already generate a broadcast op? For example, broadcastIfNeeded. I also think we have to make sure that this constant goes through vectorizeOne code as a copy is generated for those cases where the constant also has a user that still need a scalar version of it after vectorization.

mlir/test/Dialect/Linalg/vectorization.mlir
1521

I understand that supporting this case is mostly for completeness as we should generate a vector.broadcast instead of a vector.gather in the future (this access is loading the same element at every iteration, right?)

awarzynski added inline comments.Jan 11 2023, 7:31 AM
mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
636

Wish I'd known about broadcastIfNeeded before, thanks! This will significantly simplify this function.

mlir/test/Dialect/Linalg/vectorization.mlir
1521

Correct.

This is to document what's being generated "today" rather than what we should be aiming for. Not sure whether such tests are desirable. I do find them very helpful when going over the vectoriser. But I might be in a minority :)

Switch to using broadcastIfNeeded

dcaballe accepted this revision.Jan 11 2023, 6:31 PM

Thanks!

This revision is now accepted and ready to land.Jan 11 2023, 6:31 PM