[Mlir-commits] [mlir] [mlir][AMDGPU] Relax restrictions on raw_buffer_load (PR #102229)

[Krzysztof Drewniak]

https://github.com/krzysz00 created https://github.com/llvm/llvm-project/pull/102229

The buffer load operation might temporarily admit "illegal" types, like i4, or types not yet known to the rewrite pattern (like f8E5M2) and then be rewritten to legal types before lowering to LLVM. This patch removes the verifier restriction that prevents this pattern. This harmonizes the definition of `amdgpu.raw_buffer_load` with the definition of `amdgpu.raw_buffer_store`.

>From e873ea3215e5d062d5092c1073b8bc2db85f2a4c Mon Sep 17 00:00:00 2001
From: Krzysztof Drewniak <Krzysztof.Drewniak at amd.com>
Date: Thu, 18 Jul 2024 22:53:27 +0000
Subject: [PATCH] [mlir][AMDGPU] Relax restrictions on raw_buffer_load

The buffer load operation might temporarily admit "illegal" types,
like i4, or types not yet known to the rewrite pattern (like f8E5M2)
and then be rewritten to legal types before lowering to LLVM. This
patch removes the verifier restriction that prevents this pattern.
---
 mlir/include/mlir/Dialect/AMDGPU/IR/AMDGPU.td | 12 ++----------
 1 file changed, 2 insertions(+), 10 deletions(-)

diff --git a/mlir/include/mlir/Dialect/AMDGPU/IR/AMDGPU.td b/mlir/include/mlir/Dialect/AMDGPU/IR/AMDGPU.td
index aa2b4543927a7..dddbd837c1cbc 100644
--- a/mlir/include/mlir/Dialect/AMDGPU/IR/AMDGPU.td
+++ b/mlir/include/mlir/Dialect/AMDGPU/IR/AMDGPU.td
@@ -126,11 +126,7 @@ def AMDGPU_RawBufferLoadOp :
                    DefaultValuedAttr<BoolAttr, "true">:$boundsCheck,
                    OptionalAttr<I32Attr>:$indexOffset,
                    Optional<I32>:$sgprOffset)>,
-    Results<(outs AnyTypeOf<[BF16, F16, F32, I32, I8, F8E5M2FNUZ, F8E4M3FNUZ,
-                              VectorOfLengthAndType<[2, 4], [F32, I32]>,
-                              VectorOfLengthAndType<[2, 4, 8], [F16, BF16]>,
-                              VectorOfLengthAndType<[2, 4, 8, 16],
-                                [I8, F8E5M2FNUZ, F8E4M3FNUZ]>]>:$value)> {
+    Results<(outs AnyType:$value)> {
 
   let summary = "Raw Buffer load, exposing GCN features";
   let description = [{
@@ -176,11 +172,7 @@ def AMDGPU_RawBufferLoadOp :
 def AMDGPU_RawBufferStoreOp :
     AMDGPU_Op<"raw_buffer_store", [AllElementTypesMatch<["value", "memref"]>,
       AttrSizedOperandSegments]>,
-    Arguments<(ins AnyTypeOf<[BF16, F16, F32, I32, I8, F8E5M2FNUZ, F8E4M3FNUZ,
-                              VectorOfLengthAndType<[2, 4], [F32, I32]>,
-                              VectorOfLengthAndType<[2, 4, 8], [F16, BF16]>,
-                              VectorOfLengthAndType<[2, 4, 8, 16],
-                                [I8, F8E5M2FNUZ, F8E4M3FNUZ]>]>:$value,
+    Arguments<(ins AnyType:$value,
                    Arg<AnyMemRef, "buffer to store to", [MemWrite]>:$memref,
                    Variadic<I32>:$indices,
                    DefaultValuedAttr<BoolAttr, "true">:$boundsCheck,