diff --git a/mlir/include/mlir/Dialect/LLVMIR/LLVMAttrDefs.td b/mlir/include/mlir/Dialect/LLVMIR/LLVMAttrDefs.td index ede9f2d365b20..41c30b81770bc 100644 --- a/mlir/include/mlir/Dialect/LLVMIR/LLVMAttrDefs.td +++ b/mlir/include/mlir/Dialect/LLVMIR/LLVMAttrDefs.td @@ -1213,6 +1213,17 @@ def LLVM_UndefAttr : LLVM_Attr<"Undef", "undef">; /// Folded into from LLVM::PoisonOp. def LLVM_PoisonAttr : LLVM_Attr<"Poison", "poison">; +//===----------------------------------------------------------------------===// +// DSOLocalEquivalentAttr +//===----------------------------------------------------------------------===// + +/// Folded into from LLVM::DSOLocalEquivalentOp. +def LLVM_DSOLocalEquivalentAttr : LLVM_Attr<"DSOLocalEquivalent", + "dso_local_equivalent"> { + let parameters = (ins "FlatSymbolRefAttr":$sym); + let assemblyFormat = "$sym"; +} + //===----------------------------------------------------------------------===// // VecTypeHintAttr //===----------------------------------------------------------------------===// diff --git a/mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td b/mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td index 90cc851c0a3b2..5e1ab32e0a53d 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]> { + 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 @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..309cee6429af8 100644 --- a/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp +++ b/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp @@ -2122,6 +2122,57 @@ OpFoldResult LLVM::AddressOfOp::fold(FoldAdaptor) { return getGlobalNameAttr(); } +//===----------------------------------------------------------------------===// +// LLVM::DSOLocalEquivalentOp +//===----------------------------------------------------------------------===// + +LLVMFuncOp +DSOLocalEquivalentOp::getFunction(SymbolTableCollection &symbolTable) { + return dyn_cast_or_null(symbolTable.lookupSymbolIn( + parentLLVMModule(*this), getFunctionNameAttr())); +} + +AliasOp DSOLocalEquivalentOp::getAlias(SymbolTableCollection &symbolTable) { + return dyn_cast_or_null(symbolTable.lookupSymbolIn( + parentLLVMModule(*this), getFunctionNameAttr())); +} + +LogicalResult +DSOLocalEquivalentOp::verifySymbolUses(SymbolTableCollection &symbolTable) { + Operation *symbol = symbolTable.lookupSymbolIn(parentLLVMModule(*this), + getFunctionNameAttr()); + auto function = dyn_cast_or_null(symbol); + auto alias = dyn_cast_or_null(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(); +} + +/// Fold a dso_local_equivalent operation to a dedicated dso_local_equivalent +/// attribute. +OpFoldResult DSOLocalEquivalentOp::fold(FoldAdaptor) { + return DSOLocalEquivalentAttr::get(getContext(), 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..f8fa22253bea8 100644 --- a/mlir/lib/Target/LLVMIR/Dialect/LLVMIR/LLVMToLLVMIRTranslation.cpp +++ b/mlir/lib/Target/LLVMIR/Dialect/LLVMIR/LLVMToLLVMIRTranslation.cpp @@ -526,6 +526,31 @@ 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(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()); + + moduleTranslation.mapValue( + dsoLocalEquivalentOp.getResult(), + llvm::DSOLocalEquivalent::get(cast(llvmValue))); + return success(); + } + return failure(); } diff --git a/mlir/lib/Target/LLVMIR/ModuleImport.cpp b/mlir/lib/Target/LLVMIR/ModuleImport.cpp index a07189ae1323c..a1a6c001d602f 100644 --- a/mlir/lib/Target/LLVMIR/ModuleImport.cpp +++ b/mlir/lib/Target/LLVMIR/ModuleImport.cpp @@ -1248,6 +1248,18 @@ FailureOr ModuleImport::convertConstant(llvm::Constant *constant) { return builder.create(loc, type).getResult(); } + // Convert dso_local_equivalent. + if (auto *dsoLocalEquivalent = dyn_cast(constant)) { + Type type = convertType(dsoLocalEquivalent->getType()); + return builder + .create( + loc, type, + FlatSymbolRefAttr::get( + builder.getContext(), + dsoLocalEquivalent->getGlobalValue()->getName())) + .getResult(); + } + // Convert global variable accesses. if (auto *globalObj = dyn_cast(constant)) { Type type = convertType(globalObj->getType()); @@ -1347,6 +1359,15 @@ FailureOr ModuleImport::convertConstant(llvm::Constant *constant) { if (isa(constant)) error = " since blockaddress(...) is unsupported"; + if (isa(constant)) + error = " since ptrauth(...) is unsupported"; + + if (isa(constant)) + error = " since no_cfi is unsupported"; + + if (isa(constant)) + error = " since global value is unsupported"; + return emitError(loc) << "unhandled constant: " << diag(*constant) << error; } diff --git a/mlir/test/Dialect/LLVMIR/constant-folding.mlir b/mlir/test/Dialect/LLVMIR/constant-folding.mlir index a3570a11d4a2d..99f657f0aefec 100644 --- a/mlir/test/Dialect/LLVMIR/constant-folding.mlir +++ b/mlir/test/Dialect/LLVMIR/constant-folding.mlir @@ -182,3 +182,17 @@ func.func @insert_op(%arg0: index, %arg1: memref<13x13xi64>, %arg2: index) { vector.print %101 : vector<1xi64> return } + +// ----- + +// CHECK-LABEL: llvm.func @dso_local_equivalent_select +llvm.func @dso_local_equivalent_select(%arg: i1) -> !llvm.ptr { + // CHECK-NEXT: %[[DSOLOCALEQ:.+]] = llvm.dso_local_equivalent @yay + %0 = llvm.dso_local_equivalent @yay : !llvm.ptr + %1 = llvm.dso_local_equivalent @yay : !llvm.ptr + %2 = arith.select %arg, %0, %1 : !llvm.ptr + // CHECK-NEXT: llvm.return %[[DSOLOCALEQ]] + llvm.return %2 : !llvm.ptr +} + +llvm.func @yay() 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..83566d6649932 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 @y() : !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 @y : !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 { + %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()