[pytorch] rewrite MemRef container in C++

Signed-off-by: Avimitin <[email protected]>

tests/builder.nix

@@ -25,6 +25,7 @@ let
     name = pname;
 
     CC = "${stdenv.targetPlatform.config}-cc";
+    CXX = "${stdenv.targetPlatform.config}-c++";
 
     NIX_CFLAGS_COMPILE =
       let

tests/pytorch/README.md

@@ -8,7 +8,7 @@ Assuming that the new PyTorch test have project name call `demo`, let's create t
 cd tests/pytorch
 mkdir -p demo
 cd demo
-touch demo.c demo.py config.nix
+touch demo.cc demo.py config.nix
 ```
 
 Developers should put their PyTorch implementation into "<project-name>.py" file.
@@ -25,30 +25,26 @@ For each PyTorch tests, developers should call the MLIR model from "<project-nam
 In our case, here is an example "demo.c" file:
 
 ```c
-// 1. Include the MemRef wrapper
-#include "memref.h"
-
-// 2. Create corresponding MemRef struct with data type `float` and one dimension.
-NEW_MEMREF(float, 1);
-
-// 3. Declare the MLIR model interface
-extern void _mlir_ciface_forward(struct MemRef_float_dim1 *output,
-                                 struct MemRef_float_dim1 *arg1,
-                                 struct MemRef_float_dim1 *arg2);
-
-// 4. Create example data array. The ".vdata" attribute will help emulator load the data into correct memory.
-__attribute((section(".vdata"))) float input_float_0[512] = {1, 2, 3};
-struct MemRef_float_dim1 input1 = {
-    .allocatedPtr = input_float_0,
-    .alignedPtr = input_float_0,
-    .offset = 0,
-    .sizes = {512},
-    .strides = {1},
-};
-
-// 5. Declare the main entry. In t1 all tests entry should be `int test()` instead of main().
-int test() {
-  _mlir_ciface_forward(&output, &input1, &input2);
+// 1. include the memref C++ wrapper.
+#include "memref.hpp"
+
+// 2. Declare the MLIR C interface, the argument layout can be guess from the generated MLIR model.
+extern "C" void _mlir_ciface_forward(MemRef<float, 1> *output,
+                                     MemRef<float, 1> *arg1,
+                                     MemRef<float, 1> *arg2);
+
+// 3. Declare the data sizes, here we use a vector that is one-dimension, with length 512.
+static const int32_t sizes[1] = {512};
+
+// 4. Declare a static data and add ".vdata" annotation. This can indicate the emulator to put these data to correct memory.
+__attribute((section(".vdata"))) float input_float_1[512] = {1, 2, 3};
+
+// 5. Declare a one dimension MemRef with float data type.
+MemRef<float, 1> input1(input_float_1, sizes);
+
+// 6. Mark the test function as extern "C", so that the linker can link it with our main function.
+extern "C" int test() {
+  // call _mlir_ciface_forward(...)
   return 0;
 }
 ```
@@ -61,7 +57,7 @@ file to indicate our build system to find and build the test case:
   # Tell our build system to include the memref.h header.
   # Developer could add extra headers here.
   includes = [
-    ../memref.h
+    ../memref.hpp
   ];
 
   # Tell the build system to run buddy-opt with three phrase, with arguments to run in each phrase

tests/pytorch/default.nix

@@ -86,9 +86,9 @@ let
         buddy-llc forward.ll $buddyLLCArgs --filetype=obj -o forward.o
 
         echo "Using include dir $buddyIncludeDir"
+        $CXX -nostdlib -I$buddyIncludeDir -c ${caseName}.cc -o host.o
         $CC -T${linkerScript} \
-          -I$buddyIncludeDir \
-          ${caseName}.c forward.o ${t1main} \
+          host.o forward.o ${t1main} \
           -o $pname.elf
 
         runHook postBuild

tests/pytorch/demo/config.nix

@@ -1,6 +1,6 @@
 {
   includes = [
-    ../memref.h
+    ../memref.hpp
   ];
 
   buddyOptArgs = [

tests/pytorch/demo/demo.cc

@@ -0,0 +1,22 @@
+#include "memref.hpp"
+
+extern "C" void _mlir_ciface_forward(MemRef<float, 1> *output,
+                                     MemRef<float, 1> *arg1,
+                                     MemRef<float, 1> *arg2);
+
+// One-dimension, with length 512
+static const int32_t sizes[1] = {512};
+
+__attribute((section(".vdata"))) float input_float_1[512] = {1, 2, 3};
+MemRef<float, 1> input1(input_float_1, sizes);
+
+__attribute((section(".vdata"))) float input_float_2[512] = {4, 5, 6};
+MemRef<float, 1> input2(input_float_2, sizes);
+
+__attribute((section(".vdata"))) float output_float_1[512];
+MemRef<float, 1> output(output_float_1, sizes);
+
+extern "C" int test() {
+  _mlir_ciface_forward(&output, &input1, &input2);
+  return 0;
+}

tests/pytorch/memref.hpp

@@ -0,0 +1,44 @@
+#ifndef MEMREF_H
+#define MEMREF_H
+
+#include <cstddef>
+#include <cstdint>
+
+template <typename T, size_t N> class MemRef {
+public:
+  constexpr MemRef(T *data, const int32_t sizes[N]);
+
+protected:
+  inline void setStrides();
+
+  // https://github.com/llvm/llvm-project/blob/a50b9633357007ff886f3fd228ca4b8a9b9b9852/mlir/lib/Conversion/LLVMCommon/TypeConverter.cpp#L401
+  T *allocated = nullptr;
+  T *aligned = nullptr;
+  int32_t offset = 0;
+  int32_t sizes[N];
+  int32_t strides[N];
+};
+
+template <typename T, std::size_t N> constexpr
+MemRef<T, N>::MemRef(T *data, const int32_t sizes[N]) {
+  for (size_t i = 0; i < N; i++) {
+    this->sizes[i] = sizes[i];
+  }
+
+  setStrides();
+
+  allocated = data;
+  aligned = data;
+}
+
+template <typename T, std::size_t N> inline void MemRef<T, N>::setStrides() {
+  strides[N - 1] = 1;
+  if (N < 2)
+    return;
+
+  for (std::size_t i = N - 1; i > 0; i--) {
+    strides[i - 1] = strides[i] * sizes[i];
+  }
+}
+
+#endif