gaspersic/triton
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

[FRONTEND] Added comment on TensorSizeTrait::maxElement (#20)

Browse Source
This commit is contained in:
Philippe Tillet
2022-07-25 01:18:45 -07:00
committed by GitHub
parent 7eda373a12
commit 27c9f3d8cb
1 changed files with 7 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
include/triton/Dialect/Triton/IR/Traits.h
View File

@@ -17,6 +17,13 @@ class TensorSizeTrait : public TraitBase<ConcreteType, TensorSizeTrait> {
public: public:
// TODO: move impl to .cc files // TODO: move impl to .cc files
static LogicalResult verifyTrait(Operation *op) { static LogicalResult verifyTrait(Operation *op) {
// The rationale for this number is to prevent users from creating programs
// that would have catastrophic register pressure and cause the compiler to
// hang.
// Since H100 has 256KB registers, we should allow users to create tensors
// of size up to 256K elements. It will spill for datatypes wider than 1B,
// but we probably should limit number of elements (rather than bytes) to
// keep specs simple
int constexpr maxElement = 1048576; int constexpr maxElement = 1048576;
for (auto opType : op->getOperandTypes()) { for (auto opType : op->getOperandTypes()) {
if (auto tensorType = opType.dyn_cast<RankedTensorType>()) { if (auto tensorType = opType.dyn_cast<RankedTensorType>()) {