feat: partial lowering for getrf

Author: avik-pal

src/enzyme_ad/jax/Passes/LowerEnzymeXLALapack.cpp

@@ -1283,6 +1283,48 @@ struct GemqrtOpLowering : public OpRewritePattern<enzymexla::GemqrtOp> {
   }
 };
 
+Value anyNonFiniteValue(PatternRewriter &rewriter, Location loc, Type outType,
+                        Value input, int64_t inputRank) {
+  auto areFinite = stablehlo::AndOp::create(
+      rewriter, loc,
+      stablehlo::IsFiniteOp::create(
+          rewriter, loc, stablehlo::RealOp::create(rewriter, loc, input)),
+      stablehlo::IsFiniteOp::create(
+          rewriter, loc, stablehlo::ImagOp::create(rewriter, loc, input)));
+
+  SmallVector<int64_t> reductionDims;
+  for (int i = inputRank - 2; i < inputRank; i++)
+    reductionDims.push_back(i);
+
+  auto initValType = RankedTensorType::get({}, rewriter.getI1Type());
+  auto initVal = stablehlo::ConstantOp::create(
+      rewriter, loc, initValType, cast<ElementsAttr>(makeAttr(initValType, 1)));
+
+  auto allFinite = stablehlo::ReduceOp::create(
+      rewriter, loc, ValueRange{areFinite.getResult()}, ValueRange{initVal},
+      rewriter.getDenseI64ArrayAttr(reductionDims));
+
+  {
+    OpBuilder::InsertionGuard guard(rewriter);
+    auto &region = allFinite.getBody();
+    auto *block = rewriter.createBlock(&region, {}, {initValType, initValType},
+                                       {loc, loc});
+
+    rewriter.setInsertionPointToStart(block);
+    stablehlo::ReturnOp::create(
+        rewriter, loc,
+        ValueRange{stablehlo::AndOp::create(rewriter, loc,
+                                            block->getArgument(0),
+                                            block->getArgument(1))
+                       .getResult()});
+  }
+
+  // 0 is success, 1 is failure
+  return stablehlo::ConvertOp::create(
+      rewriter, loc, outType,
+      stablehlo::NotOp::create(rewriter, loc, allFinite.getResult(0)));
+}
+
 struct GetrfOpLowering : public OpRewritePattern<enzymexla::GetrfOp> {
   std::string backend;
   int64_t blasIntWidth;
@@ -1307,18 +1349,78 @@ struct GetrfOpLowering : public OpRewritePattern<enzymexla::GetrfOp> {
 
 private:
   LogicalResult matchAndRewriteCPU(enzymexla::GetrfOp op,
-                                   PatternRewriter &rewriter) const override {
+                                   PatternRewriter &rewriter) const {
     return failure();
   }
 
   LogicalResult matchAndRewriteCUDA(enzymexla::GetrfOp op,
-                                    PatternRewriter &rewriter) const override {
+                                    PatternRewriter &rewriter) const {
     return failure();
   }
 
   LogicalResult matchAndRewriteTPU(enzymexla::GetrfOp op,
-                                   PatternRewriter &rewriter) const override {
-    return failure();
+                                   PatternRewriter &rewriter) const {
+    auto input = op.getInput();
+
+    auto inputType = cast<RankedTensorType>(input.getType());
+    auto pivotType = cast<RankedTensorType>(op.getResult(1).getType());
+    auto infoType = cast<RankedTensorType>(op.getResult(3).getType());
+
+    auto inputShape = inputType.getShape();
+    auto inputRank = inputType.getRank();
+
+    SmallVector<int64_t> permutationShape(inputShape.begin(),
+                                          inputShape.end() - 2);
+    permutationShape.push_back(inputShape[inputRank - 2]);
+    auto permutationType =
+        RankedTensorType::get(permutationShape, rewriter.getI32Type());
+
+    auto pivotTPUType =
+        RankedTensorType::get(pivotType.getShape(), rewriter.getI32Type());
+
+    // TPU returns (LU, pivots, permutation). info isn't returned. based on
+    // how JAX operates, I am assuming info != 0 when there is a nan in the
+    // output.
+    auto customCall = stablehlo::CustomCallOp::create(
+        rewriter, op.getLoc(),
+        TypeRange{inputType, pivotTPUType, permutationType}, ValueRange{input},
+        rewriter.getStringAttr("LuDecomposition"),
+        /*has_side_effect*/ nullptr,
+        /*backend_config*/ nullptr,
+        /*api_version*/ nullptr,
+        /*calledcomputations*/ nullptr,
+        /*operand_layouts*/ nullptr,
+        /*result_layouts*/ nullptr,
+        /*output_operand_aliases*/ nullptr);
+
+    // LAPACK returns 1-indexed pivots, while XLA returns 0-indexed pivots.
+    // We make it consistent with LAPACK by adding 1 to the pivots.
+    auto pivots1Indexed = stablehlo::AddOp::create(
+        rewriter, op.getLoc(),
+        stablehlo::ConstantOp::create(
+            rewriter, op.getLoc(), pivotType,
+            cast<ElementsAttr>(makeAttr(pivotType, 1))),
+        stablehlo::ConvertOp::create(rewriter, op.getLoc(), pivotType,
+                                     customCall.getResult(1)));
+
+    auto permutation1Indexed = stablehlo::AddOp::create(
+        rewriter, op.getLoc(),
+        stablehlo::ConstantOp::create(
+            rewriter, op.getLoc(), permutationType,
+            cast<ElementsAttr>(makeAttr(permutationType, 1))),
+        stablehlo::ConvertOp::create(rewriter, op.getLoc(), permutationType,
+                                     customCall.getResult(2)));
+
+    auto info = anyNonFiniteValue(rewriter, op.getLoc(), infoType,
+                                  customCall.getResult(0), inputRank);
+
+    rewriter.replaceAllUsesWith(op.getResult(0), customCall.getResult(0));
+    rewriter.replaceAllUsesWith(op.getResult(1), pivots1Indexed);
+    rewriter.replaceAllUsesWith(op.getResult(2), permutation1Indexed);
+    rewriter.replaceAllUsesWith(op.getResult(3), info);
+    rewriter.eraseOp(op);
+
+    return success();
   }
 };