Skip to content

implement blas.hemm op #2581

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
mofeing wants to merge 4 commits into
base: main
Choose a base branch
from
Draft

implement blas.hemm op #2581

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
fe9aa49
Reapply "implement Hemm custom calls for CUDA"
mofeing Jun 23, 2026
411aa75
Reapply "add `HemmOp` op"
mofeing Jun 23, 2026
e402c0e
Reapply "fix compilation of Hemm custom calls"
mofeing Jun 23, 2026
3983ba0
Reapply "implement CPU lowering"
mofeing Jun 23, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
32 changes: 32 additions & 0 deletions src/enzyme_ad/jax/Dialect/EnzymeXLAOps.td
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -628,6 +628,38 @@ def SymmOp : EnzymeXLA_Op<"blas.symm", [Pure, SameOperandsAndResultElementType]>
}];
}

def HemmOp : EnzymeXLA_Op<"blas.hemm", [Pure, SameOperandsAndResultElementType]> {
let summary = "Multiplication involving an Hermitian matrix";

let description = [{
C := alpha*A*B + beta*C, or C := alpha*B*A + beta*C, where alpha and beta are scalars, A is an Hermitian matrix"
}];

let arguments = (ins
TensorFloat:$alpha,
HLO_Tensor:$A,
HLO_Tensor:$B,
TensorFloat:$beta,
HLO_Tensor:$C,
EnzymeXLA_LapackSideAttr:$side,
EnzymeXLA_LapackUploAttr:$uplo
);

let results = (outs
HLO_Tensor: $output
);

let assemblyFormat = [{
$alpha `,` $A `,` $B `,` $beta `,` $C attr-dict `:` functional-type(operands, results)
}];

let hasVerifier = 1;

let builders = [
OpBuilder<(ins "Value":$A, "Value":$B, "enzymexla::LapackSide":$side, "enzymexla::LapackUplo":$uplo)>
];
}

def SyrkOp: EnzymeXLA_Op<"blas.syrk", [Pure, SameOperandsAndResultElementType]> {
let summary = "Multiplication involving a symmetric matrix";

Expand Down
110 changes: 110 additions & 0 deletions src/enzyme_ad/jax/Dialect/Ops.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1262,6 +1262,116 @@ void GemmOp::build(OpBuilder &builder, OperationState &result, Value A, Value B,
builder.getContext(), transb));
}

LogicalResult enzymexla::HemmOp::verify() {
auto alpha = op.getAlpha();
auto a = op.getA();
auto b = op.getB();
auto beta = op.getBeta();
auto c = op.getC();
auto side = op.getSide();

auto type_alpha = cast<RankedTensorType>(alpha.getType());
auto type_A = cast<RankedTensorType>(A.getType());
auto type_B = cast<RankedTensorType>(B.getType());
auto type_beta = cast<RankedTensorType>(beta.getType());
auto type_C = cast<RankedTensorType>(C.getType());

auto shape_A = type_A.getShape();
auto shape_B = type_B.getShape();
auto shape_C = type_C.getShape();

auto type_element = type_alpha.getElementType();
auto rank = type_A.getRank();

auto inner_dim_A =
side == enzymexla::LapackSide::left ? rank - 1 : rank - 2;
auto inner_dim_B =
side == enzymexla::LapackSide::left ? rank - 2 : rank - 1;

auto outer_dim_A =
side == enzymexla::LapackSide::left ? rank - 2 : rank - 1;
auto outer_dim_B =
side == enzymexla::LapackSide::left ? rank - 1 : rank - 2;

if (type_A.getElementType() != type_element ||
type_B.getElementType() != type_element ||
type_C.getElementType() != type_element ||
type_beta.getElementType() != type_element) {
return emitOpError("Element types of alpha, A, B and C must match");
}

if (!isa<mlir::ComplexType>(type_element)) {
return emitOpError("Hemm works only with complex element type")
}

if (type_A.getRank() != type_B.getRank() ||
type_A.getRank() != type_C.getRank()) {
return emitOpError("Ranks of A, B and C must match");
}

if (shape_A.drop_back(2) != shape_B.drop_back(2) ||
shape_A.drop_back(2) != shape_C.drop_back(2)) {
return emitOpError("Batch dimensions of A, B and C must match");
}

if (shape_A[inner_dim_A] != shape_B[inner_dim_B]) {
return emitOpError("Inner dimensions of A and B must match");
}

if (shape_A[outer_dim_A] != shape_C[rank - 2] ||
shape_B[outer_dim_B] != shape_C[rank - 1]) {
return emitOpError(
"Outer dimensions of A and B must match corresponding dimensions of C");
}

if (getResult().getType() != type_C) {
return emitOpError("Result type must match C's type");
}

return success();
}

void HemmOp::build(OpBuilder &builder, OperationState &result, Value A, Value B,
enzymexla::LapackSide side, enzymexla::LapackUplo uplo) {
auto type_A = cast<RankedTensorType>(A.getType());
auto type_B = cast<RankedTensorType>(B.getType());

auto element_type = type_A.getElementType();
auto rank = type_A.getRank();

auto outer_dim_A =
side == enzymexla::LapackTranspose::left ? rank - 2 : rank - 1;
auto outer_dim_B =
uplo == enzymexla::LapackTranspose::right ? rank - 1 : rank - 2;

auto shape_a = type_A.getShape();
auto shape_b = type_B.getShape();
SmallVector<int64_t> shape_c;
for (int i = 0; i < rank - 2; i++) {
shape_c.push_back(shape_a[i]);
}
shape_c.push_back(shape_a[outer_dim_A]);
shape_c.push_back(shape_b[outer_dim_B]);

auto type_scalar = RankedTensorType::get({}, element_type);
auto alpha = stablehlo::ConstantOp::create(
builder, result.location, type_scalar,
cast<ElementsAttr>(makeAttr(type_scalar, 1)));
auto beta = stablehlo::ConstantOp::create(
builder, result.location, type_scalar,
cast<ElementsAttr>(makeAttr(type_scalar, 0)));

auto type_C = RankedTensorType::get(shape_c, element_type);
auto C =
stablehlo::ConstantOp::create(builder, result.location, type_C,
cast<ElementsAttr>(makeAttr(type_C, 0)));

result.addTypes(type_C);
result.addOperands({alpha, A, B, beta, C});
result.addAttribute("side", enzymexla::LapackSideAttr::get(builder.getContext(), side));
result.addAttribute("uplo", enzymexla::LapackUploAttr::get(builder.getContext(), uplo));
}

LogicalResult enzymexla::SyrkOp::verify() {
auto CType = cast<RankedTensorType>(getC().getType());
bool isComplex = false;
Expand Down
167 changes: 166 additions & 1 deletion src/enzyme_ad/jax/Passes/LowerEnzymeXLABLAS.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -527,6 +527,171 @@ struct SymmOpLowering : public OpRewritePattern<enzymexla::SymmOp> {
}
};

struct HemmOpLowering : public OpRewritePattern<enzymexla::HemmOp> {
using OpRewritePattern<enzymexla::HemmOp>::OpRewritePattern;

std::string backend;
int64_t blasIntWidth;

HemmOpLowering(std::string backend, int64_t blasIntWidth,
MLIRContext *context, PatternBenefit benefit = 1)
: OpRewritePattern(context, benefit), backend(backend),
blasIntWidth(blasIntWidth) {}

LogicalResult matchAndRewrite(enzymexla::HemmOp op,
PatternRewriter &rewriter) const override {
auto AType = cast<RankedTensorType>(op.getA().getType());
auto nBatchDims = AType.getRank() - 2;

if (nBatchDims == 0) {
if (backend == "cpu") {
return matchAndRewriteCPU(op, rewriter);
} else if (backend == "cuda") {
return matchAndRewriteCUDA(op, rewriter);
}
}

return matchAndRewriteFallback(op, rewriter);
}

LogicalResult matchAndRewriteCPU(enzymexla::HemmOp op,
PatternRewriter &rewriter) const {
auto ctx = op->getContext();
LLVMTypeConverter typeConverter(ctx);

auto alpha = op.getAlpha();
auto a = op.getA();
auto b = op.getB();
auto beta = op.getBeta();
auto c = op.getC();
auto side = op.getSide().getBlasChar();
auto uplo = op.getUplo().getBlasChar();

auto type_a = cast<RankedTensorType>(a.getType());
auto type_b = cast<RankedTensorType>(b.getType());
auto type_c = cast<RankedTensorType>(c.getType());
auto type_elem = type_a.getElementType();

if (!type_a || !type_b || !type_c)
return rewriter.notifyMatchFailure(
op, "expected ranked tensor types");

if (type_a.getRank() != 2 || type_b.getRank() > 2 || type_c.getRank() > 2)
return rewriter.notifyMatchFailure(op,"only 2D matrices supported for HemmOp");

std::string fn;
if (auto prefix = lapackPrecisionPrefix(elementType)) {
fn = *prefix + "hemm_";
} else {
op->emitOpError() << "Unsupported element type: " << elementType;
return rewriter.notifyMatchFailure(op, "unsupported element type");
}
std::string bind_fn = "enzymexla_blas_bind_" + fn;
std::string wrapper_fn = "enzymexla_blas_wrap_" + fn;

auto type_char = rewriter.getI8Type();
auto type_blas_int = rewriter.getIntegerType(blasIntWidth);
auto type_llvm_char = typeConverter.convertType(type_llvm_char);
auto type_llvm_blas_int = typeConverter.convertType(type_blas_int);
auto type_llvm_ptr = LLVM::LLVMPointerType::get(ctx);
auto type_llvm_void = LLVM::LLVMVoidType::get(ctx);
auto type_llvm_elem = typeConverter.convertType(type_elem);

auto moduleOp = op->getParentOfType<ModuleOp>();

if (!moduleOp.lookupSymbol<LLVM::LLVMFuncOp>(blasFn)) {
OpBuilder::InsertionGuard guard(rewriter);
rewriter.setInsertionPointToStart(moduleOp.getBody());
auto funcType = LLVM::LLVMFunctionType::get(type_llvm_void,
{
type_llvm_char, // side
type_llvm_char, // uplo
type_llvm_blas_int, // m
type_llvm_blas_int, // n
type_llvm_ptr, // alpha
type_llvm_ptr, // A
type_llvm_blas_int, // lda
type_llvm_ptr, // B
type_llvm_blas_int, // ldb
type_llvm_ptr, // beta
type_llvm_ptr, // C
type_llvm_blas_int, // ldc
},
false
);
LLVM::LLVMFuncOp::create(rewriter, op.getLoc(), bind_fn, funcType,
LLVM::Linkage::External);
}

if (!moduleOp.lookupSymbol<LLVM::LLVMFuncOp>(wrapper_fn)) {
OpBuilder::InsertionGuard guard(rewriter);
rewriter.setInsertionPointToStart(moduleOp.getBody());

auto funcType = LLVM::LLVMFunctionType::get(type_llvm_void, {/*side*/type_llvm_ptr, /*uplo*/type_llvm_ptr, /*m*/type_llvm_ptr, /*n*/type_llvm_ptr, /*alpha*/type_llvm_ptr, /*A*/type_llvm_ptr, /*lda*/type_llvm_ptr, /*B*/type_llvm_ptr, /*ldb*/type_llvm_ptr, /*beta*/type_llvm_ptr, /*C*/type_llvm_ptr, /*ldc*/type_llvm_ptr}, false);
auto funcOp = LLVM::LLVMFuncOp::create(rewriter, op.getLoc(), wrapper_fn, funcType, LLVM::Linkage::Private);
rewriter.setInsertionPointToStart(funcOp.addEntryBlock(rewriter));

auto opn_side = LLVM::LoadOp(rewriter, op.getLoc(), type_llvm_char, funcOp.getArgument(0));
auto opn_uplo = LLVM::LoadOp(rewriter, op.getLoc(), type_llvm_char, funcOp.getArgument(1));
auto opn_m = funcOp.getArgument(2);
auto opn_n = funcOp.getArgument(3);
auto opn_alpha = funcOp.getArgument(4);
auto opn_a = funcOp.getArgument(5);
auto opn_lda = LLVM::LoadOp(rewriter, op.getLoc(), type_llvm_blas_int, funcOp.getArgument(6));
auto opn_b = funcOp.getArgument(7);
auto opn_ldb = LLVM::LoadOp(rewriter, op.getLoc(), type_llvm_blas_int, funcOp.getArgument(8));
auto opn_beta = funcOp.getArgument(9);
auto opn_c = LLVM::LoadOp(rewriter, op.getLoc(), type_llvm_blas_int, funcOp.getArgument(10));
auto opn_ldc = LLVM::LoadOp(rewriter, op.getLoc(), type_llvm_blas_int, funcOp.getArgument(11));

LLVM::CallOp::create(rewriter, op.getLoc(), TypeRange{}, SymbolRefAttr::get(ctx, bind_fn), {opn_side, opn_uplo, opn_m, opn_n, opn_alpha, opn_a, opn_lda, opn_b, opn_ldb, opn_beta, opn_c, opn_ldc});
LLVM::ReturnOp::create(rewriter, op.getLoc(), ValueRange{});
}

SmallVector<bool> isColMajorArr(12, true);
SmallVector<int64_t> operandRanks = {0, 0, 0, 0, 0, type_a.getRank(), 0, type_b.getRank(), 0, 0, type_c.getRank(), 0};
SmallVector<int64_t> outputRanks = {type_c.getRank()};
auto operandLayouts = getSHLOLayout(rewriter, operandRanks, isColMajorArr, 2);
auto resultLayouts = getSHLOLayout(rewriter, outputRanks, {true}, 2);

SmallVector<Attribute> aliases = {stablehlo::OutputOperandAliasAttr::get(ctx, {}, 10, {})};

auto side_value = stablehlo::ConstantOp::create(rewriter, op.getLoc(), type_char, cast<ElementsAttr>(makeAttr(type_char, side)));
auto side_value = stablehlo::ConstantOp::create(rewriter, op.getLoc(), type_char, cast<ElementsAttr>(makeAttr(type_char, uplo)));
auto m = stablehlo::ConvertOp::create(rewriter, op.getLoc(), type_blas_int, stablehlo::GetDimensionSizeOp::create(rewriter, op.getLoc(), C, type_c.getRank() - 2));
auto n = stablehlo::ConvertOp::create(rewriter, op.getLoc(), type_blas_int, stablehlo::GetDimensionSizeOp::create(rewriter, op.getLoc(), C, type_c.getRank() - 1));
auto lda = stablehlo::ConvertOp::create(rewriter, op.getLoc(), type_blas_int, stablehlo::GetDimensionSizeOp::create(rewriter, op.getLoc(), A, type_a.getRank() - 2));
auto ldb = stablehlo::ConvertOp::create(rewriter, op.getLoc(), type_blas_int, stablehlo::GetDimensionSizeOp::create(rewriter, op.getLoc(), B, type_b.getRank() - 2));
auto ldc = stablehlo::ConvertOp::create(rewriter, op.getLoc(), type_blas_int, stablehlo::GetDimensionSizeOp::create(rewriter, op.getLoc(), C, type_c.getRank() - 2));

auto jitCall = enzymexla::JITCallOp::create(
rewriter, op.getLoc(), TypeRange{type_c},
mlir::FlatSymbolRefAttr::get(ctx, wrapper_fn),
ValueRange{side_value, uplo_value, m, n, alpha, A, lda, B, ldb, beta, C, ldc},
rewriter.getStringAttr(""),
/*operand_layouts=*/operandLayouts,
/*result_layouts=*/resultLayouts,
/*arg_attrs=*/nullptr,
/*res_attrs=*/nullptr,
/*output_operand_aliases=*/rewriter.getArrayAttr(aliases),
/*xla_side_effect_free=*/rewriter.getUnitAttr());

rewriter.replaceOp(op, jitCall);

return success();
}t

LogicalResult matchAndRewriteCUDA(enzymexla::HemmOp op,
PatternRewriter &rewriter) const {
return success();
}

LogicalResult matchAndRewriteFallback(enzymexla::HemmOp op,
PatternRewriter &rewriter) const {
return success();
}
};

struct SyrkOpLowering : public OpRewritePattern<enzymexla::SyrkOp> {
using OpRewritePattern<enzymexla::SyrkOp>::OpRewritePattern;

Expand Down Expand Up @@ -1052,7 +1217,7 @@ struct LowerEnzymeXLABLASPass
// We need to run SymmOp lowering first, since it conditionally lowers
// to a custom call that needs us to detect potential Syrk Ops
RewritePatternSet patternsSet1(context);
patternsSet1.add<SymmOpLowering>(backend, blasIntWidth, context);
patternsSet1.add<SymmOpLowering, HemmOpLowering>(backend, blasIntWidth, context);
if (failed(applyPatternsGreedily(getOperation(), std::move(patternsSet1),
config))) {
signalPassFailure();
Expand Down
Loading
Loading