mylib.c

   1
   2 #include <TH/TH.h>
   3
   4 /*
   5
   6   Example of FFI extension I started from:
   7
   8     https://github.com/pytorch/extension-ffi.git
   9
  10   There is this tutorial
  11
  12     https://github.com/pytorch/tutorials/blob/master/Creating%20Extensions%20using%20FFI.md
  13
  14   And TH's Tensor definition are here in my install:
  15
  16     anaconda3/lib/python3.5/site-packages/torch/lib/include/TH/generic/THTensor.h
  17
  18  */
  19
  20 #include "flatland.h"
  21
  22 int generate_sequence(THByteTensor *output) {
  23   long nb_sequences = 1;
  24   long nb_images_per_sequence = 5;
  25   long depth = 3;
  26   long width = 64;
  27   long height = 64;
  28   long s;
  29   unsigned char *a, *b;
  30   int c, k, i, j, st0, st1, st2, st3, st4;
  31
  32   THByteTensor_resize5d(output, nb_sequences, nb_images_per_sequence, depth, height, width);
  33
  34   st0 = THByteTensor_stride(output, 0);
  35   st1 = THByteTensor_stride(output, 1);
  36   st2 = THByteTensor_stride(output, 2);
  37   st3 = THByteTensor_stride(output, 3);
  38   st4 = THByteTensor_stride(output, 4);
  39
  40   a =
  41     THByteTensor_storage(output)->data + THByteTensor_storageOffset(output);
  42
  43   for(s = 0; s < nb_sequences; s++) {
  44     unsigned char result[nb_images_per_sequence * depth * width * height];
  45     unsigned char *r = result;
  46     fl_generate_sequences(1, nb_images_per_sequence, width, height, result);
  47     for(k = 0; k < nb_images_per_sequence; k++) {
  48       for(c = 0; c < depth; c++) {
  49         for(i = 0; i < height; i++) {
  50           b = a
  51             + s * st0 + k * st1 + c * st2 + i * st3;
  52           for(j = 0; j < width; j++) {
  53             *b = (unsigned char) (*r);
  54             r++;
  55             b += st4;
  56           }
  57         }
  58       }
  59     }
  60   }
  61
  62   return 1;
  63 }