You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@tvm.apache.org by GitBox <gi...@apache.org> on 2020/07/22 06:32:01 UTC
[GitHub] [incubator-tvm] junrushao1994 commented on a change in pull request #6103: [Ansor][AutoTVM v2.0] Phase 1: Access Analyzer
junrushao1994 commented on a change in pull request #6103:
URL: https://github.com/apache/incubator-tvm/pull/6103#discussion_r458504576
##########
File path: include/tvm/auto_scheduler/loop_state.h
##########
@@ -159,10 +159,16 @@ using IterKey = std::pair<int, int>;
*/
class AttachMapNode : public Object {
public:
+ struct key_hash : public std::function<std::size_t(IterKey)> {
Review comment:
any specific reason that we inherit `std::function`?
##########
File path: include/tvm/auto_scheduler/auto_schedule.h
##########
@@ -18,29 +18,27 @@
*/
/*!
- * \file auto_scheduler/auto_schedule.h
Review comment:
Could you elaborate why we move all of these to public headers? For cpptests?
##########
File path: include/tvm/auto_scheduler/loop_state.h
##########
@@ -159,10 +159,16 @@ using IterKey = std::pair<int, int>;
*/
class AttachMapNode : public Object {
public:
+ struct key_hash : public std::function<std::size_t(IterKey)> {
Review comment:
Also, suggest a better name: `IterKeyHash`
##########
File path: include/tvm/auto_scheduler/compute_dag.h
##########
@@ -0,0 +1,237 @@
+/*
+ * 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.
+ */
+
+/*!
+ * \file tvm/auto_scheduler/compute_dag.h
+ * \brief The auto-scheduler's computational graph and related program analyses.
+ *
+ * We convert a compute declaration described by `tvm.compute` (could be a single operator or a
+ * subgraph) to a ComputeDAG. It keeps the input/output tensors of the compute declaration,
+ * a list of all operations in the DAG as well as static analysis results for the DAG (e.g. the
+ * total float operation count, consumer/producer relations of each operation stage, whether an
+ * operation stage should be tiled/compute inlined ...). These analyses can help the search policy
+ * to make decisions during search process.
+ * ComputeDAG is also responsible for the interaction between TVM Auto-scheduler `LoopState` and
+ * TVM schedule (e.g. applying the `LoopState` transform steps to TVM schedule, providing
+ * `LoopState` with extra information got from TVM schedule ...).
+ */
+
+#ifndef TVM_AUTO_SCHEDULER_COMPUTE_DAG_H_
+#define TVM_AUTO_SCHEDULER_COMPUTE_DAG_H_
+
+#include <tvm/auto_scheduler/loop_state.h>
+#include <tvm/te/schedule.h>
+
+#include <unordered_map>
+#include <unordered_set>
+#include <utility>
+#include <vector>
+
+namespace tvm {
+namespace auto_scheduler {
+
+/*! \brief Static analysis result for a ComputeDAG */
+class AccessAnalyzerNode : public Object {
+ public:
+ template <class T>
+ using OperationMap = std::unordered_map<te::Operation, T, ObjectHash, ObjectEqual>;
+
+ /*! \brief Map an operation to all operations it reads from.
+ * For each operation pair, use a two-dimentional array to multiple multi-dimentional accesses*/
+ OperationMap<OperationMap<std::vector<std::vector<PrimExpr>>>> read_from;
+ /*! \brief Map an operation to all operations it is read by.
+ * For each operation pair, use a two-dimentional array to multiple multi-dimentional accesses*/
+ OperationMap<OperationMap<std::vector<std::vector<PrimExpr>>>> read_by;
+ /*! \brief Store the number of common outer iterators for operation pairs that have
+ * read-write relations. */
+ OperationMap<OperationMap<int>> num_common_outer_iterators;
+ /*! \brief Store whether the operation is injective */
+ OperationMap<bool> is_injective;
+ /*! \brief Store whether the operation is strictly-inlineable */
+ OperationMap<bool> is_strict_inlineable;
+ /*! \brief Store whether the operation needs multi-level tiling */
+ OperationMap<bool> needs_multi_level_tiling;
+ /*! \brief Store whether the operation is an output operation */
+ OperationMap<bool> is_output;
+ /*! \brief Store the topological order of operations */
+ Array<te::Operation> ops_topo_order;
Review comment:
What's the relationship between this array and `ComputeDAG::ops`?
##########
File path: include/tvm/auto_scheduler/compute_dag.h
##########
@@ -0,0 +1,242 @@
+/*r
+ * 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.
+ */
+
+/*!
+ * \file tvm/auto_scheduler/compute_dag.h
+ * \brief The auto-scheduler's computational graph and related program analyses.
+ *
+ * We convert a compute declaration described by `tvm.compute` (could be a single operator or a
+ * subgraph) to a ComputeDAG. It keeps the input/output tensors of the compute declaration,
+ * a list of all operations in the DAG as well as static analysis results for the DAG (e.g. the
+ * total float operation count, consumer/producer relations of each operation stage, whether an
+ * operation stage should be tiled/compute inlined ...). These analyses can help the search policy
+ * to make decisions during search process.
+ * ComputeDAG is also responsible for the interaction between TVM Auto-scheduler `LoopState` and
+ * TVM schedule (e.g. applying the `LoopState` transform steps to TVM schedule, providing
+ * `LoopState` with extra information got from TVM schedule ...).
+ */
+
+#ifndef TVM_AUTO_SCHEDULER_COMPUTE_DAG_H_
+#define TVM_AUTO_SCHEDULER_COMPUTE_DAG_H_
+
+#include <tvm/auto_scheduler/loop_state.h>
+#include <tvm/runtime/c_runtime_api.h>
+#include <tvm/te/schedule.h>
+
+#include <unordered_map>
+#include <unordered_set>
+#include <utility>
+#include <vector>
+
+namespace tvm {
+namespace auto_scheduler {
+
+/*! \brief Static analysis result for a ComputeDAG */
+class AccessAnalyzerNode : public Object {
+ public:
+ template <class T>
+ using OperationMap = std::unordered_map<te::Operation, T, ObjectHash, ObjectEqual>;
+
+ /*! \brief Map an operation to all operations it reads from.
+ * For each operation pair, use a two-dimentional array to multiple multi-dimentional accesses*/
+ OperationMap<OperationMap<std::vector<std::vector<PrimExpr>>>> read_from;
+ /*! \brief Map an operation to all operations it is read by.
+ * For each operation pair, use a two-dimentional array to multiple multi-dimentional accesses*/
+ OperationMap<OperationMap<std::vector<std::vector<PrimExpr>>>> read_by;
+ /*! \brief Store the number of common outer iterators for operation pairs that have
+ * read-write relations. */
+ OperationMap<OperationMap<int>> num_common_outer_iterators;
+ /*! \brief Store whether the operation is injective */
+ OperationMap<bool> is_injective;
+ /*! \brief Store whether the operation is strictly-inlineable */
+ OperationMap<bool> is_strict_inlineable;
+ /*! \brief Store whether the operation needs multi-level tiling */
+ OperationMap<bool> needs_multi_level_tiling;
+ /*! \brief Store whether the operation is an output operation */
+ OperationMap<bool> is_output;
+ /*! \brief Store the topological order of operations */
+ Array<te::Operation> ops_topo_order;
+
+ static constexpr const char* _type_key = "auto_scheduler.AccessAnalyzer";
+ TVM_DECLARE_FINAL_OBJECT_INFO(AccessAnalyzerNode, Object);
+};
+
+/*!
+ * \brief Managed reference to AccessAnalyzerNode.
+ * \sa AccessAnalyzerNode
+ */
+class AccessAnalyzer : public ObjectRef {
+ public:
+ explicit AccessAnalyzer(const Array<te::Tensor>& tensors);
+
+ /*!
+ * \brief Return whether this operation needs multi-level tiling
+ * \param op The operation
+ */
+ TVM_DLL bool NeedsMultiLevelTiling(const te::Operation& op) const;
+
+ /*!
+ * \brief Return whether this operation is an injective operation
+ * \param op The operation
+ */
+ TVM_DLL bool IsInjective(const te::Operation& op) const;
+
+ /*!
+ * \brief Return whether this operation is strictly inlinable
+ * \param op The operation
+ */
+ TVM_DLL bool IsStrictInlineable(const te::Operation& op) const;
+
+ /*!
+ * \brief Return whether this operation is an output op
+ * \param op The operation
+ */
+ TVM_DLL bool IsOutput(const te::Operation& op) const;
+
+ /*!
+ * \brief Get all consumers of on operation
+ * \param state The current loop state
+ * \param op The operation
+ * \param consumers The return consumer set
+ * \note This function propagates the relation for inlined ops
+ */
+ TVM_DLL void GetConsumers(
+ const State& state, const te::Operation& op,
+ std::unordered_set<te::Operation, ObjectHash, ObjectEqual>* consumers) const;
+
+ /*!
+ * \brief Get all producers of on operation
+ * \param state The current loop state
+ * \param op The operation
+ * \param producers The return producer set
+ * \note This function propagates the relation for inlined ops
+ */
+ TVM_DLL void GetProducers(
+ const State& state, const te::Operation& op,
+ std::unordered_set<te::Operation, ObjectHash, ObjectEqual>* producers) const;
Review comment:
ditto
##########
File path: include/tvm/auto_scheduler/compute_dag.h
##########
@@ -0,0 +1,237 @@
+/*
+ * 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.
+ */
+
+/*!
+ * \file tvm/auto_scheduler/compute_dag.h
+ * \brief The auto-scheduler's computational graph and related program analyses.
+ *
+ * We convert a compute declaration described by `tvm.compute` (could be a single operator or a
+ * subgraph) to a ComputeDAG. It keeps the input/output tensors of the compute declaration,
+ * a list of all operations in the DAG as well as static analysis results for the DAG (e.g. the
+ * total float operation count, consumer/producer relations of each operation stage, whether an
+ * operation stage should be tiled/compute inlined ...). These analyses can help the search policy
+ * to make decisions during search process.
+ * ComputeDAG is also responsible for the interaction between TVM Auto-scheduler `LoopState` and
+ * TVM schedule (e.g. applying the `LoopState` transform steps to TVM schedule, providing
+ * `LoopState` with extra information got from TVM schedule ...).
+ */
+
+#ifndef TVM_AUTO_SCHEDULER_COMPUTE_DAG_H_
+#define TVM_AUTO_SCHEDULER_COMPUTE_DAG_H_
+
+#include <tvm/auto_scheduler/loop_state.h>
+#include <tvm/te/schedule.h>
+
+#include <unordered_map>
+#include <unordered_set>
+#include <utility>
+#include <vector>
+
+namespace tvm {
+namespace auto_scheduler {
+
+/*! \brief Static analysis result for a ComputeDAG */
+class AccessAnalyzerNode : public Object {
+ public:
+ template <class T>
+ using OperationMap = std::unordered_map<te::Operation, T, ObjectHash, ObjectEqual>;
+
+ /*! \brief Map an operation to all operations it reads from.
+ * For each operation pair, use a two-dimentional array to multiple multi-dimentional accesses*/
+ OperationMap<OperationMap<std::vector<std::vector<PrimExpr>>>> read_from;
Review comment:
It is somewhat annoying to have such nested data structure. What about we defining some other types for `OperationMap<OperationMap<T>>`?
##########
File path: include/tvm/auto_scheduler/compute_dag.h
##########
@@ -0,0 +1,237 @@
+/*
+ * 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.
+ */
+
+/*!
+ * \file tvm/auto_scheduler/compute_dag.h
+ * \brief The auto-scheduler's computational graph and related program analyses.
+ *
+ * We convert a compute declaration described by `tvm.compute` (could be a single operator or a
+ * subgraph) to a ComputeDAG. It keeps the input/output tensors of the compute declaration,
+ * a list of all operations in the DAG as well as static analysis results for the DAG (e.g. the
+ * total float operation count, consumer/producer relations of each operation stage, whether an
+ * operation stage should be tiled/compute inlined ...). These analyses can help the search policy
+ * to make decisions during search process.
+ * ComputeDAG is also responsible for the interaction between TVM Auto-scheduler `LoopState` and
+ * TVM schedule (e.g. applying the `LoopState` transform steps to TVM schedule, providing
+ * `LoopState` with extra information got from TVM schedule ...).
+ */
+
+#ifndef TVM_AUTO_SCHEDULER_COMPUTE_DAG_H_
+#define TVM_AUTO_SCHEDULER_COMPUTE_DAG_H_
+
+#include <tvm/auto_scheduler/loop_state.h>
+#include <tvm/te/schedule.h>
+
+#include <unordered_map>
+#include <unordered_set>
+#include <utility>
+#include <vector>
+
+namespace tvm {
+namespace auto_scheduler {
+
+/*! \brief Static analysis result for a ComputeDAG */
+class AccessAnalyzerNode : public Object {
+ public:
+ template <class T>
+ using OperationMap = std::unordered_map<te::Operation, T, ObjectHash, ObjectEqual>;
+
+ /*! \brief Map an operation to all operations it reads from.
+ * For each operation pair, use a two-dimentional array to multiple multi-dimentional accesses*/
Review comment:
Elaborate "multi-dimensional access", like, the inner `std::vector` represents the multi-dimensional indices. Also use a specific type alias for `std::vector<PrimExpr>`
##########
File path: include/tvm/auto_scheduler/compute_dag.h
##########
@@ -0,0 +1,242 @@
+/*r
+ * 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.
+ */
+
+/*!
+ * \file tvm/auto_scheduler/compute_dag.h
+ * \brief The auto-scheduler's computational graph and related program analyses.
+ *
+ * We convert a compute declaration described by `tvm.compute` (could be a single operator or a
+ * subgraph) to a ComputeDAG. It keeps the input/output tensors of the compute declaration,
+ * a list of all operations in the DAG as well as static analysis results for the DAG (e.g. the
+ * total float operation count, consumer/producer relations of each operation stage, whether an
+ * operation stage should be tiled/compute inlined ...). These analyses can help the search policy
+ * to make decisions during search process.
+ * ComputeDAG is also responsible for the interaction between TVM Auto-scheduler `LoopState` and
+ * TVM schedule (e.g. applying the `LoopState` transform steps to TVM schedule, providing
+ * `LoopState` with extra information got from TVM schedule ...).
+ */
+
+#ifndef TVM_AUTO_SCHEDULER_COMPUTE_DAG_H_
+#define TVM_AUTO_SCHEDULER_COMPUTE_DAG_H_
+
+#include <tvm/auto_scheduler/loop_state.h>
+#include <tvm/runtime/c_runtime_api.h>
+#include <tvm/te/schedule.h>
+
+#include <unordered_map>
+#include <unordered_set>
+#include <utility>
+#include <vector>
+
+namespace tvm {
+namespace auto_scheduler {
+
+/*! \brief Static analysis result for a ComputeDAG */
+class AccessAnalyzerNode : public Object {
+ public:
+ template <class T>
+ using OperationMap = std::unordered_map<te::Operation, T, ObjectHash, ObjectEqual>;
+
+ /*! \brief Map an operation to all operations it reads from.
+ * For each operation pair, use a two-dimentional array to multiple multi-dimentional accesses*/
+ OperationMap<OperationMap<std::vector<std::vector<PrimExpr>>>> read_from;
+ /*! \brief Map an operation to all operations it is read by.
+ * For each operation pair, use a two-dimentional array to multiple multi-dimentional accesses*/
+ OperationMap<OperationMap<std::vector<std::vector<PrimExpr>>>> read_by;
+ /*! \brief Store the number of common outer iterators for operation pairs that have
+ * read-write relations. */
+ OperationMap<OperationMap<int>> num_common_outer_iterators;
+ /*! \brief Store whether the operation is injective */
+ OperationMap<bool> is_injective;
+ /*! \brief Store whether the operation is strictly-inlineable */
+ OperationMap<bool> is_strict_inlineable;
+ /*! \brief Store whether the operation needs multi-level tiling */
+ OperationMap<bool> needs_multi_level_tiling;
+ /*! \brief Store whether the operation is an output operation */
+ OperationMap<bool> is_output;
+ /*! \brief Store the topological order of operations */
+ Array<te::Operation> ops_topo_order;
+
+ static constexpr const char* _type_key = "auto_scheduler.AccessAnalyzer";
+ TVM_DECLARE_FINAL_OBJECT_INFO(AccessAnalyzerNode, Object);
+};
+
+/*!
+ * \brief Managed reference to AccessAnalyzerNode.
+ * \sa AccessAnalyzerNode
+ */
+class AccessAnalyzer : public ObjectRef {
+ public:
+ explicit AccessAnalyzer(const Array<te::Tensor>& tensors);
+
+ /*!
+ * \brief Return whether this operation needs multi-level tiling
+ * \param op The operation
+ */
+ TVM_DLL bool NeedsMultiLevelTiling(const te::Operation& op) const;
+
+ /*!
+ * \brief Return whether this operation is an injective operation
+ * \param op The operation
+ */
+ TVM_DLL bool IsInjective(const te::Operation& op) const;
+
+ /*!
+ * \brief Return whether this operation is strictly inlinable
+ * \param op The operation
+ */
+ TVM_DLL bool IsStrictInlineable(const te::Operation& op) const;
+
+ /*!
+ * \brief Return whether this operation is an output op
+ * \param op The operation
+ */
+ TVM_DLL bool IsOutput(const te::Operation& op) const;
+
+ /*!
+ * \brief Get all consumers of on operation
+ * \param state The current loop state
+ * \param op The operation
+ * \param consumers The return consumer set
+ * \note This function propagates the relation for inlined ops
+ */
+ TVM_DLL void GetConsumers(
+ const State& state, const te::Operation& op,
+ std::unordered_set<te::Operation, ObjectHash, ObjectEqual>* consumers) const;
Review comment:
Shall we just return consumers?
```suggestion
TVM_DLL std::unordered_set<te::Operation, ObjectHash, ObjectEqual>* GetConsumers(
const State& state, const te::Operation& op) const;
```
##########
File path: include/tvm/auto_scheduler/compute_dag.h
##########
@@ -0,0 +1,237 @@
+/*
+ * 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.
+ */
+
+/*!
+ * \file tvm/auto_scheduler/compute_dag.h
+ * \brief The auto-scheduler's computational graph and related program analyses.
+ *
+ * We convert a compute declaration described by `tvm.compute` (could be a single operator or a
+ * subgraph) to a ComputeDAG. It keeps the input/output tensors of the compute declaration,
+ * a list of all operations in the DAG as well as static analysis results for the DAG (e.g. the
+ * total float operation count, consumer/producer relations of each operation stage, whether an
+ * operation stage should be tiled/compute inlined ...). These analyses can help the search policy
+ * to make decisions during search process.
+ * ComputeDAG is also responsible for the interaction between TVM Auto-scheduler `LoopState` and
+ * TVM schedule (e.g. applying the `LoopState` transform steps to TVM schedule, providing
+ * `LoopState` with extra information got from TVM schedule ...).
+ */
+
+#ifndef TVM_AUTO_SCHEDULER_COMPUTE_DAG_H_
+#define TVM_AUTO_SCHEDULER_COMPUTE_DAG_H_
+
+#include <tvm/auto_scheduler/loop_state.h>
+#include <tvm/te/schedule.h>
+
+#include <unordered_map>
+#include <unordered_set>
+#include <utility>
+#include <vector>
+
+namespace tvm {
+namespace auto_scheduler {
+
+/*! \brief Static analysis result for a ComputeDAG */
+class AccessAnalyzerNode : public Object {
+ public:
+ template <class T>
+ using OperationMap = std::unordered_map<te::Operation, T, ObjectHash, ObjectEqual>;
Review comment:
seems like pointer hashing and equality
```suggestion
using OperationMap = std::unordered_map<te::Operation, T, ObjectPtrHash, ObjectPtrEqual>;
```
##########
File path: include/tvm/auto_scheduler/compute_dag.h
##########
@@ -0,0 +1,242 @@
+/*r
+ * 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.
+ */
+
+/*!
+ * \file tvm/auto_scheduler/compute_dag.h
+ * \brief The auto-scheduler's computational graph and related program analyses.
+ *
+ * We convert a compute declaration described by `tvm.compute` (could be a single operator or a
+ * subgraph) to a ComputeDAG. It keeps the input/output tensors of the compute declaration,
+ * a list of all operations in the DAG as well as static analysis results for the DAG (e.g. the
+ * total float operation count, consumer/producer relations of each operation stage, whether an
+ * operation stage should be tiled/compute inlined ...). These analyses can help the search policy
+ * to make decisions during search process.
+ * ComputeDAG is also responsible for the interaction between TVM Auto-scheduler `LoopState` and
+ * TVM schedule (e.g. applying the `LoopState` transform steps to TVM schedule, providing
+ * `LoopState` with extra information got from TVM schedule ...).
+ */
+
+#ifndef TVM_AUTO_SCHEDULER_COMPUTE_DAG_H_
+#define TVM_AUTO_SCHEDULER_COMPUTE_DAG_H_
+
+#include <tvm/auto_scheduler/loop_state.h>
+#include <tvm/runtime/c_runtime_api.h>
+#include <tvm/te/schedule.h>
+
+#include <unordered_map>
+#include <unordered_set>
+#include <utility>
+#include <vector>
+
+namespace tvm {
+namespace auto_scheduler {
+
+/*! \brief Static analysis result for a ComputeDAG */
+class AccessAnalyzerNode : public Object {
+ public:
+ template <class T>
+ using OperationMap = std::unordered_map<te::Operation, T, ObjectHash, ObjectEqual>;
+
+ /*! \brief Map an operation to all operations it reads from.
+ * For each operation pair, use a two-dimentional array to multiple multi-dimentional accesses*/
+ OperationMap<OperationMap<std::vector<std::vector<PrimExpr>>>> read_from;
+ /*! \brief Map an operation to all operations it is read by.
+ * For each operation pair, use a two-dimentional array to multiple multi-dimentional accesses*/
+ OperationMap<OperationMap<std::vector<std::vector<PrimExpr>>>> read_by;
+ /*! \brief Store the number of common outer iterators for operation pairs that have
+ * read-write relations. */
+ OperationMap<OperationMap<int>> num_common_outer_iterators;
+ /*! \brief Store whether the operation is injective */
+ OperationMap<bool> is_injective;
+ /*! \brief Store whether the operation is strictly-inlineable */
+ OperationMap<bool> is_strict_inlineable;
+ /*! \brief Store whether the operation needs multi-level tiling */
+ OperationMap<bool> needs_multi_level_tiling;
+ /*! \brief Store whether the operation is an output operation */
+ OperationMap<bool> is_output;
+ /*! \brief Store the topological order of operations */
+ Array<te::Operation> ops_topo_order;
+
+ static constexpr const char* _type_key = "auto_scheduler.AccessAnalyzer";
+ TVM_DECLARE_FINAL_OBJECT_INFO(AccessAnalyzerNode, Object);
+};
+
+/*!
+ * \brief Managed reference to AccessAnalyzerNode.
+ * \sa AccessAnalyzerNode
+ */
+class AccessAnalyzer : public ObjectRef {
+ public:
+ explicit AccessAnalyzer(const Array<te::Tensor>& tensors);
+
+ /*!
+ * \brief Return whether this operation needs multi-level tiling
+ * \param op The operation
+ */
+ TVM_DLL bool NeedsMultiLevelTiling(const te::Operation& op) const;
+
+ /*!
+ * \brief Return whether this operation is an injective operation
+ * \param op The operation
+ */
+ TVM_DLL bool IsInjective(const te::Operation& op) const;
+
+ /*!
+ * \brief Return whether this operation is strictly inlinable
+ * \param op The operation
+ */
+ TVM_DLL bool IsStrictInlineable(const te::Operation& op) const;
+
+ /*!
+ * \brief Return whether this operation is an output op
+ * \param op The operation
+ */
+ TVM_DLL bool IsOutput(const te::Operation& op) const;
+
+ /*!
+ * \brief Get all consumers of on operation
+ * \param state The current loop state
+ * \param op The operation
+ * \param consumers The return consumer set
+ * \note This function propagates the relation for inlined ops
+ */
+ TVM_DLL void GetConsumers(
+ const State& state, const te::Operation& op,
+ std::unordered_set<te::Operation, ObjectHash, ObjectEqual>* consumers) const;
+
+ /*!
+ * \brief Get all producers of on operation
+ * \param state The current loop state
+ * \param op The operation
+ * \param producers The return producer set
+ * \note This function propagates the relation for inlined ops
+ */
+ TVM_DLL void GetProducers(
+ const State& state, const te::Operation& op,
+ std::unordered_set<te::Operation, ObjectHash, ObjectEqual>* producers) const;
+
+ /*!
+ * \brief Get all direct producers of on operation
+ * \param op The operation
+ * \param producers The return producer set
+ * \note This function DOES NOT propagate the relation for inlined ops
+ */
+ TVM_DLL void GetDirectProducers(
+ const te::Operation& op,
+ std::unordered_set<te::Operation, ObjectHash, ObjectEqual>* producers) const;
Review comment:
ditto
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
users@infra.apache.org