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