From 8337a3fb3f44a536ba40e895dad7441b2f4a59f2 Mon Sep 17 00:00:00 2001 From: Mikhail Zolotukhin Date: Tue, 31 Aug 2021 20:27:44 -0700 Subject: [PATCH] [TensorExpr] Wrap error messages with buildErrorMessage call. (#64330) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/64330 Test Plan: Imported from OSS Reviewed By: bertmaher Differential Revision: D30687226 Pulled By: ZolotukhinM fbshipit-source-id: ade1be2ad6847c6afbba60307ef854696821b4e3 --- test/cpp/tensorexpr/test_loopnest.cpp | 2 +- torch/csrc/jit/tensorexpr/llvm_codegen.cpp | 19 ++-- torch/csrc/jit/tensorexpr/loopnest.cpp | 116 ++++++++++++++++----- .../csrc/jit/tensorexpr/mem_dependency_checker.cpp | 19 +++- torch/csrc/jit/tensorexpr/registerizer.cpp | 15 ++- 5 files changed, 129 insertions(+), 42 deletions(-) diff --git a/test/cpp/tensorexpr/test_loopnest.cpp b/test/cpp/tensorexpr/test_loopnest.cpp index c2b33e2..b1d59a1 100644 --- a/test/cpp/tensorexpr/test_loopnest.cpp +++ b/test/cpp/tensorexpr/test_loopnest.cpp @@ -3554,7 +3554,7 @@ TEST(LoopNest, DetectInlineRankMismatch) { LoopNest l({reshape}, {a, reshape}); ASSERT_THROWS_WITH( l.computeInline(l.getLoopBodyFor(a)), - "Placeholder indexed access is inconsistent with its rank"); + "Number of indices doesn't match buf rank in the fuser."); } TEST(LoopNest, CacheReadsSimple) { diff --git a/torch/csrc/jit/tensorexpr/llvm_codegen.cpp b/torch/csrc/jit/tensorexpr/llvm_codegen.cpp index b9ea708..6c212e6 100644 --- a/torch/csrc/jit/tensorexpr/llvm_codegen.cpp +++ b/torch/csrc/jit/tensorexpr/llvm_codegen.cpp @@ -903,7 +903,9 @@ void LLVMCodeGenImpl::visit(HalfImmPtr v) { } void LLVMCodeGenImpl::visit(BFloat16ImmPtr v) { - TORCH_INTERNAL_ASSERT(false, "llvm codegen does not support bfloat16"); + TORCH_INTERNAL_ASSERT( + false, + buildErrorMessage("Fuser's LLVM codegen does not support bfloat16")); } void LLVMCodeGenImpl::visit(BoolImmPtr v) { @@ -1535,7 +1537,10 @@ void LLVMCodeGenImpl::emitIsNan(IntrinsicsPtr v) { if (!v->param(0)->dtype().is_floating_point()) { value_ = toVec(llvm::ConstantInt::get(dstType, 0), v->dtype().lanes()); } else { - TORCH_INTERNAL_ASSERT(v->dtype().scalar_type() == ScalarType::Int); + TORCH_INTERNAL_ASSERT( + v->dtype().scalar_type() == ScalarType::Int, + buildErrorMessage( + "Unexpected non-Int dtype of Intrinsics' result value in the fuser.")); auto is_nan = irb_.CreateFCmpUNO( value_, llvm::ConstantFP::get(value_->getType(), 0.)); if (v->dtype().lanes() > 1) { @@ -1762,11 +1767,11 @@ void LLVMCodeGenImpl::visit(IntrinsicsPtr v) { } else { TORCH_INTERNAL_ASSERT( false, - v, - "Unimplemented lowering:", - v->op_type(), - " for input of dtype", - v->dtype().scalar_dtype()); + buildErrorMessage( + std::string("Unimplemented lowering for intrinsic '") + + std::to_string(v->op_type()) + "' for input of dtype " + + std::to_string(v->dtype().scalar_dtype()) + + " in LLVM codegen of the fuser.")); } std::vector params; diff --git a/torch/csrc/jit/tensorexpr/loopnest.cpp b/torch/csrc/jit/tensorexpr/loopnest.cpp index 11020cc..e67d094 100644 --- a/torch/csrc/jit/tensorexpr/loopnest.cpp +++ b/torch/csrc/jit/tensorexpr/loopnest.cpp @@ -536,7 +536,8 @@ class FunctionInliner : public IRMutator { // (since we don't support in-place writes). Resolves issue 52581. TORCH_INTERNAL_ASSERT( *intValue(i) == 0, - "Constant index impression should always be zero"); + buildErrorMessage( + "Unexpected non-zero constant index in inlined buffer in the fuser.")); producer_index_vars_.push_back(nullptr); } else { throw std::logic_error("cannot inline Buf with compound indices"); @@ -547,21 +548,25 @@ class FunctionInliner : public IRMutator { private: ExprPtr mutate_loads(BufPtr buf, std::vector dims) { std::vector index_vars; - TORCH_INTERNAL_ASSERT(buf->ndim() == producer_index_vars_.size()); + TORCH_INTERNAL_ASSERT( + buf->ndim() == producer_index_vars_.size(), + buildErrorMessage( + "Dimensions of producer and consumer expressions do not match in inliner in the fuser.")); for (const auto i : c10::irange(buf->ndim())) { VarPtr func_callee_arg = producer_index_vars_.at(i); ExprPtr func_caller_param = dims.at(i); if (func_callee_arg == nullptr) { TORCH_INTERNAL_ASSERT( intValue(func_caller_param) && *intValue(func_caller_param) == 0, - "We are implicitly assuming that if you have an index of 0, that must also be inlined into an index of 0"); + buildErrorMessage( + "We are implicitly assuming that if you have an index of 0, that must also be inlined into an index of 0")); continue; } if (func_callee_arg == nullptr) continue; auto iter = inline_mapping_.find(func_callee_arg); if (iter != inline_mapping_.end()) { - throw std::runtime_error( + throw std::logic_error( "Duplicated variables: " + func_callee_arg->name_hint()); } // Add a mapping for each function parameter to it's source name. @@ -603,10 +608,10 @@ class FunctionInliner : public IRMutator { return IRMutator::mutate(v); } - if (v->indices().size() != buf->ndim()) { - throw malformed_input( - "Placeholder indexed access is inconsistent with its rank", v); - } + TORCH_INTERNAL_ASSERT( + v->indices().size() == buf->ndim(), + buildErrorMessage( + "Number of indices doesn't match buf rank in the fuser.")); return mutate_loads(buf, v->indices()); } @@ -646,7 +651,10 @@ class FunctionInliner : public IRMutator { if (v == producer_ && !outputs_.count(buf_)) { in_producer_ = true; producer_ = to(IRMutator::mutate(v)); - TORCH_INTERNAL_ASSERT(producer_ != nullptr); + TORCH_INTERNAL_ASSERT( + producer_, + buildErrorMessage( + "Producer statement for output buf should remain non-null in the fuser")); in_producer_ = false; return nullptr; } else { @@ -748,7 +756,10 @@ bool LoopNest::computeInline(BufPtr b) { } } - TORCH_INTERNAL_ASSERT(relevant_store); + TORCH_INTERNAL_ASSERT( + relevant_store, + buildErrorMessage( + "Cannot find a relevant store to inline a buf in the fuser.")); GRAPH_DEBUG("ComputeInline: Def: ", std::to_string(relevant_store)); FunctionInliner inliner(relevant_store, output_bufs_); @@ -772,7 +783,11 @@ void LoopNest::inlineIntermediateBufs(bool allow_duplicated_work) { auto input_bufs = getInputBufs(); for (auto buf : intermediate_bufs) { - TORCH_INTERNAL_ASSERT(buf_load_store_uses.count(buf)); + TORCH_INTERNAL_ASSERT( + buf_load_store_uses.count(buf), + buildErrorMessage( + "Could not find uses of buf '" + buf->name_hint() + + "' in the fuser.")); std::vector& uses = buf_load_store_uses[buf]; auto stores = c10::filter( uses, [](const BufLoadOrStoreUse& use) { return use.isStore; }); @@ -789,7 +804,11 @@ void LoopNest::inlineIntermediateBufs(bool allow_duplicated_work) { } } else { // If S is not a store, it must be an ExternalCall. - TORCH_INTERNAL_ASSERT(to(stores[0].s)); + TORCH_INTERNAL_ASSERT( + to(stores[0].s), + buildErrorMessage( + "Expected stmt: " + std::to_string(stores[0].s) + + "\nto be either a Store or an ExternalCall in the fuser.")); } } @@ -1154,7 +1173,10 @@ bool LoopNest::optimizeConditionals() { ifthenelse_exprs.front(), &cond_var, &comp_values, &sub_exprs)) { continue; } - TORCH_INTERNAL_ASSERT(comp_values.size() >= 1); + TORCH_INTERNAL_ASSERT( + comp_values.size() >= 1, + buildErrorMessage( + "Expected at least one expression in optimizeConditional in the fuser.")); comp_values.insert(comp_values.begin(), immLike(comp_values[0], 0)); auto fors = getLoopStmtsFor(store); @@ -1513,7 +1535,10 @@ void LoopNest::splitWithMask(ForPtr f, int factor, ForPtr* inner) { std::vector LoopNest::distributeLoop( ForPtr loop, const std::unordered_set& pivots) { - TORCH_INTERNAL_ASSERT(loop); + TORCH_INTERNAL_ASSERT( + loop, + buildErrorMessage( + "Expected non-null loop in distributeLoop in the fuser.")); auto root = loop->get_parent(); if (root == nullptr) { throw malformed_input("Loop without parent: ", loop); @@ -1758,7 +1783,10 @@ bool LoopNest::unsafeFuseLoops( break; } } - TORCH_INTERNAL_ASSERT(it != root_block->end()); + TORCH_INTERNAL_ASSERT( + it != root_block->end(), + buildErrorMessage( + "Could not find the given loop in the root stmt in unsafeFuseLoop the fuser.")); for (auto l : loops) { if (*it != l) { return false; @@ -2032,7 +2060,10 @@ std::vector LoopNest::reorder( parent->replace_stmt(loops.front(), empty_block); for (size_t i = 1; i < loops.size(); ++i) { auto block = to(loops[i]->get_parent()); - TORCH_INTERNAL_ASSERT(block); + TORCH_INTERNAL_ASSERT( + block, + buildErrorMessage( + "Expected parent stmt to be a non-null Block in reorder transformation the fuser.")); block->remove_stmt(loops[i]); } @@ -2191,9 +2222,13 @@ std::vector LoopNest::getLoopStmtsInLoopNest(ForPtr f, size_t num) { ForPtr curr_for = f; loops[0] = curr_for; for (size_t i = 1; i < num; ++i) { - TORCH_INTERNAL_ASSERT(curr_for->body()->nstmts() == 1); + TORCH_INTERNAL_ASSERT( + curr_for->body()->nstmts() == 1, + buildErrorMessage("Expected a single stmt in the loop body.")); curr_for = to(curr_for->body()->front()); - TORCH_INTERNAL_ASSERT(curr_for); + TORCH_INTERNAL_ASSERT( + curr_for, + buildErrorMessage("Expected the only child stmt to be a For loop.")); loops[i] = curr_for; } return loops; @@ -2303,7 +2338,10 @@ void LoopNest::compressBuffer(BufPtr buf, StmtPtr stmt) { // Find the parent common to all the buffer accesses. BlockPtr parent = to(writes.front()->get_parent()); - TORCH_INTERNAL_ASSERT(parent); + TORCH_INTERNAL_ASSERT( + parent, + buildErrorMessage( + "Expected parent stmt to be a non-null block in compressBuffer in the fuser.")); for (auto w : writes) { parent = Block::getSharedParent(parent, w); } @@ -2325,7 +2363,10 @@ void LoopNest::compressBuffer(BufPtr buf, StmtPtr stmt) { // Vector to indicate which dimensions could be compressed away. std::vector dims(buf->dims().size(), true); auto check_indices = [&](const std::vector& indices) { - TORCH_INTERNAL_ASSERT(indices.size() == dims.size()); + TORCH_INTERNAL_ASSERT( + indices.size() == dims.size(), + buildErrorMessage( + "Expected ranks to match in compressBuffer in the fuser.")); for (size_t i = 0; i < indices.size(); ++i) { auto index_vars = NodeFinder::find(indices[i]); for (auto iv : index_vars) { @@ -2367,7 +2408,10 @@ void LoopNest::compressBuffer(BufPtr buf, StmtPtr stmt) { // Modify all access to reflect the removed dims. auto get_new_indices = [&](const std::vector& indices) { - TORCH_INTERNAL_ASSERT(indices.size() == dims.size()); + TORCH_INTERNAL_ASSERT( + indices.size() == dims.size(), + buildErrorMessage( + "Expected ranks to match in compressBuffer in the fuser.")); std::vector new_indices(indices); for (size_t i = 0; i < dims.size(); ++i) { if (dims[i]) { @@ -2573,7 +2617,10 @@ class CacheReplacer : public IRMutator { // Map indices to call-parameters. std::vector newIndices; - TORCH_INTERNAL_ASSERT(offsets_.size() == v->indices().size()); + TORCH_INTERNAL_ASSERT( + offsets_.size() == v->indices().size(), + buildErrorMessage( + "Expected ranks to match in CacheReplacer in the fuser.")); for (size_t i = 0; i < v->indices().size(); ++i) { ExprPtr index = v->indices()[i]->accept_mutator(this); ExprPtr offset = offsets_[i]; @@ -2595,7 +2642,10 @@ class CacheReplacer : public IRMutator { // Map indices to call-parameters. std::vector newIndices; - TORCH_INTERNAL_ASSERT(offsets_.size() == v->indices().size()); + TORCH_INTERNAL_ASSERT( + offsets_.size() == v->indices().size(), + buildErrorMessage( + "Expected ranks to match in CacheReplacer in the fuser.")); for (size_t i = 0; i < v->indices().size(); ++i) { ExprPtr index = v->indices()[i]->accept_mutator(this); ExprPtr offset = offsets_[i]; @@ -2643,7 +2693,10 @@ LoopNest::AccessResult LoopNest::cacheAccesses( return {nullptr, nullptr}; } - TORCH_INTERNAL_ASSERT(bounds_it->second.size() == 1); + TORCH_INTERNAL_ASSERT( + bounds_it->second.size() == 1, + buildErrorMessage( + "Unexpected number of bound info entries in cacheAccesses in the fuser.")); TensorAccessBoundsInfo& info = bounds_it->second[0]; bool hasReads = info.kind == kLoad || info.kind == kMutate; bool hasWrites = info.kind == kStore || info.kind == kMutate; @@ -2998,7 +3051,10 @@ class RfactorStoreRewriter : public IRMutator { return IRMutator::mutate(v); } - TORCH_INTERNAL_ASSERT(old_indices_.size() == v->indices().size()); + TORCH_INTERNAL_ASSERT( + old_indices_.size() == v->indices().size(), + buildErrorMessage( + "Expected ranks to match in RfactorStoreRewriter in the fuser.")); bool equal_indices = true; for (size_t i = 0; i < v->indices().size(); ++i) { @@ -3032,7 +3088,10 @@ class RfactorStoreRewriter : public IRMutator { return IRMutator::mutate(v); } - TORCH_INTERNAL_ASSERT(old_indices_.size() == v->indices().size()); + TORCH_INTERNAL_ASSERT( + old_indices_.size() == v->indices().size(), + buildErrorMessage( + "Expected ranks to match in RfactorStoreRewriter in the fuser.")); bool equal_indices = true; for (size_t i = 0; i < v->indices().size(); ++i) { @@ -3141,7 +3200,10 @@ bool LoopNest::rfactor( // X[*indexes] = ReduceOp(X[*indexes] + T[*indexes + {reduction_var}], // reduce_axis={reduction_var}) BlockPtr b = outer_reduction_for->body(); - TORCH_INTERNAL_ASSERT(b->nstmts() == 1); + TORCH_INTERNAL_ASSERT( + b->nstmts() == 1, + buildErrorMessage( + "Expected to have a single stmt in the block in rfactor transformation in the fuser.")); StmtPtr first_reduction_loop = b->stmts().front(); auto rfac_buf_indices = orig_buf_indices; rfac_buf_indices.emplace_back(reduction_var); diff --git a/torch/csrc/jit/tensorexpr/mem_dependency_checker.cpp b/torch/csrc/jit/tensorexpr/mem_dependency_checker.cpp index e1688e3..3f77041 100644 --- a/torch/csrc/jit/tensorexpr/mem_dependency_checker.cpp +++ b/torch/csrc/jit/tensorexpr/mem_dependency_checker.cpp @@ -76,12 +76,16 @@ std::vector AccessInfo::getIndices() const { void AccessInfo::addDependency(const std::shared_ptr& write) { auto res = dependencies_.emplace(write->id(), write); - TORCH_INTERNAL_ASSERT(res.second); + TORCH_INTERNAL_ASSERT( + res.second, + buildErrorMessage("Duplicate entry in mem dep checker in the fuser.")); } void AccessInfo::addDependent(const std::shared_ptr& read) { auto res = dependents_.emplace(read->id(), read); - TORCH_INTERNAL_ASSERT(res.second); + TORCH_INTERNAL_ASSERT( + res.second, + buildErrorMessage("Duplicate entry in mem dep checker in the fuser.")); } bool AccessInfo::hasDependency(const std::shared_ptr& info) const { @@ -590,7 +594,10 @@ bool executionSafetyCheck( if (aStrides.empty() || oStrides.empty()) { return false; } - TORCH_INTERNAL_ASSERT(info->bounds().size() == other->bounds().size()); + TORCH_INTERNAL_ASSERT( + info->bounds().size() == other->bounds().size(), + buildErrorMessage( + "Dimension mismatch for two accesses in mem dep checker in the fuser.")); for (size_t b = 0; b < info->bounds().size(); ++b) { ExprPtr aIndexStride = aStrides[b]; ExprPtr oIndexStride = oStrides[b]; @@ -1150,7 +1157,11 @@ void MemDependencyChecker::visit(FreePtr v) { VarPtr var = v->buffer_var(); auto it = intermediates_.find(var); - TORCH_INTERNAL_ASSERT(it != intermediates_.end()); + TORCH_INTERNAL_ASSERT( + it != intermediates_.end(), + buildErrorMessage( + "Expected to find '" + var->name_hint() + + "' in intermediate vars in mem dep checker in the fuser.")); IndexBounds bounds = it->second->bounds(); auto info = std::make_shared( diff --git a/torch/csrc/jit/tensorexpr/registerizer.cpp b/torch/csrc/jit/tensorexpr/registerizer.cpp index 8684f2a..c4c4957 100644 --- a/torch/csrc/jit/tensorexpr/registerizer.cpp +++ b/torch/csrc/jit/tensorexpr/registerizer.cpp @@ -43,8 +43,14 @@ void AccessInfo::addLoad( } void AccessInfo::merge(const std::shared_ptr& other) { - TORCH_INTERNAL_ASSERT(hash_ == other->hash()); - TORCH_INTERNAL_ASSERT(indices_.size() == other->indices().size()); + TORCH_INTERNAL_ASSERT( + hash_ == other->hash(), + buildErrorMessage( + "Expected hashes to match in registerizer in the fuser.")); + TORCH_INTERNAL_ASSERT( + indices_.size() == other->indices().size(), + buildErrorMessage( + "Expected ranks to match in registerizer in the fuser.")); last_usage_ = other->last_usage(); for (auto s : other->stores()) { @@ -68,7 +74,10 @@ void AccessInfo::merge(const std::shared_ptr& other) { bool AccessInfo::overlaps(const std::shared_ptr& other) { // All accesses to a buf must have the same dimensionality. - TORCH_INTERNAL_ASSERT(indices_.size() == other->indices().size()); + TORCH_INTERNAL_ASSERT( + indices_.size() == other->indices().size(), + buildErrorMessage( + "Expected ranks to match in registerizer in the fuser.")); auto& other_indices = other->indices(); -- 2.7.4