Skip to content

Commit 668c398

Browse files
vivekkhandelwal1vivekkhandelwal1
committed
Make checks generic and disable some to comply with stable torch
1 parent 644f7ab commit 668c398

File tree

9 files changed

+83
-48
lines changed

9 files changed

+83
-48
lines changed

‎lib/Dialect/Torch/Transforms/AbstractInterpLibrary.cpp

+35-6
Original file line numberDiff line numberDiff line change
@@ -15394,6 +15394,9 @@ StringRef mlir::torch::Torch::getAbstractInterpLibrary() {
1539415394
" return %0 : !torch.int\n"
1539515395
" }\n"
1539615396
" func.func @\"__torch_mlir_dtype_fn.aten.linalg_vector_norm\"(%arg0: !torch.tuple<int, int>, %arg1: !torch.number, %arg2: !torch.optional<list<int>>, %arg3: !torch.bool, %arg4: !torch.optional<int>) -> !torch.int {\n"
15397+
" %int6 = torch.constant.int 6\n"
15398+
" %int15 = torch.constant.int 15\n"
15399+
" %int5 = torch.constant.int 5\n"
1539715400
" %true = torch.constant.bool true\n"
1539815401
" %none = torch.constant.none\n"
1539915402
" %str = torch.constant.str \"AssertionError: \"\n"
@@ -15442,12 +15445,22 @@ StringRef mlir::torch::Torch::getAbstractInterpLibrary() {
1544215445
" }\n"
1544315446
" torch.prim.If.yield %9 : !torch.int\n"
1544415447
" } else {\n"
15445-
" %5 = func.call @\"__torch_mlir_dtype_fn.aten.std\"(%arg0, %true) : (!torch.tuple<int, int>, !torch.bool) -> !torch.int\n"
15446-
" torch.prim.If.yield %5 : !torch.int\n"
15448+
" %5 = torch.prim.ListConstruct %int5, %int15 : (!torch.int, !torch.int) -> !torch.list<int>\n"
15449+
" %6 = torch.aten.__contains__.int_list %5, %0#1 : !torch.list<int>, !torch.int -> !torch.bool\n"
15450+
" %7 = torch.prim.If %6 -> (!torch.int) {\n"
15451+
" torch.prim.If.yield %int6 : !torch.int\n"
15452+
" } else {\n"
15453+
" %8 = func.call @\"__torch_mlir_dtype_fn.aten.std\"(%arg0, %true) : (!torch.tuple<int, int>, !torch.bool) -> !torch.int\n"
15454+
" torch.prim.If.yield %8 : !torch.int\n"
15455+
" }\n"
15456+
" torch.prim.If.yield %7 : !torch.int\n"
1544715457
" }\n"
1544815458
" return %4 : !torch.int\n"
1544915459
" }\n"
1545015460
" func.func @\"__torch_mlir_dtype_fn.aten.linalg_norm\"(%arg0: !torch.tuple<int, int>, %arg1: !torch.optional<number>, %arg2: !torch.optional<list<int>>, %arg3: !torch.bool, %arg4: !torch.optional<int>) -> !torch.int {\n"
15461+
" %int6 = torch.constant.int 6\n"
15462+
" %int15 = torch.constant.int 15\n"
15463+
" %int5 = torch.constant.int 5\n"
1545115464
" %true = torch.constant.bool true\n"
1545215465
" %none = torch.constant.none\n"
1545315466
" %str = torch.constant.str \"AssertionError: \"\n"
@@ -15496,8 +15509,15 @@ StringRef mlir::torch::Torch::getAbstractInterpLibrary() {
1549615509
" }\n"
1549715510
" torch.prim.If.yield %9 : !torch.int\n"
1549815511
" } else {\n"
15499-
" %5 = func.call @\"__torch_mlir_dtype_fn.aten.std\"(%arg0, %true) : (!torch.tuple<int, int>, !torch.bool) -> !torch.int\n"
15500-
" torch.prim.If.yield %5 : !torch.int\n"
15512+
" %5 = torch.prim.ListConstruct %int5, %int15 : (!torch.int, !torch.int) -> !torch.list<int>\n"
15513+
" %6 = torch.aten.__contains__.int_list %5, %0#1 : !torch.list<int>, !torch.int -> !torch.bool\n"
15514+
" %7 = torch.prim.If %6 -> (!torch.int) {\n"
15515+
" torch.prim.If.yield %int6 : !torch.int\n"
15516+
" } else {\n"
15517+
" %8 = func.call @\"__torch_mlir_dtype_fn.aten.std\"(%arg0, %true) : (!torch.tuple<int, int>, !torch.bool) -> !torch.int\n"
15518+
" torch.prim.If.yield %8 : !torch.int\n"
15519+
" }\n"
15520+
" torch.prim.If.yield %7 : !torch.int\n"
1550115521
" }\n"
1550215522
" return %4 : !torch.int\n"
1550315523
" }\n"
@@ -15521,6 +15541,8 @@ StringRef mlir::torch::Torch::getAbstractInterpLibrary() {
1552115541
" }\n"
1552215542
" func.func @\"__torch_mlir_dtype_fn.aten.norm.Scalar\"(%arg0: !torch.tuple<int, int>, %arg1: !torch.number) -> !torch.int {\n"
1552315543
" %true = torch.constant.bool true\n"
15544+
" %int6 = torch.constant.int 6\n"
15545+
" %int15 = torch.constant.int 15\n"
1552415546
" %int5 = torch.constant.int 5\n"
1552515547
" %int8 = torch.constant.int 8\n"
1552615548
" %none = torch.constant.none\n"
@@ -15538,8 +15560,15 @@ StringRef mlir::torch::Torch::getAbstractInterpLibrary() {
1553815560
" %4 = torch.prim.If %3 -> (!torch.int) {\n"
1553915561
" torch.prim.If.yield %int5 : !torch.int\n"
1554015562
" } else {\n"
15541-
" %5 = func.call @\"__torch_mlir_dtype_fn.aten.std\"(%arg0, %true) : (!torch.tuple<int, int>, !torch.bool) -> !torch.int\n"
15542-
" torch.prim.If.yield %5 : !torch.int\n"
15563+
" %5 = torch.prim.ListConstruct %int5, %int15 : (!torch.int, !torch.int) -> !torch.list<int>\n"
15564+
" %6 = torch.aten.__contains__.int_list %5, %0#1 : !torch.list<int>, !torch.int -> !torch.bool\n"
15565+
" %7 = torch.prim.If %6 -> (!torch.int) {\n"
15566+
" torch.prim.If.yield %int6 : !torch.int\n"
15567+
" } else {\n"
15568+
" %8 = func.call @\"__torch_mlir_dtype_fn.aten.std\"(%arg0, %true) : (!torch.tuple<int, int>, !torch.bool) -> !torch.int\n"
15569+
" torch.prim.If.yield %8 : !torch.int\n"
15570+
" }\n"
15571+
" torch.prim.If.yield %7 : !torch.int\n"
1554315572
" }\n"
1554415573
" return %4 : !torch.int\n"
1554515574
" }\n"

‎projects/pt1/python/torch_mlir/jit_ir_importer/build_tools/abstract_interp_lib_gen.py

+6
Original file line numberDiff line numberDiff line change
@@ -5291,6 +5291,8 @@ def aten〇linalg_vector_norm〡dtype(self_rank_dtype: Tuple[int, int], ord: Uni
52915291
return aten〇std〡dtype((self_rank, dtype))
52925292
assert not is_complex_dtype(dtype)
52935293
return dtype
5294+
if self_dtype in [torch.float16, torch.bfloat16]:
5295+
return torch.float32
52945296
return aten〇std〡dtype(self_rank_dtype)
52955297

52965298
@check_dtype_function(
@@ -5314,6 +5316,8 @@ def aten〇linalg_norm〡dtype(self_rank_dtype: Tuple[int, int], ord: Optional[U
53145316
return aten〇std〡dtype((self_rank, dtype))
53155317
assert not is_complex_dtype(dtype)
53165318
return dtype
5319+
if self_dtype in [torch.float16, torch.bfloat16]:
5320+
return torch.float32
53175321
return aten〇std〡dtype(self_rank_dtype)
53185322

53195323
def aten〇binary_cross_entropy_with_logits〡dtype(self_rank_dtype: Tuple[int, int], target_rank_dtype: Tuple[int, int], weight_rank_dtype: Optional[Tuple[int, int]] = None, pos_weight_rank_dtype: Optional[Tuple[int, int]] = None, reduction: int = 1) -> int:
@@ -5347,6 +5351,8 @@ def aten〇norm〇Scalar〡dtype(self_rank_dtype: Tuple[int, int], p: Union[int,
53475351
# Should possibly be added to aten〇std〡dtype.
53485352
if self_dtype == torch.complex32:
53495353
return torch.half
5354+
if self_dtype in [torch.float16, torch.bfloat16]:
5355+
return torch.float32
53505356
return aten〇std〡dtype(self_rank_dtype)
53515357

53525358
@check_dtype_function([Invocation(0.0),

‎pytorch-hash.txt

+1-1
Original file line numberDiff line numberDiff line change
@@ -1 +1 @@
1-
881e8b6bda5909486ea9225380033b3c5fba4b37
1+
dab7e5700392e4e20626de9c367acb76187807f5

‎pytorch-requirements.txt

+1-1
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,3 @@
11
-f https://download.pytorch.org/whl/nightly/cpu/torch/
22
--pre
3-
torch==2.8.0.dev20250406
3+
torch==2.8.0.dev20250423

‎test/python/fx_importer/basic_test.py

+5-5
Original file line numberDiff line numberDiff line change
@@ -88,11 +88,11 @@ def forward(self, x):
8888
@run
8989
# CHECK-LABEL: test_import_frozen_exported_program_with_dynamic_shapes
9090
# CHECK: func.func @test_net(%[[ARG0:[a-zA-Z0-9]+]]: !torch.vtensor<[?,?,5],f32>) -> !torch.vtensor<[?,?,5],f32>
91-
# CHECK: %[[S0:.*]] = torch.symbolic_int "s35" {min_val = {{[0-9]+}}, max_val = {{[0-9]+}}} : !torch.int
92-
# CHECK: %[[S1:.*]] = torch.symbolic_int "s16" {min_val = 2, max_val = {{[0-9]+}}} : !torch.int
93-
# CHECK: torch.bind_symbolic_shape %[[ARG0]], [%[[S1]], %[[S0]]], affine_map<()[s0, s1] -> (s1, s0, 5)> : !torch.vtensor<[?,?,5],f32>
91+
# CHECK: %[[S0:.*]] = torch.symbolic_int "{{[a-z0-9]+}}" {min_val = {{[0-9]+}}, max_val = {{[0-9]+}}} : !torch.int
92+
# CHECK: %[[S1:.*]] = torch.symbolic_int "{{[a-z0-9]+}}" {min_val = 2, max_val = {{[0-9]+}}} : !torch.int
93+
# CHECK-DISABLED: torch.bind_symbolic_shape %[[ARG0]], [%[[S1]], %[[S0]]], affine_map<()[s0, s1] -> (s1, s0, 5)> : !torch.vtensor<[?,?,5],f32>
9494
# CHECK: %[[TANH:.*]] = torch.aten.tanh %[[ARG0]] : !torch.vtensor<[?,?,5],f32> -> !torch.vtensor<[?,?,5],f32>
95-
# CHECK: torch.bind_symbolic_shape %[[TANH]], [%[[S1]], %[[S0]]], affine_map<()[s0, s1] -> (s1, s0, 5)> : !torch.vtensor<[?,?,5],f32>
95+
# CHECK-DISABLED: torch.bind_symbolic_shape %[[TANH]], [%[[S1]], %[[S0]]], affine_map<()[s0, s1] -> (s1, s0, 5)> : !torch.vtensor<[?,?,5],f32>
9696
# CHECK: return %[[TANH]] : !torch.vtensor<[?,?,5],f32>
9797
def test_import_frozen_exported_program_with_dynamic_shapes():
9898
class Basic(nn.Module):
@@ -118,7 +118,7 @@ def forward(self, x):
118118
@run
119119
# CHECK-LABEL: test_broadcast_with_dynamic_shapes
120120
# CHECK: func.func @test_net(%[[ARG0:[a-zA-Z0-9]+]]: !torch.vtensor<[1,2],f32>, %[[ARG1:[a-zA-Z0-9]+]]: !torch.vtensor<[?],f32>) -> !torch.vtensor<[?,2],f32>
121-
# CHECK: %[[S0:.*]] = torch.symbolic_int "s58" {min_val = {{[0-9]+}}, max_val = {{[0-9]+}}} : !torch.int
121+
# CHECK: %[[S0:.*]] = torch.symbolic_int "{{[a-z0-9]+}}" {min_val = {{[0-9]+}}, max_val = {{[0-9]+}}} : !torch.int
122122
# CHECK: torch.bind_symbolic_shape %[[ARG1]], [%[[S0]]], affine_map<()[s0] -> (s0)> : !torch.vtensor<[?],f32>
123123
# CHECK: torch.aten.size.int
124124
# CHECK: torch.prim.ListConstruct

‎test/python/fx_importer/custom_op_test.py

+8-8
Original file line numberDiff line numberDiff line change
@@ -26,15 +26,15 @@ def run(f):
2626
# CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]: !torch.vtensor<[?,?,3],f32>,
2727
# CHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]]: !torch.vtensor<[?,?,3],f32>,
2828
# CHECK-SAME: %[[ARG2:[a-zA-Z0-9]+]]: !torch.vtensor<[?,?,3],f32>) -> !torch.vtensor<[?,?,3],f32> {
29-
# CHECK: %[[S0:.+]] = torch.symbolic_int "s35" {min_val = 5, max_val = 10} : !torch.int
30-
# CHECK: %[[S1:.+]] = torch.symbolic_int "s16" {min_val = {{[0-9]+}}, max_val = 100} : !torch.int
31-
# CHECK: %[[S2:.+]] = torch.symbolic_int "s43" {min_val = {{[0-9]+}}, max_val = 50} : !torch.int
32-
# CHECK: %[[S3:.+]] = torch.symbolic_int "s23" {min_val = {{[0-9]+}}, max_val = {{[0-9]+}}} : !torch.int
33-
# CHECK: torch.bind_symbolic_shape %[[ARG0]], [%[[S1]], %[[S0]]], affine_map<()[s0, s1] -> (s1, s0, 3)> : !torch.vtensor<[?,?,3],f32>
29+
# CHECK: %[[S0:.+]] = torch.symbolic_int "{{[a-z0-9]+}}" {min_val = 5, max_val = 10} : !torch.int
30+
# CHECK: %[[S1:.+]] = torch.symbolic_int "{{[a-z0-9]+}}" {min_val = {{[0-9]+}}, max_val = 100} : !torch.int
31+
# CHECK: %[[S2:.+]] = torch.symbolic_int "{{[a-z0-9]+}}" {min_val = {{[0-9]+}}, max_val = 50} : !torch.int
32+
# CHECK: %[[S3:.+]] = torch.symbolic_int "{{[a-z0-9]+}}" {min_val = {{[0-9]+}}, max_val = {{[0-9]+}}} : !torch.int
33+
# CHECK-DISABLED: torch.bind_symbolic_shape %[[ARG0]], [%[[S1]], %[[S0]]], affine_map<()[s0, s1] -> (s1, s0, 3)> : !torch.vtensor<[?,?,3],f32>
3434
# CHECK: torch.bind_symbolic_shape %[[ARG1]], [%[[S0]], %[[S2]]], affine_map<()[s0, s1] -> (s0, s1, 3)> : !torch.vtensor<[?,?,3],f32>
35-
# CHECK: torch.bind_symbolic_shape %[[ARG2]], [%[[S3]], %[[S0]]], affine_map<()[s0, s1] -> (s1, s0, 3)> : !torch.vtensor<[?,?,3],f32>
35+
# CHECK-DISABLED: torch.bind_symbolic_shape %[[ARG2]], [%[[S3]], %[[S0]]], affine_map<()[s0, s1] -> (s1, s0, 3)> : !torch.vtensor<[?,?,3],f32>
3636
# CHECK: %[[OP:.+]] = torch.operator "torch.my_custom_library.tanh_sigmoid_cat_op"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) : (!torch.vtensor<[?,?,3],f32>, !torch.vtensor<[?,?,3],f32>, !torch.vtensor<[?,?,3],f32>) -> !torch.vtensor<[?,?,3],f32>
37-
# CHECK: torch.bind_symbolic_shape %[[OP]], [%[[S1]], %[[S3]], %[[S0]], %[[S2]]], affine_map<()[s0, s1, s2, s3] -> (s2, s1 + s3 + s0 * 2, 3)> : !torch.vtensor<[?,?,3],f32>
37+
# CHECK-DISABLED: torch.bind_symbolic_shape %[[OP]], [%[[S1]], %[[S3]], %[[S0]], %[[S2]]], affine_map<()[s0, s1, s2, s3] -> (s2, s1 + s3 + s0 * 2, 3)> : !torch.vtensor<[?,?,3],f32>
3838
# CHECK: return %[[OP]] : !torch.vtensor<[?,?,3],f32>
3939
def test_tanh_sigmoid_cat_custom_op():
4040

@@ -89,7 +89,7 @@ def forward(self, x, y, z):
8989
@run
9090
# CHECK-LABEL: test_custom_op_array_output
9191
# CHECK: func.func @main(%[[ARG0:[a-zA-Z0-9]+]]: !torch.vtensor<[?,3],f32>)
92-
# CHECK: %[[S0:.+]] = torch.symbolic_int "s35" {min_val = {{[0-9]+}}, max_val = 10} : !torch.int
92+
# CHECK: %[[S0:.+]] = torch.symbolic_int "{{[a-z0-9]+}}" {min_val = {{[0-9]+}}, max_val = 10} : !torch.int
9393
# CHECK: %[[int:.+]] = torch.constant.int 4
9494
# CHECK: %[[V0:.+]] = torch.operator "torch.my_custom_library.array_output_op"(%[[int]], %[[ARG0]]) : (!torch.int, !torch.vtensor<[?,3],f32>) -> !torch.list<vtensor>
9595
# CHECK: %[[V1:.+]]:4 = torch.prim.ListUnpack %[[V0]] : !torch.list<vtensor> -> !torch.vtensor<[?,3],f32>, !torch.vtensor<[?,3],f32>, !torch.vtensor<[?,3],f32>, !torch.vtensor<[?,3],f32>

0 commit comments

Comments
 (0)