[MLIR][LLVMIR] Allow calling and invoking aliases

Author: frabert

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

@@ -65,14 +65,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
@@ -1184,10 +1186,12 @@ LogicalResult CallOp::verifySymbolUses(SymbolTableCollection &symbolTable) {
       fnType = fn.getFunctionType();
     } else if (auto ifunc = dyn_cast<IFuncOp>(callee)) {
       fnType = ifunc.getIFuncType();
+    } else if (auto alias = dyn_cast<AliasOp>(callee)) {
+      fnType = alias.getAliasType();
     } else {
-      return emitOpError()
-             << "'" << calleeName.getValue()
-             << "' does not reference a valid LLVM function or IFunc";
+        return emitOpError()
+               << "'" << calleeName.getValue()
+               << "' does not reference a valid LLVM function or IFunc";
     }
   }
 

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

@@ -426,13 +426,20 @@ convertOperationImpl(Operation &opInst, llvm::IRBuilderBase &builder,
               moduleTranslation.lookupFunction(attr.getValue())) {
         call = builder.CreateCall(function, operandsRef, opBundles);
       } else {
-        Operation *moduleOp = parentLLVMModule(&opInst);
-        Operation *ifuncOp =
-            moduleTranslation.symbolTable().lookupSymbolIn(moduleOp, attr);
-        llvm::GlobalValue *ifunc = moduleTranslation.lookupIFunc(ifuncOp);
-        llvm::FunctionType *calleeType = llvm::cast<llvm::FunctionType>(
-            moduleTranslation.convertType(callOp.getCalleeFunctionType()));
-        call = builder.CreateCall(calleeType, ifunc, operandsRef, opBundles);
+        auto &st = moduleTranslation.symbolTable().getSymbolTable(callOp);
+        auto alias = st.lookup<AliasOp>(attr.getValue());
+        if (alias) {
+          llvm::GlobalValue *callee = moduleTranslation.lookupAlias(alias);
+          llvm::FunctionType *ftype = llvm::cast<llvm::FunctionType>(moduleTranslation.convertType(alias.getAliasType()));
+          call = builder.CreateCall(ftype, callee, operandsRef, opBundles);
+        } else {
+          Operation *moduleOp = parentLLVMModule(&opInst);
+          Operation *ifuncOp = st.lookupSymbolIn(moduleOp, attr);
+          llvm::GlobalValue *ifunc = moduleTranslation.lookupIFunc(ifuncOp);
+          llvm::FunctionType *calleeType = llvm::cast<llvm::FunctionType>(
+              moduleTranslation.convertType(callOp.getCalleeFunctionType()));
+          call = builder.CreateCall(calleeType, ifunc, operandsRef, opBundles);
+        }
       }
     } else {
       llvm::FunctionType *calleeType = llvm::cast<llvm::FunctionType>(
@@ -554,8 +561,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