Implement IsSupportedArgument

cf9114fd · rocking · e279b402 · cf9114fd
Commit cf9114fd authored Jun 08, 2023 by rocking
Hide whitespace changes
Inline Side-by-side

Showing with 6 additions and 2 deletions

include/ck/tensor_operation/gpu/device/impl/device_put_element_impl.hpp ...sor_operation/gpu/device/impl/device_put_element_impl.hpp +6 -2

No files found.
--- a/include/ck/tensor_operation/gpu/device/impl/device_put_element_impl.hpp
+++ b/include/ck/tensor_operation/gpu/device/impl/device_put_element_impl.hpp
@@ -121,8 +121,12 @@ struct DevicePutElementImpl
    bool IsSupportedArgument(const BaseArgument* p_arg) override
    {
        const Argument* pArg = dynamic_cast<const Argument*>(p_arg);
-        // TODO
-        ignore = pArg;
+        index_t input_length = pArg->in_grid_desc_.GetTransforms()[I0].GetUpperLengths()[I0];
+
+        if(input_length % InVectorSize != 0)
+        {
+            return false;
+        }
        return true;
    }