Add option to emit stateful functions to the emitc backend.

[jacobhegna]

Summary

The usual emitc backend emits stateless cpp functions which take the
model inputs as function arguments and returns the resulting tensor.
However, it does not provide meaningful names to the arguments of the
function. Moreover, the order of the arguments is not stable in a
variety of applications (anything which comes from a tflite file, namely
MLGO models), so switching to a model where arguments are set via setter
methods avoids issues of tracking down which anonymous argument
corresponds to which model input.

Moreover, this change emits the stateful functions in two files: a
header and cpp file. This is useful for MLGO as we intend to deploy the
emitc runtime code in an anonymous namespace inside of the .cpp file of
the generated model to allow each generated model to have its own
version of the runtime (to prevent future changes to the runtime to
break previously deployed models).

[marbre]

Currently no reviewers are specified. Would you like @simon-camp and me to review? We initially pushed the code and mostly maintain it.

[jacobhegna]

Ah sorry, I got sidetracked with something else. Let me update the patch with some tests then it will be ready for review.

[marbre]

I added myself and especially @simon-camp as reviewer, since I might be oof when this is ready for review.

[jacobhegna]

Updating to include test and to update the generated code.

The delay was due to making sure the generated API was acceptable for use in
D146483.

[jacobhegna]

Ok, should be good to review.

[jacobhegna]

Including definitions of some enums.

[jacobhegna]

Bump on this review @simon-camp @marbre

[jpienaar]

mlir/include/mlir/Target/Cpp/CppEmitter.h

27

This is related only to the variable initialization right? (e.g., stateless won't be stateless if the ops being lowered mutate some global state which is not what this is controlling).

37

Could you describe these above the function? (I almost wonder if we are at point where we want a translate configuration struct for these)

mlir/lib/Target/Cpp/TranslateRegistration.cpp

41

Other descriptions here don't end with period

mlir/lib/Target/Cpp/TranslateToCpp.cpp

94

Comment?

194

Doxygen comments?

691

Prints the declaration of the stateful function class. The output of this method is something that looks like: ?

692

pseudocode

706

Could we give one concrete reason here instead of this? This feels a little bit mysterious and O(5) months down the line someone will wonder why.

711

Is this tensor specific? What happens if stateful is set and non-tensors used? (e.g., scalars)

759

pseudocode

791

API

793

This is generally useful though, I'd prefer the choices being laid out here. I don't think this is controversial, but should just be documented.

891

Nit: MLIR style is to elide trivial braces. (unrelaetd, but you can combine this into 1 conditional without loss of readability)

897

Why not switch given enum class?

1160

ml_program.identifier ?

1248

MLIR style switched to

auto tType = dyn_cast<TensorType>(type);

1250

Nit: error style is sentence fragment starting with lower case.

[simon-camp]

To me this looks like a lot of additional code that targets a very specific use case. Additionally this hardcodes the assumption that the Tensor class exposes a get method into the emitter.

Could you think of a representation that is more suitable for codegen'ing through dialect conversions? For example by lowering the model to loops we would end up with a mix of memref, arith and scf dialects. This would both be easier to model in the EmitC dialect and deal with the problem of version skew in an external reference implementation.
The issue with the anonymous function arguments would of course persist though, as well as the question about memory ownership of the arguments and results.

[jpienaar]

In D147570#4448073, @simon-camp wrote:

To me this looks like a lot of additional code that targets a very specific use case. Additionally this hardcodes the assumption that the Tensor class exposes a get method into the emitter.

Could you think of a representation that is more suitable for codegen'ing through dialect conversions? For example by lowering the model to loops we would end up with a mix of memref, arith and scf dialects. This would both be easier to model in the EmitC dialect and deal with the problem of version skew in an external reference implementation.
The issue with the anonymous function arguments would of course persist though, as well as the question about memory ownership of the arguments and results.

I agree it's too specialized, but I was thinking slightly different: the above should be able to remain as is with a transform pass on Emitc that creates classes etc. So this hardcoding would all be in a pass before translate rather than in translate.

Class construct would be needed and we'd still need a way to specify header file or body content. I think being able to not encode this into translate becomes a target then to further work here. E.g., being able to get this same functionality but as a pass with new ops ("transformToStatefulAccessors")

What do you think of the interim state where this is usable until one could do it as pass?

[simon-camp]

Can you add an error message when the module contains multiple func ops, paired with a test in test/Target/Cpp/invalid.mlir. Currently the translation succeeds but generates code that contains repeated class definitions with the same name.

If emit-cpp-only-one-fn is set and the dedicated function contains func.call ops, these are emitted as calls to free functions whose generation is skipped. This should also raise an error.

mlir/lib/Target/Cpp/TranslateToCpp.cpp

108–110

This method seems to be unused.

688

Can you call os.unindent() here so that the indentation level is not changed by the function.

783

858

This should be moved into printFuncOpBody.

1250

Please add a test for this in test/Target/Cpp/invalid.mlir.

mlir/include/mlir/Target/Cpp/CppEmitter.h

 #include "mlir/IR/Value.h"
 #include "llvm/ADT/ScopedHashTable.h"
 #include "llvm/Support/raw_ostream.h"
 #include <stack>
 
 namespace mlir {
 namespace emitc {
 
+enum class CppFileKind { Header, Cpp };
+
+enum class EmitCppKind { Stateless, Stateful };
+
 /// Translates the given operation to C++ code. The operation or operations in
 /// the region of 'op' need almost all be in EmitC dialect. The parameter
 /// 'declareVariablesAtTop' enforces that all variables for op results and block
 /// arguments are declared at the beginning of the function.
 LogicalResult translateToCpp(Operation *op, raw_ostream &os,
-                             bool declareVariablesAtTop = false);
+                             bool declareVariablesAtTop,
+                             emitc::EmitCppKind emitCppKind,
+                             CppFileKind cppFileKind, std::string argNameAttr,
+                             std::string modelName, std::string onlyOneFnName);
 } // namespace emitc
 } // namespace mlir
 
 #endif // MLIR_TARGET_CPP_CPPEMITTER_H

mlir/lib/Target/Cpp/TranslateRegistration.cpp

 //===----------------------------------------------------------------------===//
 
 void registerToCppTranslation() {
   static llvm::cl::opt<bool> declareVariablesAtTop(
       "declare-variables-at-top",
       llvm::cl::desc("Declare variables at top when emitting C/C++"),
       llvm::cl::init(false));
 
+  static llvm::cl::opt<emitc::EmitCppKind> emitCppKind(
+      "emit-cpp-kind",
+      llvm::cl::desc("Emit stateful versions of the MLIR functions"),
+      llvm::cl::init(emitc::EmitCppKind::Stateless),
+      llvm::cl::values(
+          clEnumValN(emitc::EmitCppKind::Stateless, "stateless",
+                     "Emit a stateless function."),
+          clEnumValN(emitc::EmitCppKind::Stateful, "stateful",
+                     "Emit a 'stateful function' in the form of a class")
+
+              ));
+
+  static llvm::cl::opt<std::string> argNameAttr(
+      "emit-cpp-arg-name-attr",
+      llvm::cl::desc("(Stateful only) Attribute which holds the argument names "
+                     "in the MLIR block"));
+
+  static llvm::cl::opt<std::string> modelName(
+      "emit-cpp-model-name",
+      llvm::cl::desc("(Stateful only) Name of the model. Will be exposed in a "
+                     "name() method of the class."));
+
+  static llvm::cl::opt<std::string> onlyOneFnName(
+      "emit-cpp-only-one-fn",
+      llvm::cl::desc(
+          "Only translate one function in the MLIR module to C++. This "
+          "argument is the name of that function. If empty, translate all."));
+
+  static llvm::cl::opt<emitc::CppFileKind> cppFileKind(
+      "emit-cpp-file-kind",
+      llvm::cl::desc("If emitting stateful functions, should this call emit "
+                     "the .h or .cpp file."),
+      llvm::cl::values(clEnumValN(emitc::CppFileKind::Header, "header",
+                                  "Emit the .h file in a .h/.cpp pair"),
+                       clEnumValN(emitc::CppFileKind::Cpp, "cpp",
+                                  "Emit the .cpp file in a .h/.cpp pair")));
+
   TranslateFromMLIRRegistration reg(
       "mlir-to-cpp", "translate from mlir to cpp",
       [](Operation *op, raw_ostream &output) {
         return emitc::translateToCpp(
             op, output,
-            /*declareVariablesAtTop=*/declareVariablesAtTop);
+            /*declareVariablesAtTop=*/declareVariablesAtTop, emitCppKind,
+            cppFileKind, argNameAttr, modelName, onlyOneFnName);
       },
       [](DialectRegistry &registry) {
         // clang-format off
         registry.insert<arith::ArithDialect,
                         cf::ControlFlowDialect,
                         emitc::EmitCDialect,
                         func::FuncDialect,
                         math::MathDialect,
                         scf::SCFDialect>();
         // clang-format on
       });
 }
 
 } // namespace mlir

mlir/lib/Target/Cpp/TranslateToCpp.cpp

inline LogicalResult interleaveCommaWithError(const Container &c,
                                               raw_ostream &os,
                                               UnaryFunctor eachFn) {
   return interleaveWithError(c.begin(), c.end(), eachFn, [&]() { os << ", "; });
 }
 
 namespace {
 /// Emitter that uses dialect specific emitters to emit C++ code.
 struct CppEmitter {
-  explicit CppEmitter(raw_ostream &os, bool declareVariablesAtTop);
+  explicit CppEmitter(raw_ostream &os, bool declareVariablesAtTop,
+                      emitc::EmitCppKind emitCppKind, CppFileKind cppFileKind,
+                      std::string argNameAttr, std::string modelName,
+                      std::string onlyOneFnName);
 
   /// Emits attribute or returns failure.
   LogicalResult emitAttribute(Location loc, Attribute attr);
 
   /// Emits operation 'op' with/without training semicolon or returns failure.
   LogicalResult emitOperation(Operation &op, bool trailingSemicolon);
 
+  /// Emits the type of the "underlying buffer" pointer, or returns failure
+  LogicalResult emitBufferPointerType(Location loc, Type type);
+
   /// Emits type 'type' or returns failure.
   LogicalResult emitType(Location loc, Type type);
 
   /// Emits array of types as a std::tuple of the emitted types.
   /// - emits void for an empty array;
   /// - emits the type of the only element for arrays of size one;
   /// - emits a std::tuple otherwise;
-  LogicalResult emitTypes(Location loc, ArrayRef<Type> types);
+  LogicalResult emitTypes(Location loc, ArrayRef<Type> types,
+                          bool useBufferPointerType = false);
 
   /// Emits array of types as a std::tuple of the emitted types independently of
   /// the array size.
-  LogicalResult emitTupleType(Location loc, ArrayRef<Type> types);
+  LogicalResult emitTupleType(Location loc, ArrayRef<Type> types,
+                              bool useBufferPointerType = false);
 
   /// Emits an assignment for a variable which has been declared previously.
   LogicalResult emitVariableAssignment(OpResult result);
 
   /// Emits a variable declaration for a result of an operation.
   LogicalResult emitVariableDeclaration(OpResult result,
                                         bool trailingSemicolon);
 
+  /// Emits a setter for a member variable corresponding to the variable in
+  /// result.
+  LogicalResult emitVariableSetter(OpResult result, bool trailingSemicolon);
+
   /// Emits the variable declaration and assignment prefix for 'op'.
   /// - emits separate variable followed by std::tie for multi-valued operation;
   /// - emits single type followed by variable for single result;
   /// - emits nothing if no value produced by op;
   /// Emits final '=' operator where a type is produced. Returns failure if
   /// any result type could not be converted.
   LogicalResult emitAssignPrefix(Operation &op);
 
struct CppEmitter {
 
   /// Returns the output stream.
   raw_indented_ostream &ostream() { return os; };
 
   /// Returns if all variables for op results and basic block arguments need to
   /// be declared at the beginning of a function.
   bool shouldDeclareVariablesAtTop() { return declareVariablesAtTop; };
 
+  EmitCppKind getEmitCppKind() { return emitCppKind; };
+
+  CppFileKind getCppFileKind() const { return cppFileKind; }
+
+  std::string getArgNameAttr() const { return argNameAttr; }
+
+  std::string getModelName() const { return modelName; }
+
+  std::string shouldOnlyPrintOneFn() const { return onlyOneFnName; }
+
 private:
   using ValueMapper = llvm::ScopedHashTable<Value, std::string>;
   using BlockMapper = llvm::ScopedHashTable<Block *, std::string>;
 
   /// Output stream to emit to.
   raw_indented_ostream os;
 
   /// Boolean to enforce that all variables for op results and block
   /// arguments are declared at the beginning of the function. This also
   /// includes results from ops located in nested regions.
   bool declareVariablesAtTop;
 
+  EmitCppKind emitCppKind;
+
+  CppFileKind cppFileKind;
+
+  std::string argNameAttr;
+
+  std::string modelName;
+
+  std::string onlyOneFnName;
+
   /// Map from value to name of C++ variable that contain the name.
   ValueMapper valueMapper;
 
   /// Map from block to name of C++ label.
   BlockMapper blockMapper;
 
   /// The number of values in the current scope. This is used to declare the
   /// names of values in a scope.
static LogicalResult printOperation(CppEmitter &emitter, ModuleOp moduleOp) {
 
   for (Operation &op : moduleOp) {
     if (failed(emitter.emitOperation(op, /*trailingSemicolon=*/false)))
       return failure();
   }
   return success();
 }
 
-static LogicalResult printOperation(CppEmitter &emitter,
+static LogicalResult printFuncOpBody(CppEmitter &emitter,
                                      func::FuncOp functionOp) {
-  // We need to declare variables at top if the function has multiple blocks.
-  if (!emitter.shouldDeclareVariablesAtTop() &&
-      functionOp.getBlocks().size() > 1) {
-    return functionOp.emitOpError(
-        "with multiple blocks needs variables declared at top");
-  }
-
-  CppEmitter::Scope scope(emitter);
   raw_indented_ostream &os = emitter.ostream();
-  if (failed(emitter.emitTypes(functionOp.getLoc(),
-                               functionOp.getFunctionType().getResults())))
-    return failure();
-  os << " " << functionOp.getName();
-
-  os << "(";
-  if (failed(interleaveCommaWithError(
-          functionOp.getArguments(), os,
-          [&](BlockArgument arg) -> LogicalResult {
-            if (failed(emitter.emitType(functionOp.getLoc(), arg.getType())))
-              return failure();
-            os << " " << emitter.getOrCreateName(arg);
-            return success();
-          })))
-    return failure();
-  os << ") {\n";
   os.indent();
   if (emitter.shouldDeclareVariablesAtTop()) {
     // Declare all variables that hold op results including those from nested
     // regions.
     WalkResult result =
         functionOp.walk<WalkOrder::PreOrder>([&](Operation *op) -> WalkResult {
           for (OpResult result : op->getResults()) {
             if (failed(emitter.emitVariableDeclaration(
for (Operation &op : block.getOperations()) {
       bool trailingSemicolon =
           !isa<scf::IfOp, scf::ForOp, cf::CondBranchOp>(op);
 
       if (failed(emitter.emitOperation(
               op, /*trailingSemicolon=*/trailingSemicolon)))
         return failure();
     }
   }
-  os.unindent() << "}\n";
   return success();
 }
 
-CppEmitter::CppEmitter(raw_ostream &os, bool declareVariablesAtTop)
-    : os(os), declareVariablesAtTop(declareVariablesAtTop) {
+/// The goal of this method is to print something that looks like the following
+/// psuedocode:
+///
+/// class _MyFnImpl;
+/// class MyFn {
+/// private:
+///   std::unique_ptr<_MyFnImpl> impl;
+/// public:
+///   void* get_input_buffer(std::string_view name);
+///   float* run();
+///   static std::string_view name() { return "MyFn"; }
+/// };
+///
+/// This declaration follows the pimpl design pattern. We use pimpl because
+/// internally these methods communicate via a Tensor<T, Shape...> type, but we
+/// don't want this to be a part of the public interface for various
+/// project-specific reasons.
+///
+/// Users of the generated code are intended to set inputs by acquiring a
+/// non-owning, mutable pointer to the various input buffers, and setting the
+/// raw tensor data directly. The non-owning pointer returned by run() refers
+/// to the result of the computation using the aforementioned input buffers,
+/// and the result pointer is valid until the next time run() is called.
+///
+/// Note: by design, this class is intended to be used in a single-threaded
+/// context.
+static LogicalResult printStatefulFnDecl(CppEmitter &emitter,
+                                         func::FuncOp functionOp) {
+  CppEmitter::Scope scope(emitter);
+  raw_indented_ostream &os = emitter.ostream();
+
+  // Name of the impl class
+  std::string pimpl = "_" + emitter.getModelName() + "Impl";
+
+  // Declare the impl class
+  os << "class " << pimpl << ";\n";
+
+  // Declare main class
+  os << "class " << emitter.getModelName();
+
+  os << " {\nprivate:\n  std::unique_ptr<" << pimpl << "> impl;\n";
+
+  os << "public:\n";
+  // Declare the constructor and destructor of the model. Note that we need to
+  // explicitly declare the constructor because we must allocate the memory
+  // for/construct the impl member, and we need the destructor for formal build
+  // reasons around the pimpl design pattern.
+  os << "  " << emitter.getModelName() << "();\n";
+  os << "  ~" << emitter.getModelName() << "();\n";
+
+  // Declare the get_input_buffer method
+  os << "  void* get_input_buffer(std::string_view name);";
+
+  // Declare the static name() method
+  os << "  static std::string_view name()";
+  os << " { return \"" << emitter.getModelName() << "\"; }\n  ";
+
+  // Declare run() method
+  if (failed(emitter.emitTypes(functionOp.getLoc(),
+                               functionOp.getFunctionType().getResults(),
+                               true)))
+    return failure();
+
+  os << "run();\n};\n";
+  return success();
+}
+
+/// The goal of this method is to print something that looks like the following
+/// psuedocode:
+///
+/// class _MyFnImpl {
+/// private:
+///   Tensor<float, 1> Result;
+///   Tensor<float, 1> v0;
+///   Tensor<float, 1> v1;
+///   ...
+/// public:
+///   void* get_input_buffer(std::string_name name) {
+///     if(name == "input_0") {
+///       return static_cast<void*>(v0.get());
+///     }
+///     if(name == "input_1") {
+///       return static_cast<void*>(v1.get());
+///     }
+///     assert(false && "Invalid input name!);
+///     return nullptr;
+///   }
+///   float* run() { Result = runImpl(); return Result.get(); }
+///   Tensor<float, 1> runImpl() {
+///     // insert the primary definition of the function
+///   }
+/// };
+/// void* get_input_buffer(std::string_name name) {
+///   return impl->get_input_buffer(name);
+/// }
+/// float* MyFn::run() { return impl->run(); }
+///
+/// This is an implementation of the pimpl design pattern. In the _MyFnImpl
+/// class, we need both run and runImpl because we need objects of type
+/// _MyFnImpl to own a persistent buffer containing the result of the function,
+/// so that the external api can return a non-owning pointer to the result. As
+/// with the declaration code, these choices are motivated by the needs of the
+/// github.com/google/ml-compiler-opt project.
+static LogicalResult printStatefulFnDefn(CppEmitter &emitter,
+                                         func::FuncOp functionOp) {
+  CppEmitter::Scope scope(emitter);
+  raw_indented_ostream &os = emitter.ostream();
+
+  auto args = functionOp.getArguments();
+  auto argAttrs = functionOp.getArgAttrsAttr();
+
+  // Name of the impl member
+  std::string pimpl = "_" + emitter.getModelName() + "Impl";
+
+  // Define the impl
+  os << "class " << pimpl << " {\n";
+  os << "private:\n";
+
+  // Declare the result tensor member
+  if (failed(emitter.emitTypes(functionOp.getLoc(),
+                               functionOp.getFunctionType().getResults())))
+    return failure();
+  os << " result;\n";
+
+  // Declare the tensor members for each input
+  for (int i = 0; i < functionOp.getNumArguments(); ++i) {
+    if (failed(emitter.emitType(functionOp.getLoc(), args[i].getType()))) {
+      return failure();
+    }
+    os << " " << emitter.getOrCreateName(args[i]) << ";\n";
+  }
+
+  os << "public:\n";
+
+  // Define the get_input_buffer method
+  os << "void* get_input_buffer(std::string_view name) {\n";
+  for (int i = 0; i < functionOp.getNumArguments(); ++i) {
+    auto attr = argAttrs[i].cast<DictionaryAttr>();
+    auto name = attr.get(emitter.getArgNameAttr()).cast<ArrayAttr>()[0];
+    auto varname = emitter.getOrCreateName(args[i]);
+    os << "  if (name == " << name << ") { return static_cast<void*>("
+       << varname << ".get()); }\n";
+  }
+  os << "  assert(false && \"Unknown input name!\");\n";
+  os << "  return nullptr;\n";
+  os << "}\n";
+
+  // Define the run() method
+  if (failed(emitter.emitTypes(functionOp.getLoc(),
+                               functionOp.getFunctionType().getResults(),
+                               true)))
+    return failure();
+
+  os << "run() {\n";
+  os << "   result = runImpl();\n";
+  os << "   return result.get();\n";
+  os << " }\n";
+
+  // Define the runImpl() method
+  if (failed(emitter.emitTypes(functionOp.getLoc(),
+                               functionOp.getFunctionType().getResults(),
+                               false)))
+    return failure();
+  os << " runImpl() {\n";
+  if (failed(printFuncOpBody(emitter, functionOp))) {
+    return failure();
+  }
+  os.unindent();
+  os << "}\n";
+
+  os << "};\n";
+
+  // Define the model constructor
+  os << emitter.getModelName() << "::" << emitter.getModelName()
+     << "() : impl{std::make_unique<" << pimpl << ">()} {}\n";
+  // Define the model destructor
+  os << emitter.getModelName() << "::~" << emitter.getModelName() << "() {}\n";
+
+  // Define the get_input_buffer method for the model
+  os << "void* " << emitter.getModelName()
+     << "::get_input_buffer(std::string_view name) { return "
+        "impl->get_input_buffer(name); }\n";
+
+  // Define the run() method for the model
+  if (failed(emitter.emitTypes(functionOp.getLoc(),
+                               functionOp.getFunctionType().getResults(),
+                               true)))
+    return failure();
+
+  os << " ";
+  os << emitter.getModelName() << "::run() { return impl->run(); }\n";
+
+  return success();
+}
+
+static LogicalResult printStatefulFn(CppEmitter &emitter,
+                                     func::FuncOp functionOp) {
+  // Early-exit if we are not supposed to print this function
+  auto printFnName = emitter.shouldOnlyPrintOneFn();
+  if (!printFnName.empty()) {
+    if (functionOp.getName() != printFnName) {
+      return success();
+    }
+  }
+
+  auto kind = emitter.getCppFileKind();
+  if (kind == CppFileKind::Header) {
+    return printStatefulFnDecl(emitter, functionOp);
+  }
+  if (kind == CppFileKind::Cpp) {
+    return printStatefulFnDefn(emitter, functionOp);
+  }
+  // Should be unreachable
+  return failure();
+}
+
+static LogicalResult printOperation(CppEmitter &emitter,
+                                    func::FuncOp functionOp) {
+
+  if (emitter.getEmitCppKind() == EmitCppKind::Stateful) {
+    return printStatefulFn(emitter, functionOp);
+  }
+
+  // We need to declare variables at top if the function has multiple blocks.
+  if (!emitter.shouldDeclareVariablesAtTop() &&
+      functionOp.getBlocks().size() > 1) {
+    return functionOp.emitOpError(
+        "with multiple blocks needs variables declared at top");
+  }
+
+  CppEmitter::Scope scope(emitter);
+  raw_indented_ostream &os = emitter.ostream();
+  if (failed(emitter.emitTypes(functionOp.getLoc(),
+                               functionOp.getFunctionType().getResults())))
+    return failure();
+  os << " " << functionOp.getName();
+
+  os << "(";
+  if (failed(interleaveCommaWithError(
+          functionOp.getArguments(), os,
+          [&](BlockArgument arg) -> LogicalResult {
+            if (failed(emitter.emitType(functionOp.getLoc(), arg.getType())))
+              return failure();
+            os << " " << emitter.getOrCreateName(arg);
+            return success();
+          })))
+    return failure();
+  os << ") {\n";
+
+  if (failed(printFuncOpBody(emitter, functionOp))) {
+    return failure();
+  }
+  os << "}\n";
+  return success();
+}
+
+CppEmitter::CppEmitter(raw_ostream &os, bool declareVariablesAtTop,
+                       emitc::EmitCppKind emitCppKind, CppFileKind cppFileKind,
+                       std::string argNameAttr, std::string modelName,
+                       std::string onlyOneFnName)
+    : os(os), declareVariablesAtTop(declareVariablesAtTop),
+      emitCppKind(emitCppKind), cppFileKind(cppFileKind),
+      argNameAttr(argNameAttr), modelName(modelName),
+      onlyOneFnName(onlyOneFnName) {
   valueInScopeCount.push(0);
   labelInScopeCount.push(0);
 }
 
 /// Return the existing or a new name for a Value.
 StringRef CppEmitter::getOrCreateName(Value val) {
   if (!valueMapper.count(val))
     valueMapper.insert(val, formatv("v{0}", ++valueInScopeCount.top()));
LogicalResult CppEmitter::emitVariableDeclaration(OpResult result,
   if (failed(emitType(result.getOwner()->getLoc(), result.getType())))
     return failure();
   os << " " << getOrCreateName(result);
   if (trailingSemicolon)
     os << ";\n";
   return success();
 }
 
+LogicalResult CppEmitter::emitVariableSetter(OpResult result,
+                                             bool trailingSemicolon) {
+  std::string name = "EMPTY_NAME";
+  if (result.getOwner()->getAttrDictionary().contains("iree.identifier")) {
+    name = result.getOwner()
+               ->getAttr(argNameAttr)
+               .cast<ArrayAttr>()[0]
+               .cast<StringAttr>()
+               .str();
+  }
+
+  os << "void set_" << name << "(";
+  if (failed(emitType(result.getOwner()->getLoc(), result.getType())))
+    return failure();
+  os << " x) {\n" << name << " = x;\n}\n";
+
+  return success();
+}
+
 LogicalResult CppEmitter::emitAssignPrefix(Operation &op) {
   switch (op.getNumResults()) {
   case 0:
     break;
   case 1: {
     OpResult result = op.getResult(0);
     if (shouldDeclareVariablesAtTop()) {
       if (failed(emitVariableAssignment(result)))
LogicalResult status =
           });
 
   if (failed(status))
     return failure();
   os << (trailingSemicolon ? ";\n" : "\n");
   return success();
 }
 
+LogicalResult CppEmitter::emitBufferPointerType(Location loc, Type type) {
+  auto tType = type.dyn_cast<TensorType>();
+  if (!tType) {
+    return emitError(loc, "Can only emit buffer pointer type for Tensors");
+  }
+  if (failed(emitType(loc, tType.getElementType())))
+    return failure();
+  os << "* ";
+  return success();
+}
+
 LogicalResult CppEmitter::emitType(Location loc, Type type) {
   if (auto iType = type.dyn_cast<IntegerType>()) {
     switch (iType.getWidth()) {
     case 1:
       return (os << "bool"), success();
     case 8:
     case 16:
     case 32:
if (auto pType = type.dyn_cast<emitc::PointerType>()) {
     if (failed(emitType(loc, pType.getPointee())))
       return failure();
     os << "*";
     return success();
   }
   return emitError(loc, "cannot emit type ") << type;
 }
 
-LogicalResult CppEmitter::emitTypes(Location loc, ArrayRef<Type> types) {
+LogicalResult CppEmitter::emitTypes(Location loc, ArrayRef<Type> types,
+                                    bool useBufferPointerType) {
   switch (types.size()) {
   case 0:
     os << "void";
     return success();
   case 1:
+    if (useBufferPointerType) {
+      return emitBufferPointerType(loc, types.front());
+    } else {
       return emitType(loc, types.front());
+    }
   default:
-    return emitTupleType(loc, types);
+    return emitTupleType(loc, types, useBufferPointerType);
   }
 }
 
-LogicalResult CppEmitter::emitTupleType(Location loc, ArrayRef<Type> types) {
+LogicalResult CppEmitter::emitTupleType(Location loc, ArrayRef<Type> types,
+                                        bool useBufferPointerType) {
   os << "std::tuple<";
-  if (failed(interleaveCommaWithError(
-          types, os, [&](Type type) { return emitType(loc, type); })))
+  if (failed(interleaveCommaWithError(types, os, [&](Type type) {
+        if (useBufferPointerType) {
+          return emitBufferPointerType(loc, type);
+        } else {
+          return emitType(loc, type);
+        }
+      })))
     return failure();
   os << ">";
   return success();
 }
 
-LogicalResult emitc::translateToCpp(Operation *op, raw_ostream &os,
-                                    bool declareVariablesAtTop) {
-  CppEmitter emitter(os, declareVariablesAtTop);
+LogicalResult emitc::translateToCpp(
+    Operation *op, raw_ostream &os, bool declareVariablesAtTop,
+    emitc::EmitCppKind emitCppKind, CppFileKind cppFileKind,
+    std::string argNameAttr, std::string modelName, std::string onlyOneFnName) {
+  CppEmitter emitter(os, declareVariablesAtTop, emitCppKind, cppFileKind,
+                     argNameAttr, modelName, onlyOneFnName);
   return emitter.emitOperation(*op, /*trailingSemicolon=*/false);
 }

mlir/test/Target/Cpp/emit_stateful_fn.mlir

+// RUN: mlir-translate -mlir-to-cpp --emit-cpp-kind=stateful --emit-cpp-arg-name-attr=tf_saved_model.index_path --emit-cpp-model-name=test --emit-cpp-file-kind=header %s | FileCheck %s -check-prefix=HEADER
+// RUN: mlir-translate -mlir-to-cpp --emit-cpp-kind=stateful --emit-cpp-arg-name-attr=tf_saved_model.index_path --emit-cpp-model-name=test --emit-cpp-file-kind=cpp %s | FileCheck %s -check-prefix=CPP
+
+func.func @test(%arg0: tensor<1xf32> {tf_saved_model.index_path=["a"]},
+                %arg1: tensor<1xf32> {tf_saved_model.index_path=["b"]})
+               -> (tensor<1xf32>) {
+  %0 = emitc.call "tosa::add"(%arg0, %arg1) {args = [0 : index, 1 : index]} : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xf32>
+  return %0 : tensor<1xf32>
+}
+
+// HEADER: class _testImpl;
+// HEADER: class test {
+// HEADER: private:
+// HEADER: std::unique_ptr<_testImpl> impl;
+// HEADER: public:
+// HEADER: void* get_input_buffer(std::string_view name);
+// HEADER: float* run();
+// HEADER: };
+
+// CPP: class _testImpl {
+// CPP: private:
+// CPP: Tensor<float, 1> result;
+// CPP: Tensor<float, 1> v1;
+// CPP: Tensor<float, 1> v2;
+// CPP: public:
+// CPP: void* get_input_buffer(std::string_view name) {
+// CPP:   if (name == "a") { return static_cast<void*>(v1.get()); }
+// CPP:   if (name == "b") { return static_cast<void*>(v2.get()); }
+// CPP:   assert(false && "Unknown input name!");
+// CPP:   return nullptr;
+// CPP: }
+// CPP: float* run() {
+// CPP:   result = runImpl();
+// CPP:   return result.get();
+// CPP: }
+// CPP: Tensor<float, 1> runImpl() {
+// CPP:   Tensor<float, 1> v3 = tosa::add(v1, v2);
+// CPP:   return v3;
+// CPP: }
+// CPP: };
+// CPP: test::test() : impl{std::make_unique<_testImpl>()} {}
+// CPP: test::~test() {}
+// CPP: void* test::get_input_buffer(std::string_view name) { return impl->get_input_buffer(name); }
+// CPP: float* test::run() { return impl->run(); }