/home/docs/checkouts/readthedocs.org/user_builds/advanced-micro-devices-composable-kernel/checkouts/docs-6.4.3/include/ck/tensor_operation/gpu/device/device_elementwise_scale.hpp Source File

/home/docs/checkouts/readthedocs.org/user_builds/advanced-micro-devices-composable-kernel/checkouts/docs-6.4.3/include/ck/tensor_operation/gpu/device/device_elementwise_scale.hpp Source File#

Composable Kernel: /home/docs/checkouts/readthedocs.org/user_builds/advanced-micro-devices-composable-kernel/checkouts/docs-6.4.3/include/ck/tensor_operation/gpu/device/device_elementwise_scale.hpp Source File

Go to the documentation of this file.

 // SPDX-License-Identifier: MIT
 // Copyright (c) 2018-2024, Advanced Micro Devices, Inc. All rights reserved.
  
 #pragma once
  
 #include <memory>
 #include <array>
  
 #include "ck/ck.hpp"
 #include "ck/tensor_operation/gpu/device/device_base.hpp"
  
 namespace ck {
 namespace tensor_operation {
 namespace device {
  
 template <typename InDataTypeTuple,
           typename OutDataTypeTuple,
           typename ElementwiseOperation,
           typename UnaryOperation,
           typename Scale,
           index_t NumDim>
 struct DeviceElementwise : public BaseOperator
 {
     static constexpr int NumInput  = InDataTypeTuple::Size();
     static constexpr int NumOutput = OutDataTypeTuple::Size();
  
     virtual std::unique_ptr<BaseArgument>
     MakeArgumentPointer(const std::array<index_t, NumDim> lengths,
                         const std::array<std::array<index_t, NumDim>, NumInput> inStridesArray,
                         const std::array<std::array<index_t, NumDim>, NumOutput> outStridesArray,
                         const std::array<const void*, NumInput> in_dev_buffers,
                         const std::array<void*, NumOutput> out_dev_buffers,
                         ElementwiseOperation elementwise_op,
                         UnaryOperation unary_op,
                         Scale scale_op) = 0;
  
     virtual std::unique_ptr<BaseInvoker> MakeInvokerPointer() = 0;
 }; // namespace device
  
 template <typename InDataTypeTuple,
           typename OutDataTypeTuple,
           typename ElementwiseOperation,
           typename UnaryOperation,
           typename Scale,
           index_t NumDim>
 using DeviceElementwisePtr = std::unique_ptr<DeviceElementwise<InDataTypeTuple,
                                                                OutDataTypeTuple,
                                                                ElementwiseOperation,
                                                                UnaryOperation,
                                                                Scale,
                                                                NumDim>>;
  
 } // namespace device
 } // namespace tensor_operation
 } // namespace ck