[Codegen][GPU] Allow iree_gpu.barrier_region to take multiple operands/results

[qedawkins]

The restriction to a single input and output was artificial as this op simply represents synchronization on input and output values. Additionally this removes the restriction on tensor/vector types, but for the time being this op is still only used with those types.

[MaheshRavishankar]

So what this mean to have multiple operations within the region?

I am actually not sure I fully understand why a region is needed and not just a value_barrier.

[qedawkins]

The main reason is to be able to better represent the barrier semantics of scf.forall after doing loop fusion. Currently we start with something like this

%0 = scf.forall {
  ...
}
scf.forall {
  tensor.extract_slice %0
  ...
}

And after fusion we form something like this

scf.forall {
  %0 = scf.for {
    // %0 body
  }
  iree_gpu.barrier_region %0 {
   ^bb0(%s)
    tensor.extract_slice %s
  }
  ...
}

(this actually has a bug because tensor.extract_slice is not guaranteed to bufferize to a read-effecting op, it typically just bufferizes to a memref.subview, meaning the actual read can happen later on after the barrier, but that is a bug with the loop fusion pattern).

This lowering works out ok because of some later vectorization patterns, but this requires the loop fusion pattern to find an extract_slice and for that extract_slice to actually be loading the thread local slice of data. This is too much spooky action at a distance and isn't an accurate representation of the IR before fusion. What we want to generate is this

scf.forall {
  %0 = iree_gpu.barrier_region ins(%0_dest) {
    %1 = scf.for {
      // %0 body
    }
    %2 = tensor.insert_slice %1
    iree_gpu.yield %2
  }
  tensor.extract_slice %0
  ...
}

Where the idea is that we have a barrier before and after the body of the fused forall op, instead of before and after the read of the shared memory result.

The reason this IR needs to be a region instead of an iree_gpu.value_barrier is that the value barrier would look something like this

%0 = scf.forall {
  ...
}
scf.forall {
  %alloc = bufferization.alloc_tensor()
  %0_dest = iree_gpu.value_barrier %alloc
  %1 = scf.for {
    // %0 body
  }
  %2 = tensor.insert_slice %1
  %3 = iree_gpu.value_barrier %2
  tensor.extract_slice %3
  ...
}

Where unless we made iree_gpu.value_barrier non-speculatable, would produce incorrect IR due to hoisting of the value_barrier out of parent loops.

[MaheshRavishankar]

Thanks @qedawkins for the explanation here and offline!

[qedawkins]

@MaheshRavishankar can you take a look when you get a chance? This is blocking a number of follow up PRs I have.

[MaheshRavishankar]

THis itself looks fine, but the vectorization of barrier ops is kind of annoying.

[qedawkins]

Yeah don't worry the vectorization pattern is going away (it's broken anyway). I just needed a couple more changes dependent on this change before I could drop it so I had to make it work for now.