| /* |
| * 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 src/lang/data_layout.cc |
| * \brief Data SLayout expression. |
| */ |
| #include <tvm/arith/analyzer.h> |
| #include <tvm/ffi/dtype.h> |
| #include <tvm/ffi/function.h> |
| #include <tvm/ffi/reflection/registry.h> |
| #include <tvm/ir/expr.h> |
| #include <tvm/runtime/logging.h> |
| #include <tvm/s_tir/data_layout.h> |
| #include <tvm/tirx/analysis.h> |
| #include <tvm/tirx/expr.h> |
| #include <tvm/tirx/stmt_functor.h> |
| #include <tvm/tirx/var.h> |
| |
| #include <algorithm> |
| #include <cctype> |
| |
| namespace tvm { |
| namespace tirx { |
| using tirx::IterVar; |
| using tirx::IterVarNode; |
| using tirx::Var; |
| |
| TVM_FFI_STATIC_INIT_BLOCK() { |
| SLayoutNode::RegisterReflection(); |
| SBijectiveLayoutNode::RegisterReflection(); |
| } |
| |
| const SLayoutAxis SLayoutAxis::UPPER_CASE[] = { |
| SLayoutAxis('A'), SLayoutAxis('B'), SLayoutAxis('C'), SLayoutAxis('D'), SLayoutAxis('E'), |
| SLayoutAxis('F'), SLayoutAxis('G'), SLayoutAxis('H'), SLayoutAxis('I'), SLayoutAxis('J'), |
| SLayoutAxis('K'), SLayoutAxis('L'), SLayoutAxis('M'), SLayoutAxis('N'), SLayoutAxis('O'), |
| SLayoutAxis('P'), SLayoutAxis('Q'), SLayoutAxis('R'), SLayoutAxis('S'), SLayoutAxis('T'), |
| SLayoutAxis('U'), SLayoutAxis('V'), SLayoutAxis('W'), SLayoutAxis('X'), SLayoutAxis('Y'), |
| SLayoutAxis('Z')}; |
| |
| const SLayoutAxis SLayoutAxis::LOWER_CASE[] = { |
| SLayoutAxis('a'), SLayoutAxis('b'), SLayoutAxis('c'), SLayoutAxis('d'), SLayoutAxis('e'), |
| SLayoutAxis('f'), SLayoutAxis('g'), SLayoutAxis('h'), SLayoutAxis('i'), SLayoutAxis('j'), |
| SLayoutAxis('k'), SLayoutAxis('l'), SLayoutAxis('m'), SLayoutAxis('n'), SLayoutAxis('o'), |
| SLayoutAxis('p'), SLayoutAxis('q'), SLayoutAxis('r'), SLayoutAxis('s'), SLayoutAxis('t'), |
| SLayoutAxis('u'), SLayoutAxis('v'), SLayoutAxis('w'), SLayoutAxis('x'), SLayoutAxis('y'), |
| SLayoutAxis('z')}; |
| |
| const SLayoutAxis& SLayoutAxis::Get(const char name) { |
| TVM_FFI_ICHECK((name >= 'A' && name <= 'Z') || (name >= 'a' && name <= 'z')) |
| << "Invalid layout axis name: " << name << ". Has to be A-Z or a-z."; |
| return (name >= 'A' && name <= 'Z') ? SLayoutAxis::UPPER_CASE[name - 'A'] |
| : SLayoutAxis::LOWER_CASE[name - 'a']; |
| } |
| |
| const SLayoutAxis& SLayoutAxis::Get(const IterVar& itvar) { |
| const std::string axis = itvar->var.get()->name; |
| TVM_FFI_ICHECK_EQ(axis.size(), 1) << "Invalid layout axis " << axis; |
| return SLayoutAxis::Get(axis[0]); |
| } |
| |
| const SLayoutAxis& SLayoutAxis::Get(const std::string& name) { |
| TVM_FFI_ICHECK_EQ(name.length(), 1) << "Invalid axis " << name; |
| return SLayoutAxis::Get(name[0]); |
| } |
| |
| SLayout::SLayout(const ffi::Array<IterVar>& axes) { |
| auto node = ffi::make_object<SLayoutNode>(); |
| node->axes = axes; |
| std::ostringstream repr; |
| |
| for (const IterVar& packed_axis : axes) { |
| auto unpacked_axes = UnpackIterVar(packed_axis); |
| bool is_grouped = unpacked_axes.size() > 1; |
| |
| if (is_grouped) repr << "["; |
| for (const IterVar& axis : unpacked_axes) { |
| if (const auto* factor = axis->dom->extent.as<IntImmNode>()) { |
| TVM_FFI_ICHECK_GT(factor->value, 0); |
| repr << factor->value; |
| } else { |
| TVM_FFI_ICHECK(!is_grouped) |
| << "Only Subordinate Axes with extent is allowed within a packed dim"; |
| } |
| TVM_FFI_ICHECK_EQ(axis->var.get()->name.size(), 1) |
| << "Invalid layout axis " << axis->var.get()->name; |
| char c = axis->var.get()->name.operator std::string()[0]; |
| TVM_FFI_ICHECK((c >= 'a' && c <= 'z') || (c >= 'A' && c <= 'Z')) |
| << "Invalid layout axis " << c; |
| repr << axis->var.get()->name; |
| } |
| if (is_grouped) repr << "]"; |
| } |
| |
| node->name = repr.str(); |
| data_ = std::move(node); |
| } |
| |
| SLayout::SLayout(const std::string& name, PrimType index_ty) { // NOLINT(*) |
| TVM_FFI_CHECK(index_ty.code() == DLDataTypeCode::kDLInt, TypeError) |
| << "The input dtype should be integer type"; |
| if (name == "__undef__") return; |
| |
| auto node = ffi::make_object<SLayoutNode>(); |
| node->name = name; |
| |
| if (name.empty()) return; // scalar |
| |
| // parse layout string |
| int32_t factor = 0; |
| bool in_packing = false; |
| std::vector<IterVar> unpacked_axes; |
| |
| for (char c : name) { |
| if (c >= 'A' && c <= 'Z') { |
| TVM_FFI_ICHECK_EQ(factor, 0) << "Invalid layout " << name << ": invalid factor size " |
| << factor << " before dimension " << c; |
| PrimVar axis_var(std::string(1, c), index_ty); |
| IterVar axis(Range(IntImm(index_ty, 0), axis_var), axis_var, tirx::kDataPar); |
| if (!in_packing) { |
| node->axes.push_back(axis); |
| } else { |
| unpacked_axes.push_back(axis); |
| } |
| } else if (c >= 'a' && c <= 'z') { |
| TVM_FFI_ICHECK_GT(factor, 0) << "Invalid layout " << name << ": invalid factor size " |
| << factor << " for dimension " << c; |
| std::stringstream name; |
| name << factor << c; |
| IterVar axis(Range(IntImm(index_ty, 0), IntImm(index_ty, factor)), |
| PrimVar(name.str(), index_ty), tirx::kDataPar); |
| if (!in_packing) { |
| node->axes.push_back(axis); |
| } else { |
| unpacked_axes.push_back(axis); |
| } |
| factor = 0; |
| } else if (c >= '0' && c <= '9') { |
| TVM_FFI_ICHECK(factor >= 0) << "Invalid layout " << name << ": _ is adjacent to a number."; |
| factor = factor * 10 + c - '0'; |
| } else if (c == '[') { |
| TVM_FFI_ICHECK(!in_packing) << "Invalid layout " << name << ": can't do nested packing"; |
| in_packing = true; |
| } else if (c == ']') { |
| TVM_FFI_ICHECK(in_packing) << "Invalid layout " << name |
| << ": encountered ] without matching bracket"; |
| TVM_FFI_ICHECK(unpacked_axes.size() > 1) |
| << "Invalid layout " << name << ": found empty/single packed axis"; |
| std::stringstream ss; |
| int64_t extent = 1; |
| for (auto& axis : unpacked_axes) { |
| TVM_FFI_ICHECK(axis->dom->extent.as<IntImmNode>()) |
| << "Invalid SLayout " << name << ": can't have variable sized node(" << axis->var->name |
| << ") within a packed axis"; |
| auto axis_name = axis->var->name.operator std::string(); |
| auto factor = axis->dom->extent.as<IntImm>().value(); |
| ss << axis_name; |
| extent = extent * factor->value; |
| } |
| std::string grouped_name = ss.str(); |
| IterVar grouped_axis(Range(IntImm(index_ty, 0), IntImm(index_ty, extent)), |
| PrimVar(grouped_name, index_ty), tirx::kDataPar); |
| node->axes.push_back(grouped_axis); |
| |
| in_packing = false; |
| unpacked_axes.clear(); |
| } else { |
| TVM_FFI_THROW(InternalError) << "Invalid layout " << name; |
| } |
| } |
| TVM_FFI_ICHECK(in_packing == false) |
| << "Invalid SLayout " << name << ": haven't terminated the packing sequence"; |
| |
| // validate layout |
| std::vector<int> axis_cnt(256, 0); |
| for (const IterVar& pv : node->axes) { |
| for (const IterVar& v : UnpackIterVar(pv)) { |
| auto axis_str = v->var.get()->name.operator std::string(); |
| TVM_FFI_ICHECK_EQ(axis_str.size(), 1); |
| char axis = axis_str[0]; |
| TVM_FFI_ICHECK((axis >= 'a' && axis <= 'z') || (axis >= 'A' && axis <= 'Z')); |
| axis_cnt[axis] += 1; |
| } |
| } |
| for (const IterVar& pv : node->axes) { |
| for (const IterVar& v : UnpackIterVar(pv)) { |
| char axis = v->var.get()->name.operator std::string()[0]; |
| if (axis >= 'a' && axis <= 'z') { |
| TVM_FFI_ICHECK(axis_cnt[axis - 'a' + 'A']) |
| << "Invalid layout " << name << ": missing axis " << std::toupper(axis); |
| TVM_FFI_ICHECK(axis_cnt[axis] == 1) |
| << "Invalid layout " << name << ": found more than one subordinate " |
| << std::toupper(axis); |
| } |
| } |
| } |
| |
| data_ = std::move(node); |
| } |
| |
| SLayout SLayout::SubLayout(size_t pos, size_t len) const { |
| if (!defined() || pos > ndim()) return SLayout::Undef(); |
| if (len == 0) return SLayout(ffi::Array<IterVar>()); |
| if (pos + len > ndim()) len = ndim() - pos; |
| ffi::Array<IterVar> new_layout; |
| const auto axes = operator->()->axes; |
| for (size_t i = pos; i < pos + len; ++i) { |
| new_layout.push_back(axes[i]); |
| } |
| return SLayout(new_layout); |
| } |
| |
| ffi::Array<IterVar> SLayout::UnpackIterVar(IterVar packed_iter) { |
| ffi::Array<IterVar> result; |
| int64_t factor = 0, final_factor = 1; |
| |
| std::string name(packed_iter->var->name.c_str()); |
| PrimType index_ty = packed_iter->var.ty(); |
| |
| for (auto ch : name) { |
| if (ch >= '0' && ch <= '9') { |
| factor = factor * 10 + (ch - '0'); |
| } else if (ch >= 'a' && ch <= 'z') { |
| TVM_FFI_ICHECK(factor != 0) << "Invalid Factor Size"; |
| result.push_back(IterVar(Range(IntImm(index_ty, 0), IntImm(index_ty, factor)), |
| PrimVar(std::string(1, ch), index_ty), tirx::kDataPar)); |
| final_factor *= factor; |
| factor = 0; |
| } else if (ch >= 'A' && ch <= 'Z') { |
| TVM_FFI_ICHECK(factor == 0) << "Can't have non-zero factors for primal axis"; |
| PrimVar axis_var(std::string(1, ch), index_ty); |
| result.push_back(IterVar(Range(IntImm(index_ty, 0), axis_var), axis_var, tirx::kDataPar)); |
| } |
| } |
| |
| return result; |
| } |
| |
| IterVar SLayout::PackIterVar(ffi::Array<IterVar> iter_vars) { |
| std::stringstream name; |
| size_t extent = 1; |
| |
| PrimType index_ty = iter_vars[0]->dom->extent.as<PrimExpr>().value().ty(); |
| for (auto itvar : iter_vars) { |
| TVM_FFI_ICHECK(itvar->dom->extent.as<IntImm>()) |
| << "Packed Axis can contain only Subordinate Axes"; |
| name << itvar->dom->extent.as<IntImm>().value() << itvar->var->name; |
| extent = extent * itvar->dom->extent.as<IntImm>().value()->value; |
| } |
| |
| return IterVar(Range(IntImm(index_ty, 0), IntImm(index_ty, extent)), |
| PrimVar(name.str(), index_ty), tirx::kDataPar); |
| } |
| |
| int32_t SLayout::FactorOf(const SLayoutAxis& axis) const { |
| if (!defined()) return -1; |
| const SLayoutAxis& sub = axis.ToSubordinate(); |
| |
| int32_t factor = 1; |
| bool has_sub = false; |
| for (const IterVar& packed_itvar : operator->()->axes) { |
| for (auto itvar : UnpackIterVar(packed_itvar)) { |
| if (sub == SLayoutAxis::Get(itvar)) { |
| has_sub = true; |
| int32_t val = itvar->dom->extent.as<IntImmNode>()->value; |
| factor *= val; |
| } |
| } |
| } |
| factor = has_sub ? factor : -1; |
| |
| return factor; |
| } |
| |
| TVM_FFI_STATIC_INIT_BLOCK() { |
| namespace refl = tvm::ffi::reflection; |
| refl::TypeAttrDef<SLayoutNode>().def(refl::type_attr::kRepr, |
| [](SLayout l, ffi::Function) -> ffi::String { |
| return "SLayout(" + std::string(l->name) + ")"; |
| }); |
| } |
| |
| inline bool GetStoreRule(ffi::Array<PrimExpr>* index_rule, ffi::Array<PrimExpr>* shape_rule, |
| const SLayout& src_layout, const SLayout& dst_layout) { |
| if (!src_layout.defined() || src_layout.name().empty()) { |
| LOG(WARNING) << "src layout '" << src_layout.name() << "' is invalid."; |
| return false; |
| } |
| if (!dst_layout.defined() || dst_layout.name().empty()) { |
| LOG(WARNING) << "dst layout '" << dst_layout.name() << "' is invalid."; |
| return false; |
| } |
| |
| std::vector<bool> exists(128, false); |
| PrimExpr norm_indexes[128]; |
| for (auto& it : norm_indexes) it = PrimExpr(0); |
| |
| for (size_t i = 0; i < src_layout.ndim(); i++) { |
| auto factor = src_layout.PackedAxisAt(i)->dom->extent; |
| auto src_unpacked_axes = SLayout::UnpackIterVar(src_layout.PackedAxisAt(i)); |
| |
| if (src_unpacked_axes.size() == 1 && SLayoutAxis::Get(src_unpacked_axes[0]).IsPrimal()) { |
| const auto& prim_axis = SLayoutAxis::Get(src_unpacked_axes[0]); |
| int64_t offset = src_layout.FactorOf(prim_axis); |
| if (offset == -1) |
| norm_indexes[prim_axis.name()[0] - 'A'] = |
| norm_indexes[prim_axis.name()[0] - 'A'] + src_layout.PackedAxisAt(i); |
| else |
| norm_indexes[prim_axis.name()[0] - 'A'] = |
| norm_indexes[prim_axis.name()[0] - 'A'] + |
| src_layout.PackedAxisAt(i) * src_layout.FactorOf(prim_axis); |
| exists[prim_axis.name()[0]] = true; |
| } else { |
| int64_t value = 1; |
| std::vector<int> index_divs(src_unpacked_axes.size()); |
| for (size_t j = 0; j < src_unpacked_axes.size(); j++) { |
| index_divs[j] = value; |
| const auto* extent = src_unpacked_axes[j]->dom->extent.as<IntImmNode>(); |
| TVM_FFI_ICHECK(extent) << "Expected Integer Extents for Offset Calculation"; |
| index_divs.push_back(value); |
| value = value * extent->value; |
| } |
| std::reverse(index_divs.begin(), index_divs.end()); |
| |
| for (size_t j = 0; j < src_unpacked_axes.size(); j++) { |
| const int extent = src_unpacked_axes[j]->dom->extent.as<IntImmNode>()->value; |
| const SLayoutAxis& store_axis_impl = SLayoutAxis::Get(src_unpacked_axes[j]); |
| const SLayoutAxis& sub_axis = store_axis_impl.ToSubordinate(); /* Not Needed */ |
| const SLayoutAxis& prim_axis = store_axis_impl.ToPrimal(); |
| |
| PrimExpr factor_ij = indexdiv(src_layout.PackedAxisAt(i), index_divs[j]); |
| if (j != 0) factor_ij = indexmod(factor_ij, extent); |
| |
| for (size_t k = i; k < src_layout.ndim(); k++) { |
| size_t l = 0; |
| if (k == i) l = j + 1; |
| |
| auto inter_unpacked_axes = SLayout::UnpackIterVar(src_layout.PackedAxisAt(k)); |
| for (; l < inter_unpacked_axes.size(); l++) { |
| const SLayoutAxis& axis = SLayoutAxis::Get(inter_unpacked_axes[l]); |
| if (axis == sub_axis) { |
| IntImm sub_extent = inter_unpacked_axes[l]->dom->extent.as_or_throw<IntImm>(); |
| factor_ij = factor_ij * IntImm(sub_extent.ty(), sub_extent->value); |
| } |
| } |
| } |
| |
| norm_indexes[prim_axis.name()[0] - 'A'] = |
| norm_indexes[prim_axis.name()[0] - 'A'] + factor_ij; |
| } |
| } |
| } |
| |
| arith::Analyzer ana; |
| |
| for (size_t i = 0; i < dst_layout.ndim(); i++) { |
| const auto dst_unpacked_axes = SLayout::UnpackIterVar(dst_layout.PackedAxisAt(i)); |
| |
| if (dst_unpacked_axes.size() == 1 && SLayoutAxis::Get(dst_unpacked_axes[0]).IsPrimal()) { |
| const auto& prim_axis = SLayoutAxis::Get(dst_unpacked_axes[0]); |
| if (!exists[prim_axis.name()[0]]) return false; |
| int64_t offset = dst_layout.FactorOf(prim_axis); |
| if (offset != -1) { |
| index_rule->push_back( |
| indexdiv(norm_indexes[prim_axis.name()[0] - 'A'], dst_layout.FactorOf(prim_axis))); |
| shape_rule->push_back( |
| indexdiv(norm_indexes[prim_axis.name()[0] - 'A'] + (dst_layout.FactorOf(prim_axis) - 1), |
| dst_layout.FactorOf(prim_axis))); |
| } else { |
| index_rule->push_back(norm_indexes[prim_axis.name()[0] - 'A']); |
| shape_rule->push_back(norm_indexes[prim_axis.name()[0] - 'A']); |
| } |
| } else { |
| PrimExpr factor(0); |
| for (size_t j = 0; j < dst_unpacked_axes.size(); j++) { |
| const auto& prim_axis = SLayoutAxis::Get(dst_unpacked_axes[j]).ToPrimal(); |
| const auto& sub_axis = SLayoutAxis::Get(dst_unpacked_axes[j]).ToSubordinate(); |
| const auto* extent = dst_unpacked_axes[j]->dom->extent.as<IntImmNode>(); |
| TVM_FFI_ICHECK(extent) << "Expected extent to be IntImmNode"; |
| |
| size_t divfactor = 1; |
| for (size_t k = i; k < dst_layout.ndim(); k++) { |
| size_t l = 0; |
| if (k == i) l = j + 1; |
| |
| const auto inter_unpacked_axes = SLayout::UnpackIterVar(dst_layout.PackedAxisAt(k)); |
| for (; l < inter_unpacked_axes.size(); l++) { |
| const auto& axis = SLayoutAxis::Get(inter_unpacked_axes[l]); |
| if (sub_axis == axis) { |
| const auto* sub_extent = inter_unpacked_axes[l]->dom->extent.as<IntImmNode>(); |
| TVM_FFI_ICHECK(sub_extent) << "Expected Integer Extents for Offset Calculation"; |
| divfactor = divfactor * sub_extent->value; |
| } |
| } |
| } |
| |
| factor = factor + indexmod(indexdiv(norm_indexes[prim_axis.name()[0] - 'A'], divfactor), |
| extent->value); |
| for (size_t k = j + 1; k < dst_unpacked_axes.size(); k++) { |
| factor = factor * dst_unpacked_axes[k]->dom->extent.as<IntImm>().value(); |
| } |
| } |
| ana->Simplify(factor); |
| index_rule->push_back(factor); |
| shape_rule->push_back(factor); |
| } |
| } |
| |
| std::stringstream ss; |
| ss << "index rule for " << src_layout.name() << "-->" << dst_layout.name() << ": [ "; |
| for (const auto& r : *index_rule) { |
| ss << r << ", "; |
| } |
| ss << "]" << std::endl; |
| |
| ss << "shape rule for " << src_layout.name() << "-->" << dst_layout.name() << ": [ "; |
| for (const auto& r : *shape_rule) { |
| ss << r << ", "; |
| } |
| ss << "]" << std::endl; |
| VLOG(1) << ss.str() << std::endl; |
| |
| return true; |
| } |
| |
| inline ffi::Array<PrimExpr> TransformIndex(const ffi::Array<PrimExpr>& src_index, |
| const ffi::Array<IterVar>& src_axis, |
| const ffi::Array<PrimExpr>& transform_rule) { |
| arith::Analyzer ana; |
| ffi::Array<PrimExpr> result; |
| std::unordered_map<const tirx::VarNode*, PrimExpr> bind_map; |
| for (size_t i = 0; i < src_index.size(); ++i) { |
| bind_map[src_axis[i]->var.get()] = src_index[i]; |
| } |
| for (PrimExpr rule : transform_rule) { |
| result.push_back(ana->Simplify(tirx::Substitute(rule, bind_map))); |
| } |
| return result; |
| } |
| |
| ffi::Array<PrimExpr> SBijectiveLayout::ForwardIndex(const ffi::Array<PrimExpr>& src_index) const { |
| TVM_FFI_ICHECK(defined()) << "Cannot operate on an undefined bijective layout."; |
| const SBijectiveLayoutNode* self = operator->(); |
| TVM_FFI_ICHECK_EQ(src_index.size(), self->src_layout->axes.size()) |
| << "Input mismatch with layout " << self->src_layout; |
| return TransformIndex(src_index, self->src_layout->axes, self->index_forward_rule); |
| } |
| |
| ffi::Array<PrimExpr> SBijectiveLayout::BackwardIndex(const ffi::Array<PrimExpr>& dst_index) const { |
| TVM_FFI_ICHECK(defined()) << "Cannot operate on an undefined bijective layout."; |
| const SBijectiveLayoutNode* self = operator->(); |
| TVM_FFI_ICHECK_EQ(dst_index.size(), self->dst_layout->axes.size()) |
| << "Output mismatch with layout " << self->dst_layout; |
| return TransformIndex(dst_index, self->dst_layout->axes, self->index_backward_rule); |
| } |
| |
| inline ffi::Array<PrimExpr> TransformShape(const ffi::Array<PrimExpr>& src_shape, |
| const ffi::Array<IterVar>& src_axis, |
| const ffi::Array<IterVar>& target_axis, |
| const ffi::Array<PrimExpr>& transform_rule) { |
| arith::Analyzer ana; |
| TVM_FFI_ICHECK_EQ(src_shape.size(), src_axis.size()) |
| << "Input shape size " << src_shape.size() << " mismatch with the expected shape size " |
| << src_axis.size(); |
| // bind variables for original axes |
| // for major-axis, bind the corresponding size |
| // for minor-axis, simply bind it as 0, so that we can reuse forward/backward_rule, |
| // e.g., (C * 16 + c) / 32 |
| std::unordered_map<const tirx::VarNode*, PrimExpr> bind_map; |
| for (size_t i = 0; i < src_shape.size(); ++i) { |
| PrimExpr orig_shape = src_shape[i]; |
| IterVar orig_axis = src_axis[i]; |
| auto layout = SLayout::UnpackIterVar(orig_axis); |
| if (layout.size() != 1 || !SLayoutAxis::Get(layout[0]).IsPrimal()) { |
| if (orig_shape.defined()) { |
| const auto* orig_shape_const = orig_shape.as<IntImmNode>(); |
| const auto* orig_axis_extent = orig_axis->dom->extent.as<IntImmNode>(); |
| if (orig_shape_const) { |
| TVM_FFI_ICHECK_EQ(orig_shape_const->value, orig_axis_extent->value) |
| << "Input shape mismatch at index " << i << ". Expected " << orig_axis->dom->extent |
| << ", get " << orig_shape; |
| } |
| } |
| bind_map[orig_axis->var.get()] = IntImm(orig_axis->var.ty(), 0); |
| } else { |
| bind_map[orig_axis->var.get()] = orig_axis->var.ty() == orig_shape.ty() |
| ? orig_shape |
| : cast(orig_axis->var.ty(), orig_shape); |
| } |
| } |
| // infer the target shape, |
| // for major-axis, use the forward/backward_rule directly, |
| // for minor-axis, simply use the extent. |
| ffi::Array<PrimExpr> result; |
| TVM_FFI_ICHECK_EQ(transform_rule.size(), target_axis.size()); |
| for (size_t i = 0; i < transform_rule.size(); ++i) { |
| PrimExpr rule = transform_rule[i]; |
| IterVar axis = target_axis[i]; |
| auto layout = SLayout::UnpackIterVar(axis); |
| if (layout.size() != 1 || !SLayoutAxis::Get(layout[0]).IsPrimal()) { |
| result.push_back(axis->dom->extent); |
| } else { |
| result.push_back(ana->Simplify(tirx::Substitute(rule, bind_map))); |
| } |
| } |
| |
| std::stringstream ss; |
| ss << "shape rule for " << SLayout(src_axis).name() << "-->" << SLayout(target_axis).name() |
| << ": [ "; |
| for (const auto& r : transform_rule) { |
| ss << r << ", "; |
| } |
| ss << "]" << std::endl; |
| |
| ss << "shape transform: [ "; |
| for (const auto& s : src_shape) { |
| ss << s << ", "; |
| } |
| ss << "] --> [ "; |
| for (const auto& r : result) { |
| ss << r << ", "; |
| } |
| ss << "]" << std::endl; |
| VLOG(1) << std::endl << ss.str(); |
| |
| return result; |
| } |
| |
| ffi::Array<PrimExpr> SBijectiveLayout::ForwardShape(const ffi::Array<PrimExpr>& shape) const { |
| TVM_FFI_ICHECK(defined()) << "Cannot operate on an undefined bijective layout."; |
| const SBijectiveLayoutNode* self = operator->(); |
| return TransformShape(shape, self->src_layout->axes, self->dst_layout->axes, |
| self->shape_forward_rule); |
| } |
| |
| ffi::Array<PrimExpr> SBijectiveLayout::BackwardShape(const ffi::Array<PrimExpr>& shape) const { |
| TVM_FFI_ICHECK(defined()) << "Cannot operate on an undefined bijective layout."; |
| const SBijectiveLayoutNode* self = operator->(); |
| return TransformShape(shape, self->dst_layout->axes, self->src_layout->axes, |
| self->shape_backward_rule); |
| } |
| |
| SBijectiveLayout::SBijectiveLayout(SLayout src_layout, SLayout dst_layout) { |
| auto n = ffi::make_object<SBijectiveLayoutNode>(); |
| |
| n->src_layout = std::move(src_layout); |
| n->dst_layout = std::move(dst_layout); |
| // To be consistent with previous behavior, a nullptr layout is created |
| // when argument is invalid. |
| if (GetStoreRule(&n->index_forward_rule, &n->shape_forward_rule, n->src_layout, n->dst_layout)) { |
| TVM_FFI_ICHECK(GetStoreRule(&n->index_backward_rule, &n->shape_backward_rule, n->dst_layout, |
| n->src_layout)); |
| data_ = std::move(n); |
| } |
| } |
| |
| TVM_FFI_STATIC_INIT_BLOCK() { |
| namespace refl = tvm::ffi::reflection; |
| refl::TypeAttrDef<SBijectiveLayoutNode>().def( |
| refl::type_attr::kRepr, [](SBijectiveLayout bl, ffi::Function) -> ffi::String { |
| return "SBijectiveLayout(" + std::string(bl->src_layout.name()) + "->" + |
| std::string(bl->dst_layout.name()) + ")"; |
| }); |
| } |
| |
| TVM_FFI_STATIC_INIT_BLOCK() { |
| namespace refl = tvm::ffi::reflection; |
| refl::GlobalDef() |
| .def("s_tir.SLayout", [](std::string name, PrimType dtype) { return SLayout(name, dtype); }) |
| .def("s_tir.SLayoutIndexOf", |
| [](SLayout layout, std::string axis) -> int { return layout.IndexOf(axis); }) |
| .def("s_tir.SLayoutFactorOf", |
| [](SLayout layout, std::string axis) -> int { |
| return layout.FactorOf(SLayoutAxis::Get(axis)); |
| }) |
| .def("s_tir.SLayoutNdim", [](SLayout layout) -> int { return layout.ndim(); }) |
| .def("s_tir.SLayoutGetItem", |
| [](SLayout layout, int idx) -> std::string { |
| const auto& axis = layout.PackedAxisAt(idx); |
| return axis->var->name; |
| }) |
| .def("s_tir.SBijectiveLayout", |
| [](SLayout src_layout, SLayout dst_layout) -> SBijectiveLayout { |
| return SBijectiveLayout(src_layout, dst_layout); |
| }) |
| .def_method("s_tir.SBijectiveLayoutForwardIndex", &SBijectiveLayout::ForwardIndex) |
| .def_method("s_tir.SBijectiveLayoutBackwardIndex", &SBijectiveLayout::BackwardIndex) |
| .def_method("s_tir.SBijectiveLayoutForwardShape", &SBijectiveLayout::ForwardShape) |
| .def_method("s_tir.SBijectiveLayoutBackwardShape", &SBijectiveLayout::BackwardShape); |
| } |
| } // namespace tirx |
| } // namespace tvm |
