This is an archive of the discontinued LLVM Phabricator instance.

[mlir][transform] Allow arbitrary indices to be scalable

Authored by awarzynski on Jul 3 2023, 4:50 AM.



This change lifts the limitation that only the trailing dimensions/sizes
in dynamic index lists can be scalable. It allows us to extend
MaskedVectorizeOp and TileOp from the Transform dialect so that the
following is allowed:

%1, %loops:3 = transform.structured.tile %0 [4, [4], [4]]

This is also a follow up for
that will enable the following (middle vector dimension is scalable):

transform.structured.masked_vectorize %0 vector_sizes [2, [4], 8]

To facilate this change, the hooks for parsing and printing dynamic
index lists are updated accordingly (printDynamicIndexList and
parseDynamicIndexList, respectively). MaskedVectorizeOp and TileOp
are updated to include an array of attribute of bools that captures
whether the corresponding vector dimension/tile size, respectively, are
scalable or not.

This change is a part of a larger effort to enable scalable
vectorisation in Linalg. See this RFC for more context:

Diff Detail

Event Timeline

awarzynski created this revision.Jul 3 2023, 4:50 AM
Herald added a project: Restricted Project. · View Herald TranscriptJul 3 2023, 4:50 AM
awarzynski requested review of this revision.Jul 3 2023, 4:50 AM
ftynse accepted this revision.Jul 3 2023, 4:52 AM
This revision is now accepted and ready to land.Jul 3 2023, 4:52 AM
c-rhodes accepted this revision.Jul 3 2023, 5:27 AM

I've left a few bits but otherwise LGTM, cheers


nit: I count 5 different variable names now for the same concept:

  • scalableDims
  • scalable_sizes
  • scalableVals
  • scalableIdxs
  • scalables

I think it would be clearer to use consistent naming.


If the intention here is to check there are scalableVals then not empty would be clearer.

Thanks Cullen, I will make the suggested adjustments before merging.


Meh, naming is hard. I agree, but IMHO it's a wider problem and not directly related to this patch.

  • The definition of MaskedVectorizeOp contains $vector_sizes and $static_sizes This patch adds scalable_sizes, which is consistent with the naming scheme in
  • The signature of printDynamicIndexList contains values and integers. This patch adds scalables, which is consistent with what's already there.

You are right that the names in *.td and *.cpp files are inconsistent. However, this inconsistency predates this change. Having said that, scalableIdxs should be replaced with scalables (for consistency). I will update that.

scalableVals is a local function variable and is consistent with integerVals (per-exisiting function variable serving similar purpose).

scalableDims is only used in the context of VectorType in which the concept of a "dimension" makes sense (it does not make sense in the context of e.g. tiles).

Thanks for raising this! I will make a few adjustments that hopefully will reduce the confusion.


I swear you suggested something similar recently in one of my patches. I need to pay more attention - thank you!

The bot seems broken by this, can you have a look please?

Apologies for the breakage and thanks for reverting!

I keep MLIR_ENABLE_BINDINGS_PYTHON as default (off), so that's why I didn't see that :( will be AFK most of today, so reverting was the right call, thanks!

awarzynski reopened this revision.Jul 4 2023, 1:37 PM

So it looks like the failure was due to missing changes in the python bindings - I am about to send an update.

I will go ahead and re-apply this patch sometime tomorrow unless there are new comments.

This revision is now accepted and ready to land.Jul 4 2023, 1:37 PM
awarzynski updated this revision to Diff 537169.Jul 4 2023, 1:40 PM

Update Python bindings