[tfrt:jitrt] Compute arguments layout when constructing Executable

name                                      old cpu/op  new cpu/op  delta
BM_GetExecutableUniqueShapes              82.4ns ± 0%  82.6ns ± 0%   +0.29%
BM_GetExecutableSameShapes                66.8ns ± 0%  67.0ns ± 1%   +0.35%
BM_GetExecutableKnownShapes               50.6ns ± 0%  51.0ns ± 0%   +0.80%
BM_InitializeCallFrameUniqueShapes_true   70.6ns ± 2%  68.6ns ± 0%   -2.84%
BM_InitializeCallFrameSameShapes_true     70.3ns ± 0%  68.6ns ± 0%   -2.35%
BM_InitializeCallFrameKnownShapes_true    70.8ns ± 0%  69.2ns ± 3%   -2.22%
BM_InitializeCallFrameUniqueShapes_false  35.7ns ± 0%  25.1ns ± 0%  -29.67%
BM_InitializeCallFrameSameShapes_false    35.7ns ± 1%  25.1ns ± 0%  -29.74%
BM_InitializeCallFrameKnownShapes_false   35.7ns ± 1%  25.1ns ± 0%  -29.72%

PiperOrigin-RevId: 444942223

Author: ezhulenev

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

@@ -776,6 +776,7 @@ class Executable {
   using KernelFunctionPtr = void (*)(void**);
 
   // Forward declare types defined below.
+  struct ArgumentsMemoryLayout;
   struct ResultsMemoryLayout;
   struct CallFrame;
   struct ExecuteOpts;
@@ -786,6 +787,7 @@ class Executable {
              std::unique_ptr<mlir::ExecutionEngine> engine,
              KernelFunctionPtr fptr, FunctionType signature,
              FunctionType runtime_signature,
+             ArgumentsMemoryLayout arguments_memory_layout,
              ResultsMemoryLayout results_memory_layout,
              Optional<size_t> specialization,
              std::chrono::milliseconds time_to_compile)
@@ -795,6 +797,7 @@ class Executable {
         fptr_(fptr),
         signature_(std::move(signature)),
         runtime_signature_(std::move(runtime_signature)),
+        arguments_memory_layout_(std::move(arguments_memory_layout)),
         results_memory_layout_(std::move(results_memory_layout)),
         specialization_(specialization),
         time_to_compile_(time_to_compile) {
@@ -887,6 +890,12 @@ class Executable {
     llvm::StringRef error;
   };
 
+  // Requirements for passing arguments to the compiled function.
+  struct ArgumentsMemoryLayout {
+    // Currently we always pass arguments as an array of pointers.
+    size_t num_args_ptrs;
+  };
+
   // Requirements for the contiguous block of memory to store compiled function
   // results. When we invoke a compiled fuction we allocate a block of memory,
   // and pass pointers to pre-computed offsets as output arguments to the
@@ -933,10 +942,15 @@ class Executable {
     KernelContext* kernel_context;
   };
 
-  // Verifies that all types in the entrypoint function signature are supported
-  // at runtime and we know how to pass arguments and fetch results. Returns
-  // a pre-computed layout for the function results. If some of the operands
-  // or results are not supported returns an error.
+  // Verifies that all operands types in the entrypoint function signature are
+  // supported at run time . Returns a pre-computed layout for the function
+  // arguments. If some arguments are not supported returns an error.
+  static Expected<ArgumentsMemoryLayout> GetArgumentsMemoryLayout(
+      const FunctionType& signature);
+
+  // Verifies that all results types in the entrypoint function signature are
+  // supported at run time . Returns a pre-computed layout for the function
+  // results. If some results are not supported returns an error.
   static Expected<ResultsMemoryLayout> GetResultsMemoryLayout(
       const FunctionType& signature);
 
@@ -973,6 +987,7 @@ class Executable {
   // expected by the runtime.
   FunctionType runtime_signature_;
 
+  ArgumentsMemoryLayout arguments_memory_layout_;
   ResultsMemoryLayout results_memory_layout_;
 
   // Specialization id if this executable is a specialization, or an empty

backends/jitrt/lib/jitrt.cc

@@ -114,10 +114,41 @@ llvm::orc::SymbolMap RuntimeApiSymbolMap(llvm::orc::MangleAndInterner);
 }  // namespace runtime
 
 //----------------------------------------------------------------------------//
-// Get compiled function results memory layout.
+// Get compiled function arguments and results memory layouts.
 //----------------------------------------------------------------------------//
 
-Expected<Executable::ResultsMemoryLayout> Executable::GetResultsMemoryLayout(
+using ArgumentsMemoryLayout = Executable::ArgumentsMemoryLayout;
+using ResultsMemoryLayout = Executable::ResultsMemoryLayout;
+
+Expected<ArgumentsMemoryLayout> Executable::GetArgumentsMemoryLayout(
+    const FunctionType& signature) {
+  // Size of the arguments pointers array.
+  size_t num_args_ptrs = 0;
+
+  // Verify all operands types and record memory requirements.
+  for (unsigned i = 0; i < signature.num_operands(); ++i) {
+    auto* type = signature.operand(i);
+
+    // Kernel context passed as an opaque pointer.
+    if (auto* ctx = dyn_cast<KernelContextOperandType>(type)) {
+      ++num_args_ptrs;
+      continue;
+    }
+
+    // Memref passed as: 2 pointers + offset + rank * (size + stride)
+    if (auto* memref = llvm::dyn_cast<MemrefType>(type)) {
+      num_args_ptrs += 3 + 2 * memref->rank();
+      continue;
+    }
+
+    return MakeStringError("unknown operand #", i,
+                           " type memory layout: ", *type);
+  }
+
+  return ArgumentsMemoryLayout{num_args_ptrs};
+}
+
+Expected<ResultsMemoryLayout> Executable::GetResultsMemoryLayout(
     const FunctionType& signature) {
   // Size of the memory block required for storing results, and offsets for
   // each function result.
@@ -189,21 +220,6 @@ Expected<MemrefDesc> ConvertTensorToMemrefDesc(const Tensor& tensor) {
 // Executable CallFrame initialization.
 // -------------------------------------------------------------------------- //
 
-// Returns the number of call frame arguments required to pass the `memref` to
-// the compiled kernel.
-static size_t GetArgsCount(const MemrefDesc& memref) {
-  // Memref layout: 2 pointers + offset + rank * (size + stride)
-  return 3 + 2 * memref.sizes.size();
-}
-
-// Returns the number of call frame arguments required to pass all operands
-// to the compiled kernel.
-static size_t GetArgsCount(ArrayRef<MemrefDesc> operands) {
-  size_t n = 0;
-  for (const MemrefDesc& memref : operands) n += GetArgsCount(memref);
-  return n;
-}
-
 // Unpack `memref` argument into pointers to the data to be compatible with
 // compiled MLIR function ABI.
 static void AddMemrefArgument(const MemrefDesc& memref, size_t* offset,
@@ -274,8 +290,8 @@ Error Executable::InitializeCallFrame(ArrayRef<MemrefDesc> operands,
     }
   }
 
-  size_t n_args_elems = 1 + GetArgsCount(operands);
-  call_frame->args.resize_for_overwrite(n_args_elems);
+  size_t num_args_ptrs = arguments_memory_layout_.num_args_ptrs;
+  call_frame->args.resize_for_overwrite(num_args_ptrs);
 
   // Add a placeholder for the kernel context as the first argument.
   call_frame->args[0] = nullptr;
@@ -288,7 +304,7 @@ Error Executable::InitializeCallFrame(ArrayRef<MemrefDesc> operands,
   for (const MemrefDesc& desc : operands)
     AddMemrefArgument(desc, &offset, &call_frame->args);
 
-  assert(offset == n_args_elems &&
+  assert(offset == num_args_ptrs &&
          "reserved number of args must match the argument offset");
 
   // Allocate storage for results.
@@ -778,6 +794,11 @@ JitCompilationContext::Instantiate(CompilationOptions opts,
   auto runtime_signature = FunctionType::Convert(runtime_type);
   if (auto err = runtime_signature.takeError()) return std::move(err);
 
+  // Get the memory layout fo passing function arguments.
+  auto arguments_memory_layout =
+      Executable::GetArgumentsMemoryLayout(*runtime_signature);
+  if (auto err = arguments_memory_layout.takeError()) return std::move(err);
+
   // Get the memory layout for returning function results.
   auto results_memory_layout =
       Executable::GetResultsMemoryLayout(*runtime_signature);
@@ -858,7 +879,8 @@ JitCompilationContext::Instantiate(CompilationOptions opts,
   return Executable(
       ctx->name().str(), std::move(memory_mapper), std::move(*engine),
       *kernel_fn, std::move(*signature), std::move(*runtime_signature),
-      std::move(*results_memory_layout), specialization, time_to_compile);
+      std::move(*arguments_memory_layout), std::move(*results_memory_layout),
+      specialization, time_to_compile);
 }
 
 llvm::Error JitCompilationContext::Specialize(