From 87380afc898d6e4ade7def4c89e0ac91f8523f85 Mon Sep 17 00:00:00 2001
From: Damian Rouson <rouson@lbl.gov>
Date: Thu, 8 Feb 2024 08:50:48 -0800
Subject: [PATCH] WIP: start batch inference

---
 example/concurrent-inferences.f90            | 17 +++++++++++++++++
 src/inference_engine/inference_engine_m_.f90 | 13 +++++++++++--
 src/inference_engine/inference_engine_s.f90  |  8 +++++++-
 3 files changed, 35 insertions(+), 3 deletions(-)

diff --git a/example/concurrent-inferences.f90 b/example/concurrent-inferences.f90
index 65e35d0f8..e88ea55d1 100644
--- a/example/concurrent-inferences.f90
+++ b/example/concurrent-inferences.f90
@@ -77,6 +77,23 @@ program concurrent_inferences
       call system_clock(t_finish)
       print *,"Concurrent inference time with non-type-bound procedure: ", real(t_finish - t_start, real64)/real(clock_rate, real64)
 
+      print *,"Performing batched inferences on an intrinsic array"
+      block 
+        real ,allocatable :: inputs_batch(:,:,:,:), outputs_batch(:,:,:,:)
+        integer n
+        associate(num_inputs => inputs(1,1,1)%num_inputs())
+          inputs_batch = reshape([((((inputs(i,j,k)%values(), &
+            i=1,size(inputs,1)), j=1,size(inputs,2)), k=1,size(inputs,3)), n=1,num_inputs())], &
+            shape=[shape(inputs), num_inputs])
+        end associate
+
+        call system_clock(t_start, clock_rate)
+        outputs_batch = inference_engine%infer(tensor_batch)  ! implicit allocation of outputs array
+        call system_clock(t_finish)
+        print *,"Elemental inference time: ", real(t_finish - t_start, real64)/real(clock_rate, real64)
+
+      end block
+
     end block
   end block
 
diff --git a/src/inference_engine/inference_engine_m_.f90 b/src/inference_engine/inference_engine_m_.f90
index 63a05042c..2d6c8c819 100644
--- a/src/inference_engine/inference_engine_m_.f90
+++ b/src/inference_engine/inference_engine_m_.f90
@@ -26,7 +26,9 @@ module inference_engine_m_
     integer, allocatable :: nodes_(:)
     class(activation_strategy_t), allocatable :: activation_strategy_ ! Strategy Pattern facilitates elemental activation
   contains
-    procedure :: infer
+    procedure :: single_infer
+    procedure :: batch_infer
+    generic :: infer => single_infer, batch_infer
     procedure :: to_json
     procedure :: num_inputs
     procedure :: num_outputs
@@ -105,13 +107,20 @@ elemental module subroutine assert_conformable_with(self, inference_engine)
       type(inference_engine_t), intent(in) :: inference_engine
     end subroutine
 
-    elemental module function infer(self, inputs) result(outputs)
+    elemental module function single_infer(self, inputs) result(outputs)
       implicit none
       class(inference_engine_t), intent(in) :: self
       type(tensor_t), intent(in) :: inputs
       type(tensor_t) outputs
     end function
 
+  pure module function batch_infer(self, inputs) result(outputs)
+    implicit none
+    class(inference_engine_t), intent(in) :: self
+    real, intent(in) :: inputs(:,:,:,:)
+    real, allocatable :: outputs(:,:,:)
+  end function
+
     elemental module function num_outputs(self) result(output_count)
       implicit none
       class(inference_engine_t), intent(in) :: self
diff --git a/src/inference_engine/inference_engine_s.f90 b/src/inference_engine/inference_engine_s.f90
index 17c4df927..61ce20566 100644
--- a/src/inference_engine/inference_engine_s.f90
+++ b/src/inference_engine/inference_engine_s.f90
@@ -26,7 +26,7 @@
     exchange%activation_strategy_ = self%activation_strategy_ 
   end procedure
 
-  module procedure infer
+  module procedure single_infer
 
     real(rkind), allocatable :: a(:,:)
     integer, parameter :: input_layer = 0
@@ -53,6 +53,12 @@
 
   end procedure
 
+  module procedure batch_infer
+
+    ! here is where the fun starts :) 
+    
+  end procedure
+
   pure subroutine inference_engine_consistency(self)
 
     type(inference_engine_t), intent(in) :: self