Merge pull request #535 from facebookresearch/pr/constant

check that all accesses to registers have constant index expressions

Author: skimo-openhub

tc/core/polyhedral/cuda/codegen.cc

@@ -396,10 +396,49 @@ struct LdgWrapper {
   std::ostream& out_;
 };
 
+template <typename AFF>
+isl::ast_expr buildAccess(AFF access, const CodegenStatementContext& context) {
+  return context.build().access_from(access);
+}
+
+void emitAccess(isl::ast_expr access, const CodegenStatementContext& context) {
+  context.ss << access.to_C_str();
+}
+
 template <typename AFF>
 void emitAccess(AFF access, const CodegenStatementContext& context) {
+  emitAccess(buildAccess(access, context), context);
+}
+
+// Check that the given expression is an access with constant index expressions
+void checkConstantAccess(isl::ast_expr expr) {
+  auto op = expr.as<isl::ast_expr_op>();
+  auto access = op.as<isl::ast_op_access>();
+  TC_CHECK(access);
+  for (int i = 1; i < access.get_n_arg(); ++i) {
+    auto arg = access.get_arg(i);
+    TC_CHECK(arg.as<isl::ast_expr_int>())
+        << "expected constant subscript, got " << arg.to_C_str();
+  }
+}
+
+// Print an access to a(n array of) register(s), checking that
+// the index expressions are constant.
+void emitRegisterAccess(
+    isl::pw_multi_aff access,
+    const CodegenStatementContext& context) {
+  auto expr = buildAccess(access, context);
+  checkConstantAccess(expr);
+  emitAccess(expr, context);
+}
+
+// Print an access to global memory, wrapping the access in an "__ldg()"
+// call if the accessed tensor is known to be read-only.
+void emitGlobalAccess(
+    isl::multi_pw_aff access,
+    const CodegenStatementContext& context) {
   LdgWrapper ldgWrapper(context, access.get_tuple_id(isl::dim_type::out));
-  context.ss << context.build().access_from(access).to_C_str();
+  emitAccess(access, context);
 }
 } // namespace
 
@@ -414,9 +453,9 @@ void emitCopyStmt(const CodegenStatementContext& context) {
   if (isRead) {
     emitAccess(isl::multi_pw_aff(promoted), context);
     context.ss << " = ";
-    emitAccess(isl::multi_pw_aff(original), context);
+    emitGlobalAccess(isl::multi_pw_aff(original), context);
   } else {
-    emitAccess(isl::multi_pw_aff(original), context);
+    emitGlobalAccess(isl::multi_pw_aff(original), context);
     context.ss << " = ";
     emitAccess(isl::multi_pw_aff(promoted), context);
   }
@@ -625,7 +664,8 @@ void emitMappedTensorAccess(
     return;
   }
 
-  auto tensorId = context.scop().promotedDecl(promotionInfo.groupId).tensorId;
+  auto decl = context.scop().promotedDecl(promotionInfo.groupId);
+  auto tensorId = decl.tensorId;
 
   // Here and below in comments: D = domain, O = original tensor, P = promoted
   // tensor, S = partial schedule, A = AST loops;
@@ -651,7 +691,11 @@ void emitMappedTensorAccess(
   auto astToPromoted =
       isl::pw_multi_aff(promotion).pullback(astToScheduledOriginal);
 
-  emitAccess(astToPromoted, context);
+  if (decl.kind == Scop::PromotedDecl::Kind::Register) {
+    emitRegisterAccess(astToPromoted, context);
+  } else {
+    emitAccess(astToPromoted, context);
+  }
 }
 
 } // namespace detail

test/test_cuda_mapper_memory_promotion.cc

@@ -485,19 +485,6 @@ def fun(float(N,K) A, float(K,M) B, float(N,M) C) -> (O) {
     EXPECT_TRUE(cDeclPos == std::string::npos)
         << "tensor C promoted to register but has no reuse";
   }
-
-  void expectNoSymbolicSubscript(const std::string& code) {
-    // We don't know the exact name of the iterator, but it starts with c.
-    auto oWithIteratorPos = code.find("_O_0[c");
-    auto oWithThreadPos = code.find("_O_0[t1");
-
-    EXPECT_TRUE(oWithIteratorPos == std::string::npos)
-        << "accessing local arrays with iterators in subscripts makes "
-        << "these arrays placed in local memory instead of registers";
-    EXPECT_TRUE(oWithThreadPos == std::string::npos)
-        << "expected per-thread groups to be computed, i.e. thread "
-        << "identifiers should not appear in the subscripts";
-  }
 };
 
 TEST_F(MatMulBias, RegisterPromotion) {
@@ -562,7 +549,6 @@ TEST_F(MatMulBias, RegistersAtRoot) {
       << "expected O to be promoted to registers";
 
   expectNoABCPromotion(code);
-  expectNoSymbolicSubscript(code);
 
   auto o00Pos = code.find("_O_0[0][0]");
   auto o10Pos = code.find("_O_0[1][0]");
@@ -597,7 +583,6 @@ TEST_F(MatMulBias, RegistersAtRootNotEnoughUnroll) {
       << "not expected O to be promoted to registers";
 
   expectNoABCPromotion(code);
-  expectNoSymbolicSubscript(code);
 }
 
 TEST_F(MatMulBias, RegistersBelowFirstBand) {
@@ -621,7 +606,6 @@ TEST_F(MatMulBias, RegistersBelowFirstBand) {
   EXPECT_TRUE(oDeclPos != std::string::npos)
       << "expected O to be promoted to registers";
   expectNoABCPromotion(code);
-  expectNoSymbolicSubscript(code);
 }
 
 class Strided : public TestMapper {