Merge remote-tracking branch 'origin/develop' into rewrite-fast-gelu

src/onnx/parse_convolution.cpp

@@ -47,15 +47,17 @@ struct parse_convolution : op_parser<parse_convolution>
                           onnx_parser::node_info info,
                           std::vector<instruction_ref> args) const
     {
-        auto op      = make_op(opd.op_name);
-        auto values  = op.to_value();
-        auto l0      = args[0];
-        auto weights = args[1];
-        auto in_lens = l0->get_shape().lens();
+        auto op       = make_op(opd.op_name);
+        auto values   = op.to_value();
+        auto l0       = args[0];
+        auto weights  = args[1];
+        auto l0_shape = l0->get_shape();
+        auto w_shape  = weights->get_shape();
+        auto in_lens  = l0_shape.max_lens();
         assert(in_lens.size() > 2);
         auto kdims = in_lens.size() - 2;
 
-        // ensure pads availabe only when auto_pad is "NOT_SET"
+        // ensure pads available only when auto_pad is "NOT_SET"
         check_padding_mode(info, "CONV");
 
         if(contains(info.attributes, "strides"))
@@ -79,21 +81,65 @@ struct parse_convolution : op_parser<parse_convolution>
             copy(info.attributes["pads"].ints(), std::back_inserter(padding));
             check_attr_sizes(kdims, padding.size() / 2, "PARSE_CONV: inconsistent paddings");
         }
-
         if(contains(info.attributes, "auto_pad"))
         {
-            auto weight_lens = weights->get_shape().lens();
-            std::vector<std::size_t> k_lens(weight_lens.begin() + 2, weight_lens.end());
-            cal_auto_padding_size(info,
-                                  values,
-                                  k_lens,
-                                  values["dilation"].to_vector<std::size_t>(),
-                                  in_lens,
-                                  padding);
-            auto auto_pad = info.attributes["auto_pad"].s();
+            bool is_same_padding = false;
+            auto auto_pad        = info.attributes["auto_pad"].s();
             if(auto_pad.find("SAME") != std::string::npos)
             {
-                values["padding_mode"] = to_value(op::padding_mode_t::same);
+                is_same_padding = true;
+            }
+
+            // check if image shape is dynamic
+            bool image_shape_dynamic = false;
+            if(l0_shape.dynamic())
+            {
+                auto dyn_dims = l0_shape.dyn_dims();
+                std::for_each(dyn_dims.begin() + 2, dyn_dims.end(), [&](auto dyn_dim) {
+                    if(not dyn_dim.is_fixed())
+                    {
+                        image_shape_dynamic = true;
+                    }
+                });
+            }
+
+            // check if kernel shape is dynamic
+            bool kernel_shape_dynamic = false;
+            if(w_shape.dynamic())
+            {
+                auto dyn_dims = w_shape.dyn_dims();
+                std::for_each(dyn_dims.begin() + 2, dyn_dims.end(), [&](auto dyn_dim) {
+                    if(not dyn_dim.is_fixed())
+                    {
+                        kernel_shape_dynamic = true;
+                    }
+                });
+            }
+
+            if(is_same_padding)
+            {
+                if(image_shape_dynamic or kernel_shape_dynamic)
+                {
+                    // must calculate "same" padding with input shape data
+                    bool is_same_upper     = (auto_pad.find("SAME_UPPER") != std::string::npos);
+                    values["padding_mode"] = is_same_upper
+                                                 ? to_value(op::padding_mode_t::same_upper)
+                                                 : to_value(op::padding_mode_t::same_lower);
+                    values["use_dynamic_same_auto_pad"] = true;
+                }
+                else
+                {
+                    values["padding_mode"] = to_value(op::padding_mode_t::same);
+                    // kernel shape will be fixed, so max_lens() == min_len() for kernel lengths
+                    auto weight_lens = weights->get_shape().max_lens();
+                    std::vector<std::size_t> k_lens(weight_lens.begin() + 2, weight_lens.end());
+                    cal_auto_padding_size(info,
+                                          values,
+                                          k_lens,
+                                          values["dilation"].to_vector<std::size_t>(),
+                                          in_lens,
+                                          padding);
+                }
             }
         }
         values["padding"] = std::vector<size_t>(padding.begin(), padding.end());

src/onnx/parse_if.cpp

@@ -47,7 +47,8 @@ struct parse_if : op_parser<parse_if>
 
         if(args.front()->get_shape().elements() != 1)
         {
-            MIGRAPHX_THROW("PARSE_IF: condition input can have only one element!");
+            MIGRAPHX_THROW("PARSE_IF: " + info.name +
+                           " condition input can have only one element!");
         }
 
         std::string then_name = info.name + "_if";
@@ -69,7 +70,8 @@ struct parse_if : op_parser<parse_if>
                           else_out_shapes.begin(),
                           else_out_shapes.end()))
         {
-            MIGRAPHX_THROW("PARSE_IF: then and else sub_grahps must have same output shapes!");
+            MIGRAPHX_THROW("PARSE_IF: " + info.name +
+                           " then and else sub_grahps must have same output shapes!");
         }
 
         auto if_ret = info.add_instruction(make_op("if"), args, {then_mdl, else_mdl});

src/onnx/parse_instancenorm.cpp

@@ -32,9 +32,12 @@ namespace onnx {
 
 struct parse_instancenorm : op_parser<parse_instancenorm>
 {
+    const std::set<shape::type_t> valid_types = {
+        shape::float_type, shape::half_type, shape::double_type};
+
     std::vector<op_desc> operators() const { return {{"InstanceNormalization"}}; }
 
-    instruction_ref parse(const op_desc& /*opd*/,
+    instruction_ref parse(const op_desc& opd,
                           const onnx_parser& parser,
                           onnx_parser::node_info info,
                           std::vector<instruction_ref> args) const
@@ -52,6 +55,11 @@ struct parse_instancenorm : op_parser<parse_instancenorm>
         auto scale = args[1];
         auto bias  = args[2];
         auto dims  = x->get_shape().lens();
+        auto dtype = x->get_shape().type();
+        if(not contains(valid_types, dtype))
+            MIGRAPHX_THROW(opd.op_name + ": invalid output type: " + std::to_string(dtype) +
+                           ". Valid types are 1 (float), 10 (half), and 11 (double).");
+
         auto ndims = dims.size();
         assert(ndims >= 2);
         auto kdims = ndims - 2;
@@ -65,7 +73,7 @@ struct parse_instancenorm : op_parser<parse_instancenorm>
         auto l0              = info.add_instruction(make_op("sqdiff"), x, mean_bcast);
         auto variance        = info.add_instruction(make_op("reduce_mean", {{"axes", axes}}), l0);
         auto l1              = info.add_instruction(make_op("sub"), x, mean_bcast);
-        auto epsilon_literal = info.add_literal(epsilon);
+        auto epsilon_literal = info.add_literal(literal{shape{dtype}, {epsilon}});
         auto epsilon_bcast =
             info.add_instruction(make_op("multibroadcast", {{"out_lens", dims}}), epsilon_literal);
         auto variance_bcast =

src/onnx/parse_mod.cpp

+/*
+ * The MIT License (MIT)
+ *
+ * Copyright (c) 2015-2022 Advanced Micro Devices, Inc. All rights reserved.
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a copy
+ * of this software and associated documentation files (the "Software"), to deal
+ * in the Software without restriction, including without limitation the rights
+ * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+ * copies of the Software, and to permit persons to whom the Software is
+ * furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in
+ * all copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL THE
+ * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+ * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+ * THE SOFTWARE.
+ */
+#include <migraphx/onnx/op_parser.hpp>
+#include <migraphx/ranges.hpp>
+#include <migraphx/instruction.hpp>
+#include <migraphx/make_op.hpp>
+
+namespace migraphx {
+inline namespace MIGRAPHX_INLINE_NS {
+namespace onnx {
+
+struct parse_mod : op_parser<parse_mod>
+{
+    std::vector<op_desc> operators() const { return {{"Mod"}}; }
+
+    instruction_ref parse(const op_desc& /*opd*/,
+                          const onnx_parser& parser,
+                          onnx_parser::node_info info,
+                          std::vector<instruction_ref> args) const
+    {
+        std::string mod = "mod";
+        if(is_type_float(args[0]->get_shape().type()) || is_type_float(args[1]->get_shape().type()))
+        {
+            if(!contains(info.attributes, "fmod"))
+            {
+                MIGRAPHX_THROW("Mod operator with float args and fmod=0 invalid");
+            }
+        }
+
+        if(contains(info.attributes, "fmod"))
+        {
+            if(parser.parse_value(info.attributes.at("fmod")).at<int>() == 1)
+            {
+                mod = "fmod";
+            }
+        }
+        return info.add_common_op(mod, args[0], args[1]);
+    }
+};
+
+} // namespace onnx
+} // namespace MIGRAPHX_INLINE_NS
+} // namespace migraphx

src/pad_calc.cpp

+/*
+ * The MIT License (MIT)
+ *
+ * Copyright (c) 2015-2022 Advanced Micro Devices, Inc. All rights reserved.
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a copy
+ * of this software and associated documentation files (the "Software"), to deal
+ * in the Software without restriction, including without limitation the rights
+ * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+ * copies of the Software, and to permit persons to whom the Software is
+ * furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in
+ * all copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL THE
+ * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+ * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+ * THE SOFTWARE.
+ */
+#include <migraphx/pad_calc.hpp>
+
+namespace migraphx {
+inline namespace MIGRAPHX_INLINE_NS {
+
+void calculate_padding(int64_t idx,
+                       std::vector<int64_t>& pads,
+                       int64_t input_dim,
+                       int64_t stride,
+                       int64_t dilation,
+                       int64_t weight_dim,
+                       bool is_same_upper)
+{
+    int64_t output_dim     = (input_dim + stride - 1) / stride; // round up result
+    int64_t new_weight_dim = weight_dim + (weight_dim - 1) * (dilation - 1);
+    int64_t pad =
+        std::max(static_cast<int64_t>(0), (output_dim - 1) * stride + new_weight_dim - input_dim);
+    auto pad_ndims = pads.size() / 2;
+
+    if(is_same_upper)
+    {
+        pads[idx]             = pad / 2;
+        pads[idx + pad_ndims] = pad - pad / 2;
+    }
+    else
+    {
+        pads[idx + pad_ndims] = pad / 2;
+        pads[idx]             = pad - pad / 2;
+    }
+}
+
+std::vector<std::size_t> calc_dyn_auto_pad(std::vector<std::size_t> tensor_lens,
+                                           std::vector<std::size_t> k_lens,
+                                           std::vector<std::size_t> strides,
+                                           std::vector<std::size_t> dilations,
+                                           bool use_upper)
+{
+    std::vector<std::size_t> padding;
+    padding.resize(2 * k_lens.size());
+    for(size_t i = 0; i < padding.size() / 2; i++)
+    {
+        std::ptrdiff_t input_dim      = tensor_lens[i];
+        std::ptrdiff_t stride         = strides[i];
+        std::ptrdiff_t weight_dim     = k_lens[i];
+        std::ptrdiff_t dilation       = dilations[i];
+        std::ptrdiff_t output_dim     = (input_dim + stride - 1) / stride; // round up result
+        std::ptrdiff_t new_weight_dim = weight_dim + (weight_dim - 1) * (dilation - 1);
+        std::size_t pad               = std::max(static_cast<std::ptrdiff_t>(0),
+                                   (output_dim - 1) * stride + new_weight_dim - input_dim);
+        auto pad_ndims                = padding.size() / 2;
+
+        if(use_upper)
+        {
+            padding[i]             = pad / 2;
+            padding[i + pad_ndims] = pad - pad / 2;
+        }
+        else
+        {
+            padding[i + pad_ndims] = pad / 2;
+            padding[i]             = pad - pad / 2;
+        }
+    }
+    return padding;
+}
+
+} // namespace MIGRAPHX_INLINE_NS
+} // namespace migraphx

src/program.cpp

@@ -307,9 +307,12 @@ std::vector<argument> generic_eval(const module* mod,
                     if(not contains(params, param_name))
                         MIGRAPHX_THROW("Parameter not found: " + param_name);
                     auto param = params[param_name];
-                    if(param.get_shape() != ins->get_shape())
+                    // TODO: may want to check correct number of dimensions and/or was within bounds
+                    if(not ins->get_shape().dynamic() and param.get_shape() != ins->get_shape())
+                    {
                         MIGRAPHX_THROW("Incorrect shape {" + to_string(param.get_shape()) +
                                        "} for parameter: " + param_name);
+                    }
                     return param;
                 }));
         }
@@ -352,7 +355,10 @@ std::vector<argument> generic_eval(const module* mod,
                             }));
         }
         assert(results.find(ins) != results.end());
-        assert(results.at(ins).get_shape() == ins->get_shape());
+        if(not ins->get_shape().dynamic())
+        {
+            assert(results.at(ins).get_shape() == ins->get_shape());
+        }
     }
     return {results.at(std::prev(mod->end()))};
 }

src/sqlite.cpp

+/*
+ * The MIT License (MIT)
+ *
+ * Copyright (c) 2015-2022 Advanced Micro Devices, Inc. All rights reserved.
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a copy
+ * of this software and associated documentation files (the "Software"), to deal
+ * in the Software without restriction, including without limitation the rights
+ * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+ * copies of the Software, and to permit persons to whom the Software is
+ * furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in
+ * all copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL THE
+ * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+ * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+ * THE SOFTWARE.
+ */
+
+#include <migraphx/sqlite.hpp>
+#include <migraphx/manage_ptr.hpp>
+#include <migraphx/errors.hpp>
+#include <sqlite3.h>
+#include <algorithm>
+
+namespace migraphx {
+inline namespace MIGRAPHX_INLINE_NS {
+
+using sqlite3_ptr = MIGRAPHX_MANAGE_PTR(sqlite3*, sqlite3_close);
+
+struct sqlite_impl
+{
+    sqlite3* get() const { return ptr.get(); }
+    void open(const fs::path& p, int flags)
+    {
+        sqlite3* ptr_tmp = nullptr;
+        int rc           = sqlite3_open_v2(p.string().c_str(), &ptr_tmp, flags, nullptr);
+        ptr              = sqlite3_ptr{ptr_tmp};
+        if(rc != 0)
+            MIGRAPHX_THROW("error opening " + p.string() + ": " + error_message());
+    }
+
+    template <class F>
+    void exec(const char* sql, F f)
+    {
+        auto callback = [](void* obj, auto... xs) -> int {
+            try
+            {
+                const auto* g = static_cast<const F*>(obj);
+                (*g)(xs...);
+                return 0;
+            }
+            catch(...)
+            {
+                return -1;
+            }
+        };
+        int rc = sqlite3_exec(get(), sql, callback, &f, nullptr);
+        if(rc != 0)
+            MIGRAPHX_THROW(error_message());
+    }
+
+    std::string error_message() const
+    {
+        std::string msg = "sqlite3: ";
+        return msg + sqlite3_errmsg(get());
+    }
+    sqlite3_ptr ptr;
+};
+
+sqlite sqlite::read(const fs::path& p)
+{
+    sqlite r;
+    r.impl = std::make_shared<sqlite_impl>();
+    r.impl->open(p, SQLITE_OPEN_READONLY);
+    return r;
+}
+
+sqlite sqlite::write(const fs::path& p)
+{
+    sqlite r;
+    r.impl = std::make_shared<sqlite_impl>();
+    // Using '+' instead of bitwise '|' to avoid compilation warning
+    r.impl->open(p, SQLITE_OPEN_READWRITE + SQLITE_OPEN_CREATE);
+    return r;
+}
+
+std::vector<std::unordered_map<std::string, std::string>> sqlite::execute(const std::string& s)
+{
+    std::vector<std::unordered_map<std::string, std::string>> result;
+    impl->exec(s.c_str(), [&](int n, char** texts, char** names) {
+        std::unordered_map<std::string, std::string> row;
+        row.reserve(n);
+        std::transform(
+            names,
+            names + n,
+            texts,
+            std::inserter(row, row.begin()),
+            [&](const char* name, const char* text) { return std::make_pair(name, text); });
+        result.push_back(row);
+    });
+    return result;
+}
+
+} // namespace MIGRAPHX_INLINE_NS
+} // namespace migraphx

src/targets/fpga/CMakeLists.txt

+#####################################################################################
+# The MIT License (MIT)
+#
+# Copyright (c) 2015-2022 Advanced Micro Devices, Inc. All rights reserved.
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+# THE SOFTWARE.
+#####################################################################################
+
+add_library(migraphx_fpga
+    target.cpp
+    lowering.cpp
+    subgraph.cpp
+    vitis_ai_adapter.cpp
+)
+
+set_target_properties(migraphx_fpga PROPERTIES EXPORT_NAME fpga)
+rocm_set_soversion(migraphx_fpga ${MIGRAPHX_SO_VERSION})
+
+rocm_clang_tidy_check(migraphx_fpga)
+target_link_libraries(migraphx_fpga migraphx)
+
+rocm_install_targets(
+    TARGETS migraphx_fpga
+    INCLUDE
+    ${CMAKE_CURRENT_SOURCE_DIR}/include
+)

src/targets/fpga/include/migraphx/fpga/context.hpp

+/*
+ * The MIT License (MIT)
+ *
+ * Copyright (c) 2015-2022 Advanced Micro Devices, Inc. All rights reserved.
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a copy
+ * of this software and associated documentation files (the "Software"), to deal
+ * in the Software without restriction, including without limitation the rights
+ * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+ * copies of the Software, and to permit persons to whom the Software is
+ * furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in
+ * all copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL THE
+ * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+ * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+ * THE SOFTWARE.
+ */
+
+#ifndef MIGRAPHX_GUARD_FPGA_CONTEXT_HPP
+#define MIGRAPHX_GUARD_FPGA_CONTEXT_HPP
+
+#include <migraphx/config.hpp>
+
+namespace migraphx {
+inline namespace MIGRAPHX_INLINE_NS {
+namespace fpga {
+
+struct context
+{
+    int id = 0;
+
+    void finish() const {}
+};
+
+} // namespace fpga
+} // namespace MIGRAPHX_INLINE_NS
+} // namespace migraphx
+
+#endif // MIGRAPHX_GUARD_FPGA_CONTEXT_HPP

src/targets/fpga/include/migraphx/fpga/lowering.hpp

+/*
+ * The MIT License (MIT)
+ *
+ * Copyright (c) 2015-2022 Advanced Micro Devices, Inc. All rights reserved.
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a copy
+ * of this software and associated documentation files (the "Software"), to deal
+ * in the Software without restriction, including without limitation the rights
+ * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+ * copies of the Software, and to permit persons to whom the Software is
+ * furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in
+ * all copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL THE
+ * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+ * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+ * THE SOFTWARE.
+ */
+
+#ifndef MIGRAPHX_GUARD_FPGA_LOWERING_HPP
+#define MIGRAPHX_GUARD_FPGA_LOWERING_HPP
+
+#include <migraphx/program.hpp>
+#include <migraphx/config.hpp>
+#include <migraphx/fpga/context.hpp>
+
+namespace migraphx {
+inline namespace MIGRAPHX_INLINE_NS {
+namespace fpga {
+
+struct lowering
+{
+    context* ctx = nullptr;
+    std::string name() const { return "fpga::lowering"; }
+    void apply(module& m) const;
+};
+
+} // namespace fpga
+} // namespace MIGRAPHX_INLINE_NS
+} // namespace migraphx
+
+#endif // MIGRAPHX_GUARD_FPGA_LOWERING_HPP

src/targets/fpga/include/migraphx/fpga/subgraph.hpp

+/*
+ * The MIT License (MIT)
+ *
+ * Copyright (c) 2015-2022 Advanced Micro Devices, Inc. All rights reserved.
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a copy
+ * of this software and associated documentation files (the "Software"), to deal
+ * in the Software without restriction, including without limitation the rights
+ * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+ * copies of the Software, and to permit persons to whom the Software is
+ * furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in
+ * all copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL THE
+ * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+ * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+ * THE SOFTWARE.
+ */
+
+#ifndef MIGRAPHX_GUARD_FPGA_SUBGRAPH_HPP
+#define MIGRAPHX_GUARD_FPGA_SUBGRAPH_HPP
+
+#include <migraphx/program.hpp>
+#include <migraphx/config.hpp>
+
+namespace migraphx {
+inline namespace MIGRAPHX_INLINE_NS {
+namespace fpga {
+
+struct subgraph
+{
+    std::string name() const { return "fpga::subgraph"; }
+    void apply(module_pass_manager& mpm) const;
+};
+
+} // namespace fpga
+} // namespace MIGRAPHX_INLINE_NS
+} // namespace migraphx
+
+#endif // MIGRAPHX_GUARD_FPGA_SUBGRAPH_HPP

src/targets/fpga/include/migraphx/fpga/target.hpp

+/*
+ * The MIT License (MIT)
+ *
+ * Copyright (c) 2015-2022 Advanced Micro Devices, Inc. All rights reserved.
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a copy
+ * of this software and associated documentation files (the "Software"), to deal
+ * in the Software without restriction, including without limitation the rights
+ * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+ * copies of the Software, and to permit persons to whom the Software is
+ * furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in
+ * all copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL THE
+ * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+ * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+ * THE SOFTWARE.
+ */
+
+#ifndef MIGRAPHX_GUARD_FPGA_TARGET_HPP
+#define MIGRAPHX_GUARD_FPGA_TARGET_HPP
+
+#include <migraphx/program.hpp>
+#include <migraphx/register_target.hpp>
+#include <migraphx/compile_options.hpp>
+#include <migraphx/fpga/context.hpp>
+#include <migraphx/config.hpp>
+
+namespace migraphx {
+inline namespace MIGRAPHX_INLINE_NS {
+struct pass;
+namespace fpga {
+
+struct target
+{
+    std::string name() const;
+    std::vector<pass> get_passes(migraphx::context& ctx, const compile_options&) const;
+    migraphx::context get_context() const { return context{}; }
+    float is_supported(instruction_ref ins, support_metric m);
+
+    argument copy_to(const argument& arg) const { return arg; }
+    argument copy_from(const argument& arg) const { return arg; }
+    argument allocate(const shape& s) const;
+};
+
+MIGRAPHX_REGISTER_TARGET(target);
+
+} // namespace fpga
+} // namespace MIGRAPHX_INLINE_NS
+} // namespace migraphx
+
+#endif // MIGRAPHX_GUARD_FPGA_TARGET_HPP

src/targets/fpga/include/migraphx/fpga/vitis_ai_adapter.hpp

+/*
+ * The MIT License (MIT)
+ *
+ * Copyright (c) 2015-2022 Advanced Micro Devices, Inc. All rights reserved.
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a copy
+ * of this software and associated documentation files (the "Software"), to deal
+ * in the Software without restriction, including without limitation the rights
+ * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+ * copies of the Software, and to permit persons to whom the Software is
+ * furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in
+ * all copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL THE
+ * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+ * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+ * THE SOFTWARE.
+ */
+
+#ifndef MIGRAPHX_GUARD_FPGA_VITIS_AI_ADAPTER_HPP
+#define MIGRAPHX_GUARD_FPGA_VITIS_AI_ADAPTER_HPP
+
+#include <string>
+
+#include <migraphx/instruction.hpp>
+#include <migraphx/pass_manager.hpp>
+
+namespace vitis_ai {
+
+class x_model
+{
+    migraphx::shape shape;
+
+    public:
+    migraphx::shape get_shape() const;
+    void set_shape(migraphx::shape);
+};
+
+x_model create_xmodel(migraphx::module_ref mod);
+
+migraphx::argument execute(const x_model& xmodel,
+                           const migraphx::shape& output_shape,
+                           std::vector<migraphx::argument>& args);
+
+} // namespace vitis_ai
+
+#endif // MIGRAPHX_GUARD_FPGA_VITIS_AI_ADAPTER_HPP

src/targets/fpga/lowering.cpp

+/*
+ * The MIT License (MIT)
+ *
+ * Copyright (c) 2015-2022 Advanced Micro Devices, Inc. All rights reserved.
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a copy
+ * of this software and associated documentation files (the "Software"), to deal
+ * in the Software without restriction, including without limitation the rights
+ * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+ * copies of the Software, and to permit persons to whom the Software is
+ * furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in
+ * all copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL THE
+ * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+ * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+ * THE SOFTWARE.
+ */
+
+#include <migraphx/fpga/lowering.hpp>
+#include <migraphx/instruction.hpp>
+#include <migraphx/iterator_for.hpp>
+#include <migraphx/register_op.hpp>
+#include <migraphx/stringutils.hpp>
+#include <iostream>
+
+#include "migraphx/fpga/vitis_ai_adapter.hpp"
+
+namespace migraphx {
+inline namespace MIGRAPHX_INLINE_NS {
+
+namespace fpga {
+
+struct fpga_vitis_op
+{
+    fpga_vitis_op() = default;
+    explicit fpga_vitis_op(vitis_ai::x_model model) : xmodel(std::move(model)){};
+
+    vitis_ai::x_model xmodel;
+    int dummy = 0;
+
+    template <class Self, class F>
+    static auto reflect(Self& self, F f)
+    {
+        // return pack(f(self.xmodel, "xmodel"));
+        return pack(f(self.dummy, "dummy"));
+    }
+
+    std::string name() const { return "fpga::vitis_ai"; }
+
+    shape compute_shape(const std::vector<shape>& inputs) const
+    {
+        (void)inputs;
+        return xmodel.get_shape();
+    }
+
+    argument
+    compute(const context& ctx, const shape& output_shape, std::vector<argument> args) const
+    {
+        std::cout << "The context is " << ctx.id << std::endl;
+        return ::vitis_ai::execute(xmodel, output_shape, args);
+    }
+};
+MIGRAPHX_REGISTER_OP(fpga_vitis_op)
+
+void lowering::apply(module& m) const
+{
+    auto* mod = &m;
+
+    // test modifying the context from a pass
+    ctx->id = 2;
+
+    for(auto it : iterator_for(*mod))
+    {
+        if(it->name() == "fpga::vitis_placeholder")
+        {
+            assert(it->module_inputs().size() == 1);
+            auto xmodel = ::vitis_ai::create_xmodel(it->module_inputs()[0]);
+            mod->replace_instruction(it, fpga_vitis_op{xmodel}, it->inputs());
+        }
+    }
+}
+
+} // namespace fpga
+} // namespace MIGRAPHX_INLINE_NS
+} // namespace migraphx

src/targets/fpga/subgraph.cpp

+/*
+ * The MIT License (MIT)
+ *
+ * Copyright (c) 2015-2022 Advanced Micro Devices, Inc. All rights reserved.
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a copy
+ * of this software and associated documentation files (the "Software"), to deal
+ * in the Software without restriction, including without limitation the rights
+ * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+ * copies of the Software, and to permit persons to whom the Software is
+ * furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in
+ * all copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL THE
+ * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+ * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+ * THE SOFTWARE.
+ */
+
+#include <migraphx/fpga/subgraph.hpp>
+
+#include <migraphx/instruction.hpp>
+#include "migraphx/iterator.hpp"
+#include <migraphx/iterator_for.hpp>
+#include "migraphx/make_op.hpp"
+#include "migraphx/module.hpp"
+#include "migraphx/ranges.hpp"
+#include <migraphx/register_op.hpp>
+#include <migraphx/stringutils.hpp>
+#include <migraphx/pass_manager.hpp>
+
+namespace migraphx {
+inline namespace MIGRAPHX_INLINE_NS {
+
+namespace fpga {
+
+struct fpga_placeholder_op
+{
+    fpga_placeholder_op() = default;
+
+    int dummy = 0;
+
+    template <class Self, class F>
+    static auto reflect(Self& self, F f)
+    {
+        return pack(f(self.dummy, "dummy"));
+    }
+
+    std::string name() const { return "fpga::vitis_placeholder"; }
+
+    shape compute_shape(const std::vector<shape>& inputs, std::vector<module_ref> mods) const
+    {
+        (void)inputs;
+        if(mods.size() != 1)
+        {
+            MIGRAPHX_THROW("should have one submodule.");
+        }
+        module_ref sm = mods.front();
+        if(sm->get_output_shapes().size() != 1)
+            MIGRAPHX_THROW("Only one return");
+        return sm->get_output_shapes().front();
+    }
+};
+MIGRAPHX_REGISTER_OP(fpga_placeholder_op)
+
+bool is_fpga_instr(migraphx::instruction_ref it)
+{
+    // assuming all instructions that aren't @param, @literal, or input data are fpga instrs
+    if(migraphx::starts_with(it->name(), "@"))
+    {
+        return false;
+    }
+    // no inputs to the instr means it's input data
+    if(it->inputs().empty())
+    {
+        return false;
+    }
+    return true;
+}
+
+void subgraph::apply(module_pass_manager& mpm) const
+{
+    auto& mod = mpm.get_module();
+    auto* pm  = mpm.create_module(mod.name() + ":fpga");
+    pm->set_bypass();
+
+    migraphx::instruction_ref first = mod.end();
+    migraphx::instruction_ref last;
+    std::vector<migraphx::instruction_ref> literal_inputs;
+    for(auto it : iterator_for(mod))
+    {
+        // assuming we want all the params/literals as inputs to the FPGA submodule
+        if(migraphx::starts_with(it->name(), "@param") ||
+           migraphx::starts_with(it->name(), "@literal"))
+        {
+            literal_inputs.push_back(it);
+        }
+        if(is_fpga_instr(it))
+        {
+            if(first == mod.end())
+            {
+                first = it;
+            }
+            last = it;
+        }
+    }
+
+    // TODO(varunsh): this code may be replaceable by code in the fuse_pointwise pass
+
+    // assuming all FPGA instructions are in one contiguous range
+    pm->insert_instructions(pm->end(), first, last, {});
+
+    migraphx::instruction_ref placeholder_ins;
+    for(auto it : iterator_for(mod))
+    {
+        if(migraphx::starts_with(it->name(), "@return"))
+        {
+            placeholder_ins = mod.insert_instruction(
+                it, migraphx::make_op("fpga::vitis_placeholder"), literal_inputs, {pm});
+            break;
+        }
+    }
+
+    mod.replace_return({placeholder_ins});
+}
+
+} // namespace fpga
+} // namespace MIGRAPHX_INLINE_NS
+} // namespace migraphx

src/targets/fpga/target.cpp

+/*
+ * The MIT License (MIT)
+ *
+ * Copyright (c) 2015-2022 Advanced Micro Devices, Inc. All rights reserved.
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a copy
+ * of this software and associated documentation files (the "Software"), to deal
+ * in the Software without restriction, including without limitation the rights
+ * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+ * copies of the Software, and to permit persons to whom the Software is
+ * furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in
+ * all copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL THE
+ * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+ * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+ * THE SOFTWARE.
+ */
+
+#include <migraphx/fpga/target.hpp>
+#include <migraphx/fpga/lowering.hpp>
+#include <migraphx/fpga/subgraph.hpp>
+#include <migraphx/register_target.hpp>
+#include <migraphx/pass.hpp>
+#include <migraphx/auto_contiguous.hpp>
+#include <migraphx/rewrite_rnn.hpp>
+#include <migraphx/eliminate_pad.hpp>
+#include <migraphx/insert_pad.hpp>
+#include <migraphx/dead_code_elimination.hpp>
+#include <migraphx/generate.hpp>
+#include <migraphx/normalize_ops.hpp>
+
+namespace migraphx {
+inline namespace MIGRAPHX_INLINE_NS {
+namespace fpga {
+
+std::string target::name() const { return "fpga"; }
+
+std::vector<pass> target::get_passes(migraphx::context& gctx, const compile_options&) const
+{
+    // not sure if all these passes are needed but they were copied from ref/
+    auto& ctx = any_cast<context>(gctx);
+    return {normalize_ops{},
+            eliminate_pad{},
+            dead_code_elimination{},
+            insert_pad{},
+            dead_code_elimination{},
+            rewrite_rnn{},
+            dead_code_elimination{},
+            auto_contiguous{},
+            dead_code_elimination{},
+            subgraph{},
+            dead_code_elimination{},
+            lowering{&ctx},
+            dead_code_elimination{}};
+}
+
+argument target::allocate(const shape& s) const { return fill_argument(s, 0); }
+
+float is_supported(instruction_ref ins, support_metric m)
+{
+    // for now, not using the ins and metric to return a value
+    (void)ins;
+    (void)m;
+    return 1.0;
+}
+
+MIGRAPHX_REGISTER_TARGET(target);
+
+} // namespace fpga
+} // namespace MIGRAPHX_INLINE_NS
+} // namespace migraphx

src/targets/fpga/vitis_ai_adapter.cpp

+/*
+ * The MIT License (MIT)
+ *
+ * Copyright (c) 2015-2022 Advanced Micro Devices, Inc. All rights reserved.
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a copy
+ * of this software and associated documentation files (the "Software"), to deal
+ * in the Software without restriction, including without limitation the rights
+ * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+ * copies of the Software, and to permit persons to whom the Software is
+ * furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in
+ * all copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL THE
+ * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+ * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+ * THE SOFTWARE.
+ */
+
+#include "migraphx/fpga/vitis_ai_adapter.hpp"
+
+#include "migraphx/module.hpp"
+
+#include "migraphx/stringutils.hpp"
+namespace vitis_ai {
+
+migraphx::shape x_model::get_shape() const { return shape; };
+
+void x_model::set_shape(migraphx::shape s) { shape = s; }
+
+x_model create_xmodel(const migraphx::module_ref mod)
+{
+    std::cout << "Calling an external function: create_xmodel!\n";
+    x_model xmodel;
+    xmodel.set_shape(mod->get_output_shapes());
+    return xmodel;
+}
+
+migraphx::argument execute(const x_model& xmodel,
+                           const migraphx::shape& output_shape,
+                           std::vector<migraphx::argument>& args)
+{
+    (void)xmodel;
+
+    std::cout << "Calling an external function: execute!\n";
+
+    std::cout << "Output Shape: " << output_shape << std::endl;
+    std::cout << "Args: " << args.size() << std::endl;
+    for(const auto& arg : args)
+    {
+        std::cout << "  " << arg.get_shape() << std::endl;
+    }
+    std::cout << std::endl;
+
+    migraphx::argument result{output_shape};
+
+    return result;
+}
+
+} // namespace vitis_ai

src/targets/gpu/CMakeLists.txt

@@ -184,6 +184,7 @@ add_library(migraphx_gpu
     pack_int8_args.cpp
     prefuse_ops.cpp
     pad.cpp
+    perfdb.cpp
     pooling.cpp
     quant_convolution.cpp
     reverse.cpp

src/targets/gpu/compile_hip_code_object.cpp

@@ -51,9 +51,9 @@ static const char* const make_tensor_template = R"__migraphx__(
 template<>
 struct make_tensor<${n}>
 {
-    static __device__ auto apply(void* p)
+    static __device__ auto apply(void* __restrict__ p)
     {
-        return make_tensor_view(reinterpret_cast<${type}*>(p), make_shape(${lens}, ${strides}));
+        return make_tensor_view(reinterpret_cast<${type}* __restrict__>(p), make_shape(${lens}, ${strides}));
     }
 };
 )__migraphx__";

src/targets/gpu/hip.cpp

@@ -23,13 +23,13 @@
  */
 
 #include <migraphx/gpu/hip.hpp>
-
 #include <migraphx/manage_ptr.hpp>
 #include <migraphx/register_op.hpp>
 #include <migraphx/gpu/context.hpp>
 #include <migraphx/gpu/device/contiguous.hpp>
 #include <miopen/miopen.h>
-
+#include <memory>
+#include <mutex>
 #include <vector>
 
 namespace migraphx {
@@ -77,12 +77,38 @@ void* get_device_ptr(void* hptr)
     return result;
 }
 
-hip_ptr allocate_gpu(std::size_t sz, bool host = false)
+struct host_ptr_cache
+{
+    std::unordered_map<void*, std::weak_ptr<void>> cache;
+    std::mutex m;
+    std::shared_ptr<void> get(void* ptr)
+    {
+        std::lock_guard<std::mutex> lock(m);
+        auto it = cache.find(ptr);
+        if(it != cache.end())
+            return it->second.lock();
+        return nullptr;
+    }
+
+    void put(const std::shared_ptr<void>& p)
+    {
+        std::lock_guard<std::mutex> lock(m);
+        cache[p.get()] = p;
+    }
+};
+
+static host_ptr_cache& get_host_ptr_cache()
+{
+    static host_ptr_cache cache;
+    return cache;
+}
+
+std::shared_ptr<void> allocate_gpu(std::size_t sz, bool host = false)
 {
     if(sz > get_available_gpu_memory())
         MIGRAPHX_THROW("Memory not available to allocate buffer: " + std::to_string(sz));
-    void* result = nullptr;
-    auto status  = host ? hipHostMalloc(&result, sz) : hipMalloc(&result, sz);
+    void* alloc_ptr = nullptr;
+    auto status     = host ? hipHostMalloc(&alloc_ptr, sz) : hipMalloc(&alloc_ptr, sz);
     if(status != hipSuccess)
     {
         if(host)
@@ -90,16 +116,28 @@ hip_ptr allocate_gpu(std::size_t sz, bool host = false)
         else
             return allocate_gpu(sz, true);
     }
-    assert(result != nullptr);
-    return hip_ptr{result};
+    assert(alloc_ptr != nullptr);
+    std::shared_ptr<void> result = share(hip_ptr{alloc_ptr});
+    if(host)
+    {
+        get_host_ptr_cache().put(result);
+    }
+    return result;
 }
 
-hip_host_ptr register_on_gpu(void* ptr, std::size_t sz)
+std::shared_ptr<void> register_on_gpu(void* ptr, std::size_t sz)
 {
+    std::shared_ptr<void> result = get_host_ptr_cache().get(ptr);
+    if(result)
+    {
+        return result;
+    }
     auto status = hipHostRegister(ptr, sz, hipHostRegisterMapped);
     if(status != hipSuccess)
         MIGRAPHX_THROW("Gpu register failed: " + hip_error(status));
-    return hip_host_ptr{ptr};
+    result = share(hip_host_ptr{ptr});
+    get_host_ptr_cache().put(result);
+    return result;
 }
 
 template <class T>
@@ -115,7 +153,7 @@ std::vector<T> read_from_gpu(const void* x, std::size_t sz)
     return result;
 }
 
-hip_ptr write_to_gpu(const void* x, std::size_t sz, bool host = false)
+std::shared_ptr<void> write_to_gpu(const void* x, std::size_t sz, bool host = false)
 {
     gpu_sync();
     auto result = allocate_gpu(sz, host);
@@ -137,22 +175,21 @@ hip_ptr write_to_gpu(const T& x)
 
 argument allocate_gpu(const shape& s, bool host)
 {
-    auto p = share(allocate_gpu(s.bytes() + 1, host));
+    auto p = allocate_gpu(s.bytes() + 1, host);
     return {s, [p]() mutable { return reinterpret_cast<char*>(p.get()); }};
 }
 
 argument register_on_gpu(const argument& arg)
 {
     auto arg_shared = arg.share();
-    auto p          = share(register_on_gpu(arg_shared.data(), arg_shared.get_shape().bytes()));
-    return {arg_shared.get_shape(), [p, a = std::move(arg_shared)]() mutable {
-                return get_device_ptr(p.get());
-            }}; // namespace gpu
-} // namespace MIGRAPHX_INLINE_NS
+    auto p          = register_on_gpu(arg_shared.data(), arg_shared.get_shape().bytes());
+    return {arg_shared.get_shape(),
+            [p, a = std::move(arg_shared)]() mutable { return get_device_ptr(p.get()); }};
+}
 
 argument to_gpu(const argument& arg, bool host)
 {
-    auto p = share(write_to_gpu(arg.data(), arg.get_shape().bytes(), host));
+    auto p = write_to_gpu(arg.data(), arg.get_shape().bytes(), host);
     return {arg.get_shape(), p};
 }