diff --git a/mlir/include/mlir/IR/BuiltinLocationAttributes.td b/mlir/include/mlir/IR/BuiltinLocationAttributes.td --- a/mlir/include/mlir/IR/BuiltinLocationAttributes.td +++ b/mlir/include/mlir/IR/BuiltinLocationAttributes.td @@ -15,10 +15,11 @@ include "mlir/IR/AttrTypeBase.td" include "mlir/IR/BuiltinDialect.td" +include "mlir/IR/SubElementInterfaces.td" // Base class for Builtin dialect location attributes. -class Builtin_LocationAttr - : AttrDef { +class Builtin_LocationAttr traits = []> + : AttrDef { let cppClassName = name; let mnemonic = ?; } @@ -27,7 +28,9 @@ // CallSiteLoc //===----------------------------------------------------------------------===// -def CallSiteLoc : Builtin_LocationAttr<"CallSiteLoc"> { +def CallSiteLoc : Builtin_LocationAttr<"CallSiteLoc", [ + DeclareAttrInterfaceMethods + ]> { let summary = "A callsite source location"; let description = [{ Syntax: @@ -104,7 +107,9 @@ // FusedLoc //===----------------------------------------------------------------------===// -def FusedLoc : Builtin_LocationAttr<"FusedLoc"> { +def FusedLoc : Builtin_LocationAttr<"FusedLoc", [ + DeclareAttrInterfaceMethods + ]> { let summary = "A tuple of other source locations"; let description = [{ Syntax: @@ -143,7 +148,9 @@ // NameLoc //===----------------------------------------------------------------------===// -def NameLoc : Builtin_LocationAttr<"NameLoc"> { +def NameLoc : Builtin_LocationAttr<"NameLoc", [ + DeclareAttrInterfaceMethods + ]> { let summary = "A named source location"; let description = [{ Syntax: @@ -180,7 +187,9 @@ // OpaqueLoc //===----------------------------------------------------------------------===// -def OpaqueLoc : Builtin_LocationAttr<"OpaqueLoc"> { +def OpaqueLoc : Builtin_LocationAttr<"OpaqueLoc", [ + DeclareAttrInterfaceMethods + ]> { let summary = "An opaque source location"; let description = [{ An instance of this location essentially contains a pointer to some data diff --git a/mlir/include/mlir/IR/Location.h b/mlir/include/mlir/IR/Location.h --- a/mlir/include/mlir/IR/Location.h +++ b/mlir/include/mlir/IR/Location.h @@ -15,6 +15,7 @@ #define MLIR_IR_LOCATION_H #include "mlir/IR/Attributes.h" +#include "mlir/IR/SubElementInterfaces.h" #include "llvm/Support/PointerLikeTypeTraits.h" namespace mlir { diff --git a/mlir/lib/IR/Location.cpp b/mlir/lib/IR/Location.cpp --- a/mlir/lib/IR/Location.cpp +++ b/mlir/lib/IR/Location.cpp @@ -80,6 +80,20 @@ return CallSiteLoc::get(name, caller); } +void CallSiteLoc::walkImmediateSubElements( + function_ref walkAttrsFn, + function_ref walkTypesFn) const { + walkAttrsFn(getCallee()); + walkAttrsFn(getCaller()); +} + +Attribute +CallSiteLoc::replaceImmediateSubElements(ArrayRef replAttrs, + ArrayRef replTypes) const { + return get(replAttrs[0].cast(), + replAttrs[1].cast()); +} + //===----------------------------------------------------------------------===// // FusedLoc //===----------------------------------------------------------------------===// @@ -121,3 +135,55 @@ return Base::get(context, locs, metadata); } + +void FusedLoc::walkImmediateSubElements( + function_ref walkAttrsFn, + function_ref walkTypesFn) const { + for (Attribute attr : getLocations()) + walkAttrsFn(attr); + walkAttrsFn(getMetadata()); +} + +Attribute +FusedLoc::replaceImmediateSubElements(ArrayRef replAttrs, + ArrayRef replTypes) const { + SmallVector newLocs; + newLocs.reserve(replAttrs.size() - 1); + for (Attribute attr : replAttrs.drop_back()) + newLocs.push_back(attr.cast()); + return get(getContext(), newLocs, replAttrs.back()); +} + +//===----------------------------------------------------------------------===// +// NameLoc +//===----------------------------------------------------------------------===// + +void NameLoc::walkImmediateSubElements( + function_ref walkAttrsFn, + function_ref walkTypesFn) const { + walkAttrsFn(getName()); + walkAttrsFn(getChildLoc()); +} + +Attribute NameLoc::replaceImmediateSubElements(ArrayRef replAttrs, + ArrayRef replTypes) const { + return get(replAttrs[0].cast(), + replAttrs[1].cast()); +} + +//===----------------------------------------------------------------------===// +// OpaqueLoc +//===----------------------------------------------------------------------===// + +void OpaqueLoc::walkImmediateSubElements( + function_ref walkAttrsFn, + function_ref walkTypesFn) const { + walkAttrsFn(getFallbackLocation()); +} + +Attribute +OpaqueLoc::replaceImmediateSubElements(ArrayRef replAttrs, + ArrayRef replTypes) const { + return get(getUnderlyingLocation(), getUnderlyingTypeID(), + replAttrs[0].cast()); +}