refactor: reuse batching interface for LU factorization

Author: avik-pal

src/enzyme_ad/jax/Implementations/StableHLOAutoDiffOpInterfaceImpl.cpp

@@ -31,6 +31,8 @@
 #include "stablehlo/dialect/ChloOps.h"
 #include "stablehlo/dialect/StablehloOps.h"
 
+#include "src/enzyme_ad/jax/Dialect/Dialect.h"
+#include "src/enzyme_ad/jax/Dialect/Ops.h"
 #include "src/enzyme_ad/jax/Implementations/WhileLoopInfo.h"
 #include "src/enzyme_ad/jax/Implementations/XLADerivatives.h"
 #include "src/enzyme_ad/jax/Utils.h"
@@ -1962,6 +1964,31 @@ struct SHLOConstantOpBatchInterface
   }
 };
 
+struct SHLOGetDimensionSizeOpBatchInterface
+    : public BatchOpInterface::ExternalModel<
+          SHLOGetDimensionSizeOpBatchInterface, GetDimensionSizeOp> {
+
+  mlir::LogicalResult createBatch(Operation *src, OpBuilder &builder,
+                                  IRMapping &mapper,
+                                  ArrayRef<int64_t> batchSizes) const {
+    auto getDimSizeOp = cast<GetDimensionSizeOp>(src);
+
+    auto newOp = builder.create<GetDimensionSizeOp>(
+        src->getLoc(), mapper.lookup(getDimSizeOp.getOperand()),
+        cast<IntegerAttr>(getDimSizeOp.getDimensionAttr()).getInt() +
+            batchSizes.size());
+    auto bcastOp = builder.create<BroadcastInDimOp>(
+        src->getLoc(),
+        RankedTensorType::get(
+            batchSizes,
+            cast<RankedTensorType>(newOp->getResult(0).getType())
+                .getElementType()),
+        newOp->getResult(0), builder.getDenseI64ArrayAttr({}));
+    mapper.map(src->getResult(0), bcastOp->getResult(0));
+    return success();
+  }
+};
+
 struct SHLOTransposeOpBatchInterface
     : public BatchOpInterface::ExternalModel<SHLOTransposeOpBatchInterface,
                                              TransposeOp> {
@@ -3918,6 +3945,8 @@ void mlir::enzyme::registerStableHLODialectAutoDiffInterface(
         *context);
 
     ConstantOp::attachInterface<SHLOConstantOpBatchInterface>(*context);
+    GetDimensionSizeOp::attachInterface<SHLOGetDimensionSizeOpBatchInterface>(
+        *context);
     TransposeOp::attachInterface<SHLOTransposeOpBatchInterface>(*context);
     IfOp::attachInterface<SHLOGenericBatchOpInterface<IfOp>>(*context);
     WhileOp::attachInterface<SHLOGenericBatchOpInterface<WhileOp>>(*context);
@@ -3947,6 +3976,10 @@ void mlir::enzyme::registerStableHLODialectAutoDiffInterface(
 
     AddOp::attachInterface<StablehloAddSimplifyMathInterface>(*context);
     SubtractOp::attachInterface<StablehloSubSimplifyMathInterface>(*context);
+
+    // TODO: move into its own file
+    enzymexla::JITCallOp::attachInterface<
+        SHLOGenericBatchOpInterface<enzymexla::JITCallOp>>(*context);
   });
 }
 

src/enzyme_ad/jax/Passes/LinalgUtils.cpp

@@ -38,7 +38,7 @@ mlir::ArrayAttr getSHLOLayout(PatternRewriter &rewriter,
   return rewriter.getArrayAttr(attrs);
 }
 
-std::optional<std::string> lapack_precision_prefix(Type elementType) {
+std::optional<std::string> lapackPrecisionPrefix(Type elementType) {
 
   // single-precision float
   if (elementType.isF32()) {

src/enzyme_ad/jax/Passes/LinalgUtils.h

@@ -17,6 +17,6 @@ mlir::ArrayAttr getSHLOLayout(mlir::PatternRewriter &rewriter,
                               llvm::SmallVector<bool> isColMajorArr,
                               int64_t maxNumDims);
 
-std::optional<std::string> lapack_precision_prefix(mlir::Type elementType);
+std::optional<std::string> lapackPrecisionPrefix(mlir::Type elementType);
 
 #endif // ENZYMEXLA_LINALGUTILS_H

src/enzyme_ad/jax/Passes/LowerEnzymeXLALapack.cpp

@@ -80,7 +80,7 @@ struct GeqrfOpLowering : public OpRewritePattern<enzymexla::GeqrfOp> {
     auto type_llvm_void = LLVM::LLVMVoidType::get(ctx);
 
     std::string fn = "geqrf_";
-    if (auto prefix = lapack_precision_prefix(inputElementType)) {
+    if (auto prefix = lapackPrecisionPrefix(inputElementType)) {
       fn = *prefix + fn;
     } else {
       op->emitOpError() << "Unsupported element type: " << inputElementType;
@@ -351,7 +351,7 @@ struct GeqrtOpLowering : public OpRewritePattern<enzymexla::GeqrtOp> {
     auto type_llvm_void = LLVM::LLVMVoidType::get(ctx);
 
     std::string fn = "geqrt_";
-    if (auto prefix = lapack_precision_prefix(inputElementType)) {
+    if (auto prefix = lapackPrecisionPrefix(inputElementType)) {
       fn = *prefix + fn;
     } else {
       op->emitOpError() << "Unsupported element type: " << inputElementType;
@@ -562,7 +562,7 @@ struct OrgqrOpLowering : public OpRewritePattern<enzymexla::OrgqrOp> {
     auto type_llvm_void = LLVM::LLVMVoidType::get(ctx);
 
     std::string fn = "gqr_";
-    if (auto prefix = lapack_precision_prefix(inputElementType)) {
+    if (auto prefix = lapackPrecisionPrefix(inputElementType)) {
       if (prefix == "s" || prefix == "d")
         fn = *prefix + "or" + fn;
       else
@@ -868,7 +868,7 @@ struct OrmqrOpLowering : public OpRewritePattern<enzymexla::OrmqrOp> {
     auto type_llvm_char = rewriter.getIntegerType(8);
 
     std::string fn = "mqr_";
-    if (auto prefix = lapack_precision_prefix(A_eltype)) {
+    if (auto prefix = lapackPrecisionPrefix(A_eltype)) {
       if (prefix == "s" || prefix == "d")
         fn = *prefix + "or" + fn;
       else
@@ -1136,7 +1136,7 @@ struct GemqrtOpLowering : public OpRewritePattern<enzymexla::GemqrtOp> {
     auto type_llvm_char = rewriter.getIntegerType(8);
 
     std::string fn = "gemqrt_";
-    if (auto prefix = lapack_precision_prefix(C_eltype)) {
+    if (auto prefix = lapackPrecisionPrefix(C_eltype)) {
       fn = *prefix + fn;
     } else {
       op->emitOpError() << "Unsupported element type: " << C_eltype;