mxrtc.h

Go to the documentation of this file.

/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 * KIND, either express or implied.  See the License for the
 * specific language governing permissions and limitations
 * under the License.
 */

#ifndef MXNET_MXRTC_H_
#define MXNET_MXRTC_H_
#include "./base.h"
#if ((MXNET_USE_CUDA) && (MXNET_USE_NVRTC))
#include <nvrtc.h>
#include <cuda.h>

#include <vector>
#include <string>
#include <memory>
#include <utility>
#include <unordered_map>
#include "./ndarray.h"

namespace mxnet {

class MXRtc {
 public:
  MXRtc(const std::string& name,
        std::vector<std::pair<std::string, NDArray> > const& input,
        std::vector<std::pair<std::string, NDArray> > const& output,
        const std::string& kernel);
  void push(std::vector<NDArray> const& input,
            std::vector<NDArray> const& output,
            unsigned int  grid_dim_X,
            unsigned int  grid_dim_Y,
            unsigned int  grid_dim_Z,
            unsigned int  block_dim_X,
            unsigned int  block_dim_Y,
            unsigned int  block_dim_Z);

 private:
  static const char str_type[];
  static std::unordered_map<std::string, char*> kernel_registry;

  std::string name_;
  index_t num_input_, num_output_;
  std::string code_;
  char* ptx_;
  std::unordered_map<int, CUmodule> module_;
  std::unordered_map<int, CUfunction> func_;

  std::string decorate(const std::string& name,
                       std::vector<std::pair<std::string, NDArray> > const& input,
                       std::vector<std::pair<std::string, NDArray> > const& output,
                       const std::string kernel);
  char* compile(const std::string& name, const std::string& code);
};

}  // namespace mxnet

#endif  // MXNET_USE_CUDA && MXNET_USE_NVRTC
#endif  // MXNET_MXRTC_H_