[CuTeDSL] Allow reassoc fastmath flag on vector reduction ops - #3384
[CuTeDSL] Allow reassoc fastmath flag on vector reduction ops#3384antoniofrighetto wants to merge 2 commits into
reassoc fastmath flag on vector reduction ops#3384Conversation
Let Vector.reduce() accept optional fastmath flags (e.g., `reassoc`), meant to be leveraged by callers during MLIR/LLVM lowering.
brandon-yujie-sun
left a comment
brandon-yujie-sun
left a comment
There was a problem hiding this comment.
Thanks for contributing the codes! LGTM with some minor questions :-)
| kind, self, acc=p_acc, reduction_dims=partial_dims, loc=loc, ip=ip | ||
| ) | ||
| raw = vector.reduction(elem_ty, kind, vec_1d, acc=acc, loc=loc, ip=ip) | ||
| raw = vector.reduction(elem_ty, kind, vec_1d, acc=acc, **fmf_kwargs, loc=loc, ip=ip) |
There was a problem hiding this comment.
This looks good, but for the multi_reduction case, shall we raise error to mention it's not supported?
There was a problem hiding this comment.
Added, thanks. Raising ValueError when handling multi-dimensional reductions earlier (so dropped passing the flag here too; not sure if we somehow meant to only reject partial reductions, and allow passing the flag here).
There was a problem hiding this comment.
Thanks for the update. LGTM
| module = ir.Module.create() | ||
| vec_ty = ir.VectorType.get([16], ir.F32Type.get()) | ||
| with ir.InsertionPoint(module.body): | ||
| fn = func.FuncOp("test", ir.FunctionType.get([vec_ty], [])) |
There was a problem hiding this comment.
I think we'd also need a func.ReturnOp([]) to make it a real valid func op?
| fn = func.FuncOp("test", ir.FunctionType.get([vec_ty], [])) | ||
| with ir.InsertionPoint(fn.add_entry_block()): | ||
| Vector(fn.arguments[0]).reduce("add", fastmath=fastmath) | ||
| return str(module) |
There was a problem hiding this comment.
Maybe do a module.operation.verify() before returning to make sure the constructed op is valid so that the validation on the serialized string would be meaningful?
2 participants
Let Vector.reduce() accept optional fastmath flags (e.g.,
reassoc), meant to be leveraged by callers during MLIR/LLVM lowering.