expr.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.
 */

// Acknowledgement: This file originates from incubator-tvm
#ifndef MXNET_IR_EXPR_H_
#define MXNET_IR_EXPR_H_

#include <mxnet/runtime/object.h>
#include <mxnet/node/node.h>
#include <mxnet/node/container.h>
#include <mxnet/runtime/data_type.h>
#include <string>

namespace mxnet {

class BaseExprNode : public Object {
 public:
  static constexpr const char* _type_key = "Expr";
  MXNET_DECLARE_BASE_OBJECT_INFO(BaseExprNode, Object);
};

class BaseExpr : public ObjectRef {
 public:
  BaseExpr() {}
  explicit BaseExpr(runtime::ObjectPtr<Object> ptr) : ObjectRef(ptr) {}
  using ContainerType = BaseExprNode;
};

class PrimExprNode : public BaseExprNode {
 public:
  MXNetDataType dtype;

  static constexpr const char* _type_key = "PrimExpr";
  MXNET_DECLARE_BASE_OBJECT_INFO(PrimExprNode, BaseExprNode);
};

class PrimExpr : public BaseExpr {
 public:
  PrimExpr() {}
  explicit PrimExpr(runtime::ObjectPtr<Object> ptr) : BaseExpr(ptr) {}
  MXNET_DLL PrimExpr(int32_t value);  // NOLINT(*)
  MXNET_DLL PrimExpr(float value);  // NOLINT(*)
  MXNET_DLL PrimExpr(std::string str);  // NOLINT(*)

  MXNetDataType dtype() const {
    return static_cast<const PrimExprNode*>(get())->dtype;
  }
  using ContainerType = PrimExprNode;
};

class IntImmNode : public PrimExprNode {
 public:
  int64_t value;

  static constexpr const char* _type_key = "IntImm";
  MXNET_DECLARE_FINAL_OBJECT_INFO(IntImmNode, PrimExprNode);
};

class IntImm : public PrimExpr {
 public:
  IntImm() {}
  explicit IntImm(runtime::ObjectPtr<Object> node) : PrimExpr(node) {}
  MXNET_DLL IntImm(MXNetDataType dtype, int64_t value);
  const IntImmNode* operator->() const {
    return static_cast<const IntImmNode*>(get());
  }
  using ContainerType = IntImmNode;
};

class FloatImmNode : public PrimExprNode {
 public:
  double value;

  static constexpr const char* _type_key = "FloatImm";
  MXNET_DECLARE_FINAL_OBJECT_INFO(FloatImmNode, PrimExprNode);
};

class FloatImm : public PrimExpr {
 public:
  FloatImm() {}
  explicit FloatImm(runtime::ObjectPtr<Object> node) : PrimExpr(node) {}
  MXNET_DLL FloatImm(MXNetDataType dtype, double value);
  const FloatImmNode* operator->() const {
    return static_cast<const FloatImmNode*>(get());
  }
  using ContainerType = FloatImmNode;
};

}  // namespace mxnet
#endif  // MXNET_IR_EXPR_H_