This is an archive of the discontinued LLVM Phabricator instance.

[mlir] introduce "encoding" attribute to tensor type
ClosedPublic

Authored by aartbik on Mar 29 2021, 3:41 PM.

Details

Summary

This CL introduces a generic attribute (called "encoding") on tensors.
The attribute currently does not carry any concrete information, but the type
system already correctly determines that tensor<8xi1,123> != tensor<8xi1,321>.
The attribute will be given meaning through an interface in subsequent CLs.

See ongoing discussion on discourse:

[RFC] Introduce a sparse tensor type to core MLIR
https://llvm.discourse.group/t/rfc-introduce-a-sparse-tensor-type-to-core-mlir/2944

A sparse tensor will look something like this:

// named alias with all properties we hold dear:
#CSR = {
  // individual named attributes
}

// actual sparse tensor type:
tensor<?x?xf64, #CSR>

I see the following rough 5 step plan going forward:

(1) introduce this format attribute in this CL, currently still empty
(2) introduce attribute interface that gives it "meaning", focused on sparse in first phase
(3) rewrite sparse compiler to use new type, remove linalg interface and "glue"
(4) teach passes to deal with new attribute, by rejecting/asserting on non-empty attribute as simplest solution, or doing meaningful rewrite in the longer run
(5) add FE support, document, test, publicize new features, extend "format" meaning to other domains if useful

Diff Detail

Event Timeline

aartbik created this revision.Mar 29 2021, 3:41 PM
aartbik requested review of this revision.Mar 29 2021, 3:41 PM
aartbik updated this revision to Diff 334007.Mar 29 2021, 3:53 PM

fixed typos

aartbik updated this revision to Diff 334012.Mar 29 2021, 4:23 PM

pass python bindings tests

aartbik retitled this revision from [WIP][mlir] introduce "format" attribute to tensor type to [mlir] introduce "format" attribute to tensor type.Mar 30 2021, 1:19 PM
aartbik edited the summary of this revision. (Show Details)
aartbik edited the summary of this revision. (Show Details)
aartbik edited the summary of this revision. (Show Details)

any early feedback on this? can I go ahead polishing this CL for review?

Sorry for the delay - figured someone else was more authoritative than me and would take this. Given the heavy use of the C-API function, I'd be tempted to create an extended version and possibly offer a static/inline fallback with the current name/signature. Could be convinced either way, though.

mlir/include/mlir-c/BuiltinTypes.h
193

I believe this may be the first optional argument in the CAPI. I think you are going to want to provide a static inline accessor mlirAttributeGetNull() which zero initializes the MlirAttribute properly. Note in the comments that this function is to be used to construct without a format.

We aren't providing stability guarantees on this API yet, but this will definitely be a backwards incompatible change. I'm open to it, but these are also very commonly used functions in the 3-arg variant: I might be tempted to create an mlirRankedTensorTypeGetExt if we think we are going to grow more such things in the future (beyond a vanilla tensor).

mlir/lib/Bindings/Python/IRAttributes.cpp
506

I don't think this is zero initializing. Use the mlirAttributeGetNull as requested above to initialize it.

mlir/lib/Bindings/Python/IRTypes.cpp
386

And here.

mlir/test/CAPI/ir.c
690–691

Use mlirAttributeGetNull

jpienaar added inline comments.
mlir/include/mlir/IR/BuiltinTypes.td
701

One of the open questions (perhaps it got resolved) was how this attribute plays with existing rewrites. Today we have many ops that say they supports tensor type, post this whether or not they support a tensor type is unclear (depends on the attribute) and that could also affect rewrites (e.g., lowering could now result in invalid IR as the type supported isn't uniform now, the high level op may support sparsity, the low level ones may not). But the next question may perhaps answer that.

Is this format allowed to be dropped? E.g., is this is merely describing structure of the contents?

aartbik marked 4 inline comments as done.Apr 7 2021, 2:47 PM

figured someone else was more authoritative than me and would take this.

Can't think of many others more authoritative than you on this subject!

mlir/include/mlir-c/BuiltinTypes.h
193

Okay, I added the mlirAttributeGetNull() method and mentioned this in the doc.
Are "default" values for paremeters an option in this CAPI?

mlir/include/mlir/IR/BuiltinTypes.td
701

Is this format allowed to be dropped?

The short answer here is "no".

More elaborate, the discussion in discourse (https://llvm.discourse.group/t/mlir-support-for-sparse-tensors/2020) converged on either a new tensor type (so that passes need to explicitly learn about this) or an attribute on the existing tensor type (where passes gradually need to learn about dealing with the semantics of the attribute). This CL goes into the latter direction.

I see the following rought 5 step plan going forward
(1) introduce this format attribute, currently still empty
(2) introduce attribute interface that gives it "meaning", focused on sparse in first phase
(3) rewrite sparse compiler to use new type, remove linalg interface and "glue"
(4) teach passes to deal with new attribute, by rejecting/asserting on non-empty attribute as simplest solution, or doing meaningful rewrite in the longer run
(5) add FE support, document, test, publicize new features, extend "format" meaning to other domains if useful

mlir/lib/Bindings/Python/IRAttributes.cpp
506

Ah, good point. These corners of the code are still a bit new to me. Thanks!

aartbik updated this revision to Diff 335932.Apr 7 2021, 2:47 PM
aartbik marked 2 inline comments as done.

first round of comments

mehdi_amini added inline comments.Apr 7 2021, 4:16 PM
mlir/include/mlir/IR/BuiltinTypes.td
701

Can you document this all in this patch?

aartbik marked an inline comment as done.Apr 7 2021, 7:18 PM
aartbik added inline comments.
mlir/include/mlir/IR/BuiltinTypes.td
701

You mean this five step plan above in the description? If so, of course!

aartbik edited the summary of this revision. (Show Details)Apr 7 2021, 7:19 PM
aartbik edited the summary of this revision. (Show Details)Apr 7 2021, 7:20 PM
mehdi_amini added inline comments.Apr 7 2021, 7:35 PM
mlir/include/mlir/IR/BuiltinTypes.td
701

Actually I meant more the attribute semantics, there is documentation that describes the tensor type and we should make sure everything is consistent and we think it through (for example: https://mlir.llvm.org/docs/Dialects/Builtin/#rankedtensortype )

bondhugula requested changes to this revision.Apr 7 2021, 9:00 PM
bondhugula added a subscriber: bondhugula.

This revision is missing an update to the tensor type documentation in mlir/doc/LangRef.md.

This revision now requires changes to proceed.Apr 7 2021, 9:00 PM
aartbik marked an inline comment as done.Apr 8 2021, 12:45 PM

This revision is missing an update to the tensor type documentation in mlir/doc/LangRef.md.

What part in LangRef.md do you want changed? Note that the tensor type doc is generated from the doc in "include/mlir/IR/BuiltinTypes.td", which is part of the CL already
(in the sense that I have a TODO, but without attribute interface I am not sure how much I can say)

mlir/include/mlir/IR/BuiltinTypes.td
701

Does that mean you want the attribute interface (and doc) as part of this CL?

mehdi_amini added inline comments.Apr 8 2021, 2:02 PM
mlir/include/mlir/IR/BuiltinTypes.td
701

Actually I was looking for changes in doc/ but I missed that the doc being now generated from TableGen there may not be anything to update separately (Yay!)!

So I'm fine with what you have right now!

aartbik marked an inline comment as done.Apr 8 2021, 2:28 PM

Uday, can you please verify that you are happy with the doc changes (autogen from td) that describe the format syntax.
Next are the interfaces (happy to include these here, but probably easier to review as next CL)....

mlir/include/mlir/IR/BuiltinTypes.td
701

\O/

silvas added inline comments.Apr 8 2021, 4:14 PM
mlir/include/mlir/IR/BuiltinTypes.td
658

Is there and advantage to having this be a "format" instead of a generic attribute with any user-specific meaning?

For example, I have a use case that could be served quite naturally by this feature if I could use a custom attribute of my own that keeps some per-dimension metadata such as whether a dimension is known to participate in numpy "size 1" broadcasting (which is a critical piece of information for numpy-based tensor frontends). One can imagine extending this to arbitrary static information that user ops are compatible with.

Also, historically, the current tensor type (equivalent of empty attribute after your patch) is not assumed to have a "straightforward, dense" layout. It is intentionally completely opaque, and from the beginning we have thought of tensor as potentially being implemented in many different ways (dense, sparse, hash table, tree, ...). @jpienaar for more info here. Some passes *lower* the opaque tensor type to straightforward dense layout (e.g. dense bufferization), but there is nothing *requiring* it semantically.

aartbik marked an inline comment as done.Apr 8 2021, 4:52 PM
aartbik added inline comments.
mlir/include/mlir/IR/BuiltinTypes.td
658

The way I see this, "format" is the generic attribute we were talking about (see https://llvm.discourse.group/t/rfc-introduce-a-sparse-tensor-type-to-core-mlir/2944), where the semantics are defined by attribute interfaces. So in that case, the sparse extension stuff could use something like this

#CSR = {

dim_order  = affine_map<(i,j) -> (i,j)>,
dim_format = [ "D", "S" ],
pointers = i64,
indices = i64

}

with sparse tensor type

%a : tensor<?x?xf64, #CSR>

and the attribute interface would provide ways to access dim_order, dim_format, etc.
and you could easily extend the meaning by adding a new "sean_attr" attribute in this like

#X = {

sean_attr = [ num_py_size_1_stuff ];

}

%a : tensor<?x?xf64, #X>

Did you have something else in mind?

aartbik added inline comments.Apr 8 2021, 4:58 PM
mlir/include/mlir/IR/BuiltinTypes.td
658

So in the example above, the tensor really looks like this:

tensor<?x?xf64, {dim_format = ["D", "S"], dim_order = affine_map<(d0, d1) -> (d0, d1)>, indices = i64, pointers = i64}>
silvas added inline comments.Apr 9 2021, 10:59 AM
mlir/include/mlir/IR/BuiltinTypes.td
658

Thanks Aart. Indeed, this seems like it addresses my needs.

I guess my main concern is the naming and documentation. "format" sounds very specific to tensor layout things (and so does the documentation). Maybe "metadata"? "properties"? Or perhaps just "attr"?

(would love to hear others' thoughts)

As a strawman, we could call it "attr" and the documentation would read "if attr is present, it represents domain-specific information associated with this tensor. Transformations that do not have the necessary domain-specific understanding of a tensor's attr should not modify the tensor or anything operating on it."

(for example, this wording avoids mentioning "denseness" as some sort of default characteristic of a tensor -- it just so happens that we have passes that lower tensors to dense code by default, and this attr provides a mechanism for preventing those transformations from kicking in)

aartbik added inline comments.Apr 9 2021, 11:26 AM
mlir/include/mlir/IR/BuiltinTypes.td
658

Yes, happy to change the name into something we can all agree on. I had originally "layout", did not like that, changed it to "format", but have absolutely no strong opinion on this.

Just "attr" seems a bit too vague, imho, but something like "properties" seems right.

Let's see what others say.

properties LGTM as well.

silvas added a subscriber: lattner.Apr 9 2021, 1:13 PM

properties LGTM as well.

The reason I (very mildly -- just trying to stimulate discussion here) don't like properties is that it sounds like something that can be ignored if I don't know about it. (at least attr sounds like an opaque thing I have to be conservative around).

Looking in a thesaurus, disposition and distinction seem possible candidates.

@lattner tends to have a good intuition for these naming things. Chris, what do you think the name of the attribute attached to tensors should be?

Oh awesome, I'm thrilled to see this happening.

I'd recommend going with something active sounding, perhaps layout or encoding? I agree that passive words like properties are probably not the right thing.

-Chris

I'd recommend going with something active sounding, perhaps layout or encoding?

Since layout may discourage future extensions beyond stuff related to layout, I will go with encoding since that indeed sounds active, and easily allows for future extensions....

aartbik updated this revision to Diff 336577.Apr 9 2021, 3:54 PM

format -> encoding

aartbik retitled this revision from [mlir] introduce "format" attribute to tensor type to [mlir] introduce "encoding" attribute to tensor type.Apr 9 2021, 3:54 PM
aartbik edited the summary of this revision. (Show Details)
stellaraccident accepted this revision.Apr 9 2021, 4:55 PM

API changes lgtm. Sounds like the rest of it has converged but I'm not as looped in on the semantics change.

silvas added a comment.Apr 9 2021, 5:00 PM

"encoding" SGTM.

aartbik updated this revision to Diff 336585.Apr 9 2021, 5:08 PM

renamed tests too

bondhugula accepted this revision.Apr 9 2021, 8:17 PM

This revision is missing an update to the tensor type documentation in mlir/doc/LangRef.md.

What part in LangRef.md do you want changed? Note that the tensor type doc is generated from the doc in "include/mlir/IR/BuiltinTypes.td", which is part of the CL already
(in the sense that I have a TODO, but without attribute interface I am not sure how much I can say)

Sorry, I meant the doc in BuiltinTypes.td (https://mlir.llvm.org/docs/Dialects/Builtin/#rankedtensortype), and you already have it. Just some minor typos I caught there.

mlir/include/mlir/IR/BuiltinTypes.td
664

these -> this
?

Missing period at end of sentence.

This revision is now accepted and ready to land.Apr 9 2021, 8:17 PM

I like structure too as it captures information about the structure of the values to me. encoding makes me think jpeg or Levensthein distance and the like.

mlir/include/mlir/IR/BuiltinTypes.td
658

As Sean mentioned Tensor has no specific representation today, it is just a value. It need not even reside in memory. The lowering is purely a convention of a given lowering. So tensors could all be encoded as gzip'd in memory as long as it is consistent and what the lowering expects. Which is why I asked about required, as it could be used as additional information that a given lowering strategy could use rather than an explicit restriction. E.g., additional information that lowering could be use, but any consistent lowering is valid as today.

This is often the gap where there is some assumption that the lowering from from tensor to memref is fixed, rather than a family. And the attribute here represents an approach to go from opaque type to explicit in memory representation with memory space, layout etc "step wise" IMHO.

701

My concern is more towards existing passes and transforms. Today you optimizations that says transform from foo.X to bar.Y if element type is supported by bar ops, all of those would be needed to be constrained if this is a required attribute. So too build methods that today just produce a tensor type but post this would need to be able to figure out the encoding given the input encodings. Those are correctness issues that seems like they should be handled either as part of (1) or as (2). These can be made restrictive as you mention, but I'd like to see that handled before we actually start allowing it to be populated. This could be done easily ODS and DRR side by constraining Tensor to have empty encoding (which is same as introducing new type at that level).

Making it required is akin to XLA pre- and post-layout distinction, now one is adding an explicit barrier between where certain optimization passes can only be run before the layout is assigned while others after. Perhaps we can avoid that here.

properties LGTM as well.

The reason I (very mildly -- just trying to stimulate discussion here) don't like properties is that it sounds like something that can be ignored if I don't know about it. (at least attr sounds like an opaque thing I have to be conservative around).

To me there is nothing about property that makes it discardable, quite the contrary actually, Wikipedia says: "a property is a characteristic of an object; a red object is said to have the property of redness", which matches how I see it: the "property" of an object is really about what is intrinsic to the object and everything to describe it.

I'm fine with "encoding" as well though (I am not convinced with "structure" as someone mentioned for example some "quantization properties" that could be put there which isn't as much about "structure" to me).

aartbik marked 3 inline comments as done.Apr 12 2021, 9:35 AM

To me there is nothing about property that makes it discardable, quite the contrary actually,

I was with Mehdi on this one. properties sounded very non-discardable ;-)
But based on Chris' input, I stick with encoding.

mlir/include/mlir/IR/BuiltinTypes.td
701

With (2), I can start adding some more verification. For some passes it may indeed simply require that the contents are empty for the time being.

aartbik updated this revision to Diff 336876.Apr 12 2021, 9:54 AM
aartbik marked an inline comment as done.

typos

This revision was automatically updated to reflect the committed changes.
jpienaar added inline comments.Jul 24 2021, 5:55 AM
mlir/include/mlir/IR/BuiltinTypes.td
701

Coming back to this: was verification added? I believe as is today all operations on all dialects that operate on tensor is considered to support any arbitrary encoding attribute (at least no verification that checks that) and passes expected to retain it and know how to handle it (e.g., if I call build or rewrite, then the build method needs to understand it, shape inference can't drop it, so needs to know how to handle, etc). But I don't think that is true, we would drop it/not know how to handle this multiple places. This is also akin to what Stella mentioned in original discussion on migration efforts.