[mlir][vector] Adds ToElementsToTargetShape pattern. #166476
Conversation
|
✅ With the latest revision this PR passed the C/C++ code formatter. |
amd-eochoalo
force-pushed
the
eochoa/2025-11-04/use-get-shape-for-unroll-default-implementation
branch
from
a2fb4a8 to
9def4f7
Compare
amd-eochoalo
requested review from
Groverkss,
antiagainst,
banach-space,
dcaballe,
hanhanW,
kuhar and
nicolasvasilache
as code owners
amd-eochoalo
changed the title
amd-eochoalo
force-pushed
the
eochoa/2025-11-04/use-get-shape-for-unroll-default-implementation
branch
from
3fb2649 to
5de71f2
Compare
amd-eochoalo
force-pushed
the
eochoa/2025-11-04/use-get-shape-for-unroll-default-implementation
branch
from
5de71f2 to
8fe386a
Compare
| /// // In SPIR-V's default environment vector of size 8 | ||
| /// // are not allowed. |
| for (const Value subVector : subVectors) { | ||
| auto elementsOp = |
There was a problem hiding this comment.
| for (const Value subVector : subVectors) { | |
| auto elementsOp = | |
| for (Value subVector : subVectors) { | |
| auto elementsOp = |
| // We have the following comment in VectorConvertToElementOp | ||
| // | ||
| // // Input vectors of size 1 are converted to scalars by the type converter. | ||
| // // We cannot use `spirv::CompositeExtractOp` directly in this case. | ||
| // // For a scalar source, the result is just the scalar itself. | ||
| // | ||
| // Which in this case means an unrealized conversion cast. |
There was a problem hiding this comment.
nit: maybe let's not copy comments around -- single-element vectors becoming scalars are a well known phenomenon in the other spirv conversion tests
There was a problem hiding this comment.
Thanks! Will delete. Just pointing out that what I found interesting is that run-signature-conversion=false and the signature was still changed somewhat, which is counterintuitive.
There was a problem hiding this comment.
Ah sorry, in this particular case what was interesting is not the signature changing (that happens in the first test). In this one is the builtin.unrealized_conversion_cast that happens because the TypeConverter will automatically convert vector<1xf32> into an f32.
| %0:5 = vector.to_elements %arg0 : vector<5xf32> | ||
|
|
||
| // CHECK-NEXT: spirv.ReturnValue %[[RETVAL]] : f32 | ||
| return %0#0 : f32 |
There was a problem hiding this comment.
maybe show what happens when some other result is returned -- this way we could check that we extract the collect element
| UnrollStorePattern, UnrollBroadcastPattern, UnrollFromElements, | ||
| UnrollToElements, UnrollStepPattern>(patterns.getContext(), | ||
| options, benefit); | ||
| UnrollToElements, UnrollStepPattern, ToElementsToTargetShape>( |
There was a problem hiding this comment.
Why did you chose a different name from all the other patterns? (I'm not saying this is a bad name, just that it creates some asymmetry)
There was a problem hiding this comment.
There is already an UnrollToElements pattern. I prefer small simple patterns over larger ones. I'll turn this back to a draft for the time being.
I'm thinking about maybe moving this pattern somewhere else. While working on adding the analogous to FromElementsOp I found that there's an issue. The issue is that I was attempting to break down FromElementsOp into multiple FromElementsOp with a suitable vector length and then reconstruct the target type with InsertOp, but the canonicalizer reverts this via the InsertChainFullyInitialized pattern.
This only happens when the targetShape is {1} which can happen when the vector in the IR is not divisible by one of the native sizes. For example, the native vector size is 4 and the vector in the IR is vector<5xf32> then the current getTargetShape will suggest rewriting it to 5 vector<1xf32>.
This I believe comes from:
int mlir::spirv::getComputeVectorSize(int64_t size) {
for (int i : {4, 3, 2}) {
if (size % i == 0)
return i;
}
return 1;
}When I go around this issue, I get builtin.unrealized_conversion_casts.
I'm still thinking about whether there is another option. (For example, replace vector<5xf32> with two vector<4xf32> and then extract what's needed from the other one.
| // CHECK: %[[ELEMS0:.+]]:4 = vector.to_elements %[[V0]] | ||
| // CHECK: %[[ELEMS1:.+]]:4 = vector.to_elements %[[V1]] | ||
| // CHECK: return %[[ELEMS0]]#3, %[[ELEMS1]]#0 | ||
| %0 = "test.op"() : () -> (vector<8xf32>) |
There was a problem hiding this comment.
why not make it the function argument?
| // COM: Here we are testing the pattern ToElementsToTargetShape | ||
| // COM: The pattern has a native shape of [4], which means | ||
| // COM: that vectors multiples of 4 will be split. In this | ||
| // COM: case, that will happen in the function's body, not the argument. | ||
|
|
||
| // CHECK-LABEL: func.func @unroll_vector_8xf32 | ||
| // CHECK-SAME: (%[[ARG0:.+]]: vector<8xf32>) | ||
| func.func @unroll_vector_8xf32(%arg0: vector<8xf32>) -> (f32, f32) { | ||
| %0:8 = vector.to_elements %arg0 : vector<8xf32> | ||
|
|
||
| // COM: We only return two elements, one from each of the | ||
| // COM: vectors. |
2 participants
Depends on #166462
Pattern
ToElementsToTargetShapesplits operands tovector.to_elementsoperations such that they are a multiple of the target shape. This allowsvector.to_elementsoperations of vectors sizes other than 2, 3, 4 to lower to SPIR-V.