Method Compute

oneflow/user/kernels/transpose_kernel.cpp:60–89

 private:
  void Compute(KernelComputeContext* ctx) const override {
    auto primitive = NewPermutePrimitive(ctx);
    CHECK(primitive);

    const Tensor* tensor_in = ctx->Tensor4ArgNameAndIndex("input", 0);
    Tensor* tensor_out = ctx->Tensor4ArgNameAndIndex("output", 0);
    const auto& perm = ctx->Attr<std::vector<int32_t>>("perm");
    const ShapeView& in_shape = tensor_in->shape_view();
    DataType dtype = tensor_out->data_type();
    size_t num_dims = tensor_in->shape_view().NumAxes();
    const int64_t* src_dims = in_shape.ptr();

    int64_t elem_cnt = tensor_out->shape_view().elem_cnt();

    if (elem_cnt != 0) {
      if (IsIdentity(in_shape, perm)) {
        // if permute vector is 0,1,...,n, do data copy directly
        AutoMemcpy(ctx->stream(), tensor_out->mut_dptr(), tensor_in->dptr(),
                   elem_cnt * GetSizeOfDataType(dtype), tensor_out->mem_case(),
                   tensor_in->mem_case());
      } else {
        primitive->Launch(ctx->stream(), dtype, num_dims, src_dims, tensor_in->dptr(), perm.data(),
                          tensor_out->mut_dptr());
      }

    } else {
      // For 0-d Tensor
      return;
    }
  }
  bool AlwaysComputeWhenAllOutputsEmpty() const override { return false; }
};