in src/tir/transforms/arg_binder.cc [152:325]
void ArgBinder::BindDLTensor(const Buffer& buffer, const PrimExpr& device_type,
const PrimExpr& device_id, const Var& handle,
const std::string& arg_name) {
const DataType tvm_shape_type = DataType::ShapeIndex();
const DataType tvm_ndim_type = DataType::Int(32);
const Stmt nop = Evaluate(0);
init_nest_.emplace_back(AssertStmt(
!Call(DataType::Bool(), builtin::isnullptr(), {handle}),
tvm::tir::StringImm(arg_name + " is expected to have non-NULL DLTensor* pointer"), nop));
// dimension checks
PrimExpr v_ndim = TVMArrayGet(tvm_ndim_type, handle, builtin::kArrNDim);
// Helper functions for shape/stride name formatting
auto shape_handle_name = [&]() { return arg_name + ".shape"; };
auto stride_handle_name = [&]() { return arg_name + ".strides"; };
auto array_element_name = [&](const std::string& arr_name, size_t k) {
std::stringstream ss;
ss << arr_name << '[' << k << ']';
return ss.str();
};
auto shape_element_name = [&](size_t k) { return array_element_name(shape_handle_name(), k); };
auto stride_element_name = [&](size_t k) { return array_element_name(stride_handle_name(), k); };
PrimExpr a_ndim = make_const(tvm_ndim_type, static_cast<int64_t>(buffer->shape.size()));
std::ostringstream ndim_err_msg;
ndim_err_msg << arg_name << ".ndim is expected to equal " << buffer->shape.size();
auto msg = tvm::tir::StringImm(ndim_err_msg.str());
init_nest_.emplace_back(AssertStmt(a_ndim == v_ndim, msg, nop));
// type checks
std::ostringstream type_err_msg;
type_err_msg << arg_name << ".dtype is expected to be " << buffer->dtype;
PrimExpr cond = (TVMArrayGet(DataType::UInt(8), handle, builtin::kArrTypeCode) ==
IntImm(DataType::UInt(8), buffer->dtype.code()) &&
TVMArrayGet(DataType::UInt(8), handle, builtin::kArrTypeBits) ==
IntImm(DataType::UInt(8), buffer->dtype.bits()) &&
TVMArrayGet(DataType::UInt(16), handle, builtin::kArrTypeLanes) ==
IntImm(DataType::UInt(16), buffer->dtype.lanes()));
if (!(buffer->dtype == DataType::Int(1) || buffer->dtype == DataType::Int(4) ||
buffer->dtype == DataType::UInt(4))) {
auto type_msg = tvm::tir::StringImm(type_err_msg.str());
asserts_.emplace_back(AssertStmt(cond, type_msg, nop));
}
// shape field
Buffer buf_shape = decl_buffer({IntImm(DataType::Int(32), buffer->shape.size())}, tvm_shape_type,
shape_handle_name());
Var v_shape(shape_handle_name(), DataType::Handle());
def_handle_dtype_.Set(v_shape, make_const(tvm_shape_type, 0));
init_nest_.emplace_back(
LetStmt(buf_shape->data, TVMArrayGet(DataType::Handle(), handle, builtin::kArrShape), nop));
init_nest_.emplace_back(DeclBuffer(buf_shape, nop));
for (size_t k = 0; k < buffer->shape.size(); ++k) {
if (buffer->dtype == DataType::Int(4) || buffer->dtype == DataType::UInt(4) ||
buffer->dtype == DataType::Int(1)) {
break;
}
Bind_(buffer->shape[k],
cast(buffer->shape[k].dtype(), BufferLoad(buf_shape, {IntImm(DataType::Int(32), k)})),
shape_element_name(k), true);
}
// strides field
Buffer buf_strides = decl_buffer({IntImm(DataType::Int(32), buffer->strides.size())},
tvm_shape_type, arg_name + ".strides");
def_handle_dtype_.Set(buf_strides->data, tir::TypeAnnotation(tvm_shape_type));
init_nest_.emplace_back(LetStmt(
buf_strides->data, TVMArrayGet(DataType::Handle(), handle, builtin::kArrStrides), nop));
init_nest_.emplace_back(DeclBuffer(buf_strides, nop));
PrimExpr v_strides_is_null = Call(DataType::Bool(1), builtin::isnullptr(), {buf_strides->data});
if (buffer->strides.size() == 0) {
// Assert the buffer is compact
DataType stype = buffer->DefaultIndexType();
PrimExpr expect_stride = make_const(stype, 1);
Array<PrimExpr> conds;
for (size_t i = buffer->shape.size(); i != 0; --i) {
size_t k = i - 1;
PrimExpr svalue = cast(stype, BufferLoad(buf_strides, {IntImm(DataType::Int(32), k)}));
conds.push_back(buffer->shape[k] == 1 || expect_stride == svalue);
expect_stride = expect_stride * buffer->shape[k];
}
std::ostringstream stride_err_msg;
stride_err_msg << stride_handle_name() << ": expected to be compact array";
if (conds.size() != 0) {
auto stride_msg = tvm::tir::StringImm(stride_err_msg.str());
Stmt check = AssertStmt(
foldl([](PrimExpr a, PrimExpr b, Span span) { return logical_and(a, b, span); },
const_true(1), conds),
stride_msg, Evaluate(0));
check = IfThenElse(Not(v_strides_is_null), check);
asserts_.emplace_back(SeqStmt({check, Evaluate(0)}));
}
} else if (buffer->buffer_type == kAutoBroadcast) {
DataType stype = buffer->DefaultIndexType();
PrimExpr stride = make_const(stype, 1);
for (size_t i = buffer->shape.size(); i != 0; --i) {
size_t k = i - 1;
PrimExpr value =
cast(buffer->shape[k].dtype(), BufferLoad(buf_strides, {IntImm(DataType::Int(32), k)}));
value = tvm::if_then_else(v_strides_is_null, stride, value);
value = tvm::if_then_else(buffer->shape[k] == 1, 0, value);
Bind_(buffer->strides[k], value, stride_element_name(k), true);
stride = analyzer_.Simplify(stride * buffer->shape[k]);
}
} else {
PrimExpr stride_from_shape = 1;
for (int k = buffer->strides.size() - 1; k >= 0; k--) {
PrimExpr explicit_stride =
cast(buffer->shape[k].dtype(), BufferLoad(buf_strides, {IntImm(DataType::Int(32), k)}));
Bind_(buffer->strides[k],
tvm::if_then_else(v_strides_is_null, stride_from_shape, explicit_stride),
stride_element_name(k), true);
stride_from_shape *=
cast(buffer->shape[k].dtype(), BufferLoad(buf_shape, {IntImm(DataType::Int(32), k)}));
}
}
// Byte_offset field.
int data_bytes = GetVectorBytes(buffer->dtype);
if (const auto* const_offset = buffer->elem_offset.as<IntImmNode>()) {
Bind_(make_const(DataType::UInt(64), const_offset->value * data_bytes),
TVMArrayGet(DataType::UInt(64), handle, builtin::kArrByteOffset),
arg_name + ".byte_offset", true);
} else {
if (Bind_(buffer->elem_offset,
cast(buffer->elem_offset.dtype(),
(TVMArrayGet(DataType::UInt(64), handle, builtin::kArrByteOffset) /
make_const(DataType::UInt(64), data_bytes))),
arg_name + ".elem_offset", true)) {
if (buffer->offset_factor > 1) {
PrimExpr offset = buffer->elem_offset;
PrimExpr factor = make_const(offset.dtype(), buffer->offset_factor);
PrimExpr zero = make_zero(offset.dtype());
BinderAddAssert(&analyzer_, truncmod(offset, factor) == zero, arg_name + ".elem_offset",
&asserts_);
}
}
}
// device info.
Bind_(device_type, TVMArrayGet(DataType::Int(32), handle, builtin::kArrDeviceType),
arg_name + ".device_type", true);
Bind_(device_id, TVMArrayGet(DataType::Int(32), handle, builtin::kArrDeviceId),
arg_name + ".device_id", true);
// Data field. Because the validation of the data field may depend
// on a dynamic size defined by the other DLTensor* parameters, this
// field must be generated last.
if (Bind_(buffer->data, TVMArrayGet(DataType::Handle(), handle, builtin::kArrData),
arg_name + ".data", true)) {
Var vptr(buffer->data);
// Check if the data pointer is NULL. This check is skipped for
// size-0 arrays, since CUDA provides a NULL pointer for size-zero
// allocations.
auto alloc_size = [&]() -> PrimExpr {
PrimExpr product = IntImm(buffer->DefaultIndexType(), 1);
for (const auto& dim : buffer->shape) {
product *= dim;
}
return product;
}();
asserts_.emplace_back(AssertStmt(
alloc_size == 0 || !Call(DataType::Bool(), builtin::isnullptr(), {vptr}),
tvm::tir::StringImm(arg_name + " is expected to have non-NULL data pointer"), nop));
def_handle_dtype_.Set(vptr, tir::TypeAnnotation(buffer->dtype));
// mark alignment of external bufs
init_nest_.emplace_back(AttrStmt(vptr, tir::attr::storage_alignment,
IntImm(DataType::Int(32), buffer->data_alignment), nop));
}
}