[Mlir-commits] [mlir] [mlir][XeGPU] Add optional layout attribute to LoadGather StoreScatter ops (PR #163414)

[Dmitry Chigarev]

================
@@ -97,7 +97,7 @@ func.func @extf_truncf(%arg0: !xegpu.tensor_desc<8x16xf16>, %arg1: !xegpu.tensor
 // CHECK-NEXT: %[[CST0:.*]] = arith.constant {layout_result_0 = #xegpu.layout<lane_layout = [16], lane_data = [1]>} dense<true> : vector<16xi1>
 // CHECK-NEXT: %[[T2:.*]] = xegpu.create_tdesc %[[ARG1]], %[[CST]] : memref<256xf16>, vector<16xindex> ->
 // CHECK-SAME: !xegpu.tensor_desc<16x16xf16, #xegpu.scatter_tdesc_attr<chunk_size = 16 : i64>, #xegpu.layout<lane_layout = [16, 1], lane_data = [1, 2]>>
-// CHECK-NEXT: %{{.*}} = xegpu.load %[[T2]], %[[CST0]]  {layout_result_0 = #xegpu.layout<lane_layout = [16, 1], lane_data = [1, 2]>}
+// CHECK-NEXT: %{{.*}} = xegpu.load %[[T2]], %[[CST0]]  <{layout = #xegpu.layout<lane_layout = [16, 1], lane_data = [1, 2]>}>
----------------
dchigarev wrote:

> It is checking whether the propagation set the temporarily layout attribute for the load result correct.

The propagation pass [uses `setDistributeLayoutAttr`](https://github.com/llvm/llvm-project/blob/63b83ea213878acde020fc8923ee65f42727009e/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp#L881) to set layouts for the ops. In order to ensure that "permament" layouts are automatically handled by all the passes, I've added support for the new "layout" attribute to the `get/setDistributeLayoutAttr` functions, meaning that the propagation pass now <i>always</i> sets permament layouts for `xegpu.load/store` ops.

I've added additional tests to ensure that the existing passes that work with layouts can handle both "permament" and "temporary" layout attributes, so I think there's nothing wrong with changing the behavior of the propagation pass? (the only problem I see are user-pipelines that do not use our utility functions `get/setDistributeLayoutAttr` and thus cannot handle the new attribute, but are there any?)

I could have preserved the old behavior of the propagation pass by only patching `getDistributeLayoutAttr` to recognize the "permament" layout while keeping `setDistributeLayoutAttr` limited to setting temporary layouts only. I don’t like this asymmetry though between the getter and the setter functions. However, if you have a reason to preserve the old behavior, I can do that.

https://github.com/llvm/llvm-project/pull/163414