[Static Runtime] Add tests for all aten ops (#62347)

Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/62347

This diff includes tests for all `aten` ops that did not already have test coverage.

Test Plan: `buck test //caffe2/benchmarks/static_runtime/static_runtime:static_runtime_cpptest`

Reviewed By: hlu1

Differential Revision: D29968280

fbshipit-source-id: 768655ca535f9e37422711673168dce193de45d2

benchmarks/static_runtime/test_scripts.h

@@ -591,3 +591,81 @@ graph(%a.1 : Tensor,
   %t0 : Float(2, 2, strides=[2, 1], device=cpu), %t1 : Float(3, 3, strides=[3, 1]), %type_matched : bool = prim::TypeCheck[types=[Float(2, 2, strides=[2, 1], device=cpu), Float(3, 3, strides=[3, 1])]](%a.1, %b.1)
   return (%t0, %t1, %type_matched)
 )IR";
+
+const auto index_without_none_script = R"JIT(
+  def forward(self, a: Tensor, idx: Tensor):
+      return a[idx].clone()
+)JIT";
+
+const auto index_with_none_script = R"JIT(
+  def forward(self, a: Tensor, idx: Tensor):
+      return a[idx, None].clone()
+)JIT";
+
+const auto index_with_two_tensors_script = R"JIT(
+  def forward(self, a: Tensor, idx_a: Tensor, idx_b: Tensor):
+      return a[idx_a, idx_b].clone()
+)JIT";
+
+const auto clamp_min_int_script = R"JIT(
+  def forward(self, a: Tensor, b: int):
+      return torch.clamp_min(a, b).clone()
+)JIT";
+
+const auto clamp_min_float_script = R"JIT(
+  def forward(self, a: Tensor, b: float):
+      return torch.clamp_min(a, b).clone()
+)JIT";
+
+const auto argmin_script = R"JIT(
+  def forward(self, a: Tensor):
+      return torch.argmin(a).clone()
+)JIT";
+
+const auto argmin_with_dim_script = R"JIT(
+  def forward(self, a: Tensor, dim: int):
+      return torch.argmin(a, dim).clone()
+)JIT";
+
+const auto argmin_with_keep_dim_script = R"JIT(
+  def forward(self, a: Tensor, dim: int):
+      return torch.argmin(a, dim, True).clone()
+)JIT";
+
+const auto getitem_tensor_script = R"JIT(
+  def forward(self, key: Tensor):
+      d = {key: 1}
+      return d[key]
+)JIT";
+
+const auto getitem_int_script = R"JIT(
+  def forward(self, key: int):
+      d = {key: 1}
+      return d[key]
+)JIT";
+
+const auto getitem_str_script = R"JIT(
+  def forward(self, key: str):
+      d = {key: 1}
+      return d[key]
+)JIT";
+
+const auto transpose_script = R"JIT(
+  def forward(self, a: Tensor, dim1: int, dim2: int):
+      return torch.transpose(a, dim1, dim2).clone()
+)JIT";
+
+const auto permute_script = R"JIT(
+  def forward(self, a: Tensor, dims: List[int]):
+      return torch.permute(a, dims).clone()
+)JIT";
+
+const auto slice_script = R"JIT(
+  def forward(self, a: Tensor, dim: int, start: int, end: int, step: int):
+    return a.slice(dim, start, end, step).clone()
+)JIT";
+
+const auto narrow_with_int_script = R"JIT(
+  def forward(self, a: Tensor, dim: int, start: int, length: int):
+      return a.narrow(dim, start, length).clone()
+)JIT";

benchmarks/static_runtime/test_static_runtime.cc

@@ -967,3 +967,159 @@ TEST(StaticRuntime, IndividualOps_TypeCheck) {
   testStaticRuntime(typecheck_ir, args_correct);
   testStaticRuntime(typecheck_ir, args_correct, args_incorrect);
 }
+
+TEST(StaticRuntime, IndividualOps_Index) {
+  // Index with boolean mask
+  auto a = at::rand({2, 2});
+  auto idx_a = torch::tensor({{0, 1}, {0, 0}}, at::kBool);
+  std::vector<IValue> args_a{a, idx_a};
+
+  // Index with tensor
+  auto b = at::rand({3, 3, 3});
+  auto idx_b = torch::tensor({0, 1, 2}, at::kLong);
+  std::vector<IValue> args_b{b, idx_b};
+
+  testStaticRuntime(index_without_none_script, args_a);
+  testStaticRuntime(index_without_none_script, args_a, args_b);
+
+  // Index with None
+  // When indexing with none, the shape of `a` becomes [2, 1, 2],
+  // so the mask must be reshaped appropriately.
+  auto idx_a_reshape = torch::tensor({{{0, 1}}, {{0, 0}}}, at::kBool);
+  std::vector<IValue> args_a_with_none{a, idx_a_reshape};
+
+  testStaticRuntime(index_with_none_script, args_a_with_none);
+  testStaticRuntime(index_with_none_script, args_a_with_none, args_b);
+
+  // Index with multiple tensors
+  auto c = at::randn({2, 2});
+  auto idx_c1 = torch::tensor({0, 0}, at::kLong);
+  auto idx_c2 = torch::tensor({0}, at::kLong);
+  std::vector<IValue> args_c{c, idx_c1, idx_c2};
+
+  auto d = at::randn({3, 3, 3});
+  auto idx_d1 = torch::tensor({{0, 0}, {0, 1}}, at::kLong);
+  auto idx_d2 = torch::tensor({{1, 1}, {1, 0}}, at::kLong);
+  std::vector<IValue> args_d{d, idx_d1, idx_d2};
+
+  testStaticRuntime(index_with_two_tensors_script, args_c);
+  testStaticRuntime(index_with_two_tensors_script, args_c, args_d);
+}
+
+TEST(StaticRuntime, IndividualOps_ClampMin) {
+  auto a = at::randn({2, 2});
+  auto b = at::randn({3, 3, 3});
+  int scalar_int = 1;
+  float scalar_float = 3.14;
+
+  std::vector<IValue> args_a_int{a, scalar_int};
+  std::vector<IValue> args_b_int{b, scalar_int};
+
+  testStaticRuntime(clamp_min_int_script, args_a_int);
+  testStaticRuntime(clamp_min_int_script, args_a_int, args_b_int);
+
+  std::vector<IValue> args_a_float{a, scalar_float};
+  std::vector<IValue> args_b_float{b, scalar_float};
+
+  testStaticRuntime(clamp_min_float_script, args_a_float);
+  testStaticRuntime(clamp_min_float_script, args_a_float, args_b_float);
+}
+
+TEST(StaticRuntime, IndividualOps_Argmin) {
+  auto a = at::randn({2, 2});
+  auto b = at::randn({3, 3, 3});
+
+  testStaticRuntime(argmin_script, {a});
+  testStaticRuntime(argmin_script, {a}, {b});
+
+  int dim_a = 0;
+  int dim_b = 1;
+
+  std::vector<IValue> args_a{a, dim_a};
+  std::vector<IValue> args_b{b, dim_b};
+
+  testStaticRuntime(argmin_with_dim_script, args_a);
+  testStaticRuntime(argmin_with_dim_script, args_a, args_b);
+
+  testStaticRuntime(argmin_with_keep_dim_script, args_a);
+  testStaticRuntime(argmin_with_keep_dim_script, args_a, args_b);
+}
+
+TEST(StaticRuntime, IndividualOps_GetItem) {
+  int int_key = 0;
+  std::string str_key = "str";
+
+  // No need to test these multiple times, args are not tensors
+  testStaticRuntime(getitem_int_script, {int_key});
+  testStaticRuntime(getitem_str_script, {str_key});
+
+  auto a = torch::tensor({1});
+  auto b = torch::tensor({1, 1});
+
+  testStaticRuntime(getitem_tensor_script, {a});
+  testStaticRuntime(getitem_tensor_script, {a}, {b});
+}
+
+TEST(StaticRuntime, IndividualOps_Transpose) {
+  auto a = at::randn({2, 2});
+  int dim1_a = 0;
+  int dim2_a = 1;
+  std::vector<IValue> args_a{a, dim1_a, dim2_a};
+
+  auto b = at::randn({3, 3, 3});
+  int dim1_b = 0;
+  int dim2_b = 2;
+  std::vector<IValue> args_b{b, dim1_b, dim2_b};
+
+  testStaticRuntime(transpose_script, args_a);
+  testStaticRuntime(transpose_script, args_a, args_b);
+}
+
+TEST(StaticRuntime, IndividualOps_Permute) {
+  auto a = at::randn({2, 2});
+  c10::List<int64_t> dims_a{1, 0};
+  std::vector<IValue> args_a{a, dims_a};
+
+  auto b = at::randn({3, 3, 3});
+  c10::List<int64_t> dims_b{0, 2, 1};
+  std::vector<IValue> args_b{b, dims_b};
+
+  testStaticRuntime(permute_script, args_a);
+  testStaticRuntime(permute_script, args_a, args_b);
+}
+
+TEST(StaticRuntime, IndividualOps_Slice) {
+  auto a = at::randn({2, 2});
+  int dim_a = 1;
+  int start_a = 0;
+  int end_a = 1;
+  int step_a = 1;
+  std::vector<IValue> args_a{a, dim_a, start_a, end_a, step_a};
+
+  auto b = at::randn({3, 3, 3});
+  int dim_b = 2;
+  int start_b = 0;
+  int end_b = 1;
+  int step_b = 2;
+  std::vector<IValue> args_b{b, dim_b, start_b, end_b, step_b};
+
+  testStaticRuntime(slice_script, args_a);
+  testStaticRuntime(slice_script, args_a, args_b);
+}
+
+TEST(StaticRuntime, IndividualOps_Narrow) {
+  auto a = at::randn({5, 5});
+  int dim_a = 0;
+  int start_a_int = 3;
+  int len_a = 2;
+  std::vector<IValue> args_a{a, dim_a, start_a_int, len_a};
+
+  auto b = at::randn({5, 5, 5});
+  int dim_b = 1;
+  int start_b_int = 2;
+  int len_b = 3;
+  std::vector<IValue> args_b{b, dim_b, start_b_int, len_b};
+
+  testStaticRuntime(narrow_with_int_script, args_a);
+  testStaticRuntime(narrow_with_int_script, args_a, args_b);
+}