[mlir] introduce "encoding" attribute to tensor type

[aartbik]

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

[aartbik]

fixed typos

[aartbik]

pass python bindings tests

[aartbik]

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

[stellaraccident]

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]

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]

In D99548#2672931, @stellaraccident wrote:

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]

first round of comments

[mehdi_amini]

mlir/include/mlir/IR/BuiltinTypes.td

701

Can you document this all in this patch?

[aartbik]

mlir/include/mlir/IR/BuiltinTypes.td

701

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

[mehdi_amini]

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]

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

[aartbik]

In D99548#2675771, @bondhugula wrote:

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]

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]

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]

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]

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]

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]

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]

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.

[mehdi_amini]

properties LGTM as well.

[silvas]

In D99548#2680345, @mehdi_amini wrote:

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?

[lattner]

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

[aartbik]

In D99548#2680605, @lattner wrote:

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....

[lattner]

👍

[aartbik]

format -> encoding

[stellaraccident]

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

[silvas]

"encoding" SGTM.

[aartbik]

renamed tests too

[bondhugula]

In D99548#2677630, @aartbik wrote:

In D99548#2675771, @bondhugula wrote:

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.

[jpienaar]

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.

[mehdi_amini]

In D99548#2680579, @silvas wrote:

In D99548#2680345, @mehdi_amini wrote:

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]

In D99548#2681226, @mehdi_amini wrote:

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]

typos

[jpienaar]

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.

mlir/include/mlir-c/BuiltinAttributes.h

 #include "mlir-c/AffineMap.h"
 #include "mlir-c/IR.h"
 #include "mlir-c/Support.h"
 
 #ifdef __cplusplus
 extern "C" {
 #endif
 
+/// Returns an empty attribute.
+MLIR_CAPI_EXPORTED MlirAttribute mlirAttributeGetNull();
+
 //===----------------------------------------------------------------------===//
 // Affine map attribute.
 //===----------------------------------------------------------------------===//
 
 /// Checks whether the given attribute is an affine map attribute.
 MLIR_CAPI_EXPORTED bool mlirAttributeIsAAffineMap(MlirAttribute attr);
 
 /// Creates an affine map attribute wrapping the given map. The attribute

mlir/include/mlir-c/BuiltinTypes.h

 MLIR_CAPI_EXPORTED bool mlirTypeIsATensor(MlirType type);
 
 /// Checks whether the given type is a ranked tensor type.
 MLIR_CAPI_EXPORTED bool mlirTypeIsARankedTensor(MlirType type);
 
 /// Checks whether the given type is an unranked tensor type.
 MLIR_CAPI_EXPORTED bool mlirTypeIsAUnrankedTensor(MlirType type);
 
-/// Creates a tensor type of a fixed rank with the given shape and element type
-/// in the same context as the element type. The type is owned by the context.
+/// Creates a tensor type of a fixed rank with the given shape, element type,
+/// and optional encoding in the same context as the element type. The type is
+/// owned by the context. Tensor types without any specific encoding field
+/// should assign mlirAttributeGetNull() to this parameter.
 MLIR_CAPI_EXPORTED MlirType mlirRankedTensorTypeGet(intptr_t rank,
                                                     const int64_t *shape,
-                                                    MlirType elementType);
+                                                    MlirType elementType,
+                                                    MlirAttribute encoding);
 
 /// Same as "mlirRankedTensorTypeGet" but returns a nullptr wrapping MlirType on
 /// illegal arguments, emitting appropriate diagnostics.
-MLIR_CAPI_EXPORTED MlirType
-mlirRankedTensorTypeGetChecked(MlirLocation loc, intptr_t rank,
-                               const int64_t *shape, MlirType elementType);
+MLIR_CAPI_EXPORTED MlirType mlirRankedTensorTypeGetChecked(
+    MlirLocation loc, intptr_t rank, const int64_t *shape, MlirType elementType,
+    MlirAttribute encoding);
 
 /// Creates an unranked tensor type with the given element type in the same
 /// context as the element type. The type is owned by the context.
 MLIR_CAPI_EXPORTED MlirType mlirUnrankedTensorTypeGet(MlirType elementType);
 
 /// Same as "mlirUnrankedTensorTypeGet" but returns a nullptr wrapping MlirType
 /// on illegal arguments, emitting appropriate diagnostics.
 MLIR_CAPI_EXPORTED MlirType

mlir/include/mlir/IR/BuiltinTypes.td

 //===----------------------------------------------------------------------===//
 
 def Builtin_RankedTensor : Builtin_Type<"RankedTensor", "TensorType"> {
   let summary = "Multi-dimensional array with a fixed number of dimensions";
   let description = [{
     Syntax:
 
     ```
-    tensor-type ::= `tensor` `<` dimension-list type `>`
+    tensor-type ::= `tensor` `<` dimension-list type (`,` encoding)? `>`
     dimension-list ::= (dimension `x`)*
     dimension ::= `?` | decimal-literal
+    encoding ::= attribute-value
     ```
 
     Values with tensor type represents aggregate N-dimensional data values, and
     have a known element type and a fixed rank with a list of dimensions. Each
     dimension may be a static non-negative decimal constant or be dynamically
     determined (indicated by `?`).
 
     The runtime representation of the MLIR tensor type is intentionally
     abstracted - you cannot control layout or get a pointer to the data. For
     low level buffer access, MLIR has a [`memref` type](#memref-type). This
     abstracted runtime representation holds both the tensor data values as well
     as information about the (potentially dynamic) shape of the tensor. The
     [`dim` operation](Dialects/Standard.md#dim-operation) returns the size of a
     dimension from a value of tensor type.
 
+    The `encoding` attribute provides additional information on the tensor.
+    An empty attribute denotes a straightforward tensor without any specific
+    structure. But particular properties, like sparsity or other specific
+    characteristics of the data of the tensor can be encoded through this
+    attribute. The semantics are defined by a type and attribute interface
+    and must be respected by all passes that operate on tensor types.
+    TODO: provide this interface, and document it further.
+
     Note: hexadecimal integer literals are not allowed in tensor type
     declarations to avoid confusion between `0xf32` and `0 x f32`. Zero sizes
     are allowed in tensors and treated as other sizes, e.g.,
     `tensor<0 x 1 x i32>` and `tensor<1 x 0 x i32>` are different types. Since
     zero sizes are not allowed in some other types, such tensors should be
     optimized away before lowering tensors to vectors.
 
     Examples:
let description = [{
     // Tensor with rank zero. Represents a scalar.
     tensor<f32>
 
     // Zero-element dimensions are allowed.
     tensor<0 x 42 x f32>
 
     // Zero-element tensor of f32 type (hexadecimal literals not allowed here).
     tensor<0xf32>
+
+    // Tensor with an encoding attribute (where #ENCODING is a named alias).
+    tensor<?x?xf64, #ENCODING>
     ```
   }];
   let parameters = (ins
     ArrayRefParameter<"int64_t">:$shape,
-    "Type":$elementType
+    "Type":$elementType,
+    "Attribute":$encoding
   );
 
   let builders = [
     TypeBuilderWithInferredContext<(ins
-      "ArrayRef<int64_t>":$shape, "Type":$elementType
+      "ArrayRef<int64_t>":$shape,
+      "Type":$elementType,
+      CArg<"Attribute", "{}">:$encoding
     ), [{
-      return $_get(elementType.getContext(), shape, elementType);
+      return $_get(elementType.getContext(), shape, elementType, encoding);
     }]>
   ];
   let skipDefaultBuilders = 1;
   let genVerifyDecl = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // TupleType

mlir/lib/Bindings/Python/IRAttributes.cpp

 private:
   template <typename ElementTy>
   static MlirAttribute
   bulkLoad(MlirContext context,
            MlirAttribute (*ctor)(MlirType, intptr_t, ElementTy *),
            MlirType mlirElementType, py::buffer_info &arrayInfo) {
     SmallVector<int64_t, 4> shape(arrayInfo.shape.begin(),
                                   arrayInfo.shape.begin() + arrayInfo.ndim);
-    auto shapedType =
-        mlirRankedTensorTypeGet(shape.size(), shape.data(), mlirElementType);
+    MlirAttribute encodingAttr = mlirAttributeGetNull();
+    auto shapedType = mlirRankedTensorTypeGet(shape.size(), shape.data(),
+                                              mlirElementType, encodingAttr);
     intptr_t numElements = arrayInfo.size;
     const ElementTy *contents = static_cast<const ElementTy *>(arrayInfo.ptr);
     return ctor(shapedType, numElements, contents);
   }
 
   static bool isUnsignedIntegerFormat(const std::string &format) {
     if (format.empty())
       return false;

mlir/lib/Bindings/Python/IRTypes.cpp

 //===- IRTypes.cpp - Exports builtin and standard types -------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "IRModule.h"
 
 #include "PybindUtils.h"
 
+#include "mlir-c/BuiltinAttributes.h"
 #include "mlir-c/BuiltinTypes.h"
 
 namespace py = pybind11;
 using namespace mlir;
 using namespace mlir::python;
 
 using llvm::SmallVector;
 using llvm::Twine;
public:
   static constexpr const char *pyClassName = "RankedTensorType";
   using PyConcreteType::PyConcreteType;
 
   static void bindDerived(ClassTy &c) {
     c.def_static(
         "get",
         [](std::vector<int64_t> shape, PyType &elementType,
            DefaultingPyLocation loc) {
+          MlirAttribute encodingAttr = mlirAttributeGetNull();
           MlirType t = mlirRankedTensorTypeGetChecked(
-              loc, shape.size(), shape.data(), elementType);
+              loc, shape.size(), shape.data(), elementType, encodingAttr);
           // TODO: Rework error reporting once diagnostic engine is exposed
           // in C API.
           if (mlirTypeIsNull(t)) {
             throw SetPyError(
                 PyExc_ValueError,
                 Twine("invalid '") +
                     py::repr(py::cast(elementType)).cast<std::string>() +
                     "' and expected floating point, integer, vector or "

mlir/lib/CAPI/IR/BuiltinAttributes.cpp

 #include "mlir/CAPI/AffineMap.h"
 #include "mlir/CAPI/IR.h"
 #include "mlir/CAPI/Support.h"
 #include "mlir/IR/Attributes.h"
 #include "mlir/IR/BuiltinTypes.h"
 
 using namespace mlir;
 
+MlirAttribute mlirAttributeGetNull() { return {nullptr}; }
+
 //===----------------------------------------------------------------------===//
 // Affine map attribute.
 //===----------------------------------------------------------------------===//
 
 bool mlirAttributeIsAAffineMap(MlirAttribute attr) {
   return unwrap(attr).isa<AffineMapAttr>();
 }
 

mlir/lib/CAPI/IR/BuiltinTypes.cpp

bool mlirTypeIsARankedTensor(MlirType type) {
   return unwrap(type).isa<RankedTensorType>();
 }
 
 bool mlirTypeIsAUnrankedTensor(MlirType type) {
   return unwrap(type).isa<UnrankedTensorType>();
 }
 
 MlirType mlirRankedTensorTypeGet(intptr_t rank, const int64_t *shape,
-                                 MlirType elementType) {
+                                 MlirType elementType, MlirAttribute encoding) {
   return wrap(RankedTensorType::get(
-      llvm::makeArrayRef(shape, static_cast<size_t>(rank)),
-      unwrap(elementType)));
+      llvm::makeArrayRef(shape, static_cast<size_t>(rank)), unwrap(elementType),
+      unwrap(encoding)));
 }
 
 MlirType mlirRankedTensorTypeGetChecked(MlirLocation loc, intptr_t rank,
                                         const int64_t *shape,
-                                        MlirType elementType) {
+                                        MlirType elementType,
+                                        MlirAttribute encoding) {
   return wrap(RankedTensorType::getChecked(
       unwrap(loc), llvm::makeArrayRef(shape, static_cast<size_t>(rank)),
-      unwrap(elementType)));
+      unwrap(elementType), unwrap(encoding)));
 }
 
 MlirType mlirUnrankedTensorTypeGet(MlirType elementType) {
   return wrap(UnrankedTensorType::get(unwrap(elementType)));
 }
 
 MlirType mlirUnrankedTensorTypeGetChecked(MlirLocation loc,
                                           MlirType elementType) {

mlir/lib/IR/AsmPrinter.cpp

TypeSwitch<Type>(type)
         os << "tensor<";
         for (int64_t dim : tensorTy.getShape()) {
           if (ShapedType::isDynamic(dim))
             os << '?';
           else
             os << dim;
           os << 'x';
         }
-        os << tensorTy.getElementType() << '>';
+        os << tensorTy.getElementType();
+        // Only print the encoding attribute value if set.
+        if (tensorTy.getEncoding()) {
+          os << ", ";
+          printAttribute(tensorTy.getEncoding());
+        }
+        os << '>';
       })
       .Case<UnrankedTensorType>([&](UnrankedTensorType tensorTy) {
         os << "tensor<*x";
         printType(tensorTy.getElementType());
         os << '>';
       })
       .Case<MemRefType>([&](MemRefType memrefTy) {
         os << "memref<";

mlir/lib/IR/BuiltinTypes.cpp

 }
 
 //===----------------------------------------------------------------------===//
 // RankedTensorType
 //===----------------------------------------------------------------------===//
 
 LogicalResult
 RankedTensorType::verify(function_ref<InFlightDiagnostic()> emitError,
-                         ArrayRef<int64_t> shape, Type elementType) {
+                         ArrayRef<int64_t> shape, Type elementType,
+                         Attribute encoding) {
   for (int64_t s : shape)
     if (s < -1)
       return emitError() << "invalid tensor dimension size";
+  // TODO: verify contents of encoding attribute.
   return checkTensorElementType(emitError, elementType);
 }
 
 //===----------------------------------------------------------------------===//
 // UnrankedTensorType
 //===----------------------------------------------------------------------===//
 
 LogicalResult

mlir/lib/Parser/TypeParser.cpp

if (consumeIf(Token::star)) {
     isUnranked = false;
     if (parseDimensionListRanked(dimensions))
       return nullptr;
   }
 
   // Parse the element type.
   auto elementTypeLoc = getToken().getLoc();
   auto elementType = parseType();
+
+  // Parse an optional encoding attribute.
+  Attribute encoding;
+  if (consumeIf(Token::comma))
+    encoding = parseAttribute();
+
   if (!elementType || parseToken(Token::greater, "expected '>' in tensor type"))
     return nullptr;
   if (!TensorType::isValidElementType(elementType))
     return emitError(elementTypeLoc, "invalid tensor element type"), nullptr;
 
-  if (isUnranked)
+  if (isUnranked) {
+    if (encoding)
+      return emitError("cannot apply encoding to unranked tensor"), nullptr;
     return UnrankedTensorType::get(elementType);
-  return RankedTensorType::get(dimensions, elementType);
+  }
+  return RankedTensorType::get(dimensions, elementType, encoding);
 }
 
 /// Parse a tuple type.
 ///
 ///   tuple-type ::= `tuple` `<` (type (`,` type)*)? `>`
 ///
 Type Parser::parseTupleType() {
   consumeToken(Token::kw_tuple);

mlir/test/CAPI/ir.c

static void printFirstOfEach(MlirContext ctx, MlirOperation operation) {
   // CHECK: Removed attr is null: 1
 
   // Add a large attribute to verify printing flags.
   int64_t eltsShape[] = {4};
   int32_t eltsData[] = {1, 2, 3, 4};
   mlirOperationSetAttributeByName(
       operation, mlirStringRefCreateFromCString("elts"),
       mlirDenseElementsAttrInt32Get(
-          mlirRankedTensorTypeGet(1, eltsShape, mlirIntegerTypeGet(ctx, 32)), 4,
-          eltsData));
+          mlirRankedTensorTypeGet(1, eltsShape, mlirIntegerTypeGet(ctx, 32),
+                                  mlirAttributeGetNull()), 4, eltsData));
   MlirOpPrintingFlags flags = mlirOpPrintingFlagsCreate();
   mlirOpPrintingFlagsElideLargeElementsAttrs(flags, 2);
   mlirOpPrintingFlagsPrintGenericOpForm(flags);
   mlirOpPrintingFlagsEnableDebugInfo(flags, /*prettyForm=*/0);
   mlirOpPrintingFlagsUseLocalScope(flags);
   fprintf(stderr, "Op print with all flags: ");
   mlirOperationPrintWithFlags(operation, flags, printToStderr, NULL);
   fprintf(stderr, "\n");
if (!mlirTypeEqual(mlirShapedTypeGetElementType(vector), f32) ||
       mlirShapedTypeGetDimSize(vector, 1) != 3 ||
       !mlirShapedTypeHasStaticShape(vector))
     return 15;
   mlirTypeDump(vector);
   fprintf(stderr, "\n");
   // CHECK: vector<2x3xf32>
 
   // Ranked tensor type.
-  MlirType rankedTensor =
-      mlirRankedTensorTypeGet(sizeof(shape) / sizeof(int64_t), shape, f32);
+  MlirType rankedTensor = mlirRankedTensorTypeGet(
+      sizeof(shape) / sizeof(int64_t), shape, f32, mlirAttributeGetNull());
   if (!mlirTypeIsATensor(rankedTensor) ||
       !mlirTypeIsARankedTensor(rankedTensor))
     return 16;
   mlirTypeDump(rankedTensor);
   fprintf(stderr, "\n");
   // CHECK: tensor<2x3xf32>
 
   // Unranked tensor type.
int printBuiltinAttributes(MlirContext ctx) {
 
   int bools[] = {0, 1};
   uint32_t uints32[] = {0u, 1u};
   int32_t ints32[] = {0, 1};
   uint64_t uints64[] = {0u, 1u};
   int64_t ints64[] = {0, 1};
   float floats[] = {0.0f, 1.0f};
   double doubles[] = {0.0, 1.0};
+  MlirAttribute encoding = mlirAttributeGetNull();
   MlirAttribute boolElements = mlirDenseElementsAttrBoolGet(
-      mlirRankedTensorTypeGet(2, shape, mlirIntegerTypeGet(ctx, 1)), 2, bools);
+      mlirRankedTensorTypeGet(2, shape, mlirIntegerTypeGet(ctx, 1), encoding),
+      2, bools);
   MlirAttribute uint32Elements = mlirDenseElementsAttrUInt32Get(
-      mlirRankedTensorTypeGet(2, shape, mlirIntegerTypeUnsignedGet(ctx, 32)), 2,
-      uints32);
+      mlirRankedTensorTypeGet(2, shape,
+                              mlirIntegerTypeUnsignedGet(ctx, 32), encoding),
+      2, uints32);
   MlirAttribute int32Elements = mlirDenseElementsAttrInt32Get(
-      mlirRankedTensorTypeGet(2, shape, mlirIntegerTypeGet(ctx, 32)), 2,
-      ints32);
+      mlirRankedTensorTypeGet(2, shape, mlirIntegerTypeGet(ctx, 32), encoding),
+      2, ints32);
   MlirAttribute uint64Elements = mlirDenseElementsAttrUInt64Get(
-      mlirRankedTensorTypeGet(2, shape, mlirIntegerTypeUnsignedGet(ctx, 64)), 2,
-      uints64);
+      mlirRankedTensorTypeGet(2, shape,
+                              mlirIntegerTypeUnsignedGet(ctx, 64), encoding),
+      2, uints64);
   MlirAttribute int64Elements = mlirDenseElementsAttrInt64Get(
-      mlirRankedTensorTypeGet(2, shape, mlirIntegerTypeGet(ctx, 64)), 2,
-      ints64);
+      mlirRankedTensorTypeGet(2, shape, mlirIntegerTypeGet(ctx, 64), encoding),
+      2, ints64);
   MlirAttribute floatElements = mlirDenseElementsAttrFloatGet(
-      mlirRankedTensorTypeGet(2, shape, mlirF32TypeGet(ctx)), 2, floats);
+      mlirRankedTensorTypeGet(2, shape, mlirF32TypeGet(ctx), encoding),
+      2, floats);
   MlirAttribute doubleElements = mlirDenseElementsAttrDoubleGet(
-      mlirRankedTensorTypeGet(2, shape, mlirF64TypeGet(ctx)), 2, doubles);
+      mlirRankedTensorTypeGet(2, shape, mlirF64TypeGet(ctx), encoding),
+      2, doubles);
 
   if (!mlirAttributeIsADenseElements(boolElements) ||
       !mlirAttributeIsADenseElements(uint32Elements) ||
       !mlirAttributeIsADenseElements(int32Elements) ||
       !mlirAttributeIsADenseElements(uint64Elements) ||
       !mlirAttributeIsADenseElements(int64Elements) ||
       !mlirAttributeIsADenseElements(floatElements) ||
       !mlirAttributeIsADenseElements(doubleElements))
int printBuiltinAttributes(MlirContext ctx) {
   // CHECK: dense<{{\[}}[0, 1]]> : tensor<1x2xui32>
   // CHECK: dense<{{\[}}[0, 1]]> : tensor<1x2xi32>
   // CHECK: dense<{{\[}}[0, 1]]> : tensor<1x2xui64>
   // CHECK: dense<{{\[}}[0, 1]]> : tensor<1x2xi64>
   // CHECK: dense<{{\[}}[0.000000e+00, 1.000000e+00]]> : tensor<1x2xf32>
   // CHECK: dense<{{\[}}[0.000000e+00, 1.000000e+00]]> : tensor<1x2xf64>
 
   MlirAttribute splatBool = mlirDenseElementsAttrBoolSplatGet(
-      mlirRankedTensorTypeGet(2, shape, mlirIntegerTypeGet(ctx, 1)), 1);
+      mlirRankedTensorTypeGet(2, shape,
+                              mlirIntegerTypeGet(ctx, 1), encoding), 1);
   MlirAttribute splatUInt32 = mlirDenseElementsAttrUInt32SplatGet(
-      mlirRankedTensorTypeGet(2, shape, mlirIntegerTypeGet(ctx, 32)), 1);
+      mlirRankedTensorTypeGet(2, shape,
+                              mlirIntegerTypeGet(ctx, 32), encoding), 1);
   MlirAttribute splatInt32 = mlirDenseElementsAttrInt32SplatGet(
-      mlirRankedTensorTypeGet(2, shape, mlirIntegerTypeGet(ctx, 32)), 1);
+      mlirRankedTensorTypeGet(2, shape,
+                              mlirIntegerTypeGet(ctx, 32), encoding), 1);
   MlirAttribute splatUInt64 = mlirDenseElementsAttrUInt64SplatGet(
-      mlirRankedTensorTypeGet(2, shape, mlirIntegerTypeGet(ctx, 64)), 1);
+      mlirRankedTensorTypeGet(2, shape,
+                              mlirIntegerTypeGet(ctx, 64), encoding), 1);
   MlirAttribute splatInt64 = mlirDenseElementsAttrInt64SplatGet(
-      mlirRankedTensorTypeGet(2, shape, mlirIntegerTypeGet(ctx, 64)), 1);
+      mlirRankedTensorTypeGet(2, shape,
+                              mlirIntegerTypeGet(ctx, 64), encoding), 1);
   MlirAttribute splatFloat = mlirDenseElementsAttrFloatSplatGet(
-      mlirRankedTensorTypeGet(2, shape, mlirF32TypeGet(ctx)), 1.0f);
+      mlirRankedTensorTypeGet(2, shape, mlirF32TypeGet(ctx), encoding), 1.0f);
   MlirAttribute splatDouble = mlirDenseElementsAttrDoubleSplatGet(
-      mlirRankedTensorTypeGet(2, shape, mlirF64TypeGet(ctx)), 1.0);
+      mlirRankedTensorTypeGet(2, shape, mlirF64TypeGet(ctx), encoding), 1.0);
 
   if (!mlirAttributeIsADenseElements(splatBool) ||
       !mlirDenseElementsAttrIsSplat(splatBool) ||
       !mlirAttributeIsADenseElements(splatUInt32) ||
       !mlirDenseElementsAttrIsSplat(splatUInt32) ||
       !mlirAttributeIsADenseElements(splatInt32) ||
       !mlirDenseElementsAttrIsSplat(splatInt32) ||
       !mlirAttributeIsADenseElements(splatUInt64) ||
int printBuiltinAttributes(MlirContext ctx) {
   mlirAttributeDump(mlirElementsAttrGetValue(floatElements, 2, uints64));
   mlirAttributeDump(mlirElementsAttrGetValue(doubleElements, 2, uints64));
   // CHECK: 1.000000e+00 : f32
   // CHECK: 1.000000e+00 : f64
 
   int64_t indices[] = {4, 7};
   int64_t two = 2;
   MlirAttribute indicesAttr = mlirDenseElementsAttrInt64Get(
-      mlirRankedTensorTypeGet(1, &two, mlirIntegerTypeGet(ctx, 64)), 2,
-      indices);
+      mlirRankedTensorTypeGet(1, &two, mlirIntegerTypeGet(ctx, 64), encoding),
+      2, indices);
   MlirAttribute valuesAttr = mlirDenseElementsAttrFloatGet(
-      mlirRankedTensorTypeGet(1, &two, mlirF32TypeGet(ctx)), 2, floats);
+      mlirRankedTensorTypeGet(1, &two, mlirF32TypeGet(ctx), encoding),
+      2, floats);
   MlirAttribute sparseAttr = mlirSparseElementsAttribute(
-      mlirRankedTensorTypeGet(2, shape, mlirF32TypeGet(ctx)), indicesAttr,
-      valuesAttr);
+      mlirRankedTensorTypeGet(2, shape, mlirF32TypeGet(ctx), encoding),
+      indicesAttr, valuesAttr);
   mlirAttributeDump(sparseAttr);
   // CHECK: sparse<[4, 7], [0.000000e+00, 1.000000e+00]> : tensor<1x2xf32>
 
   return 0;
 }
 
 int printAffineMap(MlirContext ctx) {
   MlirAffineMap emptyAffineMap = mlirAffineMapEmptyGet(ctx);

mlir/test/IR/invalid.mlir

 func @memref_stride_invalid_strides(memref<42x42xi8, offset: 0, strides: ()>) // expected-error {{invalid braces-enclosed stride list}}
 
 // -----
 
 func @memref_zero_stride(memref<42x42xi8, offset: ?, strides: [0, ?]>) // expected-error {{invalid memref stride}}
 
 // -----
 
+func @tensor_encoding_mismatch(%arg0: tensor<8xi32, "enc">) -> (tensor<8xi32>) { // expected-note {{prior use here}}
+  return %arg0: tensor<8xi32> // expected-error {{use of value '%arg0' expects different type than prior uses: 'tensor<8xi32>' vs 'tensor<8xi32, "enc">'}}
+}
+
+// -----
+
 func @bad_branch() {
 ^bb12:
   br ^missing  // expected-error {{reference to an undefined block}}
 }
 
 // -----
 
 func @block_redef() {

mlir/test/IR/parser.mlir

 
 // CHECK: func private @vectors(vector<1xf32>, vector<2x4xf32>)
 func private @vectors(vector<1 x f32>, vector<2x4xf32>)
 
 // CHECK: func private @tensors(tensor<*xf32>, tensor<*xvector<2x4xf32>>, tensor<1x?x4x?x?xi32>, tensor<i8>)
 func private @tensors(tensor<* x f32>, tensor<* x vector<2x4xf32>>,
                  tensor<1x?x4x?x?xi32>, tensor<i8>)
 
+// CHECK: func private @tensor_encoding(tensor<16x32xf64, "sparse">)
+func private @tensor_encoding(tensor<16x32xf64, "sparse">)
+
 // CHECK: func private @memrefs(memref<1x?x4x?x?xi32, #map{{[0-9]+}}>, memref<8xi8>)
 func private @memrefs(memref<1x?x4x?x?xi32, #map0>, memref<8xi8, #map1, #map1>)
 
 // Test memref affine map compositions.
 
 // CHECK: func private @memrefs2(memref<2x4x8xi8, 1>)
 func private @memrefs2(memref<2x4x8xi8, #map2, 1>)