[tfrt:jitrt] Add a flag to disable operands verification when executing JitRt program

PiperOrigin-RevId: 444937420

Author: ezhulenev

backends/jitrt/cpp_tests/jit_executable_test.cc

@@ -15,6 +15,7 @@
  */
 
 #include <memory>
+#include <utility>
 
 #include "benchmark/benchmark.h"
 #include "gtest/gtest.h"
@@ -47,9 +48,9 @@ using SymbolicShape = SymbolicShapesResolver::SymbolicShape;
 
 static const char* mlir_module = R"(
     func.func @compute(%arg0: memref<?x?xf32>,
-                  %arg1: memref<?x?xf32>,
-                  %arg3: memref<?x?xf32>,
-                  %arg4: memref<16x32xf32>) {
+                       %arg1: memref<?x?xf32>,
+                       %arg3: memref<?x?xf32>,
+                       %arg4: memref<16x32xf32>) {
       func.return
     })";
 
@@ -72,6 +73,7 @@ SmallVector<MemrefDesc> GetFakeMemrefs(SmallVector<SymbolicShape> shapes) {
     MemrefDesc desc;
     desc.dtype = DType::F32;
     desc.sizes.insert(desc.sizes.begin(), shape.begin(), shape.end());
+    desc.strides.append(shape.size(), 0);  // we don't need real strides
     memrefs.push_back(std::move(desc));
   }
 
@@ -111,6 +113,43 @@ void BenchmarkGetExecutable(benchmark::State& state,
   }
 }
 
+void BenchmarkInitializeCallFrame(benchmark::State& state,
+                                  SmallVector<MemrefDesc> operands,
+                                  bool verify) {
+  auto host = CreateSingleThreadedHostContext();
+
+  CompilationOptions opts;
+  opts.specialization = CompilationOptions::Specialization::kAlways;
+  opts.register_dialects = RegisterDefaultJitRtDialects;
+
+  CompilationPipelineOptions copts;
+  opts.create_compilation_pipeline = [copts](mlir::PassManager& pm) {
+    CreateDefaultJitRtCompilationPipeline(pm, copts);
+  };
+
+  llvm::Expected<JitExecutable> jit_executable =
+      JitExecutable::Instantiate(mlir_module, entrypoint, opts);
+  if (auto err = jit_executable.takeError()) TFRT_LOG(FATAL) << err;
+
+  // Get the executable.
+  Expected<AsyncValuePtr<Executable>> executable =
+      jit_executable->GetExecutable(operands);
+  if (auto err = executable.takeError()) TFRT_LOG(FATAL) << err;
+
+  // Check that compilation was successful.
+  host->Quiesce();
+  if (executable->IsError()) TFRT_LOG(FATAL) << executable->GetError();
+
+  for (auto _ : state) {
+    Executable::CallFrame call_frame;
+    auto err =
+        (*executable)->InitializeCallFrame(operands, &call_frame, verify);
+    benchmark::DoNotOptimize(call_frame);
+  }
+}
+
+// -------------------------------------------------------------------------- //
+
 #define BM_GetExecutable(NAME, OPERANDS)                        \
   static void BM_GetExecutable##NAME(benchmark::State& state) { \
     BenchmarkGetExecutable(state, OPERANDS);                    \
@@ -126,5 +165,38 @@ BM_GetExecutable(SameShapes,
 BM_GetExecutable(KnownShapes,
                  GetFakeMemrefs({{16, 32}, {16, 32}, {16, 32}, {16, 32}}));
 
+// -------------------------------------------------------------------------- //
+
+#define BM_InitializeCallFrame(NAME, OPERANDS, VERIFY)     \
+  static void BM_InitializeCallFrame##NAME##_##VERIFY(     \
+      benchmark::State& state) {                           \
+    BenchmarkInitializeCallFrame(state, OPERANDS, VERIFY); \
+  }                                                        \
+  BENCHMARK(BM_InitializeCallFrame##NAME##_##VERIFY)
+
+BM_InitializeCallFrame(UniqueShapes,
+                       GetFakeMemrefs({{10, 11}, {12, 13}, {14, 15}, {16, 32}}),
+                       true);
+
+BM_InitializeCallFrame(SameShapes,
+                       GetFakeMemrefs({{10, 11}, {10, 11}, {10, 11}, {16, 32}}),
+                       true);
+
+BM_InitializeCallFrame(KnownShapes,
+                       GetFakeMemrefs({{16, 32}, {16, 32}, {16, 32}, {16, 32}}),
+                       true);
+
+BM_InitializeCallFrame(UniqueShapes,
+                       GetFakeMemrefs({{10, 11}, {12, 13}, {14, 15}, {16, 32}}),
+                       false);
+
+BM_InitializeCallFrame(SameShapes,
+                       GetFakeMemrefs({{10, 11}, {10, 11}, {10, 11}, {16, 32}}),
+                       false);
+
+BM_InitializeCallFrame(KnownShapes,
+                       GetFakeMemrefs({{16, 32}, {16, 32}, {16, 32}, {16, 32}}),
+                       false);
+
 }  // namespace
 }  // namespace tfrt

backends/jitrt/include/tfrt/jitrt/jitrt.h

@@ -802,16 +802,22 @@ class Executable {
   }
 
   // Initializes call frame by adding all operands as pointers to the arguments
-  // vector. Also allocates storage for the returned values. Return values
-  // storage requirements inferred from the kernel function signature.
+  // vector. Also allocates storage for the returned values.
+  //
+  // If `verify_operands` is true (in debug mode it's always on, independent of
+  // the argument value) this function also verifies that operands passed at run
+  // time matches the executable entrypoint signature (e.g. all statically known
+  // dimensions of the memrefs matches the operands). Returns an error if finds
+  // a mismatch.
   //
   // This function leaves the kernel context argument (the first argument of a
   // kernel function) uninitialized. It will be initialized in the `Execute`
   // function right before the actual execution.
   //
   // See mlir::ExecutionEngine `packFunctionArguments` for the details.
   Error InitializeCallFrame(ArrayRef<MemrefDesc> operands,
-                            CallFrame* call_frame) const;
+                            CallFrame* call_frame,
+                            bool verify_operands = true) const;
 
   // Converts returned values owned by the call frame using provided value
   // converter. If result conversion fails (e.g. result type is not supported)
@@ -822,16 +828,20 @@ class Executable {
   Error ReturnResults(const ReturnValueConverterBase& results,
                       CallFrame* call_frame) const;
 
-  // Executes compiled function with given operands. If operands passed at
-  // runtime are not compatible with the compiled function signature, allocates
-  // error async values for all results.
+  // Executes compiled function with given operands.
+  //
+  // If `verify_operands` is true (in debug mode it's always on, independent of
+  // the argument value) this function also verifies that operands passed at run
+  // time matches the executable entrypoint signature. If some of the operands
+  // do not match the expected type, this function allocates error async values
+  // for all results and returns an error.
   //
   // Returns compiled function results via the user-provided results converter.
   // If compiled function execution completed in the error state, emits error
   // async value for all results.
   Error Execute(ArrayRef<MemrefDesc> operands,
                 const ReturnValueConverterBase& results,
-                const ExecuteOpts& opts) const;
+                const ExecuteOpts& opts, bool verify_operands = true) const;
 
   // Executes compiled function using user provided call frame.
   //

backends/jitrt/lib/jitrt.cc

@@ -224,41 +224,53 @@ static void AddMemrefArgument(const MemrefDesc& memref, size_t* offset,
   for (const Index& stride : memref.strides) add_arg(&stride);
 }
 
+// Always verify executable operands in debug mode.
+static bool VerifyOperands(bool verify_operands) {
+#if defined(NDEBUG)
+  return verify_operands;
+#endif
+  return true;
+}
+
 Error Executable::InitializeCallFrame(ArrayRef<MemrefDesc> operands,
-                                      CallFrame* call_frame) const {
+                                      CallFrame* call_frame,
+                                      bool verify_operands) const {
   // TODO(ezhulenev): If executable is specialized for operands shapes then
   // there is no need to verify them once more here. However currently we rely
   // on a hash code to look up specializations, and this can lead to collisions.
+  if (VerifyOperands(verify_operands)) {
+    // We verify run time operands against the run time signature.
+    const FunctionType& signature = runtime_signature_;
+
+    // Make sure that we call the kernel with the correct number of operands.
+    // We subtract one operand from the signature because it corresponds to the
+    // context that we prepend to the given operands.
+    if (LLVM_UNLIKELY(operands.size() != signature.num_operands() - 1))
+      return MakeStringError(
+          "number of operands doesn't match the function signature: ",
+          operands.size(), " vs ", signature.num_operands() - 1);
+
+    // Verify that all operands passed at runtime are compatible with compiled
+    // function signature.
+    auto kctx = dyn_cast<KernelContextOperandType>(signature.operand(0));
+    if (LLVM_UNLIKELY(!kctx)) {
+      return MakeStringError(
+          "expected KernelContext in first argument of signature, got: ",
+          signature.operand(0));
+    }
 
-  // Make sure that we call the kernel with the correct number of operands.
-  // We subtract one operand from the signature because it corresponds to the
-  // context that we prepend to the given operands.
-  if (LLVM_UNLIKELY(operands.size() != runtime_signature_.num_operands() - 1))
-    return MakeStringError(
-        "number of operands doesn't match the function signature: ",
-        operands.size(), " vs ", runtime_signature_.num_operands() - 1);
-
-  // Verify that all operands passed at runtime are compatible with compiled
-  // function signature.
-  auto kctx = dyn_cast<KernelContextOperandType>(runtime_signature_.operand(0));
-  if (LLVM_UNLIKELY(!kctx)) {
-    return MakeStringError(
-        "expected KernelContext in first argument of "
-        "signature, got: ",
-        runtime_signature_.operand(0));
-  }
-
-  // We use 0-based index for operands, because the kernel context operand is an
-  // internal implementation detail, and in case of an error users should get
-  // back operand index corresponding to the user provided signature.
-  for (unsigned i = 0; i < operands.size(); ++i) {
-    unsigned idx = i + 1;  // use 1-based index to fetch runtime operand
+    // We use 0-based index for operands, because the kernel context operand is
+    // an internal implementation detail, and in case of an error users should
+    // get back operand index corresponding to the user provided signature.
+    for (unsigned i = 0; i < operands.size(); ++i) {
+      unsigned idx = i + 1;  // use 1-based index to fetch runtime operand
 
-    if (auto* memref = dyn_cast<MemrefType>(runtime_signature_.operand(idx))) {
-      if (auto err = VerifyMemrefOperand(i, *memref, operands[i])) return err;
-    } else {
-      return MakeStringError("expected memref operand at #", i,
-                             ", got: ", *runtime_signature_.operand(i));
+      if (auto* memref = dyn_cast<MemrefType>(signature.operand(idx))) {
+        if (auto err = VerifyMemrefOperand(i, *memref, operands[i])) return err;
+      } else {
+        return MakeStringError("expected memref operand at #", i,
+                               ", got: ", *signature.operand(i));
+      }
     }
   }
 
@@ -415,7 +427,7 @@ Error ReturnErrors(const ReturnValueConverterBase& results, Error error) {
 
 Error Executable::Execute(ArrayRef<MemrefDesc> operands,
                           const ReturnValueConverterBase& results,
-                          const ExecuteOpts& opts) const {
+                          const ExecuteOpts& opts, bool verify_operands) const {
   // CallFrame can be allocated on the stack because compiled function will
   // unpack all the arguments it needs, and async regions will not access
   // the data after the initial function will return the result.
@@ -445,7 +457,7 @@ Error Executable::Execute(ArrayRef<MemrefDesc> operands,
 
   // Compiled function takes arguments and results as `void**` type erased
   // pointer. See mlir::ExecutionEngine `packFunctionArguments` for the details.
-  if (auto err = InitializeCallFrame(operands, &call_frame))
+  if (auto err = InitializeCallFrame(operands, &call_frame, verify_operands))
     return ReturnErrors(results, std::move(err));
 
   Execute(call_frame, opts);