[Mosaic] Internal change.

PiperOrigin-RevId: 831038411

Author: Google-ML-Automation

jaxlib/mlir/_mlir_libs/tpu_ext.cc

@@ -871,4 +871,17 @@ NB_MODULE(_tpu_ext, m) {
                          mlirOperationGetRegion(src, i));
     }
   });
+
+  mlir::python::nanobind_adaptors::mlir_type_subclass(m, "Float8EXMYType",
+                                                      mlirTpuIsAFloat8EXMYType)
+      .def_classmethod(
+          "get",
+          [](nb::object cls, MlirType exmy_type, MlirContext ctx) {
+            return cls(mlirTpuFloat8EXMYTypeGet(ctx, exmy_type));
+          },
+          nb::arg("self"), nb::arg("exmy_type") = nullptr,
+          nb::arg("ctx") = nullptr)
+      .def_property_readonly("underlying_type", [](MlirType self) {
+        return mlirTpuFloat8EXMYTypeGetUnderlyingType(self);
+      });
 }

jaxlib/mosaic/dialect/tpu/integrations/c/tpu_dialect.cc

@@ -418,6 +418,26 @@ MLIR_CAPI_EXPORTED void mlirTpuRegisterMosaicSerdePass() {
   mlir::tpu::registerMosaicSerdePass();
 }
 
+//===----------------------------------------------------------------------===//
+// Type API.
+//===----------------------------------------------------------------------===//
+
+// Float8EXMYType
+//===----------------------------------------------------------------------===//
+
+MlirType mlirTpuFloat8EXMYTypeGetUnderlyingType(MlirType exmy_type) {
+  return wrap(llvm::cast<mlir::tpu::Float8EXMYType>(unwrap(exmy_type))
+                  .getUnderlyingType());
+}
+
+bool mlirTpuIsAFloat8EXMYType(MlirType type) {
+  return llvm::isa<mlir::tpu::Float8EXMYType>(unwrap(type));
+}
+
+MlirType mlirTpuFloat8EXMYTypeGet(MlirContext ctx, MlirType exmy_type) {
+  return wrap(mlir::tpu::Float8EXMYType::get(unwrap(ctx), unwrap(exmy_type)));
+}
+
 #include "mlir/CAPI/Pass.h"  // IWYU pragma: keep
 #include "mlir/CAPI/Support.h"  // IWYU pragma: keep
 

jaxlib/mosaic/dialect/tpu/integrations/c/tpu_dialect.h

@@ -240,6 +240,14 @@ mlirTpuRelayout(MlirTpuInsertionPoint insertion_point, MlirValue val,
 
 MLIR_CAPI_EXPORTED void mlirTpuRegisterMosaicSerdePass();
 
+MLIR_CAPI_EXPORTED MlirType mlirTpuFloat8EXMYTypeGetUnderlyingType(
+    MlirType exmy_type);
+
+MLIR_CAPI_EXPORTED bool mlirTpuIsAFloat8EXMYType(MlirType type);
+
+MLIR_CAPI_EXPORTED MlirType mlirTpuFloat8EXMYTypeGet(
+    MlirContext ctx, MlirType exmy_type);
+
 #ifdef __cplusplus
 }
 #endif

jaxlib/mosaic/dialect/tpu/tpu.td

@@ -48,8 +48,9 @@ class TPU_Attr<string name, string mnemonic_, list<Trait> traits = []>
 // TODO(b/369418606): Find out the way to verify vreg size.
 def TPU_Vreg : Type<IsVectorOfNonZeroRankTypePred, "native-sized vreg", "::mlir::VectorType">;
 
-class TPU_Type<string name, string mnemonic_, list<Trait> traits = []>
-    : TypeDef<TPU_Dialect, name, traits> {
+class TPU_Type<string name, string mnemonic_, list<Trait> traits = [],
+               string baseCppType = "::mlir::Type">
+    : TypeDef<TPU_Dialect, name, traits, baseCppType> {
   let mnemonic = mnemonic_;
 }
 
@@ -82,6 +83,25 @@ def TPU_SemaphoreType : TPU_Type<"Semaphore", "semaphore", [MemRefElementTypeInt
 def TPU_DMASemaphoreType : TPU_Type<"DMASemaphore", "dma_semaphore", [MemRefElementTypeInterface]>;
 def TPU_SomeSemaphoreType : AnyTypeOf<[TPU_SemaphoreType, TPU_DMASemaphoreType]>;
 
+def TPU_Float8EXMYType : TPU_Type<"Float8EXMY", "float8_exmy",
+                            [DeclareTypeInterfaceMethods<FloatTypeInterface, ["getFloatSemantics"]>]> {
+    let summary = "EXMY type in a nearest power-of-2 bitwidth type container";
+    let description = [{
+      EXMY type in a nearest power-of-2 bitwidth type container. Meaningful bits
+      are aligned to LSB, and bits higher than the underlying exmy type in the
+      container are considered as ignored. See https://arxiv.org/abs/2405.13938
+      for more details.
+    }];
+
+    let parameters = (ins
+      TypeParameter<"::mlir::Type", "Underlying EXMY type">:$underlying_type
+    );
+
+    let assemblyFormat = [{
+      `<` $underlying_type `>`
+    }];
+}
+
 def TPU_DimensionSemantics : I32EnumAttr<"DimensionSemantics", "Dimension semantics", [
   I32EnumAttrCase<"parallel", 0>,
   I32EnumAttrCase<"arbitrary", 1>,

jaxlib/mosaic/dialect/tpu/tpu_dialect.cc

@@ -20,6 +20,7 @@ limitations under the License.
 
 #include "absl/hash/hash.h"
 #include "absl/log/log.h"
+#include "llvm/ADT/APFloat.h"
 #include "llvm/ADT/Hashing.h"
 #include "llvm/ADT/TypeSwitch.h"  // IWYU pragma: keep.
 #include "mlir/Dialect/Arith/IR/Arith.h"
@@ -292,4 +293,9 @@ DotDimensionNumbersAttr defaultDimensionNumbers(Builder &builder,
       /*rhs_batch_dims=*/{});
 }
 
+const ::llvm::fltSemantics& Float8EXMYType::getFloatSemantics() const {
+  // TODO(twsung): Fix this.
+  return llvm::APFloat::Float8E8M0FNU();
+}
+
 }  // namespace mlir::tpu