Added E2E support for torch::aten.__contains__int_list

include/torch-mlir/Dialect/Torch/IR/GeneratedTorchOps.td

@@ -6160,6 +6160,31 @@ def Torch_Aten__Contains__StrOp : Torch_Op<"aten.__contains__.str", [
   let hasFolder = 1;
 }
 
+def Torch_Aten__Contains__IntListOp : Torch_Op<"aten.__contains__.int_list", [
+    AllowsTypeRefinement,
+    HasValueSemantics,
+    ReadOnly
+  ]> {
+  let summary = "Generated op for `aten::__contains__.int_list : (int[], int) -> (bool)`";
+  let arguments = (ins
+    AnyTorchListOfTorchIntType:$l,
+    Torch_IntType:$item
+  );
+  let results = (outs
+    Torch_BoolType:$result
+  );
+  let hasCustomAssemblyFormat = 1;
+  let extraClassDefinition = [{
+    ParseResult Aten__Contains__IntListOp::parse(OpAsmParser &parser, OperationState &result) {
+      return parseDefaultTorchOp(parser, result, 2, 1);
+    }
+    void Aten__Contains__IntListOp::print(OpAsmPrinter &printer) {
+      printDefaultTorchOp(printer, *this, 2, 1);
+    }
+  }];
+  let hasFolder = 1;
+}
+
 def Torch_Aten__Getitem__DictStrOp : Torch_Op<"aten.__getitem__.Dict_str", [
     AllowsTypeRefinement,
     ReadOnly

lib/Dialect/Torch/IR/TorchOps.cpp

@@ -1584,6 +1584,37 @@ OpFoldResult Aten__Contains__StrOp::fold(ArrayRef<Attribute> operands) {
   return nullptr;
 }
 
+//===----------------------------------------------------------------------===//
+// Aten__Contains__IntListOp
+//===----------------------------------------------------------------------===//
+
+static bool isListConstructNotModified(Value torchList) {
+  return llvm::all_of(torchList.getUsers(), [](Operation *op) {
+        return isa<Aten__Contains__IntListOp>(op);
+      });
+}
+
+OpFoldResult Aten__Contains__IntListOp::fold(ArrayRef<Attribute> operands) {
+  auto itemConstruct = item();
+  if (!isListConstructNotModified(l()))
+    return nullptr;
+
+  int64_t item;
+  SmallVector<int64_t> list;
+
+  if (!matchPattern(itemConstruct, m_TorchConstantInt(&item)))
+    return nullptr;
+
+  if (!matchPattern(l(), m_TorchConstantIntList(list)))
+    return nullptr;
+
+  for (auto elem : list) {
+    if (elem == item)
+      return getI1IntegerAttr(getContext(), true);
+  }
+  return getI1IntegerAttr(getContext(), false);
+}
+
 using BinaryIntOperatorFn = std::function<int64_t(int64_t, int64_t)>;
 template <typename OpTy>
 static OpFoldResult atenBinaryIntOperatorFoldHelper(OpTy op,

python/torch_mlir/dialects/torch/importer/jit_ir/build_tools/torch_ods_gen.py

@@ -472,6 +472,7 @@ def emit_ops(emitter_td: TextEmitter, registry: Registry):
 
     # Dict ops.
     emit("aten::__contains__.str : (Dict(str, t), str) -> (bool)", has_folder=True)
+    emit("aten::__contains__.int_list : (int[], int) -> (bool)", has_folder=True)
     emit("aten::__getitem__.Dict_str : (Dict(str, t), str) -> (t)", has_folder=True)
     emit("aten::_set_item.str : (Dict(str, t), str, t) -> ()")
     emit("aten::keys.str : (Dict(str, t)) -> (str[])")

python/torch_mlir_e2e_test/test_suite/basic.py

@@ -108,6 +108,40 @@ def IsFloatingPointFloat_True(module, tu: TestUtils):
 # ==============================================================================
 
 
+class ContainsIntList(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+    @export
+    @annotate_args([
+        None
+    ])
+    def forward(self):
+        return torch.ops.aten.__contains__([1,2,3], 3)
+@register_test_case(module_factory=lambda: ContainsIntList())
+def ContainsIntList_True(module, tu: TestUtils):
+    module.forward()
+
+
+# ==============================================================================
+
+
+class ContainsIntListFalse(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+    @export
+    @annotate_args([
+        None
+    ])
+    def forward(self):
+        return torch.ops.aten.__contains__([1,2,3], 4)
+@register_test_case(module_factory=lambda: ContainsIntListFalse())
+def ContainsIntList_False(module, tu: TestUtils):
+    module.forward()
+        
+
+# ==============================================================================
+
+
 # A subgraph with multiple mm ops.
 class MmDagModule(torch.nn.Module):