From 242e6634ff14862789a2d3c1c0ba26817833c065 Mon Sep 17 00:00:00 2001
From: Naums Mogers <naumsmogers@google.com>
Date: Wed, 30 Oct 2024 13:22:56 -0700
Subject: [PATCH] [Mosaic] Add the core type enum

The new attribute allows differentiating compilation by target core.

PiperOrigin-RevId: 691531726
---
 jaxlib/mosaic/dialect/tpu/tpu.td         | 18 ++++++++++++++++++
 jaxlib/mosaic/dialect/tpu/tpu_dialect.cc | 13 +++++++++++++
 2 files changed, 31 insertions(+)

diff --git a/jaxlib/mosaic/dialect/tpu/tpu.td b/jaxlib/mosaic/dialect/tpu/tpu.td
index 783101e839b1..c05b22c5aa88 100644
--- a/jaxlib/mosaic/dialect/tpu/tpu.td
+++ b/jaxlib/mosaic/dialect/tpu/tpu.td
@@ -31,6 +31,11 @@ def TPU_Dialect : Dialect {
   let cppNamespace = "::mlir::tpu";
   let useDefaultAttributePrinterParser = 1;
   let useDefaultTypePrinterParser = 1;
+  let extraClassDeclaration = [{
+    static StringRef GetCoreTypeKey() { return "tpu.core_type"; }
+
+    static std::optional<CoreType> GetCoreTypeAttr(Operation *op);
+  }];
 }
 
 class TPU_Attr<string name, string mnemonic_, list<Trait> traits = []>
@@ -46,6 +51,19 @@ class TPU_Type<string name, string mnemonic_, list<Trait> traits = []>
   let mnemonic = mnemonic_;
 }
 
+def TPU_CoreType : I32EnumAttr<"CoreType", "Core type", [
+  I32EnumAttrCase<"kTc", 0, "tc">,
+  I32EnumAttrCase<"kScScalarSubcore", 1, "sc_scalar_subcore">,
+  I32EnumAttrCase<"kScVectorSubcore", 2, "sc_vector_subcore">
+]> {
+    let genSpecializedAttr = 0;
+    let cppNamespace = "::mlir::tpu";
+}
+
+def TPU_CoreTypeEnum : EnumAttr<TPU_Dialect, TPU_CoreType, "core_type"> {
+    let assemblyFormat = "`<` $value `>`";
+}
+
 def TPU_SemaphoreType : TPU_Type<"Semaphore", "semaphore", [MemRefElementTypeInterface]>;
 def TPU_DMASemaphoreType : TPU_Type<"DMASemaphore", "dma_semaphore", [MemRefElementTypeInterface]>;
 def TPU_SomeSemaphoreType : AnyTypeOf<[TPU_SemaphoreType, TPU_DMASemaphoreType]>;
diff --git a/jaxlib/mosaic/dialect/tpu/tpu_dialect.cc b/jaxlib/mosaic/dialect/tpu/tpu_dialect.cc
index df00093fabe6..d884ef197cda 100644
--- a/jaxlib/mosaic/dialect/tpu/tpu_dialect.cc
+++ b/jaxlib/mosaic/dialect/tpu/tpu_dialect.cc
@@ -19,6 +19,7 @@ limitations under the License.
 #include <array>
 #include <cstdint>
 #include <memory>
+#include <optional>
 #include <utility>
 #include <vector>
 
@@ -68,6 +69,18 @@ void TPUDialect::initialize() {
       >();
 }
 
+/* static */ std::optional<CoreType> TPUDialect::GetCoreTypeAttr(
+    Operation *op) {
+  Attribute attr = op->getAttr(GetCoreTypeKey());
+  if (attr == nullptr) {
+    return std::nullopt;
+  }
+  if (!mlir::isa<CoreTypeAttr>(attr)) {
+    return std::nullopt;
+  }
+  return mlir::cast<CoreTypeAttr>(attr).getValue();
+}
+
 void VectorLayoutAttr::print(AsmPrinter &printer) const {
   printer << '<';
   printer << getLayout();