more lit tests

cruise-automation · Sep 17, 2024 · e965953 · e965953
1 parent 17e1ad2
commit e965953
Showing 1 changed file with 19 additions and 0 deletions.
diff --git a/test/Pipeline/torch_to_tcp_pipeline.mlir b/test/Pipeline/torch_to_tcp_pipeline.mlir
@@ -116,3 +116,22 @@ func.func @torch.aten.div.Tensor$mixed_type_int(%arg0: !torch.vtensor<[?, ?],si1
   %0 = torch.aten.div.Tensor %arg0, %arg1 : !torch.vtensor<[?, ?],si16>, !torch.vtensor<[?, ?],si32> -> !torch.vtensor<[?, ?],si32>
   return %0 : !torch.vtensor<[?, ?],si32>
 }
+
+// -----
+
+// CHECK:   func.func @torch.aten.div.Tensor$mixed_type_uint(%[[ARG0:.+]]: tensor<?x?xi16>, %[[ARG1:.+]]: tensor<?x?xi32>) -> tensor<?x?xi32> {
+// CHECK:     %[[V0:.+]] = tcp.cast %[[ARG0]] {in_int_signedness = #tcp<signedness Unsigned>, out_int_signedness = #tcp<signedness Unsigned>} : tensor<?x?xi16> -> tensor<?x?xi32>
+// CHECK:     %[[V1:.+]] = tcp.divui %[[V0]], %[[ARG1]] {rounding_mode = #tcp<roundingMode Trunc>} : tensor<?x?xi32>, tensor<?x?xi32> -> tensor<?x?xi32>
+// CHECK:     return %[[V1]] : tensor<?x?xi32>
+func.func @torch.aten.div.Tensor$mixed_type_uint(%arg0: !torch.vtensor<[?, ?],ui16>, %arg1: !torch.vtensor<[?, ?],ui32>) -> !torch.vtensor<[?, ?],ui32> {
+  %0 = torch.aten.div.Tensor %arg0, %arg1 : !torch.vtensor<[?, ?],ui16>, !torch.vtensor<[?, ?],ui32> -> !torch.vtensor<[?, ?],ui32>
+  return %0 : !torch.vtensor<[?, ?],ui32>
+}
+
+// -----
+
+func.func @torch.aten.div.Tensor$mixed_signed_int_div(%arg0: !torch.vtensor<[?, ?],si16>, %arg1: !torch.vtensor<[?, ?],ui32>) -> !torch.vtensor<[?, ?],ui32> {
+  // expected-error @below {{failed to legalize operation 'torch.aten.div.Tensor' that was explicitly marked illegal}}
+  %0 = torch.aten.div.Tensor %arg0, %arg1 : !torch.vtensor<[?, ?],si16>, !torch.vtensor<[?, ?],ui32> -> !torch.vtensor<[?, ?],ui32>
+  return %0 : !torch.vtensor<[?, ?],ui32>
+}