[ROCm] Adding ROCDL Dialect.

This commit introduces the ROCDL Dialect (i.e. the ROCDL ops + the code to lower those ROCDL ops to LLWM intrinsics/functions). Think of ROCDL Dialect as analogous to the NVVM Dialect, but for AMD GPUs. This patch contains just the essentials needed to get a simple example up and running. We expect to make further additions to the ROCDL Dialect.

This is the first of 3 commits, the follow-up will be:
 * add a pass that lowers GPU Dialect to ROCDL Dialect
 * add a "mlir-rocm-runner" utility

Closes tensorflow/mlir#146

COPYBARA_INTEGRATE_REVIEW=tensorflow/mlir#146 from deven-amd:deven-rocdl-dialect e78e8005c75a78912631116c78dc844fcc4b0de9
PiperOrigin-RevId: 271511259

mlir/include/mlir/Dialect/LLVMIR/CMakeLists.txt

@@ -8,9 +8,16 @@ set(LLVM_TARGET_DEFINITIONS NVVMOps.td)
 mlir_tablegen(NVVMOps.h.inc -gen-op-decls)
 mlir_tablegen(NVVMOps.cpp.inc -gen-op-defs)
 add_public_tablegen_target(MLIRNVVMOpsIncGen)
+set(LLVM_TARGET_DEFINITIONS ROCDLOps.td)
+mlir_tablegen(ROCDLOps.h.inc -gen-op-decls)
+mlir_tablegen(ROCDLOps.cpp.inc -gen-op-defs)
+add_public_tablegen_target(MLIRROCDLOpsIncGen)
 set(LLVM_TARGET_DEFINITIONS LLVMOps.td)
 mlir_tablegen(LLVMConversions.inc -gen-llvmir-conversions)
 add_public_tablegen_target(MLIRLLVMConversionsIncGen)
 set(LLVM_TARGET_DEFINITIONS NVVMOps.td)
 mlir_tablegen(NVVMConversions.inc -gen-llvmir-conversions)
 add_public_tablegen_target(MLIRNVVMConversionsIncGen)
+set(LLVM_TARGET_DEFINITIONS ROCDLOps.td)
+mlir_tablegen(ROCDLConversions.inc -gen-llvmir-conversions)
+add_public_tablegen_target(MLIRROCDLConversionsIncGen)

mlir/include/mlir/Dialect/LLVMIR/ROCDLDialect.h

@@ -0,0 +1,54 @@
+//===- ROCDLDialect.h - MLIR ROCDL IR dialect -------------------*- C++ -*-===//
+//
+// Copyright 2019 The MLIR Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+// =============================================================================
+//
+// This file defines the ROCDL dialect in MLIR, containing ROCDL operations
+// and ROCDL specific extensions to the LLVM type system.
+//
+// Unfortunately there does not exists a formal definition of ROCDL IR that be
+// pointed to here. However the following links contain more information about
+// ROCDL (ROCm-Device-Library)
+//
+// https://github.com/RadeonOpenCompute/ROCm-Device-Libs/blob/master/doc/OCML.md
+// https://github.com/RadeonOpenCompute/ROCm-Device-Libs/blob/master/doc/OCKL.md
+// https://llvm.org/docs/AMDGPUUsage.html
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_DIALECT_LLVMIR_ROCDLDIALECT_H_
+#define MLIR_DIALECT_LLVMIR_ROCDLDIALECT_H_
+
+#include "mlir/IR/Dialect.h"
+#include "mlir/IR/OpDefinition.h"
+
+namespace mlir {
+namespace ROCDL {
+
+///// Ops /////
+#define GET_OP_CLASSES
+#include "mlir/Dialect/LLVMIR/ROCDLOps.h.inc"
+
+class ROCDLDialect : public Dialect {
+public:
+  explicit ROCDLDialect(MLIRContext *context);
+
+  static StringRef getDialectNamespace() { return "rocdl"; }
+};
+
+} // namespace ROCDL
+} // namespace mlir
+
+#endif /* MLIR_DIALECT_LLVMIR_ROCDLDIALECT_H_ */

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

@@ -0,0 +1,102 @@
+//===-- ROCDLOps.td - ROCDL IR dialect op definition file --*- tablegen -*-===//
+//
+// Copyright 2019 The MLIR Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+// =============================================================================
+//
+// This is the ROCDL IR operation definition file.
+//
+//===----------------------------------------------------------------------===//
+
+#ifdef ROCDLIR_OPS
+#else
+#define ROCDLIR_OPS
+
+include "mlir/Dialect/LLVMIR/LLVMOpBase.td"
+
+//===----------------------------------------------------------------------===//
+// ROCDL dialect definitions
+//===----------------------------------------------------------------------===//
+
+def ROCDL_Dialect : Dialect {
+  let name = "rocdl";
+  let cppNamespace = "ROCDL";
+}
+
+//===----------------------------------------------------------------------===//
+// ROCDL op definitions
+//===----------------------------------------------------------------------===//
+
+class ROCDL_Op<string mnemonic, list<OpTrait> traits = []> :
+  LLVM_OpBase<ROCDL_Dialect, mnemonic, traits> {
+}
+
+//===----------------------------------------------------------------------===//
+// ROCDL special register op definitions
+//===----------------------------------------------------------------------===//
+
+class ROCDL_SpecialRegisterOp<string mnemonic,
+    list<OpTrait> traits = []> :
+  ROCDL_Op<mnemonic, !listconcat(traits, [NoSideEffect])>,
+  Results<(outs LLVM_Type:$res)>, Arguments<(ins)> {
+  string llvmBuilder = "$res = createIntrinsicCall(builder,"
+    # "llvm::Intrinsic::amdgcn_" # !subst(".","_", mnemonic) # ");";
+  let parser = [{ return parseROCDLOp(parser, result); }];
+  let printer = [{ printROCDLOp(p, this->getOperation()); }];
+}
+
+class ROCDL_DeviceFunctionOp<string mnemonic, string device_function,
+                             int parameter, list<OpTrait> traits = []> :
+  ROCDL_Op<mnemonic, !listconcat(traits, [NoSideEffect])>,
+  Results<(outs LLVM_Type:$res)>, Arguments<(ins)> {
+  string llvmBuilder = "$res = createDeviceFunctionCall(builder, \""
+  # device_function # "\", " # parameter # ");";
+  let parser = [{ return parseROCDLOp(parser, result); }];
+  let printer = [{ printROCDLOp(p, this->getOperation()); }];
+}
+
+//===----------------------------------------------------------------------===//
+// Thread index and Block index
+
+def ROCDL_ThreadIdXOp : ROCDL_SpecialRegisterOp<"workitem.id.x">;
+def ROCDL_ThreadIdYOp : ROCDL_SpecialRegisterOp<"workitem.id.y">;
+def ROCDL_ThreadIdZOp : ROCDL_SpecialRegisterOp<"workitem.id.z">;
+
+def ROCDL_BlockIdXOp : ROCDL_SpecialRegisterOp<"workgroup.id.x">;
+def ROCDL_BlockIdYOp : ROCDL_SpecialRegisterOp<"workgroup.id.y">;
+def ROCDL_BlockIdZOp : ROCDL_SpecialRegisterOp<"workgroup.id.z">;
+
+//===----------------------------------------------------------------------===//
+// Thread range and Block range
+
+def ROCDL_BlockDimXOp : ROCDL_DeviceFunctionOp<"workgroup.dim.x",
+                                               "__ockl_get_local_size", 0>;
+
+def ROCDL_BlockDimYOp : ROCDL_DeviceFunctionOp<"workgroup.dim.y",
+                                               "__ockl_get_local_size", 1>;
+
+def ROCDL_BlockDimZOp : ROCDL_DeviceFunctionOp<"workgroup.dim.z",
+                                               "__ockl_get_local_size", 2>;
+
+def ROCDL_GridDimXOp : ROCDL_DeviceFunctionOp<"grid.dim.x",
+                                               "__ockl_get_global_size", 0>;
+
+def ROCDL_GridDimYOp : ROCDL_DeviceFunctionOp<"grid.dim.y",
+                                               "__ockl_get_global_size", 1>;
+
+def ROCDL_GridDimZOp : ROCDL_DeviceFunctionOp<"grid.dim.z",
+                                               "__ockl_get_global_size", 2>;
+
+
+#endif // ROCDLIR_OPS

mlir/include/mlir/Target/ROCDLIR.h

@@ -0,0 +1,45 @@
+//===- ROCDLIR.h - MLIR to LLVM + ROCDL IR conversion -----------*- C++ -*-===//
+//
+// Copyright 2019 The MLIR Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+// =============================================================================
+//
+// This file declares the entry point for the MLIR to LLVM + ROCDL IR
+// conversion.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_TARGET_ROCDLIR_H
+#define MLIR_TARGET_ROCDLIR_H
+
+#include <memory>
+
+// Forward-declare LLVM classses.
+namespace llvm {
+class Module;
+} // namespace llvm
+
+namespace mlir {
+class ModuleOp;
+
+/// Convert the given MLIR module into ROCDL IR. This conversion requires the
+/// registration of the LLVM IR dialect and will extract the LLVM context
+/// from the registered LLVM IR dialect.  In case of error, report it
+/// to the error handler registered with the MLIR context, if any (obtained from
+/// the MLIR module), and return `nullptr`.
+std::unique_ptr<llvm::Module> translateModuleToROCDLIR(ModuleOp m);
+
+} // namespace mlir
+
+#endif // MLIR_TARGET_ROCDLIR_H

mlir/lib/Dialect/LLVMIR/CMakeLists.txt

@@ -15,3 +15,12 @@ add_llvm_library(MLIRNVVMIR
   )
 add_dependencies(MLIRNVVMIR MLIRNVVMOpsIncGen MLIRNVVMConversionsIncGen LLVMAsmParser LLVMCore LLVMSupport)
 target_link_libraries(MLIRNVVMIR LLVMAsmParser LLVMCore LLVMSupport)
+
+add_llvm_library(MLIRROCDLIR
+  IR/ROCDLDialect.cpp
+
+  ADDITIONAL_HEADER_DIRS
+  ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/LLVMIR
+  )
+add_dependencies(MLIRROCDLIR MLIRROCDLOpsIncGen MLIRROCDLConversionsIncGen LLVMAsmParser LLVMCore LLVMSupport)
+target_link_libraries(MLIRROCDLIR LLVMAsmParser LLVMCore LLVMSupport)

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

@@ -0,0 +1,82 @@
+//===- ROCDLDialect.cpp - ROCDL IR Ops and Dialect registration -----------===//
+//
+// Copyright 2019 The MLIR Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+// =============================================================================
+//
+// This file defines the types and operation details for the ROCDL IR dialect in
+// MLIR, and the LLVM IR dialect.  It also registers the dialect.
+//
+// The ROCDL dialect only contains GPU specific additions on top of the general
+// LLVM dialect.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/LLVMIR/ROCDLDialect.h"
+
+#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
+#include "mlir/IR/Builders.h"
+#include "mlir/IR/MLIRContext.h"
+#include "mlir/IR/Operation.h"
+#include "mlir/IR/StandardTypes.h"
+#include "llvm/AsmParser/Parser.h"
+#include "llvm/IR/Attributes.h"
+#include "llvm/IR/Function.h"
+#include "llvm/IR/Type.h"
+#include "llvm/Support/SourceMgr.h"
+
+namespace mlir {
+namespace ROCDL {
+
+//===----------------------------------------------------------------------===//
+// Printing/parsing for ROCDL ops
+//===----------------------------------------------------------------------===//
+
+static void printROCDLOp(OpAsmPrinter &p, Operation *op) {
+  p << op->getName() << " ";
+  p.printOperands(op->getOperands());
+  if (op->getNumResults() > 0)
+    interleaveComma(op->getResultTypes(), p << " : ");
+}
+
+// <operation> ::= `rocdl.XYZ` : type
+static ParseResult parseROCDLOp(OpAsmParser &parser, OperationState &result) {
+  Type type;
+  return failure(parser.parseOptionalAttributeDict(result.attributes) ||
+                 parser.parseColonType(type) ||
+                 parser.addTypeToList(type, result.types));
+}
+
+//===----------------------------------------------------------------------===//
+// 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"
+      >();
+
+  // Support unknown operations because not all ROCDL operations are registered.
+  allowUnknownOperations();
+}
+
+#define GET_OP_CLASSES
+#include "mlir/Dialect/LLVMIR/ROCDLOps.cpp.inc"
+
+static DialectRegistration<ROCDLDialect> rocdlDialect;
+
+} // namespace ROCDL
+} // namespace mlir

mlir/lib/Target/CMakeLists.txt

@@ -28,3 +28,17 @@ target_link_libraries(MLIRTargetNVVMIR
   MLIRNVVMIR
   MLIRTargetLLVMIRModuleTranslation
   )
+add_llvm_library(MLIRTargetROCDLIR
+  LLVMIR/ConvertToROCDLIR.cpp
+
+  ADDITIONAL_HEADER_DIRS
+  ${MLIR_MAIN_INCLUDE_DIR}/mlir/Target/LLVMIR
+  DEPENDS
+  intrinsics_gen
+  )
+target_link_libraries(MLIRTargetROCDLIR
+  MLIRGPU
+  MLIRIR
+  MLIRROCDLIR
+  MLIRTargetLLVMIRModuleTranslation
+  )