[CUDA][shared memory allocation]fix 'ptxas error : Entry function 'fu…

…sion_##' uses too much shared data'
apache · Aug 12, 2024 · 6e09af0 · 6e09af0
1 parent b3d01c2
commit 6e09af0
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 3 deletions.
diff --git a/include/tvm/ir/type.h b/include/tvm/ir/type.h
@@ -157,7 +157,7 @@ class PointerTypeNode : public TypeNode {
   /*!
    * \brief The storage scope of the pointer
    */
-  String storage_scope;
+  mutable String storage_scope;
 
   void VisitAttrs(AttrVisitor* v) {
     v->Visit("element_type", &element_type);

diff --git a/src/target/source/codegen_cuda.cc b/src/target/source/codegen_cuda.cc
@@ -642,12 +642,14 @@ void CodeGenCUDA::PrintStorageSync(const CallNode* op) {
 }
 
 void CodeGenCUDA::PrintStorageScope(const std::string& scope, std::ostream& os) {  // NOLINT(*)
-  ICHECK_NE(scope, "global") << "Cannot allocate global memory when targeting CUDA. You must pass "
-                                "all global arrays as input instead";
+  // ICHECK_NE(scope, "global") << "Cannot allocate global memory when targeting CUDA. You must pass "
+  //                               "all global arrays as input instead";
   if (scope == "shared") {
     os << "__shared__ ";
   } else if (scope == "shared.dyn") {
     os << "extern __shared__ ";
+  } else if (scope == "global") {
+    os << "__device__ static ";
   }
 }
 

diff --git a/src/tir/transforms/storage_flatten.cc b/src/tir/transforms/storage_flatten.cc
@@ -1518,8 +1518,12 @@ class StorageFlattener : public StmtExprMutator {
       StorageScope skey = StorageScope::Create(GetPtrStorageScope(op->buffer->data));
 
       // use small alignment for small arrays
+      auto* ptr_type = op->buffer->data->type_annotation.as<PointerTypeNode>();
       auto dtype = op->buffer->dtype;
       size_t const_size = AllocateNode::ConstantAllocationSize(op->buffer->shape);
+      if (const_size > 41984) {
+        ptr_type->storage_scope = tvm::runtime::String("global");
+      }
       int align = GetTempAllocaAlignment(dtype, const_size);
       if (skey.tag.length() != 0) {
         MemoryInfo info = GetMemoryInfo(skey.to_string());