Analysis (apache#166)

Co-authored-by: Siva <quic_sivb@quicinc.com>
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>

Author: 2 people

include/tvm/relax/exec_builder.h

@@ -122,6 +122,15 @@ class ExecBuilderNode : public Object {
     rv = value;
     return ConvertConstant_(rv);
   }
+  /*!
+   * \brief update memory scopes.
+   *
+   * This function builds the memory scopes for constants.
+   *
+   * \param idx Index of the constant
+   * \param scope The memory scope.
+   */
+  void SaveMemoryScope(vm::Instruction::Arg idx, ffi::String scope);
   /*!
    * \brief Raw access to underlying executable build in progress.
    */

include/tvm/runtime/tensor.h

@@ -36,7 +36,6 @@
 
 #include <atomic>
 #include <functional>
-#include <string>
 #include <utility>
 #include <vector>
 
@@ -189,25 +188,14 @@ class Tensor : public tvm::ffi::Tensor {
    */
   TVM_DLL static void CopyFromBytes(const DLTensor* to, void* from, size_t nbytes,
                                     TVMStreamHandle stream = nullptr);
-
-  TVM_DLL void SetScope(ffi::String scope);
-  TVM_DLL ffi::String GetScope() const;
-
- protected:
-  /*!
-   * \brief The memory scope
-   * represents the underlaying scope information of device
-   */
-  ffi::String scope = "global";
 };
 
 /*!
  * \brief Save a DLTensor to stream
  * \param strm The output stream
  * \param tensor The tensor to be saved.
- * \param scope The tensor storage scope.
  */
-inline bool SaveDLTensor(dmlc::Stream* strm, const DLTensor* tensor, ffi::String scope = "global");
+inline bool SaveDLTensor(dmlc::Stream* strm, const DLTensor* tensor);
 
 inline void Tensor::CopyFrom(const DLTensor* other) {
   ICHECK(data_ != nullptr);
@@ -232,11 +220,10 @@ inline void Tensor::CopyTo(const Tensor& other) const {
 }
 
 /*! \brief Magic number for Tensor file */
-constexpr uint64_t kTVMNDArrayMagic = 0xDD5E40F096B4A13F;
-constexpr uint64_t kTVMNDArrayScopedMagic = 0xDD5E40F096B4A13E;
+constexpr uint64_t kTVMTensorMagic = 0xDD5E40F096B4A13F;
 
-inline bool SaveDLTensor(dmlc::Stream* strm, const DLTensor* tensor, ffi::String scope) {
-  uint64_t header = kTVMNDArrayScopedMagic, reserved = 0;
+inline bool SaveDLTensor(dmlc::Stream* strm, const DLTensor* tensor) {
+  uint64_t header = kTVMTensorMagic, reserved = 0;
   strm->Write(header);
   strm->Write(reserved);
   // Always save data as CPU context
@@ -256,7 +243,6 @@ inline bool SaveDLTensor(dmlc::Stream* strm, const DLTensor* tensor, ffi::String
   strm->Write(tensor->dtype);
   int ndim = tensor->ndim;
   strm->WriteArray(tensor->shape, ndim);
-  strm->Write(std::string(scope));
   int type_bytes = (tensor->dtype.bits + 7) / 8;
   int64_t num_elems = 1;
   for (int i = 0; i < ndim; ++i) {
@@ -280,14 +266,13 @@ inline bool SaveDLTensor(dmlc::Stream* strm, const DLTensor* tensor, ffi::String
   return true;
 }
 
-inline void Tensor::Save(dmlc::Stream* strm) const { SaveDLTensor(strm, operator->(), GetScope()); }
+inline void Tensor::Save(dmlc::Stream* strm) const { SaveDLTensor(strm, operator->()); }
 
 inline bool Tensor::Load(dmlc::Stream* strm) {
   uint64_t header, reserved;
   ICHECK(strm->Read(&header)) << "Invalid DLTensor file format";
   ICHECK(strm->Read(&reserved)) << "Invalid DLTensor file format";
-  ICHECK((header == kTVMNDArrayMagic) || (header == kTVMNDArrayScopedMagic))
-      << "Invalid DLTensor file format";
+  ICHECK(header == kTVMTensorMagic) << "Invalid DLTensor file format";
   Device dev;
   int ndim;
   DLDataType dtype;
@@ -305,11 +290,6 @@ inline bool Tensor::Load(dmlc::Stream* strm) {
   for (int i = 0; i < ret->ndim; ++i) {
     num_elems *= ret->shape[i];
   }
-  if (header == kTVMNDArrayScopedMagic) {
-    std::string scope;
-    strm->Read(&scope);
-    ret.SetScope(scope);
-  }
   int64_t data_byte_size;
   ICHECK(strm->Read(&data_byte_size)) << "Invalid DLTensor file format";
   ICHECK(data_byte_size == num_elems * elem_bytes) << "Invalid DLTensor file format";

include/tvm/runtime/vm/executable.h

@@ -155,6 +155,8 @@ class VMExecutable : public ffi::ModuleObj {
   std::unordered_map<std::string, Index> func_map;
   /*! \brief The global constant pool. */
   std::vector<ffi::Any> constants;
+  /*! \brief The VDevice memory scopes */
+  std::unordered_map<Index, std::string> memory_scopes;
   /*! \brief The offset of instruction. */
   std::vector<Index> instr_offset;
   /*! \brief The byte data of instruction. */
@@ -177,6 +179,11 @@ class VMExecutable : public ffi::ModuleObj {
    * \param strm The input stream.
    */
   void SaveGlobalSection(dmlc::Stream* strm) const;
+  /*!
+   * \brief Save the memory scopes.
+   * \param strm The output stream.
+   */
+  void SaveMemoryScopeSection(dmlc::Stream* strm) const;
   /*!
    * \brief Save the constant pool.
    * \param strm The input stream.
@@ -197,6 +204,11 @@ class VMExecutable : public ffi::ModuleObj {
    * \param strm The input stream.
    */
   void LoadGlobalSection(dmlc::Stream* strm);
+  /*!
+   * \brief Load the memory scopes.
+   * \param strm The input stream.
+   */
+  void LoadMemoryScopeSection(dmlc::Stream* strm);
   /*!
    * \brief Load the constant pool.
    * \param strm The input stream.

include/tvm/tir/transform.h

@@ -774,7 +774,7 @@ TVM_DLL Pass DefaultGPUSchedule();
 TVM_DLL Pass UseAssumeToReduceBranches();
 
 /*!
- * \brief Inject Texture Allocation intrensic.
+ * \brief Inject Texture Allocation intrinsic.
  * \return The pass.
  */
 TVM_DLL Pass InjectTextureAlloc();

python/tvm/dlight/adreno/__init__.py

@@ -21,5 +21,3 @@
 from .layout_transform import LayoutTransform
 from .fallback import Fallback
 from .pool import Pool2D
-
-# from .fallback import Fallback

python/tvm/dlight/adreno/convolution.py

@@ -21,8 +21,8 @@
 from tvm import tir
 from tvm.target import Target
 
-from .utils import schedule_inline_blocks, schedule_storage_annotate, schedule_default
-from ..base import analysis
+from .utils import schedule_inline_blocks, schedule_default
+from .. import analysis
 from .base import AdrenoScheduleRule
 
 
@@ -68,7 +68,7 @@ def schedule_conv2d(sch: tir.Schedule, blk: tir.schedule.BlockRV):
 
     def apply(  # pylint: disable=too-many-locals,missing-docstring
         self,
-        func: Union[tir.PrimFunc],
+        func: Union[tir.PrimFunc, tir.Schedule],
         target: Target,
         _: bool,
     ) -> Optional[tir.Schedule]:
@@ -102,6 +102,5 @@ def is_convolution(blk):
         Conv2d.schedule_conv2d(sch, conv_blk)
         remaining_blocks = schedule_inline_blocks(sch, remaining_blocks)
         schedule_default(sch, remaining_blocks)
-        schedule_storage_annotate(sch, remaining_blocks)
 
         return sch

python/tvm/dlight/adreno/fallback.py

@@ -21,9 +21,24 @@
 
 from tvm import tir
 from tvm.target import Target
-from ..base import analysis
+from .. import analysis
 from .base import AdrenoScheduleRule
-from .utils import get_texture_storage
+
+
+def _assert_gpu_target(target: Target):
+    if "gpu" not in target.keys:
+        raise ValueError(f"Expect a GPU target, but got {target}")
+
+
+def get_max_threads_per_block(target: Target) -> int:
+    _assert_gpu_target(target)
+    max_threads_per_block = None
+    for name in ["max_threads_per_block", "max_num_threads"]:
+        if max_threads_per_block is None:
+            max_threads_per_block = target.attrs.get(name, None)
+    if max_threads_per_block is None:
+        max_threads_per_block = 64
+    return int(max_threads_per_block)
 
 
 # pylint: disable=invalid-name,missing-function-docstring,unused-variable,unused-import
@@ -45,13 +60,13 @@ def schedule_inline_blocks(
         remaining_blocks = []
         for blk in blocks:
             block_info = analysis.get_block_info(sch, blk)
-            if block_info.is_injective() and not block_info.is_data_pad():
-                if len(block_info.consumers) == 1:
+            if block_info.is_injective() and not block_info.is_data_pad(sch):
+                if len(sch.get_consumers(blk)) == 1:
                     try:
                         sch.compute_inline(blk)
                     except Exception:  # pylint: disable=broad-exception-caught
                         remaining_blocks.append(blk)
-                elif len(block_info.producers) == 1:
+                elif len(sch.get_producers(blk)) == 1:
                     inlined_once = False
                     try:
                         # Would cause an issue inlining to producer with multiple consumers
@@ -71,34 +86,20 @@ def schedule_inline_blocks(
                 remaining_blocks.append(blk)
         return remaining_blocks
 
-    @staticmethod
-    def schedule_annotate_storage(sch: tir.Schedule, func=get_texture_storage):
-        """Annotates intermediate buffers to textures whenever it's possible to do so"""
-        return
-        # pylint: disable=unreachable
-        root_blk = analysis.get_root_block(sch)
-        blocks = sch.get_child_blocks(root_blk)
-
-        for blk in blocks:
-            block_info = analysis.get_block_info(sch, blk)
-            scope = func(block_info)
-            if scope is not None and len(sch.get_consumers(blk)) > 0:
-                sch.set_scope(blk, 0, scope)
-
     @staticmethod
     def schedule_default(sch: tir.Schedule, blk: tir.schedule.BlockRV):
         block_info = analysis.get_block_info(sch, blk)
 
         s_loops, r_loops, o_loops = [], [], []
-        v_loop = block_info.write_bufs[0].assoc_lps[-1]
+        v_loop = block_info.write_bufs(sch)[0].assoc_lps[-1]
 
         for iter_info in block_info.iters:
             if sch.get(iter_info.loop_rv) == sch.get(v_loop):
                 continue
             {"S": s_loops, "R": r_loops, "O": o_loops}.get(iter_info.kind).append(iter_info.loop_rv)
 
         iter_vars = analysis.collect_block_iter_vars_used_in_access_region(
-            block_info.block_stmt, block_info.write_bufs[0].buf_region.region
+            sch.get(blk), block_info.write_bufs(sch)[0].buf_region.region
         )
         o_outer = [lp for lp in o_loops if sch.get(lp).var in iter_vars]
         o_inner = [lp for lp in o_loops if sch.get(lp).var not in iter_vars]
@@ -114,7 +115,7 @@ def schedule_default(sch: tir.Schedule, blk: tir.schedule.BlockRV):
         tgt = Target.current(allow_none=True)
 
         b = sch.fuse(*s_loops)
-        tx_extent = analysis.get_max_threads_per_block(tgt) if tgt is not None else 256
+        tx_extent = get_max_threads_per_block(tgt) if tgt is not None else 256
         bx, tx = sch.split(b, [None, tx_extent])
         sch.bind(bx, "blockIdx.x")
         sch.bind(tx, "threadIdx.x")
@@ -139,7 +140,7 @@ def schedule_fallback(sch):
             blk
             for blk in blocks
             if analysis.get_block_info(sch, blk).is_reduction()
-            or analysis.get_block_info(sch, blk).is_data_pad()
+            or analysis.get_block_info(sch, blk).is_data_pad(sch)
         ]
         remaining_blocks = [blk for blk in blocks if blk not in schedule_blocks]
 
@@ -149,7 +150,6 @@ def schedule_fallback(sch):
         # TODO: Analyze unscheduled blocks to schedule instead of relying on remaining
         for blk in remaining_blocks:
             Fallback.schedule_default(sch, blk)
-        Fallback.schedule_annotate_storage(sch, schedule_blocks + remaining_blocks)
 
     def apply(  # pylint: disable=too-many-locals
         self,
@@ -170,7 +170,7 @@ def apply(  # pylint: disable=too-many-locals
             return None
 
         block_infos = [analysis.get_block_info(sch, block) for block in blocks]
-        if not any("texture" in block.write_bufs[0].get_scope() for block in block_infos):
+        if not any("texture" in block.write_bufs(sch)[0].get_scope() for block in block_infos):
             return None
 
         Fallback.schedule_fallback(sch)

python/tvm/dlight/adreno/layout_transform.py

@@ -22,7 +22,7 @@
 
 from tvm import tir
 from tvm.target import Target
-from ..base import analysis
+from .. import analysis
 
 from .base import AdrenoScheduleRule
 
@@ -61,11 +61,11 @@ def apply(  # pylint: disable=too-many-locals
         block_info = analysis.get_block_info(sch, blk)
         if not (
             (self.use_op_name and block_info.name == "te_layout_transform")
-            or (not self.use_op_name and block_info.is_layout_transform())
+            or (not self.use_op_name and block_info.is_layout_transform(sch))
         ):
             return None
 
-        read_buf, write_buf = (block_info.read_bufs[0], block_info.write_bufs[0])
+        read_buf, write_buf = (block_info.read_bufs(sch)[0], block_info.write_bufs(sch)[0])
         lps = block_info.get_loops()
         lpv_read, lpv_write = (
             read_buf.assoc_lps[-1],
@@ -124,9 +124,6 @@ def apply(  # pylint: disable=too-many-locals
         b = sch.fuse(*block_loops)
         tx_extent = min(sch.get(b).extent, 256)
         candidates = [1, 2, 4, 8, 16, 32]
-        ux = sch.sample_categorical(
-            candidates, [1 / len(candidates) for _ in range(len(candidates))]
-        )
         bx, tx = sch.split(b, [None, 256], preserve_unit_iters=True)
         sch.bind(bx, "blockIdx.x")
         sch.bind(tx, "threadIdx.x")

python/tvm/dlight/adreno/pool.py

@@ -21,7 +21,7 @@
 from tvm.target import Target
 
 from .base import AdrenoScheduleRule
-from ..base import analysis
+from .. import analysis
 
 
 # pylint: disable=invalid-name, unused-variable
@@ -87,6 +87,6 @@ def schedule_max_pool(blk: tir.schedule.BlockRV):
                         sch.reverse_compute_inline(blk)
                     else:
                         sch.compute_inline(blk)
-                except:  # pylint: disable=bare-except
+                except Exception:  # pylint: disable=broad-exception-caught
                     pass
         return sch

python/tvm/dlight/adreno/utils.py

@@ -22,7 +22,7 @@
 from tvm.target import Target
 from tvm import tir
 
-from ..base.analysis import BlockInfo
+from ..analysis import BlockInfo
 
 
 def get_texture_storage(block_info: BlockInfo):
@@ -83,13 +83,6 @@ def schedule_default(sch, blocks: List[tir.schedule.BlockRV] = None):
     return ret
 
 
-def schedule_storage_annotate(sch: tir.Schedule, func=get_texture_storage):
-    # Check the Write Buffer isn't one of input Params and is Texturizable...
-    from .fallback import Fallback
-
-    return Fallback.schedule_annotate_storage(sch)
-
-
 def schedule_fallback(sch, blk):
     from .fallback import Fallback