[MLIR][LLVM] Support dso_local_equivalent constants

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
 //===----------------------------------------------------------------------===//

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 @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,

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<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();
+}
+
+/// 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.
 //===----------------------------------------------------------------------===//

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<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());
+
+    moduleTranslation.mapValue(
+        dsoLocalEquivalentOp.getResult(),
+        llvm::DSOLocalEquivalent::get(cast<llvm::GlobalValue>(llvmValue)));
+    return success();
+  }
+
   return failure();
 }
 

mlir/lib/Target/LLVMIR/ModuleImport.cpp

@@ -1248,6 +1248,18 @@ 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,
+            FlatSymbolRefAttr::get(
+                builder.getContext(),
+                dsoLocalEquivalent->getGlobalValue()->getName()))
+        .getResult();
+  }
+
   // Convert global variable accesses.
   if (auto *globalObj = dyn_cast<llvm::GlobalObject>(constant)) {
     Type type = convertType(globalObj->getType());
@@ -1347,6 +1359,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;
 }
 

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()

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

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()

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()