emitHalideExpr: print parentheses around compound nested expressions

Despite what the original comment claims, the expression in the output
that corresponds to a variable in the input may not be an identifier
by itself.  Since the variable may appear in a Halide expression
at any position (e.g., as an argument to "-" or "*"), parentheses
need to be put around the expression if it is anything other than
an identifier or a non-negative number.

Author: Sven Verdoolaege

src/core/polyhedral/cuda/codegen.cc

@@ -495,22 +495,34 @@ 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);
       // 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 s = astBuild.expr_from(pwAff).to_C_str();
+      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) {

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);