[mlir] Extern constant via ElementsAttr.

[jpienaar]

Summary

Simple example of using ElementsAttr interface to implement a constant
value not uniqued in the context and whose lifetime is not guaranteed by
the context. This effectively results in the pointer to the data being
uniqued rather than the data. But it can be used as regular
ElementsAttr.

Example behavior is to lose "extern'ness" on round tripping through
textual format. Chose this as my expectation is that extern constant
would be used as part of workflow that manages state, but the moment one
outputs to textual format you have left that.

[mehdi_amini]

Chose this as my expectation is that extern constant would be used as part of workflow that manages state, but the moment one outputs to textual format you have left that.

I have strong concerns with this: even if there is some external state, the IR should round-trip as much as possible..
So instead I'd like to see a system where we would be able round-trip whatever ID is used to interact with the runtime to access the constant (it can be a name or other means of ID).
Of course that means that you can't just get a pointer to this attribute alone and get the data without some indirection, but that can be also handled by the Dialect class itself for example.

[jpienaar]

In D122770#3417918, @mehdi_amini wrote:

Chose this as my expectation is that extern constant would be used as part of workflow that manages state, but the moment one outputs to textual format you have left that.

I have strong concerns with this: even if there is some external state, the IR should round-trip as much as possible..

This roundtrips except for location in memory values are stored (but that's true for regular elements attributes too) and whether equal values are deduped or not (which would result in changing equality checks indeed). I could add the textual format so that equal values but with different backing storage are not equal and so preserve the non-deduped nature. The life time difference cannot appear in any valid program and is pre-requisite of use.

So instead I'd like to see a system where we would be able round-trip whatever ID is used to interact with the runtime to access the constant (it can be a name or other means of ID).

Would you want attributes fixed to particular runtimes? Or you'd want all runtimes to implement some interface? I was thinking primarily of easy interop with runtimes that are unaware of MLIR, so no changes to runtimes, no knowledge about allocating memory (e.g., when parsing), requirements are the MLIR context doesn't outlive backing storage and no mutation to the backing state during MLIR passes (which for some runtimes would require that one bumps the ref count on the "tensor", which means the MLIR importer side needs to be aware of the runtime semantics but its kept to edges at least). Materialization should be avoided during actual optimizations/transformations but is not explicitly avoided by any means but you'd have to explicitly convert the attribute to another attribute though, so it wouldn't be accidental. [and to be clear, I don't expect this actual instance to be used with any real runtime :)]

I'd expect rather this would go another way: you'd be having a framework specific input that decides what kind of attribute to create (e.g., if parsing for mobile you may avoid making any dense elementrs attributes with more than 4 elements actually in context and rather try and keep them extern referring to flatbuffer directly). The textual debugging format and its printing & parsing is just for testing as normal, and the decision of how and where to store constants is not grouped in with that (e.g., not link debugging format to runtime).

[mehdi_amini]

In D122770#3423839, @jpienaar wrote:

In D122770#3417918, @mehdi_amini wrote:

Chose this as my expectation is that extern constant would be used as part of workflow that manages state, but the moment one outputs to textual format you have left that.

I have strong concerns with this: even if there is some external state, the IR should round-trip as much as possible..

This roundtrips except for location in memory values are stored (but that's true for regular elements attributes too) and whether equal values are deduped or not (which would result in changing equality checks indeed). I could add the textual format so that equal values but with different backing storage are not equal and so preserve the non-deduped nature.

Got it: it seems that you're preserving more than I thought here. I still think that even if you're using a DenseIntElementsAttr when parsing back (for the storage), you should likely still produce a TestExtern1DI64Elements instance using the pointer to these data, otherwise an op defined as having a TestExtern1DI64Elements in ODS would fail the verifier when parsing back I think (can you add printing/parsing tests in your test?).

You also have a point about the pointer equality changing, which can be addressed by printing the pointer as extra "key" for the unification in the context?

The life time difference cannot appear in any valid program and is pre-requisite of use.

Right.

So instead I'd like to see a system where we would be able round-trip whatever ID is used to interact with the runtime to access the constant (it can be a name or other means of ID).

Would you want attributes fixed to particular runtimes? Or you'd want all runtimes to implement some interface? I was thinking primarily of easy interop with runtimes [...]

Yeah, actually I think you're actually on the right path here!

Something I've been looking in the past was going one step further by also to have the dialect having a possible callback into the runtime to allocate memory so that folding constant could also be externalized.

mlir/include/mlir/Support/DebugStringHelper.h

 
 namespace mlir {
 
 // Simple helper function that returns a string as printed from a op.
 template <typename T>
 static std::string debugString(T &&op) {
   std::string instrStr;
   llvm::raw_string_ostream os(instrStr);
-  op.print(os);
+  os << op;
   return os.str();
 }
 
 } // namespace mlir
 
 inline std::ostream &operator<<(std::ostream &out, const llvm::Twine &twine) {
   llvm::raw_os_ostream rout(out);
   rout << twine;
   return out;
 }
 
 #endif // MLIR_SUPPORT_DEBUGSTRINGHELPER_H

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

auto value_begin_impl(OverloadToken<llvm::APInt>) const {
         return llvm::APInt(/*numBits=*/64, value);
       }).begin();
     }
   }];
   let genVerifyDecl = 1;
   let hasCustomAssemblyFormat = 1;
 }
 
+// Test simple extern 1D vector using ElementsAttrInterface.
+def TestExtern1DI64ElementsAttr : Test_Attr<"TestExtern1DI64Elements", [
+    ElementsAttrInterface
+  ]> {
+  let mnemonic = "e1di64_elements";
+  let parameters = (ins
+    AttributeSelfTypeParameter<"", "::mlir::ShapedType">:$type,
+
+    // Extern ArrayRef whose pointer rather than contents is uniqued and which
+    // require the user ensures the lifetime of the array to exceed context.
+    AttrOrTypeParameter<"PtrI64Array", "pointer to I64 array">:$elements
+  );
+  let builders = [
+    AttrBuilder<(ins "llvm::ArrayRef<int64_t>":$value), [{
+      return get($_ctxt,
+        RankedTensorType::get({static_cast<long>(value.size())},
+    Builder($_ctxt).getI64Type()),
+        {&value.front(), value.size()});
+    }]>];
+  let extraClassDeclaration = [{
+    /// The set of data types that can be iterated by this attribute.
+    using ContiguousIterableTypesT = std::tuple<int64_t>;
+
+    /// Provide begin iterators for the various iterable types.
+    // * int64_t
+    auto value_begin_impl(OverloadToken<int64_t>) const {
+      return getElements().begin();
+    }
+  }];
+  let hasCustomAssemblyFormat = 1;
+}
+
 def TestSubElementsAccessAttr : Test_Attr<"TestSubElementsAccess", [
     DeclareAttrInterfaceMethods<SubElementAttrInterface,
         ["replaceImmediateSubAttribute"]>
   ]> {
 
   let mnemonic = "sub_elements_access";
 
   let parameters = (ins

mlir/test/lib/Dialect/Test/TestAttributes.h

 #include "mlir/IR/Attributes.h"
 #include "mlir/IR/Diagnostics.h"
 #include "mlir/IR/Dialect.h"
 #include "mlir/IR/DialectImplementation.h"
 
 #include "TestAttrInterfaces.h.inc"
 #include "TestOpEnums.h.inc"
 
+// Class that is effectively ArrayRef but just with new hash value only
+// considering base pointer and size rather than Value.
+class PtrI64Array : public llvm::ArrayRef<int64_t> {
+public:
+  using llvm::ArrayRef<int64_t>::ArrayRef;
+};
+
 #define GET_ATTRDEF_CLASSES
 #include "TestAttrDefs.h.inc"
 
 #endif // MLIR_TESTATTRIBUTES_H

mlir/test/lib/Dialect/Test/TestAttributes.cpp

 //===- TestAttributes.cpp - MLIR Test Dialect Attributes --------*- C++ -*-===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 //
 // This file contains attributes defined by the TestDialect for testing various
 // features of MLIR.
 //
 //===----------------------------------------------------------------------===//
 
 #include "TestAttributes.h"
 #include "TestDialect.h"
 #include "mlir/IR/Builders.h"
+#include "mlir/IR/BuiltinAttributes.h"
 #include "mlir/IR/DialectImplementation.h"
 #include "mlir/IR/Types.h"
 #include "mlir/Support/LogicalResult.h"
 #include "llvm/ADT/Hashing.h"
 #include "llvm/ADT/SetVector.h"
 #include "llvm/ADT/TypeSwitch.h"
 #include "llvm/ADT/bit.h"
 
 using namespace mlir;
 using namespace test;
 
+namespace llvm {
+hash_code hash_value(PtrI64Array s) {
+  // Hash only based on base pointer and size.
+  return hash_combine(&s.front(), s.size());
+}
+} // namespace llvm
+
 //===----------------------------------------------------------------------===//
 // AttrWithSelfTypeParamAttr
 //===----------------------------------------------------------------------===//
 
 Attribute AttrWithSelfTypeParamAttr::parse(AsmParser &parser, Type type) {
   Type selfType;
   if (parser.parseType(selfType))
     return Attribute();
return emitError()
            << elements.size() << ", but expected: " << type.getNumElements();
   }
   if (type.getRank() != 1 || !type.getElementType().isSignlessInteger(64))
     return emitError() << "expected single rank 64-bit shape type, but got: "
                        << type;
   return success();
 }
 
+Attribute TestExtern1DI64ElementsAttr::parse(AsmParser &parser, Type type) {
+  SmallVector<int64_t> elements;
+  if (parser.parseLess() || parser.parseLSquare())
+    return Attribute();
+  int64_t intVal;
+  while (succeeded(*parser.parseOptionalInteger(intVal))) {
+    elements.push_back(intVal);
+    if (parser.parseOptionalComma())
+      break;
+  }
+
+  if (parser.parseRSquare() || parser.parseGreater())
+    return Attribute();
+  // Return a ArrayAttr instead. This means that roundtripping through textual
+  // form loses the extern'ness of the attribute.
+  type = RankedTensorType::get({static_cast<long>(elements.size())},
+                               parser.getBuilder().getI64Type());
+  return DenseIntElementsAttr::get(type, elements);
+}
+
+void TestExtern1DI64ElementsAttr::print(mlir::AsmPrinter &printer) const {
+  printer << "<[";
+  llvm::interleaveComma(getElements(), printer);
+  printer << "]>";
+}
+
 LogicalResult
 TestAttrWithFormatAttr::verify(function_ref<InFlightDiagnostic()> emitError,
                                int64_t one, std::string two, IntegerAttr three,
                                ArrayRef<int> four) {
   if (four.size() != static_cast<unsigned>(one))
     return emitError() << "expected 'one' to equal 'four.size()'";
   return success();
 }

mlir/unittests/IR/AttributeTest.cpp

 //===- AttributeTest.cpp - Attribute unit tests ---------------------------===//
 //
 // 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 "../../test/lib/Dialect/Test/TestDialect.h"
+#include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
 #include "mlir/IR/BuiltinAttributes.h"
 #include "mlir/IR/BuiltinTypes.h"
+#include "mlir/Parser/Parser.h"
+#include "mlir/Support/DebugStringHelper.h"
 #include "gtest/gtest.h"
 
 using namespace mlir;
 using namespace mlir::detail;
 
 template <typename EltTy>
 static void testSplat(Type eltType, const EltTy &splatElt) {
   RankedTensorType shape = RankedTensorType::get({2, 1}, eltType);
TEST(SparseElementsAttrTest, GetZero) {
   EXPECT_EQ(zeroFloatValue.cast<FloatAttr>().getValueAsDouble(), 0.0f);
   EXPECT_TRUE(zeroFloatValue.getType() == floatTy);
 
   auto zeroStringValue = sparseString.getValues<Attribute>()[{1, 1}];
   EXPECT_TRUE(zeroStringValue.cast<StringAttr>().getValue().empty());
   EXPECT_TRUE(zeroStringValue.getType() == stringTy);
 }
 
+TEST(ExternConstantTest, Simple) {
+  MLIRContext context;
+  context.loadDialect<arith::ArithmeticDialect, test::TestDialect>();
+
+  SmallVector<int64_t> val = {10, 20};
+
+  // Create extern constant attribute.
+  auto attr = test::TestExtern1DI64ElementsAttr::get(&context, val);
+  for (auto it : llvm::zip(val, attr.getElements()))
+    EXPECT_EQ(std::get<0>(it), std::get<1>(it));
+  // Verify that it is view into SmallVector above/it is same backing data.
+  val = {30, 40};
+  for (auto it : llvm::zip(val, attr.getElements()))
+    EXPECT_EQ(std::get<0>(it), std::get<1>(it));
+
+  // Print and parse.
+  auto parsedAttr = mlir::parseAttribute(debugString(attr), &context)
+                        .cast<DenseElementsAttr>();
+  // Verify that we get the same uniqued Attribute from parsing as if
+  // constructed as DenseIntElementsAttr.
+  auto denseAttr = DenseIntElementsAttr::get(
+      RankedTensorType::get({2}, IntegerType::get(&context, 64)), val);
+  EXPECT_EQ(denseAttr, parsedAttr);
+
+  // Verify that the parsed attribute is a DenseElementsAttr and doesn't share
+  // the same memory as val.
+  val = {50, 60};
+  for (auto it : llvm::zip(parsedAttr.getValues<int64_t>(), attr.getElements()))
+    EXPECT_NE(std::get<0>(it), std::get<1>(it));
+
+  // The extern constants use the pointer to the data as key and hence two
+  // different arrays with the same value would not be considered equal.
+  SmallVector<int64_t> val2 = val;
+  EXPECT_EQ(attr, test::TestExtern1DI64ElementsAttr::get(&context, val));
+  EXPECT_EQ(test::TestExtern1DI64ElementsAttr::get(&context, val),
+            test::TestExtern1DI64ElementsAttr::get(&context, val));
+  EXPECT_NE(attr, test::TestExtern1DI64ElementsAttr::get(&context, val2));
+  EXPECT_NE(attr, test::TestExtern1DI64ElementsAttr::get(&context, val2));
+
+  // Verify building a tensor constant using extern constant.
+  OpBuilder b(&context);
+  auto constantOp = b.create<arith::ConstantOp>(b.getUnknownLoc(), attr);
+  ASSERT_TRUE(succeeded(constantOp.verify()));
+  ASSERT_TRUE(constantOp.getType().isa<TensorType>());
+}
+
 } // namespace