wehub-resource-sync
672 lines
20 KiB
Python
672 lines
20 KiB
Python
# 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.
|
|
"""
|
|
TIR expression functors in Python.
|
|
|
|
This module implements the visitor and mutator patterns for TIR expressions.
|
|
"""
|
|
|
|
from collections.abc import Callable
|
|
from typing import TypeVar
|
|
|
|
import tvm
|
|
from tvm.ir import Expr, Range
|
|
from tvm.tirx import IterVar
|
|
|
|
T = TypeVar("T")
|
|
|
|
|
|
def _visit_array(arr: list[T], callback: Callable[[T], None]) -> None:
|
|
"""Visit elements in an array using a callback function.
|
|
|
|
Parameters
|
|
----------
|
|
arr : List[T]
|
|
The array to be visited
|
|
callback : Callable[[T], None]
|
|
The callback function
|
|
"""
|
|
for item in arr:
|
|
callback(item)
|
|
|
|
|
|
class ExprFunctor:
|
|
"""An abstract visitor over Expr, with visiting function defined for each Expr type."""
|
|
|
|
def __init__(self):
|
|
self._dispatch_map = {
|
|
"tirx.Var": self.visit_var_,
|
|
"tirx.BufferLoad": self.visit_buffer_load_,
|
|
"tirx.ProducerLoad": self.visit_producer_load_,
|
|
"tirx.Let": self.visit_let_,
|
|
"tirx.Call": self.visit_call_,
|
|
"tirx.Add": self.visit_add_,
|
|
"tirx.Sub": self.visit_sub_,
|
|
"tirx.Mul": self.visit_mul_,
|
|
"tirx.Div": self.visit_div_,
|
|
"tirx.Mod": self.visit_mod_,
|
|
"tirx.FloorDiv": self.visit_floordiv_,
|
|
"tirx.FloorMod": self.visit_floormod_,
|
|
"tirx.Min": self.visit_min_,
|
|
"tirx.Max": self.visit_max_,
|
|
"tirx.EQ": self.visit_eq_,
|
|
"tirx.NE": self.visit_ne_,
|
|
"tirx.LT": self.visit_lt_,
|
|
"tirx.LE": self.visit_le_,
|
|
"tirx.GT": self.visit_gt_,
|
|
"tirx.GE": self.visit_ge_,
|
|
"tirx.And": self.visit_and_,
|
|
"tirx.Or": self.visit_or_,
|
|
"tirx.Reduce": self.visit_reduce_,
|
|
"tirx.Cast": self.visit_cast_,
|
|
"tirx.Not": self.visit_not_,
|
|
"tirx.Select": self.visit_select_,
|
|
"tirx.Ramp": self.visit_ramp_,
|
|
"tirx.Broadcast": self.visit_broadcast_,
|
|
"tirx.Shuffle": self.visit_shuffle_,
|
|
"tirx.IntImm": self.visit_int_imm_,
|
|
"tirx.FloatImm": self.visit_float_imm_,
|
|
"tirx.StringImm": self.visit_string_imm_,
|
|
}
|
|
|
|
def visit_expr(self, expr: Expr):
|
|
"""Apply the visitor to an expression.
|
|
|
|
Parameters
|
|
----------
|
|
expr : Expr
|
|
The expression to be visited.
|
|
|
|
Returns
|
|
-------
|
|
result : Any
|
|
The result of the visit.
|
|
"""
|
|
if expr is None:
|
|
return None
|
|
|
|
key = expr.__class__.__name__
|
|
if key.endswith("Node"):
|
|
key = key[:-4]
|
|
|
|
key = "tirx." + key
|
|
if key in self._dispatch_map:
|
|
return self._dispatch_map[key](expr)
|
|
|
|
return self.visit_expr_default_(expr)
|
|
|
|
def visit_var_(self, op):
|
|
"""Default visitor for Var node."""
|
|
return None
|
|
|
|
def visit_buffer_load_(self, op):
|
|
"""Default visitor for BufferLoad node."""
|
|
return self.visit_expr_default_(op)
|
|
|
|
def visit_producer_load_(self, op):
|
|
"""Default visitor for ProducerLoad node."""
|
|
return self.visit_expr_default_(op)
|
|
|
|
def visit_let_(self, op):
|
|
"""Default visitor for Let node."""
|
|
return self.visit_expr_default_(op)
|
|
|
|
def visit_call_(self, op):
|
|
"""Default visitor for Call node."""
|
|
return self.visit_expr_default_(op)
|
|
|
|
def visit_add_(self, op):
|
|
"""Default visitor for Add node."""
|
|
return self.visit_expr_default_(op)
|
|
|
|
def visit_sub_(self, op):
|
|
"""Default visitor for Sub node."""
|
|
return self.visit_expr_default_(op)
|
|
|
|
def visit_mul_(self, op):
|
|
"""Default visitor for Mul node."""
|
|
return self.visit_expr_default_(op)
|
|
|
|
def visit_div_(self, op):
|
|
"""Default visitor for Div node."""
|
|
return self.visit_expr_default_(op)
|
|
|
|
def visit_mod_(self, op):
|
|
"""Default visitor for Mod node."""
|
|
return self.visit_expr_default_(op)
|
|
|
|
def visit_floordiv_(self, op):
|
|
"""Default visitor for FloorDiv node."""
|
|
return self.visit_expr_default_(op)
|
|
|
|
def visit_floormod_(self, op):
|
|
"""Default visitor for FloorMod node."""
|
|
return self.visit_expr_default_(op)
|
|
|
|
def visit_min_(self, op):
|
|
"""Default visitor for Min node."""
|
|
return self.visit_expr_default_(op)
|
|
|
|
def visit_max_(self, op):
|
|
"""Default visitor for Max node."""
|
|
return self.visit_expr_default_(op)
|
|
|
|
def visit_eq_(self, op):
|
|
"""Default visitor for EQ node."""
|
|
return self.visit_expr_default_(op)
|
|
|
|
def visit_ne_(self, op):
|
|
"""Default visitor for NE node."""
|
|
return self.visit_expr_default_(op)
|
|
|
|
def visit_lt_(self, op):
|
|
"""Default visitor for LT node."""
|
|
return self.visit_expr_default_(op)
|
|
|
|
def visit_le_(self, op):
|
|
"""Default visitor for LE node."""
|
|
return self.visit_expr_default_(op)
|
|
|
|
def visit_gt_(self, op):
|
|
"""Default visitor for GT node."""
|
|
return self.visit_expr_default_(op)
|
|
|
|
def visit_ge_(self, op):
|
|
"""Default visitor for GE node."""
|
|
return self.visit_expr_default_(op)
|
|
|
|
def visit_and_(self, op):
|
|
"""Default visitor for And node."""
|
|
return self.visit_expr_default_(op)
|
|
|
|
def visit_or_(self, op):
|
|
"""Default visitor for Or node."""
|
|
return self.visit_expr_default_(op)
|
|
|
|
def visit_reduce_(self, op):
|
|
"""Default visitor for Reduce node."""
|
|
return self.visit_expr_default_(op)
|
|
|
|
def visit_cast_(self, op):
|
|
"""Default visitor for Cast node."""
|
|
return self.visit_expr_default_(op)
|
|
|
|
def visit_not_(self, op):
|
|
"""Default visitor for Not node."""
|
|
return self.visit_expr_default_(op)
|
|
|
|
def visit_select_(self, op):
|
|
"""Default visitor for Select node."""
|
|
return self.visit_expr_default_(op)
|
|
|
|
def visit_ramp_(self, op):
|
|
"""Default visitor for Ramp node."""
|
|
return self.visit_expr_default_(op)
|
|
|
|
def visit_broadcast_(self, op):
|
|
"""Default visitor for Broadcast node."""
|
|
return self.visit_expr_default_(op)
|
|
|
|
def visit_shuffle_(self, op):
|
|
"""Default visitor for Shuffle node."""
|
|
return self.visit_expr_default_(op)
|
|
|
|
def visit_int_imm_(self, op):
|
|
"""Default visitor for IntImm node."""
|
|
return self.visit_expr_default_(op)
|
|
|
|
def visit_float_imm_(self, op):
|
|
"""Default visitor for FloatImm node."""
|
|
return self.visit_expr_default_(op)
|
|
|
|
def visit_string_imm_(self, op):
|
|
"""Default visitor for StringImm node."""
|
|
return self.visit_expr_default_(op)
|
|
|
|
def visit_expr_default_(self, op):
|
|
"""Default visitor implementation."""
|
|
raise NotImplementedError(f"Do not have a default for {op.__class__.__name__}")
|
|
|
|
def __call__(self, expr):
|
|
"""Call visitor on expression.
|
|
|
|
Parameters
|
|
----------
|
|
expr : Expr
|
|
The expression.
|
|
|
|
Returns
|
|
-------
|
|
result : Any
|
|
The result of visiting.
|
|
"""
|
|
return self.visit_expr(expr)
|
|
|
|
|
|
class ExprVisitor(ExprFunctor):
|
|
"""A visitor over Expr.
|
|
|
|
This is a visitor that recursively traverses an expression. Subclasses can
|
|
override the visit methods to customize the behavior.
|
|
"""
|
|
|
|
def visit_var_(self, op):
|
|
"""Visitor implementation for Var."""
|
|
pass
|
|
|
|
def visit_buffer_load_(self, op):
|
|
"""Visitor implementation for BufferLoad."""
|
|
|
|
def _visit_indices(index):
|
|
self.visit_expr(index)
|
|
|
|
_visit_array(op.indices, _visit_indices)
|
|
|
|
def visit_producer_load_(self, op):
|
|
"""Visitor implementation for ProducerLoad."""
|
|
|
|
def _visit_indices(index):
|
|
self.visit_expr(index)
|
|
|
|
_visit_array(op.indices, _visit_indices)
|
|
|
|
def visit_let_(self, op):
|
|
"""Visitor implementation for Let."""
|
|
self.visit_expr(op.value)
|
|
self.visit_expr(op.body)
|
|
|
|
def visit_call_(self, op):
|
|
"""Visitor implementation for Call."""
|
|
|
|
def _visit_arg(arg):
|
|
self.visit_expr(arg)
|
|
|
|
_visit_array(op.args, _visit_arg)
|
|
|
|
def _visit_binary_op(self, op):
|
|
"""Helper to visit binary operators."""
|
|
self.visit_expr(op.a)
|
|
self.visit_expr(op.b)
|
|
|
|
def visit_add_(self, op):
|
|
"""Visitor implementation for Add."""
|
|
self._visit_binary_op(op)
|
|
|
|
def visit_sub_(self, op):
|
|
"""Visitor implementation for Sub."""
|
|
self._visit_binary_op(op)
|
|
|
|
def visit_mul_(self, op):
|
|
"""Visitor implementation for Mul."""
|
|
self._visit_binary_op(op)
|
|
|
|
def visit_div_(self, op):
|
|
"""Visitor implementation for Div."""
|
|
self._visit_binary_op(op)
|
|
|
|
def visit_mod_(self, op):
|
|
"""Visitor implementation for Mod."""
|
|
self._visit_binary_op(op)
|
|
|
|
def visit_floordiv_(self, op):
|
|
"""Visitor implementation for FloorDiv."""
|
|
self._visit_binary_op(op)
|
|
|
|
def visit_floormod_(self, op):
|
|
"""Visitor implementation for FloorMod."""
|
|
self._visit_binary_op(op)
|
|
|
|
def visit_min_(self, op):
|
|
"""Visitor implementation for Min."""
|
|
self._visit_binary_op(op)
|
|
|
|
def visit_max_(self, op):
|
|
"""Visitor implementation for Max."""
|
|
self._visit_binary_op(op)
|
|
|
|
def visit_eq_(self, op):
|
|
"""Visitor implementation for EQ."""
|
|
self._visit_binary_op(op)
|
|
|
|
def visit_ne_(self, op):
|
|
"""Visitor implementation for NE."""
|
|
self._visit_binary_op(op)
|
|
|
|
def visit_lt_(self, op):
|
|
"""Visitor implementation for LT."""
|
|
self._visit_binary_op(op)
|
|
|
|
def visit_le_(self, op):
|
|
"""Visitor implementation for LE."""
|
|
self._visit_binary_op(op)
|
|
|
|
def visit_gt_(self, op):
|
|
"""Visitor implementation for GT."""
|
|
self._visit_binary_op(op)
|
|
|
|
def visit_ge_(self, op):
|
|
"""Visitor implementation for GE."""
|
|
self._visit_binary_op(op)
|
|
|
|
def visit_and_(self, op):
|
|
"""Visitor implementation for And."""
|
|
self._visit_binary_op(op)
|
|
|
|
def visit_or_(self, op):
|
|
"""Visitor implementation for Or."""
|
|
self._visit_binary_op(op)
|
|
|
|
def visit_int_imm_(self, op):
|
|
"""Visitor implementation for IntImm."""
|
|
pass
|
|
|
|
def visit_float_imm_(self, op):
|
|
"""Visitor implementation for FloatImm."""
|
|
pass
|
|
|
|
def visit_string_imm_(self, op):
|
|
"""Visitor implementation for StringImm."""
|
|
pass
|
|
|
|
def visit_reduce_(self, op):
|
|
"""Visitor implementation for Reduce."""
|
|
|
|
def _visit_iter_var(iv):
|
|
self.visit_expr(iv.dom.min)
|
|
self.visit_expr(iv.dom.extent)
|
|
|
|
def _visit_source(source):
|
|
self.visit_expr(source)
|
|
|
|
_visit_array(op.axis, _visit_iter_var)
|
|
_visit_array(op.source, _visit_source)
|
|
|
|
if op.init:
|
|
_visit_array(op.init, _visit_source)
|
|
|
|
self.visit_expr(op.condition)
|
|
|
|
def visit_cast_(self, op):
|
|
"""Visitor implementation for Cast."""
|
|
self.visit_expr(op.value)
|
|
|
|
def visit_not_(self, op):
|
|
"""Visitor implementation for Not."""
|
|
self.visit_expr(op.a)
|
|
|
|
def visit_select_(self, op):
|
|
"""Visitor implementation for Select."""
|
|
self.visit_expr(op.condition)
|
|
self.visit_expr(op.true_value)
|
|
self.visit_expr(op.false_value)
|
|
|
|
def visit_ramp_(self, op):
|
|
"""Visitor implementation for Ramp."""
|
|
self.visit_expr(op.base)
|
|
self.visit_expr(op.stride)
|
|
self.visit_expr(op.lanes)
|
|
|
|
def visit_shuffle_(self, op):
|
|
"""Visitor implementation for Shuffle."""
|
|
|
|
def _visit_expr(expr):
|
|
self.visit_expr(expr)
|
|
|
|
_visit_array(op.indices, _visit_expr)
|
|
_visit_array(op.vectors, _visit_expr)
|
|
|
|
def visit_broadcast_(self, op):
|
|
"""Visitor implementation for Broadcast."""
|
|
self.visit_expr(op.value)
|
|
self.visit_expr(op.lanes)
|
|
|
|
|
|
class ExprMutator(ExprFunctor):
|
|
"""A mutator over Expr.
|
|
|
|
This is a mutator that recursively transforms an expression. Subclasses can
|
|
override the visit methods to customize the behavior.
|
|
"""
|
|
|
|
def visit_var_(self, op):
|
|
"""Mutator implementation for Var."""
|
|
return op
|
|
|
|
def visit_buffer_load_(self, op):
|
|
"""Mutator implementation for BufferLoad."""
|
|
indices = [self.visit_expr(index) for index in op.indices]
|
|
|
|
if all(old_index is new_index for old_index, new_index in zip(op.indices, indices)):
|
|
return op
|
|
else:
|
|
return tvm.tirx.BufferLoad(op.buffer, indices, op.predicate)
|
|
|
|
def visit_producer_load_(self, op):
|
|
"""Mutator implementation for ProducerLoad."""
|
|
indices = [self.visit_expr(index) for index in op.indices]
|
|
|
|
if all(old_index is new_index for old_index, new_index in zip(op.indices, indices)):
|
|
return op
|
|
else:
|
|
return tvm.tirx.ProducerLoad(op.producer, indices)
|
|
|
|
def visit_let_(self, op):
|
|
"""Mutator implementation for Let."""
|
|
var = self.visit_var_(op.var)
|
|
value = self.visit_expr(op.value)
|
|
body = self.visit_expr(op.body)
|
|
|
|
if var is op.var and value is op.value and body is op.body:
|
|
return op
|
|
else:
|
|
return tvm.tirx.Let(var, value, body)
|
|
|
|
def visit_call_(self, op):
|
|
"""Mutator implementation for Call."""
|
|
args = [self.visit_expr(arg) for arg in op.args]
|
|
|
|
if all(old_arg is new_arg for old_arg, new_arg in zip(op.args, args)):
|
|
return op
|
|
else:
|
|
return tvm.ir.Call(op.op, args, attrs=op.attrs, span=op.span, ret_ty=op.ty)
|
|
|
|
def _mutate_binary_op(self, op_cls, op):
|
|
"""Helper to mutate binary operators."""
|
|
a = self.visit_expr(op.a)
|
|
b = self.visit_expr(op.b)
|
|
|
|
if a is op.a and b is op.b:
|
|
return op
|
|
else:
|
|
return op_cls(a, b)
|
|
|
|
def visit_add_(self, op):
|
|
"""Mutator implementation for Add."""
|
|
return self._mutate_binary_op(tvm.tirx.Add, op)
|
|
|
|
def visit_sub_(self, op):
|
|
"""Mutator implementation for Sub."""
|
|
return self._mutate_binary_op(tvm.tirx.Sub, op)
|
|
|
|
def visit_mul_(self, op):
|
|
"""Mutator implementation for Mul."""
|
|
return self._mutate_binary_op(tvm.tirx.Mul, op)
|
|
|
|
def visit_div_(self, op):
|
|
"""Mutator implementation for Div."""
|
|
return self._mutate_binary_op(tvm.tirx.Div, op)
|
|
|
|
def visit_mod_(self, op):
|
|
"""Mutator implementation for Mod."""
|
|
return self._mutate_binary_op(tvm.tirx.Mod, op)
|
|
|
|
def visit_floordiv_(self, op):
|
|
"""Mutator implementation for FloorDiv."""
|
|
return self._mutate_binary_op(tvm.tirx.FloorDiv, op)
|
|
|
|
def visit_floormod_(self, op):
|
|
"""Mutator implementation for FloorMod."""
|
|
return self._mutate_binary_op(tvm.tirx.FloorMod, op)
|
|
|
|
def visit_min_(self, op):
|
|
"""Mutator implementation for Min."""
|
|
return self._mutate_binary_op(tvm.tirx.Min, op)
|
|
|
|
def visit_max_(self, op):
|
|
"""Mutator implementation for Max."""
|
|
return self._mutate_binary_op(tvm.tirx.Max, op)
|
|
|
|
def visit_eq_(self, op):
|
|
"""Mutator implementation for EQ."""
|
|
return self._mutate_binary_op(tvm.tirx.EQ, op)
|
|
|
|
def visit_ne_(self, op):
|
|
"""Mutator implementation for NE."""
|
|
return self._mutate_binary_op(tvm.tirx.NE, op)
|
|
|
|
def visit_lt_(self, op):
|
|
"""Mutator implementation for LT."""
|
|
return self._mutate_binary_op(tvm.tirx.LT, op)
|
|
|
|
def visit_le_(self, op):
|
|
"""Mutator implementation for LE."""
|
|
return self._mutate_binary_op(tvm.tirx.LE, op)
|
|
|
|
def visit_gt_(self, op):
|
|
"""Mutator implementation for GT."""
|
|
return self._mutate_binary_op(tvm.tirx.GT, op)
|
|
|
|
def visit_ge_(self, op):
|
|
"""Mutator implementation for GE."""
|
|
return self._mutate_binary_op(tvm.tirx.GE, op)
|
|
|
|
def visit_and_(self, op):
|
|
"""Mutator implementation for And."""
|
|
return self._mutate_binary_op(tvm.tirx.And, op)
|
|
|
|
def visit_or_(self, op):
|
|
"""Mutator implementation for Or."""
|
|
return self._mutate_binary_op(tvm.tirx.Or, op)
|
|
|
|
def visit_int_imm_(self, op):
|
|
"""Mutator implementation for IntImm."""
|
|
return op
|
|
|
|
def visit_float_imm_(self, op):
|
|
"""Mutator implementation for FloatImm."""
|
|
return op
|
|
|
|
def visit_string_imm_(self, op):
|
|
"""Mutator implementation for StringImm."""
|
|
return op
|
|
|
|
def visit_reduce_(self, op):
|
|
"""Mutator implementation for Reduce."""
|
|
|
|
def _mutate_iter_var(iv):
|
|
old_dom = iv.dom
|
|
new_min = self.visit_expr(old_dom.min)
|
|
new_extent = self.visit_expr(old_dom.extent)
|
|
|
|
if new_min is old_dom.min and new_extent is old_dom.extent:
|
|
return iv
|
|
else:
|
|
new_dom = Range.FromMinExtent(new_min, new_extent)
|
|
return IterVar(new_dom, iv.var, iv.iter_type, iv.thread_tag)
|
|
|
|
axis = [_mutate_iter_var(iv) for iv in op.axis]
|
|
source = [self.visit_expr(e) for e in op.source]
|
|
init = [self.visit_expr(e) for e in op.init] if op.init else []
|
|
condition = self.visit_expr(op.condition)
|
|
|
|
axis_unchanged = all(old_iv is new_iv for old_iv, new_iv in zip(op.axis, axis))
|
|
source_unchanged = all(old_e is new_e for old_e, new_e in zip(op.source, source))
|
|
init_unchanged = (
|
|
True if not op.init else all(old_e is new_e for old_e, new_e in zip(op.init, init))
|
|
)
|
|
condition_unchanged = condition is op.condition
|
|
|
|
if axis_unchanged and source_unchanged and init_unchanged and condition_unchanged:
|
|
return op
|
|
else:
|
|
return tvm.tirx.Reduce(op.combiner, source, axis, condition, op.value_index, init)
|
|
|
|
def visit_cast_(self, op):
|
|
"""Mutator implementation for Cast."""
|
|
value = self.visit_expr(op.value)
|
|
|
|
if value is op.value:
|
|
return op
|
|
else:
|
|
return tvm.tirx.Cast(op.ty, value)
|
|
|
|
def visit_not_(self, op):
|
|
"""Mutator implementation for Not."""
|
|
a = self.visit_expr(op.a)
|
|
|
|
if a is op.a:
|
|
return op
|
|
else:
|
|
return tvm.tirx.Not(a)
|
|
|
|
def visit_select_(self, op):
|
|
"""Mutator implementation for Select."""
|
|
condition = self.visit_expr(op.condition)
|
|
true_value = self.visit_expr(op.true_value)
|
|
false_value = self.visit_expr(op.false_value)
|
|
|
|
if (
|
|
condition is op.condition
|
|
and true_value is op.true_value
|
|
and false_value is op.false_value
|
|
):
|
|
return op
|
|
else:
|
|
return tvm.tirx.Select(condition, true_value, false_value)
|
|
|
|
def visit_ramp_(self, op):
|
|
"""Mutator implementation for Ramp."""
|
|
base = self.visit_expr(op.base)
|
|
stride = self.visit_expr(op.stride)
|
|
lanes = self.visit_expr(op.lanes)
|
|
|
|
if base is op.base and stride is op.stride and lanes is op.lanes:
|
|
return op
|
|
else:
|
|
return tvm.tirx.Ramp(base, stride, lanes)
|
|
|
|
def visit_broadcast_(self, op):
|
|
"""Mutator implementation for Broadcast."""
|
|
value = self.visit_expr(op.value)
|
|
lanes = self.visit_expr(op.lanes)
|
|
|
|
if value is op.value and lanes is op.lanes:
|
|
return op
|
|
else:
|
|
return tvm.tirx.Broadcast(value, lanes)
|
|
|
|
def visit_shuffle_(self, op):
|
|
"""Mutator implementation for Shuffle."""
|
|
vectors = [self.visit_expr(v) for v in op.vectors]
|
|
|
|
vectors_unchanged = all(old_v is new_v for old_v, new_v in zip(op.vectors, vectors))
|
|
|
|
if vectors_unchanged:
|
|
return op
|
|
else:
|
|
return tvm.tirx.Shuffle(vectors, op.indices)
|