1 files changed, 45 insertions, 0 deletions
diff --git a/src/armnn/OverrideInputRangeVisitor.hpp b/src/armnn/OverrideInputRangeVisitor.hpp
new file mode 100644
index 0000000000..a2da6c702e
--- /dev/null
+++ b/src/armnn/OverrideInputRangeVisitor.hpp
@@ -0,0 +1,45 @@
+//
+// Copyright © 2017 Arm Ltd. All rights reserved.
+// SPDX-License-Identifier: MIT
+//
+
+#pragma once
+
+#include "NetworkQuantizer.hpp"
+#include "LayerVisitorBase.hpp"
+
+#include <unordered_map>
+
+namespace armnn
+{
+
+/// Visitor object for overriding the input range of the quantized input layers in a network
+class OverrideInputRangeVisitor : public LayerVisitorBase<VisitorNoThrowPolicy>
+{
+private:
+    using MinMaxRange  = std::pair<float, float>;
+    using MinMaxRanges = std::vector<MinMaxRange>;
+
+public:
+    OverrideInputRangeVisitor(std::unordered_map<LayerGuid, MinMaxRanges>& guidToRangesMap,
+                              LayerBindingId layerId,
+                              const MinMaxRange& minMaxRange);
+    ~OverrideInputRangeVisitor() = default;
+
+    void VisitInputLayer(const IConnectableLayer *layer, LayerBindingId id, const char *name = nullptr) override;
+
+private:
+    /// Sets the range for the given input layer
+    void SetRange(const IConnectableLayer* layer);
+
+    /// Mapping from a layer Guid to an array of ranges for outputs
+    std::unordered_map<LayerGuid, MinMaxRanges>& m_GuidToRangesMap;
+
+    /// The id of the input layer of which to override the input range
+    LayerBindingId m_LayerId;
+
+    /// The new input range to be applied to the input layer
+    MinMaxRange m_MinMaxRange;
+};
+
+} // namespace armnn