[TKW] Rework vector mask generation #172
Conversation
Signed-off-by: Ivan Butygin <ivan.butygin@gmail.com>
Signed-off-by: Ivan Butygin <ivan.butygin@gmail.com>
| # CHECK-SAME: memref<1x3xf16, strided<[3, 1], offset: ?>> | ||
| # CHECK: vector.maskedstore %[[D27]][%[[D5]], %[[D8]]], %[[D25]], %[[D26]] : | ||
| # CHECK-SAME: memref<1x3xf16, strided<[3, 1], offset: ?>>, vector<4xi1>, vector<4xf16> | ||
| # CHECK-DAG: %[[CST:.*]] = arith.constant dense<0.000000e+00> : vector<4xf16> |
| ) -> Optional[list[IndexExpr]]: | ||
| bounds = [] | ||
| for constraint in constraints: | ||
| if not isinstance(constraint, (WorkgroupConstraint, TilingConstraint)): |
There was a problem hiding this comment.
Why do we ignore the WaveConstraints here? Does that mean that we assuming that the workgroup tile size is a multiple of the wave tile size?
There was a problem hiding this comment.
Yeah, I implicitly assumed WG size is divisible by wave size, do we have any potential examples when it's false?
There was a problem hiding this comment.
Good question. I was just thinking in terms of generality (so workgroup tile size = 27 and wave tile size = 17) but maybe its not such a common use case. We can ignore for now.
There was a problem hiding this comment.
As I am reading through some of the history, allow me to share a thought.
It is generally good practice to be as exhaustive and general as possible and failing loudly.
This way, in the future we can come back and immediately understand the problem and that it is NYI.
In the the current form, it seems to me the case proposed by Harsh would silently pass but generate wrong code?
|
|
||
| pos = arith_d.ConstantOp(IndexType.get(), i) | ||
| mask = vector_d.insertelement(cond, mask, position=pos) | ||
| mask_expr = functools.reduce( |
| @@ -171,6 +172,32 @@ def get_type_or_element_type(operand_type: IrType): | |||
| def gen_sympy_index(emitter: WaveEmitter, expr: sympy.Expr) -> OpResult: | |||
| stack: list[OpResult] = [] | |||
|
|
|||
| def _broadcast(a, b): | |||
There was a problem hiding this comment.
Seems like you can refactor this to avoid the duplication for a and b.
| mask_expr = functools.reduce( | ||
| lambda a, b: sympy.And(a, b), (new_index[dim] < dim for dim in bounds) | ||
| ) | ||
| mask = gen_sympy_index(emitter, mask_expr) |
There was a problem hiding this comment.
No action needed, but just putting this down for the record that at some point we should evaluate generating these masks prior to codegen and evaluate the performance impact.
Instead of generating individual element comparisons and doing `vector.insertelement` generate the whole mask using vector ops. Add support for vector codegen when generating MLIR IR from sympy expressions. Add method `IndexingContext.iota` to generate special symbols which map to `(1,2 ... n-1)` vec expressions. `gen_sympy_index` will start to generate vector ops when encountering such symbols, inserting proper `splat`'s between scalar vals when necessary. --------- Signed-off-by: Ivan Butygin <ivan.butygin@gmail.com>
Instead of generating individual element comparisons and doing `vector.insertelement` generate the whole mask using vector ops. Add support for vector codegen when generating MLIR IR from sympy expressions. Add method `IndexingContext.iota` to generate special symbols which map to `(1,2 ... n-1)` vec expressions. `gen_sympy_index` will start to generate vector ops when encountering such symbols, inserting proper `splat`'s between scalar vals when necessary. --------- Signed-off-by: Ivan Butygin <ivan.butygin@gmail.com> Signed-off-by: Ian <ian.nordeng@amd.com>
Instead of generating individual element comparisons and doing `vector.insertelement` generate the whole mask using vector ops. Add support for vector codegen when generating MLIR IR from sympy expressions. Add method `IndexingContext.iota` to generate special symbols which map to `(1,2 ... n-1)` vec expressions. `gen_sympy_index` will start to generate vector ops when encountering such symbols, inserting proper `splat`'s between scalar vals when necessary. --------- Signed-off-by: Ivan Butygin <ivan.butygin@gmail.com>
Instead of generating individual element comparisons and doing
vector.insertelementgenerate the whole mask using vector ops.Add support for vector codegen when generating MLIR IR from sympy expressions. Add method
IndexingContext.iotato generate special symbols which map to(1,2 ... n-1)vec expressions.gen_sympy_indexwill start to generate vector ops when encountering such symbols, inserting propersplat's between scalar vals when necessary.