Skip to content

[mlir][llvm] support -new-struct-path-tbaa #119698

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
Open

[mlir][llvm] support -new-struct-path-tbaa #119698

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
73 changes: 72 additions & 1 deletion mlir/include/mlir/Dialect/LLVMIR/LLVMAttrDefs.td
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1084,8 +1084,79 @@ def LLVM_TBAATagAttr : LLVM_Attr<"TBAATag", "tbaa_tag"> {
let assemblyFormat = "`<` struct(params) `>`";
}

def LLVM_TBAAStructFieldAttr : LLVM_Attr<"TBAAStructField", "tbaa_struct_field"> {
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
def LLVM_TBAAStructFieldAttr : LLVM_Attr<"TBAAStructField", "tbaa_struct_field"> {
//===----------------------------------------------------------------------===//
// TBAAStructFieldAttr
//===----------------------------------------------------------------------===//
def LLVM_TBAAStructFieldAttr : LLVM_Attr<"TBAAStructField", "tbaa_struct_field"> {

nit: let's use banners to structure the file a bit.

let parameters = (ins
"TBAANodeAttr":$typeDesc,
"int64_t":$offset,
"int64_t":$size
);
let assemblyFormat = "`<` struct(params) `>`";
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can you use let description / let summary to provide some doc/examples for the relevant attributes.

}


def LLVM_TBAAStructFieldAttrArray : ArrayRefParameter<"TBAAStructFieldAttr"> {
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I wonder why we implemented our own array attributes at the time. @vzakhari do you remember? Potentially we could just use:
OptionalArrayRefParameter<"TBAAStructFieldAttr">:$fields
In LLVM_TBAATypeNodeAttr.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Sorry, I do not remember. It would probably work either way.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@PikachuHyA It is quite a lot of custom printing / parsing that could be saved with OptionalArrayRefParameter<"TBAAStructFieldAttr">:$fields. So maybe we want to give it a shot.

let printer = [{
$_printer << '{';
llvm::interleaveComma($_self, $_printer, [&](TBAAStructFieldAttr attr) {
$_printer.printStrippedAttrOrType(attr);
});
$_printer << '}';
}];

let parser = [{
[&]() -> FailureOr<SmallVector<TBAAStructFieldAttr>> {
using Result = SmallVector<TBAAStructFieldAttr>;
if ($_parser.parseLBrace())
return failure();
FailureOr<Result> result = FieldParser<Result>::parse($_parser);
if (failed(result))
return failure();
if ($_parser.parseRBrace())
return failure();
return result;
}()
}];
}

def LLVM_TBAATypeNodeAttr : LLVM_Attr<"TBAATypeNode", "tbaa_type_node", [], "TBAANodeAttr"> {
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I believe if the class should derive from TBAANodeAttr it also needs to be added here:

llvm-project/mlir/lib/Dialect/LLVMIR/IR/LLVMAttrs.cpp

Line 129 in 9102afc

return llvm::isa<TBAATypeDescriptorAttr, TBAARootAttr>(attr);

let parameters = (ins
"TBAANodeAttr":$parent,
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Are new and old format mixed? Or in other words can this be a TBAATypeDescriptorAttr? If not we may want to have an attribute verifier that ensures the parent is either a root or another LLVM_TBAATypeNodeAttr.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Are new and old format mixed? Or in other words can this be a TBAATypeDescriptorAttr?

No.
The only other TBAANodeAttr used by LLVM_TBAATypeNodeAttr is LLVM_TBAARootAttr.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Makes sense.

If we want to verify this we have sort of two options:

  1. we add a verifier to attribute that checks that parent is not a TBAATypeDescriptorAttr.
  2. We implement a separate base class for the new TBBAA format which could be called TBAAAccessNodeAttr if we follow the naming scheme of the access tag. However, this also requires
    a TBAAAccessRootAttr since an attribute can only inherit from one base class AFAIK.

I do not have a clear preference.

"int64_t":$size,
StringRefParameter<>:$id,
LLVM_TBAAStructFieldAttrArray:$fields
);
let assemblyFormat = "`<` struct(params) `>`";
}

def LLVM_TBAAAccessTagAttr : LLVM_Attr<"TBAAAccessTag", "tbaa_access_tag"> {
let parameters = (ins
"TBAATypeNodeAttr":$base_type,
"TBAATypeNodeAttr":$access_type,
"int64_t":$offset,
"int64_t":$size
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It feels like the difference to TBAATagAttr is minimal (i.e mainly the size). However, I guess it still makes sense to separate them.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Another option is to add an optional "int64_t": $size in TBAATagAttr . However, this approach carries the risk of mixing the old and new formats.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think given LLVM proper has a counter part and since the two TBAA variants cannot be mixed, I am actually in favor of having a separate LLVM_TBAAAccessTagAttr!

);
let builders = [
AttrBuilderWithInferredContext<(ins "TBAATypeNodeAttr":$baseType,
"TBAATypeNodeAttr":$accessType,
"int64_t":$offset,
"int64_t":$size), [{
return $_get(baseType.getContext(), baseType, accessType, offset, size);
}]>
];
let assemblyFormat = "`<` struct(params) `>`";
}

def LLVM_TBAAAccessTagArrayAttr
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is this used?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

no.

The array of LLVM_TBAAAccessTagAttr is used by the following definition:

def LLVM_TBAATagArrayAttr
    : TypedArrayAttrBase<AnyAttrOf<[
  LLVM_TBAATagAttr,
  LLVM_TBAAAccessTagAttr
]>

Therefore, the LLVM_TBAAAccessTagArrayAttr is no longer needed.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

removed.

: TypedArrayAttrBase<LLVM_TBAAAccessTagAttr,
LLVM_TBAAAccessTagAttr.summary # " array"> {
let constBuilderCall = ?;
}

def LLVM_TBAATagArrayAttr
: TypedArrayAttrBase<LLVM_TBAATagAttr,
: TypedArrayAttrBase<AnyAttrOf<[
LLVM_TBAATagAttr,
LLVM_TBAAAccessTagAttr
]>,
LLVM_TBAATagAttr.summary # " array"> {
let constBuilderCall = ?;
}
Expand Down
2 changes: 1 addition & 1 deletion mlir/include/mlir/Target/LLVMIR/ModuleTranslation.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -346,7 +346,7 @@ class ModuleTranslation {

/// Returns the LLVM metadata corresponding to the given mlir LLVM dialect
/// TBAATagAttr.
llvm::MDNode *getTBAANode(TBAATagAttr tbaaAttr) const;
llvm::MDNode *getTBAANode(Attribute tbaaAttr) const;

/// Process tbaa LLVM Metadata operations and create LLVM
/// metadata nodes for them.
Expand Down
3 changes: 2 additions & 1 deletion mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3594,7 +3594,8 @@ struct LLVMOpAsmDialectInterface : public OpAsmDialectInterface {
LoopVectorizeAttr, LoopInterleaveAttr, LoopUnrollAttr,
LoopUnrollAndJamAttr, LoopLICMAttr, LoopDistributeAttr,
LoopPipelineAttr, LoopPeeledAttr, LoopUnswitchAttr, TBAARootAttr,
TBAATagAttr, TBAATypeDescriptorAttr>([&](auto attr) {
TBAATagAttr, TBAATypeDescriptorAttr, TBAAAccessTagAttr,
TBAATypeNodeAttr>([&](auto attr) {
os << decltype(attr)::getMnemonic();
return AliasResult::OverridableAlias;
})
Expand Down
10 changes: 9 additions & 1 deletion mlir/lib/Dialect/LLVMIR/IR/LLVMInterfaces.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -58,7 +58,15 @@ mlir::LLVM::detail::verifyAliasAnalysisOpInterface(Operation *op) {
ArrayAttr tags = iface.getTBAATagsOrNull();
if (!tags)
return success();

if (tags.size() > 0) {
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Shouldn't it just be:
return isArrayOf<TBAATagAttr>(op, tags) || isArrayOf<TBAAAccessTagAttr>(op, tags);

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The isArrayOf<xxx> reports an error if any element in the provided array is not type of xxx.
and the result type LogicalResult dose not support the || operator.

/// Verifies that all elements of `array` are instances of `Attr`.
template <class AttrT>
static LogicalResult isArrayOf(Operation *op, ArrayAttr array) {
  for (Attribute iter : array)
    if (!isa<AttrT>(iter))
      return op->emitOpError("expected op to return array of ")
             << AttrT::getMnemonic() << " attributes";
  return success();
}

see
https://github.com/llvm/llvm-project/blob/main/mlir/lib/Dialect/LLVMIR/IR/LLVMInterfaces.cpp#L19-L27

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Ok I understand. I would suggest to extend the AliasAnalysisOpInterface. It currently has one function that returns the TBAA tags. How about implementing two functions:

iface.getTBBATagsOrNull()

and

iface.getTBBAAccessTagsOrNull()

This could potentially be implemented by having two arrays stored on the alias operations. Or by checking if the array contains the old or the new kind of TBAA attributes. That way the verifier can use the existing logic for the two cases separately. And to users of the exiting TBAA attributes don't need to understand the new kind of TBAA metadata.

WDYT?

if (mlir::isa<TBAATagAttr>(tags[0])) {
return isArrayOf<TBAATagAttr>(op, tags);
}

if (mlir::isa<TBAAAccessTagAttr>(tags[0])) {
return isArrayOf<TBAAAccessTagAttr>(op, tags);
}
}
return isArrayOf<TBAATagAttr>(op, tags);
}

Expand Down
36 changes: 34 additions & 2 deletions mlir/lib/Target/LLVMIR/ModuleTranslation.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1882,7 +1882,7 @@ void ModuleTranslation::setAliasScopeMetadata(AliasAnalysisOpInterface op,
llvm::LLVMContext::MD_noalias);
}

llvm::MDNode *ModuleTranslation::getTBAANode(TBAATagAttr tbaaAttr) const {
llvm::MDNode *ModuleTranslation::getTBAANode(Attribute tbaaAttr) const {
return tbaaMetadataMapping.lookup(tbaaAttr);
}

Expand All @@ -1902,7 +1902,7 @@ void ModuleTranslation::setTBAAMetadata(AliasAnalysisOpInterface op,
return;
}

llvm::MDNode *node = getTBAANode(cast<TBAATagAttr>(tagRefs[0]));
llvm::MDNode *node = getTBAANode(tagRefs[0]);
inst->setMetadata(llvm::LLVMContext::MD_tbaa, node);
}

Expand All @@ -1922,6 +1922,7 @@ void ModuleTranslation::setBranchWeightsMetadata(BranchWeightOpInterface op) {
LogicalResult ModuleTranslation::createTBAAMetadata() {
llvm::LLVMContext &ctx = llvmModule->getContext();
llvm::IntegerType *offsetTy = llvm::IntegerType::get(ctx, 64);
llvm::IntegerType *sizeTy = llvm::IntegerType::get(ctx, 64);

// Walk the entire module and create all metadata nodes for the TBAA
// attributes. The code below relies on two invariants of the
Expand Down Expand Up @@ -1949,6 +1950,23 @@ LogicalResult ModuleTranslation::createTBAAMetadata() {
tbaaMetadataMapping.insert({descriptor, llvm::MDNode::get(ctx, operands)});
});

walker.addWalk([&](TBAATypeNodeAttr descriptor) {
SmallVector<llvm::Metadata *> operands;
operands.push_back(tbaaMetadataMapping.lookup(descriptor.getParent()));
operands.push_back(llvm::ConstantAsMetadata::get(
llvm::ConstantInt::get(sizeTy, descriptor.getSize())));
operands.push_back(llvm::MDString::get(ctx, descriptor.getId()));
for (auto field : descriptor.getFields()) {
operands.push_back(tbaaMetadataMapping.lookup(field.getTypeDesc()));
operands.push_back(llvm::ConstantAsMetadata::get(
llvm::ConstantInt::get(offsetTy, field.getOffset())));
operands.push_back(llvm::ConstantAsMetadata::get(
llvm::ConstantInt::get(sizeTy, field.getSize())));
}

tbaaMetadataMapping.insert({descriptor, llvm::MDNode::get(ctx, operands)});
});

walker.addWalk([&](TBAATagAttr tag) {
SmallVector<llvm::Metadata *> operands;

Expand All @@ -1964,6 +1982,20 @@ LogicalResult ModuleTranslation::createTBAAMetadata() {
tbaaMetadataMapping.insert({tag, llvm::MDNode::get(ctx, operands)});
});

walker.addWalk([&](TBAAAccessTagAttr tag) {
SmallVector<llvm::Metadata *> operands;

operands.push_back(tbaaMetadataMapping.lookup(tag.getBaseType()));
operands.push_back(tbaaMetadataMapping.lookup(tag.getAccessType()));

operands.push_back(llvm::ConstantAsMetadata::get(
llvm::ConstantInt::get(offsetTy, tag.getOffset())));
operands.push_back(llvm::ConstantAsMetadata::get(
llvm::ConstantInt::get(sizeTy, tag.getSize())));

tbaaMetadataMapping.insert({tag, llvm::MDNode::get(ctx, operands)});
});

mlirModule->walk([&](AliasAnalysisOpInterface analysisOpInterface) {
if (auto attr = analysisOpInterface.getTBAATagsOrNull())
walker.walk(attr);
Expand Down