emitHalideExpr: use isl AST expression generator to print expression

The printed expression is derived from plugging in the inverse schedule.
This inverse schedule may have multiple disjuncts, especially
if the original schedule tree has disjunctive filters, such as
is the case in #200.
The expression can therefore not be assumed to consist of
a single disjunct.

Use the AST expression generator instead of manually trying
to pick the affine expression apart.
Ideally, this should use the AST build at the point
where the user statement was created, but this requires
more invasive changes.

Closes #200

Author: Sven Verdoolaege

src/core/polyhedral/codegen_cuda.cc

@@ -471,13 +471,11 @@ std::string toString(isl::aff subscript) {
 }
 
 std::string toString(isl::pw_aff subscript) {
-  isl::aff subscriptAff = isl::null<isl::aff>();
-  subscript.foreach_piece([&](isl::set domain, isl::aff aff) {
-    CHECK(!subscriptAff.get()) << "expected one piece";
-    subscriptAff = aff;
-  });
-
-  return toString(subscriptAff);
+  // 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(

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]);
     }
   }
 }
@@ -442,7 +442,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;
 }
 
@@ -461,7 +461,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[32 + t1 + 64*b0][32 + t0 + 64*b1]");
+  std::string expected("C[t1 + 64 * b0 + 32][t0 + 64 * b1 + 32]");
   ASSERT_TRUE(code.find(expected) != std::string::npos);
 }
 

test/test_mapper_memory_promotion.cc

@@ -118,13 +118,13 @@ TEST_F(Sum4D, CodeOuterBand) {
                        "__shared__ float32 _C_0[16][16][16][16];"};
 
   auto copyA =
-      "_A_0[c4][c5][c6][c7] = A[16*b0 + c4][16*b1 + c5][c2 + c6][c3 + c7];";
+      "_A_0[c4][c5][c6][c7] = A[16 * b0 + c4][16 * b1 + c5][c2 + c6][c3 + c7];";
   auto copyB =
-      "_B_0[c4][c5][c6][c7] = B[16*b0 + c4][16*b1 + c5][c2 + c6][c3 + c7];";
+      "_B_0[c4][c5][c6][c7] = B[16 * b0 + c4][16 * b1 + c5][c2 + c6][c3 + c7];";
   auto compute =
       "_C_0[c4][c5][c6][t0] = (_A_0[c4][c5][c6][t0] + _B_0[c4][c5][c6][t0]);";
   auto copyC =
-      "C[16*b0 + c4][16*b1 + c5][c2 + c6][c3 + c7] = _C_0[c4][c5][c6][c7];";
+      "C[16 * b0 + c4][16 * b1 + c5][c2 + c6][c3 + c7] = _C_0[c4][c5][c6][c7];";
   auto sync = "__syncthreads()";
 
   auto code = emitCode({256, 128, 192, 224}, {16, 16, 16, 16}, {0, 0, 0, 0});
@@ -160,13 +160,13 @@ TEST_F(Sum4D, CodeBeforeThreadMapping) {
                        "__shared__ float32 _B_0[16][16][16][1];",
                        "__shared__ float32 _C_0[16][16][16][1];"};
   auto copyA =
-      "_A_0[c4][c5][c6][0] = A[16*b0 + c4][16*b1 + c5][c2 + c6][t0 + c3];";
+      "_A_0[c4][c5][c6][0] = A[16 * b0 + c4][16 * b1 + c5][c2 + c6][t0 + c3];";
   auto copyB =
-      "_B_0[c4][c5][c6][0] = B[16*b0 + c4][16*b1 + c5][c2 + c6][t0 + c3];";
+      "_B_0[c4][c5][c6][0] = B[16 * b0 + c4][16 * b1 + c5][c2 + c6][t0 + c3];";
   auto compute =
       "_C_0[c4][c5][c6][0] = (_A_0[c4][c5][c6][0] + _B_0[c4][c5][c6][0]);";
   auto copyC =
-      "C[16*b0 + c4][16*b1 + c5][c2 + c6][t0 + c3] = _C_0[c4][c5][c6][0];";
+      "C[16 * b0 + c4][16 * b1 + c5][c2 + c6][t0 + c3] = _C_0[c4][c5][c6][0];";
   auto sync = "__syncthreads()";
 
   auto code =
@@ -204,12 +204,12 @@ TEST_F(Sum4D, CodeInnerBand) {
                        "__shared__ float32 _A_0[1][1][1][1];",
                        "__shared__ float32 _B_0[1][1][1][1];"};
   auto copyA =
-      "_A_0[0][0][0][0] = A[16*b0 + c4][16*b1 + c5][c2 + c6][t0 + c3];";
+      "_A_0[0][0][0][0] = A[16 * b0 + c4][16 * b1 + c5][c2 + c6][t0 + c3];";
   auto copyB =
-      "_B_0[0][0][0][0] = B[16*b0 + c4][16*b1 + c5][c2 + c6][t0 + c3];";
+      "_B_0[0][0][0][0] = B[16 * b0 + c4][16 * b1 + c5][c2 + c6][t0 + c3];";
   auto compute = "_C_0[0][0][0][0] = (_A_0[0][0][0][0] + _B_0[0][0][0][0]);";
   auto copyC =
-      "C[16*b0 + c4][16*b1 + c5][c2 + c6][t0 + c3] = _C_0[0][0][0][0];";
+      "C[16 * b0 + c4][16 * b1 + c5][c2 + c6][t0 + c3] = _C_0[0][0][0][0];";
   auto sync = "__syncthreads()";
 
   auto code =

test/test_tc_mapper_bugs.cc

@@ -697,6 +697,47 @@ TEST(LayerNorm, ReferenceBelongsToTwoGroups) {
   atCompl.compile("layernorm", inputs, options);
 }
 
+// This case was observed when running the autotuner on example_MLP_model
+// (#200).  It calls code generation on a schedule tree containing a
+// disjunctive filter, which results in expression with more than one disjunct
+// that was not handed properly.
+// TODO: the disjunctive filter in the schedule is unexpected and its origin
+// should be identified and explained.
+TEST(TMM_128_1024_1000, DisjunctiveFilter) {
+  at::Tensor I = at::CUDA(at::kFloat).rand({128, 1024});
+  at::Tensor W = at::CUDA(at::kFloat).rand({1000, 1024});
+  std::vector<at::Tensor> inputs = {I, W};
+  std::vector<at::Tensor> outputs;
+
+  auto TC = std::string(R"TC(
+def tmm_naive(float(B, X) I, float(Y, X) W) -> (O) {
+  O(b, y) +=! I(b, rx) * W(y, rx)
+}
+)TC");
+  auto options =
+      tc::MappingOptions::makeNaiveMappingOptions()
+          .outerScheduleFusionStrategy(tc::FusionStrategy::Preserve3Coincident)
+          .outerScheduleAllowSkewing(false)
+          .outerSchedulePositiveOrthant(true)
+          .intraTileScheduleFusionStrategy(tc::FusionStrategy::Min)
+          .intraTileScheduleAllowSkewing(false)
+          .intraTileSchedulePositiveOrthant(true)
+          .tile(1, 32, 63)
+          .mapToThreads(2, 32)
+          .mapToBlocks(64, 128, 1024)
+          .unroll(128)
+          .tileImperfectlyNested(false)
+          .useSharedMemory(false)
+          .usePrivateMemory(false)
+          .unrollCopyShared(false)
+          .matchLibraryCalls(true);
+
+  tc::ATenCompilationUnit<tc::CudaTcExecutor> atCompl;
+  atCompl.define(TC);
+  // Expecting this to compile without dying.
+  atCompl.compile("tmm_naive", inputs, options);
+}
+
 TEST(Halide2Isl, MinInUpperBound) {
   at::Tensor mat1 = at::CUDA(at::kFloat).rand({1, 100, 184, 184});
   at::Tensor mat1_pad = at::CUDA(at::kFloat).rand({1, 100, 186, 186});