[MLIR][LLVMIR] Allow calling and invoking aliases

Author: frabert

mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp

@@ -138,14 +138,16 @@ static LogicalResult verifySymbolAttrUse(FlatSymbolRefAttr symbol,
                                          Operation *op,
                                          SymbolTableCollection &symbolTable) {
   StringRef name = symbol.getValue();
-  auto func =
-      symbolTable.lookupNearestSymbolFrom<LLVMFuncOp>(op, symbol.getAttr());
-  if (!func)
-    return op->emitOpError("'")
-           << name << "' does not reference a valid LLVM function";
-  if (func.isExternal())
-    return op->emitOpError("'") << name << "' does not have a definition";
-  return success();
+  auto *symOp = symbolTable.lookupNearestSymbolFrom(op, symbol.getAttr());
+  if (auto func = dyn_cast<LLVMFuncOp>(symOp)) {
+    if (func.isExternal())
+      return op->emitOpError("'") << name << "' does not have a definition";
+    return success();
+  }
+  if (auto alias = dyn_cast<AliasOp>(symOp))
+    return success();
+  return op->emitOpError("'")
+         << name << "' does not reference a valid LLVM function";
 }
 
 /// Returns a boolean type that has the same shape as `type`. It supports both
@@ -1241,14 +1243,16 @@ LogicalResult CallOp::verifySymbolUses(SymbolTableCollection &symbolTable) {
       return emitOpError()
              << "'" << calleeName.getValue()
              << "' does not reference a symbol in the current scope";
-    auto fn = dyn_cast<LLVMFuncOp>(callee);
-    if (!fn)
+    if (auto fn = dyn_cast<LLVMFuncOp>(callee)) {
+      if (failed(verifyCallOpDebugInfo(*this, fn)))
+        return failure();
+      fnType = fn.getFunctionType();
+    } else if (auto alias = dyn_cast<AliasOp>(callee)) {
+      fnType = alias.getAliasType();
+    } else {
       return emitOpError() << "'" << calleeName.getValue()
                            << "' does not reference a valid LLVM function";
-
-    if (failed(verifyCallOpDebugInfo(*this, fn)))
-      return failure();
-    fnType = fn.getFunctionType();
+    }
   }
 
   LLVMFunctionType funcType = llvm::dyn_cast<LLVMFunctionType>(fnType);

mlir/lib/Target/LLVMIR/Dialect/LLVMIR/LLVMToLLVMIRTranslation.cpp

@@ -302,9 +302,19 @@ convertOperationImpl(Operation &opInst, llvm::IRBuilderBase &builder,
     ArrayRef<llvm::Value *> operandsRef(operands);
     llvm::CallInst *call;
     if (auto attr = callOp.getCalleeAttr()) {
+      llvm::GlobalValue *callee;
+      llvm::FunctionType *ftype;
+      if (auto *func = moduleTranslation.lookupFunction(attr.getValue())) {
+        ftype = func->getFunctionType();
+        callee = func;
+      } else {
+        auto &st = moduleTranslation.symbolTable().getSymbolTable(callOp);
+        auto alias = st.lookup<AliasOp>(attr.getValue());
+        callee = moduleTranslation.lookupAlias(alias);
+        ftype = llvm::cast<llvm::FunctionType>(moduleTranslation.convertType(alias.getAliasType()));
+      }
       call =
-          builder.CreateCall(moduleTranslation.lookupFunction(attr.getValue()),
-                             operandsRef, opBundles);
+          builder.CreateCall(ftype, callee, operandsRef, opBundles);
     } else {
       llvm::FunctionType *calleeType = llvm::cast<llvm::FunctionType>(
           moduleTranslation.convertType(callOp.getCalleeFunctionType()));
@@ -415,8 +425,19 @@ convertOperationImpl(Operation &opInst, llvm::IRBuilderBase &builder,
     ArrayRef<llvm::Value *> operandsRef(operands);
     llvm::InvokeInst *result;
     if (auto attr = opInst.getAttrOfType<FlatSymbolRefAttr>("callee")) {
+      llvm::GlobalValue *callee;
+      llvm::FunctionType *ftype;
+      if (auto *func = moduleTranslation.lookupFunction(attr.getValue())) {
+        ftype = func->getFunctionType();
+        callee = func;
+      } else {
+        auto &st = moduleTranslation.symbolTable().getSymbolTable(invOp);
+        auto alias = st.lookup<AliasOp>(attr.getValue());
+        callee = moduleTranslation.lookupAlias(alias);
+        ftype = llvm::cast<llvm::FunctionType>(moduleTranslation.convertType(alias.getAliasType()));
+      }
       result = builder.CreateInvoke(
-          moduleTranslation.lookupFunction(attr.getValue()),
+          ftype, callee,
           moduleTranslation.lookupBlock(invOp.getSuccessor(0)),
           moduleTranslation.lookupBlock(invOp.getSuccessor(1)), operandsRef,
           opBundles);

mlir/test/Dialect/LLVMIR/alias.mlir

@@ -5,7 +5,7 @@ llvm.func internal @callee() -> !llvm.ptr attributes {dso_local} {
   llvm.return %0 : !llvm.ptr
 }
 
-llvm.mlir.alias external @foo_alias : !llvm.ptr {
+llvm.mlir.alias external @foo_alias : !llvm.func<ptr ()> {
   %0 = llvm.mlir.addressof @callee : !llvm.ptr
   llvm.return %0 : !llvm.ptr
 }
@@ -15,6 +15,24 @@ llvm.mlir.alias external @_ZTV1D : !llvm.struct<(array<3 x ptr>)> {
   llvm.return %0 : !llvm.ptr
 }
 
+llvm.func internal @caller() -> !llvm.ptr attributes {dso_local} {
+  %0 = llvm.call @foo_alias() : () -> !llvm.ptr
+  llvm.return %0 : !llvm.ptr
+}
+
+llvm.func @__gxx_personality_v0(...) -> i32
+
+llvm.func internal @invoker() -> !llvm.ptr attributes { dso_local, personality = @__gxx_personality_v0 } {
+  %0 = llvm.mlir.constant("\01") : !llvm.array<1 x i8>
+  %1 = llvm.mlir.zero : !llvm.ptr
+  %2 = llvm.invoke @foo_alias() to ^bb1 unwind ^bb2 : () -> !llvm.ptr
+^bb1:
+  llvm.return %1 : !llvm.ptr
+^bb2:
+  %3 = llvm.landingpad cleanup (catch %1 : !llvm.ptr) (catch %1 : !llvm.ptr) (filter %0 : !llvm.array<1 x i8>) : !llvm.struct<(ptr, i32)>
+  llvm.return %1 : !llvm.ptr
+}
+
 // CHECK: llvm.mlir.alias external @foo_alias : !llvm.ptr {
 // CHECK:   %[[ADDR:.*]] = llvm.mlir.addressof @callee : !llvm.ptr
 // CHECK:   llvm.return %[[ADDR]] : !llvm.ptr