Merge pull request #270 from facebookresearch/pr/minmax

Parser rules for min/max

Author: ftynse

include/tc/lang/parser.h

@@ -55,6 +55,19 @@ struct Parser {
       auto value = parseExp();
       L.expect(')');
       return Cast::create(type->range(), value, type);
+    } else if (L.cur().kind == TK_MIN || L.cur().kind == TK_MAX) {
+      // min/max are treated as operators later in the compilation pipeline.
+      // so we ensure they have precisely two arguments here so they can
+      // use the same pathways as other operators like + where argument
+      // count is ensured by parsing
+      auto range = L.cur().range;
+      auto tok = L.next().kind;
+      L.expect('(');
+      auto a = parseExp();
+      L.expect(',');
+      auto b = parseExp();
+      L.expect(')');
+      prefix = c(tok, range, {a, b});
     } else {
       prefix = parseIdent();
       auto range = L.cur().range;

test/cuda/test_corner_cases.cc

@@ -243,6 +243,48 @@ TEST(TestCornerCases, E20) {
       {F(1)});
 }
 
+TEST(TestCornerCases, E21) {
+  auto a = F(1);
+  auto b = F(1);
+  auto c = F(1);
+  Succeed(
+      "def f(float(1) a, float(1) b) -> (c) { c(i) = max(a(i), b(i)) }",
+      {a, b},
+      {c});
+  CHECK_EQ(
+      fmaxf(at::Scalar(a[0]).toFloat(), at::Scalar(b[0]).toFloat()),
+      at::Scalar(c[0]).toFloat());
+}
+
+TEST(TestCornerCases, E22) {
+  auto a = F(1);
+  auto b = F(1);
+  auto c = F(1);
+  Succeed(
+      "def f(float(1) a, float(1) b) -> (c) { c(i) = min(a(i), b(i)) }",
+      {a, b},
+      {c});
+  CHECK_EQ(
+      fminf(at::Scalar(a[0]).toFloat(), at::Scalar(b[0]).toFloat()),
+      at::Scalar(c[0]).toFloat());
+}
+
+TEST(TestCornerCases, E23) {
+  auto a = F(1);
+  auto b = F(1);
+  auto c = F(1);
+  auto d = F(1);
+  Succeed(
+      "def f(float(1) a, float(1) b, float(1) c) -> (d) { d(i) = min(a(i), max(b(i), c(i))) }",
+      {a, b, c},
+      {d});
+  CHECK_EQ(
+      fminf(
+          at::Scalar(a[0]).toFloat(),
+          fmaxf(at::Scalar(b[0]).toFloat(), at::Scalar(c[0]).toFloat())),
+      at::Scalar(d[0]).toFloat());
+}
+
 int main(int argc, char** argv) {
   ::testing::InitGoogleTest(&argc, argv);
   ::gflags::ParseCommandLineFlags(&argc, &argv, true);