[TOPI] [Hexagon] Batch flatten slice op initial version

[abhikran-quic]

This patch adds the initial python implementation batch flatten slice op for hexagon.

Slice ops are basically ops that make certain assumptions about the input and output dimensions and are expected to be called after the original op has been sliced according to those dimensions at the graph level.

cc @Lunderberg @cconvey @mehrdadh

[abhikran-quic]

HI @Lunderberg , @cconvey : Could you please review this PR ?

[Lunderberg]

Overall looks good, just some general comments and nitpicks.

[abhikran-quic]

Thank you @Lunderberg for your review. I have fixed the comments.

Right now CI is failing because of v69 support not being available in LLVM. Hopefully #11539 should address this and then the tests will pass in CI.

tests/python/contrib/test_hexagon/test_batch_flatten.py::TestBatchFlatten::test_batch_flatten[float16-input_shape3-n11c-1024c-1d-input_axis_sep0-nc-1d-output_axis_sep0] 'hexagonv69' is not a recognized processor for this target (ignoring processor)

[abhikran-quic]

CI is passing now.

[abhikran-quic]

Hi @Lunderberg @cconvey : Could you please review this PR for any more comments ?

[cconvey]

@abhikran-quic : Apologies for being slow to reply. I'm hoping to review this early next week if that's helpful.

[abhikran-quic]

Hi @Lunderberg @cconvey @mehrdadh,
Can you please review this PR ? It's ready to be merged from my side.

[cconvey]

Can you please review this PR ? It's ready to be merged from my side.

Reviewing now.

[Lunderberg]

Taking one more read-through after the latest commits, and found a couple more nitpicky changes.

[abhikran-quic]

Thank you @Lunderberg and @jverma-quic for your comments. I've fixed them in the latest commit.

[Lunderberg]

Thank you for making the changes, and LGTM!

[cconvey]

LGTM! One minor comment, but feel free to ignore it. We can revisit the issue later if necessary.

[cconvey]

I'm a bit suspicious about assuming that there's a block named "compute", as I don't see any promises in the documentation about the name and what it represents. But making assumptions like this seems somewhat idiomatic within TVM, so IMHO it's okay enough.

[jverma-quic]

I completely agree with you and don't like it either. But, there doesn't seem to be a way to specify a different block name for topi defined ops. Please correct me if I'm wrong.

[abhikran-quic]

Since I'm reusing batch_flatten compute in topi , the block name compute comes up. Sharing the schedule below

@main = primfn(var_A: handle, var_compute: handle) -> ()
  attr = {"global_symbol": "main", "tir.noalias": True}
  buffers = {A: Buffer(A_1: Pointer(global float16), float16, [1, 1, 1, 2, 1024], [], axis_separators=[4]),
             compute: Buffer(compute_1: Pointer(global float16), float16, [1, 2, 1024], [], axis_separators=[2])}
  buffer_map = {var_A: A, var_compute: compute} {
  block([], "root") {
    tir.reads([])
    tir.writes([])
    for (i0: int32, 0, 1) {
      for (i1_0_0: int32, 0, 1) {
        for (i1_0_1: int32, 0, 1) {
          for (i1_1_0: int32, 0, 2) {
            for (i1_1_1_0: int32, 0, 16) {
              for (i1_1_1_1: int32, 0, 64) "vectorized" {
                block([1, 2048], "compute") as [i, j] {
                  bind(i, 0)
                  bind(j, (((i1_1_0*1024) + (i1_1_1_0*64)) + i1_1_1_1))
                  tir.reads([A[0, 0, 0, (j / 1024), (j % 1024)]])
                  tir.writes([compute[0, (j / 1024), (j % 1024)]])
                  compute[0, (j / 1024), (j % 1024)] = A[0, 0, 0, (j / 1024), (j % 1024)]
                }
              }
            }
          }
        }
      }
    }
  }
  }

[cconvey]

👍

[abhikran-quic]

@Lunderberg , @mehrdadh , @cconvey : Could you please merge this PR ? I've fixed merge conflicts.

[abhikran-quic]

Gentle reminder! Please help in merging this PR. I want to raise another PR that's dependent on this one.

[mehrdadh]

@abhikran-quic please resolve the conflict by rebasing with main and push into this branch.

[abhikran-quic]

@abhikran-quic please resolve the conflict by rebasing with main and push into this branch.

@mehrdadh : I have resolved merge conflicts. Could you please review this ?