diff --git a/examples/dl-activations/relu3.py b/examples/dl-activations/relu3.py
index 2d97d7c88..054ae2c6a 100644
--- a/examples/dl-activations/relu3.py
+++ b/examples/dl-activations/relu3.py
@@ -5,7 +5,10 @@
 from collections import OrderedDict
 import ksc.expr as expr
 from ksc.type import Type
-from ksc.torch_frontend import ksc_string_to_autograd_function
+from ksc.torch_frontend import (
+    ksc_string_to_autograd_function,
+    cpp_string_to_autograd_function,
+)
 from ksc.utils import get_ksc_paths
 import torch._vmap_internals
 
@@ -58,6 +61,161 @@ def vrelu3_embedded_ks_checkpointed_map():
     )
 
 
+def vrelu3_embedded_cpp_inlined_map():
+    return cpp_string_to_autograd_function(
+        """
+        #include "knossos.h"
+
+        namespace ks{
+        tensor<1, ks::Float> vrelu3(ks::allocator * $alloc, tensor<1, ks::Float> t) {
+            auto tdata = t.data();
+            auto ret = tensor<1, ks::Float>::create($alloc, t.size());
+            auto retdata = ret.data();
+            for (int i = 0, ne = t.num_elements(); i != ne; ++i) {
+                ks::Float c$1;
+                ks::Float x = tdata[i];
+                if (x < 0.0) {
+                    c$1 = 0.0;
+                } else {
+                    if (x < 1.0) {
+                        c$1 = x * x * x / 3.0;
+                    } else {
+                        c$1 = x - 2.0 / 3.0;
+                    }
+                }
+                retdata[i] = c$1;
+            }
+            return ret;
+        }
+
+        tensor<1, ks::Float> sufrev_vrelu3(ks::allocator * $alloc, tensor<1, ks::Float> t, tensor<1, ks::Float> dret) {
+            auto tdata = t.data();
+            auto dretdata = dret.data();
+            auto ret = tensor<1, ks::Float>::create($alloc, t.size());
+            auto retdata = ret.data();
+            for (int i = 0, ne = t.num_elements(); i != ne; ++i) {
+                ks::Float c$1;
+                ks::Float x = tdata[i];
+                ks::Float dreti = dretdata[i];
+                if (x < 0.0) {
+                    c$1 = 0.0;
+                } else {
+                    if (x < 1.0) {
+                        c$1 = x * x;
+                    } else {
+                        c$1 = 1.0;
+                    }
+                }
+                retdata[i] = c$1 * dreti;
+            }
+            return ret;
+        }
+        }
+        """,
+        "vrelu3",
+        generate_lm=False,
+    )
+
+
+def vrelu3_embedded_cpp_mask():
+    return cpp_string_to_autograd_function(
+        """
+        #include "knossos.h"
+
+        namespace ks{
+        tensor<1, ks::Float> vrelu3(ks::allocator * $alloc, tensor<1, ks::Float> t) {
+            auto tdata = t.data();
+            auto ret = tensor<1, ks::Float>::create($alloc, t.size());
+            auto retdata = ret.data();
+            ks::Float third = 1.0 / 3.0;
+            ks::Float two_thirds = 2.0 / 3.0;
+            for (int i = 0, ne = t.num_elements(); i != ne; ++i) {
+                ks::Float x = tdata[i];
+                auto val0to1 = third * x * x * x;
+                auto val1up = x - two_thirds;
+                auto le1 = x <= 1;
+
+                retdata[i] = (x>0)*(le1*val0to1 + (!le1)*val1up);
+            }
+            return ret;
+        }
+
+        tensor<1, ks::Float> sufrev_vrelu3(ks::allocator * $alloc, tensor<1, ks::Float> t, tensor<1, ks::Float> dret) {
+            auto tdata = t.data();
+            auto dretdata = dret.data();
+            auto ret = tensor<1, ks::Float>::create($alloc, t.size());
+            auto retdata = ret.data();
+            for (int i = 0, ne = t.num_elements(); i != ne; ++i) {
+                ks::Float x = tdata[i];
+                ks::Float dreti = dretdata[i];
+                auto val0to1 = x * x;
+                auto val1up = 1.0;
+
+                auto le1 = x <= 1;
+
+                retdata[i] = (x>0)*(le1*val0to1 + (!le1)*val1up)*dreti;
+            }
+            return ret;
+        }
+        }
+        """,
+        "vrelu3",
+        generate_lm=False,
+    )
+
+
+def vrelu3_embedded_cpp_mask_bool_to_float():
+    return cpp_string_to_autograd_function(
+        """
+        #include "knossos.h"
+
+        namespace ks{
+        tensor<1, ks::Float> vrelu3(ks::allocator * $alloc, tensor<1, ks::Float> t) {
+            auto tdata = t.data();
+            auto ret = tensor<1, ks::Float>::create($alloc, t.size());
+            auto retdata = ret.data();
+            ks::Float third = 1.0 / 3.0;
+            ks::Float two_thirds = 2.0 / 3.0;
+            for (int i = 0, ne = t.num_elements(); i != ne; ++i) {
+                ks::Float x = tdata[i];
+                auto val0to1 = third * x * x * x;
+                auto val1up = x - two_thirds;
+                auto le1 = x <= 1;
+
+                retdata[i] = bool_to_float$ab($alloc, x > 0)
+                    * (bool_to_float$ab($alloc, le1) * val0to1
+                       + bool_to_float$ab($alloc, !le1) * val1up);
+            }
+            return ret;
+        }
+
+        tensor<1, ks::Float> sufrev_vrelu3(ks::allocator * $alloc, tensor<1, ks::Float> t, tensor<1, ks::Float> dret) {
+            auto tdata = t.data();
+            auto dretdata = dret.data();
+            auto ret = tensor<1, ks::Float>::create($alloc, t.size());
+            auto retdata = ret.data();
+            for (int i = 0, ne = t.num_elements(); i != ne; ++i) {
+                ks::Float x = tdata[i];
+                ks::Float dreti = dretdata[i];
+                auto val0to1 = x * x;
+                auto val1up = 1.0;
+
+                auto le1 = x <= 1;
+
+                retdata[i] = bool_to_float$ab($alloc, x > 0)
+                             * (bool_to_float$ab($alloc, le1) * val0to1
+                                + bool_to_float$ab($alloc, !le1) * val1up)
+                             * dreti;
+            }
+            return ret;
+        }
+        }
+        """,
+        "vrelu3",
+        generate_lm=False,
+    )
+
+
 def vrelu3_embedded_ks_checkpointed_map_handwritten_relu3():
     return ksc_string_to_autograd_function(
         """(def relu3 Float (x : Float)
@@ -88,6 +246,97 @@ def vrelu3_embedded_ks_checkpointed_map_handwritten_relu3():
     )
 
 
+def vrelu3_embedded_ks_checkpointed_map_handwritten_inlined_relu3():
+    return ksc_string_to_autograd_function(
+        """(def [vrelu3 (Vec Float)] (Vec Float)
+                (t : Vec Float)
+                (map (lam (x : Float)
+             (if (lt x 0.0)
+                 0.0
+             (if (lt x 1.0)
+                 (div (mul x (mul x x)) 3.0)
+             (sub x (div 2.0 3.0))))) t))
+
+           (def [sufrev [vrelu3 (Vec Float)]] (Vec Float)
+                ((t : Vec Float) (dret : Vec Float))
+                (map2 (lam (x_ddri : Tuple Float Float)
+               (let ((x ddri) x_ddri)
+               (if (lt x 0.0)
+                   0.0
+               (if (lt x 1.0)
+                   (mul x x)
+               ddri)))) t dret))
+        """,
+        expr.StructuredName(("vrelu3", Type.Tensor(1, Type.Float))),
+        generate_lm=False,
+    )
+
+
+def vrelu3_embedded_ks_checkpointed_map_mask():
+    return ksc_string_to_autograd_function(
+        """(def [vrelu3 (Vec Float)] (Vec Float)
+                (t : Vec Float)
+                (map (lam (x : Float)
+                   (let (val0to1 (mul x (mul x (div x 3.0))))
+                   (let (val1up (sub x (div 2.0 3.0)))
+                   (let (le1 (lte x 1.0))
+                      (mul (bool_to_float (gt x 0.0))
+                           (add (mul (bool_to_float le1) val0to1)
+                                (mul (bool_to_float (not le1)) val1up))))))) t))
+
+           (def [sufrev [vrelu3 (Vec Float)]] (Vec Float)
+                ((t : Vec Float) (dret : Vec Float))
+                (map2 (lam (x_ddri : Tuple Float Float)
+                   (let ((x ddri) x_ddri)
+                   (let (val0to1 (mul x x))
+                   (let (val1up 1.0)
+                   (let (le1 (lte x 1.0))
+                      (mul (mul (bool_to_float (gt x 0.0))
+                                (add (mul (bool_to_float le1) val0to1)
+                                     (mul (bool_to_float (not le1)) val1up)))
+                           ddri)))))) t dret))
+        """,
+        expr.StructuredName(("vrelu3", Type.Tensor(1, Type.Float))),
+        generate_lm=False,
+    )
+
+
+def vrelu3_embedded_INCORRECT_ks_upper_bound_via_map():
+    return ksc_string_to_autograd_function(
+        """(def relu3 Float (x : Float) 0.0)
+
+           (def [sufrev [relu3 Float]] Float ((x : Float) (ddr : Float)) ddr)
+
+           (def [vrelu3 (Vec Float)] (Vec Float)
+                (t : Vec Float)
+                (map (lam (ti : Float) (relu3 ti)) t))
+
+           (def [sufrev [vrelu3 (Vec Float)]] (Vec Float)
+                ((t : Vec Float) (dret : Vec Float))
+                ; TODO: 1.0 should be dret[i] - luckily we are called with dret==1.0
+                (map (lam (ti : Float) ([sufrev [relu3 Float]] ti 1.0)) t))
+        """,
+        expr.StructuredName(("vrelu3", Type.Tensor(1, Type.Float))),
+        generate_lm=False,
+    )
+
+
+def vrelu3_embedded_INCORRECT_ks_upper_bound():
+    return ksc_string_to_autograd_function(
+        """; These are not correct but they are as fast as a Knossos
+           ; implementation could possibly be.
+           (def [vrelu3 (Vec Float)] (Vec Float)
+                (t : Vec Float) t)
+
+           (def [sufrev [vrelu3 (Vec Float)]] (Vec Float)
+                ((t : Vec Float) (dret : Vec Float))
+                dret)
+        """,
+        expr.StructuredName(("vrelu3", Type.Tensor(1, Type.Float))),
+        generate_lm=False,
+    )
+
+
 # run-bench: PyTorch reference implementation
 def vrelu3_pytorch(x: torch.Tensor):
     mask1_inf = x > 1.0
@@ -148,10 +397,39 @@ def forward(self, input):
     return VReLu3()
 
 
+def vrelu3_aten():
+    this_dir = os.path.dirname(__file__)
+
+    vrelu3_aten = torch.utils.cpp_extension.load(
+        "vrelu3_aten_module", sources=[os.path.join(this_dir, "vrelu3_aten.cpp"),],
+    )
+
+    class VReLu3AtenFunction(torch.autograd.Function):
+        @staticmethod
+        def forward(ctx, input):
+            output = vrelu3_aten.forward(input)
+            ctx.save_for_backward(input)
+            return output
+
+        @staticmethod
+        def backward(ctx, grad):
+            return vrelu3_aten.backward(grad, *ctx.saved_variables)
+
+    class VReLu3Aten(torch.nn.Module):
+        def __init__(self):
+            super(VReLu3Aten, self).__init__()
+
+        def forward(self, input):
+            return VReLu3AtenFunction.apply(input)
+
+    return VReLu3Aten()
+
+
 # run-bench: Define a range of values at which to call the methods
 def vrelu3_bench_configs():
-    yield torch.randn((4,))
     yield torch.randn((16,))
+    yield torch.randn((255 * 255,))
+    yield torch.randn((1024 * 1024,))
 
 
 # yield torch.randn((256,256)) too slow to bench...
diff --git a/examples/dl-activations/vrelu3_aten.cpp b/examples/dl-activations/vrelu3_aten.cpp
new file mode 100644
index 000000000..f1562e9b7
--- /dev/null
+++ b/examples/dl-activations/vrelu3_aten.cpp
@@ -0,0 +1,23 @@
+#include <torch/extension.h>
+
+torch::Tensor vrelu3_forward(torch::Tensor x) {
+  torch::Tensor mask1_inf = x > 1.0;
+  torch::Tensor mask0_1 = (x > 0.0) & ~mask1_inf;
+  torch::Tensor val_0_1 = 1.0 / 3.0 * x * x * x;
+  torch::Tensor val_1_inf = x - 2.0 / 3.0;
+
+  return mask0_1 * val_0_1 + mask1_inf * val_1_inf;
+}
+
+torch::Tensor vrelu3_backward(torch::Tensor grad, torch::Tensor x) {
+  torch::Tensor mask1_inf = x > 1.0;
+  torch::Tensor mask0_1 = (x > 0.0) & ~mask1_inf;
+  torch::Tensor val_0_1 = x * x;
+
+  return (mask0_1 * val_0_1 + mask1_inf) * grad;
+}
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+  m.def("forward", &vrelu3_forward, "vrelu3 forward");
+  m.def("backward", &vrelu3_backward, "vrelu3 backward");
+}
diff --git a/src/bench/conftest.py b/src/bench/conftest.py
index 580cf8f44..0965e5e9a 100644
--- a/src/bench/conftest.py
+++ b/src/bench/conftest.py
@@ -109,6 +109,8 @@ def functions_to_benchmark(mod, benchmark_name, example_inputs):
             elif fn_name == benchmark_name + "_cuda_init":
                 if torch.cuda.is_available():
                     yield from function_to_manual_cuda_benchmarks(fn_obj)
+            elif fn_name == benchmark_name + "_aten":
+                yield BenchmarkFunction("Aten", fn_obj())
             elif fn_name.startswith(benchmark_name + "_embedded_"):
                 n = len(benchmark_name + "_embedded_")
                 benchmark_display_name = "Embedded " + fn_name[n:]
diff --git a/src/ksc/Lang.hs b/src/ksc/Lang.hs
index 591dd6750..03085b070 100644
--- a/src/ksc/Lang.hs
+++ b/src/ksc/Lang.hs
@@ -990,6 +990,7 @@ pprPrimFun = \case
   P_buildFromSparseTupled -> text "buildFromSparseTupled"
   P_fold -> text "fold"
   P_map -> text "map"
+  P_map2 -> text "map2"
   P_index -> text "index"
   P_shape -> text "shape"
   P_size -> text "size"
@@ -1432,6 +1433,7 @@ data PrimFun = P_inline
              | P_buildFromSparse
              | P_buildFromSparseTupled
              | P_map
+             | P_map2
              | P_fold
              | P_index
              | P_shape
@@ -1489,6 +1491,7 @@ toPrimFun = \case
   "buildFromSparseTupled" -> Just P_buildFromSparseTupled
   "fold" -> Just P_fold
   "map" -> Just P_map
+  "map2" -> Just P_map2
   "index" -> Just P_index
   "shape" -> Just P_shape
   "size" -> Just P_size
diff --git a/src/ksc/Prim.hs b/src/ksc/Prim.hs
index 8dc628aa5..5fd23ee65 100644
--- a/src/ksc/Prim.hs
+++ b/src/ksc/Prim.hs
@@ -954,6 +954,10 @@ primFunCallResultTy_maybe fun args
       (P_map      , TypeTuple [TypeLam t1 t2, TypeTensor i t1'])
         | t1 `eqType` t1'
         -> Just (TypeTensor i t2)
+      (P_map2     , TypeTuple [TypeLam t tr, TypeTensor i1 t1, TypeTensor i2 t2])
+        | t `eqType` TypeTuple [t1, t2]
+        , i1 == i2
+        -> Just (TypeTensor i1 tr)
       (P_index    , TypeTuple [indexType, TypeTensor d t])
         | indexType `eqType` tensorIndexType d
         -> Just t
diff --git a/src/runtime/knossos-prelude.h b/src/runtime/knossos-prelude.h
index 946af652a..71d616bab 100644
--- a/src/runtime/knossos-prelude.h
+++ b/src/runtime/knossos-prelude.h
@@ -144,8 +144,9 @@ inline Float sqrt$af(allocator *, Float d) { return sqrt(d); }
 
 inline Float to_float$ai(allocator *, Integer d) { return d; }
 
-inline bool or$abb(allocator *, Bool b1, Bool b2)  { return b1 || b2; }
-inline bool and$abb(allocator *, Bool b1, Bool b2) { return b1 && b2; }
+inline Bool or$abb(allocator *, Bool b1, Bool b2)  { return b1 || b2; }
+inline Bool and$abb(allocator *, Bool b1, Bool b2) { return b1 && b2; }
+inline Float bool_to_float$ab(allocator *, Bool b) { return b; }
 }
 
 #include "knossos-prelude-lm.h"
diff --git a/src/runtime/knossos.h b/src/runtime/knossos.h
index daf503cc6..5efeb2563 100644
--- a/src/runtime/knossos.h
+++ b/src/runtime/knossos.h
@@ -343,7 +343,7 @@ namespace ks {
 
 		void* allocate(size_t size)
 		{
-			KS_ASSERT(size < 1000000);
+			KS_ASSERT(size < 1000 * 1000000);
 			void* ret = buf_ + top_;
 			top_ += padded_size(size);
 			if (top_ > peak_) {
@@ -1387,6 +1387,23 @@ namespace ks {
 		return ret;
 	}
 
+	// f : (S, S') -> T
+	// map2 f : (Vec S, Vec S') -> Vec T
+        template <class S, class S_, class F, size_t Dim>
+	auto // tensor<Dim, T>
+	map2(allocator * alloc, F f, tensor<Dim, S> s, tensor<Dim, S_> s_)
+	{
+                // FIXME: assert they are the same size
+                using T = decltype(applyWithAllocator(alloc, f, make_Tuple(S{}, S_{})));
+		auto ret = tensor<Dim, T>::create(alloc, s.size());
+		S const* sdata = s.data();
+		S_ const* s_data = s_.data();
+		T* retdata = ret.data();
+		for (int i = 0, ne = s.num_elements(); i != ne; ++i)
+			retdata[i] = applyWithAllocator(alloc, f, make_Tuple(sdata[i], s_data[i]));
+		return ret;
+	}
+
 	// [suffwdpass f] : S -> (T, B)
 	// suffwdpass_map : (S -> (T, B), Vec S) -> (Vec T, Vec B)
 	template <class F, class S, size_t Dim>
diff --git a/src/runtime/prelude.ks b/src/runtime/prelude.ks
index 40b7d2109..da1f1c885 100644
--- a/src/runtime/prelude.ks
+++ b/src/runtime/prelude.ks
@@ -635,3 +635,5 @@
 (def [sufrevpass [erf Float]] Float ((ddr : Float) (x : Float))
      (div (mul (mul 2.0 ddr) (exp (neg (mul x x))))
           (sqrt (pi))))
+
+(edef bool_to_float Float Bool)
diff --git a/test/ksc/activations.ks b/test/ksc/activations.ks
index 713ae4bb9..6239e60d2 100644
--- a/test/ksc/activations.ks
+++ b/test/ksc/activations.ks
@@ -68,4 +68,46 @@
 (gdef sufrevpass [vrelu3 (Vec Float)])
 (gdef sufrev [vrelu3 (Vec Float)])
 
+(def [vrelu3_inlined (Vec Float)] (Vec Float)
+     (t : Vec Float)
+     (map (lam (x : Float)
+               (if (lt x 0.0)
+                   0.0
+                 (if (lt x 1.0)
+                     (div (mul x (mul x x)) 3.0)
+                   (sub x (div 2.0 3.0))))) t))
+
+(def [sufrev [vrelu3_inlined (Vec Float)]] (Vec Float)
+     ((t : Vec Float) (dret : Vec Float))
+     ; TODO: should be multiplied by dret[i] - luckily we are called with dret==1.0
+     (map (lam (x : Float)
+               (if (lt x 0.0)
+                   0.0
+                 (if (lt x 1.0)
+                     (mul x x)
+                   1.0))) t))
+
+(def [vrelu3_mask (Vec Float)] (Vec Float)
+     (t : Vec Float)
+     (map (lam (x : Float)
+        (let (val0to1 (mul x (mul x (div x 3.0))))
+        (let (val1up (sub x (div 2.0 3.0)))
+        (let (in0to1 (lte x 1.0))
+           (mul (bool_to_float (gt x 0.0))
+                (add (mul (bool_to_float in0to1) val0to1)
+                     (mul (bool_to_float (not in0to1)) val1up))))))) t))
+
+(def [sufrev [vrelu3_mask (Vec Float)]] (Vec Float)
+     ((t : Vec Float) (dret : Vec Float))
+     (map2 (lam (x_ddri : Tuple Float Float)
+        (let ((x ddri) x_ddri)
+        (let (val0to1 (mul x x))
+        (let (val1up 1.0)
+        (let (in0to1 (lte x 1.0))
+           (mul (mul (bool_to_float (gt x 0.0))
+                     (add (mul (bool_to_float in0to1) val0to1)
+                          (mul (bool_to_float (not in0to1)) val1up)))
+                ddri)))))) t dret))
+
+
 (def main Integer () 0)
diff --git a/test/ksc/map.ks b/test/ksc/map.ks
index 80ce9d047..387157e15 100644
--- a/test/ksc/map.ks
+++ b/test/ksc/map.ks
@@ -43,6 +43,9 @@
 (gdef sufrevpass [map_tuple (Vec (Tuple Float Float))])
 (gdef sufrev [map_tuple (Vec (Tuple Float Float))])
 
+(def map2_example (Vec Float) ((t1 : Vec Float) (t2 : Vec Float))
+     (map2 (lam (x : Tuple Float Float) (add x)) t1 t2))
+
 (def main Integer ()
      (let (unused_input (Vec_init 1.0 2.0))
      (print (map (lam (x : Float) (mul x 2.0))