[mlir][AMDGPU] Add lds_barrier op

[krzysz00]

Summary

The lds_barrier op allows workgroups to wait at a barrier for
operations to/from their local data store (LDS) to complete without
incurring the performance penalties of a full memory fence.

[kerbowa]

Is this really necessary if we have https://reviews.llvm.org/D120544. A barrier has no requirement to wait for LDS and VMEM it only does so currently because of a bug. Using inline asm like this seems like it will eventually cause problems although I'm not too familiar with MLIR.

[krzysz00]

This is a single op that expands to both a waitcnt on LDS and a barrier. I
can go digging tomorrow for what we used to lower this to some sort of
fence and a barrier, but I recall ( @whchung who may have more detail) that
using this bit of inline assembly to work around what I think was the lack
of an LDS-only fence gave a noticable performance increase

[kerbowa]

In D129522#3643798, @krzysz00 wrote:

This is a single op that expands to both a waitcnt on LDS and a barrier. I
can go digging tomorrow for what we used to lower this to some sort of
fence and a barrier, but I recall ( @whchung who may have more detail) that
using this bit of inline assembly to work around what I think was the lack
of an LDS-only fence gave a noticable performance increase

It's not really the lack of an LDS-only fence that was the problem. You could use the intrinsics for a barrier and a waitcnt for example. The problem with that idea is the backend would always wait for both LDS and VM at barriers when it saw them. So you would need to use inline asm to make the barrier invisible which could cause other problems in the backend with combining waitcnt for example. There is a lot of desire to reland https://reviews.llvm.org/D120544 in the near future. After that, the inline asm will not be needed for gfx90a+ and gfx10+. So then you should probably consider not using the inline asm for those architectures.

[krzysz00]

Re the above, you mentioned gfx90a+ ... what about gfx908?

And even if there'll be a better way to do it in the future, this type of partier+fence construct, however we implement it, is useful for kernel developers now. We can always swap out the implementation later once the compiler fix lands.

[nirvedhmeshram]

mlir/lib/Conversion/AMDGPUToROCDL/AMDGPUToROCDL.cpp

257

Should the op have side effects?

[krzysz00]

mlir/lib/Conversion/AMDGPUToROCDL/AMDGPUToROCDL.cpp

257

I don't think so.

[nirvedhmeshram]

mlir/lib/Conversion/AMDGPUToROCDL/AMDGPUToROCDL.cpp

257

So can we change this to /*has_side_effects=*/false then?

[krzysz00]

@kerbowa I'd like to still land this because the lds_barrier is a useful abstraction over whatever combination of (fence + barrier)/inline assembly/... you end up needing to use to implement it.

mlir/lib/Conversion/AMDGPUToROCDL/AMDGPUToROCDL.cpp

257

... right, that field. has_side_effects will get the compiler to not throw this out on account that it doesn't write any registers. So yeah, it some sense, it does have side effects (because it needs to stay put)

[kerbowa]

I'm not objecting to the change, just pointing out that you may miss out on some optimization since this is lowered to inline ASM, and that you may want to lower it to intrinsics in the future.

mlir/lib/Conversion/AMDGPUToROCDL/AMDGPUToROCDL.cpp

257

I think it should definitely be marked as having side effects. It may be scheduled around memops if it isn't. The backend cannot see what instructions are inside inline asm.

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

def AMDGPU_RawBufferAtomicFaddOp :
   let assemblyFormat = [{
     attr-dict $value `->` $memref `[` $indices `]`
       (`sgprOffset` $sgprOffset^)? `:`
       type($value) `->` type($memref) `,` type($indices)
   }];
   let hasVerifier = 1;
 }
 
+def AMDGPU_LDSBarrierOp : AMDGPU_Op<"lds_barrier"> {
+  let summary = "Barrier that includes a wait for LDS memory operations.";
+  let description = [{
+    `amdgpu.lds_barrier` is both a barrier (all workitems in a workgroup must reach
+    the barrier before any of them may proceed past it) and a wait for all
+    operations that affect the Local Data Store (LDS) issued from that wrokgroup
+    to complete before the workgroup may continue. Since the LDS is per-workgroup
+    memory, this barrier may be used, for example, to ensure all workitems have
+    written data to LDS before any workitem attempts to read from it.
+
+    Note that `lds_barrier` does **not** force reads to or from global memory
+    to complete before execution continues. Therefore, it should be used when
+    operations on global memory can be issued far in advance of when their results
+    are used (for example, by writing them to LDS).
+  }];
+  let assemblyFormat = "attr-dict";
+}
+
+
 #endif // AMDGPU

mlir/lib/Conversion/AMDGPUToROCDL/AMDGPUToROCDL.cpp

if (lowered->getNumResults() == 1) {
       rewriter.replaceOp(gpuOp, replacement);
     } else {
       rewriter.eraseOp(gpuOp);
     }
     return success();
   }
 };
 
+struct LDSBarrierOpLowering : public ConvertOpToLLVMPattern<LDSBarrierOp> {
+  using ConvertOpToLLVMPattern<LDSBarrierOp>::ConvertOpToLLVMPattern;
+
+  LogicalResult
+  matchAndRewrite(LDSBarrierOp op, LDSBarrierOp::Adaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override {
+    auto asmDialectAttr = LLVM::AsmDialectAttr::get(rewriter.getContext(),
+                                                    LLVM::AsmDialect::AD_ATT);
+    const char *asmStr = "s_waitcnt lgkmcnt(0)\ns_barrier";
+    const char *constraints = "";
+    rewriter.replaceOpWithNewOp<LLVM::InlineAsmOp>(
+        op,
+        /*resultTypes=*/TypeRange(), /*operands=*/ValueRange(),
+        /*asm_string=*/asmStr, constraints, /*has_side_effects=*/true,
+        /*is_align_stack=*/false, /*asm_dialect=*/asmDialectAttr,
+        /*operand_attrs=*/ArrayAttr());
+    return success();
+  }
+};
+
 struct ConvertAMDGPUToROCDLPass
     : public ConvertAMDGPUToROCDLBase<ConvertAMDGPUToROCDLPass> {
   ConvertAMDGPUToROCDLPass() = default;
 
   void runOnOperation() override {
     MLIRContext *ctx = &getContext();
     FailureOr<Chipset> maybeChipset = Chipset::parse(chipset);
     if (failed(maybeChipset)) {
if (failed(applyPartialConversion(getOperation(), target,
       signalPassFailure();
   }
 };
 } // namespace
 
 void mlir::populateAMDGPUToROCDLConversionPatterns(LLVMTypeConverter &converter,
                                                    RewritePatternSet &patterns,
                                                    Chipset chipset) {
+  patterns.add<LDSBarrierOpLowering>(converter);
   patterns.add<
       RawBufferOpLowering<RawBufferLoadOp, ROCDL::RawBufferLoadOp>,
       RawBufferOpLowering<RawBufferStoreOp, ROCDL::RawBufferStoreOp>,
       RawBufferOpLowering<RawBufferAtomicFaddOp, ROCDL::RawBufferAtomicFAddOp>>(
       converter, chipset);
 }
 
 std::unique_ptr<Pass> mlir::createConvertAMDGPUToROCDLPass() {
   return std::make_unique<ConvertAMDGPUToROCDLPass>();
 }

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

func.func @gpu_gcn_raw_buffer_atomic_fadd_f32(%value: f32, %buf: memref<64xf32>, %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.atomic.fadd %{{.*}} %[[resource]], %{{.*}}, %{{.*}}, %{{.*}} : f32
   amdgpu.raw_buffer_atomic_fadd {boundsCheck = true} %value -> %buf[%idx] : f32 -> memref<64xf32>, i32
   func.return
 }
+
+// CHECK-LABEL: func @lds_barrier
+func.func @lds_barrier() {
+  // CHECK: llvm.inline_asm has_side_effects asm_dialect = att "s_waitcnt lgkmcnt(0)\0As_barrier"
+  amdgpu.lds_barrier
+  func.return
+}

mlir/test/Dialect/AMDGPU/ops.mlir

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