Adding tensorContract CPU support

llvm/projects/hpvm-tensor-rt/dnn_sources/src/test_ops.cc

@@ -1148,9 +1148,41 @@ void* testTensorSelect(void* argmax_out){
 }
 
 
-void testTensorContract(){
+void testTensorContract(void* select_out){
 
+  Tensor* input = (Tensor*) create4DTensor(CUDNN_DATA_FLOAT, CUDNN_TENSOR_NCHW, 4, 4, 1, 1);
+  float* host_ptr = (float*) ((struct Tensor*) input)->host_data;
+
+  // Input 0
+  host_ptr[0] = 1;
+  host_ptr[1] = 1; 
+  host_ptr[2] = 1;
+  host_ptr[3] = 1;
+
+  // Input 1
+  host_ptr[4] = 2;
+  host_ptr[5] = 2;
+  host_ptr[6] = 2;
+  host_ptr[7] = 2;
   
+  // Input 2
+  host_ptr[8] = 3;
+  host_ptr[9] = 3;
+  host_ptr[10] = 3; 
+  host_ptr[11] = 3; 
+
+  // Input 3
+  host_ptr[12] = 4; 
+  host_ptr[13] = 4;
+  host_ptr[14] = 4;
+  host_ptr[15] = 4;
+
+
+  void* contract_out = tensorContract(input, select_out);
+  printf ("***** tensorContract output \n");
+
+  printTensorValues(contract_out);
+
 }
 
 
@@ -1158,8 +1190,8 @@ void testTensorContract(){
 void testNewTensorOps(){
 
   void* argmax_out = testTensorArgMax();
-  testTensorSelect(argmax_out);
-  testTensorContract();
+  void* select_out = testTensorSelect(argmax_out);
+  testTensorContract(select_out);
   
 }
 

llvm/projects/hpvm-tensor-rt/tensor_runtime/include/tensor_custom_ops_cpu.h

@@ -44,6 +44,7 @@ void* tensorArgMax(void* input_ptr){
 
 
 
+
 void* tensorSelect(void* input_ptr, float target_value){
 
   Tensor* input = (Tensor*) input_ptr;
@@ -72,3 +73,64 @@ void* tensorSelect(void* input_ptr, float target_value){
       
   return output;
 }
+
+
+
+long getOnesInVector(float* vector_host_ptr, long vector_length){
+
+  long ones_count = 0;
+  for(int i = 0; i < vector_length; i++){
+
+    if(vector_host_ptr[i] == 1)
+      ones_count += 1;
+  }
+
+  return ones_count;
+}
+
+
+void* tensorContract(void* input_ptr, void* bitvector_ptr){
+
+  Tensor* input = (Tensor*) input_ptr;
+  float* host_ptr = (float*) input->host_data;
+
+  Tensor* bitvector = (Tensor*) bitvector_ptr;
+  float* vector_host_ptr = (float*) bitvector->host_data;  
+  long vector_length = bitvector->dims.dim_sizes[0];
+  
+  long reduced_batch_size = getOnesInVector(vector_host_ptr, vector_length); 
+  
+  long batch_size = input->dims.dim_sizes[0]; 
+  long channels = input->dims.dim_sizes[1];
+  long height = input->dims.dim_sizes[2];
+  long width = input->dims.dim_sizes[3];
+
+  long image_size = channels * height * width; // Computing size of each image
+  
+  if (batch_size != vector_length){
+    printf("ERROR: bitvector length has to match input batch size \n");
+    abort();
+  }
+
+  Tensor* output = (Tensor *) create4DTensor(0, 0, reduced_batch_size, channels, height, width);
+  changeTensorPlacement(output, HOST);    
+  float* out_ptr = (float*) output->host_data;
+
+  long out_index = 0;
+  for(int i = 0; i < batch_size; i++){
+
+    // Include image if corresponding index in bitvector is '1'
+    if (vector_host_ptr[i] == 1){
+
+      for(int j = 0; j < image_size; j++){
+
+	out_ptr[j] = host_ptr[i * image_size + j];
+      }
+
+      out_ptr += image_size; // Update the output pointer to the next image boundary
+    }	   
+  }
+      
+  return output;
+}
+