From 922fc85d91b9ea76d9f600176dae4d6b97521c45 Mon Sep 17 00:00:00 2001 From: Avimitin Date: Fri, 9 Aug 2024 16:20:18 +0800 Subject: [PATCH] [pytorch] rewrite MemRef container in C++ Signed-off-by: Avimitin --- tests/builder.nix | 1 + tests/pytorch/README.md | 48 ++++++++++++++++------------------- tests/pytorch/default.nix | 4 +-- tests/pytorch/demo/config.nix | 2 +- tests/pytorch/demo/demo.c | 39 ---------------------------- tests/pytorch/demo/demo.cc | 22 ++++++++++++++++ tests/pytorch/memref.h | 28 -------------------- tests/pytorch/memref.hpp | 44 ++++++++++++++++++++++++++++++++ 8 files changed, 92 insertions(+), 96 deletions(-) delete mode 100644 tests/pytorch/demo/demo.c create mode 100644 tests/pytorch/demo/demo.cc delete mode 100644 tests/pytorch/memref.h create mode 100644 tests/pytorch/memref.hpp diff --git a/tests/builder.nix b/tests/builder.nix index 4730af191..0533d7915 100644 --- a/tests/builder.nix +++ b/tests/builder.nix @@ -25,6 +25,7 @@ let name = pname; CC = "${stdenv.targetPlatform.config}-cc"; + CXX = "${stdenv.targetPlatform.config}-c++"; NIX_CFLAGS_COMPILE = let diff --git a/tests/pytorch/README.md b/tests/pytorch/README.md index 5e53fb203..031d07b93 100644 --- a/tests/pytorch/README.md +++ b/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 ".py" file. @@ -25,30 +25,26 @@ For each PyTorch tests, developers should call the MLIR model from " *output, + MemRef *arg1, + MemRef *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 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 diff --git a/tests/pytorch/default.nix b/tests/pytorch/default.nix index 2a4d60cc2..6b54203d4 100644 --- a/tests/pytorch/default.nix +++ b/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 diff --git a/tests/pytorch/demo/config.nix b/tests/pytorch/demo/config.nix index d00359f79..0a97e1a09 100644 --- a/tests/pytorch/demo/config.nix +++ b/tests/pytorch/demo/config.nix @@ -1,6 +1,6 @@ { includes = [ - ../memref.h + ../memref.hpp ]; buddyOptArgs = [ diff --git a/tests/pytorch/demo/demo.c b/tests/pytorch/demo/demo.c deleted file mode 100644 index 6b8b971bb..000000000 --- a/tests/pytorch/demo/demo.c +++ /dev/null @@ -1,39 +0,0 @@ -#include "memref.h" - -NEW_MEMREF(float, 1); - -extern void _mlir_ciface_forward(struct MemRef_float_dim1 *output, - struct MemRef_float_dim1 *arg1, - struct MemRef_float_dim1 *arg2); - -__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}, -}; - -__attribute((section(".vdata"))) float input_float_1[512] = {4, 5, 6}; -struct MemRef_float_dim1 input2 = { - .allocatedPtr = input_float_1, - .alignedPtr = input_float_1, - .offset = 0, - .sizes = {512}, - .strides = {1}, -}; - -__attribute((section(".vdata"))) float output_float_0[512]; -struct MemRef_float_dim1 output = { - .allocatedPtr = output_float_0, - .alignedPtr = output_float_0, - .offset = 0, - .sizes = {512}, - .strides = {1}, -}; - -int test() { - _mlir_ciface_forward(&output, &input1, &input2); - return 0; -} diff --git a/tests/pytorch/demo/demo.cc b/tests/pytorch/demo/demo.cc new file mode 100644 index 000000000..cef5a3081 --- /dev/null +++ b/tests/pytorch/demo/demo.cc @@ -0,0 +1,22 @@ +#include "memref.hpp" + +extern "C" void _mlir_ciface_forward(MemRef *output, + MemRef *arg1, + MemRef *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 input1(input_float_1, sizes); + +__attribute((section(".vdata"))) float input_float_2[512] = {4, 5, 6}; +MemRef input2(input_float_2, sizes); + +__attribute((section(".vdata"))) float output_float_1[512]; +MemRef output(output_float_1, sizes); + +extern "C" int test() { + _mlir_ciface_forward(&output, &input1, &input2); + return 0; +} diff --git a/tests/pytorch/memref.h b/tests/pytorch/memref.h deleted file mode 100644 index c468af15f..000000000 --- a/tests/pytorch/memref.h +++ /dev/null @@ -1,28 +0,0 @@ -#ifndef MEMREF_H -#define MEMREF_H - -// Generate a new struct with T-type, N-dimension memref with name -// MemRef_T_dimN. -// -// Example: -// -// NEW_MEMREF(float, 2); -// // Equals to -// struct MemRef_float_dim2 { -// float *allocatedPtr; -// float *alignedPtr; -// int offset; -// int sizes[2]; -// int strides[2]; -// }; -// -#define NEW_MEMREF(T, N) \ - struct MemRef_##T##_dim##N { \ - T *allocatedPtr; \ - T *alignedPtr; \ - int offset; \ - int sizes[N]; \ - int strides[N]; \ - } - -#endif diff --git a/tests/pytorch/memref.hpp b/tests/pytorch/memref.hpp new file mode 100644 index 000000000..09276cfc6 --- /dev/null +++ b/tests/pytorch/memref.hpp @@ -0,0 +1,44 @@ +#ifndef MEMREF_H +#define MEMREF_H + +#include +#include + +template 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 constexpr +MemRef::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 inline void MemRef::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