Merge pull request #226 from facebookresearch/pr/fix-179

emitHalideExpr: print parentheses around compound nested expressions

Author: skimo-openhub

include/tc/core/polyhedral/cuda/codegen.h

@@ -31,12 +31,6 @@ struct CodegenStatementContext;
 
 namespace detail {
 
-void emitDirectSubscripts(
-    isl::pw_multi_aff subscripts,
-    const CodegenStatementContext& context);
-
-std::string toString(isl::pw_aff subscript);
-
 isl::pw_aff makeAffFromMappedExpr(
     const Halide::Expr& expr,
     const CodegenStatementContext& context);

src/core/polyhedral/cuda/codegen.cc

@@ -359,28 +359,33 @@ void emitReductionInit(
   context.ss << ";" << endl;
 }
 
-void emitCopyStmt(const CodegenStatementContext& context) {
-  using detail::emitDirectSubscripts;
+namespace {
+template <typename AFF>
+void emitAccess(AFF access, const CodegenStatementContext& context) {
+  // Use a temporary isl::ast_build to print the expression.
+  // Ideally, this should use the build at the point
+  // where the user statement was created.
+  auto astBuild = isl::ast_build::from_context(access.domain());
+  context.ss << astBuild.access_from(access).to_C_str();
+}
+} // namespace
 
+void emitCopyStmt(const CodegenStatementContext& context) {
   auto stmtId = context.statementId();
 
   auto iteratorMap = context.iteratorMap();
   auto promoted = iteratorMap.range_factor_range();
   auto original = iteratorMap.range_factor_domain().range_factor_range();
   auto isRead = stmtId.get_name() == kReadIdName;
-  auto originalName = original.get_tuple_id(isl::dim_type::out).get_name();
-  auto promotedName = promoted.get_tuple_id(isl::dim_type::out).get_name();
 
   if (isRead) {
-    context.ss << promotedName;
-    emitDirectSubscripts(promoted, context);
-    context.ss << " = " << originalName;
-    emitDirectSubscripts(original, context);
+    emitAccess(isl::multi_pw_aff(promoted), context);
+    context.ss << " = ";
+    emitAccess(isl::multi_pw_aff(original), context);
   } else {
-    context.ss << originalName;
-    emitDirectSubscripts(original, context);
-    context.ss << " = " << promotedName;
-    emitDirectSubscripts(promoted, context);
+    emitAccess(isl::multi_pw_aff(original), context);
+    context.ss << " = ";
+    emitAccess(isl::multi_pw_aff(promoted), context);
   }
   context.ss << ";" << std::endl;
 }
@@ -447,14 +452,6 @@ void AstPrinter::emitAst(isl::ast_node node) {
 
 namespace detail {
 
-std::string toString(isl::pw_aff subscript) {
-  // Use a temporary isl::ast_build to print the expression.
-  // Ideally, this should use the build at the point
-  // where the user statement was created.
-  auto astBuild = isl::ast_build::from_context(subscript.domain());
-  return astBuild.expr_from(subscript).to_C_str();
-}
-
 isl::pw_aff makeAffFromMappedExpr(
     const Halide::Expr& expr,
     const CodegenStatementContext& context) {
@@ -498,18 +495,35 @@ isl::multi_aff makeMultiAffAccess(
   return ma;
 }
 
+namespace {
+bool is_identifier_or_nonnegative_integer(isl::ast_expr expr) {
+  if (isl_ast_expr_get_type(expr.get()) == isl_ast_expr_id)
+    return true;
+  if (isl_ast_expr_get_type(expr.get()) != isl_ast_expr_int)
+    return false;
+  return isl::manage(isl_ast_expr_get_val(expr.get())).is_nonneg();
+}
+} // namespace
+
 void emitHalideExpr(
     const Halide::Expr& e,
     const CodegenStatementContext& context,
     const map<string, string>& substitutions) {
   class EmitHalide : public Halide::Internal::IRPrinter {
     using Halide::Internal::IRPrinter::visit;
     void visit(const Halide::Internal::Variable* op) {
-      // This is probably needlessly indirect, given that we just have
-      // a name to look up somewhere.
       auto pwAff = tc::polyhedral::detail::makeAffFromMappedExpr(
           Halide::Expr(op), context);
-      context.ss << tc::polyhedral::detail::toString(pwAff);
+      // Use a temporary isl::ast_build to print the expression.
+      // Ideally, this should use the build at the point
+      // where the user statement was created.
+      auto astBuild = isl::ast_build::from_context(pwAff.domain());
+      auto expr = astBuild.expr_from(pwAff);
+      auto s = expr.to_C_str();
+      if (!is_identifier_or_nonnegative_integer(expr)) {
+        s = "(" + s + ")";
+      }
+      context.ss << s;
     }
     void visit(const Halide::Internal::Call* op) {
       if (substitutions.count(op->name)) {
@@ -613,19 +627,7 @@ void emitMappedTensorAccess(
   auto astToPromoted =
       isl::pw_multi_aff(promotion).pullback(astToScheduledOriginal);
 
-  auto astBuild = isl::ast_build::from_context(astToPromoted.domain());
-  context.ss << astBuild.access_from(astToPromoted).to_C_str();
-}
-
-void emitDirectSubscripts(
-    isl::pw_multi_aff subscripts,
-    const CodegenStatementContext& context) {
-  auto mpa = isl::multi_pw_aff(subscripts); // this conversion is safe
-  for (auto pa : isl::MPA(mpa)) {
-    context.ss << "[";
-    context.ss << toString(pa.pa);
-    context.ss << "]";
-  }
+  emitAccess(astToPromoted, context);
 }
 
 } // namespace detail

test/test_mapper.cc

@@ -186,7 +186,7 @@ def fun(float(N, M) A, float(N, M) B) -> (C) {
   float32 (*B)[M] = reinterpret_cast<float32 (*)[M]>(pB);
   for (int c1 = 16 * b1; c1 < M; c1 += 4096) {
     if (M >= t1 + c1 + 1) {
-      C[t0 + 16 * b0][t1 + c1] = (A[t0 + 16 * b0][t1 + c1] + B[t0 + 16 * b0][t1 + c1]);
+      C[(t0 + 16 * b0)][(t1 + c1)] = (A[(t0 + 16 * b0)][(t1 + c1)] + B[(t0 + 16 * b0)][(t1 + c1)]);
     }
   }
 }
@@ -318,9 +318,9 @@ constexpr auto kExpectedMatmul_64_64_64 =
     for (int c1 = 0; c1 <= 63; c1 += 16) {
       for (int c2 = t1; c2 <= 15; c2 += 8) {
         for (int c3 = 0; c3 <= 15; c3 += 1) {
-          O[c0 + c2][c1 + c3] = 0.000000f;
+          O[(c0 + c2)][(c1 + c3)] = 0.000000f;
           for (int c4 = t0; c4 <= 63; c4 += 32) {
-            O[c0 + c2][c1 + c3] = (O[c0 + c2][c1 + c3] + (A[c0 + c2][c4]*B[c4][c1 + c3]));
+            O[(c0 + c2)][(c1 + c3)] = (O[(c0 + c2)][(c1 + c3)] + (A[(c0 + c2)][c4]*B[c4][(c1 + c3)]));
           }
         }
       }
@@ -443,7 +443,7 @@ TEST_F(PolyhedralMapperTest, Unroll1D) {
   auto mscop = MappedScop::makeWithOuterBlockInnerThreadStrategy(
       std::move(scop), mappingOptions);
   auto code = std::get<0>(mscop->codegen(specializedName));
-  std::string expected("C[64 * b0 + c2][t0 + 64 * b1]");
+  std::string expected("C[(64 * b0 + c2)][(t0 + 64 * b1)]");
   ASSERT_TRUE(code.find(expected) != std::string::npos) << code;
 }
 
@@ -462,7 +462,7 @@ TEST_F(PolyhedralMapperTest, Unroll2D) {
   auto mscop = MappedScop::makeWithOuterBlockInnerThreadStrategy(
       std::move(scop), mappingOptions);
   auto code = std::get<0>(mscop->codegen(specializedName));
-  std::string expected("C[t1 + 64 * b0 + 32][t0 + 64 * b1 + 32]");
+  std::string expected("C[(t1 + 64 * b0 + 32)][(t0 + 64 * b1 + 32)]");
   ASSERT_TRUE(code.find(expected) != std::string::npos);
 }
 
@@ -712,7 +712,7 @@ TEST_F(PolyhedralMapperTest, ReductionMM1D) {
   auto code = codegenMapped(kTcMM, mappingOptions);
   using tc::code::cuda::kCUBReductionName;
   EXPECT_TRUE(code.find(kCUBReductionName) != std::string::npos);
-  EXPECT_TRUE(code.find("C[c0 + c3][t0 + c1] = (C") != std::string::npos);
+  EXPECT_TRUE(code.find("C[(c0 + c3)][(t0 + c1)] = (C") != std::string::npos);
 }
 
 /*
@@ -730,7 +730,7 @@ TEST_F(PolyhedralMapperTest, ReductionMM2D) {
   auto code = codegenMapped(kTcMM, mappingOptions);
   using tc::code::cuda::kCUBReductionName;
   EXPECT_TRUE(code.find(kCUBReductionName) != std::string::npos);
-  EXPECT_TRUE(code.find("C[t1 + c0][t0 + c1] = (C") != std::string::npos);
+  EXPECT_TRUE(code.find("C[(t1 + c0)][(t0 + c1)] = (C") != std::string::npos);
 }
 
 int main(int argc, char** argv) {

test/test_tc_mapper_bugs.cc

@@ -761,6 +761,36 @@ TEST(Halide2Isl, MinInUpperBound) {
   atCompl.compile("graph2", inputs, options);
 }
 
+// Check that nested expressions are properly formatted.
+// In particular, as soon as the tensor size X is larger than the tile size,
+// the expression for "xp" is a sum of multiple loop iterators
+// in the generated code.  Parentheses need to be placed around
+// these expressions to ensure the end result is, say, "-(c1 + c3)"
+// rather than "-c1 + c3".
+// The actual convolution is one where the output is equal to the input.
+TEST(Convolution, NestedExpressions) {
+  auto convolution = "convolution";
+  auto TC = std::string(R"TC(
+  def convolution(float(X) A, float(Xp) K) -> (B) {
+      B(x) +=! A(xp) * K(X - 1 + x - xp) where xp in 0:X
+  }
+  )TC");
+  int X = 33;
+  at::Tensor A = at::CUDA(at::kFloat).zeros({X});
+  at::Tensor K = at::CUDA(at::kFloat).zeros({2 * X - 1});
+  A[10] = 1;
+  K[X - 1] = 1;
+  std::vector<at::Tensor> inputs = {A, K};
+  std::vector<at::Tensor> outputs;
+  tc::ATenCompilationUnit<tc::CudaTcExecutor> cu;
+  cu.define(TC);
+  auto options = tc::CudaMappingOptions::makeNaiveCudaMappingOptions();
+  auto handle = cu.compile(convolution, inputs, options);
+  cu.run(convolution, inputs, outputs, handle);
+  auto B = outputs[0];
+  CHECK_EQ(at::Scalar(B[10]).toFloat(), 1);
+}
+
 int main(int argc, char** argv) {
   ::testing::InitGoogleTest(&argc, argv);
   ::gflags::ParseCommandLineFlags(&argc, &argv, true);