Merge remote-tracking branch 'origin/develop' into migx-flash-attn

include/ck/utility/math.hpp

@@ -240,5 +240,21 @@ struct less
     __host__ __device__ constexpr bool operator()(T x, T y) const { return x < y; }
 };
 
+template <index_t X>
+__host__ __device__ constexpr auto next_power_of_two()
+{
+    // TODO: X need to be 2 ~ 0x7fffffff. 0, 1, or larger than 0x7fffffff will compile fail
+    constexpr index_t Y = 1 << (32 - __builtin_clz(X - 1));
+    return Y;
+}
+
+template <index_t X>
+__host__ __device__ constexpr auto next_power_of_two(Number<X> x)
+{
+    // TODO: X need to be 2 ~ 0x7fffffff. 0, 1, or larger than 0x7fffffff will compile fail
+    constexpr index_t Y = 1 << (32 - __builtin_clz(x.value - 1));
+    return Number<Y>{};
+}
+
 } // namespace math
 } // namespace ck

include/ck/utility/random_gen.hpp

+// SPDX-License-Identifier: MIT
+// Copyright (c) 2018-2023, Advanced Micro Devices, Inc. All rights reserved.
+
+#pragma once
+
+namespace ck {
+
+// Pseudo random number generator
+// version for fp32
+template <typename T, uint32_t seed_t, std::enable_if_t<std::is_same<float, T>{}, bool> = false>
+__host__ __device__ uint32_t prand_generator(index_t id, T val, uint32_t seed = seed_t)
+{
+    uint32_t x         = *(reinterpret_cast<uint32_t*>(&val));
+    uint32_t drop_bits = uint32_t(x) & 0xFFFFu;
+    drop_bits ^= x >> 16;
+    drop_bits = ((drop_bits & 31) << 11) | (drop_bits >> 5);
+    drop_bits *= 0x7000149;
+    // NOTE: If id is in 64 bit, we are only using lower 32 bit.
+    //       So, it can have an effect of using same id for multiple elements when the id is very
+    //       large!
+    uint32_t rng = (drop_bits ^ 0x13371337 ^ (id * 229791) ^ seed);
+    return rng;
+}
+
+// version for fp16
+template <typename T, uint32_t seed_t, std::enable_if_t<std::is_same<half_t, T>{}, bool> = false>
+__host__ __device__ uint32_t prand_generator(index_t id, T val, uint32_t seed = seed_t)
+{
+    uint16_t x         = *(reinterpret_cast<uint16_t*>(&val));
+    uint32_t drop_bits = uint32_t(x) & 0xFFFFu;
+    drop_bits          = ((drop_bits & 31) << 11) | (drop_bits >> 5);
+    drop_bits *= 0x7000149;
+    // NOTE: If id is in 64 bit, we are only using lower 32 bit.
+    //       So, it can have an effect of using same id for multiple elements when the id is very
+    //       large!
+    uint32_t rng = (drop_bits ^ 0x13371337 ^ (id * 229791) ^ seed);
+    return rng;
+}
+
+// return 0 if data is not fp16 or fp32
+template <typename T,
+          uint32_t seed_t,
+          std::enable_if_t<!(std::is_same<float, T>{} || std::is_same<half_t, T>{}), bool> = false>
+__host__ __device__ uint32_t prand_generator(int id, T val, uint32_t seed = seed_t)
+{
+    std::ignore = id;
+    std::ignore = val;
+    std::ignore = seed;
+
+    return 0;
+}
+
+} // namespace ck

include/ck/utility/reduction_common.hpp

@@ -25,16 +25,4 @@ struct float_equal_zero
     };
 };
 
-template <index_t N>
-static constexpr __device__ index_t get_shift()
-{
-    return (get_shift<N / 2>() + 1);
-};
-
-template <>
-constexpr __device__ index_t get_shift<1>()
-{
-    return (0);
-}
-
 } // namespace ck

include/ck/utility/reduction_operator.hpp

@@ -6,6 +6,7 @@
 #include "ck/ck.hpp"
 #include "ck/utility/data_type.hpp"
 #include "ck/utility/type.hpp"
+#include "ck/utility/type_convert.hpp"
 
 namespace ck {
 

include/ck/utility/type_convert.hpp

+// SPDX-License-Identifier: MIT
+// Copyright (c) 2018-2023, Advanced Micro Devices, Inc. All rights reserved.
+
+#pragma once
+
+#include "ck/utility/data_type.hpp"
+#include "ck/utility/f8_utils.hpp"
+#include "ck/utility/random_gen.hpp"
+
+namespace ck {
+
+// Convert X to Y
+template <typename Y, typename X>
+__host__ __device__ constexpr Y type_convert(X x)
+{
+    static_assert(!std::is_reference_v<Y> && !std::is_reference_v<X>);
+
+    return static_cast<Y>(x);
+}
+
+// convert bfp16 to fp32
+template <>
+inline __host__ __device__ constexpr float type_convert<float, bhalf_t>(bhalf_t x)
+{
+    union
+    {
+        uint32_t int32;
+        float fp32;
+    } u = {uint32_t(x) << 16};
+
+    return u.fp32;
+}
+
+// convert fp32 to bfp16
+template <>
+inline __host__ __device__ constexpr bhalf_t type_convert<bhalf_t, float>(float x)
+{
+    union
+    {
+        float fp32;
+        uint32_t int32;
+    } u = {x};
+
+    return uint16_t(u.int32 >> 16);
+}
+
+// convert bfp16 to fp16 via fp32
+template <>
+inline __host__ __device__ constexpr half_t type_convert<half_t, bhalf_t>(bhalf_t x)
+{
+    float x_fp32 = type_convert<float>(x);
+
+    return static_cast<half_t>(x_fp32);
+}
+
+// convert fp16 to bfp16 via fp32
+template <>
+inline __host__ __device__ constexpr bhalf_t type_convert<bhalf_t, half_t>(half_t x)
+{
+    float x_fp32 = static_cast<float>(x);
+
+    return type_convert<bhalf_t>(x_fp32);
+}
+
+// convert bfp16 to int8 via fp32
+template <>
+inline __host__ __device__ constexpr int8_t type_convert<int8_t, bhalf_t>(bhalf_t x)
+{
+    float x_fp32 = type_convert<float>(x);
+
+    return static_cast<int8_t>(x_fp32);
+}
+
+// convert int8 to bfp16 via fp32
+template <>
+inline __host__ __device__ constexpr bhalf_t type_convert<bhalf_t, int8_t>(int8_t x)
+{
+    float x_fp32 = static_cast<float>(x);
+
+    return type_convert<bhalf_t>(x_fp32);
+}
+
+// convert fp32 to fp8
+template <>
+inline __host__ __device__ f8_t type_convert<f8_t, float>(float x)
+{
+    constexpr bool negative_zero_nan = true;
+    constexpr bool clip              = true;
+    constexpr f8_rounding_mode rm    = f8_rounding_mode::standard;
+    constexpr uint32_t rng           = 0;
+    return utils::cast_to_f8<float, negative_zero_nan, clip, (rm == f8_rounding_mode::stochastic)>(
+        x, rng);
+}
+
+// convert fp8 to fp32
+template <>
+inline __host__ __device__ float type_convert<float, f8_t>(f8_t x)
+{
+    constexpr bool negative_zero_nan = true;
+    return utils::cast_from_f8<float, negative_zero_nan>(x);
+}
+
+// convert fp16 to fp8
+template <>
+inline __host__ __device__ f8_t type_convert<f8_t, half_t>(half_t x)
+{
+    constexpr bool negative_zero_nan = true;
+    constexpr bool clip              = true;
+    constexpr f8_rounding_mode rm    = f8_rounding_mode::standard;
+    constexpr uint32_t rng           = 0;
+    return utils::cast_to_f8<half_t, negative_zero_nan, clip, (rm == f8_rounding_mode::stochastic)>(
+        x, rng);
+}
+
+// convert fp8 to fp16
+template <>
+inline __host__ __device__ half_t type_convert<half_t, f8_t>(f8_t x)
+{
+    constexpr bool negative_zero_nan = true;
+    return utils::cast_from_f8<half_t, negative_zero_nan>(x);
+}
+
+// Declare a template function for bf16 conversion using RTN
+template <typename Y, typename X>
+__host__ __device__ constexpr Y bf16_convert_rtn(X x);
+
+// Convert fp32 to bf16 with RTN if higher precision is needed
+template <>
+inline __host__ __device__ constexpr bhalf_t bf16_convert_rtn<bhalf_t, float>(float x)
+{
+    union
+    {
+        float fp32;
+        uint32_t int32;
+    } u = {x};
+
+    // When the exponent bits are not all 1s, then the value is zero, normal,
+    // or subnormal. We round the bfloat16 mantissa up by adding 0x7FFF, plus
+    // 1 if the least significant bit of the bfloat16 mantissa is 1 (odd).
+    // This causes the bfloat16's mantissa to be incremented by 1 if the 16
+    // least significant bits of the float mantissa are greater than 0x8000,
+    // or if they are equal to 0x8000 and the least significant bit of the
+    // bfloat16 mantissa is 1 (odd). This causes it to be rounded to even when
+    // the lower 16 bits are exactly 0x8000. If the bfloat16 mantissa already
+    // has the value 0x7f, then incrementing it causes it to become 0x00 and
+    // the exponent is incremented by one, which is the next higher FP value
+    // to the unrounded bfloat16 value. When the bfloat16 value is subnormal
+    // with an exponent of 0x00 and a mantissa of 0x7f, it may be rounded up
+    // to a normal value with an exponent of 0x01 and a mantissa of 0x00.
+    // When the bfloat16 value has an exponent of 0xFE and a mantissa of 0x7F,
+    // incrementing it causes it to become an exponent of 0xFF and a mantissa
+    // of 0x00, which is Inf, the next higher value to the unrounded value.
+    bool flag0 = ~u.int32 & 0x7f800000;
+
+    // When all of the exponent bits are 1, the value is Inf or NaN.
+    // Inf is indicated by a zero mantissa. NaN is indicated by any nonzero
+    // mantissa bit. Quiet NaN is indicated by the most significant mantissa
+    // bit being 1. Signaling NaN is indicated by the most significant
+    // mantissa bit being 0 but some other bit(s) being 1. If any of the
+    // lower 16 bits of the mantissa are 1, we set the least significant bit
+    // of the bfloat16 mantissa, in order to preserve signaling NaN in case
+    // the bfloat16's mantissa bits are all 0.
+    bool flag1 = !flag0 && (u.int32 & 0xffff);
+
+    u.int32 += flag0 ? 0x7fff + ((u.int32 >> 16) & 1) : 0; // Round to nearest, round to even
+    u.int32 |= flag1 ? 0x10000 : 0x0;                      // Preserve signaling NaN
+
+    return uint16_t(u.int32 >> 16);
+}
+
+// convert fp16 to bfp16 via fp32 with RTN if higher precision is needed
+template <>
+inline __host__ __device__ constexpr bhalf_t bf16_convert_rtn<bhalf_t, half_t>(half_t x)
+{
+    float x_fp32 = static_cast<float>(x);
+
+    return bf16_convert_rtn<bhalf_t>(x_fp32);
+}
+
+// Declare a template function for fp8 conversion using SR
+template <typename Y, typename X>
+__host__ __device__ constexpr Y f8_convert_sr(X x);
+
+// convert fp32 to fp8 with stochastic rounding
+template <>
+inline __host__ __device__ f8_t f8_convert_sr<f8_t, float>(float x)
+{
+    constexpr bool negative_zero_nan = true;
+    constexpr bool clip              = true;
+    constexpr f8_rounding_mode rm    = f8_rounding_mode::stochastic;
+    constexpr int seed               = 42;
+    // as thread id is not available on host, use 0 for prn generation
+    uint32_t rng = prand_generator<float, seed>(reinterpret_cast<uintptr_t>(&x), x);
+    return utils::cast_to_f8<float, negative_zero_nan, clip, (rm == f8_rounding_mode::stochastic)>(
+        x, rng);
+}
+
+// convert fp16 to fp8 with stochastic rounding
+template <>
+inline __host__ __device__ f8_t f8_convert_sr<f8_t, half_t>(half_t x)
+{
+    constexpr bool negative_zero_nan = true;
+    constexpr bool clip              = true;
+    constexpr f8_rounding_mode rm    = f8_rounding_mode::stochastic;
+    constexpr int seed               = 42;
+    // as thread id is not available on host, use 0 for prn generation
+    uint32_t rng = prand_generator<half_t, seed>(reinterpret_cast<uintptr_t>(&x), x);
+    return utils::cast_to_f8<half_t, negative_zero_nan, clip, (rm == f8_rounding_mode::stochastic)>(
+        x, rng);
+}
+
+} // namespace ck

include/ck/utility/workgroup_barrier.hpp

+#pragma once
+#include <hip/hip_runtime.h>
+#include <stdint.h>
+
+namespace ck {
+struct workgroup_barrier
+{
+    __device__ workgroup_barrier(uint32_t* ptr) : base_ptr(ptr) {}
+
+    __device__ uint32_t ld(uint32_t offset)
+    {
+#if 0
+        float d = llvm_amdgcn_raw_buffer_load_fp32(
+                        amdgcn_make_buffer_resource(base_ptr),
+                        0,
+                        offset,
+                        AMDGCN_BUFFER_GLC);
+        union cvt {
+            float f32;
+            uint32_t u32;
+        };
+        cvt x;
+        x.f32 = d;
+        return x.u32;
+#endif
+        return __atomic_load_n(base_ptr + offset, __ATOMIC_RELAXED);
+    }
+
+    __device__ void wait_eq(uint32_t offset, uint32_t value)
+    {
+        if(threadIdx.x == 0)
+        {
+            while(ld(offset) != value) {}
+        }
+        __syncthreads();
+    }
+
+    __device__ void wait_lt(uint32_t offset, uint32_t value)
+    {
+        if(threadIdx.x == 0)
+        {
+            while(ld(offset) < value) {}
+        }
+        __syncthreads();
+    }
+
+    __device__ void wait_set(uint32_t offset, uint32_t compare, uint32_t value)
+    {
+        if(threadIdx.x == 0)
+        {
+            while(atomicCAS(base_ptr + offset, compare, value) != compare) {}
+        }
+        __syncthreads();
+    }
+
+    // enter critical zoon, assume buffer is zero when launch kernel
+    __device__ void aquire(uint32_t offset) { wait_set(offset, 0, 1); }
+
+    // exit critical zoon, assume buffer is zero when launch kernel
+    __device__ void release(uint32_t offset) { wait_set(offset, 1, 0); }
+
+    __device__ void inc(uint32_t offset)
+    {
+        __syncthreads();
+        if(threadIdx.x == 0)
+        {
+            atomicAdd(base_ptr + offset, 1);
+        }
+    }
+
+    uint32_t* base_ptr;
+};
+} // namespace ck

include/ck/utility/workgroup_synchronization.hpp

+// SPDX-License-Identifier: MIT
+// Copyright (c) 2018-2023, Advanced Micro Devices, Inc. All rights reserved.
+
+#pragma once
+#include "ck/host_utility/hip_check_error.hpp"
+
+namespace ck {
+
+// Initialization flag of Barrier object, can be any value except for zero
+static constexpr int BarrierInitFlag = 0x7856;
+
+// 1) only the first thread-block in the synchronizaton group is supposed to call this function. It
+// is the responsibility of the user to ensure the two integer values in p_control_bits are zeros
+// before calling gms_init().
+// 2) Aftercalling gms_reset(), the two integer values in p_control_bits will be zeros, so no
+// repetitious initialization of p_control_bits buffer is required
+static __device__ void gms_init(int NumWarps, int* p_control_bits)
+{
+    union
+    {
+        int two32[2];
+        unsigned long one64;
+    } regs;
+
+    regs.two32[0] = BarrierInitFlag;
+    regs.two32[1] = NumWarps;
+
+    if(threadIdx.x == 0)
+        atomicCAS(reinterpret_cast<unsigned long*>(p_control_bits), 0, regs.one64);
+};
+
+// all the workgroups in the synchronization group is supposed to call this function
+static __device__ void gms_barrier(int* p_control_bits)
+{
+    constexpr int mask = warpSize - 1;
+
+    if((threadIdx.x & mask) == 0)
+    {
+        // ensure the barrier object is initialized
+        do
+        {
+            const int r0 = __atomic_load_n(&p_control_bits[0], __ATOMIC_RELAXED);
+
+            if(r0 == BarrierInitFlag)
+                break;
+
+        } while(true);
+
+        // go ahead toward the barrier line
+        atomicSub(&p_control_bits[1], 1);
+
+        // wait until all warps have arrived
+        do
+        {
+            const int r1 = __atomic_load_n(&p_control_bits[1], __ATOMIC_RELAXED);
+
+            if(r1 == 0)
+                break;
+
+        } while(true);
+    };
+};
+
+// 1) Only the first thread-block in the synchronizaton group is supposed to call this function.
+// 2) Aftercalling gms_reset(), the two integer values in p_control_bits will be zeros, so no
+// repetitious initialization of p_control_bits buffer is required
+static __device__ void gms_reset(int* p_control_bits)
+{
+    // reset the barrier object
+    if(threadIdx.x == 0)
+        (void)atomicCAS(&p_control_bits[0], BarrierInitFlag, 0);
+};
+
+} // namespace ck

include/ck/version.h.in

+/*******************************************************************************
+ *
+ * MIT License
+ *
+ * Copyright (c) 2023 Advanced Micro Devices, Inc.
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a copy
+ * of this software and associated documentation files (the "Software"), to deal
+ * in the Software without restriction, including without limitation the rights
+ * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+ * copies of the Software, and to permit persons to whom the Software is
+ * furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in all
+ * copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+ * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+ * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+ * SOFTWARE.
+ *
+ *******************************************************************************/
+
+/* the configured version and settings for miopen- Composable Kernel */
+
+#ifndef CK_VERSION_H_
+#define CK_VERSION_H_
+
+// clang-format off
+#define CK_VERSION @CMAKE_PROJECT_VERSION@
+#define CK_VERSION_MAJOR @CMAKE_PROJECT_VERSION_MAJOR@
+#define CK_VERSION_MINOR @CMAKE_PROJECT_VERSION_MINOR@
+#define CK_VERSION_PATCH @CMAKE_PROJECT_VERSION_PATCH@
+#define CK_COMMIT_ID @COMMIT_ID@
+// clang-format on
+
+#endif

library/include/ck/library/reference_tensor_operation/cpu/reference_avgpool_bwd.hpp

+// SPDX-License-Identifier: MIT
+// Copyright (c) 2018-2023, Advanced Micro Devices, Inc. All rights reserved.
+
+#pragma once
+
+#include <iostream>
+#include <sstream>
+
+#include "ck/tensor_operation/gpu/device/device_base.hpp"
+
+#include "ck/library/utility/host_tensor.hpp"
+
+namespace ck {
+namespace tensor_operation {
+namespace host {
+
+// dinput descriptor in [N, C, Do, Ho, Wo] order
+// doutput descriptor in [N, C, Di, Hi, Wi] order
+// phyiscal layout is irrelavent
+template <ck::index_t NDimSpatial,
+          typename DInDataType,
+          typename DOutDataType,
+          typename std::enable_if<NDimSpatial >= 1 && NDimSpatial <= 3, bool>::type = false>
+struct ReferenceAvgPoolBwd : public device::BaseOperator
+{
+    // Argument
+    struct Argument : public device::BaseArgument
+    {
+        Argument(Tensor<DInDataType>& dinput,
+                 const Tensor<DOutDataType>& doutput,
+                 std::vector<ck::index_t> window_spatial_lengths,
+                 std::vector<ck::index_t> window_strides,
+                 std::vector<ck::index_t> window_dilations,
+                 std::vector<ck::index_t> dinput_left_pads,
+                 std::vector<ck::index_t> dinput_right_pads)
+            : dinput_{dinput},
+              doutput_{doutput},
+              window_spatial_lengths_{window_spatial_lengths},
+              window_strides_{window_strides},
+              window_dilations_{window_dilations},
+              in_left_pads_{dinput_left_pads},
+              in_right_pads_{dinput_right_pads}
+        {
+        }
+
+        Tensor<DInDataType>& dinput_;
+        const Tensor<DOutDataType>& doutput_;
+
+        std::vector<ck::index_t> window_spatial_lengths_;
+        std::vector<index_t> window_strides_;
+        std::vector<index_t> window_dilations_;
+        std::vector<index_t> in_left_pads_;
+        std::vector<index_t> in_right_pads_;
+    };
+
+    // Invoker
+    struct Invoker : public device::BaseInvoker
+    {
+        using Argument = ReferenceAvgPoolBwd::Argument;
+
+        template <ck::index_t NDimSpatial_,
+                  typename std::enable_if<NDimSpatial_ == 1, bool>::type = false>
+        float RunAvgPoolBwd(const Argument& arg)
+        {
+            // Let input = x, outpu = y
+            // shape of x = [10], y = [6]
+            // window_size = 5, pad = 0, stride = 1, dilation = 1
+            // Forward:
+            // y0 = 1/5 * (x0 + x1 + x2 + x3 + x4)
+            // y1 = 1/5 * (x1 + x2 + x3 + x4 + x5)
+            // ...
+            // y5 = 1/5 * (x5 + x6 + x7 + x8 + x9)
+            // y6 = 1/5 * (x6 + x7 + x8 + x9)
+            // ...
+            // y9 = 1/5 * (x9)
+
+            // Backward:
+            // shape of dy = [6], dx = [10]
+            // dx0 = 1/5 * dy0
+            // dx1 = 1/5 * (dy0 + dy1)
+            // dx2 = 1/5 * (dy0 + dy1 + dy2)
+            // ...
+            // dx4 = 1/5 * (dy0 + dy1 + dy2 + dy3 + dy4)
+            // dx5 = 1/5 * (dy1 + dy2 + dy3 + dy4 + dy5)
+            // ...
+            // dx9 = 1/5 * (dy5 + dy6 + dy7 + dy8 + dy9)
+
+            auto f_ncw = [&](auto n, auto c, auto wi) {
+                std::size_t X  = arg.window_spatial_lengths_[0];
+                std::size_t Wo = arg.doutput_.GetLengths()[2];
+
+                float v_acc = 0;
+
+                for(std::size_t x = 0; x < X; ++x)
+                {
+                    // Out_Position = (In_Position + pad - x * dilation) / stride
+                    auto w_tmp = static_cast<ck::long_index_t>(wi) +
+                                 static_cast<ck::long_index_t>(arg.in_left_pads_[0]) -
+                                 static_cast<ck::long_index_t>(x * arg.window_dilations_[0]);
+
+                    // Check the input pixel validity (in perspective of being affected by some
+                    // doutput pixel)
+                    if(w_tmp % arg.window_strides_[0] == 0)
+                    {
+                        auto wo = static_cast<ck::long_index_t>(w_tmp) /
+                                  static_cast<ck::long_index_t>(arg.window_strides_[0]);
+
+                        // Get the doutput pixel in valid range to accumulate the gradients for this
+                        // input pixel
+                        if(wo >= 0 && ck::type_convert<std::size_t>(wo) < Wo)
+                        {
+                            v_acc += ck::type_convert<float>(arg.doutput_(n, c, wo));
+                        }
+                    }
+                }
+
+                v_acc /= ck::type_convert<float>(X);
+                arg.dinput_(n, c, wi) = ck::type_convert<DInDataType>(v_acc);
+            };
+
+            make_ParallelTensorFunctor(f_ncw,
+                                       arg.dinput_.GetLengths()[0],
+                                       arg.dinput_.GetLengths()[1],
+                                       arg.dinput_.GetLengths()[2])(
+                std::thread::hardware_concurrency());
+
+            return 0;
+        }
+
+        template <ck::index_t NDimSpatial_,
+                  typename std::enable_if<NDimSpatial_ == 2, bool>::type = false>
+        float RunAvgPoolBwd(const Argument& arg)
+        {
+            auto f_nchw = [&](auto n, auto c, auto hi, auto wi) {
+                std::size_t Y = arg.window_spatial_lengths_[0];
+                std::size_t X = arg.window_spatial_lengths_[1];
+
+                std::size_t Ho = arg.doutput_.GetLengths()[2];
+                std::size_t Wo = arg.doutput_.GetLengths()[3];
+
+                float v_acc = 0;
+
+                for(std::size_t y = 0; y < Y; ++y)
+                {
+                    // Out_Position = (In_Position + pad - x * dilation) / stride
+                    auto h_tmp = static_cast<ck::long_index_t>(hi) +
+                                 static_cast<ck::long_index_t>(arg.in_left_pads_[0]) -
+                                 static_cast<ck::long_index_t>(y * arg.window_dilations_[0]);
+
+                    // Check the input pixel validity (in perspective of being affected by some
+                    // doutput pixel)
+                    if(h_tmp % arg.window_strides_[0] == 0)
+                    {
+                        auto ho = static_cast<ck::long_index_t>(h_tmp) /
+                                  static_cast<ck::long_index_t>(arg.window_strides_[0]);
+
+                        // Get the doutput pixel in valid range to accumulate the gradients for this
+                        // input pixel
+                        if(ho >= 0 && ck::type_convert<std::size_t>(ho) < Ho)
+                        {
+                            for(std::size_t x = 0; x < X; ++x)
+                            {
+                                auto w_tmp =
+                                    static_cast<ck::long_index_t>(wi) +
+                                    static_cast<ck::long_index_t>(arg.in_left_pads_[1]) -
+                                    static_cast<ck::long_index_t>(x * arg.window_dilations_[1]);
+                                if(w_tmp % arg.window_strides_[1] == 0)
+                                {
+                                    auto wo = static_cast<ck::long_index_t>(w_tmp) /
+                                              static_cast<ck::long_index_t>(arg.window_strides_[1]);
+                                    if(wo >= 0 && ck::type_convert<std::size_t>(wo) < Wo)
+                                    {
+                                        v_acc +=
+                                            ck::type_convert<float>(arg.doutput_(n, c, ho, wo));
+                                    }
+                                }
+                            }
+                        }
+                    }
+                }
+
+                v_acc /= ck::type_convert<float>(Y * X);
+                arg.dinput_(n, c, hi, wi) = ck::type_convert<DInDataType>(v_acc);
+            };
+
+            make_ParallelTensorFunctor(f_nchw,
+                                       arg.dinput_.GetLengths()[0],
+                                       arg.dinput_.GetLengths()[1],
+                                       arg.dinput_.GetLengths()[2],
+                                       arg.dinput_.GetLengths()[3])(
+                std::thread::hardware_concurrency());
+
+            return 0;
+        }
+
+        template <ck::index_t NDimSpatial_,
+                  typename std::enable_if<NDimSpatial_ == 3, bool>::type = false>
+        float RunAvgPoolBwd(const Argument& arg)
+        {
+            auto f_ncdhw = [&](auto n, auto c, auto di, auto hi, auto wi) {
+                std::size_t Z = arg.window_spatial_lengths_[0];
+                std::size_t Y = arg.window_spatial_lengths_[1];
+                std::size_t X = arg.window_spatial_lengths_[2];
+
+                std::size_t Do = arg.doutput_.GetLengths()[2];
+                std::size_t Ho = arg.doutput_.GetLengths()[3];
+                std::size_t Wo = arg.doutput_.GetLengths()[4];
+
+                float v_acc = 0;
+
+                for(std::size_t z = 0; z < Z; ++z)
+                {
+                    // Out_Position = (In_Position + pad - x * dilation) / stride
+                    auto d_tmp = static_cast<ck::long_index_t>(di) +
+                                 static_cast<ck::long_index_t>(arg.in_left_pads_[0]) -
+                                 static_cast<ck::long_index_t>(z * arg.window_dilations_[0]);
+
+                    // Check the input pixel validity (in perspective of being affected by some
+                    // doutput pixel)
+                    if(d_tmp % arg.window_strides_[0] == 0)
+                    {
+                        auto do_ = static_cast<ck::long_index_t>(d_tmp) /
+                                   static_cast<ck::long_index_t>(arg.window_strides_[0]);
+
+                        // Get the doutput pixel in valid range to accumulate the gradients for this
+                        // input pixel
+                        if(do_ >= 0 && ck::type_convert<std::size_t>(do_) < Do)
+                        {
+                            for(std::size_t y = 0; y < Y; ++y)
+                            {
+                                auto h_tmp =
+                                    static_cast<ck::long_index_t>(hi) +
+                                    static_cast<ck::long_index_t>(arg.in_left_pads_[1]) -
+                                    static_cast<ck::long_index_t>(y * arg.window_dilations_[1]);
+                                if(h_tmp % arg.window_strides_[1] == 0)
+                                {
+                                    auto ho = static_cast<ck::long_index_t>(h_tmp) /
+                                              static_cast<ck::long_index_t>(arg.window_strides_[1]);
+                                    if(ho >= 0 && ck::type_convert<std::size_t>(ho) < Ho)
+                                    {
+                                        for(std::size_t x = 0; x < X; ++x)
+                                        {
+                                            auto w_tmp = static_cast<ck::long_index_t>(wi) +
+                                                         static_cast<ck::long_index_t>(
+                                                             arg.in_left_pads_[2]) -
+                                                         static_cast<ck::long_index_t>(
+                                                             x * arg.window_dilations_[2]);
+
+                                            if(w_tmp % arg.window_strides_[2] == 0)
+                                            {
+                                                auto wo = static_cast<ck::long_index_t>(w_tmp) /
+                                                          static_cast<ck::long_index_t>(
+                                                              arg.window_strides_[2]);
+                                                if(wo >= 0 &&
+                                                   ck::type_convert<std::size_t>(wo) < Wo)
+                                                {
+                                                    v_acc += ck::type_convert<float>(
+                                                        arg.doutput_(n, c, do_, ho, wo));
+                                                }
+                                            }
+                                        }
+                                    }
+                                }
+                            }
+                        }
+                    }
+                }
+
+                v_acc /= ck::type_convert<float>(Z * Y * X);
+                arg.dinput_(n, c, di, hi, wi) = ck::type_convert<DInDataType>(v_acc);
+            };
+
+            make_ParallelTensorFunctor(f_ncdhw,
+                                       arg.dinput_.GetLengths()[0],
+                                       arg.dinput_.GetLengths()[1],
+                                       arg.dinput_.GetLengths()[2],
+                                       arg.dinput_.GetLengths()[3],
+                                       arg.dinput_.GetLengths()[4])(
+                std::thread::hardware_concurrency());
+
+            return 0;
+        }
+
+        float Run(const Argument& arg)
+        {
+            if(!(arg.dinput_.GetNumOfDimension() == NDimSpatial + 2 &&
+                 arg.doutput_.GetNumOfDimension() == NDimSpatial + 2))
+            {
+                throw std::runtime_error("wrong! inconsistent dimension");
+            }
+
+            return RunAvgPoolBwd<NDimSpatial>(arg);
+        }
+
+        float Run(const device::BaseArgument* p_arg,
+                  const StreamConfig& /* stream_config */ = StreamConfig{}) override
+        {
+            return Run(*dynamic_cast<const Argument*>(p_arg));
+        }
+    };
+
+    static constexpr bool IsValidCompilationParameter()
+    {
+        // TODO: properly implement this check
+        return true;
+    }
+
+    bool IsSupportedArgument(const device::BaseArgument*) override { return true; }
+
+    static auto MakeArgument(Tensor<DInDataType>& dinput,
+                             const Tensor<DOutDataType>& doutput,
+                             std::vector<ck::index_t> window_spatial_lengths,
+                             std::vector<ck::index_t> window_strides,
+                             std::vector<ck::index_t> window_dilations,
+                             std::vector<ck::index_t> dinput_left_pads,
+                             std::vector<ck::index_t> dinput_right_pads)
+    {
+        if(window_spatial_lengths.size() != NDimSpatial || window_strides.size() != NDimSpatial ||
+           window_dilations.size() != NDimSpatial || dinput_left_pads.size() != NDimSpatial ||
+           dinput_right_pads.size() != NDimSpatial)
+            throw std::runtime_error("dimension is incorrect");
+
+        return Argument{dinput,
+                        doutput,
+                        window_spatial_lengths,
+                        window_strides,
+                        window_dilations,
+                        dinput_left_pads,
+                        dinput_right_pads};
+    }
+
+    static auto MakeInvoker() { return Invoker{}; }
+
+    virtual std::unique_ptr<device::BaseInvoker> MakeInvokerPointer()
+    {
+        return std::make_unique<Invoker>(Invoker{});
+    }
+
+    std::string GetTypeString() const override
+    {
+        auto str = std::stringstream();
+
+        // clang-format off
+        str << "ReferenceAvgPoolBwd"
+            << std::endl;
+        // clang-format on
+
+        return str.str();
+    }
+};
+
+} // namespace host
+} // namespace tensor_operation
+} // namespace ck

library/include/ck/library/reference_tensor_operation/cpu/reference_conv_bwd_data.hpp

@@ -125,7 +125,7 @@ struct ReferenceConvBwdData : public device::BaseOperator
 
                     arg.in_element_op_(v_in, v_acc);
 
-                    arg.input_(g, n, c, wi) = ck::type_convert<InDataType>(v_acc);
+                    arg.input_(g, n, c, wi) = ck::type_convert<InDataType>(v_in);
                 };
 
                 make_ParallelTensorFunctor(f_ncw,
@@ -201,7 +201,7 @@ struct ReferenceConvBwdData : public device::BaseOperator
 
                     arg.in_element_op_(v_in, v_acc);
 
-                    arg.input_(g, n, c, hi, wi) = ck::type_convert<InDataType>(v_acc);
+                    arg.input_(g, n, c, hi, wi) = ck::type_convert<InDataType>(v_in);
                 };
 
                 make_ParallelTensorFunctor(f_nchw,
@@ -299,7 +299,7 @@ struct ReferenceConvBwdData : public device::BaseOperator
 
                     arg.in_element_op_(v_in, v_acc);
 
-                    arg.input_(g, n, c, di, hi, wi) = ck::type_convert<InDataType>(v_acc);
+                    arg.input_(g, n, c, di, hi, wi) = ck::type_convert<InDataType>(v_in);
                 };
 
                 make_ParallelTensorFunctor(f_ncdhw,

library/include/ck/library/reference_tensor_operation/cpu/reference_maxpool_bwd.hpp

+// SPDX-License-Identifier: MIT
+// Copyright (c) 2018-2023, Advanced Micro Devices, Inc. All rights reserved.
+
+#pragma once
+
+#include <iostream>
+#include <sstream>
+#include <vector>
+
+#include "ck/tensor_operation/gpu/device/device_base.hpp"
+#include "ck/library/utility/host_tensor.hpp"
+#include "ck/library/utility/host_tensor_generator.hpp"
+
+namespace ck {
+namespace tensor_operation {
+namespace host {
+using namespace std;
+
+template <typename DOutDataType,
+          typename IndexDataType,
+          typename ConputeDataType,
+          typename DInDataType,
+          typename ElementwiseOperation>
+struct ReferenceMaxPoolBwd : public device::BaseOperator
+{
+    // Argument
+    struct Argument : public device::BaseArgument
+    {
+        Argument(const Tensor<DOutDataType>& dout,
+                 const Tensor<IndexDataType>& indices,
+                 Tensor<DInDataType>& din,
+                 ElementwiseOperation elementwise_op)
+            : dout_(dout), indices_(indices), din_(din), elementwise_op_(elementwise_op)
+        {
+        }
+
+        const Tensor<DOutDataType>& dout_;
+        const Tensor<IndexDataType>& indices_;
+        Tensor<DInDataType>& din_;
+        ElementwiseOperation elementwise_op_;
+    };
+
+    // Invoker
+    struct Invoker : public device::BaseInvoker
+    {
+        float Run(const Argument& arg)
+        {
+            int din_length  = arg.din_.GetElementSpaceSize();
+            int dout_length = arg.dout_.GetElementSpaceSize();
+            std::vector<ConputeDataType> buf(din_length, 0);
+
+            for(int i = 0; i < dout_length; ++i)
+            {
+                int index = arg.indices_.mData[i];
+                if(index >= 0 && index < din_length)
+                    buf[index] += ck::type_convert<ConputeDataType>(arg.dout_.mData[i]);
+            }
+
+            for(int i = 0; i < din_length; ++i)
+                arg.din_.mData[i] = ck::type_convert<DInDataType>(buf[i]);
+            return 0;
+        }
+
+        float Run(const device::BaseArgument* p_arg,
+                  const StreamConfig& /* stream_config */ = StreamConfig{}) override
+        {
+            return Run(*dynamic_cast<const Argument*>(p_arg));
+        }
+    };
+
+    bool IsSupportedArgument(const device::BaseArgument*) override { return true; }
+
+    static auto MakeArgument(const Tensor<DOutDataType>& dout,
+                             const Tensor<IndexDataType>& indices,
+                             Tensor<DInDataType>& din,
+                             ElementwiseOperation elementwise_op)
+    {
+        return Argument{dout, indices, din, elementwise_op};
+    }
+
+    static auto MakeInvoker() { return Invoker{}; }
+
+    virtual std::unique_ptr<device::BaseInvoker> MakeInvokerPointer()
+    {
+        return std::make_unique<Invoker>(Invoker{});
+    }
+
+    std::string GetTypeString() const override
+    {
+        auto str = std::stringstream();
+
+        // clang-format off
+        str << "ReferenceMaxPoolBwd"
+            << std::endl;
+        // clang-format on
+
+        return str.str();
+    }
+};
+
+} // namespace host
+} // namespace tensor_operation
+} // namespace ck

library/include/ck/library/reference_tensor_operation/cpu/reference_pool_fwd.hpp

@@ -39,6 +39,7 @@ struct ReferencePoolingFwd : public device::BaseOperator
                  Tensor<IndexDataType>& out_indices,
                  const std::vector<ck::index_t>& window_spatial_lengths,
                  const std::vector<ck::index_t>& window_strides,
+                 const std::vector<ck::index_t>& window_dilations,
                  const std::vector<ck::index_t>& in_left_pads,
                  const std::vector<ck::index_t>& /*in_right_pads*/)
             : in_(in),
@@ -46,6 +47,7 @@ struct ReferencePoolingFwd : public device::BaseOperator
               out_indices_(out_indices),
               window_spatial_lengths_(window_spatial_lengths),
               window_strides_(window_strides),
+              window_dilations_(window_dilations),
               in_left_pads_(in_left_pads),
               reduceLength_(1)
         {
@@ -58,6 +60,7 @@ struct ReferencePoolingFwd : public device::BaseOperator
         Tensor<IndexDataType>& out_indices_;
         const std::vector<ck::index_t>& window_spatial_lengths_;
         const std::vector<ck::index_t>& window_strides_;
+        const std::vector<ck::index_t>& window_dilations_;
         const std::vector<ck::index_t>& in_left_pads_;
         int reduceLength_;
     };
@@ -85,14 +88,17 @@ struct ReferencePoolingFwd : public device::BaseOperator
 
                     for(ck::index_t z = 0; z < arg.window_spatial_lengths_[0]; ++z)
                     {
-                        ck::index_t di = do_ * arg.window_strides_[0] + z - arg.in_left_pads_[0];
+                        ck::index_t di = do_ * arg.window_strides_[0] +
+                                         z * arg.window_dilations_[0] - arg.in_left_pads_[0];
                         for(ck::index_t y = 0; y < arg.window_spatial_lengths_[1]; ++y)
                         {
-                            ck::index_t hi = ho * arg.window_strides_[1] + y - arg.in_left_pads_[1];
+                            ck::index_t hi = ho * arg.window_strides_[1] +
+                                             y * arg.window_dilations_[1] - arg.in_left_pads_[1];
                             for(ck::index_t x = 0; x < arg.window_spatial_lengths_[2]; ++x)
                             {
-                                ck::index_t wi =
-                                    wo * arg.window_strides_[2] + x - arg.in_left_pads_[2];
+                                ck::index_t wi = wo * arg.window_strides_[2] +
+                                                 x * arg.window_dilations_[2] -
+                                                 arg.in_left_pads_[2];
                                 if(di >= 0 &&
                                    di < static_cast<ck::index_t>(arg.in_.mDesc.GetLengths()[2]) &&
                                    hi >= 0 &&
@@ -100,8 +106,8 @@ struct ReferencePoolingFwd : public device::BaseOperator
                                    wi >= 0 &&
                                    wi < static_cast<ck::index_t>(arg.in_.mDesc.GetLengths()[4]))
                                 {
-                                    ComputeDataType currVal =
-                                        static_cast<ComputeDataType>(arg.in_(n, c, di, hi, wi));
+                                    ComputeDataType currVal = ck::type_convert<ComputeDataType>(
+                                        arg.in_(n, c, di, hi, wi));
 
                                     in_elementwise_op(currVal, currVal);
 
@@ -112,7 +118,7 @@ struct ReferencePoolingFwd : public device::BaseOperator
                     }
                     acc_elementwise_op(accuVal, accuVal);
 
-                    arg.out_(n, c, do_, ho, wo) = accuVal;
+                    arg.out_(n, c, do_, ho, wo) = ck::type_convert<OutDataType>(accuVal);
                 };
 
                 make_ParallelTensorFunctor(f_ncdhw,
@@ -136,14 +142,17 @@ struct ReferencePoolingFwd : public device::BaseOperator
 
                     for(ck::index_t z = 0; z < arg.window_spatial_lengths_[0]; ++z)
                     {
-                        ck::index_t di = do_ * arg.window_strides_[0] + z - arg.in_left_pads_[0];
+                        ck::index_t di = do_ * arg.window_strides_[0] +
+                                         z * arg.window_dilations_[0] - arg.in_left_pads_[0];
                         for(ck::index_t y = 0; y < arg.window_spatial_lengths_[1]; ++y)
                         {
-                            ck::index_t hi = ho * arg.window_strides_[1] + y - arg.in_left_pads_[1];
+                            ck::index_t hi = ho * arg.window_strides_[1] +
+                                             y * arg.window_dilations_[1] - arg.in_left_pads_[1];
                             for(ck::index_t x = 0; x < arg.window_spatial_lengths_[2]; ++x)
                             {
-                                ck::index_t wi =
-                                    wo * arg.window_strides_[2] + x - arg.in_left_pads_[2];
+                                ck::index_t wi = wo * arg.window_strides_[2] +
+                                                 x * arg.window_dilations_[2] -
+                                                 arg.in_left_pads_[2];
                                 if(di >= 0 &&
                                    di < static_cast<ck::index_t>(arg.in_.mDesc.GetLengths()[2]) &&
                                    hi >= 0 &&
@@ -151,8 +160,8 @@ struct ReferencePoolingFwd : public device::BaseOperator
                                    wi >= 0 &&
                                    wi < static_cast<ck::index_t>(arg.in_.mDesc.GetLengths()[4]))
                                 {
-                                    ComputeDataType currVal =
-                                        static_cast<ComputeDataType>(arg.in_(n, c, di, hi, wi));
+                                    ComputeDataType currVal = ck::type_convert<ComputeDataType>(
+                                        arg.in_(n, c, di, hi, wi));
                                     IndexDataType currIndex =
                                         arg.in_.GetOffsetFromMultiIndex(n, c, di, hi, wi);
 
@@ -166,7 +175,7 @@ struct ReferencePoolingFwd : public device::BaseOperator
 
                     acc_elementwise_op(accuVal, accuVal);
 
-                    arg.out_(n, c, do_, ho, wo)         = accuVal;
+                    arg.out_(n, c, do_, ho, wo)         = ck::type_convert<OutDataType>(accuVal);
                     arg.out_indices_(n, c, do_, ho, wo) = accuIndex;
                 };
 
@@ -202,17 +211,19 @@ struct ReferencePoolingFwd : public device::BaseOperator
 
                     for(ck::index_t y = 0; y < arg.window_spatial_lengths_[0]; ++y)
                     {
-                        ck::index_t hi = ho * arg.window_strides_[0] + y - arg.in_left_pads_[0];
+                        ck::index_t hi = ho * arg.window_strides_[0] +
+                                         y * arg.window_dilations_[0] - arg.in_left_pads_[0];
                         for(ck::index_t x = 0; x < arg.window_spatial_lengths_[1]; ++x)
                         {
-                            ck::index_t wi = wo * arg.window_strides_[1] + x - arg.in_left_pads_[1];
+                            ck::index_t wi = wo * arg.window_strides_[1] +
+                                             x * arg.window_dilations_[1] - arg.in_left_pads_[1];
                             if(hi >= 0 &&
                                hi < static_cast<ck::index_t>(arg.in_.mDesc.GetLengths()[2]) &&
                                wi >= 0 &&
                                wi < static_cast<ck::index_t>(arg.in_.mDesc.GetLengths()[3]))
                             {
                                 ComputeDataType currVal =
-                                    static_cast<ComputeDataType>(arg.in_(n, c, hi, wi));
+                                    ck::type_convert<ComputeDataType>(arg.in_(n, c, hi, wi));
 
                                 in_elementwise_op(currVal, currVal);
 
@@ -222,7 +233,7 @@ struct ReferencePoolingFwd : public device::BaseOperator
                     }
 
                     acc_elementwise_op(accuVal, accuVal);
-                    arg.out_(n, c, ho, wo) = accuVal;
+                    arg.out_(n, c, ho, wo) = ck::type_convert<OutDataType>(accuVal);
                 };
 
                 make_ParallelTensorFunctor(f_nchw,
@@ -255,7 +266,7 @@ struct ReferencePoolingFwd : public device::BaseOperator
                                wi < static_cast<ck::index_t>(arg.in_.mDesc.GetLengths()[3]))
                             {
                                 ComputeDataType currVal =
-                                    static_cast<ComputeDataType>(arg.in_(n, c, hi, wi));
+                                    ck::type_convert<ComputeDataType>(arg.in_(n, c, hi, wi));
 
                                 IndexDataType currIndex =
                                     arg.in_.GetOffsetFromMultiIndex(n, c, hi, wi);
@@ -268,7 +279,7 @@ struct ReferencePoolingFwd : public device::BaseOperator
                     }
 
                     acc_elementwise_op(accuVal, accuVal);
-                    arg.out_(n, c, ho, wo)         = accuVal;
+                    arg.out_(n, c, ho, wo)         = ck::type_convert<OutDataType>(accuVal);
                     arg.out_indices_(n, c, ho, wo) = accuIndex;
                 };
 
@@ -308,6 +319,7 @@ struct ReferencePoolingFwd : public device::BaseOperator
                              Tensor<IndexDataType>& out_indices,
                              const std::vector<ck::index_t>& window_spatial_lengths,
                              const std::vector<ck::index_t>& window_strides,
+                             const std::vector<ck::index_t>& window_dilations,
                              const std::vector<ck::index_t>& in_left_pads,
                              const std::vector<ck::index_t>& in_right_pads)
     {
@@ -316,6 +328,7 @@ struct ReferencePoolingFwd : public device::BaseOperator
                         out_indices,
                         window_spatial_lengths,
                         window_strides,
+                        window_dilations,
                         in_left_pads,
                         in_right_pads};
     }

library/include/ck/library/tensor_operation_instance/device_operation_instance_factory.hpp

@@ -17,6 +17,7 @@ namespace instance {
 using F64  = double;
 using F32  = float;
 using F16  = ck::half_t;
+using F8   = ck::f8_t;
 using BF16 = ck::bhalf_t;
 using I8   = int8_t;
 using I32  = int32_t;

library/include/ck/library/tensor_operation_instance/gpu/batched_gemm.hpp

@@ -16,7 +16,7 @@ namespace ck {
 namespace tensor_operation {
 namespace device {
 namespace instance {
-
+#ifdef CK_ENABLE_BF16
 void add_device_batched_gemm_xdl_bf16_bf16_bf16_gkm_gkn_gmn_instances(
     std::vector<std::unique_ptr<
         DeviceBatchedGemm<Col, Row, Row, BF16, BF16, BF16, PassThrough, PassThrough, PassThrough>>>&
@@ -36,7 +36,8 @@ void add_device_batched_gemm_xdl_bf16_bf16_bf16_gmk_gnk_gmn_instances(
     std::vector<std::unique_ptr<
         DeviceBatchedGemm<Row, Col, Row, BF16, BF16, BF16, PassThrough, PassThrough, PassThrough>>>&
         instances);
-
+#endif
+#ifdef CK_ENABLE_FP16
 void add_device_batched_gemm_xdl_f16_f16_f16_gkm_gkn_gmn_instances(
     std::vector<std::unique_ptr<
         DeviceBatchedGemm<Col, Row, Row, F16, F16, F16, PassThrough, PassThrough, PassThrough>>>&
@@ -56,7 +57,8 @@ void add_device_batched_gemm_xdl_f16_f16_f16_gmk_gnk_gmn_instances(
     std::vector<std::unique_ptr<
         DeviceBatchedGemm<Row, Col, Row, F16, F16, F16, PassThrough, PassThrough, PassThrough>>>&
         instances);
-
+#endif
+#ifdef CK_ENABLE_FP32
 void add_device_batched_gemm_xdl_f32_f32_f32_gkm_gkn_gmn_instances(
     std::vector<std::unique_ptr<
         DeviceBatchedGemm<Col, Row, Row, F32, F32, F32, PassThrough, PassThrough, PassThrough>>>&
@@ -76,7 +78,8 @@ void add_device_batched_gemm_xdl_f32_f32_f32_gmk_gnk_gmn_instances(
     std::vector<std::unique_ptr<
         DeviceBatchedGemm<Row, Col, Row, F32, F32, F32, PassThrough, PassThrough, PassThrough>>>&
         instances);
-
+#endif
+#ifdef CK_ENABLE_INT8
 void add_device_batched_gemm_xdl_int8_int8_int8_gkm_gkn_gmn_instances(
     std::vector<std::unique_ptr<DeviceBatchedGemm<Col,
                                                   Row,
@@ -120,7 +123,7 @@ void add_device_batched_gemm_xdl_int8_int8_int8_gmk_gnk_gmn_instances(
                                                   PassThrough,
                                                   PassThrough,
                                                   PassThrough>>>& instances);
-
+#endif
 template <typename ALayout,
           typename BLayout,
           typename CLayout,
@@ -151,7 +154,7 @@ struct DeviceOperationInstanceFactory<ck::tensor_operation::device::DeviceBatche
     static auto GetInstances()
     {
         std::vector<std::unique_ptr<DeviceOp>> op_ptrs;
-
+#ifdef CK_ENABLE_FP32
         if constexpr(is_same_v<ADataType, float> && is_same_v<BDataType, float> &&
                      is_same_v<CDataType, float>)
         {
@@ -176,8 +179,10 @@ struct DeviceOperationInstanceFactory<ck::tensor_operation::device::DeviceBatche
                 add_device_batched_gemm_xdl_f32_f32_f32_gkm_gnk_gmn_instances(op_ptrs);
             }
         }
-        else if constexpr(is_same_v<ADataType, half_t> && is_same_v<BDataType, half_t> &&
-                          is_same_v<CDataType, half_t>)
+#endif
+#ifdef CK_ENABLE_FP16
+        if constexpr(is_same_v<ADataType, half_t> && is_same_v<BDataType, half_t> &&
+                     is_same_v<CDataType, half_t>)
         {
             if constexpr(is_same_v<ALayout, Row> && is_same_v<BLayout, Row> &&
                          is_same_v<CLayout, Row>)
@@ -200,8 +205,10 @@ struct DeviceOperationInstanceFactory<ck::tensor_operation::device::DeviceBatche
                 add_device_batched_gemm_xdl_f16_f16_f16_gkm_gnk_gmn_instances(op_ptrs);
             }
         }
-        else if constexpr(is_same_v<ADataType, bhalf_t> && is_same_v<BDataType, bhalf_t> &&
-                          is_same_v<CDataType, bhalf_t>)
+#endif
+#ifdef CK_ENABLE_BF16
+        if constexpr(is_same_v<ADataType, bhalf_t> && is_same_v<BDataType, bhalf_t> &&
+                     is_same_v<CDataType, bhalf_t>)
         {
             if constexpr(is_same_v<ALayout, Row> && is_same_v<BLayout, Row> &&
                          is_same_v<CLayout, Row>)
@@ -224,8 +231,10 @@ struct DeviceOperationInstanceFactory<ck::tensor_operation::device::DeviceBatche
                 add_device_batched_gemm_xdl_bf16_bf16_bf16_gkm_gnk_gmn_instances(op_ptrs);
             }
         }
-        else if constexpr(is_same_v<ADataType, int8_t> && is_same_v<BDataType, int8_t> &&
-                          is_same_v<CDataType, int8_t>)
+#endif
+#ifdef CK_ENABLE_INT8
+        if constexpr(is_same_v<ADataType, int8_t> && is_same_v<BDataType, int8_t> &&
+                     is_same_v<CDataType, int8_t>)
         {
             if constexpr(is_same_v<ALayout, Row> && is_same_v<BLayout, Row> &&
                          is_same_v<CLayout, Row>)
@@ -248,7 +257,7 @@ struct DeviceOperationInstanceFactory<ck::tensor_operation::device::DeviceBatche
                 add_device_batched_gemm_xdl_int8_int8_int8_gkm_gnk_gmn_instances(op_ptrs);
             }
         }
-
+#endif
         return op_ptrs;
     }
 };

library/include/ck/library/tensor_operation_instance/gpu/batched_gemm_add_relu_gemm_add.hpp

@@ -14,7 +14,7 @@
 
 using CDE0ElementOp = ck::tensor_operation::element_wise::AddRelu;
 using CDE1ElementOp = ck::tensor_operation::element_wise::Add;
-
+#ifdef CK_ENABLE_FP16
 namespace ck {
 namespace tensor_operation {
 namespace device {
@@ -137,3 +137,4 @@ struct DeviceOperationInstanceFactory<
 } // namespace device
 } // namespace tensor_operation
 } // namespace ck
+#endif

library/include/ck/library/tensor_operation_instance/gpu/batched_gemm_bias_permute.hpp

@@ -13,7 +13,7 @@
 #include "ck/tensor_operation/gpu/element/element_wise_operation.hpp"
 
 #include "ck/library/tensor_operation_instance/device_operation_instance_factory.hpp"
-
+#ifdef CK_ENABLE_FP16
 namespace ck {
 namespace tensor_operation {
 namespace device {
@@ -91,3 +91,4 @@ struct DeviceOperationInstanceFactory<
 } // namespace device
 } // namespace tensor_operation
 } // namespace ck
+#endif

library/include/ck/library/tensor_operation_instance/gpu/batched_gemm_bias_softmax_gemm_permute.hpp

@@ -16,7 +16,7 @@ namespace ck {
 namespace tensor_operation {
 namespace device {
 namespace instance {
-
+#ifdef CK_ENABLE_FP16
 void add_device_batched_gemm_bias_masking_softmax_gemm_permute_xdl_cshuffle_f16_f16_f16_f16_gmk_gnk_gno_gmo_instances(
     std::vector<std::unique_ptr<
         DeviceBatchedGemmSoftmaxGemmPermute<2,
@@ -58,7 +58,8 @@ void add_device_batched_gemm_bias_softmax_gemm_permute_xdl_cshuffle_f16_f16_f16_
                                                             PassThrough,
                                                             MaskingSpecialization::MaskDisabled>>>&
         instances);
-
+#endif
+#ifdef CK_ENABLE_BF16
 void add_device_batched_gemm_bias_masking_softmax_gemm_permute_xdl_cshuffle_bf16_bf16_bf16_bf16_gmk_gnk_gno_gmo_instances(
     std::vector<std::unique_ptr<
         DeviceBatchedGemmSoftmaxGemmPermute<2,
@@ -100,7 +101,7 @@ void add_device_batched_gemm_bias_softmax_gemm_permute_xdl_cshuffle_bf16_bf16_bf
                                                             PassThrough,
                                                             MaskingSpecialization::MaskDisabled>>>&
         instances);
-
+#endif
 template <typename ADataType,
           typename B0DataType,
           typename B1DataType,
@@ -147,7 +148,7 @@ struct DeviceOperationInstanceFactory<
     static auto GetInstances()
     {
         std::vector<std::unique_ptr<DeviceOp>> op_ptrs;
-
+#ifdef CK_ENABLE_FP16
         if constexpr(is_same_v<ADataType, half_t> && is_same_v<B0DataType, half_t> &&
                      is_same_v<B1DataType, half_t> && is_same_v<CDataType, half_t> &&
                      Acc0BiasDataType::Size() == 1 &&
@@ -164,6 +165,8 @@ struct DeviceOperationInstanceFactory<
                     op_ptrs);
             }
         }
+#endif
+#ifdef CK_ENABLE_BF16
         else if constexpr(is_same_v<ADataType, BF16> && is_same_v<B0DataType, BF16> &&
                           is_same_v<B1DataType, BF16> && is_same_v<CDataType, BF16> &&
                           Acc0BiasDataType::Size() == 1 &&
@@ -180,6 +183,7 @@ struct DeviceOperationInstanceFactory<
                     op_ptrs);
             }
         }
+#endif
         return op_ptrs;
     }
 };

library/include/ck/library/tensor_operation_instance/gpu/batched_gemm_gemm.hpp

@@ -16,7 +16,7 @@ namespace ck {
 namespace tensor_operation {
 namespace device {
 namespace instance {
-
+#ifdef CK_ENABLE_FP16
 void add_device_batched_gemm_gemm_xdl_cshuffle_f16_f16_f16_f16_gmk_gnk_gno_gmo_instance(
     std::vector<std::unique_ptr<DeviceBatchedGemmGemm<Row,
                                                       Col,
@@ -111,3 +111,4 @@ struct DeviceOperationInstanceFactory<
 } // namespace device
 } // namespace tensor_operation
 } // namespace ck
+#endif

library/include/ck/library/tensor_operation_instance/gpu/batched_gemm_multi_d.hpp

 // SPDX-License-Identifier: MIT
-// Copyright (c) 2018-2022, Advanced Micro Devices, Inc. All rights reserved.
+// Copyright (c) 2018-2023, Advanced Micro Devices, Inc. All rights reserved.
 
 #pragma once
 
@@ -14,11 +14,12 @@
 
 #include "ck/library/tensor_operation_instance/device_operation_instance_factory.hpp"
 
+#ifdef DL_KERNELS
 namespace ck {
 namespace tensor_operation {
 namespace device {
 namespace instance {
-
+#ifdef CK_ENABLE_FP16
 void add_device_batched_gemm_multi_d_dl_f16_f16_f16_gkm_gkn_gmn_instances(
     std::vector<std::unique_ptr<DeviceBatchedGemmMultiD<Col,
                                                         Row,
@@ -122,7 +123,8 @@ void add_device_batched_gemm_multi_d_dl_f16_f16_f16_gmk_gnk_gmn_irregular_instan
                                                         PassThrough,
                                                         PassThrough,
                                                         PassThrough>>>& instances);
-
+#endif
+#ifdef CK_ENABLE_INT8
 void add_device_batched_gemm_multi_d_dl_i8_i8_i8_gkm_gkn_gmn_instances(
     std::vector<std::unique_ptr<DeviceBatchedGemmMultiD<Col,
                                                         Row,
@@ -226,7 +228,7 @@ void add_device_batched_gemm_multi_d_dl_i8_i8_i8_gmk_gnk_gmn_irregular_instances
                                                         PassThrough,
                                                         PassThrough,
                                                         PassThrough>>>& instances);
-
+#endif
 template <typename ALayout,
           typename BLayout,
           typename ELayout,
@@ -261,7 +263,7 @@ struct DeviceOperationInstanceFactory<ck::tensor_operation::device::DeviceBatche
     static auto GetInstances()
     {
         std::vector<std::unique_ptr<DeviceOp>> op_ptrs;
-
+#ifdef CK_ENABLE_FP16
         if constexpr(is_same_v<ADataType, half_t> && is_same_v<BDataType, half_t> &&
                      is_same_v<EDataType, half_t>)
         {
@@ -294,6 +296,8 @@ struct DeviceOperationInstanceFactory<ck::tensor_operation::device::DeviceBatche
                     op_ptrs);
             }
         }
+#endif
+#ifdef CK_ENABLE_INT8
         else if constexpr(is_same_v<ADataType, int8_t> && is_same_v<BDataType, int8_t> &&
                           is_same_v<EDataType, int8_t>)
         {
@@ -326,7 +330,7 @@ struct DeviceOperationInstanceFactory<ck::tensor_operation::device::DeviceBatche
                     op_ptrs);
             }
         }
-
+#endif
         return op_ptrs;
     }
 };
@@ -335,3 +339,4 @@ struct DeviceOperationInstanceFactory<ck::tensor_operation::device::DeviceBatche
 } // namespace device
 } // namespace tensor_operation
 } // namespace ck
+#endif

library/include/ck/library/tensor_operation_instance/gpu/batched_gemm_softmax_gemm.hpp

@@ -11,7 +11,7 @@
 #include "ck/tensor_operation/gpu/element/element_wise_operation.hpp"
 
 #include "ck/library/tensor_operation_instance/device_operation_instance_factory.hpp"
-
+#ifdef CK_ENABLE_FP16
 namespace ck {
 namespace tensor_operation {
 namespace device {
@@ -119,3 +119,4 @@ struct DeviceOperationInstanceFactory<
 } // namespace device
 } // namespace tensor_operation
 } // namespace ck
+#endif