From 5fc340d77bd4f07f4a99f43d9354a9ead62ae1e5 Mon Sep 17 00:00:00 2001
From: Ben Vanik <ben.vanik@gmail.com>
Date: Tue, 29 Oct 2024 22:06:53 -0700
Subject: [PATCH] Registering the ROCDL dialect in init_mlir_dialects. (#18944)

We're emitting those ops now (directly or indirectly) and in order to
parse IR that contains them (from e.g. resuming
--compile-to=executable-targets) we must have all dialects registered.
---
 compiler/src/iree/compiler/Tools/BUILD.bazel          | 1 +
 compiler/src/iree/compiler/Tools/CMakeLists.txt       | 2 ++
 compiler/src/iree/compiler/Tools/init_mlir_dialects.h | 2 ++
 3 files changed, 5 insertions(+)

diff --git a/compiler/src/iree/compiler/Tools/BUILD.bazel b/compiler/src/iree/compiler/Tools/BUILD.bazel
index 7c717d641ebf..7a813c2a51cd 100644
--- a/compiler/src/iree/compiler/Tools/BUILD.bazel
+++ b/compiler/src/iree/compiler/Tools/BUILD.bazel
@@ -114,6 +114,7 @@ iree_compiler_cc_library(
         "@llvm-project//mlir:MathDialect",
         "@llvm-project//mlir:MemRefDialect",
         "@llvm-project//mlir:QuantOps",
+        "@llvm-project//mlir:ROCDLDialect",
         "@llvm-project//mlir:SCFDialect",
         "@llvm-project//mlir:SCFToGPU",
         "@llvm-project//mlir:SCFTransforms",
diff --git a/compiler/src/iree/compiler/Tools/CMakeLists.txt b/compiler/src/iree/compiler/Tools/CMakeLists.txt
index d9033ba99cee..8ad934245c3f 100644
--- a/compiler/src/iree/compiler/Tools/CMakeLists.txt
+++ b/compiler/src/iree/compiler/Tools/CMakeLists.txt
@@ -5,6 +5,7 @@
 # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
 # Doesn't use bazel_to_cmake because of various special logic throughout.
+# That there's various special logic throughout is _bad_. Don't replicate this.
 
 # Enable compiler targets based on options.
 set(IREE_COMPILER_TARGETS "")
@@ -95,6 +96,7 @@ iree_cc_library(
     MLIRLinalgTransforms
     MLIRMLProgramDialect
     MLIRQuantDialect
+    MLIRROCDLDialect
     MLIRSCFDialect
     MLIRSCFToGPU
     MLIRSCFTransforms
diff --git a/compiler/src/iree/compiler/Tools/init_mlir_dialects.h b/compiler/src/iree/compiler/Tools/init_mlir_dialects.h
index e399e63dffe8..ce2b67134989 100644
--- a/compiler/src/iree/compiler/Tools/init_mlir_dialects.h
+++ b/compiler/src/iree/compiler/Tools/init_mlir_dialects.h
@@ -29,6 +29,7 @@
 #include "mlir/Dialect/GPU/IR/GPUDialect.h"
 #include "mlir/Dialect/GPU/TransformOps/GPUTransformOps.h"
 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
+#include "mlir/Dialect/LLVMIR/ROCDLDialect.h"
 #include "mlir/Dialect/LLVMIR/Transforms/InlinerInterfaceImpl.h"
 #include "mlir/Dialect/Linalg/IR/Linalg.h"
 #include "mlir/Dialect/Linalg/TransformOps/DialectExtension.h"
@@ -82,6 +83,7 @@ inline void registerMlirDialects(DialectRegistry &registry) {
                   pdl_interp::PDLInterpDialect,
                   scf::SCFDialect,
                   quant::QuantDialect,
+                  ROCDL::ROCDLDialect,
                   spirv::SPIRVDialect,
                   arm_neon::ArmNeonDialect,
                   arm_sve::ArmSVEDialect,