integrating auto_cont changes

30c8ff61 · Khalique Ahmed · 4ff8a292 · 7aee6388 · 30c8ff61 · 30c8ff61
Commit 30c8ff61 authored Aug 14, 2023 by Khalique Ahmed
20 changed files
--- a/Jenkinsfile
+++ b/Jenkinsfile
@@ -26,7 +26,7 @@ def rocmtestnode(Map conf) {
            rm -rf build
            mkdir build
            cd build
-            cmake -DCMAKE_C_COMPILER_LAUNCHER=ccache -DCMAKE_CXX_COMPILER_LAUNCHER=ccache -DBUILD_DEV=On ${flags} ..
+            cmake -DCMAKE_C_COMPILER_LAUNCHER=ccache -DCMAKE_CXX_COMPILER_LAUNCHER=ccache -DBUILD_DEV=On -DCMAKE_EXECUTE_PROCESS_COMMAND_ECHO=STDOUT ${flags} ..
            git diff
            git diff-index --quiet HEAD || (echo "Git repo is not clean after running cmake." && exit 1)
            make -j\$(nproc) generate VERBOSE=1
@@ -90,7 +90,7 @@ def rocmnodename(name) {
    } else if(name == "mi100+") {
        node_name = "${rocmtest_name} && (gfx908 || gfx90a) && !vm";
    } else if(name == "cdna") {
-        node_name = "${rocmtest_name} && (gfx908 || gfx90a || vega) && !vm";
+        node_name = "${rocmtest_name} && (gfx908 || gfx90a || vega20) && !vm";
    } else if(name == "nogpu") {
        node_name = "${rocmtest_name} && nogpu";
    }

--- a/cmake/Embed.cmake
+++ b/cmake/Embed.cmake
@@ -24,11 +24,7 @@
 find_program(EMBED_LD ld)
 find_program(EMBED_OBJCOPY objcopy)

-if(LINUX)
-    option(EMBED_USE_LD "Use ld to embed data files" ON)
-else()
-    option(EMBED_USE_LD "Use ld to embed data files" OFF)
-endif()
+option(EMBED_USE_LD "Use ld to embed data files" OFF)

 function(wrap_string)
    set(options)
@@ -60,8 +56,8 @@ endfunction()

 function(generate_embed_source EMBED_NAME)
    set(options)
-    set(oneValueArgs SRC HEADER)
-    set(multiValueArgs OBJECTS SYMBOLS)
+    set(oneValueArgs SRC HEADER RELATIVE)
+    set(multiValueArgs OBJECTS SYMBOLS FILES)

    cmake_parse_arguments(PARSE "${options}" "${oneValueArgs}" "${multiValueArgs}" ${ARGN})

@@ -78,6 +74,8 @@ function(generate_embed_source EMBED_NAME)
    foreach(idx RANGE ${LEN})
        list(GET PARSE_SYMBOLS ${idx} SYMBOL)
        list(GET PARSE_OBJECTS ${idx} OBJECT)
+        list(GET PARSE_FILES ${idx} FILE)
+
        set(START_SYMBOL "_binary_${SYMBOL}_start")
        set(END_SYMBOL "_binary_${SYMBOL}_end")
        if(EMBED_USE_LD)
@@ -92,9 +90,11 @@ function(generate_embed_source EMBED_NAME)
            ")
        endif()

-        # TODO: Should use NAME_WLE
-        get_filename_component(BASE_NAME "${OBJECT}" NAME)
-        string(REGEX REPLACE ".[A-Za-z0-9_]+$" "" BASE_NAME ${BASE_NAME})
+        if(PARSE_RELATIVE)
+            file(RELATIVE_PATH BASE_NAME ${PARSE_RELATIVE} "${FILE}")
+        else()
+            get_filename_component(BASE_NAME "${FILE}" NAME)
+        endif()

        string(APPEND INIT_KERNELS "
            { \"${BASE_NAME}\", { ${START_SYMBOL}, ${END_SYMBOL}} },
@@ -162,6 +162,11 @@ function(embed_file OUTPUT_FILE OUTPUT_SYMBOL FILE)
 endfunction()

 function(add_embed_library EMBED_NAME)
+    set(options)
+    set(oneValueArgs RELATIVE)
+    set(multiValueArgs)
+    cmake_parse_arguments(PARSE "${options}" "${oneValueArgs}" "${multiValueArgs}" ${ARGN})
+
    file(MAKE_DIRECTORY ${CMAKE_CURRENT_BINARY_DIR}/embed)
    file(MAKE_DIRECTORY ${CMAKE_CURRENT_BINARY_DIR}/embed/${EMBED_NAME})
    set(EMBED_DIR ${CMAKE_CURRENT_BINARY_DIR}/embed/${EMBED_NAME})
@@ -171,15 +176,26 @@ function(add_embed_library EMBED_NAME)
    set(OUTPUT_FILES)
    set(SYMBOLS)
    message(STATUS "Embedding files")
-    foreach(FILE ${ARGN})
+    foreach(FILE ${PARSE_UNPARSED_ARGUMENTS})
        embed_file(OUTPUT_FILE OUTPUT_SYMBOL ${FILE})
        list(APPEND OUTPUT_FILES ${OUTPUT_FILE})
        list(APPEND SYMBOLS ${OUTPUT_SYMBOL})
    endforeach()
    message(STATUS "Generating embedding library ${EMBED_NAME}")
-    generate_embed_source(${EMBED_NAME} SRC ${SRC_FILE} HEADER ${HEADER_FILE} OBJECTS ${OUTPUT_FILES} SYMBOLS ${SYMBOLS})
-    add_library(${EMBED_NAME} STATIC ${OUTPUT_FILES} "${SRC_FILE}")
-    target_include_directories(${EMBED_NAME} PUBLIC "${EMBED_DIR}/include")
-    target_compile_options(${EMBED_NAME} PRIVATE -Wno-reserved-identifier -Wno-extern-initializer -Wno-missing-variable-declarations)
-    set_target_properties(${EMBED_NAME} PROPERTIES POSITION_INDEPENDENT_CODE On)
+    generate_embed_source(${EMBED_NAME} SRC ${SRC_FILE} HEADER ${HEADER_FILE} OBJECTS ${OUTPUT_FILES} SYMBOLS ${SYMBOLS} RELATIVE ${PARSE_RELATIVE} FILES ${PARSE_UNPARSED_ARGUMENTS})
+    
+    set(INTERNAL_EMBED_LIB embed_lib_${EMBED_NAME})
+    add_library(${INTERNAL_EMBED_LIB} OBJECT "${SRC_FILE}")
+    target_include_directories(${INTERNAL_EMBED_LIB} PRIVATE "${EMBED_DIR}/include")
+    target_compile_options(${INTERNAL_EMBED_LIB} PRIVATE -Wno-reserved-identifier -Wno-extern-initializer -Wno-missing-variable-declarations)
+    set_target_properties(${INTERNAL_EMBED_LIB} PROPERTIES POSITION_INDEPENDENT_CODE On)
+    
+    add_library(${EMBED_NAME} INTERFACE)
+    if(EMBED_USE_LD)
+        target_sources(${EMBED_NAME} INTERFACE ${OUTPUT_FILES})
+    else()
+        target_sources(${INTERNAL_EMBED_LIB} PRIVATE ${OUTPUT_FILES})
+    endif()
+    target_sources(${EMBED_NAME} INTERFACE $<TARGET_OBJECTS:${INTERNAL_EMBED_LIB}>)
+    target_include_directories(${EMBED_NAME} INTERFACE "${EMBED_DIR}/include")
 endfunction()
--- a/cmake/PythonModules.cmake
+++ b/cmake/PythonModules.cmake
@@ -38,12 +38,22 @@ macro(find_python version)
    find_program(PYTHON_CONFIG_${version} python${version}-config)
    if(EXISTS ${PYTHON_CONFIG_${version}})
        py_exec(COMMAND ${PYTHON_CONFIG_${version}} --includes OUTPUT_VARIABLE _python_include_args)
+        execute_process(COMMAND ${PYTHON_CONFIG_${version}} --ldflags --embed OUTPUT_VARIABLE _python_ldflags_args RESULT_VARIABLE _python_ldflags_result)
+        if(NOT _python_ldflags_result EQUAL 0)
+            py_exec(COMMAND ${PYTHON_CONFIG_${version}} --ldflags OUTPUT_VARIABLE _python_ldflags_args)
+        endif()
        separate_arguments(_python_includes UNIX_COMMAND "${_python_include_args}")
+        separate_arguments(_python_ldflags UNIX_COMMAND "${_python_ldflags_args}")
        string(REPLACE "-I" "" _python_includes "${_python_includes}")
        add_library(python${version}::headers INTERFACE IMPORTED GLOBAL)
        set_target_properties(python${version}::headers PROPERTIES
            INTERFACE_INCLUDE_DIRECTORIES "${_python_includes}"
        )
+        add_library(python${version}::runtime INTERFACE IMPORTED GLOBAL)
+        set_target_properties(python${version}::runtime PROPERTIES
+            INTERFACE_LINK_OPTIONS "${_python_ldflags}"
+            INTERFACE_LINK_LIBRARIES python${version}::headers
+        )
        py_exec(COMMAND ${PYTHON_CONFIG_${version}} --prefix OUTPUT_VARIABLE _python_prefix)
        string(STRIP "${_python_prefix}" _python_prefix)
        set(PYTHON_${version}_EXECUTABLE "${_python_prefix}/bin/python${version}" CACHE PATH "")

--- a/docs/driver/read.rst
+++ b/docs/driver/read.rst
@@ -82,6 +82,10 @@ Print out program in text format.

 Print out program in binary format.

+.. option::  --py
+
+Print out program using python API.
+
 .. option::  --output, -o [std::string]

 Output to file.

--- a/src/api/CMakeLists.txt
+++ b/src/api/CMakeLists.txt
@@ -26,6 +26,7 @@ add_library(migraphx_c
    api.cpp
 )
 set_target_properties(migraphx_c PROPERTIES EXPORT_NAME c)
+migraphx_generate_export_header(migraphx_c DIRECTORY migraphx/api)

 # migraphx_c is stable API interface library. SO version of this should be 
 # bumped when binary compatibility is broken. 

--- a/src/api/api.cpp
+++ b/src/api/api.cpp
@@ -44,7 +44,7 @@ namespace migraphx {

 static thread_local bool disable_exception_catch = false; // NOLINT

-extern "C" void migraphx_test_private_disable_exception_catch(bool b)
+extern "C" MIGRAPHX_C_EXPORT void migraphx_test_private_disable_exception_catch(bool b)
 {
    disable_exception_catch = b;
 }

--- a/src/api/include/migraphx/migraphx.h
+++ b/src/api/include/migraphx/migraphx.h
--- a/src/auto_contiguous.cpp
+++ b/src/auto_contiguous.cpp
@@ -25,6 +25,7 @@
 #include <migraphx/program.hpp>
 #include <migraphx/instruction.hpp>
 #include <migraphx/make_op.hpp>
+#include <migraphx/ranges.hpp>

 #include <migraphx/iterator_for.hpp>

@@ -59,17 +60,26 @@ void auto_contiguous::apply(module& m) const
    auto last = std::prev(m.end());
    for(auto ins : iterator_for(m))
    {
-        if(ins->name() == "layout")
+        if(contains({"layout", "contiguous", "@return", "@param", "@outline"}, ins->name()))
            continue;
        // for last instruction that is NOT a return
        if(ins->outputs().empty() and ins != last)
            continue;
        shape s = ins->get_shape();
-        if(not s.dynamic() and not s.standard() and s.elements() != 0)
+        if(s.dynamic())
+            continue;
+        if(s.type() == shape::tuple_type)
+            continue;
+        if(s.standard() and ins->name() == "@literal")
+            continue;
+        if(s.scalar() and not contains(ins->name(), "broadcast"))
        {
-            auto c = m.insert_instruction(std::next(ins), make_op("contiguous"), ins);
-            m.replace_instruction(ins, c);
+            continue;
        }
+        if(ins->name() == "pointwise")
+            std::cout << "HERE" << std::endl;
+        auto c = m.insert_instruction(std::next(ins), make_op("contiguous"), ins);
+        m.replace_instruction(ins, c);
    }
 }


--- a/src/driver/CMakeLists.txt
+++ b/src/driver/CMakeLists.txt
@@ -32,18 +32,20 @@ add_executable(driver
    marker_roctx.cpp
 )
 set_target_properties(driver PROPERTIES OUTPUT_NAME migraphx-driver)
-# Copy driver for backwards compatibility
-add_custom_command(
-        TARGET driver 
+if(NOT WIN32)
+    # Copy driver for backwards compatibility (Linux only)
+    add_custom_command(
+        TARGET driver
        POST_BUILD COMMAND ${CMAKE_COMMAND} -E copy
                $<TARGET_FILE:driver>
                ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/driver
        BYPRODUCTS ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/driver
-)
-set_directory_properties(PROPERTIES ADDITIONAL_MAKE_CLEAN_FILES ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/driver)
+    )
+    set_directory_properties(PROPERTIES ADDITIONAL_CLEAN_FILES ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/driver)
+endif()
 rocm_clang_tidy_check(driver)

-target_link_libraries(driver migraphx_all_targets migraphx_onnx migraphx_tf)
+target_link_libraries(driver migraphx_all_targets migraphx_onnx migraphx_tf migraphx_py)

 rocm_install_targets(
  TARGETS driver

--- a/src/driver/main.cpp
+++ b/src/driver/main.cpp
@@ -32,6 +32,7 @@

 #include <migraphx/tf.hpp>
 #include <migraphx/onnx.hpp>
+#include <migraphx/py.hpp>
 #include <migraphx/stringutils.hpp>
 #include <migraphx/convert_to_json.hpp>
 #include <migraphx/load_save.hpp>
@@ -241,6 +242,20 @@ struct loader
        return options;
    }

+    static std::string get_file_type(const std::string& file)
+    {
+        if(ends_with(file, ".onnx"))
+            return "onnx";
+        else if(ends_with(file, ".pb"))
+            return "tf";
+        else if(ends_with(file, ".json"))
+            return "json";
+        else if(ends_with(file, ".py"))
+            return "py";
+        else
+            return "migraphx";
+    }
+
    program load()
    {
        program p;
@@ -248,14 +263,7 @@ struct loader
        {
            if(file_type.empty())
            {
-                if(ends_with(file, ".onnx"))
-                    file_type = "onnx";
-                else if(ends_with(file, ".pb"))
-                    file_type = "tf";
-                else if(ends_with(file, ".json"))
-                    file_type = "json";
-                else
-                    file_type = "migraphx";
+                file_type = get_file_type(file);
            }
            std::cout << "Reading: " << file << std::endl;
            if(file_type == "onnx")
@@ -272,6 +280,10 @@ struct loader
                options.format = "json";
                p              = migraphx::load(file, options);
            }
+            else if(file_type == "py")
+            {
+                p = migraphx::load_py(file);
+            }
            else if(file_type == "migraphx")
            {
                p = migraphx::load(file);

--- a/src/dynamic_loader.cpp
+++ b/src/dynamic_loader.cpp
@@ -48,7 +48,7 @@ struct dynamic_loader_impl
 #pragma GCC diagnostic ignored "-Wignored-attributes"
 #endif
    dynamic_loader_impl(const fs::path& p, std::shared_ptr<tmp_dir> t = nullptr)
-        : handle(dlopen(p.string().c_str(), RTLD_LAZY),
+        : handle(dlopen(p.string().c_str(), RTLD_GLOBAL | RTLD_NOW),
                 manage_deleter<decltype(&dlclose), &dlclose>{}),
          temp(std::move(t))
    {
@@ -81,6 +81,18 @@ fs::path dynamic_loader::path(void* address)
    return p;
 }

+optional<dynamic_loader> dynamic_loader::try_load(const fs::path& p)
+{
+    try
+    {
+        return dynamic_loader{p};
+    }
+    catch(const std::exception&)
+    {
+        return nullopt;
+    }
+}
+
 dynamic_loader::dynamic_loader(const fs::path& p) : impl(std::make_shared<dynamic_loader_impl>(p))
 {
 }

--- a/src/eliminate_contiguous.cpp
+++ b/src/eliminate_contiguous.cpp
@@ -30,6 +30,7 @@
 #include <migraphx/op/contiguous.hpp>
 #include <migraphx/op/identity.hpp>
 #include <migraphx/par_for.hpp>
+#include <type_traits>
 #include <utility>

 namespace migraphx {
@@ -161,6 +162,18 @@ static void remove_contiguous(const std::string& op_name, module& m, F f)
    }
 }

+static void remove_contiguous_noops(const std::string& op_name, module& m)
+{
+    for(auto ins : iterator_for(m))
+    {
+        if(ins->name() != op_name)
+            continue;
+        if(ins->inputs().front()->get_shape() != ins->get_shape())
+            continue;
+        m.replace_instruction(ins, ins->inputs().front());
+    }
+}
+
 void eliminate_contiguous::apply(module& m) const
 {
    // Skip contiguous from splits first
@@ -170,6 +183,7 @@ void eliminate_contiguous::apply(module& m) const
        return (ins->inputs().front()->outputs().size() == 1);
    });
    remove_contiguous(op_name, m, [](auto) { return true; });
+    remove_contiguous_noops(op_name, m);
 }

 } // namespace MIGRAPHX_INLINE_NS

--- a/src/eliminate_layout.cpp
+++ b/src/eliminate_layout.cpp
@@ -108,91 +108,42 @@ void remove_layout(module& m)
 //     return convs;
 // }

-// void remove_layout(module& m, const std::vector<instruction_ref>& convs)
-// {
-//     if(convs.size() < 2) return;
-//     m.debug_print();
-
-//     for(auto i = 0; i < convs.size() - 1; i++)
-//     {
-//         bool reached_start = false;
-//         for(auto ins : iterator_for(m))
-//         {
-//             if(ins == convs[i])
-//                 reached_start = true;
-//             if(reached_start)
-//             {
-//                 if(ins->name() == "gpu::pooling")
-//                     break;
-//                 if(ins == convs[i + 1])
-//                 {
-
-//                     m.debug_print(convs[i]->outputs().front());
-//                     m.debug_print(convs[i]->outputs().front()->outputs().front());
-//                     m.replace_instruction(convs[i]->outputs().front(),
-//                     convs[i]->outputs().front()->outputs().front()); std::cout << "HERE" <<
-//                     std::endl; m.debug_print(convs[i]->outputs().front());
-
-//                     // m.debug_print(convs[i]->outputs().front());
-//                     // m.debug_print(convs[i]->outputs().front()->outputs().front());
-
-//                     std::cout << std::endl;
-//                     m.debug_print(convs[i]->inputs());
-//                     std::cout << std::endl;
-//                     m.debug_print(convs[i + 1]->inputs());
-
-//                     for(auto j = 0; j < convs[i + 1]->inputs().size(); j++)
-//                     {
-//                         if(convs[i]->inputs()[j] == convs[i + 1]->inputs()[j])
-//                         {
-//                             std::cout << "HERE2" << std::endl;
-//                             continue;
-//                         }
-//                         m.replace_instruction(convs[i + 1]->inputs()[j], convs[i +
-//                         1]->inputs()[j]->inputs().front()); m.debug_print(convs[i+1]);
-//                     }
-//                     break;
-//                 }
-//             }
-
-//         }
-//     }
-// }
-
-// void remove_layout(module& m, const std::unordered_set<instruction_ref>& output_layouts)
-// {
-//     for(auto ins : iterator_for(m))
-//     {
-//         if(ins->name() != "gpu::precompile_op")
-//             continue;
+void remove_layout(module& m, const std::unordered_set<instruction_ref>& output_layouts)
+{
+    for(auto ins : iterator_for(m))
+    {
+        if(ins->name() != "gpu::precompile_op")
+            continue;

-//         auto precompile_op = ins->get_operator();
-//         auto val           = precompile_op.to_value();
+        auto precompile_op = ins->get_operator();
+        auto val           = precompile_op.to_value();

-//         if(val["op"].at("name").to<std::string>() != "layout")
-//         {
-//             // std::cout << val["op"].at("name").to<std::string>() << std::endl;
-//             continue;
-//         }
-//         m.debug_print(ins);
-//         if(ins->get_shape() != ins->inputs().front()->get_shape())
-//         {
-//             std::cout << ins->get_shape() << " " << ins->inputs().front()->get_shape() <<
-//             std::endl; continue;
-//         }
-//         if(contains(output_layouts, ins))
-//             continue;
+        if(val["op"].at("name").to<std::string>() != "layout")
+        {
+            // std::cout << val["op"].at("name").to<std::string>() << std::endl;
+            continue;
+        }
+        // m.debug_print(ins);
+        if(ins->get_shape() != ins->inputs().front()->get_shape())
+        {
+            // std::cout << ins->get_shape() << " " << ins->inputs().front()->get_shape() <<
+            // std::endl;
+            continue;
+        }
+        if(contains(output_layouts, ins))
+            continue;

-//         m.replace_instruction(ins, ins->inputs().front());
-//     }
-// }
+        m.replace_instruction(ins, ins->inputs().front());
+    }
+}

 void eliminate_layout::apply(module_pass_manager& mpm) const
 {
-    // std::unordered_set<instruction_ref> output_layouts =
-    // preserve_output_layout(mpm.get_module()); remove_layout(mpm.get_module(),
+    std::unordered_set<instruction_ref> output_layouts =
+    preserve_output_layout(mpm.get_module()); 
+    remove_layout(mpm.get_module(), output_layouts);
    // find_convs(mpm.get_module()));
-    remove_layout(mpm.get_module());
+    // remove_layout(mpm.get_module());
    mpm.run_pass(dead_code_elimination{});
 }


--- a/src/include/migraphx/builtin.hpp
+++ b/src/include/migraphx/builtin.hpp
@@ -90,7 +90,17 @@ struct param
 struct returns
 {
    std::string name() const { return "@return"; }
-    shape compute_shape(const std::vector<shape>&) const { return {}; }
+
+    shape compute_shape(const std::vector<shape>& arg) const
+    {
+        if(arg.empty())
+            return {};
+        else if(arg.size() == 1)
+            return arg[0];
+        else
+            return arg;
+    }
+
    argument compute(context&, const shape&, const std::vector<argument>&) const
    {
        MIGRAPHX_THROW("builtin");

--- a/src/include/migraphx/dynamic_loader.hpp
+++ b/src/include/migraphx/dynamic_loader.hpp
@@ -26,6 +26,7 @@

 #include <migraphx/config.hpp>
 #include <migraphx/filesystem.hpp>
+#include <migraphx/optional.hpp>
 #include <functional>
 #include <memory>
 #include <vector>
@@ -43,6 +44,9 @@ struct MIGRAPHX_EXPORT dynamic_loader
        return path(reinterpret_cast<void*>(address));
    }
    static fs::path path(void* address);
+
+    static optional<dynamic_loader> try_load(const fs::path& p);
+
    dynamic_loader() = default;

    dynamic_loader(const fs::path& p);

--- a/src/include/migraphx/module.hpp
+++ b/src/include/migraphx/module.hpp
@@ -222,7 +222,17 @@ struct MIGRAPHX_EXPORT module
    void annotate(std::ostream& os, std::function<void(instruction_ref)> a) const;

    std::vector<module_ref> get_sub_modules(bool shallow = false) const;
+    /* sorts the module in topological order aka reverse-post order (RPO) DFS order
+       it takes last instruction or @return as the root and walks back the graph and moves inputs
+       of the each instruction such that it appears before the instruction itself.
+    */
    module& sort();
+    /* Any instruction "X" can have module arguments and those modules inside them can use any other
+     * instruction "Y" from predecessor modules of the instruction "X". Such instruction "Y" inside
+     * module args are not listed as input instructions to "X". But those instructions "Y" must be
+     * evaluted before the instruction "X" can. Therefore such "Y" instructions are considered
+     * implicit dependency to "X".
+     */
    ins_dep_map calc_implicit_deps() const;

    MIGRAPHX_EXPORT friend std::ostream& operator<<(std::ostream& os, const module& m);

--- a/src/instruction.cpp
+++ b/src/instruction.cpp
@@ -64,10 +64,7 @@ void instruction::replace(const shape& r)
        result = r;
        for(auto&& ins : output)
        {
-            if(ins->name() == "@return")
-                continue;
-
-            assert(ins->name().front() != '@');
+            assert(ins->name() == "@return" or ins->name().front() != '@');
            ins->recompute_shape();
        }
    }
@@ -122,10 +119,6 @@ bool instruction::valid() const
    {
        computed = result;
    }
-    else if(op.name() == "@return")
-    {
-        computed = {};
-    }
    else
    {
        try
@@ -145,6 +138,7 @@ bool instruction::valid() const
 }

 shape instruction::get_shape() const { return result; }
+
 const literal& instruction::get_literal() const
 {
    assert(op.name() == "@literal");

--- a/src/layout_nhwc.cpp
+++ b/src/layout_nhwc.cpp
@@ -149,10 +149,10 @@ void layout_nhwc::apply(module_pass_manager& mpm) const
    // std::cout << "after layout" << std::endl;
    // mpm.get_module().debug_print();
    // if(not this->skip_elim_contiguous)
-    mpm.run_pass(eliminate_contiguous{"contiguous"});
-    mpm.run_pass(dead_code_elimination{});
-    mpm.run_pass(auto_contiguous{});
-    mpm.run_pass(dead_code_elimination{});
+    // mpm.run_pass(eliminate_contiguous{"contiguous"});
+    // mpm.run_pass(dead_code_elimination{});
+    // mpm.run_pass(auto_contiguous{});
+    // mpm.run_pass(dead_code_elimination{});
    // remove_layout(mpm.get_module(), output_layouts);
    // mpm.run_pass(dead_code_elimination{});
 }

--- a/src/module.cpp
+++ b/src/module.cpp
@@ -460,11 +460,11 @@ instruction_ref module::add_parameter(std::string name, shape s)

 instruction_ref module::add_return(std::vector<instruction_ref> args)
 {
-    impl->push_back({builtin::returns{}, {}, std::move(args)});
+    shape instr_shape = compute_shape(builtin::returns{}, args);
+    impl->push_back({builtin::returns{}, instr_shape, std::move(args)});
    auto result = std::prev(impl->instructions.end());
    instruction::backreference(result);
    assert(result->valid(begin()));
-
    return result;
 }

@@ -1011,9 +1011,17 @@ std::vector<module_ref> module::get_sub_modules(bool shallow) const

 module& module::sort()
 {
+    auto implicit_deps = calc_implicit_deps();
    fix([&](auto self, auto ins) {
        this->move_instruction(ins, this->begin());
-        for(auto child : ins->inputs())
+        auto ins_inputs = ins->inputs();
+        if(implicit_deps.find(ins) != implicit_deps.end())
+        {
+            auto ins_implict_inputs = implicit_deps.at(ins);
+            ins_inputs.insert(
+                ins_inputs.end(), ins_implict_inputs.begin(), ins_implict_inputs.end());
+        }
+        for(auto child : ins_inputs)
        {
            if(not contains(this->impl->instructions, child))
            {

--- a/src/program.cpp
+++ b/src/program.cpp
@@ -40,13 +40,14 @@
 #include <migraphx/make_op.hpp>
 #include <migraphx/marker.hpp>
 #include <migraphx/supported_segments.hpp>
+
 #include <iostream>
+#include <queue>
 #include <sstream>
 #include <algorithm>
 #include <set>
 #include <unordered_map>
 #include <utility>
-
 #include <unordered_set>
 #include <map>
 #include <cassert>
@@ -1191,11 +1192,19 @@ void program::remove_unused_modules()

 program& program::sort()
 {
-    for(auto& pp : this->impl->modules)
+    std::queue<migraphx::module_ref> mqueue;
+    mqueue.push(get_main_module());
+    while(not mqueue.empty())
    {
-        pp.second.sort();
+        module_ref current_mod = mqueue.front();
+        current_mod->sort();
+        mqueue.pop();
+        auto child_mods = current_mod->get_sub_modules(true);
+        for(auto& sub_mod : child_mods)
+        {
+            mqueue.push(sub_mod);
+        }
    }
-
    return *this;
 }