[port] permute to new api (#172)

* [port] permute to new api

* retrigger CI

* fix errors

Author: kshitij12345

functorch/csrc/BatchRulesViews.cpp

@@ -328,6 +328,24 @@ std::tuple<Tensor,optional<int64_t>> diagonal_backward_batch_rule(
   return std::make_tuple(std::move(result), 0);
 }
 
+std::tuple<Tensor, optional<int64_t>> permute_batching_rule(
+    const Tensor &self, optional<int64_t> self_bdim, IntArrayRef dims)
+{
+  if (!self_bdim.has_value()) {
+    return std::make_tuple(self.permute(dims), self_bdim);
+  }
+
+  auto self_ = moveBatchDimToFront(self, self_bdim);
+  VmapDimVector dims_;
+  dims_.reserve(dims.size() + 1);
+  dims_.emplace_back(0);
+  for (auto dim : dims) {
+    dims_.emplace_back(getPhysicalDim(self_, self_bdim.has_value(), dim));
+  }
+
+  return std::make_tuple(self_.permute(dims_), 0);
+}
+
 std::tuple<Tensor,optional<int64_t>> select_backward_batch_rule(
     const Tensor& grad_input, optional<int64_t> grad_input_bdim,
     IntArrayRef input_sizes, int64_t dim, int64_t index) {
@@ -405,6 +423,7 @@ TORCH_LIBRARY_IMPL(aten, FT_BATCHED_KEY, m) {
   VMAP_SUPPORT("squeeze.dim", squeeze_dim_batch_rule);
   VMAP_SUPPORT("_reshape_alias", _reshape_alias_batch_rule);
   VMAP_SUPPORT("roll", roll_batch_rule);
+  VMAP_SUPPORT("permute", permute_batching_rule);
   VMAP_SUPPORT("diagonal", diagonal_batching_rule);
   VMAP_SUPPORT("diagonal_backward", diagonal_backward_batch_rule);
   VMAP_SUPPORT("select_backward", select_backward_batch_rule);

functorch/csrc/BatchingRegistrations.cpp

@@ -291,28 +291,6 @@ Tensor transpose_int_batching_rule(const Tensor& self, int64_t dim0, int64_t dim
   return self_physical.getPhysicalToLogicalMap().apply(result);
 }
 
-Tensor permute_batching_rule(const Tensor& self, IntArrayRef dims) {
-  if (!participatesInCurrentLevel(self)) {
-    c10::impl::ExcludeDispatchKeyGuard guard(kBatchedKey);
-    return self.permute(dims);
-  }
-
-  auto self_physical = MultiBatchVmapTransform::logicalToPhysical(self);
-  auto dims_physical = self_physical.getPhysicalDims(dims);
-
-  VmapDimVector all_dims_physical;
-  all_dims_physical.reserve(self_physical.tensor().dim());
-  for (int64_t bdim = 0; bdim < self_physical.numBatchDims(); bdim++) {
-    all_dims_physical.push_back(bdim);
-  }
-  all_dims_physical.insert(
-      all_dims_physical.end(),
-      dims_physical.begin(),
-      dims_physical.end());
-  auto result = self_physical.tensor().permute(all_dims_physical);
-  return self_physical.getPhysicalToLogicalMap().apply(result);
-}
-
 static int64_t getGradInputPhysicalDim(int64_t dim, IntArrayRef input_sizes, int64_t num_batch_dims) {
   return maybe_wrap_dim(dim, input_sizes.size()) + num_batch_dims;
 }
@@ -963,7 +941,6 @@ TORCH_LIBRARY_IMPL(aten, FT_BATCHED_KEY, m) {
   // NB: static_cast because there's another variant of narrow. However, we don't
   // want to support the other variant yet bc it isn't documented...
   m.impl("numpy_T", native::numpy_T); // composite wrt autograd
-  m.impl("permute", permute_batching_rule);
   m.impl("reshape_as", native::reshape_as); // composite wrt autograd
   m.impl("slice.Tensor", slice_batching_rule);
   m.impl("split.Tensor", split_batching_rule);