[ROCm] Add ROCDL Dialect.

include/mlir/Conversion/GPUToROCDL/GPUToROCDLPass.h

@@ -0,0 +1,28 @@
+//===- GPUToNVMMPass.h - Convert GPU kernel to ROCDL 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.
+// =============================================================================
+#ifndef MLIR_CONVERSION_GPUTOROCDL_GPUTOROCDLPASS_H_
+#define MLIR_CONVERSION_GPUTOROCDL_GPUTOROCDLPASS_H_
+
+namespace mlir {
+struct FunctionPassBase;
+
+/// Creates a pass that lowers GPU dialect operations to ROCDL counterparts.
+FunctionPassBase *createLowerGpuOpsToROCDLOpsPass();
+
+} // namespace mlir
+
+#endif // MLIR_CONVERSION_GPUTOROCDL_GPUTOROCDLPASS_H_

include/mlir/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)

include/mlir/LLVMIR/ROCDLDialect.h

@@ -0,0 +1,43 @@
+//===- 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 IR dialect in MLIR, containing ROCDL operations and
+// ROCDL specific extensions to the LLVM type system.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_LLVMIR_ROCDLDIALECT_H_
+#define MLIR_LLVMIR_ROCDLDIALECT_H_
+
+#include "mlir/IR/Dialect.h"
+#include "mlir/IR/OpDefinition.h"
+namespace mlir {
+namespace ROCDL {
+
+///// Ops /////
+#define GET_OP_CLASSES
+#include "mlir/LLVMIR/ROCDLOps.h.inc"
+
+class ROCDLDialect : public Dialect {
+public:
+  explicit ROCDLDialect(MLIRContext *context);
+};
+
+} // namespace ROCDL
+} // namespace mlir
+
+#endif /* MLIR_LLVMIR_ROCDLDIALECT_H_ */

include/mlir/LLVMIR/ROCDLOps.td

@@ -0,0 +1,70 @@
+//===-- 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/LLVMIR/LLVMOpBase.td"
+
+def ROCDL_Dialect : Dialect {
+  let name = "rocdl";
+  let cppNamespace = "ROCDL";
+}
+
+class ROCDL_Op<string mnemonic, list<OpTrait> traits = []> :
+  LLVM_OpBase<ROCDL_Dialect, mnemonic, traits> {
+}
+
+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()); }];
+}
+
+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">;
+
+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()); }];
+}
+
+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

include/mlir/Target/ROCDLIR.h

@@ -0,0 +1,44 @@
+//===- 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

lib/Conversion/CMakeLists.txt

@@ -2,6 +2,7 @@ add_subdirectory(LoopsToGPU)
 add_subdirectory(ControlFlowToCFG)
 add_subdirectory(GPUToCUDA)
 add_subdirectory(GPUToNVVM)
+add_subdirectory(GPUToROCDL)
 add_subdirectory(GPUToSPIRV)
 add_subdirectory(StandardToLLVM)
 add_subdirectory(StandardToSPIRV)

lib/Conversion/GPUToROCDL/CMakeLists.txt

@@ -0,0 +1,10 @@
+add_llvm_library(MLIRGPUtoROCDLTransforms
+  LowerGpuOpsToROCDLOps.cpp
+  )
+target_link_libraries(MLIRGPUtoROCDLTransforms
+  LLVMSupport
+  MLIRGPU
+  MLIRLLVMIR
+  MLIRROCDLIR
+  MLIRPass
+  )

lib/Conversion/GPUToROCDL/LowerGpuOpsToROCDLOps.cpp

@@ -0,0 +1,142 @@
+//===- LowerGpuOpsToROCDLOps.cpp - MLIR GPU to ROCDL lowering passes --------===//
+//
+// 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 implements a pass to generate ROCDLIR operations for higher-level
+// GPU operations.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/GPU/GPUDialect.h"
+#include "mlir/IR/Builders.h"
+#include "mlir/IR/StandardTypes.h"
+#include "mlir/LLVMIR/LLVMDialect.h"
+#include "mlir/LLVMIR/ROCDLDialect.h"
+#include "mlir/Pass/Pass.h"
+#include "mlir/Pass/PassRegistry.h"
+
+#include "llvm/ADT/StringSwitch.h"
+
+namespace mlir {
+namespace {
+
+// A pass that replaces all occurences of GPU operations with their
+// corresponding ROCDL equivalent.
+//
+// This pass does not handle launching of kernels. Instead, it is meant to be
+// used on the body region of a launch or the body region of a kernel
+// function.
+class LowerGpuOpsToROCDLOpsPass : public FunctionPass<LowerGpuOpsToROCDLOpsPass> {
+private:
+  enum dimension { X = 0, Y = 1, Z = 2, invalid };
+
+  template <typename T> dimension dimensionToIndex(T op) {
+    return llvm::StringSwitch<dimension>(op.dimension())
+        .Case("x", X)
+        .Case("y", Y)
+        .Case("z", Z)
+        .Default(invalid);
+  }
+
+  // Helper that replaces Op with XOp, YOp, or ZOp dependeing on the dimension
+  // that Op operates on.  Op is assumed to return an `std.index` value and
+  // XOp, YOp and ZOp are assumed to return an `llvm.i32` value.  Depending on
+  // `indexBitwidth`, sign-extend or truncate the resulting value to match the
+  // bitwidth expected by the consumers of the value.
+  template <typename XOp, typename YOp, typename ZOp, class Op>
+  void replaceWithIntrinsicOrDeviceFunction(Op operation, LLVM::LLVMDialect *dialect,
+                            unsigned indexBitwidth) {
+    assert(operation.getType().isIndex() &&
+           "expected an operation returning index");
+    OpBuilder builder(operation);
+    auto loc = operation.getLoc();
+    Value *newOp;
+    switch (dimensionToIndex(operation)) {
+    case X:
+      newOp = builder.create<XOp>(loc, LLVM::LLVMType::getInt32Ty(dialect));
+      break;
+    case Y:
+      newOp = builder.create<YOp>(loc, LLVM::LLVMType::getInt32Ty(dialect));
+      break;
+    case Z:
+      newOp = builder.create<ZOp>(loc, LLVM::LLVMType::getInt32Ty(dialect));
+      break;
+    default:
+      operation.emitError("Illegal dimension: " + operation.dimension());
+      signalPassFailure();
+      return;
+    }
+
+    if (indexBitwidth > 32) {
+      newOp = builder.create<LLVM::SExtOp>(
+          loc, LLVM::LLVMType::getIntNTy(dialect, indexBitwidth), newOp);
+    } else if (indexBitwidth < 32) {
+      newOp = builder.create<LLVM::TruncOp>(
+          loc, LLVM::LLVMType::getIntNTy(dialect, indexBitwidth), newOp);
+    }
+    operation.replaceAllUsesWith(newOp);
+    operation.erase();
+  }
+
+public:
+  void runOnFunction() {
+    LLVM::LLVMDialect *llvmDialect =
+        getContext().getRegisteredDialect<LLVM::LLVMDialect>();
+    unsigned indexBitwidth =
+        llvmDialect->getLLVMModule().getDataLayout().getPointerSizeInBits();
+    getFunction().walk([&](Operation *opInst) {
+      if (auto threadId = dyn_cast<gpu::ThreadId>(opInst)) {
+        replaceWithIntrinsicOrDeviceFunction<ROCDL::ThreadIdXOp, ROCDL::ThreadIdYOp,
+                             ROCDL::ThreadIdZOp>(threadId, llvmDialect,
+                                                indexBitwidth);
+        return;
+      }
+      if (auto blockId = dyn_cast<gpu::BlockId>(opInst)) {
+        replaceWithIntrinsicOrDeviceFunction<ROCDL::BlockIdXOp, ROCDL::BlockIdYOp,
+                             ROCDL::BlockIdZOp>(blockId, llvmDialect,
+                                               indexBitwidth);
+        return;
+      }
+
+      // BlockDimX/Y/Z are implemented as device functions on ROCDL.
+      if (auto blockDim = dyn_cast<gpu::BlockDim>(opInst)) {
+        replaceWithIntrinsicOrDeviceFunction<
+            ROCDL::BlockDimXOp, ROCDL::BlockDimYOp, ROCDL::BlockDimZOp>(
+            blockDim, llvmDialect, indexBitwidth);
+        return;
+      }
+      // GridDimX/Y/Z are implemented as device functions on ROCDL.
+      if (auto gridDim = dyn_cast<gpu::GridDim>(opInst)) {
+        replaceWithIntrinsicOrDeviceFunction<
+            ROCDL::GridDimXOp, ROCDL::GridDimYOp, ROCDL::GridDimZOp>(
+            gridDim, llvmDialect, indexBitwidth);
+        return;
+      }
+    });
+  }
+};
+
+} // anonymous namespace
+
+FunctionPassBase *createLowerGpuOpsToROCDLOpsPass() {
+  return new LowerGpuOpsToROCDLOpsPass();
+}
+
+static PassRegistration<LowerGpuOpsToROCDLOpsPass>
+    pass("lower-gpu-ops-to-rocdl-ops",
+         "Generate ROCDL operations for gpu operations");
+
+} // namespace mlir

lib/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/LLVMIR
+  )
+add_dependencies(MLIRROCDLIR MLIRROCDLOpsIncGen MLIRROCDLConversionsIncGen LLVMAsmParser LLVMCore LLVMSupport)
+target_link_libraries(MLIRROCDLIR LLVMAsmParser LLVMCore LLVMSupport)