[TritonGEN] Add triton_gen.cache_controls operation

[victor-eds]

Add triton_gen.cache_controls operation to represent SPV_INTEL_cache_controls decorations in MLIR.

This operation does not convert to any operation in a different dialect, as it supports translation straight to LLVM IR as metadata.

[victor-eds]

This is a highly experimental PR. Not ready to merge at all, as this is just a POC. Missing:

• Cache control specification via special attributes instead of just i32 attribute
• Check this works in the pipeline (translation interface is properly registered) (should be working as the translation interface is being used)
• Support more than one annotation by not overriding MD if present

[whitneywhtsang]

Instead of expressing cache control at the pointer, the new proposal is to express at the memory operation (KhronosGroup/SPIRV-LLVM-Translator#2587).

[victor-eds]

Instead of expressing cache control at the pointer, the new proposal is to express at the memory operation (KhronosGroup/SPIRV-LLVM-Translator#2587).

Updated implementation. This is ready to review now.

[whitneywhtsang]

Given that the new design has cache control on the memory operation instead of the pointer, would it make sense for TritonGEN to add an attribute instead of a new operation that express the cache control on the pointer?

[etiotto]

Given that the new design has cache control on the memory operation instead of the pointer, would it make sense for TritonGEN to add an attribute instead of a new operation that express the cache control on the pointer?

The cache control is a property of the load/store operations rather than the ptr used by those operations. So I think the attribute is a better way to go.

[victor-eds]

The cache control is a property of the load/store operations rather than the ptr used by those operations. So I think the attribute is a better way to go.

My fear is that MLIR transform passes may drop these attributes, as these are not guaranteed to be preserved. That's why I went with this design.

[whitneywhtsang]

My fear is that MLIR transform passes may drop these attributes, as these are not guaranteed to be preserved. That's why I went with this design.

It is the same concern for metadata on memory operations too. Do you think there is a high chance that attributes would be dropped in practice?

[victor-eds]

My fear is that MLIR transform passes may drop these attributes, as these are not guaranteed to be preserved. That's why I went with this design.

It is the same concern for metadata on memory operations too. Do you think there is a high chance that attributes would be dropped in practice?

That will depend on how late in the pipeline these are added. The only way we can make sure these are added is if the cache control information is to be added as part of a conversion to LLVM dialect pattern. Otherwise, even if we have an operation foo correctly annotated and converting to llvm.store, we cannot guarantee the attributes will be carried over in the conversion to llvm.store.

[etiotto]

The cache control is a property of the load/store operations rather than the ptr used by those operations. So I think the attribute is a better way to go.

My fear is that MLIR transform passes may drop these attributes, as these are not guaranteed to be preserved. That's why I went with this design.

Got it, however the chance is small because we aren't running any transformations on the TritonGEN code

[victor-eds]

The cache control is a property of the load/store operations rather than the ptr used by those operations. So I think the attribute is a better way to go.

My fear is that MLIR transform passes may drop these attributes, as these are not guaranteed to be preserved. That's why I went with this design.

Got it, however the chance is small because we aren't running any transformations on the TritonGEN code

Yeah, I was thinking more on:

That will depend on how late in the pipeline these are added. The only way we can make sure these are added is if the cache control information is to be added as part of a conversion to LLVM dialect pattern. Otherwise, even if we have an operation foo correctly annotated and converting to llvm.store, we cannot guarantee the attributes will be carried over in the conversion to llvm.store.

If the cache control is to be set only when creating the corresponding llvm.load/store operation, the attribute approach might work.

[victor-eds]

Are we fine with current approach or do we wanna switch to using discardable attributes in operations? @whitneywhtsang @etiotto

If this cache info is to be inserted as part of the conversion to LLVM dialect, I think we can go with the attribute design.

[whitneywhtsang]

Are we fine with current approach or do we wanna switch to using discardable attributes in operations? @whitneywhtsang @etiotto

My preference is attributes in operations.

[etiotto]

My preference is attributes in operations.

The 2D block load operation in the TritonGen dialect already has an attribute for the cache control. When we lower that operation into a call we cannot pass the cache control as an argument because the "builtin" call we need to generate does not have that argument. Instead of taking the cache control as an argument we need to generate metadata and attach it to the function call.
Can we do that cleanly ?

[victor-eds]

The 2D block load operation in the TritonGen dialect already has an attribute for the cache control. When we lower that operation into a call we cannot pass the cache control as an argument because the "builtin" call we need to generate does not have that argument. Instead of taking the cache control as an argument we need to generate metadata and attach it to the function call.

With the changes I've just pushed, that'd involve attaching an attribute to the call operation as in the examples.