added compilation of shared library and multiple instances for gemm, cleaned up code design

python/ait_impl/samples/model_container.h

-#pragma once
-
-#include "model-generated.h"
-#include "model_interface.h"
-#include "raii_wrapper.h"
-
-#include <condition_variable>
-#include <cstring>
-#include <future>
-#include <mutex>
-#include <numeric>
-#include <unordered_map>
-
-namespace ait {
-
-// ModelContainer inherits from this class; its implementation is
-// generated at compilation time. Most of the ModelContainer
-// logic does not need codegen; anything that does should be put
-// into this class instead.
-class ModelContainerBase {
- public:
-  ModelContainerBase(
-      size_t num_inputs,
-      size_t num_outputs,
-      size_t num_unbound_constants,
-      size_t params_size,
-      AITemplateAllocator& allocator);
-
- protected:
-  // The set of unbound constants/weights/parameters. These are constants which
-  // have no value at compile time and do not participate in constant folding.
-  // They must be set via SetConstant prior to inference.
-  std::unordered_map<std::string, size_t> unbound_constant_name_to_idx_;
-
-  // a single piece of memory for all constants
-  GPUPtr constants_;
-
-  // size of the containers below: # inputs + # outputs + # unbound constants.
-  size_t num_params_;
-  // These entries correspond to inputs/outputs/unbound constants in order;
-  // inputs first, then outputs, then constants.
-  std::vector<const char*> param_names_;
-  std::vector<std::vector<int64_t>> max_param_shapes_;
-  std::vector<AITemplateDtype> param_dtypes_;
-
-  // NB: technically these could be derived from both the max shape and
-  // the dytpe, but it's easier to just cache them.
-  std::vector<size_t> max_param_storage_bytes_;
-  std::vector<size_t> max_param_numel_;
-};
-
-// This creates a new ModelContainer; its implementation is also
-// codegened (the parameters passed to the ctor are determined
-// at compilation time)
-class ModelContainer;
-ModelContainer* CreateModelContainer(
-    size_t num_runtimes,
-    AITemplateAllocator& allocator);
-
-// Each ModelContainer contains num_models Models. Inference runs
-// can be started by invoking Run() with lists of pre-allocated
-// input/output tensors. GetOutputMaximumShape() can be used to
-// determine how much memory is required for each output.
-//
-// If there are N tensors marked with is_output=True,
-// the user will always be expected to pass N output pointers -
-// extra copies will occur if the outputs are views of constants,
-// inputs, or other outputs in this case to avoid surprises.
-//
-// Use stream = nullptr for default stream. ModelContainer/Model does not
-// create or own any stream. The user is expected to create and manage streams.
-//
-// We can support at most num_models concurrent inferences.
-// Run() takes a stream to run the inference on. For example,
-// to start up two inferences on different streams concurrently,
-// we can do this:
-//
-// model_container.Run(inputs0, num_inputs, outputs0, num_ouputs, stream0, ...);
-// model_container.Run(inputs1, num_inputs, outputs1, num_ouputs, stream1, ...);
-// StreamSynchronize(stream0);
-// StreamSynchronize(stream1);
-//
-// Note that if there are no models available for inference, Run() will block
-// until one becomes available.
-//
-// ModelContainer optionally takes an allocator argument, which it will use to
-// allocate the space for the buffers used for intermediate tensors and
-// constants. If it is nullptr, the default allocator will be used (e.g. just
-// {cuda/hip}{Malloc/Free}).
-// Important: we assume that the allocator lives until the ModelContainer is
-// destroyed. The default allocator has a static lifetime.
-class ModelContainer : ModelContainerBase {
- public:
-  ModelContainer(
-      size_t num_models,
-      size_t blob_size,
-      size_t workspace_size,
-      size_t num_inputs,
-      size_t num_outputs,
-      size_t num_unbound_constants,
-      size_t params_size,
-      AITemplateAllocator& allocator);
-
-  void Run(
-      const AITData* inputs,
-      size_t num_inputs,
-      AITData* outputs,
-      size_t num_outputs,
-      StreamType stream,
-      bool sync,
-      bool graph_mode,
-      int64_t** output_shapes_out);
-
-  void RunWithOutputsOnHost(
-      const AITData* inputs,
-      size_t num_inputs,
-      AITData* outputs,
-      size_t num_outputs,
-      StreamType stream,
-      bool graph_mode,
-      int64_t** output_shapes_out);
-
-  float Benchmark(
-      const AITData* inputs,
-      size_t num_inputs,
-      AITData* outputs,
-      size_t num_outputs,
-      StreamType stream,
-      bool graph_mode,
-      size_t count,
-      size_t num_threads,
-      bool use_unique_stream_per_thread,
-      int64_t** output_shapes_out);
-
-  void SetConstant(const char* name, const AITData& tensor);
-
-  size_t NumInputs() const;
-  size_t NumOutputs() const;
-
-  const char* InputName(size_t input_idx) const;
-  const char* OutputName(size_t output_idx) const;
-
-  AITemplateParamShape MaxOutputShape(size_t output_idx) const;
-  AITemplateDtype OutputDtype(size_t output_idx) const;
-  size_t MaxOutputStorageBytes(size_t output_idx) const;
-
-  size_t GetNumRuntimes() const {
-    return models_.size();
-  }
-
- private:
-  void PrepareForRun(
-      Model* model,
-      const AITData* inputs,
-      size_t num_inputs,
-      AITData* outputs,
-      size_t num_outputs);
-      Model* GetAvailableModel();
-  void ReclaimFinishedModels(std::unique_lock<std::mutex>& lk);
-  void ValidateDtype(AITemplateDtype dtype, size_t idx) const;
-
-  float BenchmarkImpl(
-      const AITData* inputs,
-      size_t num_inputs,
-      AITData* outputs,
-      size_t num_outputs,
-      StreamType stream,
-      bool graph_mode,
-      size_t count,
-      int64_t** output_shapes_out);
-
-  AITemplateAllocator& allocator_;
-
-  std::vector<Model> models_;
-  std::vector<Model*> available_models_;
-  std::deque<Model*> pending_models_;
-
-  // Guards accesses to available/pending models.
-  std::mutex models_mutex_;
-  // Notified whenever a model is put into pending_models_.
-  std::condition_variable pending_models_available_;
-
-  size_t num_inputs_;
-  size_t num_outputs_;
-};
-
-} // namespace ait