Add support for __ldg

CUDA supports a faster read-only cache for data that is never written.
This commits adds support in cuda/codegen to generate such accesses.
This is achieved by simply inspecting the schedule tree for all
read-only tensor references and saving their isl::id in a set that is
passed in the context.

The corresponding CudaMappingOptions are also added to control triggering of
the option.

A simple test is also added.

Author: ftynse

docs/source/mapping_options.rst

@@ -75,6 +75,8 @@ The following options are currently available:
 
 * :code:`.unrollCopyShared(<boolean>)`: Also unroll the copies to and from shared memory introduced by the :code:`TC` mapper. If :code:`unroll` value is not provided, has no effect.
 
+* :code:`.useReaOnlyCache(<boolean>)`: Emit loads to the readonly cache when appropriate.
+
 * :code:`.matchLibraryCalls(<boolean>)`: Replace computation patterns with calls to highly optimized libraries (such as CUB, CUTLASS) when possible.
 
 * :code:`.fixParametersBeforeScheduling(<boolean>)`: Perform automatic loop scheduling taking into account specific tensor sizes. May produce faster kernels but significantly increases compilation time. Note that the *mapping* will be performed for specific tensor sizes anyway.

tc/autotuner/parameters.cc

@@ -235,6 +235,7 @@ void TuningConfiguration::applyToParameters(
   useSharedMemory.apply(f);
   usePrivateMemory.apply(f);
   unrollCopyShared.apply(f);
+  useReadOnlyCache.apply(f);
   matchLibraryCalls.apply(f);
 }
 
@@ -268,6 +269,7 @@ std::vector<ParameterView> TuningConfiguration::collectParameters() {
   params.emplace_back(useSharedMemory);
   params.emplace_back(usePrivateMemory);
   params.emplace_back(unrollCopyShared);
+  params.emplace_back(useReadOnlyCache);
   params.emplace_back(matchLibraryCalls);
 
   return params;
@@ -298,6 +300,7 @@ void TuningConfiguration::fromCudaMappingOptions(
   useSharedMemory.selectValue(options.proto().use_shared_memory());
   usePrivateMemory.selectValue(options.proto().use_private_memory());
   unrollCopyShared.selectValue(options.proto().unroll_copy_shared());
+  useReadOnlyCache.selectValue(options.proto().use_readonly_cache());
 }
 
 void TuningConfiguration::fromCpuMappingOptions(
@@ -325,6 +328,7 @@ void TuningConfiguration::applyToCudaMappingOptions(
   options.useSharedMemory(useSharedMemory.value());
   options.usePrivateMemory(usePrivateMemory.value());
   options.unrollCopyShared(unrollCopyShared.value());
+  options.useReadOnlyCache(useReadOnlyCache.value());
 }
 
 void TuningConfiguration::applyToCpuMappingOptions(
@@ -338,6 +342,7 @@ TuningConfiguration::TuningConfiguration()
       useSharedMemory("use shared memory"),
       usePrivateMemory("use private memory"),
       unrollCopyShared("unroll copy shared"),
+      useReadOnlyCache("use readonly cache (i.e. emit __ldg loads)"),
       matchLibraryCalls("match library calls") {
   addValidator([](const TuningConfiguration& conf) {
     auto b0v = conf.blockParams.dims.at(0).value();
@@ -419,6 +424,7 @@ void TuningConfiguration::fixParameters(
   maybeFixScalar(fixedParams.useSharedMemory, useSharedMemory);
   maybeFixScalar(fixedParams.usePrivateMemory, usePrivateMemory);
   maybeFixScalar(fixedParams.unrollCopyShared, unrollCopyShared);
+  maybeFixScalar(fixedParams.useReadOnlyCache, useReadOnlyCache);
   maybeFixScalar(fixedParams.matchLibraryCalls, matchLibraryCalls);
 }
 
@@ -568,6 +574,11 @@ TuningParameterFixer& TuningParameterFixer::fixUnrollCopyShared(bool val) {
   return *this;
 }
 
+TuningParameterFixer& TuningParameterFixer::fixUseReadOnlyCache(bool val) {
+  useReadOnlyCache = val;
+  return *this;
+}
+
 TuningParameterFixer& TuningParameterFixer::fixMatchLibraryCalls(bool val) {
   matchLibraryCalls = val;
   return *this;

tc/autotuner/parameters.h

@@ -188,6 +188,7 @@ class TuningConfiguration {
   BoolParameter useSharedMemory;
   BoolParameter usePrivateMemory;
   BoolParameter unrollCopyShared;
+  BoolParameter useReadOnlyCache;
   BoolParameter matchLibraryCalls;
 
  private:
@@ -209,6 +210,7 @@ class TuningParameterFixer {
   TuningParameterFixer& fixUseSharedMemory(bool val);
   TuningParameterFixer& fixUsePrivateMemory(bool val);
   TuningParameterFixer& fixUnrollCopyShared(bool val);
+  TuningParameterFixer& fixUseReadOnlyCache(bool val);
   TuningParameterFixer& fixMatchLibraryCalls(bool val);
 
  private:
@@ -223,6 +225,7 @@ class TuningParameterFixer {
   llvm::Optional<bool> useSharedMemory;
   llvm::Optional<bool> usePrivateMemory;
   llvm::Optional<bool> unrollCopyShared;
+  llvm::Optional<bool> useReadOnlyCache;
   llvm::Optional<bool> matchLibraryCalls;
 
   friend class TuningConfiguration;

tc/core/cuda/cuda_mapping_options.cc

@@ -311,6 +311,11 @@ CudaMappingOptions& CudaMappingOptions::unrollCopyShared(bool b) {
   return *this;
 }
 
+CudaMappingOptions& CudaMappingOptions::useReadOnlyCache(bool b) {
+  ownedProto_.set_use_readonly_cache(b);
+  return *this;
+}
+
 CudaMappingOptions& CudaMappingOptions::mapToThreads(
     const std::string& commaSeparatedSizes) {
   auto sizes = parseCommaSeparatedIntegers<uint64_t>(commaSeparatedSizes);
@@ -341,7 +346,8 @@ CudaMappingOptions CudaMappingOptions::makeUnmappedMappingOptions() {
   mo.genericMappingOptions(MappingOptions::makeUnmappedMappingOptions())
       .useSharedMemory(false)
       .usePrivateMemory(false)
-      .unrollCopyShared(false);
+      .unrollCopyShared(false)
+      .useReadOnlyCache(false);
   return mo;
 }
 

tc/core/cuda/cuda_mapping_options.h

@@ -194,6 +194,7 @@ class CudaMappingOptions {
   CudaMappingOptions& usePrivateMemory(bool b);
   CudaMappingOptions& maxSharedMemory(uint64_t size);
   CudaMappingOptions& unrollCopyShared(bool b);
+  CudaMappingOptions& useReadOnlyCache(bool b);
   ///@}
 
   /// Static constructors for predefined strategies.

tc/core/cuda/cuda_mapping_options_cpp_printer.cc

@@ -32,6 +32,8 @@ CudaMappingOptionsCppPrinter& operator<<(
       "usePrivateMemory", cudaOptions.proto().use_private_memory());
   prn.printBooleanOption(
       "unrollCopyShared", cudaOptions.proto().unroll_copy_shared());
+  prn.printBooleanOption(
+      "useReadOnlyCache", cudaOptions.proto().use_readonly_cache());
   if (cudaOptions.proto().has_max_shared_memory()) {
     prn.printValueOption(
         "maxSharedMemory", cudaOptions.proto().max_shared_memory());

tc/core/polyhedral/cuda/codegen.cc

@@ -365,8 +365,32 @@ void emitReductionInit(
 }
 
 namespace {
+// RAII-style wrapper around CodegenStatementContext that wraps the output
+// streamed to context.ss into an "__ldg()" intrinsic, if the tensor
+// with the given identifier is known to be read-only.
+struct LdgWrapper {
+ public:
+  LdgWrapper(const CodegenStatementContext& context, isl::id id)
+      : readOnly_(context.readOnlySet.count(id) > 0), out_(context.ss) {
+    if (readOnly_) {
+      out_ << "__ldg(&";
+    }
+  }
+
+  ~LdgWrapper() {
+    if (readOnly_) {
+      out_ << ")";
+    }
+  }
+
+ private:
+  bool readOnly_;
+  std::ostream& out_;
+};
+
 template <typename AFF>
 void emitAccess(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();
 }
 } // namespace
@@ -584,6 +608,8 @@ void emitMappedTensorAccess(
 
   // Not promoted, emitting just the mapped subscript.
   if (!promotionInfo.groupId) {
+    auto ctx = context.scop().domain().get_ctx();
+    LdgWrapper ldgWrapper(context, isl::id(ctx, name));
     context.ss << name;
     for (auto e : subscripts) {
       context.ss << "[";
@@ -681,6 +707,29 @@ size_t& nAstNodes() {
   return n;
 }
 
+namespace {
+// Collect ids of tensors that are only read on the Scop.
+std::unordered_set<isl::id, isl::IslIdIslHash> gatherReadOnlySet(
+    const MappedScop& mscop) {
+  std::unordered_set<isl::id, isl::IslIdIslHash> readOnlySet;
+
+  if (!mscop.useReadOnlyCache) {
+    return readOnlySet;
+  }
+
+  const auto& scop = mscop.scop();
+
+  auto read = scop.reads.universe().range();
+  auto written = scop.writes.universe().range();
+  auto readOnly = read.subtract(written);
+  for (auto s : readOnly.get_set_list()) {
+    readOnlySet.emplace(s.get_tuple_id());
+  }
+
+  return readOnlySet;
+}
+} // namespace
+
 string emitCudaKernel(
     const std::string& specializedName,
     const MappedScop& mscop) {
@@ -745,7 +794,9 @@ string emitCudaKernel(
   auto root = mscop.schedule();
   astBuild = astBuild.set_iterators(Codegen::makeLoopIterators(root));
   auto astNode = astBuild.node_from(schedule);
-  AstPrinter(CodegenContext(ss, mscop, nodeInfoMap)).emit(astNode);
+
+  AstPrinter(CodegenContext(ss, mscop, nodeInfoMap, gatherReadOnlySet(mscop)))
+      .emit(astNode);
   ss << "}" << endl;
 
   return ss.str();

tc/core/polyhedral/cuda/codegen.h

@@ -83,10 +83,14 @@ struct CodegenContext {
   CodegenContext(
       std::stringstream& ss_,
       const MappedScop& s,
-      const NodeInfoMapType& i)
-      : ss(ss_), mappedScop(s), nodeInfoMap(i) {}
+      const NodeInfoMapType& i,
+      const std::unordered_set<isl::id, isl::IslIdIslHash>& ros)
+      : ss(ss_), mappedScop(s), nodeInfoMap(i), readOnlySet(ros) {}
   CodegenContext(const CodegenContext& c)
-      : ss(c.ss), mappedScop(c.mappedScop), nodeInfoMap(c.nodeInfoMap) {}
+      : ss(c.ss),
+        mappedScop(c.mappedScop),
+        nodeInfoMap(c.nodeInfoMap),
+        readOnlySet(c.readOnlySet) {}
 
   const Scop& scop() const {
     return mappedScop.scop();
@@ -95,6 +99,7 @@ struct CodegenContext {
   std::stringstream& ss;
   const MappedScop& mappedScop;
   const NodeInfoMapType& nodeInfoMap;
+  const std::unordered_set<isl::id, isl::IslIdIslHash>& readOnlySet;
 };
 
 struct CodegenStatementContext : CodegenContext {

tc/core/polyhedral/cuda/mapped_scop.cc

@@ -908,7 +908,11 @@ std::unique_ptr<MappedScop> makeSpecializedMappedScop(
   tc::Block block = mappedScop.numThreads;
   std::tie(grid, block) = tightenLaunchBounds(*scop, grid, block);
   auto res = MappedScop::makeMappedScop(
-      std::move(scop), grid, block, mappedScop.unroll);
+      std::move(scop),
+      grid,
+      block,
+      mappedScop.unroll,
+      mappedScop.useReadOnlyCache);
   res->insertMappingContext();
 
   LOG_IF(INFO, FLAGS_debug_tc_mapper)
@@ -985,7 +989,8 @@ std::unique_ptr<MappedScop> MappedScop::makeWithOuterBlockInnerThreadStrategy(
       std::move(scopUPtr),
       ::tc::Grid(cudaOptions.grid),
       ::tc::Block(cudaOptions.block),
-      generic.proto.unroll()));
+      generic.proto.unroll(),
+      cudaOptions.proto().use_readonly_cache()));
   auto& scop = mappedScop->scop_;
 
   // 1a. Optionally specialize before scheduling...

tc/core/polyhedral/cuda/mapped_scop.h

@@ -61,25 +61,28 @@ class MappedScop {
       std::unique_ptr<Scop>&& scop,
       ::tc::Grid grid,
       ::tc::Block block,
-      uint64_t unroll_)
+      uint64_t unroll_,
+      bool useReadOnlyCache_)
       : scop_(std::move(scop)),
         numBlocks(grid),
         numThreads(block),
-        unroll(unroll_) {}
+        unroll(unroll_),
+        useReadOnlyCache(useReadOnlyCache_) {}
 
  public:
   static inline std::unique_ptr<MappedScop> makeOneBlockOneThread(
       std::unique_ptr<Scop>&& scop) {
     return std::unique_ptr<MappedScop>(new MappedScop(
-        std::move(scop), ::tc::Grid{1, 1, 1}, ::tc::Block{1, 1, 1}, 1));
+        std::move(scop), ::tc::Grid{1, 1, 1}, ::tc::Block{1, 1, 1}, 1, false));
   }
   static inline std::unique_ptr<MappedScop> makeMappedScop(
       std::unique_ptr<Scop>&& scop,
       ::tc::Grid grid,
       ::tc::Block block,
-      uint64_t unroll) {
+      uint64_t unroll,
+      bool useReadOnlyCache) {
     return std::unique_ptr<MappedScop>(
-        new MappedScop(std::move(scop), grid, block, unroll));
+        new MappedScop(std::move(scop), grid, block, unroll, useReadOnlyCache));
   }
 
   // Apply the hand-written OuterBlockInnerThread mapping strategy.
@@ -206,6 +209,7 @@ class MappedScop {
   const ::tc::Grid numBlocks;
   const ::tc::Block numThreads;
   const uint64_t unroll;
+  const bool useReadOnlyCache;
 
  private:
   // Information about a detected reduction that can potentially