Implement initial support for function attributes, including parser, printer,
resolver support.
Still TODO are verifier support (to make sure you don't use an attribute for a
function in another module) and the TODO in ModuleParser::finalizeModule that I
will handle in the next patch.
PiperOrigin-RevId: 209361648
diff --git a/include/mlir/IR/Attributes.h b/include/mlir/IR/Attributes.h
index 16b99fb..8ccc73e 100644
--- a/include/mlir/IR/Attributes.h
+++ b/include/mlir/IR/Attributes.h
@@ -22,8 +22,9 @@
#include "llvm/ADT/ArrayRef.h"
namespace mlir {
-class MLIRContext;
class AffineMap;
+class Function;
+class MLIRContext;
class Type;
/// Instances of the Attribute class are immutable, uniqued, immortal, and owned
@@ -38,7 +39,7 @@
Type,
Array,
AffineMap,
- // TODO: Function references.
+ Function,
};
/// Return the classification for this attribute.
@@ -192,6 +193,34 @@
Type *value;
};
+/// A function attribute represents a reference to a function object.
+///
+/// When working with IR, it is important to know that a function attribute can
+/// exist with a null Function inside of it, which occurs when a function object
+/// is deleted that had an attribute which referenced it. No references to this
+/// attribute should persist across the transformation, but that attribute will
+/// remain in MLIRContext.
+class FunctionAttr : public Attribute {
+public:
+ static FunctionAttr *get(Function *value, MLIRContext *context);
+
+ Function *getValue() const { return value; }
+
+ /// Methods for support type inquiry through isa, cast, and dyn_cast.
+ static bool classof(const Attribute *attr) {
+ return attr->getKind() == Kind::Function;
+ }
+
+ /// This function is used by the internals of the Function class to null out
+ /// attributes refering to functions that are about to be deleted.
+ static void dropFunctionReference(Function *value);
+
+private:
+ FunctionAttr(Function *value) : Attribute(Kind::Function), value(value) {}
+ ~FunctionAttr() = delete;
+ Function *value;
+};
+
} // end namespace mlir.
#endif
diff --git a/include/mlir/IR/Builders.h b/include/mlir/IR/Builders.h
index 3324d4b..62c0dcb 100644
--- a/include/mlir/IR/Builders.h
+++ b/include/mlir/IR/Builders.h
@@ -39,6 +39,7 @@
class StringAttr;
class TypeAttr;
class ArrayAttr;
+class FunctionAttr;
class AffineMapAttr;
class AffineMap;
class AffineExpr;
@@ -85,6 +86,7 @@
ArrayAttr *getArrayAttr(ArrayRef<Attribute *> value);
AffineMapAttr *getAffineMapAttr(AffineMap *value);
TypeAttr *getTypeAttr(Type *type);
+ FunctionAttr *getFunctionAttr(Function *value);
// Affine Expressions and Affine Map.
AffineMap *getAffineMap(unsigned dimCount, unsigned symbolCount,
diff --git a/include/mlir/IR/Function.h b/include/mlir/IR/Function.h
index 15e08e3..f3b7aa0 100644
--- a/include/mlir/IR/Function.h
+++ b/include/mlir/IR/Function.h
@@ -66,7 +66,7 @@
protected:
Function(StringRef name, FunctionType *type, Kind kind);
- ~Function() {}
+ ~Function();
private:
Kind kind;
diff --git a/include/mlir/IR/Identifier.h b/include/mlir/IR/Identifier.h
index 70528e3..af16636 100644
--- a/include/mlir/IR/Identifier.h
+++ b/include/mlir/IR/Identifier.h
@@ -40,10 +40,13 @@
Identifier &operator=(const Identifier &other) = default;
/// Return a StringRef for the string.
- StringRef ref() const { return StringRef(pointer, size()); }
+ StringRef strref() const { return StringRef(pointer, size()); }
+
+ /// Identifiers implicitly convert to StringRefs.
+ operator StringRef() const { return strref(); }
/// Return an std::string.
- std::string str() const { return ref().str(); }
+ std::string str() const { return strref().str(); }
/// Return a null terminated C string.
const char *c_str() const { return pointer; }
@@ -59,7 +62,7 @@
}
/// Return true if this identifier is the specified string.
- bool is(StringRef string) const { return ref().equals(string); }
+ bool is(StringRef string) const { return strref().equals(string); }
const char *begin() const { return pointer; }
const char *end() const { return pointer + size(); }
@@ -100,7 +103,7 @@
// Make identifiers hashable.
inline llvm::hash_code hash_value(Identifier arg) {
- return llvm::hash_value(arg.ref());
+ return llvm::hash_value(arg.strref());
}
} // end namespace mlir
diff --git a/include/mlir/IR/Module.h b/include/mlir/IR/Module.h
index f6c543b..38c4ef8 100644
--- a/include/mlir/IR/Module.h
+++ b/include/mlir/IR/Module.h
@@ -60,14 +60,14 @@
// Interfaces for working with the symbol table.
/// Look up a function with the specified name, returning null if no such
- /// name exists.
+ /// name exists. Function names never include the @ on them.
Function *getNamedFunction(StringRef name);
const Function *getNamedFunction(StringRef name) const {
return const_cast<Module *>(this)->getNamedFunction(name);
}
/// Look up a function with the specified name, returning null if no such
- /// name exists.
+ /// name exists. Function names never include the @ on them.
Function *getNamedFunction(Identifier name);
const Function *getNamedFunction(Identifier name) const {