Add more type support

This commit adds missing type support to the lexer, semantic analyzer,
and cuda implementation.
This also adds language and end-to-end functional tests for all types we support.

An annoying type issue comes from the inability to include system dependencies
with NVRTC so half support is explicitly disabled.
Maybe it is time to think about moving to NVCC in a separate process, which
we have been discussing for some time.

The following commit adds an issue submitted by @mdouze which now runs
properly thanks to this commit.

Author: nicolasvasilache

tc/core/libraries.h

@@ -31,9 +31,20 @@ namespace code {
 namespace c {
 
 constexpr auto types = R"C(
+// Can't include system dependencies with NVRTC
+// Can't include cuda_fp16.h with NVRTC due to transitive system dependencies
+// #include <cuda_fp16.h>
+
 // Halide type handling
+typedef char int8;
+typedef short int16;
 typedef int int32;
 typedef long int64;
+typedef unsigned char uint8;
+typedef unsigned short uint16;
+typedef unsigned int uint32;
+typedef unsigned long uint64;
+// typedef half float16;
 typedef float float32;
 typedef double float64;
 )C";

tc/core/tc2halide.cc

@@ -53,10 +53,17 @@ Type translateScalarType(int tcType) {
       return Int(32);
     case lang::TK_INT64:
       return Int(64);
+    case lang::TK_FLOAT16:
+      return Float(16);
+    case lang::TK_FLOAT32:
+      return Float(32);
+    case lang::TK_FLOAT64:
+      return Float(64);
     case lang::TK_FLOAT:
       return Float(32);
     case lang::TK_DOUBLE:
       return Float(64);
+
     default:
       LOG(FATAL) << "Unhandled TC scalar type: " << tcType << '\n';
       return Type();

tc/lang/lexer.h

@@ -41,8 +41,6 @@ namespace lang {
   _(TK_MIN, "min", "min")                        \
   _(TK_MAX, "max", "max")                        \
   _(TK_WHERE, "where", "where")                  \
-  _(TK_FLOAT, "float", "float")                  \
-  _(TK_DOUBLE, "double", "double")               \
   _(TK_DEF, "def", "def")                        \
   _(TK_ARROW, "arrow", "->")                     \
   _(TK_EQUIVALENT, "equivalent", "<=>")          \
@@ -67,15 +65,21 @@ namespace lang {
   _(TK_TIMES_EQ_B, "times_eq_b", "*=!")          \
   _(TK_MIN_EQ_B, "min_eq_b", "min=!")            \
   _(TK_MAX_EQ_B, "max_eq_b", "max=!")            \
-  _(TK_INT8, "int8", "int8")                     \
-  _(TK_INT16, "int16", "int16")                  \
-  _(TK_INT32, "int32", "int32")                  \
-  _(TK_INT64, "int64", "int64")                  \
+                                                 \
+  _(TK_BOOL, "bool", "bool")                     \
   _(TK_UINT8, "uint8", "uint8")                  \
   _(TK_UINT16, "uint16", "uint16")               \
   _(TK_UINT32, "uint32", "uint32")               \
   _(TK_UINT64, "uint64", "uint64")               \
-  _(TK_BOOL, "bool", "bool")                     \
+  _(TK_INT8, "int8", "int8")                     \
+  _(TK_INT16, "int16", "int16")                  \
+  _(TK_INT32, "int32", "int32")                  \
+  _(TK_INT64, "int64", "int64")                  \
+  _(TK_FLOAT16, "float16", "float16")            \
+  _(TK_FLOAT32, "float32", "float32")            \
+  _(TK_FLOAT64, "float64", "float64")            \
+  _(TK_FLOAT, "float", "float")                  \
+  _(TK_DOUBLE, "double", "double")               \
   _(TK_CAST, "cast", "")                         \
   _(TK_IN, "in", "in")                           \
   _(TK_GE, "ge", ">=")                           \
@@ -271,15 +275,18 @@ struct SharedParserData {
   }
   bool isScalarType(int kind) {
     switch (kind) {
-      case TK_INT8:
-      case TK_INT16:
-      case TK_INT32:
-      case TK_INT64:
+      case TK_BOOL:
       case TK_UINT8:
       case TK_UINT16:
       case TK_UINT32:
       case TK_UINT64:
-      case TK_BOOL:
+      case TK_INT8:
+      case TK_INT16:
+      case TK_INT32:
+      case TK_INT64:
+      case TK_FLOAT16:
+      case TK_FLOAT32:
+      case TK_FLOAT64:
       case TK_FLOAT:
       case TK_DOUBLE:
         return true;

tc/lang/sema.h

@@ -46,13 +46,23 @@ struct TypeInfo {
       TYPE_INFO_OPTION(TK_INT16, Int, 16)
       TYPE_INFO_OPTION(TK_INT32, Int, 32)
       TYPE_INFO_OPTION(TK_INT64, Int, 64)
+      TYPE_INFO_OPTION(TK_FLOAT16, Float, 16)
+      TYPE_INFO_OPTION(TK_FLOAT32, Float, 32)
+      TYPE_INFO_OPTION(TK_FLOAT64, Float, 64)
       TYPE_INFO_OPTION(TK_FLOAT, Float, 32)
       TYPE_INFO_OPTION(TK_DOUBLE, Float, 64)
+
 #undef TYPE_INFO_OPTION
       default:
         throw ErrorReport(scalar_type)
             << "Unhandled TC scalar type: " << scalar_type;
     }
+
+    if (code_ == Code::Float && bits_ == 16) {
+      throw ErrorReport(scalar_type)
+          << "Half precision floating point not supported "
+          << "until we can make NVRTC include system headers";
+    }
   }
   int toScalarToken() const {
     switch (code()) {
@@ -82,12 +92,15 @@ struct TypeInfo {
         }
       case Float:
         switch (bits()) {
+          case 16:
+            return TK_FLOAT16;
           case 32:
             return TK_FLOAT;
           case 64:
             return TK_DOUBLE;
         }
     }
+
     throw std::runtime_error("Unknown type info?");
   }
   Code code() const {

test/cuda/test_compile_and_run.cc

@@ -275,6 +275,37 @@ def cast(float(M,N) A, int32 four) -> (int32(M,N) output) {
   TC_CHECK_EQ(r, 0);
 }
 
+TEST_F(CompilationTest, Types) {
+  struct TypeMatch {
+    std::string s;
+    at::ScalarType a;
+  };
+  for (auto type :
+       {// TypeMatch{"bool", at::ScalarType::Bool},    // no aten version
+        TypeMatch{"uint8", at::ScalarType::Byte},
+        // TypeMatch{"uint16", at::ScalarType::Short}, // no aten version
+        // TypeMatch{"uint32", at::ScalarType::Int},   // no aten version
+        // TypeMatch{"uint64", at::ScalarType::Long},  // no aten version
+        TypeMatch{"int8", at::ScalarType::Char},
+        TypeMatch{"int16", at::ScalarType::Short},
+        TypeMatch{"int32", at::ScalarType::Int},
+        TypeMatch{"int64", at::ScalarType::Long},
+        // NVRTC include transitive dependencies issue
+        // TypeMatch{"float16", at::ScalarType::Half},
+        TypeMatch{"float32", at::ScalarType::Float},
+        TypeMatch{"float64", at::ScalarType::Double},
+        TypeMatch{"float", at::ScalarType::Float},
+        TypeMatch{"double", at::ScalarType::Double}}) {
+    std::string tc = std::string("def test_type(") + std::string(type.s) +
+        std::string("(N) A) -> (B) { B(k) +=! A(i) where k in 0:1 }");
+    std::vector<at::Tensor> outputs = Check(
+        tc,
+        "test_type",
+        tc::CudaMappingOptions::makeNaiveMappingOptions(),
+        {at::CUDA(type.a).ones({100})});
+  }
+}
+
 int main(int argc, char** argv) {
   ::testing::InitGoogleTest(&argc, argv);
   ::gflags::ParseCommandLineFlags(&argc, &argv, true);

test/test_tc2halide.cc

@@ -197,6 +197,29 @@ def foo(float(N) A) -> (B) {
 )TC";
   EXPECT_THROW(Check(tc), ::lang::ErrorReport);
 }
+
+TEST_F(TC2Isl, Types) {
+  for (auto type : {"bool",
+                    "uint8",
+                    "uint16",
+                    "uint32",
+                    "uint64",
+                    "int8",
+                    "int16",
+                    "int32",
+                    "int64",
+                    // NVRTC include transitive dependencies issue
+                    // "float16",
+                    "float32",
+                    "float64",
+                    "float",
+                    "double"}) {
+    string tc = string("def test_type(") + string(type) +
+        string("(N) A) -> (B) { B(k) +=! A(i) where k in 0:1 }");
+    Check(tc);
+  }
+}
+
 int main(int argc, char** argv) {
   ::testing::InitGoogleTest(&argc, argv);
   ::gflags::ParseCommandLineFlags(&argc, &argv, true);