[all-commits] [llvm/llvm-project] e7dcf1: [mlir][xegpu] Add SIMT distribution patterns for U...

[Charitha Saumya via All-commits]

  Branch: refs/heads/main
  Home:   https://github.com/llvm/llvm-project
  Commit: e7dcf1b7e5574d03d1ce6e7520d5683cfea37706
      https://github.com/llvm/llvm-project/commit/e7dcf1b7e5574d03d1ce6e7520d5683cfea37706
  Author: Charitha Saumya <136391709+charithaintc at users.noreply.github.com>
  Date:   2025-05-08 (Thu, 08 May 2025)

  Changed paths:
    M mlir/include/mlir/Dialect/XeGPU/IR/XeGPUOps.td
    M mlir/include/mlir/Dialect/XeGPU/IR/XeGPUTypes.td
    M mlir/lib/Dialect/XeGPU/Transforms/XeGPUSubgroupDistribute.cpp
    M mlir/test/Dialect/XeGPU/subgroup-distribution.mlir
    M mlir/test/Dialect/XeGPU/subgroup-map-propagation.mlir

  Log Message:
  -----------
  [mlir][xegpu] Add SIMT distribution patterns for UpdateNdOffset and PrefetchNd ops.  (#138033)

This PR adds support for SIMT distribution of UpdateNdOffset and
PrefetchNd ops.

For both these ops distribution will remove the layout attribute from
the tensor descriptor type. Everything else remains unchanged.

Example 1:

 ```
   #lo0 = #xegpu.layout<wi_layout = [1, 8], wi_data = [1, 1]>
   gpu.warp_execute_on_lane_0(%laneid) -> () {
     ...
     xegpu.prefetch_nd %arg0 : !xegpu.tensor_desc<4x8xf32, #lo0>
   }
 ```
 To
 ```
   %r:2 = gpu.warp_execute_on_lane_0(%laneid) -> (
   !xegpu.tensor_desc<4x8xf32, #lo0>) {
     gpu.yield %arg0: !xegpu.tensor_desc<4x8xf32, #lo0>
   }
   %1 = unrealized_conversion_cast %r#0: !xegpu.tensor_desc<4x8xf32,
     #lo0> -> !xegpu.tensor_desc<4x8xf32>
   xegpu.prefetch_nd %0 : !xegpu.tensor_desc<4x8xf32>

 ```
Example 2:
 ```
   #lo0 = #xegpu.layout<wi_layout = [1, 8], wi_data = [1, 1]>
   %r = gpu.warp_execute_on_lane_0(%laneid) ->
                   (!xegpu.tensor_desc<4x8xf32, #lo0>) {
     ...
     %update = xegpu.update_nd_offset %arg0, [%c32, %c16]:
       !xegpu.tensor_desc<4x8xf32, #lo0>
     gpu.yield %update
   }
   ...
 ```
 To
 ```
   %r:2 = gpu.warp_execute_on_lane_0(%laneid) -> (vector<4x1xf32>,
   !xegpu.tensor_desc<4x8xf32, #lo0>) {
     ...
     %dead = xegpu.update_nd_offset %arg0, [%c32, %c16]:
       !xegpu.tensor_desc<4x8xf32, #lo0> gpu.yield %dead, %arg0
     gup.yield %dead, %arg0, %c32, %c16
   }
%0 = xegpu.unrealized_conversion_cast %r#1: !xegpu.tensor_desc<4x8xf32,
        #lo0> -> !xegpu.tensor_desc<4x8xf32>
   %1 = xegpu.update_nd_offset %0, [%c32, %c16]:
     !xegpu.tensor_desc<4x8xf32>
   ...
 ```



To unsubscribe from these emails, change your notification settings at https://github.com/llvm/llvm-project/settings/notifications