[MLIR][LLVM] Support dso_local_equivalent constants #132131
Conversation
This is modeled after a new operation since its constraints differ from AddressOfOp: only support functions and function aliases (no globals) and extern_weak linkage is not allowed. An alternative approach is to use a UnitAttr in AddressOfOp and use the presence of that attribute to enforce specific semantics in the verifiers, but that sounds worse IMO. Note that this is similar in nature to no_cfi, ptrauth and blockaddress, which when introduced in the future will likely be implenented in similar fashion. While here, improve the error message for other missing constants.
|
@llvm/pr-subscribers-mlir @llvm/pr-subscribers-mlir-llvm Author: Bruno Cardoso Lopes (bcardosolopes) ChangesCreate a new operation This is similar in a way to An alternative approach is to use a new While here, improve the error message for other missing constants. Full diff: https://github.com/llvm/llvm-project/pull/132131.diff 7 Files Affected:
diff --git a/mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td b/mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td
index 90cc851c0a3b2..a4493b0d0da34 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td
@@ -1575,6 +1575,44 @@ def LLVM_AliasOp : LLVM_Op<"mlir.alias",
let hasRegionVerifier = 1;
}
+def LLVM_DSOLocalEquivalentOp : LLVM_Op<"dso_local_equivalent",
+ [Pure, ConstantLike, DeclareOpInterfaceMethods<SymbolUserOpInterface>]> {
+ let arguments = (ins FlatSymbolRefAttr:$function_name);
+ let results = (outs LLVM_AnyPointer:$res);
+
+ let summary = "Creates a LLVM dso_local_equivalent ptr";
+
+ let description = [{
+ Creates an SSA value containing a pointer to a global value (function or
+ alias to function). It represents a function which is functionally
+ equivalent to a given function, but is always defined in the current
+ linkage unit. The target function may not have `extern_weak` linkage.
+
+ Examples:
+
+ ```mlir
+ llvm.mlir.global external constant @const() : i64 {
+ %0 = llvm.mlir.addressof @const : !llvm.ptr
+ %1 = llvm.ptrtoint %0 : !llvm.ptr to i64
+ %2 = llvm.dso_local_equivalent @extern_func : !llvm.ptr
+ %4 = llvm.ptrtoint %2 : !llvm.ptr to i64
+ llvm.return %4 : i64
+ }
+ ```
+ }];
+
+ let extraClassDeclaration = [{
+ /// Return the llvm.func operation that is referenced here.
+ LLVMFuncOp getFunction(SymbolTableCollection &symbolTable);
+ /// Return the llvm.mlir.alias operation that defined the value referenced
+ /// here.
+ AliasOp getAlias(SymbolTableCollection &symbolTable);
+ }];
+
+ let assemblyFormat = "$function_name attr-dict `:` qualified(type($res))";
+ let hasFolder = 1;
+}
+
def LLVM_ComdatSelectorOp : LLVM_Op<"comdat_selector", [Symbol]> {
let arguments = (ins
SymbolNameAttr:$sym_name,
diff --git a/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp b/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp
index 5370de501a85c..1e9d4d350e9f2 100644
--- a/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp
+++ b/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp
@@ -2122,6 +2122,56 @@ OpFoldResult LLVM::AddressOfOp::fold(FoldAdaptor) {
return getGlobalNameAttr();
}
+//===----------------------------------------------------------------------===//
+// LLVM::DSOLocalEquivalentOp
+//===----------------------------------------------------------------------===//
+
+LLVMFuncOp
+DSOLocalEquivalentOp::getFunction(SymbolTableCollection &symbolTable) {
+ return dyn_cast_or_null<LLVMFuncOp>(symbolTable.lookupSymbolIn(
+ parentLLVMModule(*this), getFunctionNameAttr()));
+}
+
+AliasOp DSOLocalEquivalentOp::getAlias(SymbolTableCollection &symbolTable) {
+ return dyn_cast_or_null<AliasOp>(symbolTable.lookupSymbolIn(
+ parentLLVMModule(*this), getFunctionNameAttr()));
+}
+
+LogicalResult
+DSOLocalEquivalentOp::verifySymbolUses(SymbolTableCollection &symbolTable) {
+ Operation *symbol = symbolTable.lookupSymbolIn(parentLLVMModule(*this),
+ getFunctionNameAttr());
+ auto function = dyn_cast_or_null<LLVMFuncOp>(symbol);
+ auto alias = dyn_cast_or_null<AliasOp>(symbol);
+
+ if (!function && !alias)
+ return emitOpError(
+ "must reference a global defined by 'llvm.func' or 'llvm.mlir.alias'");
+
+ if (alias) {
+ if (alias.getInitializer()
+ .walk([&](AddressOfOp addrOp) {
+ if (addrOp.getGlobal(symbolTable))
+ return WalkResult::interrupt();
+ return WalkResult::advance();
+ })
+ .wasInterrupted())
+ return emitOpError("must reference an alias to a function");
+ }
+
+ if ((function && function.getLinkage() == LLVM::Linkage::ExternWeak) ||
+ (alias && alias.getLinkage() == LLVM::Linkage::ExternWeak))
+ return emitOpError(
+ "target function with 'extern_weak' linkage not allowed");
+
+ return success();
+}
+
+// DSOLocalEquivalentOp constant-folds to the global symbol name.
+OpFoldResult DSOLocalEquivalentOp::fold(FoldAdaptor) {
+ return getFunctionNameAttr();
+}
+
//===----------------------------------------------------------------------===//
// Verifier for LLVM::ComdatOp.
//===----------------------------------------------------------------------===//
diff --git a/mlir/lib/Target/LLVMIR/Dialect/LLVMIR/LLVMToLLVMIRTranslation.cpp b/mlir/lib/Target/LLVMIR/Dialect/LLVMIR/LLVMToLLVMIRTranslation.cpp
index 833b19c1bece2..3bf5f7b3196fc 100644
--- a/mlir/lib/Target/LLVMIR/Dialect/LLVMIR/LLVMToLLVMIRTranslation.cpp
+++ b/mlir/lib/Target/LLVMIR/Dialect/LLVMIR/LLVMToLLVMIRTranslation.cpp
@@ -526,6 +526,32 @@ convertOperationImpl(Operation &opInst, llvm::IRBuilderBase &builder,
return success();
}
+ // Emit dso_local_equivalent. We need to look up the global value referenced
+ // by the operation and store it in the MLIR-to-LLVM value mapping.
+ if (auto dsoLocalEquivalentOp =
+ dyn_cast<LLVM::DSOLocalEquivalentOp>(opInst)) {
+ LLVM::LLVMFuncOp function =
+ dsoLocalEquivalentOp.getFunction(moduleTranslation.symbolTable());
+ LLVM::AliasOp alias =
+ dsoLocalEquivalentOp.getAlias(moduleTranslation.symbolTable());
+
+ // The verifier should not have allowed this.
+ assert((function || alias) &&
+ "referencing an undefined function, or alias");
+
+ llvm::Value *llvmValue = nullptr;
+ if (alias)
+ llvmValue = moduleTranslation.lookupAlias(alias);
+ else
+ llvmValue = moduleTranslation.lookupFunction(function.getName());
+
+ auto gv = dyn_cast_or_null<llvm::GlobalValue>(llvmValue);
+ assert(gv && "expected LLVM IR global value");
+ moduleTranslation.mapValue(dsoLocalEquivalentOp.getResult(),
+ llvm::DSOLocalEquivalent::get(gv));
+ return success();
+ }
+
return failure();
}
diff --git a/mlir/lib/Target/LLVMIR/ModuleImport.cpp b/mlir/lib/Target/LLVMIR/ModuleImport.cpp
index a07189ae1323c..3af93a15f1623 100644
--- a/mlir/lib/Target/LLVMIR/ModuleImport.cpp
+++ b/mlir/lib/Target/LLVMIR/ModuleImport.cpp
@@ -1248,6 +1248,15 @@ FailureOr<Value> ModuleImport::convertConstant(llvm::Constant *constant) {
return builder.create<UndefOp>(loc, type).getResult();
}
+ // Convert dso_local_equivalent.
+ if (auto *dsoLocalEquivalent = dyn_cast<llvm::DSOLocalEquivalent>(constant)) {
+ Type type = convertType(dsoLocalEquivalent->getType());
+ return builder
+ .create<DSOLocalEquivalentOp>(
+ loc, type, dsoLocalEquivalent->getGlobalValue()->getName())
+ .getResult();
+ }
+
// Convert global variable accesses.
if (auto *globalObj = dyn_cast<llvm::GlobalObject>(constant)) {
Type type = convertType(globalObj->getType());
@@ -1347,6 +1356,15 @@ FailureOr<Value> ModuleImport::convertConstant(llvm::Constant *constant) {
if (isa<llvm::BlockAddress>(constant))
error = " since blockaddress(...) is unsupported";
+ if (isa<llvm::ConstantPtrAuth>(constant))
+ error = " since ptrauth(...) is unsupported";
+
+ if (isa<llvm::NoCFIValue>(constant))
+ error = " since no_cfi is unsupported";
+
+ if (isa<llvm::GlobalValue>(constant))
+ error = " since global value is unsupported";
+
return emitError(loc) << "unhandled constant: " << diag(*constant) << error;
}
diff --git a/mlir/test/Target/LLVMIR/Import/constant.ll b/mlir/test/Target/LLVMIR/Import/constant.ll
index 3c46f5b20c31c..3c5f5825d47ee 100644
--- a/mlir/test/Target/LLVMIR/Import/constant.ll
+++ b/mlir/test/Target/LLVMIR/Import/constant.ll
@@ -236,3 +236,41 @@ define i64 @const_exprs_with_duplicate() {
; CHECK: %[[VAL0:.+]] = llvm.ptrtoint %[[ADDR]]
; CHECK: %[[VAL1:.+]] = llvm.add %[[VAL0]], %[[VAL0]]
; CHECK: llvm.return %[[VAL1]]
+
+; // -----
+
+declare void @extern_func()
+@const = dso_local constant i32 trunc (i64 sub (i64 ptrtoint (ptr dso_local_equivalent @extern_func to i64), i64 ptrtoint (ptr @const to i64)) to i32)
+
+; CHECK: llvm.mlir.global external constant @const()
+; CHECK: %[[ADDR:.+]] = llvm.mlir.addressof @const : !llvm.ptr
+; CHECK: llvm.ptrtoint %[[ADDR]] : !llvm.ptr to i64
+; CHECK: llvm.dso_local_equivalent @extern_func : !llvm.ptr
+
+; // -----
+
+declare i32 @extern_func()
+
+define void @call_extern_func() {
+ call noundef i32 dso_local_equivalent @extern_func()
+ ret void
+}
+
+; CHECK-LABEL: @call_extern_func()
+; CHECK: %[[DSO_EQ:.+]] = llvm.dso_local_equivalent @extern_func : !llvm.ptr
+; CHECK: llvm.call %[[DSO_EQ]]() : !llvm.ptr, () -> (i32 {llvm.noundef})
+
+; // -----
+
+define void @aliasee_func() {
+ ret void
+}
+
+@alias_func = alias void (), ptr @aliasee_func
+define void @call_alias_func() {
+ call void dso_local_equivalent @alias_func()
+ ret void
+}
+
+; CHECK-LABEL: @call_alias_func()
+; CHECK: llvm.dso_local_equivalent @alias_func : !llvm.ptr
diff --git a/mlir/test/Target/LLVMIR/llvmir-invalid.mlir b/mlir/test/Target/LLVMIR/llvmir-invalid.mlir
index f755c4e508c22..273be96b42771 100644
--- a/mlir/test/Target/LLVMIR/llvmir-invalid.mlir
+++ b/mlir/test/Target/LLVMIR/llvmir-invalid.mlir
@@ -411,3 +411,45 @@ module @no_known_conversion_innermost_eltype {
}
}
#-}
+
+// -----
+
+llvm.mlir.global external @zed(42 : i32) : i32
+
+llvm.mlir.alias external @foo : i32 {
+ %0 = llvm.mlir.addressof @zed : !llvm.ptr
+ llvm.return %0 : !llvm.ptr
+}
+
+llvm.func @call_alias_func() {
+ // expected-error @below{{'llvm.dso_local_equivalent' op must reference an alias to a function}}
+ %0 = llvm.dso_local_equivalent @foo : !llvm.ptr
+ llvm.call %0() : !llvm.ptr, () -> (i32)
+ llvm.return
+}
+
+// -----
+
+llvm.mlir.global external @zed() : !llvm.ptr
+
+llvm.func @call_alias_func() {
+ // expected-error @below{{op must reference a global defined by 'llvm.func' or 'llvm.mlir.alias'}}
+ %0 = llvm.dso_local_equivalent @foo : !llvm.ptr
+ llvm.call %0() : !llvm.ptr, () -> (i32)
+ llvm.return
+}
+
+// -----
+
+llvm.mlir.global external constant @const() {addr_space = 0 : i32, dso_local} : i32 {
+ %0 = llvm.mlir.addressof @const : !llvm.ptr
+ %1 = llvm.ptrtoint %0 : !llvm.ptr to i64
+ // expected-error @below{{'llvm.dso_local_equivalent' op target function with 'extern_weak' linkage not allowed}}
+ %2 = llvm.dso_local_equivalent @extern_func : !llvm.ptr
+ %3 = llvm.ptrtoint %2 : !llvm.ptr to i64
+ %4 = llvm.sub %3, %1 : i64
+ %5 = llvm.trunc %4 : i64 to i32
+ llvm.return %5 : i32
+}
+
+llvm.func extern_weak @extern_func()
diff --git a/mlir/test/Target/LLVMIR/llvmir.mlir b/mlir/test/Target/LLVMIR/llvmir.mlir
index 5a1f43ba1d018..2476b1b15faaa 100644
--- a/mlir/test/Target/LLVMIR/llvmir.mlir
+++ b/mlir/test/Target/LLVMIR/llvmir.mlir
@@ -2793,3 +2793,52 @@ module {
// CHECK: !llvm.module.flags = !{![[#DBG:]]}
// CHECK: ![[#DBG]] = !{i32 2, !"Debug Info Version", i32 3}
+
+// -----
+
+llvm.mlir.global external constant @const() {addr_space = 0 : i32, dso_local} : i32 {
+ %0 = llvm.mlir.addressof @const : !llvm.ptr
+ %1 = llvm.ptrtoint %0 : !llvm.ptr to i64
+ %2 = llvm.dso_local_equivalent @extern_func : !llvm.ptr
+ %3 = llvm.ptrtoint %2 : !llvm.ptr to i64
+ %4 = llvm.sub %3, %1 : i64
+ %5 = llvm.trunc %4 : i64 to i32
+ llvm.return %5 : i32
+}
+
+llvm.func @extern_func()
+
+// CHECK: @const = dso_local constant i32 trunc
+// CHECK-SAME: (i64 sub (i64 ptrtoint
+// CHECK-SAME: (ptr dso_local_equivalent @extern_func to i64),
+// CHECK-SAME: i64 ptrtoint (ptr @const to i64)) to i32)
+
+// -----
+
+llvm.func @extern_func() -> i32
+llvm.func @call_extern_func() {
+ %0 = llvm.dso_local_equivalent @extern_func : !llvm.ptr
+ %1 = llvm.call %0() : !llvm.ptr, () -> (i32 {llvm.noundef})
+ llvm.return
+}
+
+// CHECK-LABEL: @call_extern_func()
+// CHECK: call noundef i32 dso_local_equivalent @extern_func()
+
+// -----
+
+llvm.mlir.alias external @alias_func : !llvm.func<void ()> {
+ %0 = llvm.mlir.addressof @aliasee_func : !llvm.ptr
+ llvm.return %0 : !llvm.ptr
+}
+llvm.func @aliasee_func() {
+ llvm.return
+}
+llvm.func @call_alias_func() {
+ %0 = llvm.dso_local_equivalent @alias_func : !llvm.ptr
+ llvm.call %0() : !llvm.ptr, () -> ()
+ llvm.return
+}
+
+// CHECK-LABEL: @call_alias_func
+// CHECK: call void dso_local_equivalent @alias_func()
|
There was a problem hiding this comment.
Thanks! Code looks good to me.
There is a slight issue with constant folding this. Since it produces a FlatSymbolRefAttr like AddressOfOp it currently would be constant folded to an AddressOfOp if I understand correctly (see my comment on the PR). So maybe this is a hint that we should model this using AddressOfOp.
Is LLVM always generating a DSOLocalEquivalent if the symbol references a DSOLocalEquivalent? If this can be reliably computed from the IR we may just add a helper function to AddressOfOp that computes if it references a "DSOLocalEquivalent" and use this in the export to create the correct constant type. That way we avoid having two ops that model basically the same thing.
Alternatively, we may have to introduce a DSOLocalSymbolRefAttr or similar so that we can reliable materialize the correct op type. Even then it may make sense to still use AddressOfOp to have one way of referencing a symbol.
I think I favor computing this property from the IR if this is possible.
mlir/lib/Target/LLVMIR/Dialect/LLVMIR/LLVMToLLVMIRTranslation.cpp
Outdated
Show resolved
Hide resolved
|
|
||
| // DSOLocalEquivalentOp constant-folds to the global symbol name. | ||
| OpFoldResult DSOLocalEquivalentOp::fold(FoldAdaptor) { | ||
| return getFunctionNameAttr(); |
There was a problem hiding this comment.
For this to work we also need to update materializeConstant (
).However, I wonder how we can decide there if something is an AddressOf or a DSOLocalEquivalentOp? I guess we would needs some other attribute type to achieve this?
There was a problem hiding this comment.
Good point, that explains the comments for Undef as others in LLVMAttrs.td, I agree we need an attribute.
There was a problem hiding this comment.
I was searching for ways to test materializeConstant, but when I comment out most of this function and run ninja check-mlir, it all passes, which makes this either non-tested or not being used anymore?
Operation *LLVMDialect::materializeConstant(OpBuilder &builder, Attribute value,
Type type, Location loc) {
// If this was folded from an operation other than llvm.mlir.constant, it
// should be materialized as such. Note that an llvm.mlir.zero may fold into
// a builtin zero attribute and thus will materialize as a llvm.mlir.constant.
// if (auto symbol = dyn_cast<FlatSymbolRefAttr>(value))
// if (isa<LLVM::LLVMPointerType>(type))
// return builder.create<LLVM::AddressOfOp>(loc, type, symbol);
// if (isa<LLVM::UndefAttr>(value))
// return builder.create<LLVM::UndefOp>(loc, type);
// if (isa<LLVM::PoisonAttr>(value))
// return builder.create<LLVM::PoisonOp>(loc, type);
// if (isa<LLVM::ZeroAttr>(value))
// return builder.create<LLVM::ZeroOp>(loc, type);
// Otherwise try materializing it as a regular llvm.mlir.constant op.
return LLVM::ConstantOp::materialize(builder, value, type, loc);
}
Any ideas on how to test any modifications of this function?
True, thanks for pointing this out!
If we go this way, I'm curious if should we do the same in the future for
Can you elaborate a bit? Not sure I understand what you mean. My understanding is that Example, this is valid: This is invalid:
Gotcha, let me explore a bit more and see what can be done. |
I did read up a bit more what this dso_local_equivalent means. I do not claim I fully understand it but it is essentially a symbol reference that maybe replaced by some other value at runtime (in a dynamic linking scenario). That means it is not computable in IR and indeed needs to be represented in some way.
I think for the current problem I would probably keep the op you introduced (passes that consume addressof likely would not work with a dso_local_equivalent symbol anyways) and add some special attribute to make materialization work. |
It does hehe, we have uses of it, though not as many (hopefully not hidden other currently failures) - I have a few ideas but let's wait until the time comes! |
|
Address all reviews. @gysit I added a testcase to |
I didn't have much success (see inline comments), but did one last round of improvements by using |
Create a new operation
DSOLocalEquivalentOp, following the steps of other constants.This is similar in a way to
AddressOfOpbut with specific semantics: only support functions and function aliases (no globals) and extern_weak linkage is not allowed.An alternative approach is to use a new
UnitAttrinAddressOfOpand check that attribute to enforce specific semantics in the verifiers. The drawback is going against what other constants do and having to add more attributes in the future when we introduceno_cfi,blockaddress, etc.While here, improve the error message for other missing constants.