[GitHub] [incubator-tvm] ANSHUMAN87 commented on a change in pull request #6580: Faster sparse_dense on GPUs

[GitBox <gi...@apache.org>]

ANSHUMAN87 commented on a change in pull request #6580:
URL: https://github.com/apache/incubator-tvm/pull/6580#discussion_r499212868



##########
File path: python/tvm/tir/ir_builder.py
##########
@@ -74,8 +75,23 @@ def asobject(self):
     def dtype(self):
         return self._content_type
 
+    def _linear_index(self, index):
+        if not isinstance(index, tuple):
+            return index
+        assert len(index) == len(self._shape), "Index size (%s) does not match shape size (%s)" % (
+            len(index),
+            len(self._shape),
+        )
+        dim_size = 1
+        lidx = 0
+        for dim, idx in zip(reversed(self._shape), reversed(index)):
+            lidx += idx * dim_size
+            dim_size *= dim
+        return lidx
+
     def __getitem__(self, index):
         t = DataType(self._content_type)
+        index = self._linear_index(index)

Review comment:
       @tkonolige : I am just trying to understand why this change is required for your requirement? 
   Would you please help me understand. Thanks!




----------------------------------------------------------------
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