Merge pull request #915 from andresovela/xformer-weights-on-external-…

…memory Improve support for writing weights in array on source files
xmos · Jul 24, 2024 · a52a2e1 · a52a2e1
2 parents 468784d + d569b34
commit a52a2e1
Show file tree

Hide file tree

Showing 5 changed files with 48 additions and 18 deletions.
diff --git a/xformer/Transforms/Options.h b/xformer/Transforms/Options.h
@@ -14,7 +14,8 @@ extern llvm::cl::opt<bool> enableBetaFloatOption;
 extern llvm::cl::opt<unsigned> threadCountOption;
 extern llvm::cl::opt<std::string> weightsFilenameOption;
 extern llvm::cl::opt<unsigned> loadExternallyIfLargerOption;
-extern llvm::cl::opt<bool> tileLoadOption;
+extern llvm::cl::opt<bool> weightsAsArrayOption;
+extern llvm::cl::opt<bool> weightsInExternalMemory;
 extern llvm::cl::opt<unsigned> maxLoadExternalSizeOption;
 extern llvm::cl::opt<double> convQuantErrorThresholdOption;
 extern llvm::cl::opt<bool> convForceErrorCheckOption;

diff --git a/xformer/Transforms/WriteWeights.cpp b/xformer/Transforms/WriteWeights.cpp
@@ -145,9 +145,10 @@ void WriteWeights::runOnOperation() {
   patterns.insert<WriteWeightsPattern>(&tensorsVec, ctx);
   (void)applyPatternsAndFoldGreedily(func, std::move(patterns));
 
-  if (tileLoadOption) {
+  if (weightsAsArrayOption) {
     if (failed(utils::writeTileServerDataToFile(weightsFilenameOption,
-                                                tensorsVec))) {
+                                                tensorsVec,
+                                                weightsInExternalMemory))) {
       f.emitError("Failed to write tile data!");
       signalPassFailure();
       return;

diff --git a/xformer/Utils/FileIO.cpp b/xformer/Utils/FileIO.cpp
@@ -36,35 +36,51 @@ LogicalResult writeWeightsToFile(const std::string &filename,
 
 LogicalResult
 writeTileServerDataToFile(const std::string &filename,
-                          std::vector<std::vector<char>> tensorsVec) {
+                          std::vector<std::vector<char>> tensorsVec,
+                          bool placeInExternalMemory) {
   // Add header
   auto tileHeader = utils::tileRamHeader();
   tensorsVec.insert(tensorsVec.begin(), tileHeader);
 
-  std::ostringstream out;
-  out << R"(#ifndef TILESERVERGEN_H
-#define TILESERVERGEN_H
+  std::ostringstream cOut;
+  cOut << R"(#include <stdint.h>)";
 
-const int8_t tile_server_weights[] = {
-)";
+  if (placeInExternalMemory) {
+    cOut << "\n\n" << R"(__attribute__ ((section(".ExtMem.data"))))" << "\n";
+  }
+
+  cOut << "const int8_t tile_server_weights[] = {\n";
   int lineEnding = 0;
+  int weightsSize = 0;
   for (auto const &tensor : tensorsVec) {
     for (auto const &i : tensor) {
-      out << (int)i << ", ";
+      cOut << (int)i << ", ";
       lineEnding++;
+      weightsSize++;
       if (lineEnding > 80) {
-        out << "\n";
+        cOut << "\n";
         lineEnding = 0;
       }
     }
   }
 
-  out << R"(};
+  cOut << R"(};
+)";
+
+  if (failed(utils::writeDataToFile(filename + ".c", cOut.str()))) {
+    return failure();
+  }
+
+  std::ostringstream hOut;
+    hOut << R"(#ifndef TILESERVERGEN_H
+#define TILESERVERGEN_H
+
+#define TILE_SERVER_WEIGHTS_SIZE ()" << weightsSize << R"(U)
 
 #endif // TILESERVERGEN_H
 )";
 
-  return utils::writeDataToFile(filename, out.str());
+  return utils::writeDataToFile(filename + ".h", hOut.str());
 }
 
 LogicalResult getFlatBufferStringFromMLIR(

diff --git a/xformer/Utils/FileIO.h b/xformer/Utils/FileIO.h
@@ -15,7 +15,8 @@ LogicalResult writeWeightsToFile(const std::string &filename,
 
 LogicalResult
 writeTileServerDataToFile(const std::string &filename,
-                          std::vector<std::vector<char>> tensorsVec);
+                          std::vector<std::vector<char>> tensorsVec,
+                          bool placeInExternalMemory);
 
 LogicalResult getFlatBufferStringFromMLIR(
     mlir::ModuleOp module, std::map<std::string, std::string> metadata,

diff --git a/xformer/XCoreOptMain.cpp b/xformer/XCoreOptMain.cpp
@@ -72,8 +72,12 @@ cl::alias aliasWeightsFilenameOption("f",
                                      cl::desc("Alias for --xcore-weights-file"),
                                      cl::aliasopt(weightsFilenameOption));
 
-cl::opt<bool> tileLoadOption("xcore-load-tile",
-                             cl::desc("Enable loading weights from a tile."),
+cl::opt<bool> weightsAsArrayOption("xcore-write-weights-as-array",
+                             cl::desc("Write the weights in the form of an array in a source file (creates .c/.h files with <xcore-weights-file> as the file name)."),
+                             cl::init(false), cl::cat(XformerCategory));
+
+cl::opt<bool> weightsInExternalMemory("xcore-weights-in-external-memory",
+                             cl::desc("Annotate the generated weights array with an attribute to place it in external memory."),
                              cl::init(false), cl::cat(XformerCategory));
 
 cl::opt<unsigned> loadExternallyIfLargerOption(
@@ -453,11 +457,18 @@ int main(int argc, char **argv) {
         "Please specify an output filename using the -o option!");
   }
 
-  if (mlir::xcore::tileLoadOption.getNumOccurrences() > 0 &&
+  if (mlir::xcore::weightsAsArrayOption.getNumOccurrences() > 0 &&
       mlir::xcore::threadCountOption < 4) {
+    // TODO: This feels wrong considering this is just about the export format of the weights
     return failedMessage("Please specify at least four threads using "
                          "xcore-thread-count option when using the "
-                         "xcore-load-tile option!");
+                         "xcore-write-weights-as-array option!");
+  }
+
+  if (mlir::xcore::weightsInExternalMemory.getNumOccurrences() > 0 &&
+      mlir::xcore::weightsAsArrayOption.getNumOccurrences() == 0) {
+    return failedMessage("Please specify the xcore-write-weights-as-array"
+                         "when using the xcore-weights-in-external-memory option!");
   }
 
   if (mlir::xcore::loadExternallyIfLargerOption.getNumOccurrences() > 0 &&