[stablehlo] Implement product of householder reflectors (#15555)
The stablehlo.custom call for product of householder reflectors can be
directly generated via multiple operations. Implemented the non-blocked
variant as a lowering directly to `linalg`. Testing we have a simplified
lit test as the algorithm is not friendly for lit tests. Instead we use a
numerical validation in stablehlo_ops e2e tests.
diff --git a/compiler/plugins/input/StableHLO/stablehlo-iree/Conversion/BUILD.bazel b/compiler/plugins/input/StableHLO/stablehlo-iree/Conversion/BUILD.bazel
index 5da79e9..22fcb46 100644
--- a/compiler/plugins/input/StableHLO/stablehlo-iree/Conversion/BUILD.bazel
+++ b/compiler/plugins/input/StableHLO/stablehlo-iree/Conversion/BUILD.bazel
@@ -72,6 +72,7 @@
"LegalizeToLinalgUtils.cpp",
"LegalizeToLinalgUtils.h",
"MapStableHLOToScalarOp.h",
+ "StableHLOCustomCalls.cpp",
"StableHLOToArith.cpp",
"StableHLOToIREEInputDialects.cpp",
"StableHLOToLinalg.cpp",
diff --git a/compiler/plugins/input/StableHLO/stablehlo-iree/Conversion/CMakeLists.txt b/compiler/plugins/input/StableHLO/stablehlo-iree/Conversion/CMakeLists.txt
index f615ae2..5cf31a3 100644
--- a/compiler/plugins/input/StableHLO/stablehlo-iree/Conversion/CMakeLists.txt
+++ b/compiler/plugins/input/StableHLO/stablehlo-iree/Conversion/CMakeLists.txt
@@ -56,6 +56,7 @@
"LegalizeToLinalgUtils.cpp"
"LegalizeToLinalgUtils.h"
"MapStableHLOToScalarOp.h"
+ "StableHLOCustomCalls.cpp"
"StableHLOToArith.cpp"
"StableHLOToIREEInputDialects.cpp"
"StableHLOToLinalg.cpp"
diff --git a/compiler/plugins/input/StableHLO/stablehlo-iree/Conversion/Passes.cpp b/compiler/plugins/input/StableHLO/stablehlo-iree/Conversion/Passes.cpp
index d735622..34c61da 100644
--- a/compiler/plugins/input/StableHLO/stablehlo-iree/Conversion/Passes.cpp
+++ b/compiler/plugins/input/StableHLO/stablehlo-iree/Conversion/Passes.cpp
@@ -42,6 +42,7 @@
passManager.addNestedPass<func::FuncOp>(mlir::createCanonicalizerPass());
passManager.addNestedPass<func::FuncOp>(createStableHLOCanonicalize());
passManager.addNestedPass<func::FuncOp>(mlir::createCSEPass());
+ passManager.addNestedPass<func::FuncOp>(createLegalizeStableHLOCustomCalls());
passManager.addNestedPass<func::FuncOp>(
stablehlo::createLegalizeControlFlow());
diff --git a/compiler/plugins/input/StableHLO/stablehlo-iree/Conversion/Passes.td b/compiler/plugins/input/StableHLO/stablehlo-iree/Conversion/Passes.td
index 4f19875..02dc578 100644
--- a/compiler/plugins/input/StableHLO/stablehlo-iree/Conversion/Passes.td
+++ b/compiler/plugins/input/StableHLO/stablehlo-iree/Conversion/Passes.td
@@ -51,6 +51,11 @@
let summary = "Legalizes from CHLO ops flow to StableHLO and Shape ops";
}
+def LegalizeStableHLOCustomCalls :
+ Pass<"iree-stablehlo-legalize-custom-calls", "func::FuncOp"> {
+ let summary = "Legalizes specialized custom calls to decomposed implementations";
+}
+
def LegalizeShapeComputations :
Pass<"iree-stablehlo-legalize-shape-computations", "func::FuncOp"> {
let summary = "Legalizes StableHLO shape operations to core-mlir operations";
diff --git a/compiler/plugins/input/StableHLO/stablehlo-iree/Conversion/Preprocessing/StableHLOToStableHLO.cpp b/compiler/plugins/input/StableHLO/stablehlo-iree/Conversion/Preprocessing/StableHLOToStableHLO.cpp
index fff4726..7fb89d1 100644
--- a/compiler/plugins/input/StableHLO/stablehlo-iree/Conversion/Preprocessing/StableHLOToStableHLO.cpp
+++ b/compiler/plugins/input/StableHLO/stablehlo-iree/Conversion/Preprocessing/StableHLOToStableHLO.cpp
@@ -1532,6 +1532,7 @@
}
};
+// TODO(suderman): Move this to solve code.
struct CustomCallIsTopK final
: OpRewritePattern<mlir::stablehlo::CustomCallOp> {
using OpRewritePattern::OpRewritePattern;
@@ -1751,6 +1752,7 @@
}
};
+// TODO(suderman): Move this to solve code.
struct ApproxTopK final : OpRewritePattern<mlir::stablehlo::CustomCallOp> {
using OpRewritePattern::OpRewritePattern;
diff --git a/compiler/plugins/input/StableHLO/stablehlo-iree/Conversion/StableHLOCustomCalls.cpp b/compiler/plugins/input/StableHLO/stablehlo-iree/Conversion/StableHLOCustomCalls.cpp
new file mode 100644
index 0000000..3ad03ea
--- /dev/null
+++ b/compiler/plugins/input/StableHLO/stablehlo-iree/Conversion/StableHLOCustomCalls.cpp
@@ -0,0 +1,247 @@
+// Copyright 2023 The IREE Authors
+//
+// Licensed under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+// Implements logic for lowering CHLO ops to StableHLO and Shape dialect ops,
+// taking care of CHLO's broadcasting semantics
+
+#include "llvm/ADT/STLExtras.h"
+#include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/Dialect/Linalg/IR/Linalg.h"
+#include "mlir/Dialect/SCF/IR/SCF.h"
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/IR/BuiltinTypes.h"
+#include "mlir/IR/ImplicitLocOpBuilder.h"
+#include "mlir/IR/TypeUtilities.h"
+#include "mlir/Support/LogicalResult.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+#include "stablehlo-iree/Conversion/Passes.h"
+#include "stablehlo-iree/Conversion/Preprocessing/Rewriters.h"
+#include "stablehlo-iree/Conversion/Rewriters.h"
+#include "stablehlo/dialect/BroadcastUtils.h"
+#include "stablehlo/dialect/StablehloOps.h"
+
+#include <algorithm>
+
+namespace mlir::iree_compiler::stablehlo {
+
+#define GEN_PASS_DEF_LEGALIZESTABLEHLOCUSTOMCALLS
+#include "stablehlo-iree/Conversion/Passes.h.inc"
+
+namespace {
+
+// Computes householder using vector `v` and a `tau` matrice
+// with the k-th element. See householder transformation as below:
+// https://en.wikipedia.org/wiki/Householder_transformation
+static Value computeHouseholder(Value v, Value tau, Value k,
+ ImplicitLocOpBuilder b) {
+ auto vTy = cast<ShapedType>(v.getType());
+
+ SmallVector<int64_t> hShape(vTy.getShape());
+ hShape.push_back(hShape.back());
+
+ auto hTy = RankedTensorType::get(hShape, vTy.getElementType());
+ Value empty = b.create<tensor::EmptyOp>(hShape, vTy.getElementType());
+
+ auto outMap = b.getMultiDimIdentityMap(hShape.size());
+
+ SmallVector<AffineExpr> exprs;
+ for (int i = 0, s = vTy.getRank(); i < s; ++i) {
+ exprs.push_back(b.getAffineDimExpr(i));
+ }
+
+ AffineMap vMap = AffineMap::get(hTy.getRank(), 0, exprs, b.getContext());
+
+ exprs.back() = b.getAffineDimExpr(vTy.getRank());
+ AffineMap vTMap = AffineMap::get(hTy.getRank(), 0, exprs, b.getContext());
+
+ SmallVector<AffineMap> affineMaps = {vMap, vTMap, outMap};
+
+ SmallVector<utils::IteratorType> iterTypes(hShape.size(),
+ utils::IteratorType::parallel);
+
+ Value zero = b.create<arith::ConstantOp>(b.getZeroAttr(vTy.getElementType()));
+ Value one =
+ b.create<arith::ConstantOp>(b.getFloatAttr(vTy.getElementType(), 1.0));
+
+ return b
+ .create<linalg::GenericOp>(
+ hTy, ValueRange{v, v}, empty, affineMaps, iterTypes,
+ [&](OpBuilder &bb, Location loc, ValueRange args) {
+ ImplicitLocOpBuilder b(loc, bb);
+ SmallVector<Value> indices;
+ for (int i = 0, s = hTy.getRank(); i < s; ++i) {
+ indices.push_back(b.create<linalg::IndexOp>(loc, i));
+ }
+
+ SmallVector<Value> tauIndices(indices.begin(), indices.end() - 2);
+ tauIndices.push_back(k);
+ Value t = b.create<tensor::ExtractOp>(tau, tauIndices);
+
+ // Generates the lower triangularization of the matrix with
+ // one values on the diagonal.
+ auto tri = [&](Value v, Value i) {
+ Value eq =
+ b.create<arith::CmpIOp>(arith::CmpIPredicate::eq, i, k);
+ Value lt =
+ b.create<arith::CmpIOp>(arith::CmpIPredicate::ult, i, k);
+ Value sel = b.create<arith::SelectOp>(eq, one, v);
+ return b.create<arith::SelectOp>(lt, zero, sel);
+ };
+
+ Value v = tri(args[0], indices[indices.size() - 2]);
+ Value vT = tri(args[1], indices[indices.size() - 1]);
+
+ Value h = b.create<arith::MulFOp>(v, vT);
+ h = b.create<arith::MulFOp>(h, t);
+
+ Value isDiag = b.create<arith::CmpIOp>(arith::CmpIPredicate::eq,
+ indices[indices.size() - 2],
+ indices[indices.size() - 1]);
+ Value diag = b.create<arith::SelectOp>(isDiag, one, zero);
+ Value sub = b.create<arith::SubFOp>(diag, h);
+
+ b.create<linalg::YieldOp>(sub);
+ })
+ .getResult(0);
+}
+
+// Slices the k-th column of matrix and computes the householder transformation
+// for using the `tau` value.
+static Value computeHouseholderSlice(Value matrix, Value tau, Value k,
+ ImplicitLocOpBuilder b) {
+ auto matrixTy = cast<ShapedType>(matrix.getType());
+ int rank = matrixTy.getRank();
+
+ SmallVector<OpFoldResult> vStrides(rank, b.getIndexAttr(1));
+ SmallVector<int64_t> vShape(matrixTy.getShape());
+ vShape[vShape.size() - 1] = 1;
+
+ SmallVector<OpFoldResult> vOffsets(rank, b.getIndexAttr(0));
+ vOffsets[vOffsets.size() - 1] = k;
+
+ SmallVector<OpFoldResult> vSizes;
+ for (auto v : vShape) {
+ vSizes.push_back(b.getIndexAttr(v));
+ }
+
+ auto sliceTy = RankedTensorType::get(vShape, matrixTy.getElementType());
+ Value v = b.create<tensor::ExtractSliceOp>(sliceTy, matrix, vOffsets, vSizes,
+ vStrides);
+
+ SmallVector<ReassociationIndices> reass;
+ for (int i = 0; i < rank - 2; ++i) {
+ reass.push_back({i});
+ }
+ reass.push_back({rank - 2, rank - 1});
+
+ ArrayRef<int64_t> collapseVShape(vShape.begin(), vShape.end() - 1);
+ auto collapseVTy =
+ RankedTensorType::get(collapseVShape, matrixTy.getElementType());
+ Value collapseV = b.create<tensor::CollapseShapeOp>(collapseVTy, v, reass);
+
+ Value householder = computeHouseholder(collapseV, tau, k, b);
+ return householder;
+}
+
+struct HouseholderReflectorRewriter final
+ : OpRewritePattern<mlir::stablehlo::CustomCallOp> {
+ using OpRewritePattern<mlir::stablehlo::CustomCallOp>::OpRewritePattern;
+ using OpAdaptor = mlir::stablehlo::CustomCallOp::Adaptor;
+
+ LogicalResult matchAndRewrite(mlir::stablehlo::CustomCallOp op,
+ PatternRewriter &rewriter) const final {
+ if (op.getCallTargetName() != "ProductOfElementaryHouseholderReflectors") {
+ return rewriter.notifyMatchFailure(
+ op, "not ProductOfElementaryHouseholderReflectors");
+ }
+
+ ImplicitLocOpBuilder b(op.getLoc(), rewriter);
+ auto matrix = op.getOperand(0);
+ auto tau = op.getOperand(1);
+ auto matrixTy = cast<ShapedType>(matrix.getType());
+ auto tauTy = cast<ShapedType>(tau.getType());
+ auto rank = matrixTy.getRank();
+
+ if (isa<ComplexType>(matrixTy.getElementType())) {
+ return rewriter.notifyMatchFailure(op, "complex types not supported");
+ }
+
+ if (rank < 2) {
+ return rewriter.notifyMatchFailure(op, "requires minimum rank 2 matrix");
+ }
+
+ // Implementation needs to be checked to work with variable dimension
+ // lengths. Should be relatively straightforward.
+ if (!matrixTy.hasStaticShape() || !tauTy.hasStaticShape()) {
+ return rewriter.notifyMatchFailure(op,
+ "not supported for dynamic shapes");
+ }
+
+ Value zero = b.create<arith::ConstantIndexOp>(0);
+ Value one = b.create<arith::ConstantIndexOp>(1);
+ Value k = b.create<arith::ConstantIndexOp>(tauTy.getShape().back());
+ Value householder0 = computeHouseholderSlice(matrix, tau, zero, b);
+ auto scf = b.create<scf::ForOp>(
+ one, k, one, ValueRange{householder0},
+ [&](OpBuilder &bb, Location loc, Value iv, ValueRange args) {
+ ImplicitLocOpBuilder b(loc, bb);
+ Value householder = computeHouseholderSlice(matrix, tau, iv, b);
+
+ std::vector<int64_t> batch(rank - 2);
+ for (int i = 0; i < rank - 2; ++i)
+ batch[i] = i;
+ std::vector<int64_t> lhsContract = {rank - 1};
+ std::vector<int64_t> rhsContract = {rank - 2};
+
+ auto dotNums = mlir::stablehlo::DotDimensionNumbersAttr::get(
+ b.getContext(), batch, batch, lhsContract, rhsContract);
+ Value dot = b.create<mlir::stablehlo::DotGeneralOp>(
+ householder0.getType(), args[0], householder, dotNums, nullptr);
+ b.create<scf::YieldOp>(loc, dot);
+ });
+
+ SmallVector<OpFoldResult> vOffsets(rank, b.getIndexAttr(0));
+ SmallVector<OpFoldResult> vStrides(rank, b.getIndexAttr(1));
+ SmallVector<int64_t> vShape(matrixTy.getShape());
+ SmallVector<OpFoldResult> vSizes;
+ for (auto v : vShape) {
+ vSizes.push_back(b.getIndexAttr(v));
+ }
+
+ auto sliceTy = RankedTensorType::get(vShape, matrixTy.getElementType());
+ Value v = b.create<tensor::ExtractSliceOp>(sliceTy, scf.getResult(0),
+ vOffsets, vSizes, vStrides);
+
+ rewriter.replaceOp(op, v);
+ return success();
+ }
+};
+
+//===----------------------------------------------------------------------===//
+// Pass Definition.
+//===----------------------------------------------------------------------===//
+
+struct LegalizeStableHLOCustomCalls final
+ : impl::LegalizeStableHLOCustomCallsBase<LegalizeStableHLOCustomCalls> {
+ void getDependentDialects(DialectRegistry ®istry) const override {
+ registry.insert<arith::ArithDialect, linalg::LinalgDialect, scf::SCFDialect,
+ mlir::stablehlo::StablehloDialect, tensor::TensorDialect>();
+ }
+
+ void runOnOperation() override {
+ func::FuncOp f = getOperation();
+ MLIRContext *ctx = f.getContext();
+
+ RewritePatternSet patterns(ctx);
+ patterns.add<HouseholderReflectorRewriter>(ctx);
+ if (failed(applyPatternsAndFoldGreedily(f, std::move(patterns)))) {
+ signalPassFailure();
+ }
+ }
+};
+} // namespace
+} // namespace mlir::iree_compiler::stablehlo
diff --git a/compiler/plugins/input/StableHLO/stablehlo-iree/Conversion/test/BUILD.bazel b/compiler/plugins/input/StableHLO/stablehlo-iree/Conversion/test/BUILD.bazel
index c0c0b41..c42d6f8 100644
--- a/compiler/plugins/input/StableHLO/stablehlo-iree/Conversion/test/BUILD.bazel
+++ b/compiler/plugins/input/StableHLO/stablehlo-iree/Conversion/test/BUILD.bazel
@@ -23,6 +23,7 @@
"legalize_chlo_with_broadcast.mlir",
"legalize_control_flow.mlir",
"legalize_shape_computations.mlir",
+ "stablehlo_custom_calls.mlir",
"stablehlo_to_iree_input_dialects.mlir",
"stablehlo_to_linalg_convolution.mlir",
"stablehlo_to_linalg_dot_prod.mlir",
diff --git a/compiler/plugins/input/StableHLO/stablehlo-iree/Conversion/test/CMakeLists.txt b/compiler/plugins/input/StableHLO/stablehlo-iree/Conversion/test/CMakeLists.txt
index 7400707..3faf84a 100644
--- a/compiler/plugins/input/StableHLO/stablehlo-iree/Conversion/test/CMakeLists.txt
+++ b/compiler/plugins/input/StableHLO/stablehlo-iree/Conversion/test/CMakeLists.txt
@@ -21,6 +21,7 @@
"legalize_chlo_with_broadcast.mlir"
"legalize_control_flow.mlir"
"legalize_shape_computations.mlir"
+ "stablehlo_custom_calls.mlir"
"stablehlo_to_iree_input_dialects.mlir"
"stablehlo_to_linalg.mlir"
"stablehlo_to_linalg_convolution.mlir"
diff --git a/compiler/plugins/input/StableHLO/stablehlo-iree/Conversion/test/stablehlo_custom_calls.mlir b/compiler/plugins/input/StableHLO/stablehlo-iree/Conversion/test/stablehlo_custom_calls.mlir
new file mode 100644
index 0000000..b756500
--- /dev/null
+++ b/compiler/plugins/input/StableHLO/stablehlo-iree/Conversion/test/stablehlo_custom_calls.mlir
@@ -0,0 +1,12 @@
+// RUN: iree-opt --split-input-file --iree-stablehlo-legalize-custom-calls \
+// RUN: --cse %s | FileCheck %s
+
+// CHECK-LABEL: @householder
+func.func public @householder(%arg0: tensor<4x3xf32>, %arg1: tensor<2xf32>) -> (tensor<4x3xf32>) {
+ // CHECK: linalg.generic
+ // CHECK: scf.for
+ // CHECK: linalg.generic
+ // CHECK: stablehlo.dot_general
+ %0 = stablehlo.custom_call @ProductOfElementaryHouseholderReflectors(%arg0, %arg1) : (tensor<4x3xf32>, tensor<2xf32>) -> tensor<4x3xf32>
+ return %0 : tensor<4x3xf32>
+}
diff --git a/tests/e2e/stablehlo_ops/BUILD.bazel b/tests/e2e/stablehlo_ops/BUILD.bazel
index 5b769cb..b7f5d4c 100644
--- a/tests/e2e/stablehlo_ops/BUILD.bazel
+++ b/tests/e2e/stablehlo_ops/BUILD.bazel
@@ -51,6 +51,7 @@
"finite.mlir",
"floor.mlir",
"gather.mlir",
+ "householder.mlir",
"iota.mlir",
"log.mlir",
"log_plus_one.mlir",
@@ -123,6 +124,7 @@
"finite.mlir",
"floor.mlir",
"gather.mlir",
+ "householder.mlir",
"iota.mlir",
"log.mlir",
"log_plus_one.mlir",
@@ -198,6 +200,7 @@
"finite.mlir",
"floor.mlir",
"gather.mlir",
+ "householder.mlir",
"iota.mlir",
"log.mlir",
"log_plus_one.mlir",
@@ -276,6 +279,7 @@
"finite.mlir",
"floor.mlir",
"gather.mlir",
+ "householder.mlir",
"iota.mlir",
"log.mlir",
"log_plus_one.mlir",
@@ -369,6 +373,7 @@
"finite.mlir",
"floor.mlir",
"gather.mlir",
+ "householder.mlir",
"iota.mlir",
"log.mlir",
"log_plus_one.mlir",
@@ -457,6 +462,7 @@
"finite.mlir",
"floor.mlir",
"gather.mlir",
+ "householder.mlir",
"iota.mlir",
"log.mlir",
"log_plus_one.mlir",
@@ -537,6 +543,7 @@
"finite.mlir",
"floor.mlir",
"gather.mlir",
+ "householder.mlir",
"iota.mlir",
"log.mlir",
"log_plus_one.mlir",
diff --git a/tests/e2e/stablehlo_ops/CMakeLists.txt b/tests/e2e/stablehlo_ops/CMakeLists.txt
index 714e2a8..1397a9d 100644
--- a/tests/e2e/stablehlo_ops/CMakeLists.txt
+++ b/tests/e2e/stablehlo_ops/CMakeLists.txt
@@ -42,6 +42,7 @@
"finite.mlir"
"floor.mlir"
"gather.mlir"
+ "householder.mlir"
"iota.mlir"
"log.mlir"
"log_plus_one.mlir"
@@ -112,6 +113,7 @@
"finite.mlir"
"floor.mlir"
"gather.mlir"
+ "householder.mlir"
"iota.mlir"
"log.mlir"
"log_plus_one.mlir"
@@ -182,6 +184,7 @@
"finite.mlir"
"floor.mlir"
"gather.mlir"
+ "householder.mlir"
"iota.mlir"
"log.mlir"
"log_plus_one.mlir"
@@ -253,6 +256,7 @@
"finite.mlir"
"floor.mlir"
"gather.mlir"
+ "householder.mlir"
"iota.mlir"
"log.mlir"
"log_plus_one.mlir"
@@ -330,6 +334,7 @@
"finite.mlir"
"floor.mlir"
"gather.mlir"
+ "householder.mlir"
"iota.mlir"
"log.mlir"
"log_plus_one.mlir"
@@ -412,6 +417,7 @@
"finite.mlir"
"floor.mlir"
"gather.mlir"
+ "householder.mlir"
"iota.mlir"
"log.mlir"
"log_plus_one.mlir"
@@ -492,6 +498,7 @@
"finite.mlir"
"floor.mlir"
"gather.mlir"
+ "householder.mlir"
"iota.mlir"
"log.mlir"
"log_plus_one.mlir"
@@ -575,6 +582,7 @@
"finite.mlir"
"floor.mlir"
"gather.mlir"
+ "householder.mlir"
"iota.mlir"
"log.mlir"
"log_plus_one.mlir"
diff --git a/tests/e2e/stablehlo_ops/householder.mlir b/tests/e2e/stablehlo_ops/householder.mlir
new file mode 100644
index 0000000..91c713e
--- /dev/null
+++ b/tests/e2e/stablehlo_ops/householder.mlir
@@ -0,0 +1,15 @@
+func.func public @householder_test() -> () {
+ %m = util.unfoldable_constant dense<[
+ [0.19151945, 0.62210877, 0.43772774],
+ [0.78535858, 0.77997581, 0.27259261],
+ [0.27646426, 0.80187218, 0.95813935],
+ [0.87593263, 0.35781727, 0.50099513]]> : tensor<4x3xf32>
+ %t = util.unfoldable_constant dense<[0.68346294, 0.71270203]> : tensor<2xf32>
+ %result = stablehlo.custom_call @ProductOfElementaryHouseholderReflectors(%m, %t) : (tensor<4x3xf32>, tensor<2xf32>) -> tensor<4x3xf32>
+ check.expect_almost_eq_const(%result, dense<[
+ [ 0.31653708, 0.10644531, 0.32681817],
+ [-0.53676350, 0.37089570, -0.31482652],
+ [-0.18895307, -0.54206765, 0.63208690],
+ [-0.59866750, -0.16177818, 0.08177957]]> : tensor<4x3xf32>): tensor<4x3xf32>
+ return
+}
