[ExecuTorch][Vulkan] Tests for et_vk.adamw_step

backends/vulkan/test/op_tests/CMakeLists.txt

@@ -116,6 +116,9 @@ if(TARGET vulkan_backend AND LIB_TORCH)
     quantized_linear_test ${CMAKE_CURRENT_SOURCE_DIR}/quantized_linear_test.cpp
     test_utils
   )
+  vulkan_op_test(
+    adamw_step_test ${CMAKE_CURRENT_SOURCE_DIR}/adamw_step_test.cpp test_utils
+  )
 
   # Only build generated op tests if a path to tags.yaml and
   # native_functions.yaml is provided. These files are required for codegen.

backends/vulkan/test/op_tests/adamw_step_test.cpp

@@ -0,0 +1,191 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#include <gtest/gtest.h>
+
+#include <ATen/ATen.h>
+
+#include <executorch/backends/vulkan/runtime/api/api.h>
+#include <executorch/backends/vulkan/runtime/graph/ComputeGraph.h>
+#include <executorch/backends/vulkan/runtime/graph/ops/OperatorRegistry.h>
+
+#include "test_utils.h"
+
+#include <cmath>
+#include <tuple>
+
+//
+// Reference implementation: an ATen transcription of the CPU-eager
+// adamw_step_impl golden in backends/vulkan/custom_ops_lib.py.
+//
+
+std::tuple<at::Tensor, at::Tensor, at::Tensor> adamw_reference_impl(
+    const at::Tensor& param_in,
+    const at::Tensor& m_in,
+    const at::Tensor& v_in,
+    const at::Tensor& grad,
+    const double lr,
+    const double beta1,
+    const double beta2,
+    const double eps,
+    const double weight_decay,
+    const double bias_correction1,
+    const double bias_correction2) {
+  at::Tensor param = param_in.clone();
+  at::Tensor m = m_in.clone();
+  at::Tensor v = v_in.clone();
+
+  param.mul_(1.0 - lr * weight_decay);
+  m.mul_(beta1).add_(grad, 1.0 - beta1);
+  v.mul_(beta2).addcmul_(grad, grad, 1.0 - beta2);
+  at::Tensor mhat = m / bias_correction1;
+  at::Tensor denom = (v / bias_correction2).sqrt() + eps;
+  param.addcdiv_(mhat, denom, -lr);
+
+  return std::make_tuple(param, m, v);
+}
+
+//
+// Test harness
+//
+
+void test_vulkan_adamw_step(
+    const std::vector<int64_t>& sizes,
+    const double lr,
+    const double beta1,
+    const double beta2,
+    const double eps,
+    const double weight_decay,
+    const int step) {
+  torch::manual_seed(0);
+
+  const double bias_correction1 = 1.0 - std::pow(beta1, step);
+  const double bias_correction2 = 1.0 - std::pow(beta2, step);
+
+  const auto options = at::device(at::kCPU).dtype(at::kFloat);
+  at::Tensor param = at::randn(sizes, options);
+  at::Tensor m = at::randn(sizes, options);
+  at::Tensor v = at::rand(sizes, options); // second moment is non-negative
+  at::Tensor grad = at::randn(sizes, options);
+
+  at::Tensor ref_param;
+  at::Tensor ref_m;
+  at::Tensor ref_v;
+  std::tie(ref_param, ref_m, ref_v) = adamw_reference_impl(
+      param,
+      m,
+      v,
+      grad,
+      lr,
+      beta1,
+      beta2,
+      eps,
+      weight_decay,
+      bias_correction1,
+      bias_correction2);
+
+  using namespace vkcompute;
+
+  GraphConfig config;
+  ComputeGraph graph(config);
+
+  ValueRef r_param = graph.add_tensor(sizes, vkapi::kFloat, utils::kBuffer);
+  ValueRef r_m = graph.add_tensor(sizes, vkapi::kFloat, utils::kBuffer);
+  ValueRef r_v = graph.add_tensor(sizes, vkapi::kFloat, utils::kBuffer);
+  ValueRef r_grad = graph.add_tensor(sizes, vkapi::kFloat, utils::kBuffer);
+
+  ValueRef staging_param = graph.set_input_tensor(r_param);
+  ValueRef staging_m = graph.set_input_tensor(r_m);
+  ValueRef staging_v = graph.set_input_tensor(r_v);
+  ValueRef staging_grad = graph.set_input_tensor(r_grad);
+
+  VK_GET_OP_FN("et_vk.adamw_step.default")
+  (graph,
+   {
+       r_param,
+       r_m,
+       r_v,
+       r_grad,
+       graph.add_scalar<double>(lr),
+       graph.add_scalar<double>(beta1),
+       graph.add_scalar<double>(beta2),
+       graph.add_scalar<double>(eps),
+       graph.add_scalar<double>(weight_decay),
+       graph.add_scalar<double>(bias_correction1),
+       graph.add_scalar<double>(bias_correction2),
+   });
+
+  ValueRef staging_param_out = graph.set_output_tensor(r_param);
+  ValueRef staging_m_out = graph.set_output_tensor(r_m);
+  ValueRef staging_v_out = graph.set_output_tensor(r_v);
+
+  graph.prepare();
+  graph.prepack();
+
+  graph.maybe_cast_and_copy_into_staging(
+      staging_param, param.const_data_ptr(), param.numel(), vkapi::kFloat);
+  graph.maybe_cast_and_copy_into_staging(
+      staging_m, m.const_data_ptr(), m.numel(), vkapi::kFloat);
+  graph.maybe_cast_and_copy_into_staging(
+      staging_v, v.const_data_ptr(), v.numel(), vkapi::kFloat);
+  graph.maybe_cast_and_copy_into_staging(
+      staging_grad, grad.const_data_ptr(), grad.numel(), vkapi::kFloat);
+
+  graph.execute();
+
+  at::Tensor vk_param = at::empty_like(param);
+  at::Tensor vk_m = at::empty_like(m);
+  at::Tensor vk_v = at::empty_like(v);
+
+  graph.maybe_cast_and_copy_from_staging(
+      staging_param_out,
+      vk_param.mutable_data_ptr(),
+      vk_param.numel(),
+      vkapi::kFloat);
+  graph.maybe_cast_and_copy_from_staging(
+      staging_m_out, vk_m.mutable_data_ptr(), vk_m.numel(), vkapi::kFloat);
+  graph.maybe_cast_and_copy_from_staging(
+      staging_v_out, vk_v.mutable_data_ptr(), vk_v.numel(), vkapi::kFloat);
+
+  ASSERT_TRUE(at::allclose(ref_param, vk_param, /*rtol=*/1e-4, /*atol=*/1e-4));
+  ASSERT_TRUE(at::allclose(ref_m, vk_m, /*rtol=*/1e-4, /*atol=*/1e-4));
+  ASSERT_TRUE(at::allclose(ref_v, vk_v, /*rtol=*/1e-4, /*atol=*/1e-4));
+}
+
+TEST(VulkanAdamwStepTest, test_adamw_step_no_weight_decay) {
+  test_vulkan_adamw_step(
+      /*sizes=*/{4, 8},
+      /*lr=*/1e-3,
+      /*beta1=*/0.9,
+      /*beta2=*/0.999,
+      /*eps=*/1e-8,
+      /*weight_decay=*/0.0,
+      /*step=*/1);
+}
+
+TEST(VulkanAdamwStepTest, test_adamw_step_with_weight_decay) {
+  test_vulkan_adamw_step(
+      /*sizes=*/{3, 5, 7},
+      /*lr=*/1e-2,
+      /*beta1=*/0.9,
+      /*beta2=*/0.999,
+      /*eps=*/1e-8,
+      /*weight_decay=*/0.01,
+      /*step=*/10);
+}
+
+TEST(VulkanAdamwStepTest, test_adamw_step_large_1d) {
+  test_vulkan_adamw_step(
+      /*sizes=*/{1000},
+      /*lr=*/1e-3,
+      /*beta1=*/0.9,
+      /*beta2=*/0.999,
+      /*eps=*/1e-8,
+      /*weight_decay=*/0.05,
+      /*step=*/5);
+}

backends/vulkan/test/op_tests/targets.bzl

@@ -180,6 +180,12 @@ def define_common_targets(is_fbcode = False):
             ":test_utils",
         ]
     )
+    define_test_targets(
+        "adamw_step_test",
+        extra_deps = [
+            ":test_utils",
+        ]
+    )
     define_test_targets(
         "rms_norm_test",
         extra_deps = [