[mlir] Implement SymbolUserOpInterface in LLVM::CallOp

[ezhulenev]

Summary

Avoid expensive calls to SymbolTable::lookupNearestSymbolFrom in verifier

[ezhulenev]

Rebased

[ezhulenev]

Rebase

[ezhulenev]

Rebased

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

def LLVM_LandingpadOp : LLVM_Op<"landingpad"> {
   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<CallOpInterface>]> {
+                           DeclareOpInterfaceMethods<CallOpInterface>,
+                           DeclareOpInterfaceMethods<SymbolUserOpInterface>]> {
   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.
OpBuilder<(ins "TypeRange":$results, "StringAttr":$callee,
       build($_builder, $_state, results, SymbolRefAttr::get(callee), operands);
     }]>,
     OpBuilder<(ins "TypeRange":$results, "StringRef":$callee,
                    CArg<"ValueRange", "{}">:$operands), [{
       build($_builder, $_state, results,
             StringAttr::get($_builder.getContext(), callee), operands);
     }]>];
   let hasCustomAssemblyFormat = 1;
-  let hasVerifier = 1;
 }
+
 def LLVM_ExtractElementOp : LLVM_Op<"extractelement", [NoSideEffect]> {
   let arguments = (ins LLVM_AnyVector:$vector, AnyInteger:$position);
   let results = (outs LLVM_Type:$res);
   string llvmBuilder = [{
     $res = builder.CreateExtractElement($vector, $position);
   }];
   let builders = [
     OpBuilder<(ins "Value":$vector, "Value":$position,

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

CallInterfaceCallable CallOp::getCallableForCallee() {
   // Indirect call, callee Value is the first operand.
   return getOperand(0);
 }
 
 Operation::operand_range CallOp::getArgOperands() {
   return getOperands().drop_front(getCallee().has_value() ? 0 : 1);
 }
 
-LogicalResult CallOp::verify() {
+LogicalResult CallOp::verifySymbolUses(SymbolTableCollection &symbolTable) {
   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;
 
   bool isIndirect = false;
 
   // If this is an indirect call, the callee attribute is missing.
   FlatSymbolRefAttr calleeName = getCalleeAttr();
   if (!calleeName) {
     isIndirect = true;
     if (!getNumOperands())
       return emitOpError(
           "must have either a `callee` attribute or at least an operand");
     auto ptrType = getOperand(0).getType().dyn_cast<LLVMPointerType>();
     if (!ptrType)
       return emitOpError("indirect call expects a pointer as callee: ")
              << ptrType;
     fnType = ptrType.getElementType();
   } else {
     Operation *callee =
-        SymbolTable::lookupNearestSymbolFrom(*this, calleeName.getAttr());
+        symbolTable.lookupNearestSymbolFrom(*this, calleeName.getAttr());
     if (!callee)
       return emitOpError()
              << "'" << calleeName.getValue()
              << "' does not reference a symbol in the current scope";
     auto fn = dyn_cast<LLVMFuncOp>(callee);
     if (!fn)
       return emitOpError() << "'" << calleeName.getValue()
                            << "' does not reference a valid LLVM function";

mlir/test/Dialect/LLVMIR/invalid.mlir

func.func @store_malformed_elem_type(%foo: !llvm.ptr, %bar: f32) {
   llvm.store %bar, %foo : !llvm.ptr, "f32"
 }
 
 // -----
 
 func.func @call_non_function_type(%callee : !llvm.func<i8 (i8)>, %arg : i8) {
   // expected-error@+1 {{expected function type}}
   llvm.call %callee(%arg) : !llvm.func<i8 (i8)>
+  llvm.return
 }
 
 // -----
 
 func.func @invalid_call() {
   // expected-error@+1 {{'llvm.call' op must have either a `callee` attribute or at least an operand}}
   "llvm.call"() : () -> ()
+  llvm.return
 }
 
 // -----
 
 func.func @call_non_function_type(%callee : !llvm.func<i8 (i8)>, %arg : i8) {
   // expected-error@+1 {{expected function type}}
   llvm.call %callee(%arg) : !llvm.func<i8 (i8)>
+  llvm.return
 }
 
 // -----
 
 func.func @call_unknown_symbol() {
   // expected-error@+1 {{'llvm.call' op 'missing_callee' does not reference a symbol in the current scope}}
   llvm.call @missing_callee() : () -> ()
+  llvm.return
 }
 
 // -----
 
 func.func private @standard_func_callee()
 
 func.func @call_non_llvm() {
   // expected-error@+1 {{'llvm.call' op 'standard_func_callee' does not reference a valid LLVM function}}
   llvm.call @standard_func_callee() : () -> ()
+  llvm.return
 }
 
 // -----
 
 func.func @call_non_llvm_indirect(%arg0 : tensor<*xi32>) {
   // expected-error@+1 {{'llvm.call' op operand #0 must be LLVM dialect-compatible type}}
   "llvm.call"(%arg0) : (tensor<*xi32>) -> ()
+  llvm.return
 }
 
 // -----
 
 llvm.func @callee_func(i8) -> ()
 
 func.func @callee_arg_mismatch(%arg0 : i32) {
   // expected-error@+1 {{'llvm.call' op operand type mismatch for operand 0: 'i32' != 'i8'}}
   llvm.call @callee_func(%arg0) : (i32) -> ()
+  llvm.return
 }
 
 // -----
 
 func.func @indirect_callee_arg_mismatch(%arg0 : i32, %callee : !llvm.ptr<func<void(i8)>>) {
   // expected-error@+1 {{'llvm.call' op operand type mismatch for operand 0: 'i32' != 'i8'}}
   "llvm.call"(%callee, %arg0) : (!llvm.ptr<func<void(i8)>>, i32) -> ()
+  llvm.return
 }
 
 // -----
 
 llvm.func @callee_func() -> (i8)
 
 func.func @callee_return_mismatch() {
   // expected-error@+1 {{'llvm.call' op result type mismatch: 'i32' != 'i8'}}
   %res = llvm.call @callee_func() : () -> (i32)
+  llvm.return
 }
 
 // -----
 
 func.func @indirect_callee_return_mismatch(%callee : !llvm.ptr<func<i8()>>) {
   // expected-error@+1 {{'llvm.call' op result type mismatch: 'i32' != 'i8'}}
   "llvm.call"(%callee) : (!llvm.ptr<func<i8()>>) -> (i32)
+  llvm.return
 }
 
 // -----
 
 func.func @call_too_many_results(%callee : () -> (i32,i32)) {
   // expected-error@+1 {{expected function with 0 or 1 result}}
   llvm.call %callee() : () -> (i32, i32)
+  llvm.return
 }
 
 // -----
 
 func.func @call_non_llvm_result(%callee : () -> (tensor<*xi32>)) {
   // expected-error@+1 {{expected result to have LLVM type}}
   llvm.call %callee() : () -> (tensor<*xi32>)
+  llvm.return
 }
 
 // -----
 
 func.func @call_non_llvm_input(%callee : (tensor<*xi32>) -> (), %arg : tensor<*xi32>) {
   // expected-error@+1 {{expected LLVM types as inputs}}
   llvm.call %callee(%arg) : (tensor<*xi32>) -> ()
+  llvm.return
 }
 
 // -----
 
 llvm.func @void_func_result(%arg0: i32) {
   // expected-error@below {{expected no operands}}
   // expected-note@above {{when returning from function}}
   llvm.return %arg0: i32