Address review: use make_tensor_ptr and store defaults by value

Replace manual TensorImpl construction with make_tensor_ptr from extension/tensor, removing the #ifdef USE_ATEN_LIB block and simplifying the destructor. Store defaults_ by value since it is always initialized.

Author: BryanBradfo

extension/training/optimizer/adamw.cpp

@@ -8,13 +8,15 @@
 
 #include <executorch/extension/training/optimizer/adamw.h>
 
+#include <executorch/extension/tensor/tensor_ptr.h>
 #include <executorch/runtime/core/error.h>
 
 #include <cmath>
 #include <cstring>
 
 using executorch::aten::Tensor;
-using executorch::aten::TensorImpl;
+using executorch::extension::make_tensor_ptr;
+using executorch::extension::TensorPtr;
 using ::executorch::runtime::Error;
 
 namespace executorch {
@@ -116,16 +118,15 @@ AdamWParamGroup::named_parameters() const {
 void AdamW::add_param_group(const AdamWParamGroup& param_group) {
   AdamWParamGroup param_group_(param_group.named_parameters());
   if (!param_group.has_options()) {
-    param_group_.set_options(defaults_->clone());
+    param_group_.set_options(defaults_.clone());
   } else {
     param_group_.set_options(param_group.options().clone());
   }
   param_groups_.emplace_back(std::move(param_group_));
 }
 
-Error AdamW::step(
-    const std::map<std::string_view, executorch::aten::Tensor>&
-        named_gradients) {
+Error AdamW::step(const std::map<std::string_view, executorch::aten::Tensor>&
+                      named_gradients) {
   for (auto& group : param_groups_) {
     auto& options = static_cast<AdamWOptions&>(group.options());
     const double lr = options.lr();
@@ -162,29 +163,23 @@ Error AdamW::step(
         std::memset(m_buf_ptr, 0, g.nbytes());
         std::memset(v_buf_ptr, 0, g.nbytes());
 
-        Tensor m_buf(nullptr);
-        Tensor v_buf(nullptr);
-#ifdef USE_ATEN_LIB
-        std::vector<int64_t> sizes(g.sizes().begin(), g.sizes().end());
-        m_buf = torch::from_blob(m_buf_ptr, sizes, g.scalar_type());
-        v_buf = torch::from_blob(v_buf_ptr, sizes, g.scalar_type());
-#else
-        TensorImpl* m_impl = new TensorImpl(
-            g.scalar_type(),
-            g.sizes().size(),
-            const_cast<TensorImpl::SizesType*>(g.sizes().data()),
+        std::vector<executorch::aten::SizesType> sizes(
+            g.sizes().begin(), g.sizes().end());
+        auto m_ptr = make_tensor_ptr(
+            sizes,
             m_buf_ptr,
-            const_cast<TensorImpl::DimOrderType*>(g.dim_order().data()));
-        TensorImpl* v_impl = new TensorImpl(
             g.scalar_type(),
-            g.sizes().size(),
-            const_cast<TensorImpl::SizesType*>(g.sizes().data()),
+            executorch::aten::TensorShapeDynamism::STATIC,
+            [](void* p) { free(p); });
+        auto v_ptr = make_tensor_ptr(
+            sizes,
             v_buf_ptr,
-            const_cast<TensorImpl::DimOrderType*>(g.dim_order().data()));
-        m_buf = Tensor(m_impl);
-        v_buf = Tensor(v_impl);
-#endif
-        auto state = std::make_unique<AdamWParamState>(m_buf, v_buf);
+            g.scalar_type(),
+            executorch::aten::TensorShapeDynamism::STATIC,
+            [](void* p) { free(p); });
+
+        auto state = std::make_unique<AdamWParamState>(
+            std::move(m_ptr), std::move(v_ptr));
         state_ptr = state.get();
         state_[p.unsafeGetTensorImpl()] = std::move(state);
       } else {
@@ -214,17 +209,7 @@ Error AdamW::step(
   return Error::Ok;
 }
 
-AdamW::~AdamW() {
-  for (const auto& state_kv : state_) {
-    auto& state = *state_kv.second;
-    free(state.exp_avg().unsafeGetTensorImpl()->mutable_data());
-    free(state.exp_avg_sq().unsafeGetTensorImpl()->mutable_data());
-#ifndef USE_ATEN_LIB
-    delete state.exp_avg().unsafeGetTensorImpl();
-    delete state.exp_avg_sq().unsafeGetTensorImpl();
-#endif
-  }
-}
+AdamW::~AdamW() = default;
 
 } // namespace optimizer
 } // namespace training

extension/training/optimizer/adamw.h

@@ -15,6 +15,7 @@
  */
 #pragma once
 
+#include <executorch/extension/tensor/tensor_ptr.h>
 #include <executorch/runtime/core/error.h>
 #include <executorch/runtime/core/exec_aten/exec_aten.h>
 #include <cstdint>
@@ -41,16 +42,18 @@ class ET_EXPERIMENTAL AdamWParamState {
    * @param[in] exp_avg_sq The second moment (EMA of squared gradients) buffer.
    */
   AdamWParamState(
-      executorch::aten::Tensor& exp_avg,
-      executorch::aten::Tensor& exp_avg_sq)
-      : exp_avg_(exp_avg), exp_avg_sq_(exp_avg_sq), step_count_(0) {}
+      executorch::extension::TensorPtr exp_avg,
+      executorch::extension::TensorPtr exp_avg_sq)
+      : exp_avg_(std::move(exp_avg)),
+        exp_avg_sq_(std::move(exp_avg_sq)),
+        step_count_(0) {}
 
   executorch::aten::Tensor& exp_avg() {
-    return exp_avg_;
+    return *exp_avg_;
   }
 
   executorch::aten::Tensor& exp_avg_sq() {
-    return exp_avg_sq_;
+    return *exp_avg_sq_;
   }
 
   int64_t step_count() const {
@@ -62,8 +65,8 @@ class ET_EXPERIMENTAL AdamWParamState {
   }
 
  private:
-  executorch::aten::Tensor exp_avg_;
-  executorch::aten::Tensor exp_avg_sq_;
+  executorch::extension::TensorPtr exp_avg_;
+  executorch::extension::TensorPtr exp_avg_sq_;
   int64_t step_count_;
 };
 
@@ -179,7 +182,7 @@ class ET_EXPERIMENTAL AdamW {
   explicit AdamW(
       const std::vector<AdamWParamGroup>& param_groups,
       AdamWOptions defaults)
-      : defaults_(std::make_unique<AdamWOptions>(defaults)) {
+      : defaults_(defaults) {
     for (const auto& param_group : param_groups) {
       add_param_group(param_group);
     }
@@ -209,7 +212,7 @@ class ET_EXPERIMENTAL AdamW {
  private:
   std::vector<AdamWParamGroup> param_groups_;
   std::unordered_map<void*, std::unique_ptr<AdamWParamState>> state_;
-  std::unique_ptr<AdamWOptions> defaults_;
+  AdamWOptions defaults_;
 };
 
 } // namespace optimizer

extension/training/optimizer/targets.bzl

@@ -48,6 +48,7 @@ def define_common_targets():
                 "adamw.h",
             ],
             exported_deps = [
+                "//executorch/extension/tensor:tensor" + aten_suffix,
                 "//executorch/runtime/core:core",
                 "//executorch/runtime/core/exec_aten:lib" + aten_suffix,
             ],

extension/training/optimizer/test/adamw_test.cpp

@@ -6,6 +6,7 @@
  * LICENSE file in the root directory of this source tree.
  */
 
+#include <executorch/extension/tensor/tensor_ptr.h>
 #include <executorch/extension/training/optimizer/adamw.h>
 #include <executorch/runtime/core/exec_aten/exec_aten.h>
 #include <executorch/runtime/core/exec_aten/testing_util/tensor_factory.h>
@@ -33,10 +34,11 @@ class AdamWOptimizerTest : public ::testing::Test {
 };
 
 TEST_F(AdamWOptimizerTest, AdamWParamStateTest) {
-  TensorFactory<ScalarType::Float> tf;
-  Tensor exp_avg = tf.make({2, 2}, {0, 0, 0, 0});
-  Tensor exp_avg_sq = tf.make({2, 2}, {0, 0, 0, 0});
-  AdamWParamState state(exp_avg, exp_avg_sq);
+  auto exp_avg =
+      executorch::extension::make_tensor_ptr({2, 2}, {0.f, 0.f, 0.f, 0.f});
+  auto exp_avg_sq =
+      executorch::extension::make_tensor_ptr({2, 2}, {0.f, 0.f, 0.f, 0.f});
+  AdamWParamState state(std::move(exp_avg), std::move(exp_avg_sq));
 
   EXPECT_EQ(state.step_count(), 0);
   state.increment_step_count();
@@ -78,8 +80,8 @@ TEST_F(AdamWOptimizerTest, AdamWOptimizerSimple) {
     optimizer.step(named_gradients);
   }
 
-  auto p1 = static_cast<const float*>(
-      named_parameters.at("param1").const_data_ptr());
+  auto p1 =
+      static_cast<const float*>(named_parameters.at("param1").const_data_ptr());
   // With a constant gradient of -1 and no weight decay, the bias-corrected
   // m_hat / sqrt(v_hat) is ~= -1 at every step, so each step shifts p by
   // +lr. After 10 steps of lr=0.1, p should be near 2.0.
@@ -103,8 +105,8 @@ TEST_F(AdamWOptimizerTest, AdamWOptimizerDecoupledWeightDecay) {
   named_gradients.insert({"param1", tf.make({1, 1}, {0.0})});
   optimizer.step(named_gradients);
 
-  auto p1 = static_cast<const float*>(
-      named_parameters.at("param1").const_data_ptr());
+  auto p1 =
+      static_cast<const float*>(named_parameters.at("param1").const_data_ptr());
   EXPECT_NEAR(p1[0], 0.95, 1e-5);
 }
 
@@ -124,10 +126,10 @@ TEST_F(AdamWOptimizerTest, AdamWOptimizerMultipleParams) {
     optimizer.step(named_gradients);
   }
 
-  auto p1 = static_cast<const float*>(
-      named_parameters.at("param1").const_data_ptr());
-  auto p2 = static_cast<const float*>(
-      named_parameters.at("param2").const_data_ptr());
+  auto p1 =
+      static_cast<const float*>(named_parameters.at("param1").const_data_ptr());
+  auto p2 =
+      static_cast<const float*>(named_parameters.at("param2").const_data_ptr());
   // Each param sees a constant gradient of +/- 1 for 5 steps -> p shifts by
   // roughly +/- 5 * lr = +/- 0.5. State is tracked independently per param.
   EXPECT_NEAR(p1[0], 1.5, 0.1);