Add an argument to enable mixed precision

32880bcb · Rostyslav Geyyer · 30d8a363 · 32880bcb
Commit 32880bcb authored Jul 31, 2023 by Rostyslav Geyyer
Show whitespace changes
Inline Side-by-side

Showing with 6 additions and 3 deletions

include/ck/tensor_operation/gpu/device/impl/device_gemm_xdl_cshuffle.hpp ...or_operation/gpu/device/impl/device_gemm_xdl_cshuffle.hpp +6 -3

No files found.
--- a/include/ck/tensor_operation/gpu/device/impl/device_gemm_xdl_cshuffle.hpp
+++ b/include/ck/tensor_operation/gpu/device/impl/device_gemm_xdl_cshuffle.hpp
@@ -65,7 +65,8 @@ template <typename ALayout,
          typename CShuffleBlockTransferClusterLengths_MBlock_MPerBlock_NBlock_NPerBlock,
          index_t CShuffleBlockTransferScalarPerVector_NPerBlock,
          LoopScheduler LoopSched     = make_default_loop_scheduler(),
-          PipelineVersion PipelineVer = PipelineVersion::v1>
+          PipelineVersion PipelineVer = PipelineVersion::v1,
+          bool EnableMixedPrecision   = false>
 struct DeviceGemm_Xdl_CShuffle : public DeviceGemm<ALayout,
                                                   BLayout,
                                                   CLayout,
@@ -87,7 +88,8 @@ struct DeviceGemm_Xdl_CShuffle : public DeviceGemm<ALayout,
        ALayout,
        BLayout,
        CLayout,
-        ADataType, // TODO: distinguish A/B datatype
+        ADataType,
+        BDataType,
        GemmAccDataType,
        CShuffleDataType,
        CDataType,
@@ -128,7 +130,8 @@ struct DeviceGemm_Xdl_CShuffle : public DeviceGemm<ALayout,
        CShuffleBlockTransferClusterLengths_MBlock_MPerBlock_NBlock_NPerBlock,
        CShuffleBlockTransferScalarPerVector_NPerBlock,
        LoopSched,
-        PipelineVer>;
+        PipelineVer,
+        EnableMixedPrecision>;

    using Argument = typename GridwiseGemm::Argument;