Merge pull request #115 from facebookresearch/emit-inifnities

[Hotfix] emit macros for inf/inff

Author: nicolasvasilache

docs/source/framework/pytorch_integration/layers_database.rst

@@ -139,7 +139,7 @@ Softmax
 .. code::
 
     def softmax(float(N, D) I) -> (O, maxVal, expDistance, expSum) {
-        maxVal(n) max= I(n, d)
+        maxVal(n) max=! I(n, d)
         expDistance(n, d) = exp(I(n, d) - maxVal(n))
         expSum(n) +=! expDistance(n, d)
         O(n, d) = expDistance(n, d) / expSum(n)

include/tc/core/libraries.h

@@ -36,7 +36,11 @@ typedef int int32;
 typedef long int64;
 typedef float float32;
 typedef double float64;
+)C";
 
+constexpr auto defines = R"C(
+#define inff __int_as_float(0x7f800000)
+#define inf __longlong_as_double(0x7ff0000000000000LL)
 )C";
 
 constexpr auto mathFunctionDecl = R"C(

src/core/polyhedral/mapped_scop.cc

@@ -564,7 +564,8 @@ std::tuple<std::string, tc::Grid, tc::Block> MappedScop::codegen(
   auto mappedScopForCodegen = makeSpecializedMappedScop(*this);
 
   std::stringstream code;
-  code << code::cpp::boundsAsTemplate << code::c::types << std::endl;
+  code << code::cpp::boundsAsTemplate << code::c::types << code::c::defines
+       << std::endl;
   if (mappedScopForCodegen->scop().treeSyncUpdateMap.size() != 0) {
     code << code::cuda::common;
     code << code::cuda::cubBlockReduce;

test/test_execution_engine.cc

@@ -87,7 +87,7 @@ TEST_F(ATenCompilationUnitTest, DISABLED_SoftmaxB) {
       outputs);
 }
 
-TEST_F(ATenCompilationUnitTest, DISABLED_SoftmaxC) {
+TEST_F(ATenCompilationUnitTest, SoftmaxC) {
   at::Tensor a = at::CUDA(at::kFloat).rand({32, 16});
   std::vector<at::Tensor> inputs = {a};
   std::vector<at::Tensor> outputs;
@@ -114,7 +114,7 @@ TEST_F(ATenCompilationUnitTest, SoftmaxD) {
   Check(
       R"(
       def softmax(float(N, D) I) -> (O, maxVal, expDistance, expSum) {
-        maxVal(n) max= I(n, d)
+        maxVal(n) max=! I(n, d)
         expDistance(n, d) = exp(I(n, d) - maxVal(n))
         expSum(n) +=! expDistance(n, d)
         O(n, d) = expDistance(n, d) / expSum(n)

test_python/layers/test_softmax.py

@@ -27,7 +27,7 @@ class TestSoftmax(unittest.TestCase):
     def test_softmax(self):
         LANG = """
         def softmax(float(N, D) I) -> (O, maxVal, expDistance, expSum) {
-          maxVal(n) max= I(n, d)
+          maxVal(n) max=! I(n, d)
           expDistance(n, d) = exp(I(n, d) - maxVal(n))
           expSum(n) +=! expDistance(n, d)
           O(n, d) = expDistance(n, d) / expSum(n)