def within_index(b, e, s, i): """Return a boolean value that indicates if i is within the given index. Parameter --------- b : Expr beginning of the index e : Expr end of the index s : Expr strides of index i : Expr array position Returns ------- selected: Expr bool expression that is True is the array position would be selected by the index and False otherwise """ bc = tvm.expr.Select(s < 0, i <= e, i < b) ec = tvm.expr.Select(s < 0, i > b, i >= e) ss = tvm.if_then_else(s < 0, ((i - e) + (e % tvm.abs(s)) + 1) % tvm.abs(s), (i - b) % s) return tvm.expr.Select(tvm.expr.Or(bc, ec), tvm.const(False), ss.equal(0))
def abs(x): """Take absolute value of the input of x, element-wise. Parameters ---------- x : tvm.Tensor Input argument. Returns ------- y : tvm.Tensor The result. """ return tvm.compute(x.shape, lambda *i: tvm.abs(x(*i)))
def make_idx(b, e, s, z, i): """Return the array position in the selection that corresponds to an array position in the full array. The returned value is only meaningful if within_index() returns True for the same set of parameters. Parameter --------- b : Expr beginning of the index e : Expr end of the index s : Expr strides of index z : Expr size of the indexed dimension i : Expr array position Returns ------- postion: Expr int expression that corresponds to an array position in the selection. """ bc = tvm.expr.Select(s < 0, i <= e, i < b) ec = tvm.expr.Select(s < 0, i > b, i >= e) # Clamp to array size b = tvm.expr.Select(z < b, z - 1, b) ss = tvm.if_then_else(s < 0, (b - i) // tvm.abs(s), (i - b) // s) return tvm.if_then_else(tvm.expr.Or(bc, ec), 88, ss)