[MLIR][TORCH] Add TorchToTosa lowering for aten.where.self op

e2e_testing/xfail_sets.py

@@ -463,6 +463,7 @@
     "ArangeStartIntModule_basic",
     "ArangeStartNegativeStepIntModule_basic",
     "ArangeZeroElementOutputModule_basic",
+    "ElementwiseAtenWhereSelfModule_basic",
 }
 
 LTC_XFAIL_SET = {

lib/Conversion/TorchToTosa/TorchToTosa.cpp

@@ -3005,6 +3005,30 @@ LogicalResult ConvertAtenOp<AtenBroadcastToOp>::matchAndRewrite(
       "unimplemented: broadcasts other than same rank or zero ranked tensor.");
 }
 
+
+template <>
+LogicalResult ConvertAtenOp<AtenWhereSelfOp>::matchAndRewrite(
+    AtenWhereSelfOp op, OpAdaptor adaptor,
+    ConversionPatternRewriter &rewriter) const {
+
+  // Not a tensor type.
+  auto selfType = adaptor.self().getType().dyn_cast<TensorType>();
+  if (!selfType)
+    return rewriter.notifyMatchFailure(
+        op, "Only tensor types input are currently supported");
+  auto condType = adaptor.condition().getType().dyn_cast<TensorType>();
+  if (!condType)
+    return rewriter.notifyMatchFailure(
+        op, "Only tensor types condition are currently supported");
+
+  auto outType = getTypeConverter()->convertType(op.getType());
+  rewriter.replaceOpWithNewOp<tosa::SelectOp>(op, outType, adaptor.condition(),
+                                              adaptor.self(), adaptor.other());
+
+  return success();
+}
+
+
 template <>
 LogicalResult ConvertAtenOp<AtenArangeStartStepOp>::matchAndRewrite(
     AtenArangeStartStepOp op, OpAdaptor adaptor,
@@ -3703,6 +3727,7 @@ class ConvertTorchToTosa : public ConvertTorchToTosaBase<ConvertTorchToTosa> {
     INSERT_ATENOP_PATTERN(AtenMaxDimOp);
     INSERT_ATENOP_PATTERN(AtenSliceTensorOp);
     INSERT_ATENOP_PATTERN(AtenBroadcastToOp);
+    INSERT_ATENOP_PATTERN(AtenWhereSelfOp);
     INSERT_ATENOP_PATTERN(AtenArangeStartStepOp);
 #undef INSERT_ATENOP_PATTERN
 

python/torch_mlir_e2e_test/test_suite/elementwise.py

@@ -134,6 +134,30 @@ def ElementwiseTernaryModule_basic(module, tu: TestUtils):
 # ==============================================================================
 
 
+class ElementwiseAtenWhereSelfModule(torch.nn.Module):
+
+    def __init__(self):
+        super().__init__()
+
+    @export
+    @annotate_args([
+        None,
+        ([1, 1, 5, 5], torch.bool, True),
+        ([1, 12, 5, 5], torch.float32, True),
+        ([], torch.float32, True),
+    ])
+    def forward(self, a, b, c):
+        return torch.ops.aten.where(a, b, c)
+
+
+@register_test_case(module_factory=lambda: ElementwiseAtenWhereSelfModule())
+def ElementwiseAtenWhereSelfModule_basic(module, tu: TestUtils):
+    module.forward(torch.zeros(1, 1, 5, 5, dtype=torch.bool), torch.rand(1, 12, 5, 5), torch.rand(()))
+
+
+# ==============================================================================
+
+
 class ElementwiseWhereSelfModule(torch.nn.Module):
 
     def __init__(self):