Merge pull request #28 from facebookresearch/pr/let_bindings

Implement where clause variables and exists

Author: zdevito

docs/source/inference.rst

@@ -271,3 +271,20 @@ Strided indexing with dynamic stride
 The value of :code:`S(0)` is not fixed until runtime, so we can't resolve the
 size of :code:`A` or the range of the loop. This case throws a compile-time
 error. A :code:`where` clause that defines the range of :code:`i` is required.
+
+Constant fill using an exists clause
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. code::
+
+    def constant_fill(float(N) A, float c) -> B {
+      B(i) = c where exists A(i)
+    }
+
+An :code:`exists` clause allows you to add additional expressions to the range
+inference process without having the expressions affect the actual computation.
+In this example, it allows you to say that :code:`B(i)` should have the same size as
+:code:`A(i)`, but be filled with a constant value :code:`c`.  That is, you should infer the
+range of :code:`B(i)` to exist at all the places where :code:`A(i)` exists.
+It is equivalent to writing the expression :code:`true ? c : A(i)`, but with
+clearer intentions.

include/tc/core/libraries.h

@@ -140,6 +140,7 @@ constexpr auto boundsAsTemplate = R"C(
 template<typename T> inline __device__ T floord(T n, T d) {
   return n < 0 ? - (-n + d - 1)/d : n / d;
 }
+#define if_then_else(cond,a,b) (cond) ? (a) : (b);
 )C";
 } // namespace cpp
 

include/tc/lang/lexer.h

@@ -77,9 +77,17 @@ namespace lang {
   _(TK_UINT64, "uint64", "uint64")               \
   _(TK_BOOL, "bool", "bool")                     \
   _(TK_CAST, "cast", "")                         \
-  _(TK_IN, "in", "in")
+  _(TK_IN, "in", "in")                           \
+  _(TK_GE, "ge", ">=")                           \
+  _(TK_LE, "le", "<=")                           \
+  _(TK_EQ, "eq", "==")                           \
+  _(TK_NE, "neq", "!=")                          \
+  _(TK_AND, "and", "&&")                         \
+  _(TK_OR, "or", "||")                           \
+  _(TK_LET, "let", "")                           \
+  _(TK_EXISTS, "exists", "exists")
 
-static const char* valid_single_char_tokens = "+-*/()[]?:,={}>";
+static const char* valid_single_char_tokens = "+-*/()[]?:,={}><!";
 
 enum TokenKind {
   // we use characters to represent themselves so skip all valid characters
@@ -121,11 +129,14 @@ struct SharedParserData {
     // listed in increasing order of precedence
     std::vector<std::vector<int>> binary_ops = {
         {'?'},
+        {TK_OR},
+        {TK_AND},
+        {'>', '<', TK_LE, TK_GE, TK_EQ, TK_NE},
         {'+', '-'},
         {'*', '/'},
     };
     std::vector<std::vector<int>> unary_ops = {
-        {'-'},
+        {'-', '!'},
     };
 
     std::stringstream ss;
@@ -371,9 +382,20 @@ struct Lexer {
     next();
     return true;
   }
+  Token lookahead() {
+    if (!lookahead_) {
+      lookahead_.reset(new Token(lex()));
+    }
+    return *lookahead_;
+  }
   Token next() {
     auto r = cur_;
-    cur_ = lex();
+    if (lookahead_) {
+      cur_ = *lookahead_;
+      lookahead_.reset();
+    } else {
+      cur_ = lex();
+    }
     return r;
   }
   void reportError(const std::string& what, const Token& t);
@@ -407,6 +429,7 @@ struct Lexer {
   }
   size_t pos;
   Token cur_;
+  std::unique_ptr<Token> lookahead_;
   SharedParserData& shared;
 };
 } // namespace lang

include/tc/lang/parser.h

@@ -149,6 +149,24 @@ struct Parser {
     auto r = parseExp();
     return RangeConstraint::create(id->range(), id, l, r);
   }
+  TreeRef parseLetBinding() {
+    auto ident = parseIdent();
+    L.expect('=');
+    auto exp = parseExp();
+    return Let::create(ident->range(), ident, exp);
+  }
+  TreeRef parseWhereClause() {
+    auto lookahead = L.lookahead();
+    if (lookahead.kind == '=') {
+      return parseLetBinding();
+    } else if (lookahead.kind == TK_IN) {
+      return parseRangeConstraint();
+    } else {
+      L.expect(TK_EXISTS);
+      auto exp = parseExp();
+      return Exists::create(exp->range(), {exp});
+    }
+  }
   TreeRef parseParam() {
     if (L.cur().kind == TK_IDENT) {
       auto ident = parseIdent();
@@ -159,10 +177,9 @@ struct Parser {
     auto ident = parseIdent();
     return Param::create(typ->range(), ident, typ);
   }
-  TreeRef parseRangeConstraints() {
+  TreeRef parseWhereClauses() {
     if (L.nextIf(TK_WHERE)) {
-      return parseNonEmptyList(
-          ',', [&](int i) { return parseRangeConstraint(); });
+      return parseNonEmptyList(',', [&](int i) { return parseWhereClause(); });
     }
     return List::create(L.cur().range, {});
   }
@@ -200,7 +217,7 @@ struct Parser {
     auto assign = parseAssignment();
     auto rhs = parseExp();
     TreeRef equivalent_statement = parseEquivalent();
-    TreeRef range_statements = parseRangeConstraints();
+    TreeRef range_statements = parseWhereClauses();
     TreeRef empty_reduction_variables = c(TK_LIST, ident->range(), {});
     return Comprehension::create(
         ident->range(),

include/tc/lang/sema.h

@@ -194,6 +194,27 @@ struct Sema {
     }
     return e;
   }
+  void expectBool(TreeRef anchor, int token) {
+    if (token != TK_BOOL) {
+      throw ErrorReport(anchor)
+          << "expected boolean but found " << kindToString(token);
+    }
+  }
+  TreeRef expectBool(TreeRef exp) {
+    expectBool(exp, typeOfExpr(exp)->kind());
+    return exp;
+  }
+  TreeRef lookupVarOrCreateIndex(Ident ident) {
+    TreeRef type = lookup(ident, false);
+    if (!type) {
+      // variable exp is not defined, so a reduction variable is created
+      // a reduction variable index i
+      type = indexType(ident);
+      insert(index_env, ident, type, true);
+      reduction_variables.push_back(ident);
+    }
+    return type;
+  }
   TreeRef checkExp(TreeRef exp, bool allow_access) {
     switch (exp->kind()) {
       case TK_APPLY: {
@@ -205,6 +226,7 @@ struct Sema {
           throw ErrorReport(exp)
               << "tensor accesses cannot be used in this context";
         }
+
         // also handle built-in functions log, exp, etc.
         auto ident = a.name();
         if (builtin_functions.count(ident.name()) > 0) {
@@ -239,14 +261,7 @@ struct Sema {
       } break;
       case TK_IDENT: {
         auto ident = Ident(exp);
-        TreeRef type = lookup(ident, false);
-        if (!type) {
-          // variable exp is not defined, so a reduction variable is created
-          // a reduction variable index i
-          type = indexType(exp);
-          insert(index_env, ident, type, true);
-          reduction_variables.push_back(exp);
-        }
+        auto type = lookupVarOrCreateIndex(ident);
         if (type->kind() == TK_TENSOR_TYPE) {
           auto tt = TensorType(type);
           if (tt.dims().size() != 0) {
@@ -276,6 +291,35 @@ struct Sema {
             exp->map([&](TreeRef c) { return checkExp(c, allow_access); });
         return withType(nexp, matchAllTypes(nexp));
       } break;
+      case TK_EQ:
+      case TK_NE:
+      case TK_GE:
+      case TK_LE:
+      case '<':
+      case '>': {
+        auto nexp =
+            exp->map([&](TreeRef c) { return checkExp(c, allow_access); });
+        // make sure the types match but the return type
+        // is always bool
+        matchAllTypes(nexp);
+        return withType(nexp, boolType(exp));
+      } break;
+      case TK_AND:
+      case TK_OR:
+      case '!': {
+        auto nexp =
+            exp->map([&](TreeRef c) { return checkExp(c, allow_access); });
+        expectBool(exp, matchAllTypes(nexp)->kind());
+        return withType(nexp, boolType(exp));
+      } break;
+      case '?': {
+        auto nexp =
+            exp->map([&](TreeRef c) { return checkExp(c, allow_access); });
+        expectBool(nexp->tree(0));
+        auto rtype =
+            match_types(typeOfExpr(nexp->tree(1)), typeOfExpr(nexp->tree(2)));
+        return withType(nexp, rtype);
+      }
       case TK_CONST: {
         auto c = Const(exp);
         return withType(exp, c.type());
@@ -322,7 +366,10 @@ struct Sema {
   TreeRef floatType(TreeRef anchor) {
     return c(TK_FLOAT, anchor->range(), {});
   }
-  void checkDim(const Ident& dim) {
+  TreeRef boolType(TreeRef anchor) {
+    return c(TK_BOOL, anchor->range(), {});
+  }
+  void checkDim(Ident dim) {
     insert(env, dim, dimType(dim), false);
   }
   TreeRef checkTensorType(TreeRef type) {
@@ -354,6 +401,33 @@ struct Sema {
     }
     return List::create(list->range(), std::move(r));
   }
+  TreeRef checkRangeConstraint(RangeConstraint rc) {
+    // RCs are checked _before_ the rhs of the TC, so
+    // it is possible the index is not in the environment yet
+    // calling lookupOrCreate ensures it exists
+    lookupVarOrCreateIndex(rc.ident());
+    // calling looking directly in the index_env ensures that
+    // we are actually constraining an index and not some other variable
+    lookup(index_env, rc.ident(), true);
+    auto s = expectIntegral(checkExp(rc.start(), false));
+    auto e = expectIntegral(checkExp(rc.end(), false));
+    return RangeConstraint::create(rc.range(), rc.ident(), s, e);
+  }
+  TreeRef checkLet(Let l) {
+    auto rhs = checkExp(l.rhs(), true);
+    insert(let_env, l.name(), typeOfExpr(rhs), true);
+    return Let::create(l.range(), l.name(), rhs);
+  }
+  TreeRef checkWhereClause(TreeRef ref) {
+    if (ref->kind() == TK_LET) {
+      return checkLet(Let(ref));
+    } else if (ref->kind() == TK_EXISTS) {
+      auto exp = checkExp(Exists(ref).exp(), true);
+      return Exists::create(ref->range(), exp);
+    } else {
+      return checkRangeConstraint(RangeConstraint(ref));
+    }
+  }
   TreeRef checkStmt(TreeRef stmt_) {
     auto stmt = Comprehension(stmt_);
 
@@ -374,6 +448,11 @@ struct Sema {
       output_indices.push_back(new_var);
     }
 
+    // where clauses are checked _before_ the rhs because they
+    // introduce let bindings that are in scope for the rhs
+    auto where_clauses_ = stmt.whereClauses().map(
+        [&](TreeRef rc) { return checkWhereClause(rc); });
+
     TreeRef rhs_ = checkExp(stmt.rhs(), true);
     TreeRef scalar_type = typeOfExpr(rhs_);
 
@@ -408,20 +487,11 @@ struct Sema {
     // if we redefined an input, it is no longer valid for range expressions
     live_input_names.erase(stmt.ident().name());
 
-    auto range_constraints =
-        stmt.rangeConstraints().map([&](const RangeConstraint& rc) {
-          lookup(index_env, rc.ident(), true);
-          auto s = expectIntegral(checkExp(rc.start(), false));
-          auto e = expectIntegral(checkExp(rc.end(), false));
-          return RangeConstraint::create(rc.range(), rc.ident(), s, e);
-        });
-
-    auto equivalent_statement_ =
-        stmt.equivalent().map([&](const Equivalent& eq) {
-          auto indices_ = eq.accesses().map(
-              [&](TreeRef index) { return checkExp(index, true); });
-          return Equivalent::create(eq.range(), eq.name(), indices_);
-        });
+    auto equivalent_statement_ = stmt.equivalent().map([&](Equivalent eq) {
+      auto indices_ = eq.accesses().map(
+          [&](TreeRef index) { return checkExp(index, true); });
+      return Equivalent::create(eq.range(), eq.name(), indices_);
+    });
 
     TreeRef assignment = stmt.assignment();
     // For semantic consistency we allow overwriting reductions like +=!
@@ -446,13 +516,16 @@ struct Sema {
         stmt.indices(),
         stmt.assignment(),
         rhs_,
-        range_constraints,
+        where_clauses_,
         equivalent_statement_,
         reduction_variable_list);
+    // clear the per-statement environments to get ready for the next statement
     index_env.clear();
+    let_env.clear();
+
     return result;
   }
-  bool isNotInplace(const TreeRef& assignment) {
+  bool isNotInplace(TreeRef assignment) {
     switch (assignment->kind()) {
       case TK_PLUS_EQ_B:
       case TK_TIMES_EQ_B:
@@ -493,13 +566,15 @@ struct Sema {
       throw ErrorReport(ident) << name << " already defined";
     }
   }
-  TreeRef lookup(const Ident& ident, bool required) {
+  TreeRef lookup(Ident ident, bool required) {
     TreeRef v = lookup(index_env, ident, false);
+    if (!v)
+      v = lookup(let_env, ident, false);
     if (!v)
       v = lookup(env, ident, required);
     return v;
   }
-  TreeRef lookup(Env& the_env, const Ident& ident, bool required) {
+  TreeRef lookup(Env& the_env, Ident ident, bool required) {
     std::string name = ident.name();
     auto it = the_env.find(name);
     if (required && it == the_env.end()) {
@@ -517,6 +592,7 @@ struct Sema {
 
   std::vector<TreeRef> reduction_variables; // per-statement
   Env index_env; // per-statement
+  Env let_env; // per-statement, used for where i = <exp>
 
   Env env; // name -> type
   Env annotated_output_types; // name -> type, for all annotated returns types