Ejemplo n.º 1
0
    static void instantiate(const callable_type_data *DYND_UNUSED(self), const callable_type *DYND_UNUSED(self_tp),
                            kernel_builder *ckb, intptr_t ckb_offset, const ndt::type &DYND_UNUSED(dst_tp),
                            const char *dst_arrmeta, intptr_t DYND_UNUSED(nsrc), const ndt::type *src_tp,
                            const char *const *src_arrmeta, kernel_request_t kernreq,
                            const eval::eval_context *DYND_UNUSED(ectx), const nd::array &kwds,
                            const std::map<std::string, ndt::type> &DYND_UNUSED(tp_vars))
    {
      const size_stride_t *dst_size_stride = reinterpret_cast<const size_stride_t *>(dst_arrmeta);
      const size_stride_t *src_size_stride = reinterpret_cast<const size_stride_t *>(src_arrmeta[0]);

      array axes = kwds.p("axes");
      array shape = kwds.p("shape");

      int ndim = static_cast<int>(src_tp[0].get_ndim());

      int rank = static_cast<int>(axes.is_missing() ? ndim : (ndim - 1));
      int istride = static_cast<int>(src_size_stride[ndim - 1].stride / sizeof(src_type));
      int idist = static_cast<int>(src_size_stride[0].stride / sizeof(src_type));
      int ostride = static_cast<int>(dst_size_stride[ndim - 1].stride / sizeof(dst_type));
      int odist = static_cast<int>(dst_size_stride[0].stride / sizeof(dst_type));

      std::vector<int> n(rank), inembed(rank), onembed(rank);
      for (int i = 0, j = axes.is_missing() ? 0 : 1; j < ndim; ++i, ++j) {
        n[i] = static_cast<int>(src_size_stride[j].dim_size);
        inembed[i] = static_cast<int>(src_size_stride[j].dim_size);
        onembed[i] = static_cast<int>(dst_size_stride[j].dim_size);
      }

      int batch = static_cast<int>(axes.is_missing() ? 1 : src_size_stride[0].dim_size);

      self_type *self = self_type::create(ckb, kernreq, ckb_offset);
      cufftPlanMany(&self->plan, rank, n.data(), inembed.data(), istride, idist, onembed.data(), ostride, odist,
                    CUFFT_Z2Z, batch);
    }
Ejemplo n.º 2
0
    static int resolve_dst_type(const callable_type_data *DYND_UNUSED(self), const callable_type *DYND_UNUSED(self_tp),
                                char *DYND_UNUSED(data), intptr_t DYND_UNUSED(nsrc), const ndt::type *src_tp,
                                int DYND_UNUSED(throw_on_error), ndt::type &dst_tp, const nd::array &kwds,
                                const std::map<std::string, ndt::type> &DYND_UNUSED(tp_vars))
    {
      nd::array shape = kwds.p("shape");
      //      if (shape.is_missing()) {
      dst_tp = src_tp[0];
      //    }

      return 1;
    }