gemini-code-assist[bot] commented on code in PR #19875:
URL: https://github.com/apache/tvm/pull/19875#discussion_r3461433341
##########
src/relax/op/nn/nn.cc:
##########
@@ -186,10 +188,14 @@ Type InferTypeSoftmax(const Call& call, const
BlockBuilder& ctx) {
if (data_ty->IsUnknownNdim()) {
return data_ty;
}
- if (!data_ty->IsUnknownDtype() && !data_ty->dtype.is_float() &&
!data_ty->dtype.is_bfloat()) {
- TVM_FFI_VISIT_THROW(TypeError, call) << "Softmax requires the input tensor
to have float "
- "dtype. However, the given input
dtype is "
- << data_ty->dtype;
+ if (!data_ty->IsUnknownDtype()) {
+ PrimType data_dtype = data_ty->dtype;
+ // Softmax only requires a floating element kind; lane encoding is
irrelevant to the check.
+ if (data_dtype.code() != kDLFloat && data_dtype.code() != kDLBfloat) {
Review Comment:
We can use the more idiomatic `MatchesCode` helper on `PrimType` instead of
directly comparing `code()` against `kDLFloat` and `kDLBfloat`.
```suggestion
if (!data_dtype.MatchesCode(kDLFloat, kDLBfloat)) {
```
##########
src/relax/op/nn/nn.cc:
##########
@@ -380,10 +386,14 @@ bool NormCheckDtypeAndShape(const Call& call, const
BlockBuilder& ctx,
axes_non_neg = NormalizeAxes(call, ctx, data_ty->ndim, axes);
}
int n_axis = axes.size();
- if (!data_ty->IsUnknownDtype() && (!data_ty->dtype.is_float() &&
!data_ty->dtype.is_bfloat())) {
- TVM_FFI_VISIT_THROW(TypeError, call)
- << op << " requires the input data to have float dtype. However, the
given data dtype is "
- << data_ty->dtype;
+ if (!data_ty->IsUnknownDtype()) {
+ PrimType data_dtype = data_ty->dtype;
+ // Norm ops only require a floating element kind; lane encoding is
irrelevant to the check.
+ if (data_dtype.code() != kDLFloat && data_dtype.code() != kDLBfloat) {
Review Comment:
We can use the more idiomatic `MatchesCode` helper on `PrimType` instead of
directly comparing `code()` against `kDLFloat` and `kDLBfloat`.
```suggestion
if (!data_dtype.MatchesCode(kDLFloat, kDLBfloat)) {
```
##########
src/relax/op/tensor/index.cc:
##########
@@ -84,11 +84,15 @@ Type InferTypeTake(const Call& call, const BlockBuilder&
ctx) {
if (indices_ty->IsUnknownDtype()) {
LOG(WARNING) << "Data type of indices has not been specified. Assume it
has an integer type.";
- } else if (!(indices_ty->dtype.is_int() || indices_ty->dtype.is_uint())) {
- TVM_FFI_VISIT_THROW(TypeError, call)
- << "Take op requires the input indices to have integer dtype. However,
the "
- "given indices dtype is "
- << indices_ty->dtype;
+ } else {
+ PrimType indices_dtype = indices_ty->dtype;
+ if (!indices_dtype.MatchesCode(DLDataTypeCode::kDLInt) &&
+ !indices_dtype.MatchesCode(DLDataTypeCode::kDLUInt)) {
Review Comment:
Instead of calling `MatchesCode` twice separately for `kDLInt` and
`kDLUInt`, you can simplify this check by passing both codes to a single
`MatchesCode` call. `MatchesCode` is a variadic template designed exactly for
this purpose.
```c
if (!indices_dtype.MatchesCode(DLDataTypeCode::kDLInt,
DLDataTypeCode::kDLUInt)) {
```
##########
src/relax/op/nn/nn.cc:
##########
@@ -122,7 +122,9 @@ Type InferTypePRelu(const Call& call, const BlockBuilder&
ctx) {
if (data_ty->IsUnknownNdim()) {
return data_ty;
}
- if (!data_ty->IsUnknownDtype() && !data_ty->dtype.is_float()) {
+ PrimType data_dtype = data_ty->dtype;
+ // PRelu preserves the old float-kind check; vector lanes are irrelevant to
this check.
+ if (!data_ty->IsUnknownDtype() && data_dtype.code() !=
DLDataTypeCode::kDLFloat) {
Review Comment:
We can use the more idiomatic `MatchesCode` helper on `PrimType` instead of
directly comparing `code()` against `kDLFloat`.
```suggestion
if (!data_ty->IsUnknownDtype() &&
!data_dtype.MatchesCode(DLDataTypeCode::kDLFloat)) {
```
##########
src/relax/op/distributed/nn.cc:
##########
@@ -33,7 +33,9 @@ Type InferDistTypeSoftmax(const Call& call, const
BlockBuilder& ctx) {
if (input_tensor_ty->IsUnknownNdim()) {
TVM_FFI_VISIT_THROW(ValueError, call) << "Input of distributed operator
must have known ndim";
}
- if (!input_tensor_ty->IsUnknownDtype() &&
!input_tensor_ty->dtype.is_float()) {
+ PrimType input_dtype = input_tensor_ty->dtype;
+ // Softmax validation preserves the old float-kind check; lanes do not
affect this policy.
+ if (!input_tensor_ty->IsUnknownDtype() && input_dtype.code() !=
DLDataTypeCode::kDLFloat) {
Review Comment:
We can use the more idiomatic `MatchesCode` helper on `PrimType` instead of
directly comparing `code()` against `kDLFloat`.
```suggestion
if (!input_tensor_ty->IsUnknownDtype() &&
!input_dtype.MatchesCode(DLDataTypeCode::kDLFloat)) {
```
##########
src/relax/op/distributed/unary.h:
##########
@@ -40,15 +40,22 @@ Type InferDistTypeUnary(const Call& call, const
BlockBuilder& ctx, FType f_compu
distributed::DTensorType input_dtensor_ty = input_dtensor_tys[0];
TensorType input_tensor_ty = input_dtensor_ty->tensor_ty;
+ PrimType input_dtype = input_tensor_ty->dtype;
+ // Unary op validation preserves the old float-kind check; lanes do not
affect this policy.
if (require_float_dtype && !input_tensor_ty->IsUnknownDtype() &&
- !input_tensor_ty->dtype.is_float()) {
+ input_dtype.code() != DLDataTypeCode::kDLFloat) {
Review Comment:
We can use the more idiomatic `MatchesCode` helper on `PrimType` instead of
directly comparing `code()` against `kDLFloat`.
```suggestion
if (require_float_dtype && !input_tensor_ty->IsUnknownDtype() &&
!input_dtype.MatchesCode(DLDataTypeCode::kDLFloat)) {
```
##########
src/arith/const_fold.h:
##########
@@ -72,18 +72,29 @@ inline ffi::Optional<PrimExpr> TryConstFold(PrimExpr a);
* \param type The type to represent index.
* \return the checked result.
*/
-inline bool IsIndexType(const DataType& type) {
- return type.is_int() && !type.is_scalable_or_fixed_length_vector() &&
- (type.bits() == 32 || type.bits() == 64);
+inline bool IsIndexType(DLDataType type) {
+ return type.code == static_cast<uint8_t>(DLDataTypeCode::kDLInt) &&
+ (type.bits == 32 || type.bits == 64) && type.lanes == 1;
+}
+
+inline bool IsIndexTypedExpr(const PrimExprNode* expr) {
+ TVM_FFI_DCHECK(expr != nullptr);
+ TVM_FFI_DCHECK(expr->BaseExprNode::ty.defined());
+ const auto* prim_ty = expr->BaseExprNode::ty.as<PrimTypeNode>();
+ TVM_FFI_DCHECK(prim_ty != nullptr);
+ return IsIndexType(prim_ty->dtype);
+}
+
+inline bool IsIndexTypedExpr(const PrimExpr& expr) {
+ return IsIndexTypedExpr(static_cast<const PrimExprNode*>(expr.get()));
}
/*! \brief Helper to get const folding result repr in int64. */
-inline int64_t GetFoldResultInt64Repr(int64_t x, const DataType& dtype) {
+inline int64_t GetFoldResultInt64Repr(int64_t x, const PrimType& dtype) {
if (dtype.bits() < 64) {
x &= (1LL << dtype.bits()) - 1;
}
- if (dtype.is_int()) {
- // get sign extended value of integer with specified bits
+ if (dtype.code() == DLDataTypeCode::kDLInt) {
Review Comment:
We can use the more idiomatic `MatchesCode` helper on `PrimType` instead of
directly comparing `code()` against `kDLInt`.
```suggestion
if (dtype.MatchesCode(DLDataTypeCode::kDLInt)) {
```
##########
src/arith/transitive_comparison_analyzer.cc:
##########
@@ -615,7 +615,8 @@ CompareResult
TransitiveComparisonAnalyzer::Impl::TryCompare(const PrimExpr& lhs
const PrimExpr&
rhs_expr,
bool
propagate_inequalities) const {
// Currently only supports integer checks
- if (!lhs_expr.dtype().is_int() || !rhs_expr.dtype().is_int()) {
+ if (lhs_expr.ty().code() != DLDataTypeCode::kDLInt ||
+ rhs_expr.ty().code() != DLDataTypeCode::kDLInt) {
Review Comment:
We can use the more idiomatic `MatchesCode` helper on `PrimType` instead of
directly comparing `code()` against `kDLInt`.
```suggestion
if (!lhs_expr.ty().MatchesCode(DLDataTypeCode::kDLInt) ||
!rhs_expr.ty().MatchesCode(DLDataTypeCode::kDLInt)) {
```
##########
src/relax/op/tensor/search.cc:
##########
@@ -119,13 +118,15 @@ Type InferTypeWhere(const Call& call, const BlockBuilder&
ctx) {
}
}
- if (!cond_ty->dtype.is_bool()) {
+ PrimType cond_dtype = cond_ty->dtype;
+ // Where condition validation only checks the boolean element kind; lanes
are irrelevant here.
+ if (cond_dtype.code() != DLDataTypeCode::kDLBool) {
Review Comment:
We can use the more idiomatic `MatchesCode` helper on `PrimType` instead of
directly comparing `code()` against `kDLBool`.
```suggestion
if (!cond_dtype.MatchesCode(DLDataTypeCode::kDLBool)) {
```
##########
src/relax/transform/gradient.cc:
##########
@@ -707,7 +707,8 @@ class GradientMutator : private ExprMutator {
static bool IsFloatTensorType(const Type& ty) {
auto* tensor_ty = ty.as<TensorTypeNode>();
- return tensor_ty && tensor_ty->dtype.is_float();
+ // Gradient eligibility preserves the old float-kind check; lanes do not
affect this policy.
+ return tensor_ty && tensor_ty->dtype.code() == DLDataTypeCode::kDLFloat;
Review Comment:
We can use the more idiomatic `MatchesCode` helper on `PrimType` instead of
directly comparing `code()` against `kDLFloat`.
```suggestion
return tensor_ty &&
tensor_ty->dtype.MatchesCode(DLDataTypeCode::kDLFloat);
```
##########
src/arith/int_set.cc:
##########
@@ -72,8 +72,10 @@ TVM_FFI_STATIC_INIT_BLOCK() {
IntervalSet Intersect(AnalyzerObj* analyzer, IntervalSet a, IntervalSet b) {
PrimExpr max_value = min(a->max_value, b->max_value);
PrimExpr min_value = max(a->min_value, b->min_value);
- if ((max_value.dtype().is_int() || max_value.dtype().is_uint()) &&
- (min_value.dtype().is_int() || min_value.dtype().is_uint()) &&
+ PrimType max_ty = max_value.ty();
+ PrimType min_ty = min_value.ty();
+ if ((max_ty.code() == DLDataTypeCode::kDLInt || max_ty.code() ==
DLDataTypeCode::kDLUInt) &&
+ (min_ty.code() == DLDataTypeCode::kDLInt || min_ty.code() ==
DLDataTypeCode::kDLUInt) &&
Review Comment:
We can use the more idiomatic `MatchesCode` helper on `PrimType` instead of
directly comparing `code()` against `kDLInt` and `kDLUInt`.
```suggestion
if (max_ty.MatchesCode(DLDataTypeCode::kDLInt, DLDataTypeCode::kDLUInt) &&
min_ty.MatchesCode(DLDataTypeCode::kDLInt, DLDataTypeCode::kDLUInt) &&
```
##########
src/arith/int_set.cc:
##########
@@ -569,18 +572,19 @@ class IntervalSetEvaluator : public
ExprFunctor<IntervalSet(const PrimExpr&)> {
// short cut for the int set.
if (value_set->min_value.same_as(value_set->max_value)) {
if (value_set->IsEmpty()) return value_set;
- return IntervalSet::SinglePoint(cast(op->dtype, value_set->min_value));
+ return IntervalSet::SinglePoint(cast(op->ty(), value_set->min_value));
}
PrimExpr min_value =
- value_set->HasLowerBound() ? cast(op->dtype, value_set->min_value) :
neg_inf();
+ value_set->HasLowerBound() ? cast(op->ty(), value_set->min_value) :
neg_inf();
PrimExpr max_value =
- value_set->HasUpperBound() ? cast(op->dtype, value_set->max_value) :
pos_inf();
+ value_set->HasUpperBound() ? cast(op->ty(), value_set->max_value) :
pos_inf();
return IntervalSet(min_value, max_value);
}
IntervalSet VisitExpr_(const BufferLoadNode* op) final {
- if (!(op->dtype.is_int() || op->dtype.is_uint())) {
- DLOG(WARNING) << "cannot evaluate set BufferLoad which loads from a " <<
op->dtype
+ PrimType op_ty = op->ty();
+ if (!(op_ty.code() == DLDataTypeCode::kDLInt || op_ty.code() ==
DLDataTypeCode::kDLUInt)) {
Review Comment:
We can use the more idiomatic `MatchesCode` helper on `PrimType` instead of
directly comparing `code()` against `kDLInt` and `kDLUInt`.
```suggestion
if (!op_ty.MatchesCode(DLDataTypeCode::kDLInt, DLDataTypeCode::kDLUInt))
{
```
##########
src/arith/int_constraints.cc:
##########
@@ -74,7 +74,9 @@ ffi::Array<PrimExpr> AsConditions(const ffi::Array<Var>&
variables,
IntGroupBounds::IntGroupBounds(PrimExpr coef, ffi::Array<PrimExpr> lower,
ffi::Array<PrimExpr> equal,
ffi::Array<PrimExpr> upper) {
- TVM_FFI_ICHECK(coef.dtype().is_int() || coef.dtype().is_uint())
+ PrimType coef_ty = coef.ty();
+ TVM_FFI_ICHECK(coef_ty.code() == DLDataTypeCode::kDLInt ||
+ coef_ty.code() == DLDataTypeCode::kDLUInt)
Review Comment:
We can use the more idiomatic `MatchesCode` helper on `PrimType` instead of
directly comparing `code()` against `kDLInt` and `kDLUInt`.
```c
TVM_FFI_ICHECK(coef_ty.MatchesCode(DLDataTypeCode::kDLInt,
DLDataTypeCode::kDLUInt))
```
##########
src/arith/bound_deducer.cc:
##########
@@ -96,7 +96,8 @@ class BoundDeducer : public ExprFunctor<void(const
PrimExpr&)> {
void VisitExprDefault_(const ffi::Object* op) final { success_ = false; }
SignType GetSignType(const PrimExpr& e) {
- if (e.dtype().is_uint()) {
+ PrimType e_ty = e.ty();
+ if (e_ty.code() == DLDataTypeCode::kDLUInt) {
Review Comment:
We can use the more idiomatic `MatchesCode` helper on `PrimType` instead of
directly comparing `code()` against `kDLUInt`.
```suggestion
if (e_ty.MatchesCode(DLDataTypeCode::kDLUInt)) {
```
##########
src/arith/int_constraints.cc:
##########
@@ -232,7 +234,9 @@ IntConstraints::IntConstraints(ffi::Array<Var> variables,
ffi::Map<Var, Range> r
}
TVM_FFI_ICHECK(relations.defined());
for (const auto& var : variables) {
- TVM_FFI_ICHECK(var.dtype().is_int() || var.dtype().is_uint())
+ PrimType var_ty = var.ty();
+ TVM_FFI_ICHECK(var_ty.code() == DLDataTypeCode::kDLInt ||
+ var_ty.code() == DLDataTypeCode::kDLUInt)
Review Comment:
We can use the more idiomatic `MatchesCode` helper on `PrimType` instead of
directly comparing `code()` against `kDLInt` and `kDLUInt`.
```c
TVM_FFI_ICHECK(var_ty.MatchesCode(DLDataTypeCode::kDLInt,
DLDataTypeCode::kDLUInt))
```
--
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]
