Add a builtin DimensionListAttribute intended to handle the many ops that just want a small array of int64_t as an attribute

[mehdi_amini]

Summary

Right now the default for such case is DenseIntElementAttr, however because it
is so generic its API is quite hard to use and it is easy to either materialize
individual elements as individual IntegerAttr or as APInt. It is also not
possible to easily get an ArrayRef<int64>.

D116394 is an illustration of the effect of this new attribute on a recent part of
the codebase (the TOSA dialect). There is saves multiple traversals of ArrayAttr
and local creation of SmallVector<int64_t> in various situation.

[mehdi_amini]

Fix CMakeLists.txt

[Mogball]

mlir/include/mlir/IR/BuiltinAttributes.td

355–357

Maybe a concrete example? dims<1, 2, 3>

360

You should also document that embedding the attribute in an op's assembly format results in a different format [1, 2, 3] since the interaction with the builtin attribute parser/printer isn't obvious.

mlir/lib/Parser/AttributeParser.cpp

841–857

Also, can you separate out the lambda definition?

[jpienaar]

mlir/lib/Parser/TokenKinds.def

85

Do we really need a new keyword? This seems a very specific attribute and added to built-in (I'm not saying it isn't useful, just this is very privileged compared to just wanting a dense attribute with specialized storage)

[mehdi_amini]

mlir/lib/Parser/TokenKinds.def

85

The main question is how would you print it: #builtin.dims<1, 2 ,3>?

[mehdi_amini]

Add a trait to help for defining ArrayRef-like Attributes in general and use it for this new attribute, also address Jeff's comments.

[mehdi_amini]

Rebase

[Mogball]

mlir/include/mlir/IR/AttributeSupport.h

280 ↗

(On Diff #396751)

I'm not against having this trait, but is remapping literally everything necessary?

mlir/include/mlir/IR/BuiltinAttributes.td

364

Copy paste error?

370

[Mogball]

Some of the relevant changes in the Tosa patch need to be moved here

[mehdi_amini]

Backport changes the were in the child revision by mistake

[mehdi_amini]

mlir/include/mlir/IR/AttributeSupport.h

280 ↗

(On Diff #396751)

If we don't forward the APIs, then when you can do something like for (int64_t dim : op.getDimensions()), you have to always do for (int64_t dim : op.getDimensions().getArray()) or cast explicitly: for (int64_t dim : ArrayRef<int64_t>(op.getDimensions()))

[mehdi_amini]

Fix copy/paste error in doc and fully qualified namespace

[rriddle]

Blocking until after holiday so I can take a look. I'm a bit apprehensive about how some of this is setup right now, namely why we need a specialized builtin attribute just for dimensions. Also the ArrayRefAttr API looks like a more general range mixin (likely living in ADT).

[mehdi_amini]

Rebase: remove all the ArrayRef methods after fixing the accessors on operands in D116466

mlir/include/mlir/IR/Builders.h

public:
   ArrayAttr getI32ArrayAttr(ArrayRef<int32_t> values);
   ArrayAttr getI64ArrayAttr(ArrayRef<int64_t> values);
   ArrayAttr getIndexArrayAttr(ArrayRef<int64_t> values);
   ArrayAttr getF32ArrayAttr(ArrayRef<float> values);
   ArrayAttr getF64ArrayAttr(ArrayRef<double> values);
   ArrayAttr getStrArrayAttr(ArrayRef<StringRef> values);
   ArrayAttr getTypeArrayAttr(TypeRange values);
 
+  // Dimension list is a Dense array of int64_t values.
+  DimensionListAttr getDimListAttr(ArrayRef<int64_t> values);
+
   // Affine expressions and affine maps.
   AffineExpr getAffineDimExpr(unsigned position);
   AffineExpr getAffineSymbolExpr(unsigned position);
   AffineExpr getAffineConstantExpr(int64_t constant);
 
   // Special cases of affine maps and integer sets
   /// Returns a zero result affine map with no dimensions or symbols: () -> ().
   AffineMap getEmptyAffineMap();

mlir/include/mlir/IR/BuiltinAttributes.td

def Builtin_DenseStringElementsAttr : Builtin_Attr<
 
   public:
   }];
   let genAccessors = 0;
   let genStorageClass = 0;
   let skipDefaultBuilders = 1;
 }
 
+
+//===----------------------------------------------------------------------===//
+// DimensionListAttr
+//===----------------------------------------------------------------------===//
+
+def Builtin_DimensionListAttr : Builtin_Attr<"DimensionList"> {
+  let mnemonic = "dims";
+  let summary = "An Attribute containing a dense array of in64_t, suitable to "
+    "model a list of dimensions or a subscript for example.";
+  let description = [{
+    Syntax:
+
+    ```
+    `dims` `<` <comma separate list of integers> `>`
+    ```
+
+    For example: `dims<42, -5>` or `dims<>`.
+
+    When used embedded in an operation, it is displayed with an array syntax:
+
+    ```
+    `dims` `[` <comma separate list of integers> `]`
+    ```
+
+    For example: `test.op [42, -5]`.
+  }];
+  let parameters = (ins ArrayRefParameter<"int64_t">:$array);
+  let returnType = "::llvm::ArrayRef<int64_t>";
+  code convertFromStorage = "$_self.getArray()";
+  let assemblyFormat = "`[` $array `]`";
+  let extraClassDeclaration = [{
+    // Convenient implicit conversion to allow using this attribute in APIs
+    // taking an ArrayRef<int64_t>.
+    operator ArrayRef<int64_t>() { return getArray(); }
+  }];
+}
+
+class DimensionListAttrOfSize<int n> : Confined<Builtin_DimensionListAttr,
+   [AttrConstraint<
+      CPred<"$_self.cast<::mlir::DimensionListAttr>().getArray().size() == " #n>,
+            "with exactly " # n # " elements">
+   ]>;
+
+class DimensionListAttrOfMaxSize<int n> : Confined<Builtin_DimensionListAttr,
+   [AttrConstraint<
+      CPred<"$_self.cast<::mlir::DimensionListAttr>().getArray().size() <= " #n>,
+            "with at most " # n # " elements">
+   ]>;
+
+
 //===----------------------------------------------------------------------===//
 // DictionaryAttr
 //===----------------------------------------------------------------------===//
 
 def Builtin_DictionaryAttr : Builtin_Attr<"Dictionary", [
     DeclareAttrInterfaceMethods<SubElementAttrInterface,
         ["replaceImmediateSubAttribute"]>
   ]> {

mlir/lib/IR/AsmPrinter.cpp

protected:
   void printNamedAttribute(NamedAttribute attr);
   void printTrailingLocation(Location loc, bool allowAlias = true);
   void printLocationInternal(LocationAttr loc, bool pretty = false);
 
   /// Print a dense elements attribute. If 'allowHex' is true, a hex string is
   /// used instead of individual elements when the elements attr is large.
   void printDenseElementsAttr(DenseElementsAttr attr, bool allowHex);
 
+  void printDimensionListAttr(DimensionListAttr attr);
+
   /// Print a dense string elements attribute.
   void printDenseStringElementsAttr(DenseStringElementsAttr attr);
 
   /// Print a dense elements attribute. If 'allowHex' is true, a hex string is
   /// used instead of individual elements when the elements attr is large.
   void printDenseIntOrFPElementsAttr(DenseIntOrFPElementsAttr attr,
                                      bool allowHex);
 
if (auto opaqueAttr = attr.dyn_cast<OpaqueAttr>()) {
     if (printerFlags.shouldElideElementsAttr(intOrFpEltAttr)) {
       printElidedElementsAttr(os);
     } else {
       os << "dense<";
       printDenseIntOrFPElementsAttr(intOrFpEltAttr, /*allowHex=*/true);
       os << '>';
     }
 
+  } else if (auto dimsList = attr.dyn_cast<DimensionListAttr>()) {
+    os << "dims<";
+    printDimensionListAttr(dimsList);
+    os << '>';
+
   } else if (auto strEltAttr = attr.dyn_cast<DenseStringElementsAttr>()) {
     if (printerFlags.shouldElideElementsAttr(strEltAttr)) {
       printElidedElementsAttr(os);
     } else {
       os << "dense<";
       printDenseStringElementsAttr(strEltAttr);
       os << '>';
     }
if (ComplexType complexTy = elementType.dyn_cast<ComplexType>()) {
     assert(elementType.isa<FloatType>() && "unexpected element type");
     auto valueIt = attr.value_begin<APFloat>();
     printDenseElementsAttrImpl(attr.isSplat(), type, os, [&](unsigned index) {
       printFloatValue(*(valueIt + index), os);
     });
   }
 }
 
+void AsmPrinter::Impl::printDimensionListAttr(DimensionListAttr attr) {
+  llvm::interleaveComma(attr.getArray(), os);
+}
+
 void AsmPrinter::Impl::printDenseStringElementsAttr(
     DenseStringElementsAttr attr) {
   ArrayRef<StringRef> data = attr.getRawStringData();
   auto printFn = [&](unsigned index) {
     os << "\"";
     printEscapedString(data[index], os);
     os << "\"";
   };

mlir/lib/IR/Builders.cpp

ArrayAttr Builder::getI32ArrayAttr(ArrayRef<int32_t> values) {
   return getArrayAttr(attrs);
 }
 ArrayAttr Builder::getI64ArrayAttr(ArrayRef<int64_t> values) {
   auto attrs = llvm::to_vector<8>(llvm::map_range(
       values, [this](int64_t v) -> Attribute { return getI64IntegerAttr(v); }));
   return getArrayAttr(attrs);
 }
 
+DimensionListAttr Builder::getDimListAttr(ArrayRef<int64_t> values) {
+  return DimensionListAttr::get(getContext(), values);
+}
+
 ArrayAttr Builder::getIndexArrayAttr(ArrayRef<int64_t> values) {
   auto attrs = llvm::to_vector<8>(
       llvm::map_range(values, [this](int64_t v) -> Attribute {
         return getIntegerAttr(IndexType::get(getContext()), v);
       }));
   return getArrayAttr(attrs);
 }
 

mlir/lib/IR/BuiltinAttributes.cpp

 //===- BuiltinAttributes.cpp - MLIR Builtin Attribute Classes -------------===//
 //
 // 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 "mlir/IR/BuiltinAttributes.h"
 #include "AttributeDetail.h"
 #include "mlir/IR/AffineMap.h"
 #include "mlir/IR/BuiltinDialect.h"
 #include "mlir/IR/Dialect.h"
+#include "mlir/IR/DialectImplementation.h"
 #include "mlir/IR/IntegerSet.h"
+#include "mlir/IR/OpImplementation.h"
 #include "mlir/IR/Operation.h"
 #include "mlir/IR/SymbolTable.h"
 #include "mlir/IR/Types.h"
 #include "mlir/Interfaces/DecodeAttributesInterfaces.h"
 #include "llvm/ADT/APSInt.h"
 #include "llvm/ADT/Sequence.h"
+#include "llvm/ADT/TypeSwitch.h"
 #include "llvm/Support/Endian.h"
 
 using namespace mlir;
 using namespace mlir::detail;
 
 //===----------------------------------------------------------------------===//
 /// Tablegen Attribute Definitions
 //===----------------------------------------------------------------------===//
 
 #define GET_ATTRDEF_CLASSES
 #include "mlir/IR/BuiltinAttributes.cpp.inc"
 
 //===----------------------------------------------------------------------===//
 // BuiltinDialect
 //===----------------------------------------------------------------------===//
 
 void BuiltinDialect::registerAttributes() {
   addAttributes<AffineMapAttr, ArrayAttr, DenseIntOrFPElementsAttr,
-                DenseStringElementsAttr, DictionaryAttr, FloatAttr,
-                SymbolRefAttr, IntegerAttr, IntegerSetAttr, OpaqueAttr,
-                OpaqueElementsAttr, SparseElementsAttr, StringAttr, TypeAttr,
-                UnitAttr>();
+                DimensionListAttr, DenseStringElementsAttr, DictionaryAttr,
+                FloatAttr, SymbolRefAttr, IntegerAttr, IntegerSetAttr,
+                OpaqueAttr, OpaqueElementsAttr, SparseElementsAttr, StringAttr,
+                TypeAttr, UnitAttr>();
 }
 
 //===----------------------------------------------------------------------===//
 // ArrayAttr
 //===----------------------------------------------------------------------===//
 
 void ArrayAttr::walkImmediateSubElements(
     function_ref<void(Attribute)> walkAttrsFn,

mlir/lib/Parser/AttributeParser.cpp

Attribute Parser::parseAttribute(Type type) {
   case Token::kw_true:
     consumeToken(Token::kw_true);
     return builder.getBoolAttr(true);
 
   // Parse a dense elements attribute.
   case Token::kw_dense:
     return parseDenseElementsAttr(type);
 
+  // Parse a dimension lists attribute.
+  case Token::kw_dims:
+    return parseDimensionListAttr();
+
   // Parse a dictionary attribute.
   case Token::l_brace: {
     NamedAttrList elements;
     if (parseAttributeDict(elements))
       return nullptr;
     return elements.getDictionary(getContext());
   }
 
Attribute Parser::parseDenseElementsAttr(Type attrType) {
   // case.
   auto loc = attrType ? attribLoc : getToken().getLoc();
   auto type = parseElementsLiteralType(attrType);
   if (!type)
     return nullptr;
   return literalParser.getAttr(loc, type);
 }
 
+Attribute Parser::parseDimensionListAttr() {
+  consumeToken(Token::kw_dims);
+  SmallVector<int64_t> dims;
+  auto parseElement = [&]() -> LogicalResult {
+    if (getToken().is(Token::greater))
+      return success();
+    bool negative = false;
+    if (getToken().is(Token::minus)) {
+      negative = true;
+      consumeToken();
+    }
+    if (getToken().isNot(Token::integer))
+      return emitError("expected integer in dimensions list");
+    Optional<uint64_t> value = getToken().getUInt64IntegerValue();
+    if (!value || *value > std::numeric_limits<int64_t>::max())
+      return emitError("integer overflow in dimensions list");
+    dims.emplace_back(negative ? -*value : *value);
+    consumeToken();
+    return success();
+  };
+  if (parseCommaSeparatedList(Delimiter::LessGreater, parseElement))
+    return nullptr;
+  return DimensionListAttr::get(getContext(), dims);
+}
+
 /// Parse an opaque elements attribute.
 Attribute Parser::parseOpaqueElementsAttr(Type attrType) {
   consumeToken(Token::kw_opaque);
   if (parseToken(Token::less, "expected '<' after 'opaque'"))
     return nullptr;
 
   if (getToken().isNot(Token::string))
     return (emitError("expected dialect namespace"), nullptr);

mlir/lib/Parser/Parser.h

public:
 
   /// Parse an opaque elements attribute.
   Attribute parseOpaqueElementsAttr(Type attrType);
 
   /// Parse a dense elements attribute.
   Attribute parseDenseElementsAttr(Type attrType);
   ShapedType parseElementsLiteralType(Type type);
 
+  /// Parse a dimension lists attribute.
+  Attribute parseDimensionListAttr();
+
   /// Parse a sparse elements attribute.
   Attribute parseSparseElementsAttr(Type attrType);
 
   //===--------------------------------------------------------------------===//
   // Location Parsing
   //===--------------------------------------------------------------------===//
 
   /// Parse a raw location instance.

mlir/lib/Parser/TokenKinds.def

 // this list and to cater to OCD.
 TOK_KEYWORD(affine_map)
 TOK_KEYWORD(affine_set)
 TOK_KEYWORD(attributes)
 TOK_KEYWORD(bf16)
 TOK_KEYWORD(ceildiv)
 TOK_KEYWORD(complex)
 TOK_KEYWORD(dense)
+TOK_KEYWORD(dims)
 TOK_KEYWORD(f16)
 TOK_KEYWORD(f32)
 TOK_KEYWORD(f64)
 TOK_KEYWORD(f80)
 TOK_KEYWORD(f128)
 TOK_KEYWORD(false)
 TOK_KEYWORD(floordiv)
 TOK_KEYWORD(for)

mlir/test/IR/attribute.mlir

"test.ranked_int_elements_attr"() {
     matrix_i64_attr = dense<6> : tensor<4x8xi64>,
     vector_i32_attr = dense<5> : tensor<2xi64>
   } : () -> ()
   return
 }
 
 // -----
 
+//===----------------------------------------------------------------------===//
+// Test DimensionListAttr
+//===----------------------------------------------------------------------===//
+
+// CHECK-LABEL: func @dim_lists
+func @dim_lists() {
+  // CHECK: test.dim_list_attr
+  // CHECK-SAME: attr1 = dims<6>
+  // CHECK-SAME: attr2 = dims<-1, 2, 3>
+  // CHECK-SAME: unregistered_attr = dims<>
+  "test.dim_list_attr"() {
+    attr1 = dims<6>,
+    attr2 = dims<-1, 2, 3>,
+    unregistered_attr = dims<>
+  } : () -> ()
+  return
+}
+
+// -----
+
+// CHECK-LABEL: func @dim_list_attr_custom
+func @dim_list_attr_custom() {
+  // CHECK: test.dim_list_attr_custom
+  // CHECK-SAME: theDims : [6]
+  test.dim_list_attr_custom theDims : [6]
+  return
+}
+
+
+// -----
+
 func @wrong_shape_fail() {
   // expected-error @+1 {{failed to satisfy constraint: 64-bit signless int elements attribute of shape [4, 8]}}
   "test.ranked_int_elements_attr"() {
     matrix_i64_attr = dense<6> : tensor<4xi64>,
     vector_i32_attr = dense<5> : tensor<2xi32>
   } : () -> ()
   return
 }

mlir/test/lib/Dialect/Test/CMakeLists.txt

 set(LLVM_OPTIONAL_SOURCES
   TestDialect.cpp
   TestPatterns.cpp
   TestTraits.cpp
 )
 
 set(LLVM_TARGET_DEFINITIONS TestInterfaces.td)
 mlir_tablegen(TestAttrInterfaces.h.inc -gen-attr-interface-decls)
 mlir_tablegen(TestAttrInterfaces.cpp.inc -gen-attr-interface-defs)
 mlir_tablegen(TestTypeInterfaces.h.inc -gen-type-interface-decls)
 mlir_tablegen(TestTypeInterfaces.cpp.inc -gen-type-interface-defs)
 mlir_tablegen(TestOpInterfaces.h.inc -gen-op-interface-decls)
 mlir_tablegen(TestOpInterfaces.cpp.inc -gen-op-interface-defs)
 add_public_tablegen_target(MLIRTestInterfaceIncGen)
 
 set(LLVM_TARGET_DEFINITIONS TestOps.td)
-mlir_tablegen(TestAttrDefs.h.inc -gen-attrdef-decls)
-mlir_tablegen(TestAttrDefs.cpp.inc -gen-attrdef-defs)
+mlir_tablegen(TestAttrDefs.h.inc -gen-attrdef-decls -attrdefs-dialect=test)
+mlir_tablegen(TestAttrDefs.cpp.inc -gen-attrdef-defs -attrdefs-dialect=test)
 add_public_tablegen_target(MLIRTestAttrDefIncGen)
 
 set(LLVM_TARGET_DEFINITIONS TestTypeDefs.td)
 mlir_tablegen(TestTypeDefs.h.inc -gen-typedef-decls -typedefs-dialect=test)
 mlir_tablegen(TestTypeDefs.cpp.inc -gen-typedef-defs -typedefs-dialect=test)
 add_public_tablegen_target(MLIRTestTypeDefIncGen)
 
 

mlir/test/lib/Dialect/Test/TestOps.td

 //===-- TestOps.td - Test dialect operation definitions ----*- tablegen -*-===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef TEST_OPS
 #define TEST_OPS
 
 include "TestDialect.td"
 include "mlir/Dialect/DLTI/DLTIBase.td"
 include "mlir/IR/EnumAttr.td"
+include "mlir/IR/BuiltinAttributes.td"
 include "mlir/IR/OpBase.td"
 include "mlir/IR/OpAsmInterface.td"
 include "mlir/IR/RegionKindInterface.td"
 include "mlir/IR/SymbolInterfaces.td"
 include "mlir/Interfaces/CallInterfaces.td"
 include "mlir/Interfaces/ControlFlowInterfaces.td"
 include "mlir/Interfaces/CopyOpInterface.td"
 include "mlir/Interfaces/DataLayoutInterfaces.td"

 
 def RankedIntElementsAttrOp : TEST_Op<"ranked_int_elements_attr"> {
   let arguments = (ins
       RankedI32ElementsAttr<[2]>:$vector_i32_attr,
       RankedI64ElementsAttr<[4, 8]>:$matrix_i64_attr
   );
 }
 
+def DimListAttrOp : TEST_Op<"dim_list_attr"> {
+  let arguments = (ins
+      Builtin_DimensionListAttr:$attr1,
+      Builtin_DimensionListAttr:$attr2
+  );
+}
+
+def DimListAttrCustomAsmOp : TEST_Op<"dim_list_attr_custom"> {
+  let arguments = (ins
+      Builtin_DimensionListAttr:$attr1
+  );
+  let assemblyFormat = "`theDims` `:` $attr1 attr-dict";
+}
+
 def DerivedTypeAttrOp : TEST_Op<"derived_type_attr", []> {
   let results = (outs AnyTensor:$output);
   DerivedTypeAttr element_dtype =
     DerivedTypeAttr<"return getElementTypeOrSelf(getOutput().getType());">;
   DerivedAttr size = DerivedAttr<"int",
     "return getOutput().getType().cast<ShapedType>().getSizeInBits();",
     "$_builder.getI32IntegerAttr($_self)">;
 }

utils/bazel/llvm-project-overlay/mlir/test/BUILD.bazel

gentbl_cc_library(
     ],
 )
 
 gentbl_cc_library(
     name = "TestAttrDefsIncGen",
     strip_include_prefix = "lib/Dialect/Test",
     tbl_outs = [
         (
-            ["-gen-attrdef-decls"],
+            [
+                "-gen-attrdef-decls",
+                "--attrdefs-dialect=test",
+            ],
             "lib/Dialect/Test/TestAttrDefs.h.inc",
         ),
         (
-            ["-gen-attrdef-defs"],
+            [
+                "-gen-attrdef-defs",
+                "--attrdefs-dialect=test",
+            ],
             "lib/Dialect/Test/TestAttrDefs.cpp.inc",
         ),
     ],
     tblgen = "//mlir:mlir-tblgen",
     td_file = "lib/Dialect/Test/TestOps.td",
     test = True,
     deps = [
         ":TestOpTdFiles",