Tensor layout options

tools/inference_engine/include/inference_engine_tensor.hpp

@@ -6,20 +6,26 @@
 
 namespace inference_engine
 {
+    enum class Layout
+    {
+        NCHW,
+        NHWC,        
+        UNKNOWN
+    };
+
     struct Tensor
     {
         inference_engine_data_type_t data_type = INFERENCE_ENGINE_DATA_TYPE_UNKNOWN;
-        std::vector<std::uint64_t> dims;
-        std::vector<std::uint64_t> strides;
+        std::vector<std::uint64_t> dims;    // specified in NCHW layout as DirectML
+        std::vector<std::uint64_t> strides; // specified in NCHW layout as DirectML
 
         Tensor() = default;
-        Tensor(inference_engine_data_type_t dt, std::vector<std::uint64_t>&& dimensions)
+        Tensor(inference_engine_data_type_t dt, std::vector<std::uint64_t>&& dimensions, Layout layout = Layout::NCHW)
             : data_type(dt)
             , dims(std::move(dimensions))
         {
-            // ToDo: Decide if should we allow for default strides?
-            // if Yes: NCHW strides calculation should be default
             strides.resize(dims.size());
+            calculate_strides(layout);
         }
         Tensor(const inference_engine_tensor_t& tensor_desc)
             : data_type(tensor_desc.data_type)
@@ -31,6 +37,11 @@ namespace inference_engine
             }
         }
 
+        bool operator==(const Tensor& other) const
+        {
+            return data_type == other.data_type && dims == other.dims && strides == other.strides;
+        }
+
         operator inference_engine_tensor_t() const
         {
             inference_engine_tensor_t ret{};
@@ -61,5 +72,35 @@ namespace inference_engine
             }
             return 1;
         }
+
+    protected:
+        void calculate_strides(Layout layout)
+        {
+            if (layout == Layout::NCHW && dims.size() > 1)
+            {
+                strides[dims.size() - 1] = 1;
+                for (size_t i = dims.size() - 2; i < dims.size(); --i)
+                    strides[i] = strides[i + 1] * dims[i + 1];
+            }
+            else if (layout == Layout::NHWC && dims.size() == 4)
+            {
+                // dims are also given in NCHW order, as DML
+                strides[0] = dims[1] * dims[2] * dims[3]; // n stride = H * W * C
+                strides[1] = 1;                           // c stride = 1
+                strides[2] = dims[3] * dims[1];           // h stride = W * C
+                strides[3] = dims[1];                     // w stride = C
+            }
+            else
+            {
+                assert(!"unsupported layout");
+            }
+        }
     };
+
+    struct IdToTensor
+    {
+        std::size_t id;
+        Tensor tensor;
+    };
+
 }  // namespace inference_engine

tools/inference_engine/src/impl/model.cpp

@@ -14,9 +14,9 @@ namespace inference_engine
 {
     namespace
     {
-        std::vector<TensorMapping> build_output_mapping(const std::vector<std::unique_ptr<GpuNode>>& nodes)
+        std::vector<IdToTensor> build_output_mapping(const std::vector<std::unique_ptr<GpuNode>>& nodes)
         {
-            std::vector<TensorMapping> ret{};
+            std::vector<IdToTensor> ret{};
             for (const auto& n : nodes)
             {
                 if (n->get_outputs().empty())
@@ -34,7 +34,7 @@ namespace inference_engine
         nodes_.reserve(1024);
     }
 
-    std::vector<std::unique_ptr<inference_engine::GpuNode>> DAG::compile(std::span<TensorMapping> input_mappings)
+    std::vector<std::unique_ptr<inference_engine::GpuNode>> DAG::compile(std::span<IdToTensor> input_mappings)
     {
         create_adjacency_list();
 
@@ -62,7 +62,7 @@ namespace inference_engine
 
             // maybe we can move it to a separate function?
             // set input tensor if this is port (important: we have topological sorted, so we assume here that all inputs are traversed first)!
-            auto it = std::find_if(std::begin(input_mappings), std::end(input_mappings), [&](const TensorMapping& im)
+            auto it = std::find_if(std::begin(input_mappings), std::end(input_mappings), [&](const IdToTensor& im)
                 {
                     return im.id == ret[i]->get_id();
                 });
@@ -172,7 +172,7 @@ void ExecutableModel::set_resource(inference_engine_node_id_t id, GpuResource::P
     }
     }
 
-const std::vector<inference_engine::TensorMapping>& ExecutableModel::get_outputs() const
+const std::vector<inference_engine::IdToTensor>& ExecutableModel::get_outputs() const
 {
     return output_mappings_;
 }
@@ -268,7 +268,7 @@ class FusionVisitor : public GpuVisitor
     }
 };
 
-inference_engine::ExecutableModel ModelDescriptor::compile(GpuContext& ctx, GpuStream& stream, std::span<TensorMapping> input_mappings)
+inference_engine::ExecutableModel ModelDescriptor::compile(GpuContext& ctx, GpuStream& stream, std::span<IdToTensor> input_mappings)
 {
     //ToDo: we need some data structure to represent graph (random order of example features below)
     // 1) Sorting graph

tools/inference_engine/src/impl/model.h

@@ -1,8 +1,7 @@
 #pragma once
 #include "gpu_context.h"
-#include "inference_engine_tensor.h"
 #include "inference_engine_operators.h"
-#include "tensor.h"
+#include "inference_engine_tensor.hpp"
 #include "node.h"
 
 #include <vector>
@@ -49,7 +48,7 @@ class DAG
     }
 
 public:
-    std::vector<std::unique_ptr<GpuNode>> compile(std::span<TensorMapping> input_mappings);
+    std::vector<std::unique_ptr<GpuNode>> compile(std::span<IdToTensor> input_mappings);
 
 
 private:
@@ -73,11 +72,11 @@ struct ExecutableModel
     void execute(GpuStream& stream);
     void set_resource(inference_engine_node_id_t id, GpuResource::Ptr rsc);
 
-    const std::vector<TensorMapping>& get_outputs() const;
+    const std::vector<IdToTensor>& get_outputs() const;
 
 private:
     const std::vector<std::unique_ptr<GpuNode>> nodes_;
-    const std::vector<TensorMapping> output_mappings_;
+    const std::vector<IdToTensor> output_mappings_;
 };
 
 class ModelDescriptor
@@ -92,7 +91,7 @@ class ModelDescriptor
         return dag_.add_node<T>(desc, name);
     }
 
-    ExecutableModel compile(GpuContext& ctx, GpuStream& stream, std::span<TensorMapping> input_mappings);
+    ExecutableModel compile(GpuContext& ctx, GpuStream& stream, std::span<IdToTensor> input_mappings);
 
 private:
     DAG dag_;

tools/inference_engine/src/impl/node.h

@@ -1,7 +1,7 @@
 #pragma once
 #include "inference_engine_operators.h"
 
-#include "tensor.h"
+#include "inference_engine_tensor.hpp"
 #include "gpu_context.h"
 #include <string>
 #include <variant>

tools/inference_engine/src/inference_engine_model.cpp

@@ -39,7 +39,7 @@ INFERENCE_ENGINE_API inference_engine_model_t inferenceEngineCompileModelDescrip
             std::cout << "Wrong param input_mapping_list is nullptr or input_mapping_size is 0 " << std::endl;
             return nullptr;
         }
-        std::vector<inference_engine::TensorMapping> im{};
+        std::vector<inference_engine::IdToTensor> im{};
         for (auto i = 0; i < input_mapping_size; i++)
         {
             im.push_back({ input_mapping_list[i].id, inference_engine::Tensor(input_mapping_list[i].tensor) });