Find 2.0 fusion

src/fusion.cpp

@@ -120,14 +120,12 @@ miopenStatus_t ConvBiasActivFusion(Handle& handle,
 }
 
 static auto
-AllocateBuffersAndMakeFusionInvokeParams(const FusionContext& context,
+AllocateBuffersAndMakeFusionInvokeParams(Handle& handle,
                                          const FusionDescription& problem,
                                          std::vector<Allocator::ManageDataPtr>& invoke_bufs,
                                          miopen::OperatorArgs& params,
                                          const FusionPlanDescriptor& plan)
 {
-    auto& handle = context.GetStream();
-
     const auto allocate_buffer = [&](std::size_t size) {
         auto ptr = handle.Create(size);
         auto ret = ptr.get();
@@ -707,6 +705,29 @@ static auto GetFusedWinogradSolvers()
                                    solver::fusion::ConvBinWinogradRxSf2x3g1Fused>{};
 }
 
+static auto GetAllFusionSolvers()
+{
+    return GetFusedNonConvSolvers() + GetFusedDirectSolvers() + GetFusedIGemmSolvers() +
+           GetFusedWinogradSolvers();
+}
+
+solver::ConvSolution MakeFusedSolution(const FusionContext& ctx,
+                                       solver::Id id,
+                                       const std::optional<std::string>& perf_cfg_override,
+                                       const FusionDescription& problem,
+                                       const AnyInvokeParams& invoke_params)
+{
+    decltype(auto) db = GetDb(ctx);
+    solver::ConvSolution solution{miopenStatusInternalError};
+
+    GetAllFusionSolvers().FindById(id, [&](auto solver) {
+        solution = miopen::solver::FindSolution(
+            solver, ctx, problem, db, invoke_params, perf_cfg_override.value_or(""));
+    });
+
+    return solution;
+}
+
 struct FusionFindParameters : PrimitiveFindParameters
 {
 };
@@ -761,47 +782,48 @@ static const std::vector<std::unique_ptr<ISolversFinder>>& GetFusionSolverFinder
     return finders;
 }
 
-miopenStatus_t FusionPlanDescriptor::Compile(Handle& handle)
+static std::vector<PerfField>
+FindFusion(const ExecutionContext& ctx,
+           const FusionDescription& fusion_problem,
+           const std::function<fusion::FusionInvokeParams()>& invoke_params)
 {
-    auto fusion_ctx     = FusionContext{handle};
-    auto fusion_problem = FusionDescription{this};
-    const FindEnforce enforce;
-
-    // sols is a collection of ConvSolutions that have been returned from Find for the
-    // fusion_problem. These ConvSolutions store instructions on how to build kernels and an invoker
-    // factory.
-    std::vector<miopen::solver::ConvSolution> sols;
-
-    auto find_results = UserFindDbRecord::TryLoad(
-        handle,
+    return UserFindDbRecord::TryLoad(
+        ctx.GetStream(),
         fusion_problem,
         [&](DbRecord& record) {
             // fusion_ctx.use_dynamic_solutions_only = findMode.IsDynamicHybrid(fusion_ctx);
 
-            // We need buffers for find, thus we allocate them.
-            miopen::OperatorArgs params;
-            std::vector<Allocator::ManageDataPtr> invoke_bufs;
-            const auto invoke_params = AllocateBuffersAndMakeFusionInvokeParams(
-                fusion_ctx, fusion_problem, invoke_bufs, params, *this);
-
-            FindCore(invoke_params,
+            // We need buffers for find, thus we lazily get them, possibly allocating.
+            auto fusion_ctx = FusionContext(ctx.GetStream());
+            FindCore(invoke_params(),
                      record,
                      fusion_ctx,
                      fusion_problem,
                      FusionFindParameters{},
                      GetFusionSolverFinders());
         },
         "fusion");
+}
+
+miopenStatus_t FusionPlanDescriptor::Compile(Handle& handle)
+{
+    std::vector<Allocator::ManageDataPtr> invoke_bufs;
+    miopen::OperatorArgs params;
+
+    const auto find_results = Find(handle, [&]() {
+        return AllocateBuffersAndMakeFusionInvokeParams(
+            handle, FusionDescription{this}, invoke_bufs, params, *this);
+    });
 
-    const auto network_config = fusion_problem.MakeNetworkConfig();
+    const auto network_config = FusionDescription{this}.MakeNetworkConfig();
 
     for(const auto& result : find_results)
     {
         if(conv_fwd_algo && result.algorithm != "fusion" &&
            miopen::StringToConvolutionFwdAlgo(result.algorithm) != *conv_fwd_algo)
             continue;
-        const auto id = solver::Id{result.solver_id};
 
+        const auto id      = solver::Id{result.solver_id};
         const auto invoker = handle.GetInvoker(network_config, id);
 
         if(!invoker)
@@ -823,6 +845,13 @@ miopenStatus_t FusionPlanDescriptor::Compile(Handle& handle)
     return miopenStatusSuccess;
 }
 
+std::vector<struct PerfField>
+FusionPlanDescriptor::Find(Handle& handle,
+                           const std::function<fusion::FusionInvokeParams()>& invoke_params) const
+{
+    return FindFusion(&handle, this, invoke_params);
+}
+
 miopenStatus_t FusionPlanDescriptor::Execute(const Handle& handle,
                                              const TensorDescriptor& inputDesc,
                                              ConstData_t input,

src/include/miopen/find_solution.hpp

@@ -173,6 +173,29 @@ ConvSolution FindSolution(Solver s,
 template <class... Solvers>
 struct SolverContainer
 {
+    template <class... SolversRight>
+    auto operator+(SolverContainer<SolversRight...>) const
+    {
+        return SolverContainer<Solvers..., SolversRight...>{};
+    }
+
+    ///\todo: remove when AnySolver would be able to work with non-conv solvers
+    template <class Functor>
+    void FindById(solver::Id id, Functor&& receiver)
+    {
+        bool found = false;
+
+        miopen::each_args(
+            [&](auto solver) {
+                if(found || id != solver::Id{solver.SolverDbId()})
+                    return;
+
+                found = true;
+                receiver(solver);
+            },
+            Solvers{}...);
+    }
+
     // Search for all applicable solutions among many solvers
     template <class Context, class Problem, class Db, class Solution = miopen::solver::ConvSolution>
     std::vector<Solution>

src/include/miopen/fusion.hpp

@@ -253,6 +253,12 @@ miopenStatus_t ConvBiasActivFusion(Handle& handle,
                                    const TensorDescriptor& yDesc,
                                    Data_t y);
 
+solver::ConvSolution MakeFusedSolution(const struct FusionContext& ctx,
+                                       solver::Id id,
+                                       const std::optional<std::string>& perf_cfg_override,
+                                       const struct FusionDescription& problem,
+                                       const AnyInvokeParams& invoke_params);
+
 } // namespace miopen
 MIOPEN_DEFINE_OBJECT(miopenFusionOpDescriptor, miopen::FusionOpDescriptor);
 MIOPEN_DEFINE_OBJECT(miopenOperatorArgs, miopen::OperatorArgs);

src/include/miopen/fusion_plan.hpp

@@ -57,6 +57,8 @@ struct FusionPlanDescriptor : miopenFusionPlanDescriptor
                            Data_t output,
                            const OperatorArgs& op_args);
     miopenStatus_t Compile(Handle& handle);
+    std::vector<struct PerfField>
+    Find(Handle& handle, const std::function<fusion::FusionInvokeParams()>& invoke_params) const;
     friend std::ostream& operator<<(std::ostream& stream, const FusionPlanDescriptor& fpd);
 
     miopenStatus_t

src/include/miopen/problem.hpp

@@ -41,6 +41,9 @@
 
 #include <cstring>
 #include <unordered_map>
+#include "miopen/fusion/fusion_op_args.hpp"
+#include "miopen/fusion/fusion_invoke_params.hpp"
+#include "fusion_plan.hpp"
 
 namespace miopen {
 
@@ -138,15 +141,14 @@ struct Problem
     std::unordered_map<miopenTensorArgumentId_t, TensorDescriptor> tensor_descriptors;
     OperatorDescriptor operator_descriptor;
 
-    using AllocatedBuffers = std::unordered_map<miopenTensorArgumentId_t, Data_t>;
+    using Buffers = std::unordered_map<miopenTensorArgumentId_t, Data_t>;
 
     std::vector<Solution> FindSolutionsImpl(Handle& handle,
                                             const FindOptions& options,
                                             std::size_t max_solutions,
-                                            const AllocatedBuffers& buffers,
+                                            const Buffers& buffers,
                                             const ConvolutionDescriptor& conv_desc) const;
 
-    void TransposeImpl(const ConvolutionDescriptor& conv_desc);
     void LogDriverCommand(const ConvolutionDescriptor& conv_desc) const;
     void LogDriverCommand(const ActivationDescriptor& descriptor) const;
 };
@@ -160,19 +162,51 @@ struct FusedProblem
         // Not implemented, but silently
     }
 
-    std::vector<Solution> FindSolutions(Handle& /*handle*/,
-                                        const FindOptions& /*options*/,
-                                        std::size_t /*max_solutions*/) const
+    [[nodiscard]] std::vector<Solution>
+    FindSolutions(Handle& handle, const FindOptions& options, std::size_t max_solutions) const;
+
+    void PropagateDescriptors();
+
+    [[nodiscard]] miopenTensorArgumentId_t GetInputId() const
     {
-        MIOPEN_THROW(miopenStatusNotImplemented);
+        return problems.front().GetInputId();
     }
 
-    void PropagateDescriptors();
+    [[nodiscard]] miopenTensorArgumentId_t GetOutputId() const
+    {
+        return problems.back().GetOutputId();
+    }
+
+    [[nodiscard]] const TensorDescriptor& GetInput() const { return problems.front().GetInput(); }
+    [[nodiscard]] const TensorDescriptor& GetOutput() const { return problems.back().GetOutput(); }
+
+    [[nodiscard]] FusionPlanDescriptor AsFusionPlan() const;
+
+    friend void to_json(nlohmann::json& j, const FusedProblem& problem);
+    friend void from_json(const nlohmann::json& j, FusedProblem& problem);
+
+    [[nodiscard]] fusion::FusionInvokeParams
+    MakeInvokeParams(const std::function<Data_t(miopenTensorArgumentId_t, const TensorDescriptor&)>&
+                         buffer_getter,
+                     OperatorArgs& operator_args) const;
+
+private:
+    static void AddProblemToPlan(struct FusionPlanDescriptor& plan, const Problem& problem);
 };
 
 struct ProblemContainer : miopenProblem
 {
-    boost::variant<Problem, FusedProblem> item;
+    using Item = boost::variant<Problem, FusedProblem>;
+    Item item;
+
+    ProblemContainer() = default;
+    ProblemContainer(Item item_) // NOLINT(*-explicit-constructor)
+        : item(std::move(item_))
+    {
+    }
+
+    friend void to_json(nlohmann::json& j, const ProblemContainer& problem);
+    friend void from_json(const nlohmann::json& j, ProblemContainer& problem);
 };
 
 } // namespace miopen

src/include/miopen/search_options.hpp

@@ -28,6 +28,7 @@
 
 #include <miopen/miopen.h>
 
+#include <miopen/common.hpp>
 #include <miopen/object.hpp>
 
 #include <limits>

src/include/miopen/solution.hpp

@@ -74,6 +74,8 @@ struct Solution : miopenSolution
             if(argument.descriptor != nullptr)
                 descriptor = miopen::deref(*argument.descriptor);
         }
+
+        inline RunInput(Data_t buffer_) : buffer(buffer_) {}
     };
 
     float GetTime() const { return time; }
@@ -83,8 +85,8 @@ struct Solution : miopenSolution
     const solver::Id& GetSolver() const { return solver; }
     void SetSolver(solver::Id value) { solver = value; }
     void SetPerfConfig(const std::optional<std::string>& cfg) { perf_cfg = cfg; }
-    const Problem& GetProblem() const { return problem; }
-    void SetProblem(Problem value) { problem = std::move(value); }
+    const ProblemContainer& GetProblem() const { return problem; }
+    void SetProblem(ProblemContainer value) { problem = std::move(value); }
 
     void Run(Handle& handle,
              const std::unordered_map<miopenTensorArgumentId_t, RunInput>& inputs,
@@ -100,7 +102,7 @@ struct Solution : miopenSolution
     float time                     = 0;
     std::size_t workspace_required = 0;
     solver::Id solver;
-    Problem problem;
+    ProblemContainer problem;
     std::optional<std::string> perf_cfg = std::nullopt;
 
     void RunImpl(Handle& handle,
@@ -109,8 +111,19 @@ struct Solution : miopenSolution
                  std::size_t workspace_size,
                  const ConvolutionDescriptor& conv_desc);
 
+    void RunImpl(Handle& handle,
+                 const std::unordered_map<miopenTensorArgumentId_t, RunInput>& inputs,
+                 Data_t workspace,
+                 std::size_t workspace_size,
+                 const FusedProblem& problem_);
+
     static Problem Transpose(const Problem& problem, RunInput* x, const RunInput& w, RunInput* y);
-    void LogDriverCommand(const ConvolutionDescriptor& conv_desc) const;
+
+    void LogDriverCommand(const ConvolutionDescriptor& desc) const;
+    void LogDriverCommand(const ActivationDescriptor& desc) const;
+
+    void LogDriverCommand(const Problem& problem_) const;
+    void LogDriverCommand(const FusedProblem& problem_) const;
 };
 
 } // namespace miopen