[mlir][LLVMIR] Apply CallOp/CallableInterface on suitable operations

[myhsu]

Summary

• Applying CallOpInterface on CallOp and InvokeOp.
• Applying CallableInterface on LLVMFuncOp.

We're testing the changes using CallGraph, which uses both interfaces.

[myhsu]

Add comments on some of the interface methods. NFC.

[Mogball]

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

903

Spell out the auto

[myhsu]

Spell out auto types. NFC.
@Mogball thanks for the swift review btw.

mlir/include/mlir/Dialect/LLVMIR/LLVMDialect.h

 #include "mlir/Dialect/LLVMIR/LLVMTypes.h"
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/Dialect.h"
 #include "mlir/IR/FunctionInterfaces.h"
 #include "mlir/IR/OpDefinition.h"
 #include "mlir/IR/OpImplementation.h"
 #include "mlir/IR/TypeSupport.h"
 #include "mlir/IR/Types.h"
+#include "mlir/Interfaces/CallInterfaces.h"
 #include "mlir/Interfaces/ControlFlowInterfaces.h"
 #include "mlir/Interfaces/InferTypeOpInterface.h"
 #include "mlir/Interfaces/SideEffectInterfaces.h"
 #include "mlir/Support/ThreadLocalCache.h"
 #include "llvm/IR/DerivedTypes.h"
 #include "llvm/IR/LLVMContext.h"
 #include "llvm/IR/Module.h"
 #include "llvm/IR/Type.h"

mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td

 
 #ifndef LLVMIR_OPS
 #define LLVMIR_OPS
 
 include "mlir/Dialect/LLVMIR/LLVMOpBase.td"
 include "mlir/Dialect/LLVMIR/LLVMOpsInterfaces.td"
 include "mlir/IR/FunctionInterfaces.td"
 include "mlir/IR/SymbolInterfaces.td"
+include "mlir/Interfaces/CallInterfaces.td"
 include "mlir/Interfaces/ControlFlowInterfaces.td"
 include "mlir/Interfaces/InferTypeOpInterface.td"
 include "mlir/Interfaces/SideEffectInterfaces.td"
 
 def FMFnnan     : I32BitEnumAttrCaseBit<"nnan", 0>;
 def FMFninf     : I32BitEnumAttrCaseBit<"ninf", 1>;
 def FMFnsz      : I32BitEnumAttrCaseBit<"nsz", 2>;
 def FMFarcp     : I32BitEnumAttrCaseBit<"arcp", 3>;

 def LLVM_FPTruncOp : LLVM_CastOp<"fptrunc", "CreateFPTrunc",
                                  LLVM_ScalarOrVectorOf<LLVM_AnyFloat>,
                                  LLVM_ScalarOrVectorOf<LLVM_AnyFloat>>;
 
 // Call-related operations.
 def LLVM_InvokeOp : LLVM_Op<"invoke", [
                       AttrSizedOperandSegments,
                       DeclareOpInterfaceMethods<BranchOpInterface>,
-                      Terminator]> {
+                      DeclareOpInterfaceMethods<CallOpInterface>, Terminator]> {
   let arguments = (ins OptionalAttr<FlatSymbolRefAttr>:$callee,
                    Variadic<LLVM_Type>:$callee_operands,
                    Variadic<LLVM_Type>:$normalDestOperands,
                    Variadic<LLVM_Type>:$unwindDestOperands);
   let results = (outs Variadic<LLVM_Type>);
   let successors = (successor AnySuccessor:$normalDest,
                               AnySuccessor:$unwindDest);
 
def LLVM_LandingpadOp : LLVM_Op<"landingpad"> {
   let arguments = (ins UnitAttr:$cleanup, Variadic<LLVM_Type>);
   let results = (outs LLVM_Type:$res);
   let builders = [LLVM_OneResultOpBuilder];
   let hasCustomAssemblyFormat = 1;
   let hasVerifier = 1;
 }
 
 def LLVM_CallOp : LLVM_Op<"call",
-                          [DeclareOpInterfaceMethods<FastmathFlagsInterface>]> {
+                          [DeclareOpInterfaceMethods<FastmathFlagsInterface>,
+                           DeclareOpInterfaceMethods<CallOpInterface>]> {
   let summary = "Call to an LLVM function.";
   let description = [{
 
 
     In LLVM IR, functions may return either 0 or 1 value. LLVM IR dialect
     implements this behavior by providing a variadic `call` operation for 0- and
     1-result functions. Even though MLIR supports multi-result functions, LLVM
     IR dialect disallows them.
let description = [{
     ```
 
   }];
   let assemblyFormat = "attr-dict";
   let hasVerifier = 1;
 }
 
 def LLVM_LLVMFuncOp : LLVM_Op<"func", [
-    AutomaticAllocationScope, IsolatedFromAbove, FunctionOpInterface, Symbol
+    AutomaticAllocationScope, IsolatedFromAbove, FunctionOpInterface,
+    CallableOpInterface, Symbol
   ]> {
   let summary = "LLVM dialect function.";
 
   let description = [{
     MLIR functions are defined by an operation that is not built into the IR
     itself. The LLVM dialect provides an `llvm.func` operation to define
     functions compatible with LLVM IR. These functions have LLVM dialect
     function type but use MLIR syntax to express it. They are required to have
let extraClassDeclaration = [{
 
     bool isVarArg() { return getFunctionType().isVarArg(); }
 
     /// Returns the argument types of this function.
     ArrayRef<Type> getArgumentTypes() { return getFunctionType().getParams(); }
 
     /// Returns the result types of this function.
     ArrayRef<Type> getResultTypes() { return getFunctionType().getReturnTypes(); }
+
+    Region *getCallableRegion() { return &getBody(); }
+
+    ArrayRef<Type> getCallableResults() { return getFunctionType().getReturnType(); }
+
   }];
 
   let hasCustomAssemblyFormat = 1;
   let hasVerifier = 1;
   let hasRegionVerifier = 1;
 }
 
 def LLVM_NullOp

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

 ///===---------------------------------------------------------------------===//
 
 SuccessorOperands InvokeOp::getSuccessorOperands(unsigned index) {
   assert(index < getNumSuccessors() && "invalid successor index");
   return SuccessorOperands(index == 0 ? getNormalDestOperandsMutable()
                                       : getUnwindDestOperandsMutable());
 }
 
+CallInterfaceCallable InvokeOp::getCallableForCallee() {
+  // Direct call.
+  if (auto calleeAttr = getCalleeAttr())
+    return calleeAttr;
+  // Indirect call, callee Value is the first operand.
+  return getOperand(0);
+}
+
+Operation::operand_range InvokeOp::getArgOperands() {
+  return getOperands().drop_front(getCallee().hasValue() ? 0 : 1);
+}
+
 LogicalResult InvokeOp::verify() {
   if (getNumResults() > 1)
     return emitOpError("must have 0 or 1 result");
 
   Block *unwindDest = getUnwindDest();
   if (unwindDest->empty())
     return emitError("must have at least one operation in unwind destination");
 
ParseResult LandingpadOp::parse(OpAsmParser &parser, OperationState &result) {
   result.addTypes(type);
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // Verifying/Printing/parsing for LLVM::CallOp.
 //===----------------------------------------------------------------------===//
 
+CallInterfaceCallable CallOp::getCallableForCallee() {
+  // Direct call.
+  if (auto calleeAttr = getCalleeAttr())
+    return calleeAttr;
+  // Indirect call, callee Value is the first operand.
+  return getOperand(0);
+}
+
+Operation::operand_range CallOp::getArgOperands() {
+  return getOperands().drop_front(getCallee().hasValue() ? 0 : 1);
+}
+
 LogicalResult CallOp::verify() {
   if (getNumResults() > 1)
     return emitOpError("must have 0 or 1 result");
 
   // Type for the callee, we'll get it differently depending if it is a direct
   // or indirect call.
   Type fnType;
 

mlir/test/Dialect/LLVMIR/callgraph.mlir

+// RUN: mlir-opt --test-print-callgraph --split-input-file %s 2>&1 | FileCheck %s
+
+// CHECK: Testing : "Normal function call"
+module attributes {"test.name" = "Normal function call"} {
+  // CHECK-LABEL: ---- CallGraph ----
+  // CHECK: - Node : 'llvm.func' {{.*}} sym_name = "foo"
+  // CHECK-NEXT: -- Call-Edge : 'llvm.func' {{.*}} sym_name = "bar"
+
+  // CHECK: - Node : 'llvm.func' {{.*}} sym_name = "bar"
+  // CHECK-NEXT: -- Call-Edge : 'llvm.func' {{.*}} sym_name = "foo"
+
+  // CHECK: - Node : 'llvm.func' {{.*}} sym_name = "entry"
+  // CHECK-DAG: -- Call-Edge : 'llvm.func' {{.*}} sym_name = "foo"
+  // CHECK-DAG: -- Call-Edge : 'llvm.func' {{.*}} sym_name = "bar"
+
+  // CHECK-LABEL: -- SCCs --
+  // CHECK: - SCC :
+  // CHECK-DAG: -- Node :'llvm.func' {{.*}} sym_name = "foo"
+  // CHECK-DAG: -- Node :'llvm.func' {{.*}} sym_name = "bar"
+
+  // CHECK: - SCC :
+  // CHECK-DAG: -- Node :'llvm.func' {{.*}} sym_name = "entry"
+
+  llvm.func @foo(%arg0: i32) -> i32 {
+    %0 = llvm.mlir.constant(2 : i32) : i32
+    %1 = llvm.sub %arg0, %0  : i32
+    %2 = llvm.call @bar(%arg0, %1) : (i32, i32) -> i32
+    llvm.return %2 : i32
+  }
+  llvm.func @bar(%arg0: i32, %arg1: i32) -> i32 {
+    %0 = llvm.add %arg0, %arg1  : i32
+    %1 = llvm.call @foo(%0) : (i32) -> i32
+    llvm.return %1 : i32
+  }
+  llvm.func @entry(%arg0: i32) -> i32 {
+    %0 = llvm.mlir.constant(2 : i32) : i32
+    %1 = llvm.mlir.constant(0 : i32) : i32
+    %2 = llvm.icmp "sgt" %arg0, %1 : i32
+    llvm.cond_br %2, ^bb1, ^bb2
+  ^bb1:  // pred: ^bb0
+    %3 = llvm.call @foo(%arg0) : (i32) -> i32
+    llvm.br ^bb3(%3 : i32)
+  ^bb2:  // pred: ^bb0
+    %4 = llvm.add %arg0, %0  : i32
+    %5 = llvm.call @bar(%arg0, %4) : (i32, i32) -> i32
+    llvm.br ^bb3(%5 : i32)
+  ^bb3(%6: i32):  // 2 preds: ^bb1, ^bb2
+    llvm.return %6 : i32
+  }
+}
+
+// -----
+
+// CHECK: Testing : "Invoke call"
+module attributes {"test.name" = "Invoke call"} {
+  // CHECK-LABEL: ---- CallGraph ----
+  // CHECK: - Node : 'llvm.func' {{.*}} sym_name = "invokeLandingpad"
+  // CHECK-DAG: -- Call-Edge : 'llvm.func' {{.*}} sym_name = "foo"
+  // CHECK-DAG: -- Call-Edge : 'llvm.func' {{.*}} sym_name = "bar"
+
+  // CHECK: -- SCCs --
+  llvm.mlir.global external constant @_ZTIi() : !llvm.ptr<i8>
+  llvm.func @foo(%arg0: i32) -> !llvm.struct<(i32, f64, i32)>
+  llvm.func @bar(!llvm.ptr<i8>, !llvm.ptr<i8>, !llvm.ptr<i8>)
+  llvm.func @__gxx_personality_v0(...) -> i32
+
+  llvm.func @invokeLandingpad() -> i32 attributes { personality = @__gxx_personality_v0 } {
+    %0 = llvm.mlir.constant(0 : i32) : i32
+    %1 = llvm.mlir.constant(3 : i32) : i32
+    %2 = llvm.mlir.constant("\01") : !llvm.array<1 x i8>
+    %3 = llvm.mlir.null : !llvm.ptr<ptr<i8>>
+    %4 = llvm.mlir.null : !llvm.ptr<i8>
+    %5 = llvm.mlir.addressof @_ZTIi : !llvm.ptr<ptr<i8>>
+    %6 = llvm.bitcast %5 : !llvm.ptr<ptr<i8>> to !llvm.ptr<i8>
+    %7 = llvm.mlir.constant(1 : i32) : i32
+    %8 = llvm.alloca %7 x i8 : (i32) -> !llvm.ptr<i8>
+    %9 = llvm.invoke @foo(%7) to ^bb2 unwind ^bb1 : (i32) -> !llvm.struct<(i32, f64, i32)>
+
+  ^bb1:
+    %10 = llvm.landingpad cleanup (catch %3 : !llvm.ptr<ptr<i8>>) (catch %6 : !llvm.ptr<i8>) (filter %2 : !llvm.array<1 x i8>) : !llvm.struct<(ptr<i8>, i32)>
+    %11 = llvm.intr.eh.typeid.for %6 : i32
+    llvm.resume %10 : !llvm.struct<(ptr<i8>, i32)>
+
+  ^bb2:
+    llvm.return %7 : i32
+
+  ^bb3:
+    llvm.invoke @bar(%8, %6, %4) to ^bb2 unwind ^bb1 : (!llvm.ptr<i8>, !llvm.ptr<i8>, !llvm.ptr<i8>) -> ()
+
+  ^bb4:
+    llvm.return %0 : i32
+  }
+}