gemini-code-assist[bot] commented on code in PR #19875:
URL: https://github.com/apache/tvm/pull/19875#discussion_r3461058247
##########
src/relax/op/memory/view.cc:
##########
@@ -167,16 +167,15 @@ Type InferTypeView(const Call& call, const BlockBuilder&
ctx) {
output_ndim = data_ty->ndim;
}
- DataType output_dtype = view_dtype.value_or(data_ty->dtype);
+ DLDataType output_raw_dtype = view_dtype.value_or(data_ty->dtype->dtype);
+ PrimType output_dtype(output_raw_dtype);
- // Helper function, returns the number of bytes per vectorized
- // element. Cannot use `DataType::bytes`, as it returns the
- // number of bytes per scalar element.
- auto get_size_bytes = [](const DataType& dtype) -> ffi::Optional<IntImm> {
- if (dtype.is_void()) {
+ // Helper function returns the number of bytes per vectorized element.
+ auto get_size_bytes = [](DLDataType dtype) -> ffi::Optional<IntImm> {
+ if ((((dtype).code == kDLOpaqueHandle) && ((dtype).bits == 0) &&
((dtype).lanes == 0))) {
return std::nullopt;
} else {
- auto size_bits = dtype.bits() * dtype.lanes();
+ auto size_bits = ((dtype).bits) * static_cast<int16_t>((dtype).lanes);
return IntImm::Int64((size_bits + 7) / 8);
}
};
Review Comment:
We should use the `PrimType` wrapper to safely retrieve the lanes and bits,
and handle scalable vectors by returning `std::nullopt` since their size cannot
be determined at compile time.
```c
auto get_size_bytes = [](DLDataType dtype) -> ffi::Optional<IntImm> {\n
PrimType prim_type(dtype);\n if (prim_type.IsVoid() ||
prim_type.IsScalableVector()) {\n return std::nullopt;\n } else {\n
auto size_bits = prim_type.bits() * prim_type.lanes();\n return
IntImm::Int64((size_bits + 7) / 8);\n }\n };
```
##########
src/relax/transform/lower_alloc_tensor.cc:
##########
@@ -72,7 +72,8 @@ class Mutator : public ExprMutator {
}();
PrimExpr nbytes = [&]() -> PrimExpr {
- PrimExpr nbytes = IntImm::Int64(dtype->value.bytes());
+ PrimExpr nbytes = IntImm::Int64(
+ ((((dtype->value).bits *
static_cast<int16_t>((dtype->value).lanes)) + 7) / 8));
for (const auto& dim : shape) {
nbytes *= dim;
}
Review Comment:
When calculating `nbytes` for the tensor allocation, using raw
`static_cast<int16_t>((dtype->value).lanes)` directly bypasses the safety
checks in `PrimType`. If the tensor has a scalable vector dtype, this will
silently produce a negative or overflowed size at compile time.\n\nWe should
use the `PrimType` wrapper to safely retrieve the lanes and bits, and
explicitly check that it is not a scalable vector.
```c
PrimExpr nbytes = [&]() -> PrimExpr {\n PrimType
prim_type(dtype->value);\n
TVM_FFI_ICHECK(!prim_type.IsScalableVector())\n << "Cannot allocate
tensor with scalable vector dtype at compile time.";\n PrimExpr nbytes =
IntImm::Int64((prim_type.bits() * prim_type.lanes() + 7) / 8);\n for
(const auto& dim : shape) {\n nbytes *= dim;\n }
```
##########
src/relax/transform/utils.h:
##########
@@ -319,39 +319,39 @@ class FunctionCopier : public SymbolicVarRenewMutator {
* \return A Constant.
*/
template <typename T>
-inline Constant MakeConstantScalar(T value, DataType dtype) {
+inline Constant MakeConstantScalar(T value, DLDataType dtype) {
runtime::Tensor arr = runtime::Tensor::Empty({}, dtype, {kDLCPU, 0});
- if (dtype == DataType::Float(32)) {
+ if (dtype == DLDataType{kDLFloat, 32, 1}) {
*static_cast<float*>(arr->data) = static_cast<float>(value);
- } else if (dtype == DataType::Float(64)) {
+ } else if (dtype == DLDataType{kDLFloat, 64, 1}) {
*static_cast<double*>(arr->data) = static_cast<double>(value);
- } else if (dtype == DataType::Int(32)) {
+ } else if (dtype == DLDataType{kDLInt, 32, 1}) {
*static_cast<int32_t*>(arr->data) = static_cast<int32_t>(value);
- } else if (dtype == DataType::Int(64)) {
+ } else if (dtype == DLDataType{kDLInt, 64, 1}) {
*static_cast<int64_t*>(arr->data) = static_cast<int64_t>(value);
- } else if (dtype == DataType::Bool()) {
+ } else if (dtype == DLDataType{kDLBool, 1, 1}) {
Review Comment:
In TVM, `kDLBool` is standardly represented with 8 bits (not 1 bit).
Comparing `dtype` against `DLDataType{kDLBool, 1, 1}` will fail to match
standard boolean constants, causing them to fall through or fail.\n\nWe should
change the bit width check to `8` to match the standard boolean representation.
```c
} else if (dtype == DLDataType{kDLBool, 8, 1}) {
```
##########
src/s_tir/analysis/verify_gpu_code.cc:
##########
@@ -202,38 +204,50 @@ class GPUCodeVerifier : public StmtExprVisitor {
void CheckBufferIndicesVectorizable(const ffi::Array<PrimExpr> indices) {
for (const auto index : indices) {
if (const auto* ramp = index.as<RampNode>()) {
- if (!is_one(ramp->stride) &&
- static_cast<size_t>(ramp->dtype.lanes() * ramp->dtype.bytes()) >
max_vector_bytes_) {
+ PrimType ramp_ty = ramp->ty();
+ DLDataType ramp_dtype = ramp_ty->dtype;
+ if (!is_one(ramp->stride) && ramp_ty.IsFixedLengthVector() &&
+ static_cast<size_t>(static_cast<int16_t>((ramp_dtype).lanes) *
+ (((ramp_dtype).bits + 7) / 8)) >
max_vector_bytes_) {
std::stringstream s;
- s << "Number of lanes (" << ramp->dtype.lanes() << ") times number
of bytes ("
- << ramp->dtype.bytes() << ") for dtype " << ramp->dtype
- << " is greater than the maximum number of vector bytes (" <<
max_vector_bytes_ << ")";
+ s << "Number of lanes (" << static_cast<int16_t>((ramp_dtype).lanes)
+ << ") times number of bytes (" << (((ramp_dtype).bits + 7) / 8) <<
") for dtype "
+ << ramp_dtype << " is greater than the maximum number of vector
bytes ("
+ << max_vector_bytes_ << ")";
errors_.push_back(s.str());
}
}
}
}
void VisitExpr_(const CastNode* op) {
- if (op->dtype.is_vector()) {
- if (static_cast<size_t>(op->dtype.lanes() * op->dtype.bytes()) >
max_vector_bytes_) {
+ PrimType op_ty = op->ty();
+ DLDataType op_dtype = op_ty->dtype;
+ if (op_ty.IsFixedLengthVector()) {
+ if (static_cast<size_t>(static_cast<int16_t>((op_dtype).lanes) *
+ (((op_dtype).bits + 7) / 8)) >
max_vector_bytes_) {
Review Comment:
We can simplify this by using `op_ty.lanes()` and `op_ty.bits()` directly
instead of unpacking to `DLDataType` and performing manual bit/lane arithmetic.
```c
PrimType op_ty = op->ty();\n if (op_ty.IsFixedLengthVector()) {\n
if (static_cast<size_t>(op_ty.lanes() * ((op_ty.bits() + 7) / 8)) >
max_vector_bytes_) {
```
##########
src/s_tir/analysis/verify_gpu_code.cc:
##########
@@ -76,20 +76,22 @@ class GPUCodeVerifier : public StmtExprVisitor {
break;
}
}
+ DLDataType dtype = op->buffer->dtype->dtype;
if (storage_scope.rank == runtime::StorageRank::kLocal) {
- local_memory_per_block_ +=
- static_cast<size_t>(const_size) * op->buffer->dtype.bytes() *
op->buffer->dtype.lanes();
+ local_memory_per_block_ += static_cast<size_t>(const_size) *
(((dtype).bits + 7) / 8) *
+ static_cast<int16_t>((dtype).lanes);
} else if (storage_scope.rank == runtime::StorageRank::kShared) {
- shared_memory_per_block_ +=
- static_cast<size_t>(const_size) * op->buffer->dtype.bytes() *
op->buffer->dtype.lanes();
+ shared_memory_per_block_ += static_cast<size_t>(const_size) *
(((dtype).bits + 7) / 8) *
+ static_cast<int16_t>((dtype).lanes);
}
- if (op->buffer->dtype.is_vector()) {
- if (static_cast<size_t>(op->buffer->dtype.lanes() *
op->buffer->dtype.bytes()) >
+ if ((static_cast<int16_t>((dtype).lanes) > 1)) {
+ if (static_cast<size_t>(static_cast<int16_t>((dtype).lanes) *
(((dtype).bits + 7) / 8)) >
Review Comment:
Since `op->buffer->dtype` is already a `PrimType`, we can use its member
functions `bits()` and `lanes()` directly instead of unpacking it to
`DLDataType` and performing manual bit/lane arithmetic. This is safer and more
readable.
```c
PrimType dtype = op->buffer->dtype;\n if (storage_scope.rank ==
runtime::StorageRank::kLocal) {\n local_memory_per_block_ +=
static_cast<size_t>(const_size) * ((dtype.bits() + 7) / 8) *\n
dtype.lanes();\n } else if (storage_scope.rank ==
runtime::StorageRank::kShared) {\n shared_memory_per_block_ +=
static_cast<size_t>(const_size) * ((dtype.bits() + 7) / 8) *\n
dtype.lanes();\n }\n if (dtype.IsFixedLengthVector()) {\n
if (static_cast<size_t>(dtype.lanes() * ((dtype.bits() + 7) / 8)) >
```
##########
src/relax/transform/static_plan_block_memory.cc:
##########
@@ -135,10 +135,10 @@ class StorageTokenNode : public ffi::Object {
*/
class StorageToken : public ffi::ObjectRef {
public:
- explicit StorageToken(ffi::Array<PrimExpr> shape, DataType dtype,
std::string storage_scope,
+ explicit StorageToken(ffi::Array<PrimExpr> shape, DLDataType dtype,
std::string storage_scope,
ffi::Optional<VDevice> vdevice = std::nullopt) {
// Compute the tensor size from the shape.
- int64_t const_coeff = dtype.bytes() * dtype.lanes();
+ int64_t const_coeff = ((((dtype).bits *
static_cast<int16_t>((dtype).lanes)) + 7) / 8);
Review Comment:
We should use the `PrimType` wrapper to safely retrieve the lanes and bits,
and explicitly check that it is not a scalable vector to prevent
negative/overflowed size calculations.
```c
explicit StorageToken(ffi::Array<PrimExpr> shape, DLDataType dtype,
std::string storage_scope,\n ffi::Optional<VDevice>
vdevice = std::nullopt) {\n // Compute the tensor size from the shape.\n
PrimType prim_type(dtype);\n TVM_FFI_ICHECK(!prim_type.IsScalableVector())\n
<< "Cannot allocate storage for scalable vector dtype at compile
time.";\n int64_t const_coeff = (prim_type.bits() * prim_type.lanes() + 7) /
8;
```
##########
src/s_tir/analysis/verify_gpu_code.cc:
##########
@@ -202,38 +204,50 @@ class GPUCodeVerifier : public StmtExprVisitor {
void CheckBufferIndicesVectorizable(const ffi::Array<PrimExpr> indices) {
for (const auto index : indices) {
if (const auto* ramp = index.as<RampNode>()) {
- if (!is_one(ramp->stride) &&
- static_cast<size_t>(ramp->dtype.lanes() * ramp->dtype.bytes()) >
max_vector_bytes_) {
+ PrimType ramp_ty = ramp->ty();
+ DLDataType ramp_dtype = ramp_ty->dtype;
+ if (!is_one(ramp->stride) && ramp_ty.IsFixedLengthVector() &&
Review Comment:
We can simplify this by using `ramp_ty.lanes()` and `ramp_ty.bits()`
directly instead of unpacking to `DLDataType` and performing manual bit/lane
arithmetic.
```c
PrimType ramp_ty = ramp->ty();\n if (!is_one(ramp->stride) &&
ramp_ty.IsFixedLengthVector() &&\n
static_cast<size_t>(ramp_ty.lanes() * ((ramp_ty.bits() + 7) / 8)) >
max_vector_bytes_) {
```
##########
src/s_tir/analysis/verify_gpu_code.cc:
##########
@@ -202,38 +204,50 @@ class GPUCodeVerifier : public StmtExprVisitor {
void CheckBufferIndicesVectorizable(const ffi::Array<PrimExpr> indices) {
for (const auto index : indices) {
if (const auto* ramp = index.as<RampNode>()) {
- if (!is_one(ramp->stride) &&
- static_cast<size_t>(ramp->dtype.lanes() * ramp->dtype.bytes()) >
max_vector_bytes_) {
+ PrimType ramp_ty = ramp->ty();
+ DLDataType ramp_dtype = ramp_ty->dtype;
+ if (!is_one(ramp->stride) && ramp_ty.IsFixedLengthVector() &&
+ static_cast<size_t>(static_cast<int16_t>((ramp_dtype).lanes) *
+ (((ramp_dtype).bits + 7) / 8)) >
max_vector_bytes_) {
std::stringstream s;
- s << "Number of lanes (" << ramp->dtype.lanes() << ") times number
of bytes ("
- << ramp->dtype.bytes() << ") for dtype " << ramp->dtype
- << " is greater than the maximum number of vector bytes (" <<
max_vector_bytes_ << ")";
+ s << "Number of lanes (" << static_cast<int16_t>((ramp_dtype).lanes)
+ << ") times number of bytes (" << (((ramp_dtype).bits + 7) / 8) <<
") for dtype "
+ << ramp_dtype << " is greater than the maximum number of vector
bytes ("
+ << max_vector_bytes_ << ")";
errors_.push_back(s.str());
}
}
}
}
void VisitExpr_(const CastNode* op) {
- if (op->dtype.is_vector()) {
- if (static_cast<size_t>(op->dtype.lanes() * op->dtype.bytes()) >
max_vector_bytes_) {
+ PrimType op_ty = op->ty();
+ DLDataType op_dtype = op_ty->dtype;
+ if (op_ty.IsFixedLengthVector()) {
+ if (static_cast<size_t>(static_cast<int16_t>((op_dtype).lanes) *
+ (((op_dtype).bits + 7) / 8)) >
max_vector_bytes_) {
std::stringstream s;
- s << "Number of lanes (" << op->dtype.lanes() << ") times number of
bytes ("
- << op->dtype.bytes() << ") for dtype " << op->dtype
- << " is greater than the maximum number of vector bytes (" <<
max_vector_bytes_ << ")";
+ s << "Number of lanes (" << static_cast<int16_t>((op_dtype).lanes)
+ << ") times number of bytes (" << (((op_dtype).bits + 7) / 8) << ")
for dtype "
+ << op_dtype << " is greater than the maximum number of vector bytes
("
+ << max_vector_bytes_ << ")";
errors_.push_back(s.str());
}
}
ExprVisitor::VisitExpr_(op);
}
void VisitExpr_(const BufferLoadNode* op) {
- if (op->dtype.is_vector()) {
- if (static_cast<size_t>(op->dtype.lanes() * op->dtype.bytes()) >
max_vector_bytes_) {
+ PrimType op_ty = op->ty();
+ DLDataType op_dtype = op_ty->dtype;
+ if (op_ty.IsFixedLengthVector()) {
+ if (static_cast<size_t>(static_cast<int16_t>((op_dtype).lanes) *
+ (((op_dtype).bits + 7) / 8)) >
max_vector_bytes_) {
Review Comment:
We can simplify this by using `op_ty.lanes()` and `op_ty.bits()` directly
instead of unpacking to `DLDataType` and performing manual bit/lane arithmetic.
```c
PrimType op_ty = op->ty();\n if (op_ty.IsFixedLengthVector()) {\n
if (static_cast<size_t>(op_ty.lanes() * ((op_ty.bits() + 7) / 8)) >
max_vector_bytes_) {
```
##########
src/s_tir/analysis/verify_gpu_code.cc:
##########
@@ -242,13 +256,16 @@ class GPUCodeVerifier : public StmtExprVisitor {
}
void VisitStmt_(const BufferStoreNode* op) {
- if (op->value->dtype.is_vector()) {
- if (static_cast<size_t>(op->value->dtype.lanes() *
op->value->dtype.bytes()) >
- max_vector_bytes_) {
+ PrimType value_ty = op->value.ty();
+ DLDataType value_dtype = value_ty->dtype;
+ if (value_ty.IsFixedLengthVector()) {
+ if (static_cast<size_t>(static_cast<int16_t>((value_dtype).lanes) *
+ (((value_dtype).bits + 7) / 8)) >
max_vector_bytes_) {
Review Comment:
We can simplify this by using `value_ty.lanes()` and `value_ty.bits()`
directly instead of unpacking to `DLDataType` and performing manual bit/lane
arithmetic.
```c
PrimType value_ty = op->value.ty();\n if
(value_ty.IsFixedLengthVector()) {\n if
(static_cast<size_t>(value_ty.lanes() * ((value_ty.bits() + 7) / 8)) >
max_vector_bytes_) {
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
