Compiler is now complete

include/tl2cgen/detail/compiler/codegen/codegen.h

@@ -7,16 +7,27 @@
 #ifndef TL2CGEN_DETAIL_COMPILER_CODEGEN_CODEGEN_H_
 #define TL2CGEN_DETAIL_COMPILER_CODEGEN_CODEGEN_H_
 
+#include <filesystem>
 #include <map>
 #include <ostream>
 #include <string>
 #include <vector>
 
-namespace tl2cgen::compiler::detail {
+/* Forward declarations */
+namespace treelite {
 
-namespace ast {
+class Model;
+
+}  // namespace treelite
+
+namespace tl2cgen::compiler {
+
+struct CompilerParam;
+
+}  // namespace tl2cgen::compiler
+
+namespace tl2cgen::compiler::detail::ast {
 
-// Forward declarations
 class ASTNode;
 class MainNode;
 class FunctionNode;
@@ -26,13 +37,16 @@ class TranslationUnitNode;
 class QuantizerNode;
 class ModelMeta;
 
-}  // namespace ast
+}  // namespace tl2cgen::compiler::detail::ast
 
-namespace codegen {
+namespace tl2cgen::compiler::detail::codegen {
 
 class CodeCollection;  // forward declaration
 
 void GenerateCodeFromAST(ast::ASTNode const* node, CodeCollection& gencode);
+void WriteCodeToDisk(std::filesystem::path const& dirpath, CodeCollection const& collection);
+void WriteBuildRecipeToDisk(std::filesystem::path const& dirpath,
+    std::string const& native_lib_name, CodeCollection const& collection);
 
 // Codegen implementation for each AST node type
 void HandleMainNode(ast::MainNode const* node, CodeCollection& gencode);
@@ -70,6 +84,9 @@ class SourceFile {
   void ChangeIndent(int n_tabs_delta);  // Add or remove indent
   void PushFragment(std::string content);
   friend std::ostream& operator<<(std::ostream&, CodeCollection const&);
+  friend void WriteCodeToDisk(std::filesystem::path const& dirpath, CodeCollection const&);
+  friend void WriteBuildRecipeToDisk(
+      std::filesystem::path const&, std::string const&, CodeCollection const&);
   friend class CodeCollection;
 };
 
@@ -88,11 +105,12 @@ class CodeCollection {
   void PushFragment(std::string content);
 
   friend std::ostream& operator<<(std::ostream&, CodeCollection const&);
+  friend void WriteCodeToDisk(std::filesystem::path const&, CodeCollection const&);
+  friend void WriteBuildRecipeToDisk(
+      std::filesystem::path const&, std::string const&, CodeCollection const&);
 };
 std::ostream& operator<<(std::ostream& os, CodeCollection const& collection);
 
-}  // namespace codegen
-
-}  // namespace tl2cgen::compiler::detail
+}  // namespace tl2cgen::compiler::detail::codegen
 
 #endif  // TL2CGEN_DETAIL_COMPILER_CODEGEN_CODEGEN_H_

include/tl2cgen/detail/filesystem.h

@@ -19,21 +19,6 @@ namespace tl2cgen::detail::filesystem {
  */
 void CreateDirectoryIfNotExist(std::filesystem::path const& dirpath);
 
-/*!
- * \brief Write a sequence of strings to a text file, with newline character (\n) inserted between
- *        strings. This function is suitable for creating multi-line text files.
- * \param path Path to text file
- * \param content A sequence of strings to be written.
- */
-void WriteToFile(std::filesystem::path const& path, std::string const& content);
-
-/*!
- * \brief Write a sequence of bytes to a text file
- * \param path Path to text file
- * \param content A sequence of bytes to be written.
- */
-void WriteToFile(std::filesystem::path const& path, std::vector<char> const& content);
-
 }  // namespace tl2cgen::detail::filesystem
 
 #endif  // TL2CGEN_DETAIL_FILESYSTEM_H_

src/c_api/c_api.cc

@@ -87,7 +87,6 @@ int TL2cgenGenerateCode(
 
   std::filesystem::path dirpath_
       = std::filesystem::weakly_canonical(std::filesystem::u8path(std::string(dirpath)));
-  detail::filesystem::CreateDirectoryIfNotExist(dirpath_);
 
   /* Compile model */
   auto param = compiler::CompilerParam::ParseFromJSON(compiler_params_json_str);

src/compiler/codegen/codegen.cc

@@ -5,11 +5,15 @@
  * \author Hyunsu Cho
  */
 
+#include <rapidjson/ostreamwrapper.h>
+#include <rapidjson/prettywriter.h>
 #include <tl2cgen/detail/compiler/ast/ast.h>
 #include <tl2cgen/detail/compiler/codegen/codegen.h>
 #include <tl2cgen/detail/compiler/codegen/format_util.h>
+#include <tl2cgen/detail/filesystem.h>
 #include <tl2cgen/logging.h>
 
+#include <fstream>
 #include <string>
 #include <type_traits>
 #include <variant>
@@ -56,6 +60,49 @@ void GenerateCodeFromAST(ast::ASTNode const* node, CodeCollection& gencode) {
   }
 }
 
+void WriteCodeToDisk(std::filesystem::path const& dirpath, CodeCollection const& collection) {
+  namespace fs = tl2cgen::detail::filesystem;
+  for (auto const& [file_name, source_file] : collection.sources_) {
+    std::ofstream of(dirpath / file_name);
+    for (auto const& fragment : source_file.fragments_) {
+      of << IndentMultiLineString(fragment.content_, fragment.indent_) << "\n";
+    }
+    of << "\n";
+  }
+}
+
+void WriteBuildRecipeToDisk(std::filesystem::path const& dirpath,
+    std::string const& native_lib_name, CodeCollection const& collection) {
+  std::ofstream ofs(dirpath / "recipe.json");
+  rapidjson::OStreamWrapper ofs_wrapped(ofs);
+  rapidjson::PrettyWriter<rapidjson::OStreamWrapper> writer(ofs_wrapped);
+  writer.SetFormatOptions(rapidjson::PrettyFormatOptions::kFormatSingleLineArray);
+
+  writer.StartObject();
+  writer.Key("target");
+  writer.String(native_lib_name);
+  writer.Key("sources");
+  writer.StartArray();
+  for (auto const& [file_name, source_file] : collection.sources_) {
+    if (file_name.compare(file_name.length() - 2, 2, ".c") == 0) {
+      std::size_t line_count = 0;
+      for (auto const& fragment : source_file.fragments_) {
+        line_count += std::count(fragment.content_.begin(), fragment.content_.end(), '\n');
+      }
+      writer.StartObject();
+      writer.Key("name");
+      std::string name = file_name.substr(0, file_name.length() - 2);
+      writer.String(name);
+      writer.Key("length");
+      writer.Uint64(line_count);
+      writer.EndObject();
+    }
+  }
+  writer.EndArray();
+  writer.EndObject();
+  ofs << "\n";  // Add newline at the end, for convention's sake
+}
+
 std::string GetThresholdCType(ast::ASTNode const* node) {
   return GetThresholdCType(*node->meta_);
 }

src/compiler/codegen/postprocessor.cc

@@ -192,10 +192,10 @@ static void postprocess_impl({threshold_type}* target_result, int num_class) {{
   for (int k = 0; k < num_class; ++k) {{
     t = {exp}(target_result[k] - max_margin);
     norm_const += t;
-    pred[k] = t;
+    target_result[k] = t;
   }}
   for (int k = 0; k < num_class; ++k) {{
-    pred[k] /= ({threshold_type})norm_const;
+    target_result[k] /= ({threshold_type})norm_const;
   }}
 }}
 
@@ -207,7 +207,7 @@ void postprocess({threshold_type}* result) {{
   auto const max_num_class
       = *std::max_element(model_meta.num_class_.begin(), model_meta.num_class_.end());
   for (std::int32_t target_id = 0; target_id < model_meta.num_target_; ++target_id) {
-    fmt::print(oss, "  postprocess_impl(result[{offset}], {num_class});\n",
+    fmt::print(oss, "  postprocess_impl(&result[{offset}], {num_class});\n",
         "offset"_a = target_id * max_num_class, "num_class"_a = model_meta.num_class_[target_id]);
   }
   oss << "}\n";
@@ -239,7 +239,7 @@ void postprocess({threshold_type}* result) {{
   auto const max_num_class
       = *std::max_element(model_meta.num_class_.begin(), model_meta.num_class_.end());
   for (std::int32_t target_id = 0; target_id < model_meta.num_target_; ++target_id) {
-    fmt::print(oss, "  postprocess_impl(result[{offset}], {num_class});\n",
+    fmt::print(oss, "  postprocess_impl(&result[{offset}], {num_class});\n",
         "offset"_a = target_id * max_num_class, "num_class"_a = model_meta.num_class_[target_id]);
   }
   oss << "}\n";

src/compiler/compiler.cc

@@ -10,6 +10,7 @@
 #include <tl2cgen/detail/compiler/ast/builder.h>
 #include <tl2cgen/detail/compiler/codegen/codegen.h>
 #include <tl2cgen/detail/compiler/codegen/format_util.h>
+#include <tl2cgen/detail/filesystem.h>
 #include <treelite/tree.h>
 
 #include <filesystem>
@@ -51,11 +52,15 @@ namespace tl2cgen::compiler {
 
 void CompileModel(treelite::Model const& model, CompilerParam const& param,
     std::filesystem::path const& dirpath) {
+  tl2cgen::detail::filesystem::CreateDirectoryIfNotExist(dirpath);
   auto builder = LowerToAST(model, param);
   /* Generate C code */
   detail::codegen::CodeCollection gencode;
   detail::codegen::GenerateCodeFromAST(builder.GetRootNode(), gencode);
-  TL2CGEN_LOG(INFO) << "\n" << gencode;
+  // Write C code to disk
+  detail::codegen::WriteCodeToDisk(dirpath, gencode);
+  // Write recipe.json
+  detail::codegen::WriteBuildRecipeToDisk(dirpath, param.native_lib_name, gencode);
 }
 
 std::string DumpAST(treelite::Model const& model, CompilerParam const& param) {

src/filesystem.cc

@@ -24,14 +24,4 @@ void CreateDirectoryIfNotExist(std::filesystem::path const& dirpath) {
   }
 }
 
-void WriteToFile(std::filesystem::path const& path, std::string const& content) {
-  std::ofstream of(path);
-  of << content;
-}
-
-void WriteToFile(std::filesystem::path const& path, std::vector<char> const& content) {
-  std::ofstream of(path, std::ios::out | std::ios::binary);
-  of.write(content.data(), content.size());
-}
-
 }  // namespace tl2cgen::detail::filesystem