[Mlir-commits] [mlir] 9546d8b - [mlir][core] Add IndexElementsAttr helpers.

[Sean Silva]

Author: Sean Silva
Date: 2020-05-27T13:39:48-07:00
New Revision: 9546d8b108dce03e03e0448cebbca5fa0fe4be21

URL: https://github.com/llvm/llvm-project/commit/9546d8b108dce03e03e0448cebbca5fa0fe4be21
DIFF: https://github.com/llvm/llvm-project/commit/9546d8b108dce03e03e0448cebbca5fa0fe4be21.diff

LOG: [mlir][core] Add IndexElementsAttr helpers.

Summary:
In a follow-up, I'll update the Shape dialect to use this instead of
I64ElementsAttr.

Differential Revision: https://reviews.llvm.org/D80601

Added: 
    

Modified: 
    mlir/include/mlir/IR/Builders.h
    mlir/include/mlir/IR/OpBase.td
    mlir/lib/IR/Attributes.cpp
    mlir/lib/IR/Builders.cpp
    mlir/test/lib/Dialect/Test/TestOps.td
    mlir/test/mlir-tblgen/types.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/IR/Builders.h b/mlir/include/mlir/IR/Builders.h
index 4ade6bb1e439..424eb980cd33 100644
--- a/mlir/include/mlir/IR/Builders.h
+++ b/mlir/include/mlir/IR/Builders.h
@@ -128,6 +128,7 @@ class Builder {
   /// as attributes.
   DenseIntElementsAttr getI32TensorAttr(ArrayRef<int32_t> values);
   DenseIntElementsAttr getI64TensorAttr(ArrayRef<int64_t> values);
+  DenseIntElementsAttr getIndexTensorAttr(ArrayRef<int64_t> values);
 
   ArrayAttr getAffineMapArrayAttr(ArrayRef<AffineMap> values);
   ArrayAttr getBoolArrayAttr(ArrayRef<bool> values);

diff  --git a/mlir/include/mlir/IR/OpBase.td b/mlir/include/mlir/IR/OpBase.td
index 6a7542c7127c..5ffb1727ee35 100644
--- a/mlir/include/mlir/IR/OpBase.td
+++ b/mlir/include/mlir/IR/OpBase.td
@@ -1218,6 +1218,13 @@ class IntElementsAttrBase<Pred condition, string description> :
   let convertFromStorage = "$_self";
 }
 
+def IndexElementsAttr
+    : IntElementsAttrBase<CPred<[{$_self.cast<DenseIntElementsAttr>()
+                                      .getType()
+                                      .getElementType()
+                                      .isIndex()}]>,
+                          "index elements attribute">;
+
 class AnyIntElementsAttr<int width> : IntElementsAttrBase<
   CPred<"$_self.cast<DenseIntElementsAttr>().getType()."
         "getElementType().isInteger(" # width # ")">,

diff  --git a/mlir/lib/IR/Attributes.cpp b/mlir/lib/IR/Attributes.cpp
index 540c3c6258e2..12fd08787fa7 100644
--- a/mlir/lib/IR/Attributes.cpp
+++ b/mlir/lib/IR/Attributes.cpp
@@ -624,6 +624,8 @@ Attribute DenseElementsAttr::AttributeElementIterator::operator*() const {
                            owner.getContext());
     return IntegerAttr::get(eltTy, *IntElementIterator(owner, index));
   }
+  if (eltTy.isa<IndexType>())
+    return IntegerAttr::get(eltTy, *IntElementIterator(owner, index));
   if (auto floatEltTy = eltTy.dyn_cast<FloatType>()) {
     IntElementIterator intIt(owner, index);
     FloatElementIterator floatIt(floatEltTy.getFloatSemantics(), intIt);

diff  --git a/mlir/lib/IR/Builders.cpp b/mlir/lib/IR/Builders.cpp
index a72e03c739e3..064889724f09 100644
--- a/mlir/lib/IR/Builders.cpp
+++ b/mlir/lib/IR/Builders.cpp
@@ -130,6 +130,13 @@ DenseIntElementsAttr Builder::getI64TensorAttr(ArrayRef<int64_t> values) {
       values);
 }
 
+DenseIntElementsAttr Builder::getIndexTensorAttr(ArrayRef<int64_t> values) {
+  return DenseIntElementsAttr::get(
+      RankedTensorType::get(static_cast<int64_t>(values.size()),
+                            getIndexType()),
+      values);
+}
+
 IntegerAttr Builder::getI32IntegerAttr(int32_t value) {
   return IntegerAttr::get(getIntegerType(32), APInt(32, value));
 }

diff  --git a/mlir/test/lib/Dialect/Test/TestOps.td b/mlir/test/lib/Dialect/Test/TestOps.td
index 997d8eb44ae5..8e5b380dff45 100644
--- a/mlir/test/lib/Dialect/Test/TestOps.td
+++ b/mlir/test/lib/Dialect/Test/TestOps.td
@@ -454,6 +454,10 @@ def I32ElementsAttrOp : TEST_Op<"i32ElementsAttr"> {
   let arguments = (ins I32ElementsAttr:$attr);
 }
 
+def IndexElementsAttrOp : TEST_Op<"indexElementsAttr"> {
+  let arguments = (ins IndexElementsAttr:$attr);
+}
+
 def OpWithInferTypeInterfaceOp : TEST_Op<"op_with_infer_type_if", [
     DeclareOpInterfaceMethods<InferTypeOpInterface>]> {
   let arguments = (ins AnyTensor, AnyTensor);

diff  --git a/mlir/test/mlir-tblgen/types.mlir b/mlir/test/mlir-tblgen/types.mlir
index 6a0a80ca5e5f..5e4dac33012b 100644
--- a/mlir/test/mlir-tblgen/types.mlir
+++ b/mlir/test/mlir-tblgen/types.mlir
@@ -489,3 +489,18 @@ func @elements_attr_i32(%arg0: tensor<1x2xi32>) {
   "test.i32ElementsAttr"() {attr = dense<[1, 2]>:tensor<2xi32>} : () -> ()
   return
 }
+
+// -----
+
+func @elements_attr_index() {
+  "test.indexElementsAttr"() {attr = dense<[1, 2]>:tensor<2xindex>} : () -> ()
+  return
+}
+
+// -----
+
+func @elements_attr_not_index() {
+  // expected-error at +1 {{index elements attribute}}
+  "test.indexElementsAttr"() {attr = dense<[1, 2]>:tensor<2xi32>} : () -> ()
+  return
+}