Add scalar support for amdgpu.raw_buffer_{load,store}

[giuseros]

Summary

Introduce the possibility to load/store scalars via amdgpu.raw_buffer_{load,store}

[giuseros]

Correct a typo stemmed from rebase

[krzysz00]

Looks good to me! I can land this if you'd like

[giuseros]

In D146413#4206234, @krzysz00 wrote:

Looks good to me! I can land this if you'd like

Yes, please! Thanks @krzysz00 !

mlir/include/mlir/Dialect/AMDGPU/AMDGPU.td

let description = [{
     - The thread ID addition bit is off
     - If `boundsCheck` is false and the target chipset is RDNA, OOB_SELECT is set
       to 2 to disable bounds checks, otherwise it is 3
     - The cache coherency bits are off
   }];
   let assemblyFormat = [{
     attr-dict $memref `[` $indices `]`
       (`sgprOffset` $sgprOffset^)? `:`
-      type($memref) `,` type($indices) `->` type($value)
+      type($memref) (`,` type($indices)^)? `->` type($value)
   }];
   let hasCanonicalizer = 1;
   let hasVerifier = 1;
 }
 
 /// Raw buffer store
 def AMDGPU_RawBufferStoreOp :
     AMDGPU_Op<"raw_buffer_store", [AllElementTypesMatch<["value", "memref"]>,
let description = [{
     components is partically completed is chipset-dependent.
 
     See `amdgpu.raw_buffer_load` for a description of how the underlying
     instruction is constructed.
   }];
   let assemblyFormat = [{
     attr-dict $value `->` $memref `[` $indices `]`
       (`sgprOffset` $sgprOffset^)? `:`
-      type($value) `->` type($memref) `,` type($indices)
+      type($value) `->` type($memref) (`,` type($indices)^)?
   }];
   let hasCanonicalizer = 1;
   let hasVerifier = 1;
 }
 
 // Raw buffer atomic floating point add
 def AMDGPU_RawBufferAtomicFaddOp :
     AMDGPU_Op<"raw_buffer_atomic_fadd", [AllElementTypesMatch<["value", "memref"]>,

mlir/lib/Conversion/AMDGPUToROCDL/AMDGPUToROCDL.cpp

matchAndRewrite(GpuOp gpuOp, typename GpuOp::Adaptor adaptor,
     }
     Value word3Const = createI32Constant(rewriter, loc, word3);
     resource = rewriter.create<LLVM::InsertElementOp>(
         loc, llvm4xI32, resource, word3Const,
         this->createIndexConstant(rewriter, loc, 3));
     args.push_back(resource);
 
     // Indexing (voffset)
-    Value voffset;
+    Value voffset = createI32Constant(rewriter, loc, 0);
     for (auto pair : llvm::enumerate(adaptor.getIndices())) {
       size_t i = pair.index();
       Value index = pair.value();
       Value strideOp;
       if (ShapedType::isDynamic(strides[i])) {
         strideOp = rewriter.create<LLVM::MulOp>(
             loc, memrefDescriptor.stride(rewriter, loc, i), byteWidthConst);
       } else {
         strideOp =
             createI32Constant(rewriter, loc, strides[i] * elementByteWidth);
       }
       index = rewriter.create<LLVM::MulOp>(loc, index, strideOp);
-      voffset =
-          voffset ? rewriter.create<LLVM::AddOp>(loc, voffset, index) : index;
+      voffset = rewriter.create<LLVM::AddOp>(loc, voffset, index);
     }
     if (adaptor.getIndexOffset()) {
       int32_t indexOffset = *gpuOp.getIndexOffset() * elementByteWidth;
       Value extraOffsetConst = createI32Constant(rewriter, loc, indexOffset);
       voffset =
           voffset ? rewriter.create<LLVM::AddOp>(loc, voffset, extraOffsetConst)
                   : extraOffsetConst;
     }

mlir/test/Conversion/AMDGPUToROCDL/amdgpu-to-rocdl.mlir

 // RUN: mlir-opt %s -convert-amdgpu-to-rocdl=chipset=gfx908 | FileCheck %s
 // RUN: mlir-opt %s -convert-amdgpu-to-rocdl=chipset=gfx1030 | FileCheck %s --check-prefix=RDNA
 
+// CHECK-LABEL: func @gpu_gcn_raw_buffer_load_scalar_i32
+func.func @gpu_gcn_raw_buffer_load_scalar_i32(%buf: memref<i32>) -> i32 {
+  // CHECK: %[[ptr:.*]] = llvm.ptrtoint
+  // CHECK: %[[lowHalf:.*]] = llvm.trunc %[[ptr]] : i64 to i32
+  // CHECK: %[[resource_1:.*]] = llvm.insertelement %[[lowHalf]]
+  // CHECK: %[[highHalfI64:.*]] = llvm.lshr %[[ptr]]
+  // CHECK: %[[highHalfI32:.*]] = llvm.trunc %[[highHalfI64]] : i64 to i32
+  // CHECK: %[[highHalf:.*]] = llvm.and %[[highHalfI32]], %{{.*}} : i32
+  // CHECK: %[[resource_2:.*]] = llvm.insertelement %[[highHalf]], %[[resource_1]]
+  // CHECK: %[[numRecords:.*]] = llvm.mlir.constant(4 : i32)
+  // CHECK: %[[resource_3:.*]] = llvm.insertelement %[[numRecords]], %[[resource_2]]
+  // CHECK: %[[word3:.*]] = llvm.mlir.constant(159744 : i32)
+  // RDNA: %[[word3:.*]] = llvm.mlir.constant(822243328 : i32)
+  // CHECK: %[[resource:.*]] = llvm.insertelement %[[word3]], %[[resource_3]]
+  // CHECK: %[[ret:.*]] = rocdl.raw.buffer.load %[[resource]], %{{.*}}, %{{.*}}, %{{.*}} : i32
+  // CHECK: return %[[ret]]
+  %0 = amdgpu.raw_buffer_load {boundsCheck = true} %buf[] : memref<i32> -> i32
+  func.return %0 : i32
+}
+
 // CHECK-LABEL: func @gpu_gcn_raw_buffer_load_i32
 func.func @gpu_gcn_raw_buffer_load_i32(%buf: memref<64xi32>, %idx: i32) -> i32 {
   // CHECK: %[[ptr:.*]] = llvm.ptrtoint
   // CHECK: %[[lowHalf:.*]] = llvm.trunc %[[ptr]] : i64 to i32
   // CHECK: %[[resource_1:.*]] = llvm.insertelement %[[lowHalf]]
   // CHECK: %[[highHalfI64:.*]] = llvm.lshr %[[ptr]]
   // CHECK: %[[highHalfI32:.*]] = llvm.trunc %[[highHalfI64]] : i64 to i32
   // CHECK: %[[highHalf:.*]] = llvm.and %[[highHalfI32]], %{{.*}} : i32
func.func @gpu_gcn_raw_buffer_load_4xf8E4M3FNUZ(%buf: memref<64xf8E4M3FNUZ>, %idx: i32) -> vector<4xf8E4M3FNUZ> {
   // CHECK: %[[cast:.*]] = llvm.bitcast %[[loaded]] : i32 to vector<4xi8>
   // CHECK: %[[ret:.*]] = builtin.unrealized_conversion_cast %[[cast]] : vector<4xi8> to vector<4xf8E4M3FNUZ>
   // CHECK: return %[[ret]]
   %0 = amdgpu.raw_buffer_load {boundsCheck = true} %buf[%idx] : memref<64xf8E4M3FNUZ>, i32 -> vector<4xf8E4M3FNUZ>
   func.return %0 : vector<4xf8E4M3FNUZ>
 }
 
 // Since the lowering logic is shared with loads, only bitcasts need to be rechecked
+// CHECK-LABEL: func @gpu_gcn_raw_buffer_store_scalar_i32
+func.func @gpu_gcn_raw_buffer_store_scalar_i32(%value: i32, %buf: memref<i32>) {
+  // CHECK: %[[numRecords:.*]] = llvm.mlir.constant(4 : i32)
+  // CHECK: llvm.insertelement{{.*}}%[[numRecords]]
+  // CHECK: %[[word3:.*]] = llvm.mlir.constant(159744 : i32)
+  // CHECK: %[[resource:.*]] = llvm.insertelement{{.*}}%[[word3]]
+  // CHECK: rocdl.raw.buffer.store %{{.*}} %[[resource]], %{{.*}}, %{{.*}}, %{{.*}} : i32
+  amdgpu.raw_buffer_store {boundsCheck = true} %value -> %buf[] : i32 -> memref<i32>
+  func.return
+}
+
 // CHECK-LABEL: func @gpu_gcn_raw_buffer_store_i32
 func.func @gpu_gcn_raw_buffer_store_i32(%value: i32, %buf: memref<64xi32>, %idx: i32) {
   // CHECK: %[[numRecords:.*]] = llvm.mlir.constant(256 : i32)
   // CHECK: llvm.insertelement{{.*}}%[[numRecords]]
   // CHECK: %[[word3:.*]] = llvm.mlir.constant(159744 : i32)
   // CHECK: %[[resource:.*]] = llvm.insertelement{{.*}}%[[word3]]
   // CHECK: rocdl.raw.buffer.store %{{.*}} %[[resource]], %{{.*}}, %{{.*}}, %{{.*}} : i32
   amdgpu.raw_buffer_store {boundsCheck = true} %value -> %buf[%idx] : i32 -> memref<64xi32>, i32

mlir/test/Dialect/AMDGPU/ops.mlir

 
 // CHECK-LABEL: func @raw_buffer_load_f32_from_rank_4
 func.func @raw_buffer_load_f32_from_rank_4(%src : memref<128x64x32x16xf32>, %offset : i32, %idx0 : i32, %idx1 : i32, %idx2 : i32, %idx3 : i32) -> f32 {
   // CHECK: amdgpu.raw_buffer_load {indexOffset = 1 : i32} %{{.*}}[%{{.*}}, %{{.*}}, %{{.*}}] sgprOffset %{{.*}} : memref<128x64x32x16xf32>, i32, i32, i32, i32 -> f32
   %0 = amdgpu.raw_buffer_load {boundsCheck = true, indexOffset = 1 : i32} %src[%idx0, %idx1, %idx2, %idx3] sgprOffset %offset : memref<128x64x32x16xf32>, i32, i32, i32, i32 -> f32
   func.return %0 : f32
 }
 
+// CHECK-LABEL: func @raw_buffer_load_scalar
+func.func @raw_buffer_load_scalar(%src : memref<f32>) -> f32 {
+  // CHECK: amdgpu.raw_buffer_load {indexOffset = 1 : i32} %{{.*}}[] : memref<f32> -> f32
+  %0 = amdgpu.raw_buffer_load {indexOffset = 1 : i32} %src[] : memref<f32> -> f32
+  func.return %0 : f32
+}
+
 // CHECK-LABEL: func @raw_buffer_load_4xf32_from_rank_4
 func.func @raw_buffer_load_4xf32_from_rank_4(%src : memref<128x64x32x16xf32>, %offset : i32, %idx0 : i32, %idx1 : i32, %idx2 : i32, %idx3 : i32) -> vector<4xf32> {
   // CHECK: amdgpu.raw_buffer_load {indexOffset = 1 : i32} %{{.*}}[%{{.*}}, %{{.*}}, %{{.*}}] sgprOffset %{{.*}} : memref<128x64x32x16xf32>, i32, i32, i32, i32 -> vector<4xf32>
   %0 = amdgpu.raw_buffer_load {boundsCheck = true, indexOffset = 1 : i32} %src[%idx0, %idx1, %idx2, %idx3] sgprOffset %offset : memref<128x64x32x16xf32>, i32, i32, i32, i32 -> vector<4xf32>
   func.return %0 : vector<4xf32>
 }
 
 // CHECK-LABEL: func @raw_buffer_store_f32_to_rank_1

 
 // CHECK-LABEL: func @raw_buffer_store_4xf32_to_rank_4
 func.func @raw_buffer_store_4xf32_to_rank_4(%value : vector<4xf32>, %dst : memref<128x64x32x16xf32>, %offset : i32, %idx0 : i32, %idx1 : i32, %idx2 : i32, %idx3 : i32) {
   // CHECK: amdgpu.raw_buffer_store {indexOffset = 1 : i32} %{{.*}} -> %{{.*}}[%{{.*}}, %{{.*}}, %{{.*}}] sgprOffset %{{.*}} : vector<4xf32> -> memref<128x64x32x16xf32>, i32, i32, i32, i32
   amdgpu.raw_buffer_store {boundsCheck = true, indexOffset = 1 : i32} %value -> %dst[%idx0, %idx1, %idx2, %idx3] sgprOffset %offset : vector<4xf32> -> memref<128x64x32x16xf32>, i32, i32, i32, i32
   func.return
 }
 
+// CHECK-LABEL: func @raw_buffer_store_scalar
+func.func @raw_buffer_store_scalar(%value : f32, %dst : memref<f32>) {
+  // CHECK: amdgpu.raw_buffer_store {indexOffset = 1 : i32} %{{.*}} -> %{{.*}}[] : f32 -> memref<f32>
+  amdgpu.raw_buffer_store {indexOffset = 1 : i32} %value -> %dst[] : f32 -> memref<f32>
+  func.return
+}
+
 // CHECK-LABEL: func @raw_buffer_atomic_fadd_f32_to_rank_1
 func.func @raw_buffer_atomic_fadd_f32_to_rank_1(%value : f32, %dst : memref<128xf32>, %offset : i32, %idx0 : i32) {
   // CHECK: amdgpu.raw_buffer_atomic_fadd {indexOffset = 1 : i32} %{{.*}} -> %{{.*}}[{{.*}}] sgprOffset %{{.*}} : f32 -> memref<128xf32>, i32
   amdgpu.raw_buffer_atomic_fadd {boundsCheck = true, indexOffset = 1 : i32} %value -> %dst[%idx0] sgprOffset %offset : f32 -> memref<128xf32>, i32
   func.return
 }
 
 // CHECK-LABEL: func @raw_buffer_atomic_fadd_f32_to_rank_4