[mlir][rocdl] Exposing buffer load/store intrinsic

[jerryyin]

Summary

• Updated ROCDLOps tablegen
• Added parsing and printing function for new intrinsic
• Added unit tests

[jerryyin]

Fix lint error

[whchung]

mlir/lib/Dialect/LLVMIR/IR/ROCDLDialect.cpp

37

@jerryyin It seems this printing function is only used by ROCDL_MubufLoadOp, but no one else.

[jerryyin]

Inline printIntrinsic function

[jerryyin]

mlir/lib/Dialect/LLVMIR/IR/ROCDLDialect.cpp

37

I have inlined it in the tablegen.

[whchung]

Thanks!

[jerryyin]

The CL has landed.

mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td

 def ROCDL_mfma_f32_4x4x1f32 : ROCDL_Mfma_IntrOp<"mfma.f32.4x4x1f32">;
 def ROCDL_mfma_f32_4x4x4f16 : ROCDL_Mfma_IntrOp<"mfma.f32.4x4x4f16">;
 def ROCDL_mfma_i32_32x32x4i8 : ROCDL_Mfma_IntrOp<"mfma.i32.32x32x4i8">;
 def ROCDL_mfma_i32_16x16x4i8 : ROCDL_Mfma_IntrOp<"mfma.i32.16x16x4i8">;
 def ROCDL_mfma_i32_4x4x4i8 : ROCDL_Mfma_IntrOp<"mfma.i32.4x4x4i8">;
 def ROCDL_mfma_i32_32x32x8i8 : ROCDL_Mfma_IntrOp<"mfma.i32.32x32x8i8">;
 def ROCDL_mfma_i32_16x16x16i8 : ROCDL_Mfma_IntrOp<"mfma.i32.16x16x16i8">;
 
+//===---------------------------------------------------------------------===//
+// Vector buffer load/store intrinsics
+
+def ROCDL_MubufLoadOp :
+  ROCDL_Op<"buffer.load">, 
+  Results<(outs LLVM_Type:$res)>,
+  Arguments<(ins LLVM_Type:$rsrc,
+                 LLVM_Type:$vindex,
+                 LLVM_Type:$offset,
+                 LLVM_Type:$glc,
+                 LLVM_Type:$slc)>{
+  string llvmBuilder = [{
+      $res = createIntrinsicCall(builder,
+          llvm::Intrinsic::amdgcn_buffer_load, {$rsrc, $vindex, $offset, $glc, 
+          $slc}, {$_resultType});
+  }];
+  let parser = [{ return parseROCDLMubufLoadOp(parser, result); }];
+  let printer = [{ 
+    Operation *op = this->getOperation();
+    p << op->getName() << " " << op->getOperands()
+      << " : " << op->getResultTypes();
+  }];
+}
+
+def ROCDL_MubufStoreOp :
+  ROCDL_Op<"buffer.store">, 
+  Arguments<(ins LLVM_Type:$vdata,
+                 LLVM_Type:$rsrc,
+                 LLVM_Type:$vindex,
+                 LLVM_Type:$offset,
+                 LLVM_Type:$glc,
+                 LLVM_Type:$slc)>{
+  string llvmBuilder = [{
+    auto vdataType = op.vdata().getType().cast<LLVM::LLVMType>()
+                       .getUnderlyingType();
+    createIntrinsicCall(builder,
+          llvm::Intrinsic::amdgcn_buffer_store, {$vdata, $rsrc, $vindex, 
+          $offset, $glc, $slc}, {vdataType});
+  }];
+  let parser = [{ return parseROCDLMubufStoreOp(parser, result); }];
+  let printer = [{ 
+    Operation *op = this->getOperation();
+    p << op->getName() << " " << op->getOperands()
+      << " : " << vdata().getType();
+  }];
+}
 
 #endif // ROCDLIR_OPS

mlir/lib/Dialect/LLVMIR/IR/ROCDLDialect.cpp

 #include "llvm/IR/Function.h"
 #include "llvm/IR/Type.h"
 #include "llvm/Support/SourceMgr.h"
 
 using namespace mlir;
 using namespace ROCDL;
 
 //===----------------------------------------------------------------------===//
+// Parsing for ROCDL ops
+//===----------------------------------------------------------------------===//
+
+static LLVM::LLVMDialect *getLlvmDialect(OpAsmParser &parser) {
+  return parser.getBuilder()
+      .getContext()
+      ->getRegisteredDialect<LLVM::LLVMDialect>();
+}
+
+// <operation> ::=
+//     `llvm.amdgcn.buffer.load.* %rsrc, %vindex, %offset, %glc, %slc :
+//     result_type`
+static ParseResult parseROCDLMubufLoadOp(OpAsmParser &parser,
+                                         OperationState &result) {
+  SmallVector<OpAsmParser::OperandType, 8> ops;
+  Type type;
+  if (parser.parseOperandList(ops, 5) || parser.parseColonType(type) ||
+      parser.addTypeToList(type, result.types))
+    return failure();
+
+  auto int32Ty = LLVM::LLVMType::getInt32Ty(getLlvmDialect(parser));
+  auto int1Ty = LLVM::LLVMType::getInt1Ty(getLlvmDialect(parser));
+  auto i32x4Ty = LLVM::LLVMType::getVectorTy(int32Ty, 4);
+  return parser.resolveOperands(ops,
+                                {i32x4Ty, int32Ty, int32Ty, int1Ty, int1Ty},
+                                parser.getNameLoc(), result.operands);
+}
+
+// <operation> ::=
+//     `llvm.amdgcn.buffer.store.* %vdata, %rsrc, %vindex, %offset, %glc, %slc :
+//     result_type`
+static ParseResult parseROCDLMubufStoreOp(OpAsmParser &parser,
+                                          OperationState &result) {
+  SmallVector<OpAsmParser::OperandType, 8> ops;
+  Type type;
+  if (parser.parseOperandList(ops, 6) || parser.parseColonType(type))
+    return failure();
+
+  auto int32Ty = LLVM::LLVMType::getInt32Ty(getLlvmDialect(parser));
+  auto int1Ty = LLVM::LLVMType::getInt1Ty(getLlvmDialect(parser));
+  auto i32x4Ty = LLVM::LLVMType::getVectorTy(int32Ty, 4);
+
+  if (parser.resolveOperands(ops,
+                             {type, i32x4Ty, int32Ty, int32Ty, int1Ty, int1Ty},
+                             parser.getNameLoc(), result.operands))
+    return failure();
+  return success();
+}
+
+//===----------------------------------------------------------------------===//
 // ROCDLDialect initialization, type parsing, and registration.
 //===----------------------------------------------------------------------===//
 
 // TODO(herhut): This should be the llvm.rocdl dialect once this is supported.
 ROCDLDialect::ROCDLDialect(MLIRContext *context) : Dialect("rocdl", context) {
   addOperations<
 #define GET_OP_LIST
 #include "mlir/Dialect/LLVMIR/ROCDLOps.cpp.inc"

mlir/lib/Target/LLVMIR/ConvertToROCDLIR.cpp

 #include "llvm/IR/Module.h"
 #include "llvm/Support/ToolOutputFile.h"
 
 using namespace mlir;
 
 // Create a call to llvm intrinsic
 static llvm::Value *createIntrinsicCall(llvm::IRBuilder<> &builder,
                                         llvm::Intrinsic::ID intrinsic,
-                                        ArrayRef<llvm::Value *> args = {}) {
+                                        ArrayRef<llvm::Value *> args = {},
+                                        ArrayRef<llvm::Type *> tys = {}) {
   llvm::Module *module = builder.GetInsertBlock()->getModule();
-  llvm::Function *fn = llvm::Intrinsic::getDeclaration(module, intrinsic);
+  llvm::Function *fn = llvm::Intrinsic::getDeclaration(module, intrinsic, tys);
   return builder.CreateCall(fn, args);
 }
 
 // Create a call to ROCm-Device-Library function
 //   Currently this routine will work only for calling ROCDL functions that
 // take a single int32 argument. It is likely that the interface of this
 // function will change to make it more generic.
 static llvm::Value *createDeviceFunctionCall(llvm::IRBuilder<> &builder,

mlir/test/Dialect/LLVMIR/rocdl.mlir

func @rocdl.xdlops(%arg0 : !llvm.float, %arg1 : !llvm.float,
 
   // CHECK: rocdl.mfma.f32.16x16x8bf16 {{.*}} : (!llvm<"<2 x i16>">, !llvm<"<2 x i16>">, !llvm<"<4 x float>">, !llvm.i32, !llvm.i32, !llvm.i32) -> !llvm<"<4 x float>">
   %r19 = rocdl.mfma.f32.16x16x8bf16 %arg10, %arg10, %arg5, %arg3, %arg3, %arg3 :
                             (!llvm<"<2 x i16>">, !llvm<"<2 x i16>">, !llvm<"<4 x float>">,
                             !llvm.i32, !llvm.i32, !llvm.i32) -> !llvm<"<4 x float>">
 
   llvm.return %r0 : !llvm<"<32 x float>">
 }
+
+llvm.func @rocdl.mubuf(%rsrc : !llvm<"<4 x i32>">, %vindex : !llvm.i32,
+                       %offset : !llvm.i32, %glc : !llvm.i1,
+                       %slc : !llvm.i1, %vdata1 : !llvm<"<1 x float>">,
+                       %vdata2 : !llvm<"<2 x float>">, %vdata4 : !llvm<"<4 x float>">) {
+  // CHECK-LABEL: rocdl.mubuf
+  // CHECK: %{{.*}} = rocdl.buffer.load %{{.*}} %{{.*}} %{{.*}} %{{.*}} %{{.*}} : !llvm<"<1 x float>">
+  %r1 = rocdl.buffer.load %rsrc, %vindex, %offset, %glc, %slc : !llvm<"<1 x float>">
+  // CHECK: %{{.*}} = rocdl.buffer.load %{{.*}} %{{.*}} %{{.*}} %{{.*}} %{{.*}} : !llvm<"<2 x float>">
+  %r2 = rocdl.buffer.load %rsrc, %vindex, %offset, %glc, %slc : !llvm<"<2 x float>">
+  // CHECK: %{{.*}} = rocdl.buffer.load %{{.*}} %{{.*}} %{{.*}} %{{.*}} %{{.*}} : !llvm<"<4 x float>">
+  %r4 = rocdl.buffer.load %rsrc, %vindex, %offset, %glc, %slc : !llvm<"<4 x float>">
+
+  // CHECK: rocdl.buffer.store %{{.*}} %{{.*}} %{{.*}} %{{.*}} %{{.*}} %{{.*}} : !llvm<"<1 x float>">
+  rocdl.buffer.store %vdata1, %rsrc, %vindex, %offset, %glc, %slc : !llvm<"<1 x float>">
+  // CHECK: rocdl.buffer.store %{{.*}} %{{.*}} %{{.*}} %{{.*}} %{{.*}} %{{.*}} : !llvm<"<2 x float>">
+  rocdl.buffer.store %vdata2, %rsrc, %vindex, %offset, %glc, %slc : !llvm<"<2 x float>">
+  // CHECK: rocdl.buffer.store %{{.*}} %{{.*}} %{{.*}} %{{.*}} %{{.*}} %{{.*}} : !llvm<"<4 x float>">
+  rocdl.buffer.store %vdata4, %rsrc, %vindex, %offset, %glc, %slc : !llvm<"<4 x float>">
+
+  llvm.return
+}
+

mlir/test/Target/rocdl.mlir

llvm.func @rocdl.xdlops(%arg0 : !llvm.float, %arg1 : !llvm.float,
 
   // CHECK: call <4 x float> @llvm.amdgcn.mfma.f32.16x16x8bf16(<2 x i16> %{{.*}}, <2 x i16> %{{.*}}, <4 x float> %{{.*}}, i32 %{{.*}}, i32 %{{.*}}, i32 %{{.*}})
   %r19 = rocdl.mfma.f32.16x16x8bf16 %arg10, %arg10, %arg5, %arg3, %arg3, %arg3 :
                             (!llvm<"<2 x i16>">, !llvm<"<2 x i16>">, !llvm<"<4 x float>">,
                             !llvm.i32, !llvm.i32, !llvm.i32) -> !llvm<"<4 x float>">
 
   llvm.return %r0 : !llvm<"<32 x float>">
 }
+
+llvm.func @rocdl.mubuf(%rsrc : !llvm<"<4 x i32>">, %vindex : !llvm.i32,
+                       %offset : !llvm.i32, %glc : !llvm.i1,
+                       %slc : !llvm.i1, %vdata1 : !llvm<"<1 x float>">,
+                       %vdata2 : !llvm<"<2 x float>">, %vdata4 : !llvm<"<4 x float>">) {
+  // CHECK-LABEL: rocdl.mubuf
+  // CHECK: call <1 x float> @llvm.amdgcn.buffer.load.v1f32(<4 x i32> %{{.*}}, i32 %{{.*}}, i32 %{{.*}}, i1 %{{.*}}, i1 %{{.*}})
+  %r1 = rocdl.buffer.load %rsrc, %vindex, %offset, %glc, %slc : !llvm<"<1 x float>">
+  // CHECK: call <2 x float> @llvm.amdgcn.buffer.load.v2f32(<4 x i32> %{{.*}}, i32 %{{.*}}, i32 %{{.*}}, i1 %{{.*}}, i1 %{{.*}})
+  %r2 = rocdl.buffer.load %rsrc, %vindex, %offset, %glc, %slc : !llvm<"<2 x float>">
+  // CHECK: call <4 x float> @llvm.amdgcn.buffer.load.v4f32(<4 x i32> %{{.*}}, i32 %{{.*}}, i32 %{{.*}}, i1 %{{.*}}, i1 %{{.*}})
+  %r4 = rocdl.buffer.load %rsrc, %vindex, %offset, %glc, %slc : !llvm<"<4 x float>">
+
+  // CHECK: call void @llvm.amdgcn.buffer.store.v1f32(<1 x float> %{{.*}}, <4 x i32> %{{.*}}, i32 %{{.*}}, i32 %{{.*}}, i1 %{{.*}}, i1 %{{.*}})
+  rocdl.buffer.store %vdata1, %rsrc, %vindex, %offset, %glc, %slc : !llvm<"<1 x float>">
+  // CHECK: call void @llvm.amdgcn.buffer.store.v2f32(<2 x float> %{{.*}}, <4 x i32> %{{.*}}, i32 %{{.*}}, i32 %{{.*}}, i1 %{{.*}}, i1 %{{.*}})
+  rocdl.buffer.store %vdata2, %rsrc, %vindex, %offset, %glc, %slc : !llvm<"<2 x float>">
+  // CHECK: call void @llvm.amdgcn.buffer.store.v4f32(<4 x float> %{{.*}}, <4 x i32> %{{.*}}, i32 %{{.*}}, i32 %{{.*}}, i1 %{{.*}}, i1 %{{.*}})
+  rocdl.buffer.store %vdata4, %rsrc, %vindex, %offset, %glc, %slc : !llvm<"<4 x float>">
+
+  llvm.return
+}
+