Merge pull request PaddlePaddle#21 from Superjomn/fea/init-tensor

init tensor

Author: Superjomn

cinn/common/shared.h

@@ -50,7 +50,10 @@ struct Shared {
   inline bool operator<(const Shared& other) const { return p_ < other.p_; }
   inline bool operator==(const Shared& other) const { return p_ == other.p_; }
 
-  ~Shared() { DesRef(p_); }
+  ~Shared() {
+    DesRef(p_);
+    p_ = nullptr;
+  }
 
  private:
   //! Increase the share count.

cinn/ir/CMakeLists.txt

@@ -6,14 +6,14 @@ cc_library(ir SRCS
         function_definition.cc
         ir_operators.cc
         buffer.cc
-        tensor.cc
+        #tensor.cc
         function_base.cc
-        operation.cc
+        #operation.cc
         DEPS common boost
         )
 
 cc_test(test_ir SRCS ir_test.cc DEPS ir)
 cc_test(test_ir_printer SRCS ir_printer_test.cc DEPS ir)
 cc_test(test_ir_operators SRCS ir_operators_test.cc DEPS ir)
-cc_test(test_tensor SRCS tensor_test.cc DEPS ir)
+#cc_test(test_tensor SRCS tensor_test.cc DEPS ir)
 cc_test(test_function SRCS function_test.cc DEPS ir)

cinn/ir/ir.cc

@@ -1,6 +1,7 @@
 #include "cinn/ir/ir.h"
 #include "cinn/common/pod_value.h"
 #include "cinn/ir/ir_visitor.h"
+#include "cinn/lang/tensor.h"
 
 namespace cinn {
 
@@ -219,6 +220,22 @@ Expr Call::Make(Type type,
   node->set_type(type);
   return Expr(node);
 }
+
+void _Tensor_::Accept(IrVisitor *v) const { v->Visit(this); }
+
+lang::Tensor _Tensor_::Make(const std::vector<Expr> &shape,
+                            const std::vector<Var> &iterators,
+                            Type dtype,
+                            ir::Expr expr) {
+  CHECK_EQ(shape.size(), iterators.size()) << "dimension of the shape and the iterators should match";
+  auto n       = common::make_shared<_Tensor_>();
+  n->dtype     = dtype;
+  n->shape     = shape;
+  n->expr      = expr;
+  n->iterators = iterators;
+  return lang::Tensor(n);
+}
+
 }  // namespace ir
 
 namespace common {
@@ -247,5 +264,4 @@ Value ToValue<ir::Var>(ir::Var v) {
 }
 
 }  // namespace common
-
 }  // namespace cinn

cinn/ir/ir.h

@@ -14,6 +14,15 @@
 #include "cinn/ir/node.h"
 
 namespace cinn {
+
+namespace poly {
+class Element;
+}  // namespace poly
+
+namespace lang {
+class Tensor;
+}  // namespace lang
+
 namespace ir {
 
 using common::Object;
@@ -505,6 +514,27 @@ class _IterVar_ : public IrNode {
   static const IrNodeTy _node_type_ = IrNodeTy::_Range_;
 };
 
+class _Tensor_ : public IrNode {
+ public:
+  //! Shape of this tensor.
+  std::vector<Expr> shape;
+  //! Data type of this tensor.
+  Type dtype;
+  //! The expression that generate this tensor.
+  ir::Expr expr;
+  //! The iterators, we store the iterators to name the dimensions for better readability.
+  std::vector<Var> iterators;
+  //! Polyhedral element for analysis and schedule.
+  std::unique_ptr<poly::Element> poly_element;
+
+  static lang::Tensor Make(const std::vector<Expr>& shape,
+                           const std::vector<Var>& iterators,
+                           Type dtype,
+                           ir::Expr expr);
+
+  void Accept(ir::IrVisitor* v) const override;
+};
+
 static IterVar thread_axis(Range dom, const std::string& tag) {
   return _IterVar_::Make(dom, Var(tag), IterVarType::kThreadIndex, tag);
 }

cinn/ir/tensor.h

@@ -105,6 +105,7 @@ class Operation : public FunctionRef {
   std::string name;
 };
 
+/*
 class _Tensor_ : public IrNode {
  public:
   //! The shape of the tensor.
@@ -121,6 +122,7 @@ class _Tensor_ : public IrNode {
 
   static const IrNodeTy _node_type_ = IrNodeTy::_Tensor_;
 };
+ */
 
 class _Operation_ : public ir::FunctionBase {
  public:

cinn/ir/tensor_test.cc

@@ -39,4 +39,4 @@ TEST(Tensor, func) {
 }
 
 }  // namespace ir
-}  // namespace cinn
+}  // namespace cinn

cinn/lang/CMakeLists.txt

@@ -1 +1,8 @@
-cc_library(lang SRCS buffer.cc tensor.cc DEPS common ir)
+cc_library(lang SRCS
+        buffer.cc
+        tensor.cc
+        compute.cc
+        DEPS ir poly)
+
+cc_test(test_compute SRCS compute_test.cc DEPS lang)
+cc_test(test_tensor2 SRCS tensor_test.cc DEPS lang)

cinn/lang/compute.cc

@@ -0,0 +1,26 @@
+#include "cinn/lang/compute.h"
+#include "cinn/poly/dim.h"
+#include "cinn/utils/functional.h"
+
+namespace cinn {
+namespace lang {
+
+using ir::Expr;
+
+template <>
+Tensor Compute<compute_handle_1_t>(const std::vector<int>& dims, compute_handle_1_t handle) {
+  CHECK_EQ(dims.size(), 1);
+
+  poly::Dim dim("i", 0, dims[0] - 1);
+
+  Var i("i", Int(32));
+  auto expr = handle(i);
+  std::vector<Expr> shape;
+  for (int v : dims) shape.emplace_back(v);
+
+  Tensor tensor(shape, {i}, expr.type(), expr);
+  return std::move(tensor);
+}
+
+}  // namespace lang
+}  // namespace cinn

cinn/lang/compute.h

@@ -0,0 +1,27 @@
+#pragma once
+#include <functional>
+#include <utility>
+#include <vector>
+#include "cinn/ir/ir.h"
+#include "cinn/lang/tensor.h"
+
+namespace cinn {
+namespace lang {
+
+using ir::Var;
+using compute_handle_1_t = std::function<ir::Expr(Var i)>;
+using compute_handle_2_t = std::function<ir::Expr(Var i0, Var i1)>;
+using compute_handle_3_t = std::function<ir::Expr(Var i0, Var i1, Var i2)>;
+using compute_handle_4_t = std::function<ir::Expr(Var i0, Var i1, Var i2, Var i3)>;
+
+/**
+ * Compute a Tensor.
+ * @param dims Dimensions.
+ * @param iterators
+ * @param handle
+ */
+template <typename Fn>
+Tensor Compute(const std::vector<int>& dims, Fn handle);
+
+}  // namespace lang
+}  // namespace cinn

cinn/lang/compute_test.cc

@@ -0,0 +1,7 @@
+#include "cinn/lang/compute.h"
+#include <gtest/gtest.h>
+#include "cinn/lang/tensor.h"
+
+namespace cinn {
+namespace lang {}  // namespace lang
+}  // namespace cinn