Adding assert requiring shmem when tmem needs to be allocated (#46)

triton-lang · Jul 29, 2024 · 3ede94f · 3ede94f
1 parent 787552a
commit 3ede94f
Showing 1 changed file with 7 additions and 0 deletions.
diff --git a/lib/Dialect/TritonNvidiaGPU/Transforms/TensorMemoryAllocation.cpp b/lib/Dialect/TritonNvidiaGPU/Transforms/TensorMemoryAllocation.cpp
@@ -72,6 +72,13 @@ class TritionTensorMemoryAllocationPass
       }
     }
 
+    if (totalMemorySize > 0) {
+      assert(mod->getAttr("triton_gpu.shared") != nullptr &&
+             cast<IntegerAttr>(mod->getAttr("triton_gpu.shared")).getInt() !=
+                 0 &&
+             "Shared memory is required for allocation of Tensor Core memory.");
+    }
+
     mod->setAttr("triton_gpu.tensor_memory_size",
                  mlir::IntegerAttr::get(mlir::IntegerType::get(ctx, 32),
                                         totalMemorySize));