Lower exp with the existing unary kernel

backends/native/runtime/engine/ggml/GgmlGraphBuilder.cpp

@@ -684,6 +684,8 @@ void GraphBuilder::lower(const Node& node) {
     lower_constant_pad(node);
   } else if (key == "aten.log.default") {
     lower_unary(node, ggml_log);
+  } else if (key == "aten.exp.default") {
+    lower_unary(node, ggml_exp);
   } else if (key == "aten.neg.default") {
     lower_unary(node, ggml_neg);
   } else if (key == "aten.relu.default" || key == "aten.relu_.default") {

backends/native/test/runtime/ggml/test_ggml_engine.cpp

@@ -615,6 +615,89 @@ TEST(GgmlEngineTest, RejectsIntegerGelu) {
   expect_compile_error(program.finish("y"), "floating-point");
 }
 
+// exp of x [4, 3], or of its transpose.
+std::vector<uint8_t> make_exp_program(
+    bool transposed,
+    fbs::ScalarType dtype = fbs::ScalarType::FLOAT) {
+  TestProgram program;
+  program.input("x", dtype, {4, 3});
+  std::string self = "x";
+  if (transposed) {
+    program.call(
+        "xt",
+        "torch.ops.aten.permute_copy.default",
+        {{"self", std::string("x")}, {"dims", std::vector<int64_t>{1, 0}}},
+        dtype,
+        {3, 4});
+    self = "xt";
+  }
+  program.call(
+      "y",
+      "torch.ops.aten.exp.default",
+      {{"self", self}},
+      dtype,
+      transposed ? std::vector<int64_t>{3, 4} : std::vector<int64_t>{4, 3});
+  return program.finish("y");
+}
+
+// Within 1e-6 relative of exp in double; NaN, infinities and zero exactly.
+TEST(GgmlEngineTest, ExpMatchesLibmIncludingNonFiniteValues) {
+  constexpr float kInf = std::numeric_limits<float>::infinity();
+  const std::vector<float> x = {
+      -0.0f,
+      0.0f,
+      1.5f,
+      -2.25f,
+      30.0f,
+      -30.0f,
+      100.0f,
+      -80.0f,
+      kInf,
+      -kInf,
+      std::nanf(""),
+      1e-20f};
+  for (const bool transposed : {false, true}) {
+    for (const size_t cached : {size_t{0}, size_t{1}}) {
+      SCOPED_TRACE(std::to_string(transposed) + std::to_string(cached));
+      auto host = GgmlEngineHost::create(
+          GgmlEngineOptions{.max_cached_graphs = cached});
+      auto context = make_context(*host, make_exp_program(transposed));
+      auto executable = context->compile("forward");
+      for (const float scale : {1.0f, -0.5f}) {
+        std::vector<float> input = x;
+        for (float& value : input) {
+          value *= scale;
+        }
+        executable->set_input(0, input.data(), input.size(), ScalarType::Float);
+        executable->execute();
+        std::vector<float> y(input.size());
+        executable->get_output(0, y.data(), y.size(), ScalarType::Float);
+        for (size_t i = 0; i < y.size(); ++i) {
+          // y[r][c] reads x[c][r] when transposed.
+          const size_t source = transposed ? (i % 4) * 3 + i / 4 : i;
+          const auto expected =
+              static_cast<float>(std::exp(static_cast<double>(input[source])));
+          SCOPED_TRACE(input[source]);
+          if (std::isnan(expected)) {
+            EXPECT_TRUE(std::isnan(y[i])) << y[i];
+          } else if (std::isinf(expected) || expected == 0.0f) {
+            EXPECT_EQ(y[i], expected);
+            EXPECT_EQ(std::signbit(y[i]), std::signbit(expected));
+          } else {
+            EXPECT_NEAR(y[i], expected, 1e-6 * std::abs(expected));
+          }
+        }
+      }
+    }
+  }
+}
+
+TEST(GgmlEngineTest, RejectsIntegerExp) {
+  expect_compile_error(
+      make_exp_program(false, fbs::ScalarType::INT),
+      "supports floating-point tensors only");
+}
+
 TEST(GgmlEngineTest, RejectsUserInputMutation) {
   TestProgram program;
   program.input("x", fbs::ScalarType::FLOAT, {2, 3});

backends/native/test/runtime/ggml/test_ggml_ops.py

@@ -617,6 +617,9 @@ def test_gelu_tanh(self):
     def test_log(self):
         self._check(Lambda(torch.log), (torch.rand(2, 3, 4) + 0.1,))
 
+    def test_exp(self):
+        self._check(Lambda(torch.exp), (torch.randn(3, 1, 1).clamp(max=4.6),))
+
     def test_neg(self):
         self._check(Lambda(torch.neg), (torch.randn(2, 3, 4),))