[MLIR][LLVM] Support dso_local_equivalent constants

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.

Author: bcardosolopes

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,

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

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

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;
 }
 

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

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