[MLIR][AMDGPU] Add amdgpu.sched_barrier

This commit adds sched_barrier operator
to AMDGPU dialect that lowers to rocdl.sched.barrier.

Author: manupak

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

@@ -433,6 +433,47 @@ def AMDGPU_LDSBarrierOp : AMDGPU_Op<"lds_barrier"> {
   let assemblyFormat = "attr-dict";
 }
 
+def AMDGPU_SchedBarrierOpOpt : I32EnumAttr<"sched_barrier_opt_enum",
+    "The possible options for scheduling barriers",
+    [
+      I32EnumAttrCase<"allow_none",            0x0000>,
+      I32EnumAttrCase<"allow_non_mem_non_sideffect",  0x0001>,
+      I32EnumAttrCase<"allow_valu", 0x0002>,
+      I32EnumAttrCase<"allow_salu", 0x0004>,
+      I32EnumAttrCase<"allow_mfma_wmma",  0x0008>,
+      I32EnumAttrCase<"allow_all_vmem",  0x0010>,
+      I32EnumAttrCase<"allow_vmem_read",  0x0020>,
+      I32EnumAttrCase<"allow_vmem_write", 0x0040>,
+      I32EnumAttrCase<"allow_all_ds", 0x0080>,
+      I32EnumAttrCase<"allow_ds_read", 0x0100>,
+      I32EnumAttrCase<"allow_ds_write", 0x0200>,
+      I32EnumAttrCase<"allow_transcendental", 0x0400>
+    ]> {
+  let genSpecializedAttr = 0;
+  let cppNamespace = "::mlir::amdgpu";
+}
+
+def AMDGPU_SchedBarrierOpOptAttr : EnumAttr<AMDGPU_Dialect, AMDGPU_SchedBarrierOpOpt,
+  "sched_barrier_opt">;
+
+def AMDGPU_SchedBarrierOpOptArrayAttr : TypedArrayAttrBase<AMDGPU_SchedBarrierOpOptAttr,
+                                        "sched_barrier_opt array attribute">;
+
+def AMDGPU_SchedBarrierOp :
+  AMDGPU_Op<"sched_barrier">,
+  Arguments<(ins  AMDGPU_SchedBarrierOpOptArrayAttr:$opts)>
+  {
+  let summary = "Barrier that limits the backend scheduler of instruction movement";
+  let description = [{
+    `amdgpu.sched_barrier` serves as a barrier that could be
+    configured to restrict movements of instructions through it as
+    defined by sched_barrier_opts.
+  }];
+  let assemblyFormat = [{
+    $opts attr-dict
+  }];
+}
+
 def AMDGPU_MFMAPermB : I32EnumAttr<"MFMAPermB",
     "The possible permutations of the lanes storing B available in an MFMA",
     [

mlir/lib/Conversion/AMDGPUToROCDL/AMDGPUToROCDL.cpp

@@ -321,6 +321,26 @@ struct LDSBarrierOpLowering : public ConvertOpToLLVMPattern<LDSBarrierOp> {
     return success();
   }
 };
+
+struct SchedBarrierOpLowering : public ConvertOpToLLVMPattern<SchedBarrierOp> {
+  SchedBarrierOpLowering(LLVMTypeConverter &converter, Chipset chipset)
+      : ConvertOpToLLVMPattern<SchedBarrierOp>(converter), chipset(chipset) {}
+
+  Chipset chipset;
+
+  LogicalResult
+  matchAndRewrite(SchedBarrierOp op, SchedBarrierOp::Adaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override {
+    uint32_t combinedOpt = 0;
+    for (Attribute opt : op.getOpts()) {
+      combinedOpt |=
+          (uint32_t)cast<amdgpu::sched_barrier_opt_enumAttr>(opt).getValue();
+    }
+    rewriter.replaceOpWithNewOp<ROCDL::SchedBarrier>(op, combinedOpt);
+    return success();
+  }
+};
+
 } // namespace
 
 /// If `input` is a vector of bytes, concatentate those bytes in little-endian
@@ -879,8 +899,8 @@ void mlir::populateAMDGPUToROCDLConversionPatterns(LLVMTypeConverter &converter,
                                ROCDL::RawPtrBufferAtomicUminOp>,
            RawBufferOpLowering<RawBufferAtomicCmpswapOp,
                                ROCDL::RawPtrBufferAtomicCmpSwap>,
-           LDSBarrierOpLowering, MFMAOpLowering, WMMAOpLowering,
-           ExtPackedFp8OpLowering, PackedTrunc2xFp8OpLowering,
+           LDSBarrierOpLowering, SchedBarrierOpLowering, MFMAOpLowering,
+           WMMAOpLowering, ExtPackedFp8OpLowering, PackedTrunc2xFp8OpLowering,
            PackedStochRoundFp8OpLowering>(converter, chipset);
 }
 

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

@@ -226,3 +226,34 @@ func.func @lds_barrier() {
   amdgpu.lds_barrier
   func.return
 }
+
+// CHECK-LABEL: func @sched_barrier
+func.func @sched_barrier() {
+  // rocdl.sched.barrier 0
+  amdgpu.sched_barrier [#amdgpu<sched_barrier_opt allow_none>]
+  // rocdl.sched.barrier 1
+  amdgpu.sched_barrier [#amdgpu<sched_barrier_opt allow_non_mem_non_sideffect>]
+  // rocdl.sched.barrier 2
+  amdgpu.sched_barrier [#amdgpu<sched_barrier_opt allow_valu>]
+  // rocdl.sched.barrier 4
+  amdgpu.sched_barrier [#amdgpu<sched_barrier_opt allow_salu>]
+  // rocdl.sched.barrier 8
+  amdgpu.sched_barrier [#amdgpu<sched_barrier_opt allow_mfma_wmma>]
+  // rocdl.sched.barrier 16
+  amdgpu.sched_barrier [#amdgpu<sched_barrier_opt allow_all_vmem>]
+  // rocdl.sched.barrier 32
+  amdgpu.sched_barrier [#amdgpu<sched_barrier_opt allow_vmem_read>]
+  // rocdl.sched.barrier 64
+  amdgpu.sched_barrier [#amdgpu<sched_barrier_opt allow_vmem_write>]
+  // rocdl.sched.barrier 128
+  amdgpu.sched_barrier [#amdgpu<sched_barrier_opt allow_all_ds>]
+  // rocdl.sched.barrier 256
+  amdgpu.sched_barrier [#amdgpu<sched_barrier_opt allow_ds_read>]
+  // rocdl.sched.barrier 512
+  amdgpu.sched_barrier [#amdgpu<sched_barrier_opt allow_ds_write>]
+  // rocdl.sched.barrier 1024
+  amdgpu.sched_barrier [#amdgpu<sched_barrier_opt allow_transcendental>]
+  // rocdl.sched.barrier 18
+  amdgpu.sched_barrier [#amdgpu<sched_barrier_opt allow_valu>, #amdgpu<sched_barrier_opt allow_all_vmem>]
+  func.return
+}

mlir/test/Dialect/AMDGPU/ops.mlir

@@ -109,6 +109,15 @@ func.func @lds_barrier() {
   func.return
 }
 
+// CHECK-LABEL: func @sched_barrier
+func.func @sched_barrier() {
+  // CHECK: amdgpu.sched_barrier [#amdgpu<sched_barrier_opt allow_none>]
+  amdgpu.sched_barrier [#amdgpu<sched_barrier_opt allow_none>]
+  // CHECK: amdgpu.sched_barrier [#amdgpu<sched_barrier_opt allow_valu>, #amdgpu<sched_barrier_opt allow_all_vmem>]
+  amdgpu.sched_barrier [#amdgpu<sched_barrier_opt allow_valu>, #amdgpu<sched_barrier_opt allow_all_vmem>]
+  func.return
+}
+
 // CHECK-LABEL: func @mfma
 func.func @mfma(%arg0 : f32, %arg1 : vector<32xf32>) -> vector<32xf32> {
   // CHECK: amdgpu.mfma