This is an archive of the discontinued LLVM Phabricator instance.

[mlir][Vector] Add support for Value indices to vector.extract/insert
ClosedPublic

Authored by dcaballe on Jul 11 2023, 5:10 PM.

Details

Diff Detail

Event Timeline

dcaballe created this revision.Jul 11 2023, 5:10 PM
Herald added a reviewer: aartbik. · View Herald Transcript
Herald added a reviewer: ftynse. · View Herald Transcript
Herald added a reviewer: kuhar. · View Herald Transcript
Herald added a project: Restricted Project. · View Herald Transcript
dcaballe requested review of this revision.Jul 11 2023, 5:10 PM
dcaballe planned changes to this revision.Jul 11 2023, 5:10 PM
dcaballe planned changes to this revision.Jul 12 2023, 10:21 AM

Hi Diego,

Thanks for this. As a start, I'll take care of the tests I see failing under Dialect/Vector:

MLIR :: Dialect/Vector/canonicalize.mlir
MLIR :: Dialect/Vector/constant-fold.mlir
MLIR :: Dialect/Vector/invalid.mlir
MLIR :: Dialect/Vector/ops.mlir
MLIR :: Dialect/Vector/scalar-vector-transfer-to-memref.mlir
MLIR :: Dialect/Vector/vector-broadcast-lowering-transforms.mlir
MLIR :: Dialect/Vector/vector-contract-matvec-transforms.mlir
MLIR :: Dialect/Vector/vector-contract-to-dot-transforms.mlir
MLIR :: Dialect/Vector/vector-contract-to-matrix-intrinsics-transforms.mlir
MLIR :: Dialect/Vector/vector-contract-to-outerproduct-transforms.mlir
MLIR :: Dialect/Vector/vector-contract-to-parallel-arith-transforms.mlir
MLIR :: Dialect/Vector/vector-dropleadunitdim-transforms.mlir
MLIR :: Dialect/Vector/vector-emulate-narrow-type.mlir
MLIR :: Dialect/Vector/vector-extract-strided-slice-lowering.mlir
MLIR :: Dialect/Vector/vector-gather-lowering.mlir
MLIR :: Dialect/Vector/vector-mask-lowering-transforms.mlir
MLIR :: Dialect/Vector/vector-multi-reduction-lowering.mlir
MLIR :: Dialect/Vector/vector-multi-reduction-outer-lowering.mlir
MLIR :: Dialect/Vector/vector-outerproduct-lowering-transforms.mlir
MLIR :: Dialect/Vector/vector-scan-transforms.mlir
MLIR :: Dialect/Vector/vector-shape-cast-lowering-transforms.mlir
MLIR :: Dialect/Vector/vector-transfer-to-vector-load-store.mlir
MLIR :: Dialect/Vector/vector-transforms.mlir
MLIR :: Dialect/Vector/vector-transpose-lowering.mlir
MLIR :: Dialect/Vector/vector-warp-distribute.mlir

Hey Diego, happy to help! I can update the Vector integration tests:

MLIR :: Integration/Dialect/Vector/CPU/test-broadcast.mlir
MLIR :: Integration/Dialect/Vector/CPU/test-compress.mlir
MLIR :: Integration/Dialect/Vector/CPU/test-contraction.mlir
MLIR :: Integration/Dialect/Vector/CPU/test-expand.mlir
MLIR :: Integration/Dialect/Vector/CPU/test-extract-strided-slice.mlir
MLIR :: Integration/Dialect/Vector/CPU/test-flat-transpose-col.mlir
MLIR :: Integration/Dialect/Vector/CPU/test-flat-transpose-row.mlir
MLIR :: Integration/Dialect/Vector/CPU/test-gather.mlir
MLIR :: Integration/Dialect/Vector/CPU/test-index-vectors.mlir
MLIR :: Integration/Dialect/Vector/CPU/test-maskedload.mlir
MLIR :: Integration/Dialect/Vector/CPU/test-maskedstore.mlir
MLIR :: Integration/Dialect/Vector/CPU/test-matrix-multiply-col.mlir
MLIR :: Integration/Dialect/Vector/CPU/test-matrix-multiply-row.mlir
MLIR :: Integration/Dialect/Vector/CPU/test-outerproduct-f32.mlir
MLIR :: Integration/Dialect/Vector/CPU/test-outerproduct-i64.mlir
MLIR :: Integration/Dialect/Vector/CPU/test-reductions-f32-reassoc.mlir
MLIR :: Integration/Dialect/Vector/CPU/test-reductions-f32.mlir
MLIR :: Integration/Dialect/Vector/CPU/test-reductions-f64-reassoc.mlir
MLIR :: Integration/Dialect/Vector/CPU/test-reductions-f64.mlir
MLIR :: Integration/Dialect/Vector/CPU/test-reductions-i32.mlir
MLIR :: Integration/Dialect/Vector/CPU/test-reductions-i64.mlir
MLIR :: Integration/Dialect/Vector/CPU/test-scan.mlir
MLIR :: Integration/Dialect/Vector/CPU/test-scatter.mlir
MLIR :: Integration/Dialect/Vector/CPU/test-shape-cast.mlir
MLIR :: Integration/Dialect/Vector/CPU/test-sparse-dot-matvec.mlir
MLIR :: Integration/Dialect/Vector/CPU/test-sparse-saxpy-jagged-matvec.mlir
MLIR :: Integration/Dialect/Vector/CPU/test-transfer-read-1d.mlir
MLIR :: Integration/Dialect/Vector/CPU/test-transfer-read-2d.mlir
MLIR :: Integration/Dialect/Vector/CPU/test-transfer-read-3d.mlir
MLIR :: Integration/Dialect/Vector/CPU/test-transfer-to-loops.mlir
MLIR :: Integration/Dialect/Vector/CPU/test-transfer-write.mlir
MLIR :: Integration/Dialect/Vector/CPU/test-transpose.mlir

-Andrzej

jsetoain added inline comments.Jul 13 2023, 8:45 AM
mlir/include/mlir/Dialect/Vector/IR/VectorOps.td
722

Why not use AnySignlessIntegerOrIndex here?

jsetoain added inline comments.Jul 13 2023, 9:05 AM
mlir/include/mlir/Dialect/Vector/IR/VectorOps.td
722

I chose the wrong type. What I meant is, why not allow for attributes as well as values? That should get rid of the issues with not accepting things like [3, 3, 3] as indices, no?

Thanks a lot for the help! Please, before you spend more time on this, please, give me some time to try mixing attributes and values, as @jsetoain suggested. We can do something similar to what tensor.extract_slice/insert_slice do.

mlir/include/mlir/Dialect/Vector/IR/VectorOps.td
722

I made that decision based on other ops like tensor.extract to be consistent. Using only the index type also simplifies things. We don't have to specify the type use for the indices in the assembly format, we don't have to verify that all the indices have the same type, we don't have to check for the type in transformations that are manipulating/replacing/adding more indices, etc.

Regarding mixing attributes and values, that's actually a very good point. I see that tensor.extract_slice is actually doing something like that:

let arguments = (ins
  AnyRankedTensor:$source,
  Variadic<Index>:$offsets,
  Variadic<Index>:$sizes,
  Variadic<Index>:$strides,
  DenseI64ArrayAttr:$static_offsets,
  DenseI64ArrayAttr:$static_sizes,
  DenseI64ArrayAttr:$static_strides
);

Let me give it a try!

Matt added a subscriber: Matt.Jul 19 2023, 3:13 PM
dcaballe updated this revision to Diff 546690.Aug 2 2023, 9:25 PM

Progress... Not ready yet!

dcaballe planned changes to this revision.Aug 2 2023, 9:25 PM
dcaballe updated this revision to Diff 556976.Sep 18 2023, 3:01 PM

Final rework

This should be ready to review and land. I added support for values and constant attributes to extract and insert ops as requested.

Thanks for this work Diego! this will useful for ArmSME. We have ops to load/store slices of a tile (2-D scalable vector) and this support should enable us to lower vector.insert / vector.extract to these ops. #66691 adds support to lower vector.print to ArmSME, this could be done generically in VectorToSCF once this lands, after which we could add our target specific lowerings :) Also be useful for places where we can't fully unroll.

I've left a few comments/questions, mostly minor. I also noticed there's no tests, it would be good to have some in mlir/test/Dialect/Vector/ops.mlir. I also tried lowering a simple example to LLVM but hit an assert, is this expected to work yet? Example below

; mlir-opt -convert-vector-to-llvm

func.func @extract_element_from_vec_2d(%arg0: vector<2x16xf32>) -> vector<16xf32> {
  %c1 = arith.constant 1 : index
  %0 = vector.extract %arg0[%c1]: vector<2x16xf32>
  return %0 : vector<16xf32>
}
mlir/lib/Dialect/Vector/IR/VectorOps.cpp
1246–1247

nit: formatting, these can go on the same line

1478–1480

move above valueToExtractFrom

mlir/lib/Dialect/Vector/Transforms/LowerVectorContract.cpp
98–99

remove cast?

124–125

remove cast?

mlir/lib/Dialect/Vector/Transforms/VectorDistribute.cpp
1047–1048

it's not clear to me why is this an assert? (And below for insert)

mlir/lib/Dialect/Vector/Transforms/VectorTransferOpTransforms.cpp
711

spelling: Unexpected

+1 to adding some tests in mlir/test/Dialect/Vector/ops.mlir. Otherwise a few nits. Thanks a ton for resurrecting this work :)

mlir/include/mlir/Dialect/Vector/IR/VectorOps.td
729

Could you add examples with static sizes?

mlir/lib/Conversion/VectorToSPIRV/VectorToSPIRV.cpp
38–39

This comment is out of date :)

mlir/lib/Dialect/Vector/IR/VectorOps.cpp
226–227

[nit] There's an identical comment in the header file - I'd only keep one (otherwise they might get out of sync).

1395–1396

What about dynamic positions? If unsupported, then perhaps add an assert?

dcaballe updated this revision to Diff 557091.Sep 19 2023, 11:42 PM
dcaballe marked 6 inline comments as done.

Addressed feedback

Thanks for the feedback! I only addressed @c-rhodes comments for now. I'll take a look at the remaining ones tomorrow.

I also tried lowering a simple example to LLVM but hit an assert, is this expected to work yet? Example below

Sorry about that. My main goal with this patch was about having the representation working at vector level (and adjust all the impacted tests, which were a lot at the beginning) but not so much about adding support for the new value-based functionality end to end. I just added support for 1-D cases to the LLVM conversion to demonstrate that it works but there are quite a few TODOs, esp. wrt canonicalizations. Hopefully that makes sense!

mlir/lib/Dialect/Vector/Transforms/VectorDistribute.cpp
1047–1048

Replaced with failure. It's more like a personal preference: I usually lean towards asserting instead of just silently bailing out, esp. when I don't know the pass or the consequences of bailing out very well.

I also tried lowering a simple example to LLVM but hit an assert, is this expected to work yet? Example below

Sorry about that. My main goal with this patch was about having the representation working at vector level (and adjust all the impacted tests, which were a lot at the beginning) but not so much about adding support for the new value-based functionality end to end. I just added support for 1-D cases to the LLVM conversion to demonstrate that it works but there are quite a few TODOs, esp. wrt canonicalizations. Hopefully that makes sense!

No worries, thanks for clarifying

mlir/include/mlir/Dialect/Vector/IR/VectorOps.td
729

could you also add an example with mixed static/dynamic indices?

mlir/lib/Dialect/Vector/Transforms/VectorDistribute.cpp
1047–1048

Replaced with failure. It's more like a personal preference: I usually lean towards asserting instead of just silently bailing out, esp. when I don't know the pass or the consequences of bailing out very well.

That's fair, an assert to me indicates a bug so I was wondering why a rank-1 extract with dynamic index is considered a bug in this particular rewrite, but I understand this change touches a lot of code and it's difficult to understand the implication for all scenarios.

mlir/test/Dialect/Vector/ops.mlir
223 ↗(On Diff #557091)

test with mixed static/dynamic indices?

dcaballe updated this revision to Diff 557140.Sep 20 2023, 12:22 PM
dcaballe marked 5 inline comments as done.

Addressed feedback

I think I addressed everything. Thanks!

mlir/lib/Dialect/Vector/IR/VectorOps.cpp
226–227

I think that's common practice in MLIR. It has been pointed out to me that we should actually add doc to both decl and def. Extra doc is not harmful, I guess.

1395–1396

We can't add an assert here because this constructor is used by all kind of extract ops but the checks below should gracefully fail when they have dynamic operands.

c-rhodes accepted this revision.Sep 21 2023, 12:41 AM

LGTM cheers

mlir/include/mlir/Dialect/Vector/IR/VectorOps.td
590

%5

This revision is now accepted and ready to land.Sep 21 2023, 12:41 AM