ViewOp isn't expressive enough?

MLIR
1

Hi,

To set up the example, I have two allocations memref<3xi8> and memref<2x3xf32>. For reducing allocation cost, I allocated the two buffers together as memref<27xi8>, in the hope to slice them up later.

How do I exactly get the second slice? My observation is that, there is no way to express “skip the first 3 bytes and take the rest 24 bytes as a slice” with view/subviews.

2

Isn’t the affine map of the view providing an offset?

I’m surprised by your question here, because you reviewed and commented on this block of code just last week I think: https://github.com/tensorflow/tensorflow/blob/master/tensorflow/compiler/mlir/xla/transforms/xla_hlo_to_lhlo_with_xla.cc#L178-L186 ; isn’t this doing just what you describe?

3

@timshen, the following lowers to LLVM.

// RUN: mlir-opt %s -convert-std-to-llvm

func @views() {
  %0 = alloc(): memref<27 x i8>
  %1 = view %0[][]: memref<27 x i8> to memref<3xi8>
  %2 = view %0[][]: memref<27 x i8> to memref<2x3xf32, offset: 3, strides: [3, 1]>
  std.return
}

However you are right that there is a bug here: the offset 3 is going to be applied in multiples of f32 but what we want is that the view pointer is shifted by 3 bytes.

This calls for some fixes and should probably look like this instead:

func @views() {
  %0 = alloc(): memref<27 x i8>
  %1 = view %0[][]: memref<27 x i8> to memref<3xi8>
  %2 = view %0[%c3][%c2, %c3]: memref<27 x i8> to memref<2x3xf32>
  std.return
}

The lowering of a view op should apply the shift to the underlying buffer and can only produce contiguous memref without offset in the type (i.e. offset within the type should not refer to something outside the type ).

A consequence would be that the canonicalizations of the offsets would not be possible (cc @andydavis).

4

It intends to, but fails the ViewOp verifier. This code is what I’m currently investigating. Notice that ViewOp doesn’t take inputs with non-identity affine maps at [1].

[1] tensorflow/tensorflow/compiler/mlir/xla/transforms/xla_hlo_to_lhlo_with_xla.cc at db9c68b9441ab1732d086a21bda981582b73012c · tensorflow/tensorflow · GitHub

5

That’s because the view with the offset is applied a few lines above: tensorflow/xla_hlo_to_lhlo_with_xla.cc at db9c68b9441ab1732d086a21bda981582b73012c · tensorflow/tensorflow · GitHub ; the one you’re pointing at is just giving it a different shape in-place.

6

Yes, the problem is that L203 returns an op with a non-identity affine map. An op with an affine map is not suitable to be passed into L207.

7

There is a possibility to incorporate strides as well. I imagine the same conclusion for strides, which calls for all offsets, sizes, and strides to be completely decoupled from affine maps. I’ll prototype a fix like this since it’s blocking my work.

8

That’s… surprising to me! This does not seem documented: https://mlir.llvm.org/docs/Dialects/Standard/#stdview-viewop
Is this intentional or just an implementation limitation/issue to fix?

9

Please see: https://reviews.llvm.org/D79541

Note the tests are very much WIP but early comments on the update to Ops.td is what I am looking for.