Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add arange, reduction, scalar-tensor, constant alloc, type conversion dtype functions #2057

Merged
merged 5 commits into from
Apr 26, 2023
Merged

Add arange, reduction, scalar-tensor, constant alloc, type conversion dtype functions #2057

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
7b1c93f
Add dtype functions for "Arange ops"
ramiro050 Mar 24, 2023
9c7c7b2
Add dtype functions for reduction ops
ramiro050 Apr 18, 2023
bf91c22
Add dtype functions for scalar to tensor ops
ramiro050 Apr 24, 2023
f5e9d5f
Add dtype functions for constant tensor allocation ops
ramiro050 Apr 24, 2023
384626f
Add dtype functions for type conversion ops
ramiro050 Apr 24, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 0 additions & 3 deletions e2e_testing/xfail_sets.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -93,9 +93,6 @@
"ElementwiseAddScalar_NumToTensorFloat_Module_basic",
# ERROR: assert isinstance(e, FakeTensor)
"RsubInt0d_NumToTensor_Module_basic",

# Dtype function transition failures
"MobilenetV3Module_basic",
}

STABLEHLO_PASS_SET = {
Expand Down
540 changes: 540 additions & 0 deletions lib/Dialect/Torch/Transforms/AbstractInterpLibrary.cpp
View file
Open in desktop

Large diffs are not rendered by default.

246 changes: 0 additions & 246 deletions lib/Dialect/Torch/Transforms/RefineTypes.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -429,12 +429,6 @@ class TypeAnalysis : public dataflow::SparseDataFlowAnalysis<

void visitAtenLinearOp(AtenLinearOp op,
ArrayRef<const ValueState *> operands);
void visitAtenArangeStartStepOp(AtenArangeStartStepOp op);
void visitAtenArangeStartOp(AtenArangeStartOp op);
void visitAtenArangeOp(AtenArangeOp op);
void visitAtenArangeLikeOpHelper(Operation *op, std::optional<Value> start,
Value end, std::optional<Value> step,
Value dtype);
void visitReductionAlongAllDimsOp(Operation *op, Type dtype,
ArrayRef<const ValueState *> operands);
void visitReductionAlongDimIntListOp(Operation *op, Value dim, Value keepdim,
Expand Down Expand Up @@ -649,199 +643,11 @@ void TypeAnalysis::visitOperation(Operation *op,
return;
}

if (auto arange = dyn_cast<AtenArangeOp>(op)) {
visitAtenArangeOp(arange);
return;
}
if (auto arangeStart = dyn_cast<AtenArangeStartOp>(op)) {
visitAtenArangeStartOp(arangeStart);
return;
}
if (auto arangeStartStep = dyn_cast<AtenArangeStartStepOp>(op)) {
visitAtenArangeStartStepOp(arangeStartStep);
return;
}

if (auto sum = dyn_cast<AtenSumOp>(op)) {
Type defaultDtype = operands[0]->getValue().dtype;
// If the input dtype is bool, the result type should be i64.
if (defaultDtype.isInteger(1))
defaultDtype =
IntegerType::get(op->getContext(), 64, IntegerType::Signed);
Type dtype = getDtypeOrDefault(sum.getContext(), sum.getDtype(), defaultDtype);
auto knowledge =
ValueKnowledge::getTensorPessimisticValueState(op->getContext());
knowledge.dtype = dtype;
incorporateKnowledge(op->getResult(0), knowledge);
return;
}
if (auto sumDimIntList = dyn_cast<AtenSumDimIntListOp>(op)) {
Type defaultDtype = operands[0]->getValue().dtype;
if (!defaultDtype) {
incorporateKnowledge(
sumDimIntList.getResult(),
ValueKnowledge::getTensorPessimisticValueState(op->getContext()));
return;
}
// If the input dtype is bool, the result type should be i64.
if (defaultDtype.isInteger(1))
defaultDtype =
IntegerType::get(op->getContext(), 64, IntegerType::Signed);
Type dtype = getDtypeOrDefault(sumDimIntList.getContext(),
sumDimIntList.getDtype(), defaultDtype);
visitReductionAlongDimIntListOp(sumDimIntList, sumDimIntList.getDim(),
sumDimIntList.getKeepdim(), dtype, operands);
return;
}
if (auto meanDim = dyn_cast<AtenMeanDimOp>(op)) {
Type defaultDtype = operands[0]->getValue().dtype;
Type dtype =
getDtypeOrDefault(meanDim.getContext(), meanDim.getDtype(), defaultDtype);
visitReductionAlongDimIntListOp(meanDim, meanDim.getDim(), meanDim.getKeepdim(),
dtype, operands);
return;
}
if (auto argmax = dyn_cast<AtenArgmaxOp>(op)) {
Value dim = argmax.getDim();
Type dtype = IntegerType::get(op->getContext(), 64, IntegerType::Signed);
if (dim.getType().isa<Torch::NoneType>()) {
visitReductionAlongAllDimsOp(op, dtype, operands);
return;
}
if (dim.getType().isa<Torch::IntType>()) {
visitReductionAlongDimIntOp(argmax, argmax.getDim(), argmax.getKeepdim(), dtype,
operands);
return;
}
}
if (auto anyDim = dyn_cast<AtenAnyDimOp>(op)) {
Type dtype = operands[0]->getValue().dtype;
visitReductionAlongDimIntOp(anyDim, anyDim.getDim(), anyDim.getKeepdim(), dtype,
operands);
return;
}
if (auto maxDim = dyn_cast<AtenMaxDimOp>(op)) {
Type firstResDtype = operands[0]->getValue().dtype;
Type secondResDtype =
IntegerType::get(op->getContext(), 64, IntegerType::Signed);
visitReductionAlongDimIntOp(maxDim, maxDim.getDim(), maxDim.getKeepdim(),
firstResDtype, operands);
visitReductionAlongDimIntOp(maxDim, maxDim.getDim(), maxDim.getKeepdim(),
secondResDtype, operands, /*resNum=*/1);
return;
}
if (auto mean = dyn_cast<AtenMeanOp>(op)) {
Type defaultDtype = operands[0]->getValue().dtype;
Type dtype =
getDtypeOrDefault(mean.getContext(), mean.getDtype(), defaultDtype);
visitReductionAlongAllDimsOp(mean, dtype, operands);
return;
} else if (isa<AtenMaxOp, AtenAmaxOp>(op)) {
Type dtype = operands[0]->getValue().dtype;
visitReductionAlongAllDimsOp(op, dtype, operands);
return;
} else if (isa<AtenStdOp, AtenStdDimOp, AtenStdCorrectionOp, AtenVarOp,
AtenVarDimOp, AtenVarCorrectionOp, PrimsVarOp>(op)) {
auto input = operands[0]->getValue();
visitReductionAlongAllDimsOp(op, input.dtype, operands);
return;
}

if (auto tensorFloat = dyn_cast<AtenTensorFloatOp>(op)) {
visitScalarToTensorConversionOp<AtenTensorFloatOp>(tensorFloat);
return;
} else if (auto tensorInt = dyn_cast<AtenTensorIntOp>(op)) {
visitScalarToTensorConversionOp<AtenTensorIntOp>(tensorInt);
return;
} else if (auto tensorBool = dyn_cast<AtenTensorBoolOp>(op)) {
visitScalarToTensorConversionOp<AtenTensorBoolOp>(tensorBool);
return;
}

if (auto tensor = dyn_cast<AtenTensorOp>(op)) {
visitAtenTensorOp(tensor);
return;
}

if (auto zeros = dyn_cast<AtenZerosOp>(op)) {
visitConstantTensorAllocOp<AtenZerosOp>(zeros, /*dataType=*/{});
return;
} else if (auto ones = dyn_cast<AtenOnesOp>(op)) {
visitConstantTensorAllocOp<AtenOnesOp>(ones, /*dataType=*/{});
return;
} else if (auto emptyMemoryFormat = dyn_cast<AtenEmptyMemoryFormatOp>(op)) {
visitConstantTensorAllocOp<AtenEmptyMemoryFormatOp>(emptyMemoryFormat,
/*dataType=*/{});
return;
} else if (auto full = dyn_cast<AtenFullOp>(op)) {
visitConstantTensorAllocOp<AtenFullOp>(
full, /*dataType=*/full.getFillValue().getType());
return;
} else if (auto zerosLike = dyn_cast<AtenZerosLikeOp>(op)) {
visitConstantTensorAllocLikeOp<AtenZerosLikeOp>(zerosLike, operands);
return;
} else if (auto onesLike = dyn_cast<AtenOnesLikeOp>(op)) {
visitConstantTensorAllocLikeOp<AtenOnesLikeOp>(onesLike, operands);
return;
} else if (auto emptyLike = dyn_cast<AtenEmptyLikeOp>(op)) {
visitConstantTensorAllocLikeOp<AtenEmptyLikeOp>(emptyLike, operands);
return;
} else if (auto fullLike = dyn_cast<AtenFullLikeOp>(op)) {
visitConstantTensorAllocLikeOp<AtenFullLikeOp>(fullLike, operands);
return;
} else if (auto newZeros = dyn_cast<AtenNewZerosOp>(op)) {
visitConstantTensorNewLikeOp<AtenNewZerosOp>(newZeros, operands);
return;
} else if (auto newOnes = dyn_cast<AtenNewOnesOp>(op)) {
visitConstantTensorNewLikeOp<AtenNewOnesOp>(newOnes, operands);
return;
} else if (auto newEmpty = dyn_cast<AtenNewEmptyOp>(op)) {
visitConstantTensorNewLikeOp<AtenNewEmptyOp>(newEmpty, operands);
return;
} else if (auto newEmptyStrided = dyn_cast<AtenNewEmptyStridedOp>(op)) {
visitConstantTensorNewLikeOp<AtenNewEmptyStridedOp>(newEmptyStrided,
operands);
return;
} else if (auto randLike = dyn_cast<AtenRandLikeOp>(op)) {
visitConstantTensorAllocLikeOp<AtenRandLikeOp>(randLike, operands);
return;
} else if (auto randLike = dyn_cast<AtenRandnLikeOp>(op)) {
visitConstantTensorAllocLikeOp<AtenRandnLikeOp>(randLike, operands);
return;
} else if (auto toCopy = dyn_cast<Aten_ToCopyOp>(op)) {
visitConstantTensorAllocLikeOp<Aten_ToCopyOp>(toCopy, operands);
return;
}

if (auto toDtype = dyn_cast<AtenToDtypeOp>(op)) {
visitAtenToDtypeLikeOp<AtenToDtypeOp>(toDtype, operands);
return;
}

if (auto primsConvertElementType = dyn_cast<PrimsConvertElementTypeOp>(op)) {
visitAtenToDtypeLikeOp<PrimsConvertElementTypeOp>(primsConvertElementType,
operands);
return;
}

if (auto toDtypeLayout = dyn_cast<AtenToDtypeLayoutOp>(op)) {
visitAtenToDtypeLikeOp<AtenToDtypeLayoutOp>(toDtypeLayout, operands);
return;
}

if (auto toDtype = dyn_cast<AtenToDeviceOp>(op)) {
visitAtenToDtypeLikeOp<AtenToDeviceOp>(toDtype, operands);
return;
}

if (auto toOther = dyn_cast<AtenToOtherOp>(op)) {
visitTypeConversionOp<AtenToOtherOp>(toOther, operands);
return;
} else if (auto typeAs = dyn_cast<AtenTypeAsOp>(op)) {
visitTypeConversionOp<AtenTypeAsOp>(typeAs, operands);
return;
}

if (auto cat = dyn_cast<AtenCatOp>(op)) {
visitAtenCatLikeOp<AtenCatOp>(cat, operands);
return;
Expand Down Expand Up @@ -902,15 +708,6 @@ void TypeAnalysis::visitOperation(Operation *op,
return;
}

if (auto vectorNorm = dyn_cast<AtenLinalgVectorNormOp>(op)) {
Type defaultDtype = operands[0]->getValue().dtype;
Type dtype = getDtypeOrDefault(vectorNorm.getContext(), vectorNorm.getDtype(),
defaultDtype);
visitReductionAlongDimIntListOp(vectorNorm, vectorNorm.getDim(),
vectorNorm.getKeepdim(), dtype, operands);
return;
}

if (auto randIntLow = dyn_cast<AtenRandintLowOp>(op)) {
auto knowledge =
ValueKnowledge::getTensorPessimisticValueState(op->getContext());
Expand Down Expand Up @@ -1021,49 +818,6 @@ void TypeAnalysis::visitAtenEmbeddingBagOp(Operation *op) {
return;
}

// Arange like ops returns a 1-D tensor of size ceil(end - start).
void TypeAnalysis::visitAtenArangeLikeOpHelper(Operation *op,
std::optional<Value> start,
Value end,
std::optional<Value> step,
Value dtype) {
auto knowledge =
ValueKnowledge::getTensorPessimisticValueState(op->getContext());
int64_t dtypeInt;
if (matchPattern(dtype, m_TorchConstantInt(&dtypeInt))) {
knowledge.dtype = getTypeForDTypeInteger(op->getContext(), dtypeInt);
} else if (dtype.getType().isa<Torch::NoneType>()) {
// From torch/_torch_docs.py:
// If `dtype` is not given, infer the data type from the other input
// arguments. If any of `start`, `end`, or `step` are floating-point, the
// `dtype` is inferred to be the default dtype, see
// `torch.get_default_dtype`. Otherwise, the `dtype` is inferred to
// be `torch.int64`
if ((start.has_value() && (*start).getType().isa<Torch::FloatType>()) ||
end.getType().isa<Torch::FloatType>() ||
(step.has_value() && (*step).getType().isa<Torch::FloatType>())) {
// TODO: Should get the dtype from torch.get_default_dtype().
// For now, use float32 which is the initial default dtype.
knowledge.dtype = Float32Type::get(op->getContext());
} else
knowledge.dtype =
IntegerType::get(op->getContext(), 64, IntegerType::Signed);
}
incorporateKnowledge(op->getResult(0), knowledge);
}

void TypeAnalysis::visitAtenArangeStartStepOp(AtenArangeStartStepOp op) {
visitAtenArangeLikeOpHelper(op, op.getStart(), op.getEnd(), op.getStep(), op.getDtype());
}

void TypeAnalysis::visitAtenArangeStartOp(AtenArangeStartOp op) {
visitAtenArangeLikeOpHelper(op, op.getStart(), op.getEnd(), {}, op.getDtype());
}

void TypeAnalysis::visitAtenArangeOp(AtenArangeOp op) {
visitAtenArangeLikeOpHelper(op, {}, op.getEnd(), {}, op.getDtype());
}

void TypeAnalysis::visitReductionAlongAllDimsOp(
Operation *op, Type dtype, ArrayRef<const ValueState *> operands) {
auto knowledge =
Expand Down
2 changes: 1 addition & 1 deletion lib/Dialect/Torch/Transforms/ReifyAbstractInterpCalculationsUtils.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -182,7 +182,7 @@ FailureOr<Value> Torch::adjustFunctionArg(
// If the operand is NoneType, then we just need to derefine it to the
// optional type in the function signature.
if (operandType.isa<Torch::NoneType>()) {
assert(desiredType.isa<Torch::OptionalType>() &&
assert(!desiredType.isa<Torch::NoneType>() &&
"Don't expect library functions to have NoneType parameters");
return b.create<DerefineOp>(loc, desiredType, operand).getResult();
}
Expand Down
Loading