mirror of https://github.com/llvm/torch-mlir
[TORCH][MLIR] Add E2E support for [`aten.gt.Scalar`|`aten.where.self`]
This commit adds lowering of `aten.gt.Scalar` and `aten.where.self` as a part of element-wise ops lowering. Signed-Off-by: Gaurav Shukla <gaurav@nod-labs.com>pull/469/head snapshot-20211209.133
parent
2414bdb1f0
commit
f34eb66124
|
@ -85,6 +85,29 @@ def ElementwiseTernaryModule_basic(module, tu: TestUtils):
|
|||
# ==============================================================================
|
||||
|
||||
|
||||
class ElementwiseWhereSelfModule(torch.nn.Module):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
|
||||
@export
|
||||
@annotate_args([
|
||||
None,
|
||||
([-1, -1, -1], torch.float32, True),
|
||||
([-1, -1], torch.float32, True),
|
||||
([-1], torch.float32, True),
|
||||
])
|
||||
def forward(self, a, b, c):
|
||||
return torch.where(a > 0.5, b, c)
|
||||
|
||||
|
||||
@register_test_case(module_factory=lambda: ElementwiseWhereSelfModule())
|
||||
def ElementwiseWhereSelfModule_basic(module, tu: TestUtils):
|
||||
module.forward(tu.rand(3, 4, 5), tu.rand(4, 5), tu.rand(5))
|
||||
|
||||
|
||||
# ==============================================================================
|
||||
|
||||
|
||||
# Addition is an interesting special case of a binary op, because under the hood
|
||||
# it carries a third scalar "alpha" parameter, which needs special handling.
|
||||
class ElementwiseAddModule(torch.nn.Module):
|
||||
|
@ -303,6 +326,26 @@ def ElementwiseMaximumModule_basic(module, tu: TestUtils):
|
|||
# ==============================================================================
|
||||
|
||||
|
||||
class ElementwiseGtScalarModule(torch.nn.Module):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
|
||||
@export
|
||||
@annotate_args([
|
||||
None,
|
||||
([-1, -1], torch.float32, True),
|
||||
])
|
||||
def forward(self, x):
|
||||
return torch.gt(x, 0.6)
|
||||
|
||||
|
||||
@register_test_case(module_factory=lambda: ElementwiseGtScalarModule())
|
||||
def ElementwiseGtScalarModule_basic(module, tu: TestUtils):
|
||||
module.forward(tu.rand(3, 5))
|
||||
|
||||
# ==============================================================================
|
||||
|
||||
|
||||
class ElementwiseClampModule(torch.nn.Module):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
|
|
|
@ -1071,6 +1071,22 @@ def Torch_AtenMaximumOp : Torch_Op<"aten.maximum", [
|
|||
let assemblyFormat = "$self `,` $other attr-dict `:` type($self) `,` type($other) `->` type($result)";
|
||||
}
|
||||
|
||||
def Torch_AtenWhereSelfOp : Torch_Op<"aten.where.self", [
|
||||
AllowsTypeRefinement,
|
||||
HasValueSemantics
|
||||
]> {
|
||||
let summary = "Generated op for `aten::where.self : (Tensor, Tensor, Tensor) -> (Tensor)`";
|
||||
let arguments = (ins
|
||||
AnyTorchTensorType:$condition,
|
||||
AnyTorchTensorType:$self,
|
||||
AnyTorchTensorType:$other
|
||||
);
|
||||
let results = (outs
|
||||
AnyTorchTensorType:$result
|
||||
);
|
||||
let assemblyFormat = "$condition `,` $self `,` $other attr-dict `:` type($condition) `,` type($self) `,` type($other) `->` type($result)";
|
||||
}
|
||||
|
||||
def Torch_AtenMinimumOp : Torch_Op<"aten.minimum", [
|
||||
AllowsTypeRefinement,
|
||||
HasValueSemantics
|
||||
|
|
|
@ -1684,6 +1684,27 @@ static Value createLinalgPayloadCalculationForElementwiseOp(
|
|||
Value expPromoted = convertScalarToDtype(b, loc, operands[1], dtype);
|
||||
return b.create<math::PowFOp>(loc, payloadArgs[0], expPromoted);
|
||||
}
|
||||
|
||||
if (auto gtScalar = dyn_cast<AtenGtScalarOp>(op)) {
|
||||
Type dtype = gtScalar.self().getType().cast<ValueTensorType>().getDtype();
|
||||
if (!dtype.isa<mlir::FloatType>()) {
|
||||
gtScalar.emitError("unimplemented: non-floating point operand dtype");
|
||||
return nullptr;
|
||||
}
|
||||
Value otherPromoted = convertScalarToDtype(b, loc, operands[1], dtype);
|
||||
return b.create<arith::CmpFOp>(loc, arith::CmpFPredicate::UGT,
|
||||
payloadArgs[0], otherPromoted);
|
||||
}
|
||||
|
||||
if (auto whereSelf = dyn_cast<AtenWhereSelfOp>(op)) {
|
||||
Type dtype = converter->convertType(whereSelf.getType())
|
||||
.cast<RankedTensorType>()
|
||||
.getElementType();
|
||||
Value lhs = convertScalarToDtype(b, loc, payloadArgs[1], dtype);
|
||||
Value rhs = convertScalarToDtype(b, loc, payloadArgs[2], dtype);
|
||||
return b.create<SelectOp>(loc, payloadArgs[0], lhs, rhs);
|
||||
}
|
||||
|
||||
if (auto lerp = dyn_cast<AtenLerpTensorOp>(op)) {
|
||||
if (!lerp.getType()
|
||||
.cast<ValueTensorType>()
|
||||
|
@ -2040,7 +2061,7 @@ struct ConvertElementwiseOp : ConversionPattern {
|
|||
AtenClampOp, AtenRsubScalarOp, AtenMulScalarOp, AtenLogOp,
|
||||
AtenSqrtOp, AtenFloorOp, AtenPowTensorScalarOp, AtenLog2Op,
|
||||
AtenRsqrtOp, AtenDivScalarOp, AtenAbsOp, AtenReciprocalOp,
|
||||
AtenBitwiseAndTensorOp>(op))
|
||||
AtenBitwiseAndTensorOp, AtenGtScalarOp, AtenWhereSelfOp>(op))
|
||||
return rewriter.notifyMatchFailure(op, "not a supported elementwise op");
|
||||
|
||||
if (failed(verifyLinalgCompatibleTypes(op, rewriter)))
|
||||
|
@ -3461,7 +3482,8 @@ public:
|
|||
AtenLerpTensorOp, AtenSigmoidOp, AtenMinimumOp, AtenMaximumOp,
|
||||
AtenToDtypeOp, AtenClampOp, AtenRsubScalarOp, AtenLogOp, AtenSqrtOp,
|
||||
AtenFloorOp, AtenPowTensorScalarOp, AtenLog2Op, AtenRsqrtOp, AtenAbsOp,
|
||||
AtenReciprocalOp, AtenBitwiseAndTensorOp>();
|
||||
AtenReciprocalOp, AtenBitwiseAndTensorOp, AtenGtScalarOp,
|
||||
AtenWhereSelfOp>();
|
||||
patterns.add<ConvertElementwiseOp>(typeConverter, context);
|
||||
target.addIllegalOp<AtenSqueezeOp>();
|
||||
patterns.add<ConvertAtenSqueezeOp>(typeConverter, context);
|
||||
|
|
|
@ -235,9 +235,8 @@ public:
|
|||
ArrayRef<LatticeElement<ValueKnowledge> *> operands) final {
|
||||
if (isa<TensorStaticInfoCastOp, CopyToValueTensorOp, CopyToNonValueTensorOp,
|
||||
AtenTanhOp, AtenBatchNormOp, AtenReluOp, AtenGeluOp,
|
||||
AtenGeluBackwardOp, AtenEqScalarOp, AtenGeScalarOp, AtenGtScalarOp,
|
||||
AtenNeScalarOp, AtenBitwiseNotOp, AtenExpOp, AtenSinOp, AtenCosOp,
|
||||
AtenSigmoidOp, DerefineOp, AtenToPrimDeviceOp, AtenCpuOp,
|
||||
AtenGeluBackwardOp, AtenBitwiseNotOp, AtenExpOp, AtenSinOp,
|
||||
AtenCosOp, AtenSigmoidOp, DerefineOp, AtenToPrimDeviceOp, AtenCpuOp,
|
||||
AtenContiguousOp, AtenFill_ScalarOp, AtenDetachOp,
|
||||
AtenMaskedFill_ScalarOp, AtenCopy_Op, AtenIndexPut_Op, AtenCumsumOp,
|
||||
AtenLayerNormOp, AtenClampOp, AtenLogOp, AtenNegOp, AtenSqrtOp,
|
||||
|
@ -247,6 +246,20 @@ public:
|
|||
return getLatticeElement(op->getResult(0)).join(*operands[0]);
|
||||
}
|
||||
|
||||
// These comparison ops return a tensor with 1-bit integer dtype.
|
||||
if (isa<AtenEqScalarOp, AtenGeScalarOp, AtenGtScalarOp, AtenNeScalarOp>(
|
||||
op)) {
|
||||
auto operand = operands[0]->getValue();
|
||||
auto knowledge =
|
||||
ValueKnowledge::getNotNonePessimisticValueState(op->getContext());
|
||||
if (operand.hasSizes) {
|
||||
knowledge.hasSizes = true;
|
||||
knowledge.sizes = operand.sizes;
|
||||
}
|
||||
knowledge.dtype = IntegerType::get(op->getContext(), 1);
|
||||
return getLatticeElement(op->getResult(0)).join(knowledge);
|
||||
}
|
||||
|
||||
// Resize to [1, 1] with integer dtype.
|
||||
if (isa<AtenAnyOp, AtenAllOp>(op)) {
|
||||
auto input = operands[0]->getValue();
|
||||
|
@ -307,6 +320,8 @@ public:
|
|||
AtenDivTensorOp, Aten__And__TensorOp, AtenEqTensorOp,
|
||||
AtenMinimumOp, AtenMaximumOp, AtenBitwiseAndTensorOp>(op)) {
|
||||
return visitBinaryBroadcastingOp(op, operands);
|
||||
} else if (auto whereSelf = llvm::dyn_cast<AtenWhereSelfOp>(op)) {
|
||||
return visitAtenWhereSelfOp(whereSelf, operands);
|
||||
} else if (auto lerpTensor = llvm::dyn_cast<AtenLerpTensorOp>(op)) {
|
||||
return visitAtenLerpTensorOp(lerpTensor, operands);
|
||||
} else if (auto flatten = dyn_cast<AtenFlattenUsingIntsOp>(op)) {
|
||||
|
@ -487,6 +502,9 @@ private:
|
|||
ChangeResult visitBinaryBroadcastingOp(
|
||||
Operation *op, ArrayRef<LatticeElement<ValueKnowledge> *> operands);
|
||||
ChangeResult
|
||||
visitAtenWhereSelfOp(AtenWhereSelfOp op,
|
||||
ArrayRef<LatticeElement<ValueKnowledge> *> operands);
|
||||
ChangeResult
|
||||
visitAtenLerpTensorOp(AtenLerpTensorOp op,
|
||||
ArrayRef<LatticeElement<ValueKnowledge> *> operands);
|
||||
ChangeResult visitAtenFlattenUsingIntsOp(
|
||||
|
@ -856,6 +874,25 @@ ChangeResult TypeAnalyzer::visitBinaryBroadcastingOp(
|
|||
return getLatticeElement(op->getResult(0)).join(knowledge);
|
||||
}
|
||||
|
||||
ChangeResult TypeAnalyzer::visitAtenWhereSelfOp(
|
||||
AtenWhereSelfOp op, ArrayRef<LatticeElement<ValueKnowledge> *> operands) {
|
||||
auto condition = operands[0]->getValue();
|
||||
auto lhs = operands[1]->getValue();
|
||||
auto rhs = operands[2]->getValue();
|
||||
auto knowledge =
|
||||
ValueKnowledge::getNotNonePessimisticValueState(getContext());
|
||||
if (condition.hasSizes && lhs.hasSizes && rhs.hasSizes) {
|
||||
knowledge.hasSizes = true;
|
||||
knowledge.sizes.resize(
|
||||
std::max(condition.sizes.size(),
|
||||
std::max(lhs.sizes.size(), rhs.sizes.size())),
|
||||
kUnknownSize);
|
||||
}
|
||||
|
||||
knowledge.dtype = getPromotedResultType(getContext(), {&lhs, &rhs});
|
||||
return getLatticeElement(op->getResult(0)).join(knowledge);
|
||||
}
|
||||
|
||||
ChangeResult TypeAnalyzer::visitAtenLerpTensorOp(
|
||||
AtenLerpTensorOp op, ArrayRef<LatticeElement<ValueKnowledge> *> operands) {
|
||||
// This is a general broadcasting shape transfer function.
|
||||
|
|
|
@ -479,6 +479,7 @@ def emit_aten_ops(torch_ir_dir: str, registry: Registry):
|
|||
emit("aten::addcmul : (Tensor, Tensor, Tensor, Scalar) -> (Tensor)")
|
||||
emit("aten::addcdiv : (Tensor, Tensor, Tensor, Scalar) -> (Tensor)")
|
||||
emit("aten::maximum : (Tensor, Tensor) -> (Tensor)")
|
||||
emit("aten::where.self : (Tensor, Tensor, Tensor) -> (Tensor)")
|
||||
emit("aten::minimum : (Tensor, Tensor) -> (Tensor)")
|
||||
emit("aten::rsub.Scalar : (Tensor, Scalar, Scalar) -> (Tensor)")
|
||||
emit("aten::gelu : (Tensor) -> (Tensor)")
|
||||
|
|
Loading…
Reference in New Issue