code: def test_buffer_index_merge_mult_mod():
      m = te.size_var("m")
      n = te.size_var("n")
      s = te.size_var("s")
      k0 = te.size_var("k0")
      k1 = te.size_var("k1")
      A = tvm.tir.decl_buffer((m, n), "float32")
      A_stride = tvm.tir.decl_buffer((m, n), "float32", strides=(s, 1))
  
      def assert_simplified_equal(index_simplified, index_direct):
          assert tvm.ir.structural_equal(
              index_simplified, index_direct
          ), "index_simplified=%s, index_direct=%s" % (
              index_simplified,
              index_direct,
          )
  
      idxd = tvm.tir.indexdiv
      idxm = tvm.tir.indexmod
  
      # Test Case1
      index_simplified = A_stride.offset_of(
          (
              idxd(idxm(k0, k1), s),
              idxm(idxm(k0, k1), s) + idxd(k0, k1) * k1,
          )
      )
      index_direct = A_stride.offset_of((0, k0))
      assert_simplified_equal(index_simplified, index_direct)
  
      # Test Case2
      index_simplified = A.offset_of(
          (
              idxd(idxm(k0, idxd(k1, s)), n),
              idxm(idxm(k0, idxd(k1, s)), n) + idxm(k0, k1),
          )
      )
      index_direct = A.offset_of(
          (0, idxm(k0, idxd(k1, s)) + idxm(k0, k1))
      )
      assert_simplified_equal(index_simplified, index_direct)
      # Test Case3
      index_simplified = A.offset_of(
          (
              idxd((idxd(k0, idxd(k1, s)) * idxd(k1, s)), n)
              + idxd(idxm(k0, idxd(k1, s)), n),
              idxm((idxd(k0, idxd(k1, s)) * idxd(k1, s)), n)
              + idxm(idxm(k0, idxd(k1, s)), n),
          )
      )
      index_direct = A.offset_of((0, k0))
      assert_simplified_equal(index_simplified, index_direct)
      # Test Case4 (not able to simplify)
      index_simplified = A.offset_of(
          (
              idxd(idxm(k0, idxd(k1, s)), n),
              idxm(idxm(k0, idxd(k1, n)), n) + idxm(k0, k1),
          )
      )
      index_direct = A.offset_of(
          (
              0,
              idxd(idxm(k0, idxd(k1, s)), n) * n
              + (idxm(idxm(k0, idxd(k1, n)), n) + idxm(k0, k1)),
          )
      )
      assert_simplified_equal(index_simplified, index_direct)
  
      # Test Case5
      B = tvm.tir.decl_buffer((1, 14, 14, 1024))
      i = te.size_var("i")
      j = te.size_var("j")
      k = te.size_var("k")
  
      index_simplified1 = B.offset_of(
          (
              idxd(
                  idxd(idxd((i * 50176 + j * 28672 + k), 1024), 14), 14
              ),
              idxm(
                  idxd(idxd((i * 50176 + j * 28672 + k), 1024), 14), 14
              ),
              idxm(idxd((i * 50176 + j * 28672 + k), 1024), 14),
              idxm((i * 50176 + j * 28672 + k), 1024),
          )
      )
      index_simplified2 = B.offset_of(
          (
              idxd(idxd(i * 49 + j * 28 + idxd(k, 1024), 14), 14),
              idxm(idxd(i * 49 + j * 28 + idxd(k, 1024), 14), 14),
              idxm(i * 7 + idxd(k, 1024), 14),
              idxm(k, 1024),
          )
      )
      index_direct = B.offset_of((0, 0, 0, (i * 50176 + j * 28672 + k)))
      assert_simplified_equal(index_simplified1, index_direct)
      assert_simplified_equal(index_simplified2, index_direct)
  
file_path: tvm/tests/python/tir-base/test_tir_buffer.py
imports: import pytest
  import tvm
  import tvm.testing
  from tvm import te
  from tvm.tir import Buffer
  import numpy
label_desc: Use image node_id 7 for calling a global function with example usage: test_buffer_index_merge_mult_mod() without return types