wehub-resource-sync
138 lines
5.3 KiB
Python
138 lines
5.3 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.
|
|
# pylint: disable=invalid-name
|
|
"""Legalization functions for DLTensor inspection."""
|
|
|
|
import enum
|
|
|
|
from tvm.ir import Call
|
|
from tvm.script import tirx as T
|
|
|
|
from ... import op
|
|
from ...block_builder import BlockBuilder
|
|
from ...expr import Expr
|
|
from .common import register_legalize
|
|
|
|
|
|
class TVMStructFieldKind(enum.IntEnum):
|
|
"""Equivalent to tvm::tirx::builtin::TVMStructFieldKind
|
|
|
|
This does not use `enum.auto()` to define the values, because
|
|
`enum.auto()` starts from 1, and this must match the C++
|
|
definition which starts from 0.
|
|
"""
|
|
|
|
kDLTensorAddr = 0
|
|
kDLTensorData = 1
|
|
kDLTensorShape = 2
|
|
kDLTensorStrides = 3
|
|
kDLTensorNDim = 4
|
|
kDLTensorTypeCode = 5
|
|
kDLTensorTypeBits = 6
|
|
kDLTensorTypeLanes = 7
|
|
kDLTensorByteOffset = 8
|
|
kDLTensorDeviceId = 9
|
|
kDLTensorDeviceType = 10
|
|
kDLTensorKindBound_ = 11
|
|
kTVMValueContent = 12
|
|
kTVMValueKindBound_ = 13
|
|
|
|
|
|
@register_legalize("relax.inspect.tensor_stride_i")
|
|
def _tensor_stride_i(bb: BlockBuilder, call: Call) -> Expr:
|
|
@T.prim_func(private=True, s_tir=True)
|
|
def _get_tensor_stride_i(dlpack_handle: T.handle, axis: T.int64) -> T.int64:
|
|
T.func_attr({"tirx.is_host_func": True, "tirx.is_scheduled": True})
|
|
assert T.int64(0) <= axis, "Specified axis may not be negative"
|
|
ndim: T.let[T.int32] = T.tvm_struct_get(
|
|
dlpack_handle, 0, int(TVMStructFieldKind.kDLTensorNDim), "int32"
|
|
)
|
|
assert axis < T.Cast("int64", ndim), (
|
|
"Specified axis may not be larger than the tensor's dimensionality"
|
|
)
|
|
stride_ptr: T.let[T.handle("int64")] = T.tvm_struct_get(
|
|
dlpack_handle, 0, int(TVMStructFieldKind.kDLTensorStrides), T.handle("int64").ty
|
|
)
|
|
|
|
if T.isnullptr(stride_ptr):
|
|
shape_ptr: T.let[T.handle("int64")] = T.tvm_struct_get(
|
|
dlpack_handle, 0, int(TVMStructFieldKind.kDLTensorShape), T.handle("int64").ty
|
|
)
|
|
shape = T.decl_buffer(ndim, "int64", data=shape_ptr)
|
|
|
|
product = T.decl_buffer([], "int64")
|
|
product[()] = 1
|
|
|
|
# TODO(Lunderberg): Add a TIR lowering pass to allow
|
|
# ranges to start somewhere other than zero. This loop
|
|
# could then iterate on `range(axis+1, ndim)`.
|
|
for dim_offset in range(ndim - (axis + 1)):
|
|
dim: T.let[T.int64] = dim_offset + (axis + 1)
|
|
product[()] = product[()] * shape[dim]
|
|
|
|
return product[()]
|
|
else:
|
|
strides = T.decl_buffer(ndim, "int64", data=stride_ptr)
|
|
stride: T.let[T.int64] = strides[axis]
|
|
return stride
|
|
|
|
gvar = bb.add_func(_get_tensor_stride_i, "_get_tensor_stride_i")
|
|
return Call(gvar, call.args)
|
|
|
|
|
|
@register_legalize("relax.inspect.tensor_byte_offset")
|
|
def _tensor_byte_offset(bb: BlockBuilder, call: Call) -> Expr:
|
|
@T.prim_func(private=True, s_tir=True)
|
|
def _get_tensor_byte_offset(dlpack_handle: T.handle) -> T.int64:
|
|
T.func_attr({"tirx.is_host_func": True, "tirx.is_scheduled": True})
|
|
byte_offset: T.let[T.uint64] = T.tvm_struct_get(
|
|
dlpack_handle, 0, int(TVMStructFieldKind.kDLTensorByteOffset), "uint64"
|
|
)
|
|
return byte_offset
|
|
|
|
gvar = bb.add_func(_get_tensor_byte_offset, "_get_tensor_byte_offset")
|
|
return Call(gvar, call.args)
|
|
|
|
|
|
@register_legalize("relax.inspect.tensor_elem_offset")
|
|
def _tensor_elem_offset(bb: BlockBuilder, call: Call) -> Expr:
|
|
@T.prim_func(private=True, s_tir=True)
|
|
def _get_tensor_elem_offset(dlpack_handle: T.handle) -> T.int64:
|
|
T.func_attr({"tirx.is_host_func": True, "tirx.is_scheduled": True})
|
|
byte_offset: T.let[T.uint64] = T.tvm_struct_get(
|
|
dlpack_handle, 0, int(TVMStructFieldKind.kDLTensorByteOffset), "uint64"
|
|
)
|
|
scalar_bits: T.let[T.uint8] = T.tvm_struct_get(
|
|
dlpack_handle, 0, int(TVMStructFieldKind.kDLTensorTypeBits), "uint8"
|
|
)
|
|
lanes: T.let[T.uint16] = T.tvm_struct_get(
|
|
dlpack_handle, 0, int(TVMStructFieldKind.kDLTensorTypeLanes), "uint16"
|
|
)
|
|
bytes_per_element: T.let[T.uint64] = T.ceildiv(
|
|
scalar_bits.astype("uint64") * lanes.astype("uint64"), 8
|
|
)
|
|
elem_offset: T.let[T.uint64] = byte_offset // bytes_per_element
|
|
return elem_offset
|
|
|
|
gvar = bb.add_func(_get_tensor_elem_offset, "_get_tensor_elem_offset")
|
|
return Call(gvar, call.args)
|
|
|
|
|
|
@register_legalize("relax.size")
|
|
def _size(_bb: BlockBuilder, call: Call) -> Expr:
|
|
return op.prod(op.shape_to_tensor(op.shape_of(call.args[0])))
|