[AMD] Add BufferOp interface (#8600)

This PR introduces a common interface for buffer ops.

---------

Co-authored-by: Alexander Efimov <[email protected]>

Author: alefimov-amd

third_party/amd/include/Dialect/TritonAMDGPU/IR/CMakeLists.txt

@@ -15,4 +15,9 @@ mlir_tablegen(TritonAMDGPUEnums.h.inc -gen-enum-decls)
 mlir_tablegen(TritonAMDGPUEnums.cpp.inc -gen-enum-defs)
 mlir_tablegen(TritonAMDGPUAttrDefs.h.inc -gen-attrdef-decls)
 mlir_tablegen(TritonAMDGPUAttrDefs.cpp.inc -gen-attrdef-defs)
+
+set(LLVM_TARGET_DEFINITIONS TritonAMDGPUOpInterfaces.td)
+mlir_tablegen(TritonAMDGPUOpInterfaces.h.inc -gen-op-interface-decls)
+mlir_tablegen(TritonAMDGPUOpInterfaces.cpp.inc -gen-op-interface-defs)
+
 add_public_tablegen_target(TritonAMDGPUAttrDefsIncGen)

third_party/amd/include/Dialect/TritonAMDGPU/IR/Dialect.h

@@ -41,6 +41,7 @@
 #define GET_ATTRDEF_CLASSES
 #include "amd/include/Dialect/TritonAMDGPU/IR/TritonAMDGPUAttrDefs.h.inc"
 
+#include "amd/include/Dialect/TritonAMDGPU/IR/TritonAMDGPUOpInterfaces.h.inc"
 #define GET_OP_CLASSES
 #include "amd/include/Dialect/TritonAMDGPU/IR/Ops.h.inc"
 

third_party/amd/include/Dialect/TritonAMDGPU/IR/TritonAMDGPUOpInterfaces.td

@@ -0,0 +1,42 @@
+#ifndef TRITON_AMDGPU_OP_INTERFACES
+#define TRITON_AMDGPU_OP_INTERFACES
+
+include "mlir/IR/OpBase.td"
+
+def BufferOpInterface : OpInterface<"BufferOpInterface"> {
+  let description = [{
+    This interface is implemented by buffer load/store operations.
+    It provides methods to access common properties such base pointer, offset, mask and others.
+  }];
+
+  let cppNamespace = "::mlir::triton::amdgpu";
+
+  let methods = [
+    InterfaceMethod<
+      /*desc=*/"Get operation base ptr.",
+      /*retType=*/"::mlir::TypedValue<::mlir::triton::PointerType>",
+      /*methodName=*/"getPtr">,
+    InterfaceMethod<
+      /*desc=*/"Get mutable operation base ptr.",
+      /*retType=*/"::mlir::OpOperand &",
+      /*methodName=*/"getPtrMutable">,
+    InterfaceMethod<
+      /*desc=*/"Get operation offset tensor.",
+      /*retType=*/"::mlir::TypedValue<::mlir::TensorType>",
+      /*methodName=*/"getOffsets">,
+    InterfaceMethod<
+      /*desc=*/"Get mutable operation offset tensor.",
+      /*retType=*/"::mlir::OpOperand &",
+      /*methodName=*/"getOffsetsMutable">,
+    InterfaceMethod<
+      /*desc=*/"Get operation stride.",
+      /*retType=*/"::mlir::TypedValue<::mlir::IntegerType>",
+      /*methodName=*/"getStride">,
+    InterfaceMethod<
+      /*desc=*/"Get mutable operation stride.",
+      /*retType=*/"::mlir::MutableOperandRange ",
+      /*methodName=*/"getStrideMutable">
+  ];
+}
+
+#endif // TRITON_AMDGPU_OP_INTERFACES

third_party/amd/include/Dialect/TritonAMDGPU/IR/TritonAMDGPUOps.td

@@ -41,6 +41,7 @@ include "mlir/Interfaces/InferTypeOpInterface.td"
 include "mlir/Interfaces/SideEffectInterfaces.td" // Pure
 include "TritonAMDGPUDialect.td"
 include "TritonAMDGPUAttrDefs.td"
+include "TritonAMDGPUOpInterfaces.td"
 
 
 class TT_AMDGPU_Op<string mnemonic, list<Trait> traits = []> :
@@ -283,6 +284,7 @@ def CondBarrierOp : TT_AMDGPU_Op<"cond_barrier"> {
 def BufferLoadOp : TT_AMDGPU_Op<"buffer_load", [
   SameLoadStoreOperandsAndResultEncoding,
   AttrSizedOperandSegments,
+  BufferOpInterface,
   TypesMatchWith<"result element type matches the pointed type of ptr", "result", "ptr", "getPointerTypeToElement($_self)">,
   TypesMatchWith<"result and offsets have the same shape", "result", "offsets", "getI32SameShape($_self)">,
   TypesMatchWith<"result and mask have the same shape", "result", "mask", "getI1SameShape($_self)",
@@ -328,6 +330,7 @@ def BufferLoadOp : TT_AMDGPU_Op<"buffer_load", [
 
 def BufferLoadToLocalOp : TT_AMDGPU_Op<"buffer_load_to_local", [
   AttrSizedOperandSegments,
+  BufferOpInterface,
   TypesMatchWith<"dest element type matches pointee type of ptr", "dest", "ptr", "getPointerTypeToElement($_self)">,
   TypesMatchWith<"infer mask shape from offsets",
                  "offsets", "mask", "getI1SameShape($_self)",
@@ -364,6 +367,7 @@ def BufferLoadToLocalOp : TT_AMDGPU_Op<"buffer_load_to_local", [
 def BufferAtomicRMWOp : TT_AMDGPU_Op<"buffer_atomic_rmw", [
   AttrSizedOperandSegments,
   SameLoadStoreOperandsAndResultEncoding,
+  BufferOpInterface,
   TypesMatchWith<"result element type matches the value type", "result", "value", "$_self">,
   TypesMatchWith<"result element type matches the pointed type of ptr", "result", "ptr", "getPointerTypeToElement($_self)">,
   TypesMatchWith<"result and offsets have the same shape", "result", "offsets", "getI32SameShape($_self)">,
@@ -410,6 +414,7 @@ def BufferAtomicRMWOp : TT_AMDGPU_Op<"buffer_atomic_rmw", [
 //===----------------------------------------------------------------------===//
 def BufferAtomicCASOp : TT_AMDGPU_Op<"buffer_atomic_cas", [
   SameLoadStoreOperandsAndResultEncoding,
+  BufferOpInterface,
   TypesMatchWith<"result element type matches the val type", "result", "val", "$_self">,
   TypesMatchWith<"result element type matches the cmp type", "result", "cmp", "$_self">,
   TypesMatchWith<"result element type matches the pointed type of ptr", "result", "ptr", "getPointerTypeToElement($_self)">,
@@ -452,6 +457,7 @@ def BufferAtomicCASOp : TT_AMDGPU_Op<"buffer_atomic_cas", [
 def BufferStoreOp : TT_AMDGPU_Op<"buffer_store", [
   AttrSizedOperandSegments,
   SameLoadStoreOperandsEncoding,
+  BufferOpInterface,
   TypesMatchWith<"value element type matches the pointed type of ptr", "value", "ptr", "getPointerTypeToElement($_self)">,
   TypesMatchWith<"value and offsets have the same shape", "value", "offsets", "getI32SameShape($_self)">,
   TypesMatchWith<"value and mask have the same shape", "value", "mask", "getI1SameShape($_self)",

third_party/amd/lib/Dialect/TritonAMDGPU/IR/Dialect.cpp

@@ -62,6 +62,7 @@ void mlir::triton::amdgpu::TritonAMDGPUDialect::initialize() {
 
 #define GET_OP_CLASSES
 #include "Dialect/TritonAMDGPU/IR/Ops.cpp.inc"
+#include "Dialect/TritonAMDGPU/IR/TritonAMDGPUOpInterfaces.cpp.inc"
 
 namespace mlir::triton::amdgpu {