hipudding's Blog

#include <cstring>
#include <iostream>

#ifndef __CCE_KT_TEST__
#include "acl/acl.h"
#else
#include "tikicpulib.h"
#endif
#include "kernel_operator.h"
#include "data_loader.h"

using namespace AscendC;

#ifdef __CCE_KT_TEST__
#define __aicore__
#else
#define __aicore__ [aicore]
#endif

constexpr int BUFFER_NUM = 2;
constexpr int BLOCK_DIM = 16;

/*****************************Copy scalar to ubuf*****************************/
struct FlipTilingData {
  uint32_t height;
  uint32_t width;
  uint32_t channel;
};

inline __aicore__ int32_t align32(int32_t n) { return ((n + 31) & ~31); }
inline __aicore__ int32_t AlignDiv32(int32_t n) { return align32(n) / 32; }

#define CONVERT_TILING_DATA(tilingStruct, tilingDataPointer, tilingPointer) \
  __ubuf__ tilingStruct* tilingDataPointer =                                \
      reinterpret_cast<__ubuf__ tilingStruct*>(                             \
          (__ubuf__ uint8_t*)(tilingPointer));

#ifdef __CCE_KT_TEST__
#define INIT_TILING_DATA(tilingStruct, tilingDataPointer, tilingPointer) \
  CONVERT_TILING_DATA(tilingStruct, tilingDataPointer, tilingPointer);
#else
#define INIT_TILING_DATA(tilingStruct, tilingDataPointer, tilingPointer) \
  __ubuf__ uint8_t* tilingUbPointer = (__ubuf__ uint8_t*)get_imm(0);     \
  copy_gm_to_ubuf(((__ubuf__ uint8_t*)(tilingUbPointer)),                \
                  ((__gm__ uint8_t*)(tilingPointer)), 0, 1,              \
                  AlignDiv32(sizeof(tilingStruct)), 0, 0);               \
  CONVERT_TILING_DATA(tilingStruct, tilingDataPointer, tilingUbPointer); \
  pipe_barrier(PIPE_ALL);
#endif

#define GET_TILING_DATA(tilingData, tilingPointer) \
  INIT_TILING_DATA(FlipTilingData, tilingData, tilingPointer);

#define CHECK_ACL(x)                                                    \
  do {                                                                  \
    aclError __ret = x;                                                 \
    if (__ret != ACL_ERROR_NONE) {                                      \
      std::cerr << __FILE__ << ":" << __LINE__ << " aclError:" << __ret \
                << std::endl;                                           \
    }                                                                   \
  } while (0);

/*******************************Kernel function*******************************/
class KernelFlip {
 public:
  __aicore__ inline KernelFlip() {}
  __aicore__ inline void Init(GM_ADDR input, GM_ADDR output, uint32_t _height,
                              uint32_t _width, uint32_t _channel) {
    uint32_t blockNum = GetBlockNum();
    uint32_t blockIdx = GetBlockIdx();

    rowLength = _height / blockNum;
    startRowIdx = blockIdx * rowLength;
    if (startRowIdx + rowLength > _height) {
      rowLength = _height - startRowIdx;
    }
    width = _width;
    height = _height;
    channel = _channel;
    rowSize = width * channel;
    uint32_t bufferSize = align32(rowSize * sizeof(uint8_t));

    inputGM.SetGlobalBuffer((__gm__ uint8_t*)input + startRowIdx * rowSize,
                            rowLength * rowSize);
    outputGM.SetGlobalBuffer((__gm__ uint8_t*)output + startRowIdx * rowSize,
                             rowLength * rowSize);
    pipe.InitBuffer(inQueue, BUFFER_NUM, bufferSize);
    pipe.InitBuffer(outQueue, BUFFER_NUM, bufferSize);
  }

  __aicore__ inline void Process() {
    for (int32_t i = 0; i < rowLength; i++) {
      CopyIn(i);
      Compute(i);
      CopyOut(i);
    }
  }

 private:
  __aicore__ inline void CopyIn(int32_t loop) {
    LocalTensor<uint8_t> local = inQueue.AllocTensor<uint8_t>();
    DataCopy(local, inputGM[loop * rowSize], rowSize);
    inQueue.EnQue(local);
  }

  __aicore__ inline void Compute(int32_t loop) {
    LocalTensor<uint8_t> inputLocal = inQueue.DeQue<uint8_t>();
    LocalTensor<uint8_t> outputLocal = outQueue.AllocTensor<uint8_t>();
    for (int32_t i = 0; i < width; i++) {
      for (int32_t c = 0; c < channel; c++) {
        outputLocal.SetValue(
            i * channel + c,
            inputLocal.GetValue((width - i - 1) * channel + c));
      }
    }
    outQueue.EnQue<uint8_t>(outputLocal);
    inQueue.FreeTensor(inputLocal);
  }

  __aicore__ inline void CopyOut(int32_t loop) {
    LocalTensor<uint8_t> local = outQueue.DeQue<uint8_t>();
    DataCopy(outputGM[loop * rowSize], local, rowSize);
    outQueue.FreeTensor(local);
  }

 private:
  TPipe pipe;
  TQue<QuePosition::VECIN, BUFFER_NUM> inQueue;
  TQue<QuePosition::VECOUT, BUFFER_NUM> outQueue;
  GlobalTensor<uint8_t> inputGM, outputGM;
  uint32_t startRowIdx, rowLength, rowSize, height, width, channel;
};

/*******************************kernel interface******************************/
extern "C" __global__ __aicore__ void flip(GM_ADDR input, GM_ADDR output,
                                           GM_ADDR tiling) {
  GET_TILING_DATA(tilingData, tiling);
  KernelFlip op;
  op.Init(input, output, tilingData->height, tilingData->width,
          tilingData->channel);
  op.Process();
}

#ifndef __CCE_KT_TEST__
void flip_do(uint32_t blockDim, void* l2ctrl, void* stream, uint8_t* input,
             uint8_t* output, uint8_t* tiling) {
  flip<<<blockDim, l2ctrl, stream>>>(input, output, tiling);
}
#endif

/***********************************caller************************************/
int32_t main(int32_t argc, char* argv[]) {
  if (argc != 2) {
    std::cerr << "usage: " << argv[0] << " path/to/datafile" << std::endl;
    exit(-1);
  }

  uint32_t blockDim = BLOCK_DIM;
  uint32_t height, width, channel;
  uint8_t* data = readFile(argv[1], height, width, channel);
  const char* resultFile = std::string(argv[1]).append(".ret").c_str();

  uint32_t dataSize = width * height * channel * sizeof(uint8_t);
  size_t inputByteSize = dataSize;
  size_t outputByteSize = dataSize;
  size_t tilingSize = sizeof(FlipTilingData);

  uint8_t *inputHost, *outputHost, *tilingHost;
  uint32_t shape[]{height, width, channel};

#ifdef __CCE_KT_TEST__
  inputHost = (uint8_t*)AscendC::GmAlloc(inputByteSize);
  outputHost = (uint8_t*)AscendC::GmAlloc(outputByteSize);
  tilingHost = (uint8_t*)AscendC::GmAlloc(tilingSize);
  memcpy(tilingHost, shape, tilingSize);
  memcpy(inputHost, data, dataSize);

  AscendC::SetKernelMode(KernelMode::AIV_MODE);
  ICPU_RUN_KF(flip, blockDim, inputHost, outputHost, tilingHost);

  writeFile(resultFile, height, width, channel, outputHost);

  AscendC::GmFree((void*)inputHost);
  AscendC::GmFree((void*)outputHost);
  AscendC::GmFree((void*)tilingHost);
#else
  CHECK_ACL(aclInit(nullptr));
  aclrtContext context;
  int32_t deviceId = 0;
  CHECK_ACL(aclrtSetDevice(deviceId));
  CHECK_ACL(aclrtCreateContext(&context, deviceId));
  aclrtStream stream = nullptr;
  CHECK_ACL(aclrtCreateStream(&stream));

  uint8_t *inputDevice, *outputDevice, *tilingDevice;
  CHECK_ACL(aclrtMallocHost((void**)(&tilingHost), tilingSize));
  CHECK_ACL(aclrtMallocHost((void**)(&inputHost), inputByteSize));
  CHECK_ACL(aclrtMallocHost((void**)(&outputHost), outputByteSize));
  CHECK_ACL(aclrtMalloc((void**)&inputDevice, inputByteSize,
                        ACL_MEM_MALLOC_HUGE_FIRST));
  CHECK_ACL(aclrtMalloc((void**)&outputDevice, outputByteSize,
                        ACL_MEM_MALLOC_HUGE_FIRST));
  CHECK_ACL(aclrtMalloc((void**)&tilingDevice, tilingSize,
                        ACL_MEM_MALLOC_HUGE_FIRST));

  memcpy(tilingHost, shape, tilingSize);
  memcpy(inputHost, data, dataSize);

  CHECK_ACL(aclrtMemcpy(inputDevice, inputByteSize, inputHost, inputByteSize,
                        ACL_MEMCPY_HOST_TO_DEVICE));
  CHECK_ACL(aclrtMemcpy(tilingDevice, tilingSize, tilingHost, tilingSize,
                        ACL_MEMCPY_HOST_TO_DEVICE));

  flip_do(blockDim, nullptr, stream, inputDevice, outputDevice, tilingDevice);

  CHECK_ACL(aclrtSynchronizeStream(stream));
  CHECK_ACL(aclrtMemcpy(outputHost, outputByteSize, outputDevice,
                        outputByteSize, ACL_MEMCPY_DEVICE_TO_HOST));
  writeFile(resultFile, height, width, channel, outputHost);

  CHECK_ACL(aclrtFree(inputDevice));
  CHECK_ACL(aclrtFree(outputDevice));
  CHECK_ACL(aclrtFree(tilingDevice));
  CHECK_ACL(aclrtFreeHost(inputHost));
  CHECK_ACL(aclrtFreeHost(outputHost));
  CHECK_ACL(aclrtFreeHost(tilingHost));

  CHECK_ACL(aclrtDestroyStream(stream));
  CHECK_ACL(aclrtDestroyContext(context));
  CHECK_ACL(aclrtResetDevice(deviceId));
  CHECK_ACL(aclFinalize());
#endif
  free(data);
  return 0;
}

#include <iostream>
#include "data_loader.h"

__global__ void flip(uint8_t* input, uint8_t* output, uint32_t height,
                     uint32_t width, uint32_t channel) {
  int rowIdx = threadIdx.x + blockIdx.x * blockDim.x;
  int stride = blockDim.x * gridDim.x;
  int rowSize = width * channel;

  for (int row = rowIdx; row < height; row += stride) {
    int startOffset = row * rowSize;
    for (int idx = 0; idx < width; idx++) {
      for (int c = 0; c < channel; c++) {
        output[startOffset + idx * channel + c] =
            input[startOffset + (width - idx - 1) * channel + c];
      }
    }
  }
}

int main(int32_t argc, char* argv[]) {
  if (argc != 2) {
    std::cerr << "usage: " << argv[0] << " path/to/datafile" << std::endl;
    exit(-1);
  }
  uint32_t height, width, channel;
  char fileName[256], resultFile[256];
  memset(fileName, 0, 256);
  memset(resultFile, 0, 256);
  strcpy(fileName, argv[1]);
  strcat(resultFile, fileName);
  strcat(resultFile, ".ret");
  uint8_t* data = readFile(fileName, height, width, channel);

  uint32_t dataSize = width * height * channel * sizeof(uint8_t);
  size_t inputByteSize = dataSize;
  size_t outputByteSize = dataSize;
  uint8_t *input, *output;
  cudaMallocManaged((void**)&input, inputByteSize);
  cudaMallocManaged((void**)&output, outputByteSize);

  memcpy(input, data, inputByteSize);

  dim3 blockSize(256);
  dim3 gridSize((height + blockSize.x - 1) / blockSize.x);

  flip<<<gridSize, blockSize>>>(input, output, height, width, channel);
  cudaDeviceSynchronize();

  writeFile(resultFile, height, width, channel, output);

  cudaFree(input);
  cudaFree(output);

  return 0;
}

#include <Python.h>

// 实际业务函数体
int sum(int a, int b) {
    return a + b;
}

/**
 *   业务函数的包装函数，该函数注册到PyMethodDef中，供Python虚拟机调用
 *   传入两个参数，self和args，args中的参数需要从对象中解析出来。
 *   PyArg_ParseTuple接收一个格式串，根据格式传将参数从PyObject中解析出来
 *   返回一个整数，需要包装成PyObject类型。
 **/
static PyObject* sum_wrapper(PyObject* self, PyObject* args) {
    int a, b;
    if (!PyArg_ParseTuple(args, "ii", &a, &b)) {
        return NULL;
    }
    int result = sum(a, b);
    return Py_BuildValue("i", result);
}

/**
 *   描述需要注册到模块中的函数，每个函数定义一行，以{NULL, NULL, 0, NULL}结束
 *   每行结构有4个参数，分别为：python调用的函数名，实际调用的包装函数；传入参数的类型；
 *                          函数的描述。
 **/
static PyMethodDef SumMethods[] = {
    {"sum", sum_wrapper, METH_VARARGS, "Calculate the sum of two integers."},
    {NULL, NULL, 0, NULL}
};

/**
 *   描述模块属性，参数分别为：固定参数；模块名，就是import后的名字；模块描述；？？？；
 *                         函数数组。
 **/
static struct PyModuleDef summodule = {
    PyModuleDef_HEAD_INIT,
    "libsum",
    "A module that adds two numbers",
    -1,
    SumMethods
};

// 初始化模块，以PyInit_模块名 命名
PyMODINIT_FUNC PyInit_libsum(void) {
    return PyModule_Create(&summodule);
}

double avg(double *arr, size_t len) {
    double sum = 0;
    for (size_t i = 0; i < len; i++) {
        sum += arr[i];
    }
    return (sum/len);
}

static PyObject *avg_wrapper(PyObject *self, PyObject *args) {
  PyObject *bufobj;
  Py_buffer view;
  double result;
  if (!PyArg_ParseTuple(args, "O", &bufobj)) {
    return NULL;
  }

  // 数组内存是连续的
  if (PyObject_GetBuffer(bufobj, &view,
      PyBUF_ANY_CONTIGUOUS | PyBUF_FORMAT) == -1) {
    return NULL;
  }
  
  // 判断是一维数组
  if (view.ndim != 1) {
    PyErr_SetString(PyExc_TypeError, "Expected a 1-dimensional array");
    PyBuffer_Release(&view);
    return NULL;
  }

  // 检查是否是double类型的数组
  if (strcmp(view.format,"d") != 0) {
    PyErr_SetString(PyExc_TypeError, "Expected an array of doubles");
    PyBuffer_Release(&view);
    return NULL;
  }

  // View.buf是double数组的首指针，view.shape[0]是这一维度的长度
  result = avg((double*)view.buf, view.shape[0]);

  // 需要显式释放Py_buffer
  PyBuffer_Release(&view);
  return Py_BuildValue("d", result);
}

// 定义C++类型
class DoubleWarpper {
 private:
  double m_value;

 public:
  DoubleWarpper(double value) : m_value(value) {}
  virtual ~DoubleWarpper() { m_value = 0; }

  double get() { return m_value; }
  void set(double value) { m_value = value; }

  DoubleWarpper operator+(const DoubleWarpper &other) const {
    return DoubleWarpper(m_value + other.m_value);
  }
  DoubleWarpper operator-(const DoubleWarpper &other) const {
    return DoubleWarpper(m_value - other.m_value);
  }
  DoubleWarpper operator*(const DoubleWarpper &other) const {
    return DoubleWarpper(m_value * other.m_value);
  }
  DoubleWarpper operator/(const DoubleWarpper &other) const {
    return DoubleWarpper(m_value / other.m_value);
  }
};

// 包装成Python C API需要的结构
typedef struct {
  PyObject_HEAD;
  DoubleWarpper *warpper = nullptr;
} PyDoubleWarpper;

// Python对象构造时调用，需要返回一个PyDoubleWarpper类型的对象。 kwds？？？
static PyObject *pycal_new(PyTypeObject *type, PyObject *args, PyObject *kwds) {
  PyDoubleWarpper *self;
  self = (PyDoubleWarpper *)type->tp_alloc(type, 0);
  char *kwlist[] = {"value", 0};
  double value = 0;
  if (!PyArg_ParseTupleAndKeywords(args, kwds, "d", kwlist, &value)) {
    Py_DECREF(self);
    return nullptr;
  }
  self->warpper = new DoubleWarpper(value);
  return (PyObject *)self;
}

// Python对象销毁时调用
static void *pycal_dealloc(PyObject *py_cal) {
  delete ((PyDoubleWarpper *)py_cal)->warpper;
  Py_TYPE(py_cal)->tp_free(py_cal);
  return (void *)0;
}

// 从PyObject中获取Calculator指针
static DoubleWarpper *get_cal(PyObject *obj) {
  return ((PyDoubleWarpper *)obj)->warpper;
}

// 通过Calculator指针构造PyObject
static PyObject *return_cal(DoubleWarpper *cal, PyTypeObject *type) {
  PyDoubleWarpper *obj = PyObject_NEW(PyDoubleWarpper, type);
  obj->warpper = cal;
  return (PyObject *)obj;
}

// Calculator类提供的函数的包装
static PyObject *pycal_set(PyObject *self, PyObject *args) {
  DoubleWarpper *cal = get_cal(self);
  double value = 0;
  if (!PyArg_ParseTuple(args, "d", &value)) {
    return nullptr;
  }

  cal->set(value);
  return Py_BuildValue("i", 0);
}

// 打印对象
static PyObject *pycal_str(PyObject *self) {
  DoubleWarpper *cal = get_cal(self);
  std::stringstream ss;
  ss<<cal->get();
  return Py_BuildValue("s", ss.str().c_str());
}

// 实现加减乘除
static PyObject *pycal_add(PyObject *a, PyObject *b) {
  DoubleWarpper *cal_a = get_cal(a);
  DoubleWarpper *cal_b = get_cal(b);
  DoubleWarpper *ret = new DoubleWarpper(*cal_a + *cal_b);
  return return_cal(ret, a->ob_type);
}

static PyObject *pycal_minus(PyObject *a, PyObject *b) {
  DoubleWarpper *cal_a = get_cal(a);
  DoubleWarpper *cal_b = get_cal(b);
  DoubleWarpper *ret = new DoubleWarpper(*cal_a - *cal_b);
  return return_cal(ret, a->ob_type);
}

static PyObject *pycal_multipy(PyObject *a, PyObject *b) {
  DoubleWarpper *cal_a = get_cal(a);
  DoubleWarpper *cal_b = get_cal(b);
  DoubleWarpper *ret = new DoubleWarpper(*cal_a * *cal_b);
  return return_cal(ret, a->ob_type);
}

static PyObject *pycal_divide(PyObject *a, PyObject *b) {
  DoubleWarpper *cal_a = get_cal(a);
  DoubleWarpper *cal_b = get_cal(b);
  DoubleWarpper *ret = new DoubleWarpper(*cal_a / *cal_b);
  return return_cal(ret, a->ob_type);
}

// 对象的数字属性，这里仅实现了加减乘除
static PyNumberMethods numberMethods = {
    pycal_add,      // nb_add
    pycal_minus,    // nb_subtract;
    pycal_multipy,  // nb_multiply
    nullptr,        // nb_remainder;
    nullptr,        // nb_divmod;
    nullptr,        // nb_power;
    nullptr,        // nb_negative;
    nullptr,        // nb_positive;
    nullptr,        // nb_absolute;
    nullptr,        // nb_bool;
    nullptr,        // nb_invert;
    nullptr,        // nb_lshift;
    nullptr,        // nb_rshift;
    nullptr,        // nb_and;
    nullptr,        // nb_xor;
    nullptr,        // nb_or;
    nullptr,        // nb_int;
    nullptr,        // nb_reserved;
    nullptr,        // nb_float;
    nullptr,        // nb_inplace_add;
    nullptr,        // nb_inplace_subtract;
    nullptr,        // nb_inplace_multiply;
    nullptr,        // nb_inplace_remainder;
    nullptr,        // nb_inplace_power;
    nullptr,        // nb_inplace_lshift;
    nullptr,        // nb_inplace_rshift;
    nullptr,        // nb_inplace_and;
    nullptr,        // nb_inplace_xor;
    nullptr,        // nb_inplace_or;
    nullptr,        // nb_floor_divide;
    pycal_divide,   // nb_true_divide;
    nullptr,        // nb_inplace_floor_divide;
    nullptr,        // nb_inplace_true_divide;
    nullptr,        // nb_index;
    nullptr,        // nb_matrix_multiply;
    nullptr         // nb_inplace_matrix_multiply;

};

// PyCalculator对象的成员函数
static PyMethodDef pycal_methods[] = {
    {"set", (PyCFunction)pycal_set, METH_VARARGS, "set DoubleWarpper value."},
    {nullptr}};

// PyCalculator对象内容
static PyTypeObject DoubleWarpperType = {
    PyVarObject_HEAD_INIT(nullptr, 0) "libsum.DoubleWarpper", /* tp_name */
    sizeof(PyDoubleWarpper),                                  /* tp_basicsize */
    0,                                                        /* tp_itemsize */
    (destructor)pycal_dealloc,                                /* tp_dealloc */
    0,                                        /* tp_vectorcall_offset */
    nullptr,                                  /* tp_getattr */
    nullptr,                                  /* tp_setattr */
    nullptr,                                  /* tp_reserved */
    nullptr,                                  /* tp_repr */
    &numberMethods,                           /* tp_as_number */
    nullptr,                                  /* tp_as_sequence */
    nullptr,                                  /* tp_as_mapping */
    nullptr,                                  /* tp_hash  */
    nullptr,                                  /* tp_call */
    pycal_str,                                /* tp_str */
    nullptr,                                  /* tp_getattro */
    nullptr,                                  /* tp_setattro */
    nullptr,                                  /* tp_as_buffer */
    Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE, /* tp_flags */
    "Coustom DoubleWarpper class.",           /* tp_doc */
    nullptr,                                  /* tp_traverse */
    nullptr,                                  /* tp_clear */
    nullptr,                                  /* tp_richcompare */
    0,                                        /* tp_weaklistoffset */
    nullptr,                                  /* tp_iter */
    nullptr,                                  /* tp_iternext */
    pycal_methods,                            /* tp_methods */
    nullptr,                                  /* tp_members */
    nullptr,                                  /* tp_getset */
    nullptr,                                  /* tp_base */
    nullptr,                                  /* tp_dict */
    nullptr,                                  /* tp_descr_get */
    nullptr,                                  /* tp_descr_set */
    0,                                        /* tp_dictoffset */
    nullptr,                                  /* tp_init */
    nullptr,                                  /* tp_alloc */
    pycal_new                                 /* tp_new */
};


PyMODINIT_FUNC PyInit_libsum(void) {
  if (PyType_Ready(&DoubleWarpperType) < 0) {
    return nullptr;
  }

  PyObject *module = PyModule_Create(&summodule);
  if (module == nullptr) {
    return nullptr;
  }

  // 注册对象，这里为啥要引用+1？？
  Py_INCREF(&DoubleWarpperType);
  if (PyModule_AddObject(module, "DoubleWarpper",
                         (PyObject *)&DoubleWarpperType) < 0) {
    Py_DECREF(&DoubleWarpperType);
    Py_DECREF(module);
    return nullptr;
  }

  return module;
}

#include <Python.h>
#include <pybind11/pybind11.h>
#include <pybind11/stl.h>
#include <vector>

namespace py = pybind11;

int sum(int a, int b) { return a + b; }

double avg(std::vector<double> &arr) {
  double sum = 0;
  for (auto it = arr.begin(); it != arr.end(); it++) {
    sum += (*it);
  }
  return (sum / arr.size());
}

PYBIND11_MODULE(libsum, m) {
    m.doc() = "Py module example.";
    m.def("sum", &sum, "Calculate the sum of two integers.");
    m.def("avg", &avg, "alculate the avg of a double array.");
}