Add bitwise operators and propagate them from parser -> halide -> isl

Author: prigoyal

tc/core/libraries.h

@@ -145,6 +145,12 @@ 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))
+#define shift_left(a,b) ((a) << (b))
+#define shift_right(a,b) ((a) >> (b))
+#define bitwise_and(a,b) ((a) & (b))
+#define bitwise_xor(a,b) ((a) ^ (b))
+#define bitwise_or(a,b) ((a) | (b))
+#define bitwise_not(a) (~(a))
 )C";
 } // namespace cpp
 

tc/core/tc2halide.cc

@@ -190,6 +190,18 @@ Expr translateExpr(
       return t(0) && t(1);
     case lang::TK_OR:
       return t(0) || t(1);
+    case lang::TK_LS:
+      return t(0) << t(1);
+    case lang::TK_RS:
+      return t(0) >> t(1);
+    case '|':
+      return t(0) | t(1);
+    case '^':
+      return t(0) ^ t(1);
+    case '&':
+      return t(0) & t(1);
+    case '~':
+      return ~t(0);
     case lang::TK_BUILT_IN: {
       auto b = lang::BuiltIn(expr);
       vector<Expr> exprs;

tc/lang/lexer.h

@@ -84,10 +84,12 @@ namespace lang {
   _(TK_NE, "neq", "!=")                          \
   _(TK_AND, "and", "&&")                         \
   _(TK_OR, "or", "||")                           \
+  _(TK_LS, "ls", "<<")                           \
+  _(TK_RS, "rs", ">>")                           \
   _(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
@@ -135,12 +137,16 @@ struct SharedParserData {
         {'?'},
         {TK_OR},
         {TK_AND},
+        {'|'},
+        {'^'},
+        {'&'},
         {'>', '<', TK_LE, TK_GE, TK_EQ, TK_NE},
+        {TK_LS, TK_RS},
         {'+', '-'},
         {'*', '/', '%'},
     };
     std::vector<std::vector<int>> unary_ops = {
-        {'-', '!'},
+        {'-', '!', '~'},
     };
 
     std::stringstream ss;

tc/lang/sema.h

@@ -294,6 +294,12 @@ struct Sema {
       case '*':
       case '/':
       case '%':
+      case '~':
+      case '|':
+      case '^':
+      case '&':
+      case TK_LS:
+      case TK_RS:
       case TK_MIN:
       case TK_MAX: {
         auto nexp =

tc/lang/test_expected/bitwise.expected

@@ -0,0 +1,11 @@
+(|
+  (^
+    (&
+      (~ (const 3 (int32)))
+      (const 4 (int32)))
+    (const 5 (int32)))
+  (ls
+    (rs
+      (const 6 (int32))
+      (const 8 (int32)))
+    (const 2 (int32))))

test/cuda/test_corner_cases.cc

@@ -296,6 +296,33 @@ TEST(TestCornerCases, E24) {
   CHECK_EQ(at::Scalar(r[0]).toInt(), 0);
 }
 
+TEST(TestCornerCases, E25){
+#define GEN_BITWISE(op)                                        \
+  {                                                            \
+    auto a = 2 * I();                                          \
+    auto b = 2 * I();                                          \
+    auto r = I(0);                                             \
+    Succeed(                                                   \
+        "def f(int32 a, int32 b) -> (c) { c(i) = int32(a " #op \
+        " b) where i in 0:1 }",                                \
+        {a, b},                                                \
+        {r});                                                  \
+    auto e = at::Scalar(a).toInt() op at::Scalar(b).toInt();   \
+    CHECK_EQ(e, at::Scalar(r[0]).toInt());                     \
+  }
+
+    GEN_BITWISE(<<) GEN_BITWISE(>>) GEN_BITWISE(&) GEN_BITWISE(|)
+        GEN_BITWISE (^)}
+
+TEST(TestCornerCases, E26) {
+  auto a = I();
+  auto r = I(0);
+  Succeed(
+      "def f(int32 a) -> (c) { c(i) = int32(~a) where i in 0:1 }", {a}, {r});
+  auto e = ~at::Scalar(a).toInt();
+  CHECK_EQ(at::Scalar(r[0]).toInt(), e);
+}
+
 int main(int argc, char** argv) {
   ::testing::InitGoogleTest(&argc, argv);
   ::gflags::ParseCommandLineFlags(&argc, &argv, true);

test/test_lang.cc

@@ -244,6 +244,14 @@ int main(int argc, char** argv) {
     ss2 << p2.parseExp();
     assertEqual("function.expected", ss2.str());
   }
+  {
+    std::string bitOps = "~3&4^5|6>>8<<2";
+    Parser p(bitOps);
+    auto r = p.parseExp();
+    std::stringstream ss;
+    ss << r;
+    assertEqual("bitwise.expected", ss.str());
+  }
   assertParseEqual("trinary.expected", "a ? 3 : b ? 3 : 4", [&](Parser& p) {
     return p.parseExp();
   });