Merge remote-tracking branch 'origin/develop' into jakpiase/smfmac_test

codegen/CMakeLists.txt

@@ -1,5 +1,5 @@
 cmake_minimum_required(VERSION 3.16)
-project(composable_kernel_host)
+project(composable_kernel_host LANGUAGES CXX HIP)
 
 set(CMAKE_EXPORT_COMPILE_COMMANDS ON)
 
@@ -12,24 +12,38 @@ find_package(ROCM)
 include(ROCMInstallTargets)
 include(ROCMTest)
 
+add_compile_options(-std=c++17)
+find_package(hip)
+## HIP
+set(CMAKE_HIP_PLATFORM amd)
+set(CMAKE_HIP_COMPILER ${CMAKE_CXX_COMPILER})
+set(CMAKE_HIP_EXTENSIONS ON)
+message("CMAKE_HIP_COMPILER: ${CMAKE_HIP_COMPILER}")
+
+# add include directories
+include_directories(BEFORE
+            ${PROJECT_BINARY_DIR}/include
+            ${PROJECT_SOURCE_DIR}/include
+            ${PROJECT_SOURCE_DIR}/library/include
+            ${HIP_INCLUDE_DIRS}
+            )
+
 list(APPEND CMAKE_MODULE_PATH ${CK_ROOT}/cmake)
 include(Embed)
 file(GLOB_RECURSE KERNEL_FILES CONFIGURE_DEPENDS
-    ${CK_ROOT}/include/ck/*.hpp)
+	${CK_ROOT}/include/ck/*.hpp)
 message(STATUS "KERNEL_FILES: ${KERNEL_FILES}")
 message(STATUS "RELATIVE: ${CK_ROOT}/include")
 add_embed_library(ck_headers ${KERNEL_FILES} RELATIVE ${CK_ROOT}/include)
 
-add_definitions(-std=c++17)
-
 file(GLOB SOURCES CONFIGURE_DEPENDS src/*.cpp)
 # TODO: Use object library
 add_library(ck_host STATIC ${SOURCES})
 target_link_libraries(ck_host PRIVATE ck_headers)
 
-set_target_properties(ck_host PROPERTIES 
-    LINKER_LANGUAGE CXX
-    POSITION_INDEPENDENT_CODE ON)
+set_target_properties(ck_host PROPERTIES
+	LINKER_LANGUAGE CXX
+        POSITION_INDEPENDENT_CODE ON)
 
 target_include_directories(ck_host PUBLIC
     $<BUILD_INTERFACE:${CMAKE_CURRENT_SOURCE_DIR}/include>

codegen/driver/main.cpp

@@ -5,24 +5,27 @@
 #include <unordered_map>
 #include <vector>
 #include "ck/host/device_gemm_multiple_d/operation.hpp"
+#include "ck/host/device_grouped_conv_fwd_multiple_d/conv_fwd_op.hpp"
 #include "ck/host/stringutils.hpp"
 
 using ck::host::Transform;
 
 struct Emitters
 {
+    // retrieve the hard-coded instances provided, template them, and then store them in a map
     std::unordered_map<std::string, std::function<std::vector<std::string>()>> m;
 
     template <class T>
-    void Register(const std::string& name)
+    void Register(const std::string& name, const std::string& prologue, const std::string& epilogue)
     {
-        m[name] = [] {
-            auto configs = T::CreateOperations();
+        m[name] = [&] {
+            auto configs = T::CreateOperations(prologue, epilogue);
 
             return Transform(configs, [](const auto& ops) { return ToTuple(ops); });
         };
     }
 
+    // takes in an operation instance and uses it to substitute the correct values into the template
     template <class T>
     static std::string ToTuple(const T& ops)
     {
@@ -31,6 +34,7 @@ struct Emitters
         return "std::tuple<\n" + ck::host::JoinStrings(templates, ",\n") + ">";
     }
 
+    // Join together all the strings in the map
     std::string Emit(const std::string& name) { return ck::host::JoinStrings(m.at(name)(), "\n"); }
 
     std::vector<std::string> List() const
@@ -43,9 +47,38 @@ int main(int argc, const char* argv[])
 {
     std::string prog = argv[0];
     std::vector<std::string> args(argv + 1, argv + argc);
+
+    // Specify problem type and problem size
+    ck::host::device_gemm_multiple_d::Problem prob;
+    prob.M = 1024;
+    prob.N = 1024;
+    prob.K = 1024;
+
+    // user provided fusion
+    std::string prologue = "";
+    std::string epilogue = R"(
+struct Epilogue
+{
+    __host__ __device__ Epilogue(float alpha, float beta) : alpha_(alpha), beta_(beta){};
+
+    template <typename E, typename D>
+    __host__ __device__ constexpr void operator()(E& e, const D& d) const;
+
+    template <>
+    __host__ __device__ constexpr void operator()<ck::half_t, ck::half_t>(ck::half_t& e,
+                                                                          const ck::half_t& d) const
+    {
+        e = ck::type_convert<ck::half_t>(alpha_ * e + beta_ * ck::type_convert<float>(d));
+    }
+
+    float alpha_;
+    float beta_;
+};)";
+
+    // Load in operations into the Register
     Emitters e;
     e.Register<ck::host::device_gemm_multiple_d::Operation_Xdl_CShuffle>(
-        "DeviceGemmMultipleD_Xdl_CShuffle");
+        "DeviceGemmMultipleD_Xdl_CShuffle", prologue, epilogue);
 
     if(args.empty() or std::any_of(args.begin(), args.end(), [](auto arg) {
            return arg == "-h" or arg == "--help";
@@ -64,6 +97,7 @@ int main(int argc, const char* argv[])
         return 0;
     }
 
+    // print out all the instances for the operation that was chosen at the command line
     for(auto name : args)
         std::cout << e.Emit(name) << std::endl;
 

codegen/include/ck/host/device_gemm_multiple_d.hpp

@@ -1,5 +1,5 @@
 // SPDX-License-Identifier: MIT
-// Copyright (c) 2018-2023, Advanced Micro Devices, Inc. All rights reserved.
+// Copyright (c) 2018-2024, Advanced Micro Devices, Inc. All rights reserved.
 
 #pragma once
 

codegen/include/ck/host/device_gemm_multiple_d/operation.hpp

@@ -14,10 +14,15 @@ namespace ck {
 namespace host {
 namespace device_gemm_multiple_d {
 
+// defines all values need for an instance of fwd conv
 struct Operation_Xdl_CShuffle
 {
-    static std::vector<std::vector<Operation_Xdl_CShuffle>> CreateOperations();
-    static std::vector<Operation_Xdl_CShuffle> CreateOperations(const Problem& prob);
+    // returns a vector of instances, only given fusion operators: will use default problem spec
+    static std::vector<std::vector<Operation_Xdl_CShuffle>>
+    CreateOperations(const std::string& prologue, const std::string& epilogue);
+    // returns a vector of instances, given a problem spec and fusion operators
+    static std::vector<Operation_Xdl_CShuffle>
+    CreateOperations(const Problem& prob, const std::string& prologue, const std::string& epilogue);
     TensorDesc A{};
     TensorDesc B{};
     DataType acc               = DataType::Float;
@@ -27,13 +32,21 @@ struct Operation_Xdl_CShuffle
     std::string a_elem_op           = PassThrough;
     std::string b_elem_op           = PassThrough;
     std::string cde_elem_op         = Bilinear;
+    std::string prologue            = "";
+    std::string epilogue            = "";
     std::string gemm_specialization = "ck::tensor_operation::device::GemmSpecialization::Default";
+    // tuning parameters
     operation::TileDesc tile_desc{};
     operation::BlockTransferDesc a_block_transfer{};
     operation::BlockTransferDesc b_block_transfer{};
     operation::CShuffleDesc cshuffle{};
     operation::CBlockTransferDesc c_block_transfer{};
 
+    // functions to update fusion operators if provided
+    void update_prologue(const std::string& prologue);
+    void update_epilogue(const std::string& epilogue);
+    /**constexpr**/ bool IsSupported(std::size_t MRaw_, std::size_t NRaw_, std::size_t KRaw_);
+    // returns a templated instance
     Solution ToSolution() const;
 };
 

codegen/include/ck/host/device_gemm_multiple_d/problem.hpp

@@ -1,5 +1,5 @@
 // SPDX-License-Identifier: MIT
-// Copyright (c) 2018-2023, Advanced Micro Devices, Inc. All rights reserved.
+// Copyright (c) 2018-2024, Advanced Micro Devices, Inc. All rights reserved.
 
 #pragma once
 
@@ -12,11 +12,14 @@ namespace ck {
 namespace host {
 namespace device_gemm_multiple_d {
 
+// defines the problem specification for a GEMM operation
 struct Problem
 {
-    std::size_t M                    = 0;
-    std::size_t N                    = 0;
-    std::size_t K                    = 0;
+    // dimensions for GEMM operation
+    std::size_t M = 0;
+    std::size_t N = 0;
+    std::size_t K = 0;
+    // layouts for tensors
     bool TransA                      = false;
     bool TransB                      = false;
     bool TransE                      = false;
@@ -29,9 +32,13 @@ struct Problem
     std::string BElementOp           = PassThrough;
     std::string CDEElementOp         = PassThrough;
 
+    // returns the correct device op file for the operation
     std::string GetIncludeHeader() const;
 
-    std::vector<Solution> GetSolutions(const std::string& arch) const;
+    // returns a list of instances based on the problem spec and provided fusion operations
+    std::vector<Solution> GetSolutions(const std::string& arch,
+                                       const std::string& prologue,
+                                       const std::string& epilogue) const;
 };
 
 } // namespace device_gemm_multiple_d

codegen/include/ck/host/device_grouped_conv_fwd_multiple_d/conv_fwd_op.hpp

@@ -0,0 +1,60 @@
+// SPDX-License-Identifier: MIT
+// Copyright (c) 2018-2024, Advanced Micro Devices, Inc. All rights reserved.
+
+#pragma once
+
+#include <cstdlib>
+#include <vector>
+#include <string>
+#include "ck/host/types.hpp"
+#include "ck/host/operation/gemm.hpp"
+#include "ck/host/device_grouped_conv_fwd_multiple_d/conv_fwd_problem.hpp"
+
+namespace ck {
+namespace host {
+namespace conv {
+
+// defines the values needed for an instance of forward convolution and functions to return
+// (templated) instances
+struct Operation_Conv_Fwd_Xdl_Cshuffle
+{
+    // returns a vector of instances given the fusion operations, uses default values for problem
+    // spec
+    static std::vector<Operation_Conv_Fwd_Xdl_Cshuffle>
+    CreateOperations(const std::string& prologue, const std::string& epilogue);
+    // returns a vector of instances, provided with a problem spec and fusion operations
+    static std::vector<Operation_Conv_Fwd_Xdl_Cshuffle> CreateOperations(
+        const Problem_Conv_Fwd& prob, const std::string& prologue, const std::string& epilogue);
+    std::size_t NumDim;
+    TensorDesc A{};
+    TensorDesc B{};
+    DataType acc               = DataType::Float;
+    DataType cs_type           = DataType::Half;
+    std::vector<TensorDesc> Ds = {};
+    TensorDesc E{};
+    std::string a_elem_op   = PassThrough;
+    std::string b_elem_op   = PassThrough;
+    std::string cde_elem_op = PassThrough;
+    std::string prologue    = "";
+    std::string epilogue    = "";
+    std::string conv_specialization =
+        "ck::tensor_operation::device::ConvolutionForwardSpecialization::Default";
+    std::string gemm_specialization =
+        "ck::tensor_operation::device::GemmSpecialization::MNKPadding";
+    // tuning parameters
+    operation::TileDesc tile_desc{};
+    operation::BlockTransferDesc a_block_transfer{};
+    operation::BlockTransferDesc b_block_transfer{};
+    operation::CShuffleDesc cshuffle{};
+    operation::CBlockTransferDesc c_block_transfer{};
+
+    // functions to update fusion operations if they are provided
+    void update_prologue(const std::string& prologue);
+    void update_epilogue(const std::string& epilogue);
+    // returns a templated instance
+    Solution ToSolution() const;
+};
+
+} // namespace conv
+} // namespace host
+} // namespace ck

codegen/include/ck/host/device_grouped_conv_fwd_multiple_d/conv_fwd_problem.hpp

@@ -0,0 +1,56 @@
+// SPDX-License-Identifier: MIT
+// Copyright (c) 2024, Advanced Micro Devices, Inc. All rights reserved.
+
+#pragma once
+
+#include <cstdlib>
+#include <vector>
+#include <memory>
+#include <sstream>
+#include <iterator>
+#include <numeric>
+#include "ck/host/types.hpp"
+
+namespace ck {
+namespace host {
+namespace conv {
+
+// defines the problem specification for a forward convolution operation
+struct Problem_Conv_Fwd
+{
+    std::size_t NumDim = 0;
+    // size of a forward convolution operation
+    std::size_t G                    = 0;
+    std::size_t N                    = 0;
+    std::size_t C                    = 0;
+    std::size_t Hi                   = 0;
+    std::size_t Wi                   = 0;
+    std::size_t Ho                   = 0;
+    std::size_t Wo                   = 0;
+    std::size_t K                    = 0;
+    std::size_t Y                    = 0;
+    std::size_t X                    = 0;
+    Layout ALayout                   = Layout::NHWGC;
+    Layout BLayout                   = Layout::GKYXC;
+    Layout ELayout                   = Layout::NHWGK;
+    std::vector<Layout> DsLayout     = {};
+    DataType ADataType               = DataType::Half;
+    DataType BDataType               = DataType::Half;
+    DataType EDataType               = DataType::Half;
+    std::vector<DataType> DsDataType = {};
+    std::string AElementOp           = "ck::tensor_operation::element_wise::PassThrough";
+    std::string BElementOp           = "ck::tensor_operation::element_wise::PassThrough";
+    std::string CDEElementOp         = "ck::tensor_operation::element_wise::PassThrough";
+
+    // returns the correct device op file for the operation
+    std::string GetIncludeHeader() const;
+
+    // returns a list of instances based on the problem spec and provided fusion operations
+    std::vector<Solution> GetSolutions(const std::string& arch,
+                                       const std::string& prologue,
+                                       const std::string& epilogue) const;
+};
+
+} // namespace conv
+} // namespace host
+} // namespace ck

codegen/include/ck/host/headers.hpp

@@ -4,7 +4,6 @@
 #pragma once
 
 #include <string>
-#include <string_view>
 #include <utility>
 #include <unordered_map>
 #include <vector>

codegen/include/ck/host/operation/gemm.hpp

@@ -1,5 +1,5 @@
 // SPDX-License-Identifier: MIT
-// Copyright (c) 2018-2023, Advanced Micro Devices, Inc. All rights reserved.
+// Copyright (c) 2018-2024, Advanced Micro Devices, Inc. All rights reserved.
 
 #pragma once
 

codegen/include/ck/host/stringutils.hpp

@@ -1,5 +1,5 @@
 // SPDX-License-Identifier: MIT
-// Copyright (c) 2018-2023, Advanced Micro Devices, Inc. All rights reserved.
+// Copyright (c) 2018-2024, Advanced Micro Devices, Inc. All rights reserved.
 
 #pragma once