diff --git a/csrc/mma_type.h b/csrc/mma_type.h index 60bf608df4b..5fbbd1c8a99 100644 --- a/csrc/mma_type.h +++ b/csrc/mma_type.h @@ -33,14 +33,12 @@ enum class MatmulDomain { M = 0, N, K, Batch }; //! INPUT_B - a producer of MMA input B //! INPUT_C - a producer of a tensor used in fusion epilogue, //! for example tensor used in beta scaling fusion -//! OUTPUT_D - the main consumer of MMA op results -//! OUTPUT_AUX - fusion outputs that are consumers of OUTPUT_D +//! OUTPUT_D - fusion outputs that have the matmul as a dependency //! //! Naming convention is based on the following formula: //! D = alpha * A x B + beta * C -//! AUX = relu(D) //! Note: bias vector tensors will be assigned to INPUT_C role. -enum class MatmulRole { INPUT_A = 0, INPUT_B, OUTPUT_D, INPUT_C, OUTPUT_AUX }; +enum class MatmulRole { INPUT_A = 0, INPUT_B, INPUT_C, OUTPUT_D }; //! The expected number of occurances of core TensorView roles in fusion static constexpr size_t MATMUL_CORE_ROLES_EXPECTED_COUNT = 1; diff --git a/csrc/scheduler/matmul_utils.cpp b/csrc/scheduler/matmul_utils.cpp index 015937046a0..894e776b594 100644 --- a/csrc/scheduler/matmul_utils.cpp +++ b/csrc/scheduler/matmul_utils.cpp @@ -273,12 +273,6 @@ std::string isMatmulFusionDefinitionSupported( tvs_with_roles.insert(entry->second.begin(), entry->second.end()); } - // Non-core output roles are optional, no requirements for definitions - entry = roles_map.find(MatmulRole::OUTPUT_AUX); - if (entry != roles_map.end()) { - tvs_with_roles.insert(entry->second.begin(), entry->second.end()); - } - const auto in_out_tvs_count = fusion_inputs_tvs.size() + fusion_outputs_tvs.size(); if (in_out_tvs_count != tvs_with_roles.size()) { diff --git a/csrc/scheduler/mma_utils.cpp b/csrc/scheduler/mma_utils.cpp index 1b05d80bccf..e7ad546d1c7 100644 --- a/csrc/scheduler/mma_utils.cpp +++ b/csrc/scheduler/mma_utils.cpp @@ -1231,27 +1231,12 @@ RolesMapOpt getTensorRoles( } } - // NOTE: sort output roles in descending order by uses() size, and - // if equal then by name() to ensure the stable ordering of tensor - // views in collections assigned to the supported roles - std::sort(storage.begin(), storage.end(), [](TensorView* a, TensorView* b) { - return (a->uses().size() == b->uses().size()) - ? (a->name() < b->name()) - : (a->uses().size() > b->uses().size()); - }); - if (!storage.empty()) { - // NOTE: currently, we pick as a reference tensor one with `m` and `n` - // IterDomains and the most uses - auto pos = storage.begin(); - tensor_roles[MatmulRole::OUTPUT_D].push_back(*pos); - for (++pos; pos != storage.end(); ++pos) { - tensor_roles[MatmulRole::OUTPUT_AUX].push_back(*pos); - } + tensor_roles[MatmulRole::OUTPUT_D] = storage; } for (auto& [role, tvs] : tensor_roles) { - // NOTE: sort input roles in descending order by uses() size, and + // NOTE: sort role tvs in descending order by uses() size, and // if equal then by name() to ensure the stable ordering of tensor // views in collections assigned to the supported roles std::sort(tvs.begin(), tvs.end(), [](TensorView* a, TensorView* b) {