Lunderberg commented on code in PR #13947: URL: https://github.com/apache/tvm/pull/13947#discussion_r1124808234
########## src/tir/analysis/identify_memcpy.cc: ########## @@ -0,0 +1,317 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +/*! + * \file tir/analysis/identify_memcpy.cc + * \brief Check if a loopnest is equivalent to memcpy + */ + +#include <tvm/arith/bound.h> +#include <tvm/arith/iter_affine_map.h> +#include <tvm/runtime/container/optional.h> +#include <tvm/tir/analysis.h> +#include <tvm/tir/buffer.h> +#include <tvm/tir/stmt.h> + +#include <optional> +#include <sstream> +#include <string> +#include <variant> + +#include "../../arith/ir_visitor_with_analyzer.h" + +namespace tvm { +namespace tir { + +std::variant<MemCpyDetails, std::string> IdentifyMemCpyImpl(const For& loop, + arith::Analyzer* analyzer) { + Map<Var, arith::IntSet> loop_intervals; + Map<Var, Range> loop_ranges; + PrimExpr total_loop_iterations = 1; + + // Walk through the loopnest, stopping at the first loop whose body + // is not a loop. + Stmt stmt = loop; + while (auto* for_node = stmt.as<ForNode>()) { + loop_ranges.Set(for_node->loop_var, Range::FromMinExtent(for_node->min, for_node->extent)); + loop_intervals.Set(for_node->loop_var, + arith::IntSet::FromMinExtent(for_node->min, for_node->extent)); + total_loop_iterations = total_loop_iterations * for_node->extent; + + stmt = for_node->body; + } + + BufferStore store; + if (auto* ptr = stmt.as<BufferStoreNode>()) { + store = GetRef<BufferStore>(ptr); + } else { + return static_cast<const std::stringstream&>( + std::stringstream() + << "Expected innermost loop to have BufferStore body, but instead found " << stmt) + .str(); + } + + BufferLoad load; + if (auto* ptr = store->value.as<BufferLoadNode>()) { + load = GetRef<BufferLoad>(ptr); + } else { + return static_cast<const std::stringstream&>( + std::stringstream() + << "Expected BufferStore's value to be BufferLoad, but instead found " + << store->value) + .str(); + } + + // Now, we have a BufferStore whose value is a BufferLoad. Because + // non-flat physical indices are target-dependent, only handle cases + // where the buffer will be flattened to a 1-d physical buffer. + Array<PrimExpr> flattened_dst = store->buffer.OffsetOf(store->indices); + Array<PrimExpr> flattened_src = load->buffer.OffsetOf(load->indices); + + if (flattened_dst.size() != 1 || flattened_src.size() != 1) { + return static_cast<const std::stringstream&>( + std::stringstream() + << "Expected flattened dimension of src/dest to be 1, but found" + << flattened_src.size() << "-d src and " << flattened_dst.size() << "-d dst") + .str(); + } + PrimExpr src_index = flattened_src[0]; + PrimExpr dst_index = flattened_dst[0]; + + // First check, do the input/output form affine subsets of their + // respective buffers? + // + // For example, should exclude the following, indices are not affine + // + // for i in T.serial(16): + // B[i] = A[T.abs(i-8)] + + auto src_iter_map = arith::DetectIterMap({src_index}, loop_ranges, Bool(true), + arith::IterMapLevel::Bijective, analyzer); + if (src_iter_map->errors.size()) { + return static_cast<const std::stringstream&>(std::stringstream() + << "arith::DetectIterMap(src) returned " + << src_iter_map->errors.size() << " errors: [" + << src_iter_map->errors << "]" + << " for src_index = " << src_index) + .str(); + } + auto dst_iter_map = arith::DetectIterMap({dst_index}, loop_ranges, Bool(true), + arith::IterMapLevel::Bijective, analyzer); + if (dst_iter_map->errors.size()) { + return static_cast<const std::stringstream&>(std::stringstream() + << "arith::DetectIterMap(dst) returned " + << dst_iter_map->errors.size() << " errors: [" + << dst_iter_map->errors << "]" + << " for dst_index = " << dst_index) + .str(); + } + + // Second check, are those affine subsets contiguous? If so, then + // the index expressions will visit every location between the min + // and the max. This checks surjectivity over a linear region, + // which may not be the same as DetectIterMap's check of + // surjectivity over the affine subset. + // + // For example, should exclude the following, doesn't touch all + // output locations within the output region touched. + // + // for i in T.serial(16): + // B[2*i] = A[i] + // + // Similarly, should exclude the following, doesn't touch all + // input locations within the input region touched. + // + // for i in T.serial(16): + // B[i] = A[2*i] + total_loop_iterations = analyzer->Simplify(total_loop_iterations); + auto src_interval = analyzer->int_set(src_index, loop_intervals); + auto dst_interval = analyzer->int_set(dst_index, loop_intervals); + + if (!src_interval.HasLowerBound() || !src_interval.HasUpperBound()) { + return static_cast<const std::stringstream&>(std::stringstream() + << "Expected known bounds for src, but found " + << src_interval << " for expression " << src_index) + .str(); + } + if (!dst_interval.HasLowerBound() || !dst_interval.HasUpperBound()) { + return static_cast<const std::stringstream&>(std::stringstream() + << "Expected known bounds for dst, but found " + << dst_interval << " for expression " << dst_index) + .str(); + } + + { + PrimExpr must_prove = total_loop_iterations == src_interval.max() - src_interval.min() + 1; + PrimExpr simplified = analyzer->Simplify(must_prove); + if (!analyzer->CanProve(simplified)) { + return static_cast<const std::stringstream&>( + std::stringstream() + << "Mismatch between loop iterations (" << total_loop_iterations + << ") and number of src indices touched (" << src_interval + << ". Equality to prove simplified to " << simplified) + .str(); + } + } + { + PrimExpr must_prove = total_loop_iterations == dst_interval.max() - dst_interval.min() + 1; + PrimExpr simplified = analyzer->Simplify(must_prove); + if (!analyzer->CanProve(simplified)) { + return static_cast<const std::stringstream&>( + std::stringstream() + << "Mismatch between loop iterations (" << total_loop_iterations + << ") and number of dst indices touched (" << dst_interval + << ". Equality to prove simplified to " << simplified) + .str(); + } + } + + // Thrid check, is there a transformation applied between the input Review Comment: Thank you, and fix applied. ########## src/tir/analysis/identify_memcpy.cc: ########## @@ -0,0 +1,317 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +/*! + * \file tir/analysis/identify_memcpy.cc + * \brief Check if a loopnest is equivalent to memcpy + */ + +#include <tvm/arith/bound.h> +#include <tvm/arith/iter_affine_map.h> +#include <tvm/runtime/container/optional.h> +#include <tvm/tir/analysis.h> +#include <tvm/tir/buffer.h> +#include <tvm/tir/stmt.h> + +#include <optional> +#include <sstream> +#include <string> +#include <variant> + +#include "../../arith/ir_visitor_with_analyzer.h" + +namespace tvm { +namespace tir { + +std::variant<MemCpyDetails, std::string> IdentifyMemCpyImpl(const For& loop, + arith::Analyzer* analyzer) { + Map<Var, arith::IntSet> loop_intervals; + Map<Var, Range> loop_ranges; + PrimExpr total_loop_iterations = 1; + + // Walk through the loopnest, stopping at the first loop whose body Review Comment: Thank you, and fix applied. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
