Merge branch 'develop' into amd-develop

include/ck_tile/ops/gemm/kernel/gemm_kernel.hpp

@@ -11,20 +11,12 @@
 
 namespace ck_tile {
 
-template <typename TilePartitioner_,
-          typename GemmPipeline_,
-          typename EpiloguePipeline_,
-          typename LayoutA_,
-          typename LayoutB_,
-          typename LayoutC_>
+template <typename TilePartitioner_, typename GemmPipeline_, typename EpiloguePipeline_>
 struct GemmKernel
 {
     using TilePartitioner                    = remove_cvref_t<TilePartitioner_>;
     using GemmPipeline                       = remove_cvref_t<GemmPipeline_>;
     using EpiloguePipeline                   = remove_cvref_t<EpiloguePipeline_>;
-    using LayoutA                            = remove_cvref_t<LayoutA_>;
-    using LayoutB                            = remove_cvref_t<LayoutB_>;
-    using LayoutC                            = remove_cvref_t<LayoutC_>;
     static constexpr index_t KernelBlockSize = GemmPipeline::kBlockSize;
 
     using ADataType    = remove_cvref_t<typename GemmPipeline::ADataType>;
@@ -32,6 +24,10 @@ struct GemmKernel
     using CAccDataType = remove_cvref_t<typename GemmPipeline::CDataType>;
     using CODataType   = remove_cvref_t<typename EpiloguePipeline::ODataType>;
 
+    using LayoutA = remove_cvref_t<typename GemmPipeline::LayoutA>;
+    using LayoutB = remove_cvref_t<typename GemmPipeline::LayoutB>;
+    using LayoutC = remove_cvref_t<typename GemmPipeline::LayoutC>;
+
     __host__ static constexpr auto GridSize(index_t M_size, index_t N_size, index_t Batch_size)
     {
         return TilePartitioner::GridSize(M_size, N_size, Batch_size);
@@ -184,6 +180,7 @@ struct GemmKernel
             c_pad_view,
             make_tuple(number<TilePartitioner::kM>{}, number<TilePartitioner::kN>{}),
             {i_m, i_n});
+
         EpiloguePipeline{}(CBlockWindow_pad, acc);
     }
 };

include/ck_tile/ops/gemm/pipeline/block_gemm_pipeline_agmem_bgmem_creg_v1.hpp → include/ck_tile/ops/gemm/pipeline/gemm_pipeline_agmem_bgmem_creg_v1.hpp

@@ -4,15 +4,15 @@
 #pragma once
 
 #include "ck_tile/core.hpp"
-#include "ck_tile/ops/gemm/pipeline/block_gemm_pipeline_agmem_bgmem_creg_v1_default_policy.hpp"
+#include "ck_tile/ops/gemm/pipeline/gemm_pipeline_agmem_bgmem_creg_v1_default_policy.hpp"
 
 namespace ck_tile {
 
 //  A Tile Window: global memory
 //  B Tile Window: global memory
 //  C Distributed tensor: register
-template <typename Problem, typename Policy = BlockGemmPipelineAGmemBGmemCRegV1DefaultPolicy>
-struct BlockGemmPipelineAGmemBGmemCRegV1
+template <typename Problem, typename Policy = GemmPipelineAGmemBGmemCRegV1DefaultPolicy>
+struct GemmPipelineAGmemBGmemCRegV1
 {
     using ADataType      = remove_cvref_t<typename Problem::ADataType>;
     using BDataType      = remove_cvref_t<typename Problem::BDataType>;
@@ -33,6 +33,10 @@ struct BlockGemmPipelineAGmemBGmemCRegV1
     static constexpr bool kPadB = Problem::kPadB;
     static constexpr bool kPadC = Problem::kPadC;
 
+    using LayoutA = remove_cvref_t<typename Problem::LayoutA>;
+    using LayoutB = remove_cvref_t<typename Problem::LayoutB>;
+    using LayoutC = remove_cvref_t<typename Problem::LayoutC>;
+
     CK_TILE_HOST_DEVICE static constexpr ck_tile::index_t GetStaticLdsSize()
     {
         return ck_tile::integer_divide_ceil(

include/ck_tile/ops/gemm/pipeline/block_gemm_pipeline_agmem_bgmem_creg_v1_default_policy.hpp → include/ck_tile/ops/gemm/pipeline/gemm_pipeline_agmem_bgmem_creg_v1_default_policy.hpp

@@ -7,9 +7,9 @@
 
 namespace ck_tile {
 
-// Default policy for BlockGemmPipelineAGmemBGmemCRegV1
+// Default policy for GemmPipelineAGmemBGmemCRegV1
 // Default policy class should not be templated, put template on member functions instead
-struct BlockGemmPipelineAGmemBGmemCRegV1DefaultPolicy
+struct GemmPipelineAGmemBGmemCRegV1DefaultPolicy
 {
 #if 0
     // 2d

include/ck_tile/ops/gemm/pipeline/block_gemm_pipeline_agmem_bgmem_creg_v2.hpp → include/ck_tile/ops/gemm/pipeline/gemm_pipeline_agmem_bgmem_creg_v2.hpp

@@ -4,15 +4,15 @@
 #pragma once
 
 #include "ck_tile/core.hpp"
-#include "ck_tile/ops/gemm/pipeline/block_gemm_pipeline_agmem_bgmem_creg_v2_default_policy.hpp"
+#include "ck_tile/ops/gemm/pipeline/gemm_pipeline_agmem_bgmem_creg_v2_default_policy.hpp"
 
 namespace ck_tile {
 
 //  A Tile Window: global memory
 //  B Tile Window: global memory
 //  C Distributed tensor: register
-template <typename Problem, typename Policy = BlockGemmPipelineAGmemBGmemCRegV2DefaultPolicy>
-struct BlockGemmPipelineAGmemBGmemCRegV2
+template <typename Problem, typename Policy = GemmPipelineAGmemBGmemCRegV2DefaultPolicy>
+struct GemmPipelineAGmemBGmemCRegV2
 {
     using ADataType      = remove_cvref_t<typename Problem::ADataType>;
     using BDataType      = remove_cvref_t<typename Problem::BDataType>;

include/ck_tile/ops/gemm/pipeline/block_gemm_pipeline_agmem_bgmem_creg_v2_default_policy.hpp → include/ck_tile/ops/gemm/pipeline/gemm_pipeline_agmem_bgmem_creg_v2_default_policy.hpp

@@ -7,12 +7,11 @@
 
 namespace ck_tile {
 
-// Default policy for BlockGemmPipelineAGmemBGmemCRegV2
+// Default policy for GemmPipelineAGmemBGmemCRegV2
 // Default policy class should not be templated, put template on member functions instead
 // NOTE: policy should be binded to its corresponding operation. It's just a coincidence that
-//   BlockGemmPipelineAGmemBGmemCRegV2DefaultPolicy is the same as
-//   BlockGemmPipelineAGmemBGmemCRegV1DefaultPolicy
-using BlockGemmPipelineAGmemBGmemCRegV2DefaultPolicy =
-    BlockGemmPipelineAGmemBGmemCRegV1DefaultPolicy;
+//   GemmPipelineAGmemBGmemCRegV2DefaultPolicy is the same as
+//   GemmPipelineAGmemBGmemCRegV1DefaultPolicy
+using GemmPipelineAGmemBGmemCRegV2DefaultPolicy = GemmPipelineAGmemBGmemCRegV1DefaultPolicy;
 
 } // namespace ck_tile

include/ck_tile/ops/gemm/pipeline/block_gemm_pipeline_problem.hpp → include/ck_tile/ops/gemm/pipeline/gemm_pipeline_problem.hpp

@@ -13,20 +13,23 @@ template <typename ADataType_,
           typename BDataType_,
           typename CDataType_,
           typename BlockGemmShape_,
-          bool kPadA_ = false,
-          bool kPadB_ = false,
-          bool kPadC_ = false>
-struct BlockGemmPipelineProblem
+          typename TileGemmTraits_>
+struct GemmPipelineProblem
 {
     using ADataType      = remove_cvref_t<ADataType_>;
     using BDataType      = remove_cvref_t<BDataType_>;
     using CDataType      = remove_cvref_t<CDataType_>;
     using BlockGemmShape = remove_cvref_t<BlockGemmShape_>;
+    using GemmTraits     = remove_cvref_t<TileGemmTraits_>;
 
     static constexpr index_t kBlockSize = BlockGemmShape::NumWarps * get_warp_size();
-    static constexpr bool kPadA         = kPadA_;
-    static constexpr bool kPadB         = kPadB_;
-    static constexpr bool kPadC         = kPadC_;
+    static constexpr bool kPadA         = GemmTraits::kPadA;
+    static constexpr bool kPadB         = GemmTraits::kPadB;
+    static constexpr bool kPadC         = GemmTraits::kPadC;
+
+    using LayoutA = remove_cvref_t<typename GemmTraits::LayoutA>;
+    using LayoutB = remove_cvref_t<typename GemmTraits::LayoutB>;
+    using LayoutC = remove_cvref_t<typename GemmTraits::LayoutC>;
 
     static constexpr index_t AlignmentA = kPadA ? 1 : VectorLoadSize / sizeof(ADataType);
     static constexpr index_t AlignmentB = kPadB ? 1 : VectorLoadSize / sizeof(BDataType);

include/ck_tile/ops/gemm/pipeline/tile_gemm_traits.hpp

+// SPDX-License-Identifier: MIT
+// Copyright (c) 2018-2023, Advanced Micro Devices, Inc. All rights reserved.
+
+#pragma once
+
+#include "ck_tile/core.hpp"
+
+namespace ck_tile {
+
+template <bool kPadA_,
+          bool kPadB_,
+          bool kPadC_,
+          typename LayoutA_,
+          typename LayoutB_,
+          typename LayoutC_>
+struct TileGemmTraits
+{
+    static constexpr bool kPadA = kPadA_;
+    static constexpr bool kPadB = kPadB_;
+    static constexpr bool kPadC = kPadC_;
+
+    using LayoutA = LayoutA_;
+    using LayoutB = LayoutB_;
+    using LayoutC = LayoutC_;
+};
+
+} // namespace ck_tile

test/data_type/CMakeLists.txt

@@ -18,4 +18,9 @@ if(result EQUAL 0)
   target_link_libraries(test_bf8 PRIVATE utility)
 endif()
 
+add_gtest_executable(test_custom_type test_custom_type.cpp)
+if(result EQUAL 0)
+  target_link_libraries(test_custom_type PRIVATE utility)
+endif()
+
 add_gtest_executable(test_type_convert_const type_convert_const.cpp)