Merge branch 'develop' into update-python-packages

CMakeLists.txt

@@ -57,6 +57,12 @@ else()
 option(MIGRAPHX_ENABLE_PYTHON "Enable python bindings" ON)
 endif()
 
+if(WIN32) # CK is not yet ported to Windows
+option(MIGRAPHX_USE_COMPOSABLEKERNEL "Enable MIGraphX to use composable kernel JIT library" OFF)
+else()
+option(MIGRAPHX_USE_COMPOSABLEKERNEL "Enable MIGraphX to use composable kernel JIT library" ON)
+endif()
+
 find_path(HALF_INCLUDE_DIR half.hpp PATH_SUFFIXES half)
 if (NOT HALF_INCLUDE_DIR)
     message(FATAL_ERROR "Could not find half.hpp - Please check that the install path of half.hpp has been added to CMAKE_PREFIX_PATH")

src/CMakeLists.txt

@@ -261,9 +261,8 @@ find_package(nlohmann_json 3.8.0 REQUIRED)
 target_link_libraries(migraphx PRIVATE nlohmann_json::nlohmann_json)
 migraphx_generate_export_header(migraphx)
 
-find_package(PkgConfig)
-pkg_check_modules(SQLITE3 REQUIRED IMPORTED_TARGET sqlite3)
-target_link_libraries(migraphx PRIVATE PkgConfig::SQLITE3)
+find_package(SQLite3 REQUIRED)
+target_link_libraries(migraphx PRIVATE SQLite::SQLite3)
 
 find_package(msgpackc-cxx QUIET)
 if(NOT msgpackc-cxx_FOUND)

src/include/migraphx/op/allocate.hpp

@@ -1,7 +1,7 @@
 /*
  * The MIT License (MIT)
  *
- * Copyright (c) 2015-2022 Advanced Micro Devices, Inc. All rights reserved.
+ * Copyright (c) 2015-2023 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
@@ -41,16 +41,16 @@ namespace op {
  * Dynamic allocate:
  * One input: `allocate(output_dims)`
  * `output_dims` are the output buffer dimensions and has a static shape.
- * Either `this.s` or `this.buf_type` must be set to calculate the dynamic output shape at compute
- * time. If `this.buf_type` is set, the compute_shape() of allocate at compile time will have
- * dynamic_dimensions from {0, max_int} with rank = output_dims.ndim(). If `this.s` is set then the
- * compute_shape() will output `this.s`; `this.s` should be a dynamic shape.
+ * Either `this.s` or `this.buf_type` (but not both) must be set to calculate the dynamic output
+ * shape at compute time. If `this.buf_type` is set, the compute_shape() of allocate at compile time
+ * will have dynamic_dimensions from {0, max_int} with rank = output_dims.ndim(). If `this.s` is set
+ * then the compute_shape() will output `this.s`; `this.s` should be a dynamic shape.
  */
 struct allocate
 {
-    shape s{};
+    optional<shape> s;
     // for dynamic allocate to set the buffer type
-    shape::type_t buf_type = shape::half_type;
+    optional<shape::type_t> buf_type;
 
     template <class Self, class F>
     static auto reflect(Self& self, F f)
@@ -62,26 +62,38 @@ struct allocate
 
     shape compute_shape(const std::vector<shape>& inputs) const
     {
-        if(s != shape())
+        if(s.has_value())
         {
+            if(buf_type.has_value())
+            {
+                MIGRAPHX_THROW("ALLOCATE: shape and buf_type attributes both set");
+            }
             if(inputs.size() == 1)
             {
                 migraphx::check_shapes{inputs, *this, false}.only_dims(1);
             }
             else
             {
+                if(s->dynamic())
+                {
+                    MIGRAPHX_THROW("ALLOCATE: dynamic shape attribute and no input");
+                }
                 migraphx::check_shapes{inputs, *this, false}.has(0);
             }
-            return s;
+            return s.value();
         }
         else
         {
+            if(not buf_type.has_value())
+            {
+                MIGRAPHX_THROW("ALLOCATE: shape and buf_type attributes both not set");
+            }
             migraphx::check_shapes{inputs, *this, false}.has(1).only_dims(1);
             const auto& out_dims = inputs.at(0);
             std::size_t max_val = std::numeric_limits<std::size_t>::max();
             std::vector<shape::dynamic_dimension> dyn_dims(out_dims.lens().at(0),
                                                            shape::dynamic_dimension{0, max_val});
-            return {buf_type, dyn_dims};
+            return {buf_type.value(), dyn_dims};
         }
     }
     argument compute(const shape& output_shape, const std::vector<argument>& args) const
@@ -94,7 +106,11 @@ struct allocate
         {
             std::vector<std::size_t> output_dims(output_shape.ndim());
             args.at(0).visit([&](auto a) { output_dims.assign(a.begin(), a.end()); });
-            return argument{shape{buf_type, output_dims}};
+            if(s)
+            {
+                return argument{shape{s->type(), output_dims}};
+            }
+            return argument{shape{buf_type.value(), output_dims}};
         }
     }
 };

src/include/migraphx/streamutils.hpp

@@ -30,6 +30,7 @@
 #include <migraphx/rank.hpp>
 #include <migraphx/requires.hpp>
 #include <migraphx/config.hpp>
+#include <migraphx/optional.hpp>
 #include <vector>
 
 namespace migraphx {
@@ -68,6 +69,19 @@ auto stream_write_value_impl(rank<1>, std::ostream& os, const T& x) -> decltype(
     os << x;
 }
 
+template <class T>
+auto stream_write_value_impl(rank<1>, std::ostream& os, const optional<T>& x)
+{
+    if(x.has_value())
+    {
+        os << *x;
+    }
+    else
+    {
+        os << "nullopt";
+    }
+}
+
 template <class T>
 void stream_write_value_impl(rank<1>, std::ostream& os, const std::vector<T>& r)
 {

src/targets/gpu/CMakeLists.txt

@@ -37,8 +37,7 @@ if(NOT TARGET MIOpen)
     message(SEND_ERROR "Cant find miopen")
 endif()
 
-if(NOT WIN32)
-    # TODO: re-enable when CK is ported to Windows
+if(MIGRAPHX_USE_COMPOSABLEKERNEL)
     find_package(composable_kernel 1.0.0 REQUIRED COMPONENTS jit_library)
 endif()
 
@@ -52,10 +51,10 @@ file(GLOB KERNEL_FILES CONFIGURE_DEPENDS
     ${CMAKE_CURRENT_SOURCE_DIR}/kernels/include/migraphx/kernels/*.hpp)
 message(STATUS "KERNEL_FILES: ${KERNEL_FILES}")
 
-if(WIN32)
-    # TODO: re-enable when CK is ported to Windows
+if(NOT MIGRAPHX_USE_COMPOSABLEKERNEL)
     list(REMOVE_ITEM KERNEL_FILES
         ${CMAKE_CURRENT_SOURCE_DIR}/kernels/include/migraphx/kernels/ck_gemm.hpp
+        ${CMAKE_CURRENT_SOURCE_DIR}/kernels/include/migraphx/kernels/ck_gemm_softmax_gemm.hpp
         ${CMAKE_CURRENT_SOURCE_DIR}/kernels/include/migraphx/kernels/ck.hpp)
 endif()
 
@@ -103,9 +102,10 @@ rocm_clang_tidy_check(kernel_file_check)
 
 file(GLOB JIT_GPU_SRCS CONFIGURE_DEPENDS ${CMAKE_CURRENT_SOURCE_DIR}/jit/*.cpp)
 
-if(WIN32)
-    # TODO: re-enable when CK is ported to Windows
-    list(REMOVE_ITEM JIT_GPU_SRCS ${CMAKE_CURRENT_SOURCE_DIR}/jit/ck_gemm.cpp)
+if(NOT MIGRAPHX_USE_COMPOSABLEKERNEL)
+    list(REMOVE_ITEM JIT_GPU_SRCS
+            ${CMAKE_CURRENT_SOURCE_DIR}/jit/ck_gemm.cpp
+            ${CMAKE_CURRENT_SOURCE_DIR}/jit/ck_gemm_softmax_gemm.cpp)
 endif()
 
 add_library(migraphx_gpu
@@ -281,8 +281,7 @@ endif()
 
 target_link_libraries(migraphx_gpu PUBLIC migraphx MIOpen roc::rocblas)
 target_link_libraries(migraphx_gpu PRIVATE migraphx_device migraphx_kernels)
-if(NOT WIN32)
-    # TODO: re-enable when CK is ported to Windows
+if(MIGRAPHX_USE_COMPOSABLEKERNEL)
     target_link_libraries(migraphx_gpu PRIVATE composable_kernel::jit_library)
 endif()
 

test/CMakeLists.txt

@@ -25,7 +25,7 @@
 cmake_policy(SET CMP0057 NEW)
 
 find_package(Threads REQUIRED)
-rocm_test_link_libraries(Threads::Threads migraphx migraphx_ref migraphx_onnx migraphx_tf)
+rocm_test_link_libraries(Threads::Threads migraphx migraphx_onnx migraphx_tf)
 rocm_test_include_directories(include)
 
 set(MIGRAPHX_DISABLE_LARGE_BUFFER_TESTS Off CACHE BOOL "")
@@ -146,7 +146,10 @@ endfunction()
 
 function(test_headers PREFIX)
     file(GLOB HEADERS CONFIGURE_DEPENDS ${ARGN})
-
+    if(NOT MIGRAPHX_USE_COMPOSABLEKERNEL)
+        list(REMOVE_ITEM HEADERS
+              ${CMAKE_SOURCE_DIR}/src/targets/gpu/include/migraphx/gpu/ck.hpp)
+    endif()
     foreach(HEADER ${HEADERS})
         file(RELATIVE_PATH HEADER_REL ${CMAKE_SOURCE_DIR} ${HEADER})
         string(MAKE_C_IDENTIFIER ${HEADER_REL} TEST_NAME)

test/gpu/fuse_mlir.cpp

@@ -152,6 +152,9 @@ TEST_CASE(int_quant_dot_tanh_fails)
 
 int main(int argc, const char* argv[])
 {
-    test::run(argc, argv);
+    if(migraphx::gpu::mlir_enabled())
+    {
+        test::run(argc, argv);
+    }
     return 0;
 }

test/op_shape_test.cpp

@@ -88,7 +88,7 @@ TEST_CASE(allocate_static)
     expect_shape(out_shape, migraphx::make_op("allocate", {{"shape", to_value(out_shape)}}));
 }
 
-TEST_CASE(allocate_static_input_error)
+TEST_CASE(allocate_static_input)
 {
     migraphx::shape input{migraphx::shape::int64_type, {3}};
     migraphx::shape out_shape{migraphx::shape::float_type, {2, 3, 4}};
@@ -120,8 +120,22 @@ TEST_CASE(allocate_dyn_no_input_error)
 {
     migraphx::shape shape_attr{migraphx::shape::float_type,
                                {{1, 4}, {3, 3}, {4, 8, {4, 6}}, {4, 8}, {4, 6}}};
-    expect_shape(shape_attr,
-                 migraphx::make_op("allocate", {{"shape", migraphx::to_value(shape_attr)}}));
+    throws_shape(migraphx::make_op("allocate", {{"shape", migraphx::to_value(shape_attr)}}));
+}
+
+TEST_CASE(allocate_shape_and_buf_type_error)
+{
+    migraphx::shape shape_attr{migraphx::shape::float_type,
+                               {{1, 4}, {3, 3}, {4, 8, {4, 6}}, {4, 8}, {4, 6}}};
+    throws_shape(migraphx::make_op(
+        "allocate",
+        {{"shape", migraphx::to_value(shape_attr)}, {"buf_type", migraphx::shape::half_type}}));
+}
+
+TEST_CASE(allocate_no_attr_error)
+{
+    migraphx::shape input{migraphx::shape::int64_type, {4}};
+    throws_shape(migraphx::make_op("allocate"), input);
 }
 
 TEST_CASE(argmax_axis0)

test/ref/allocate.cpp

@@ -30,7 +30,7 @@
 
 #include <test.hpp>
 
-TEST_CASE(allocate_dyn)
+TEST_CASE(allocate_dyn0)
 {
     migraphx::program p;
     auto* mm = p.get_main_module();
@@ -47,3 +47,21 @@ TEST_CASE(allocate_dyn)
     migraphx::shape sresult{migraphx::shape::float_type, {2, 3, 4, 4}};
     result.visit([&](auto output) { EXPECT(output.get_shape() == sresult); });
 }
+
+TEST_CASE(allocate_dyn1)
+{
+    migraphx::program p;
+    auto* mm = p.get_main_module();
+    migraphx::shape s{migraphx::shape::int64_type, {4}};
+    migraphx::shape out_shape{migraphx::shape::float_type, {2, 3, 4, 4}};
+    auto out_dims = mm->add_parameter("out_dims", s);
+    mm->add_instruction(migraphx::make_op("allocate", {{"shape", migraphx::to_value(out_shape)}}),
+                        out_dims);
+    p.compile(migraphx::make_target("ref"));
+
+    migraphx::parameter_map params;
+    std::vector<int64_t> data = {2, 3, 4, 4};
+    params["out_dims"]        = migraphx::argument(s, data.data());
+    auto result               = p.eval(params).back();
+    result.visit([&](auto output) { EXPECT(output.get_shape() == out_shape); });
+}

test/targets.cpp

@@ -41,11 +41,7 @@ TEST_CASE(make_invalid_target)
 
 TEST_CASE(targets)
 {
-    // GCC doesn't load libmigraphx_ref unless necesssary even though it is linked to the test.
-    // Force it to load by making ref target
-#if defined(__GNUC__) && !defined(__clang__)
     auto ref_target = migraphx::make_target("ref");
-#endif
     auto ts = migraphx::get_targets();
     EXPECT(ts.size() >= 1);
 }