[CK Tile] Unification work - mma transformations pipeline #5508
+883
−285
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,161 @@ | ||
| // Copyright (c) Advanced Micro Devices, Inc., or its affiliates. | ||
| // SPDX-License-Identifier: MIT | ||
| #pragma once | ||
| #include "ck_tile/core/arch/arch.hpp" | ||
| #include "ck_tile/core/numeric/vector_type.hpp" | ||
|
|
||
| #include "amdgcn_mma.hpp" | ||
| #include "mma_selector.hpp" | ||
| #include "mma_traits.hpp" | ||
| #include "mma_transforms.hpp" | ||
|
|
||
| namespace ck_tile::core::arch::mma { | ||
|
|
||
| enum struct MmaPipelineOptionFlag | ||
| { | ||
| NONE = 0x0, | ||
| C_TRANSPOSE = 0x1, | ||
| COMPRESS_A = 0x2, | ||
| }; | ||
|
|
||
| struct MmaPipelineOptionFlags | ||
|
Contributor
There was a problem hiding this comment. Very verbose but I guess necessary if we really don't want to allow raw enums... |
||
| { | ||
| using Type = std::underlying_type<MmaPipelineOptionFlag>::type; | ||
|
|
||
| explicit constexpr MmaPipelineOptionFlags() : mFlags(0) {} | ||
| explicit constexpr MmaPipelineOptionFlags(Type value) : mFlags(value) {} | ||
| constexpr MmaPipelineOptionFlags(MmaPipelineOptionFlag singleFlag) : mFlags(toType(singleFlag)) | ||
| { | ||
| } | ||
| constexpr MmaPipelineOptionFlags(const MmaPipelineOptionFlags& original) | ||
| : mFlags(original.mFlags) | ||
| { | ||
| } | ||
|
|
||
| constexpr MmaPipelineOptionFlags& operator|=(MmaPipelineOptionFlag addValue) | ||
| { | ||
| mFlags |= toType(addValue); | ||
| return *this; | ||
| } | ||
| constexpr MmaPipelineOptionFlags operator|(MmaPipelineOptionFlag addValue) const | ||
| { | ||
| MmaPipelineOptionFlags result(*this); | ||
| result |= addValue; | ||
| return result; | ||
| } | ||
| constexpr MmaPipelineOptionFlags& operator&=(MmaPipelineOptionFlag maskValue) | ||
| { | ||
| mFlags &= toType(maskValue); | ||
| return *this; | ||
| } | ||
| constexpr MmaPipelineOptionFlags operator&(MmaPipelineOptionFlag maskValue) const | ||
| { | ||
| MmaPipelineOptionFlags result(*this); | ||
| result &= maskValue; | ||
| return result; | ||
| } | ||
| constexpr MmaPipelineOptionFlags operator~() const | ||
| { | ||
| MmaPipelineOptionFlags result(*this); | ||
| result.mFlags = ~result.mFlags; | ||
| return result; | ||
| } | ||
| constexpr bool testFlag(MmaPipelineOptionFlag flag) const | ||
| { | ||
| return (flag == MmaPipelineOptionFlag::NONE) ? mFlags == toType(flag) : *this & flag; | ||
| } | ||
| constexpr operator bool() const { return mFlags != toType(MmaPipelineOptionFlag::NONE); } | ||
| constexpr bool operator==(Type rhs) const { return mFlags == rhs; } | ||
|
|
||
| private: | ||
| Type mFlags; | ||
| static constexpr Type toType(MmaPipelineOptionFlag f) { return static_cast<Type>(f); } | ||
| }; | ||
|
|
||
| constexpr bool operator==(MmaPipelineOptionFlags::Type lhs, const MmaPipelineOptionFlags& rhs) | ||
| { | ||
| return rhs == lhs; | ||
| } | ||
|
|
||
| // TODO: c++20: use MmaPipelineOptionFlags directly | ||
| template <MmaPipelineOptionFlags::Type Flags_, typename Derived> | ||
| struct MmaPipelineBase | ||
| { | ||
| static constexpr auto Flags = MmaPipelineOptionFlags(Flags_); | ||
|
|
||
| private: | ||
| template <typename DstT, typename SrcT> | ||
| CK_TILE_DEVICE static auto formatBuffer(SrcT&& inputBuffer) | ||
| { | ||
| // TODO: Implement formatting logic as needed. | ||
| // This is intended to convert input fragments to the native vector types | ||
| // required by the BlockWiseMma operation for iteration | ||
| static_assert(sizeof(DstT) == sizeof(std::remove_reference_t<SrcT>), | ||
| "Size mismatch in formatBuffer"); | ||
|
|
||
| using QualifiedDstT = | ||
| std::conditional_t<std::is_const_v<std::remove_reference_t<SrcT>>, DstT const, DstT>; | ||
|
|
||
| return reinterpret_cast<QualifiedDstT&>(inputBuffer); | ||
| } | ||
|
|
||
| protected: | ||
| template <MmaPipelineOptionFlag Flag> | ||
| constexpr CK_TILE_DEVICE static bool hasFlag() | ||
| { | ||
| return Flags.testFlag(Flag); | ||
| } | ||
|
|
||
| template <typename DstT, typename Transform, typename... Args> | ||
| CK_TILE_DEVICE static auto preApplyTransform(Args&&... args) | ||
| { | ||
| return formatBuffer<DstT>(Transform::exec(std::forward<Args>(args)...)); | ||
| } | ||
|
|
||
| template <typename DstT, typename Transform, typename... Args> | ||
| CK_TILE_DEVICE static auto postApplyTransform(Args&&... args) | ||
| { | ||
| return Transform::exec(formatBuffer<DstT>(std::forward<Args>(args)...)); | ||
| } | ||
|
|
||
| public: | ||
| template <typename VecTA, typename VecTB, typename VecTC> | ||
| CK_TILE_DEVICE static decltype(auto) exec(VecTA&& a, VecTB&& b, VecTC&& accum) | ||
| { | ||
| if constexpr(MmaOpTraits<typename Derived::FragWiseMmaOp>::IsSupported) | ||
| { | ||
| // TODO: c++20: Call template functions with MmaPipelineOptionFlags directly | ||
| auto pre = Derived::template preApply<Flags_>( | ||
| hasFlag<MmaPipelineOptionFlag::C_TRANSPOSE>() ? std::forward<VecTB>(b) | ||
|
Contributor
There was a problem hiding this comment. this works only if A and B have the same type and size
Contributor
There was a problem hiding this comment. Yeah, CTranspose will not be available for all intrinsics. Also I don't think CTranspose is possible for sparse intrinsics.
Contributor
Author
There was a problem hiding this comment. I'll disable them for sparse then! |
||
| : std::forward<VecTA>(a), | ||
| hasFlag<MmaPipelineOptionFlag::C_TRANSPOSE>() ? std::forward<VecTA>(a) | ||
| : std::forward<VecTB>(b), | ||
| std::forward<VecTC>(accum)); | ||
| Derived::execImpl(pre); | ||
| return Derived::template postApply<Flags_>(std::move(pre)); | ||
| } | ||
| else | ||
| { | ||
| // Return the unsupported exec. This should print a runtime warning. (amdgcn_mma.hpp) | ||
| // Code should not reach here, but HOST/DEVICE compile passes are | ||
| // weirdly intertwined and instead of having constexpr in the calling | ||
| // site (tests) we do this. See also changes by this commit. | ||
| return Derived::FragWiseMmaOp::exec({}, {}, {}); | ||
| } | ||
| } | ||
|
Comment on lines
+122
to
+145
Contributor
There was a problem hiding this comment. Aha we are making a second-order wrapper for the intrinsic just like in CK Tile, making more sense to me now. |
||
| }; | ||
|
|
||
| #if CK_TILE_CONCEPTS && CK_TILE_CONCEPTS_HEADER | ||
|
|
||
| #include <concepts> | ||
|
|
||
| /** | ||
| * @concept MmaPipelineI | ||
| * @brief Expresses the meta-data interface required for a CRTP MmaPipeline. | ||
| */ | ||
| template <typename Derived, MmaPipelineOptionFlags::Type Flags> | ||
| concept MmaPipelineInterface = std::derived_from<Derived, MmaPipelineBase<Flags, Derived>>; | ||
|
|
||
| #endif // CK_TILE_CONCEPTS && CK_TILE_CONCEPTS_HEADER | ||
|
|
||
| } // namespace ck_tile::core::arch::mma | ||
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Uh oh!
There was an error while loading. Please reload this page.