From 014404cd3d0f0ea95f0ac0c39107afe1da7f9015 Mon Sep 17 00:00:00 2001 From: Vimarsh Sathia Date: Sun, 4 Oct 2026 12:10:50 -0500 Subject: [PATCH 1/5] Canonicalize XLA wrapper inputs Remove unused input/result pairs and merge duplicate buffer inputs when their results agree. Share one reduced function across compatible callers when removing unused inputs. Preserve callers that need the original signature. Register both rewrites for ordinary canonicalization and add LIT coverage. Assisted-By: OpenAI Codex --- src/enzyme_ad/jax/Dialect/EnzymeXLAOps.td | 6 + src/enzyme_ad/jax/Dialect/Ops.cpp | 233 ++++++++ test/lit_tests/xla_wrapper_canonicalize.mlir | 508 ++++++++++++++++++ ...xla_wrapper_canonicalize_unknown_uses.mlir | 38 ++ 4 files changed, 785 insertions(+) create mode 100644 test/lit_tests/xla_wrapper_canonicalize.mlir create mode 100644 test/lit_tests/xla_wrapper_canonicalize_unknown_uses.mlir diff --git a/src/enzyme_ad/jax/Dialect/EnzymeXLAOps.td b/src/enzyme_ad/jax/Dialect/EnzymeXLAOps.td index f839c8e930..8ffee3814d 100644 --- a/src/enzyme_ad/jax/Dialect/EnzymeXLAOps.td +++ b/src/enzyme_ad/jax/Dialect/EnzymeXLAOps.td @@ -300,6 +300,11 @@ def XLAWrapperOp : EnzymeXLA_Op<"xla_wrapper", [ DeclareOpInterfaceMethods ]> { let summary = "XLA Call operation"; + let description = [{ + Calls a raised function whose input/output pairs represent device buffers. + The corresponding function result replaces each buffer's contents. Trailing + specialized scalar inputs are compile-time values rather than device buffers. + }]; let arguments = (ins SymbolRefAttr:$fn, Variadic:$inputs, OptionalAttr:$arg_attrs, @@ -312,6 +317,7 @@ def XLAWrapperOp : EnzymeXLA_Op<"xla_wrapper", [ $fn ` ` `(` $inputs `)` prop-dict attr-dict `:` functional-type($inputs, results) }]; + let hasCanonicalizer = 1; } def GPUErrorOp : EnzymeXLA_Op<"gpu_error", [ diff --git a/src/enzyme_ad/jax/Dialect/Ops.cpp b/src/enzyme_ad/jax/Dialect/Ops.cpp index b81b77b84a..c37798686d 100644 --- a/src/enzyme_ad/jax/Dialect/Ops.cpp +++ b/src/enzyme_ad/jax/Dialect/Ops.cpp @@ -2200,6 +2200,239 @@ void XLAWrapperOp::getEffects( effects.emplace_back(MemoryEffects::Effect::get()); } +/// Remove a state slot if the function only returns its input unchanged. +/// Change a private function in place if every use supports the rewrite. +/// Otherwise, use one clone for compatible wrappers and keep other calls. +class RemoveUnusedXLAWrapperInputs final + : public OpRewritePattern { + static bool canRewrite(XLAWrapperOp wrapper) { + return !wrapper.getArgAttrsAttr() && !wrapper.getResAttrsAttr() && + !wrapper.getNumSpecialized(); + } + +public: + using OpRewritePattern::OpRewritePattern; + + LogicalResult matchAndRewrite(XLAWrapperOp op, + PatternRewriter &rewriter) const override { + if (!canRewrite(op)) + return failure(); + + auto function = dyn_cast_or_null( + SymbolTable::lookupNearestSymbolFrom(op, op.getFnAttr())); + if (!function || !function.getFunctionBody().hasOneBlock() || + function.getNumArguments() != op.getInputs().size() || + function.getArgumentTypes() != function.getResultTypes() || + function.getAllArgAttrs() || function.getAllResultAttrs()) + return failure(); + + // Keep relative symbol references in the cloned body in the same scope. + Operation *scope = SymbolTable::getNearestSymbolTable(op); + if (function->getParentOp() != scope) + return failure(); + + Block &body = function.getFunctionBody().front(); + Operation *returnOp = body.getTerminator(); + if (!returnOp->hasTrait() || + returnOp->getNumOperands() != function.getNumArguments()) + return failure(); + + llvm::BitVector unused(function.getNumArguments()); + for (BlockArgument argument : body.getArguments()) { + unsigned index = argument.getArgNumber(); + if (argument.hasOneUse() && returnOp->getOperand(index) == argument) + unused.set(index); + } + if (unused.none()) + return failure(); + + Type type = function.getTypeWithoutArgsAndResults(unused, unused); + if (!type) + return failure(); + + llvm::SmallSetVector callers; + callers.insert(op); + auto uses = SymbolTable::getSymbolUses(function.getOperation(), scope); + bool modifyInPlace = + uses.has_value() && SymbolTable::getSymbolVisibility(function) == + SymbolTable::Visibility::Private; + if (uses) { + for (const auto &use : *uses) { + auto wrapper = dyn_cast(use.getUser()); + if (wrapper && canRewrite(wrapper) && + wrapper.getFnAttr() == op.getFnAttr() && + wrapper.getInputs().size() == function.getNumArguments() && + SymbolTable::getNearestSymbolTable(wrapper) == scope) + callers.insert(wrapper); + else + modifyInPlace = false; + } + } + + FunctionOpInterface specialized = function; + std::string name = function.getName().str(); + if (!modifyInPlace) { + SymbolTable symbols(scope); + std::string stem = (function.getName() + "_without_unused").str(); + name = stem; + for (unsigned suffix = 0; symbols.lookup(name); ++suffix) + name = stem + "_" + std::to_string(suffix); + + Operation *clone = function->clone(); + specialized = cast(clone); + SymbolTable::setSymbolName(clone, name); + SymbolTable::setSymbolVisibility(clone, SymbolTable::Visibility::Private); + OpBuilder::InsertionGuard guard(rewriter); + rewriter.setInsertionPointAfter(function); + rewriter.insert(clone); + } + + rewriter.startOpModification(specialized); + Block &specializedBody = specialized.getFunctionBody().front(); + Operation *specializedReturn = specializedBody.getTerminator(); + rewriter.startOpModification(specializedReturn); + specializedReturn->eraseOperands(unused); + specializedBody.eraseArguments(unused); + specialized.setType(type); + rewriter.finalizeOpModification(specializedReturn); + rewriter.finalizeOpModification(specialized); + + for (Operation *caller : callers) { + auto wrapper = cast(caller); + rewriter.startOpModification(wrapper); + wrapper->eraseOperands(unused); + wrapper.setFnAttr(FlatSymbolRefAttr::get(op.getContext(), name)); + rewriter.finalizeOpModification(wrapper); + } + return success(); + } +}; + +/// Remove casts that preserve the buffer address. Keep views that change it. +static Value getXLAWrapperBufferIdentity(Value value) { + while (Operation *definingOp = value.getDefiningOp()) { + if (auto pointerToMemref = dyn_cast(definingOp)) { + value = pointerToMemref.getSource(); + continue; + } + if (auto cast = dyn_cast(definingOp)) { + value = cast.getSource(); + continue; + } + break; + } + return value; +} + +/// Merge equal inputs only when their returned values also agree. +/// Clone the function because other calls can pass distinct buffers. +class DeduplicateXLAWrapperInputs final + : public OpRewritePattern { +public: + using OpRewritePattern::OpRewritePattern; + + LogicalResult matchAndRewrite(XLAWrapperOp op, + PatternRewriter &rewriter) const override { + if (op.getArgAttrsAttr() || op.getResAttrsAttr() || op.getNumSpecialized()) + return failure(); + + auto function = dyn_cast_or_null( + SymbolTable::lookupNearestSymbolFrom(op, op.getFnAttr())); + auto hasMetadata = [](ArrayAttr attributes) { + return attributes && llvm::any_of(attributes, [](Attribute attribute) { + return !cast(attribute).empty(); + }); + }; + if (!function || !function.getFunctionBody().hasOneBlock() || + function.getNumArguments() != op.getInputs().size() || + function.getArgumentTypes() != function.getResultTypes() || + hasMetadata(function.getAllArgAttrs()) || + hasMetadata(function.getAllResultAttrs())) + return failure(); + + Operation *scope = SymbolTable::getNearestSymbolTable(op); + if (function->getParentOp() != scope) + return failure(); + + Block &body = function.getFunctionBody().front(); + Operation *returnOp = body.getTerminator(); + if (!returnOp->hasTrait() || + returnOp->getNumOperands() != function.getNumArguments()) + return failure(); + + DenseMap identities; + SmallVector representatives; + llvm::BitVector duplicates(function.getNumArguments()); + for (auto [index, input] : llvm::enumerate(op.getInputs())) { + auto [entry, inserted] = + identities.try_emplace(getXLAWrapperBufferIdentity(input), index); + unsigned representative = entry->second; + representatives.push_back(representative); + if (inserted) + continue; + if (input.getType() != op.getInputs()[representative].getType() || + function.getArgumentTypes()[index] != + function.getArgumentTypes()[representative]) + return failure(); + duplicates.set(index); + } + if (duplicates.none()) + return failure(); + + auto canonicalResult = [&](Value result) { + if (auto argument = dyn_cast(result); + argument && argument.getOwner() == &body) + return Value( + function.getArgument(representatives[argument.getArgNumber()])); + return result; + }; + for (auto [index, representative] : llvm::enumerate(representatives)) + if (canonicalResult(returnOp->getOperand(index)) != + canonicalResult(returnOp->getOperand(representative))) + return failure(); + + Type type = function.getTypeWithoutArgsAndResults(duplicates, duplicates); + if (!type) + return failure(); + + SymbolTable symbols(scope); + std::string stem = (function.getName() + "_without_duplicates").str(); + std::string name = stem; + for (unsigned suffix = 0; symbols.lookup(name); ++suffix) + name = stem + "_" + std::to_string(suffix); + + Operation *clone = function->clone(); + auto specialized = cast(clone); + Block &specializedBody = specialized.getFunctionBody().front(); + for (auto [index, representative] : llvm::enumerate(representatives)) + if (duplicates.test(index)) + specializedBody.getArgument(index).replaceAllUsesWith( + specializedBody.getArgument(representative)); + specializedBody.getTerminator()->eraseOperands(duplicates); + specializedBody.eraseArguments(duplicates); + specialized.setType(type); + SymbolTable::setSymbolName(clone, name); + SymbolTable::setSymbolVisibility(clone, SymbolTable::Visibility::Private); + { + OpBuilder::InsertionGuard guard(rewriter); + rewriter.setInsertionPointAfter(function); + rewriter.insert(clone); + } + + rewriter.modifyOpInPlace(op, [&] { + op->eraseOperands(duplicates); + op.setFnAttr(FlatSymbolRefAttr::get(op.getContext(), name)); + }); + return success(); + } +}; + +void XLAWrapperOp::getCanonicalizationPatterns(RewritePatternSet &results, + MLIRContext *context) { + results.insert( + context); +} + //===----------------------------------------------------------------------===// // AlternativesOp //===----------------------------------------------------------------------===// diff --git a/test/lit_tests/xla_wrapper_canonicalize.mlir b/test/lit_tests/xla_wrapper_canonicalize.mlir new file mode 100644 index 0000000000..29adeba845 --- /dev/null +++ b/test/lit_tests/xla_wrapper_canonicalize.mlir @@ -0,0 +1,508 @@ +// RUN: enzymexlamlir-opt %s --split-input-file --canonicalize --symbol-dce | FileCheck %s + +// For buffer inputs, the function result at each position replaces that buffer. +// Remove a position only if the function returns its argument unchanged there +// and has no other use of the argument. +// +// One wrapper is the only user of this private function. Remove spare from +// both signatures. Keep the original function name and body. +// Before: wrapper @negate(data, spare); negate(data, spare) -> (-data, spare) +// After: wrapper @negate(data); negate(data) -> -data +module { + func.func @single_private_call(%data: memref<4xf32>, %spare: memref<4xf32>) { + enzymexla.xla_wrapper @negate (%data, %spare) : + (memref<4xf32>, memref<4xf32>) -> () + return + } + func.func private @negate(%data: tensor<4xf32>, %spare: tensor<4xf32>) + -> (tensor<4xf32>, tensor<4xf32>) { + %negative = stablehlo.negate %data : tensor<4xf32> + return %negative, %spare : tensor<4xf32>, tensor<4xf32> + } +} + +// CHECK-LABEL: func.func @single_private_call( +// CHECK-SAME: %[[DATA:[^:]+]]: memref<4xf32>, %[[SPARE:[^:]+]]: memref<4xf32>) { +// CHECK-NEXT: enzymexla.xla_wrapper @negate (%[[DATA]]) : (memref<4xf32>) -> () +// CHECK-NEXT: return +// CHECK-NEXT: } +// CHECK-LABEL: func.func private @negate( +// CHECK-SAME: %[[DATA:[^:]+]]: tensor<4xf32>) -> tensor<4xf32> { +// CHECK-NEXT: %[[NEGATIVE:[^ ]+]] = stablehlo.negate %[[DATA]] : tensor<4xf32> +// CHECK-NEXT: return %[[NEGATIVE]] : tensor<4xf32> +// CHECK-NEXT: } +// CHECK-NEXT: } + +// ----- + +// Both users are wrappers. They pass different buffers, but they can remove +// the same unused position. Update both calls and the original private function. +// Before: negate(data, spare) -> (-data, spare) +// After: negate(data) -> -data +module { + func.func @first_private_call(%data: memref<4xf32>, %spare: memref<4xf32>) { + enzymexla.xla_wrapper @negate (%data, %spare) : + (memref<4xf32>, memref<4xf32>) -> () + return + } + func.func @second_private_call(%other: memref<4xf32>, %other_spare: memref<4xf32>) { + enzymexla.xla_wrapper @negate (%other, %other_spare) : + (memref<4xf32>, memref<4xf32>) -> () + return + } + func.func private @negate(%data: tensor<4xf32>, %spare: tensor<4xf32>) + -> (tensor<4xf32>, tensor<4xf32>) { + %negative = stablehlo.negate %data : tensor<4xf32> + return %negative, %spare : tensor<4xf32>, tensor<4xf32> + } +} + +// CHECK-LABEL: func.func @first_private_call( +// CHECK-SAME: %[[DATA:[^:]+]]: memref<4xf32>, %[[SPARE:[^:]+]]: memref<4xf32>) { +// CHECK-NEXT: enzymexla.xla_wrapper @negate (%[[DATA]]) : (memref<4xf32>) -> () +// CHECK-NEXT: return +// CHECK-NEXT: } +// CHECK-LABEL: func.func @second_private_call( +// CHECK-SAME: %[[OTHER:[^:]+]]: memref<4xf32>, %[[OTHER_SPARE:[^:]+]]: memref<4xf32>) { +// CHECK-NEXT: enzymexla.xla_wrapper @negate (%[[OTHER]]) : (memref<4xf32>) -> () +// CHECK-NEXT: return +// CHECK-NEXT: } +// CHECK-LABEL: func.func private @negate( +// CHECK-SAME: %[[DATA:[^:]+]]: tensor<4xf32>) -> tensor<4xf32> { +// CHECK-NEXT: %[[NEGATIVE:[^ ]+]] = stablehlo.negate %[[DATA]] : tensor<4xf32> +// CHECK-NEXT: return %[[NEGATIVE]] : tensor<4xf32> +// CHECK-NEXT: } +// CHECK-NEXT: } + +// ----- + +// This public function can have callers outside the module. Keep its signature +// even though only one wrapper is visible. Give that wrapper a reduced clone. +// Before: wrapper @negate(data, spare) +// After: wrapper @negate_without_unused(data) +module { + func.func @single_public_call(%data: memref<4xf32>, %spare: memref<4xf32>) { + enzymexla.xla_wrapper @negate (%data, %spare) : + (memref<4xf32>, memref<4xf32>) -> () + return + } + func.func @negate(%data: tensor<4xf32>, %spare: tensor<4xf32>) + -> (tensor<4xf32>, tensor<4xf32>) { + %negative = stablehlo.negate %data : tensor<4xf32> + return %negative, %spare : tensor<4xf32>, tensor<4xf32> + } +} + +// CHECK-LABEL: func.func @single_public_call( +// CHECK-SAME: %[[DATA:[^:]+]]: memref<4xf32>, %[[SPARE:[^:]+]]: memref<4xf32>) { +// CHECK-NEXT: enzymexla.xla_wrapper @negate_without_unused (%[[DATA]]) : (memref<4xf32>) -> () +// CHECK-NEXT: return +// CHECK-NEXT: } +// CHECK-LABEL: func.func @negate( +// CHECK-SAME: %[[DATA:[^:]+]]: tensor<4xf32>, %[[SPARE:[^:]+]]: tensor<4xf32>) -> (tensor<4xf32>, tensor<4xf32>) { +// CHECK-NEXT: %[[NEGATIVE:[^ ]+]] = stablehlo.negate %[[DATA]] : tensor<4xf32> +// CHECK-NEXT: return %[[NEGATIVE]], %[[SPARE]] : tensor<4xf32>, tensor<4xf32> +// CHECK-NEXT: } +// CHECK-LABEL: func.func private @negate_without_unused( +// CHECK-SAME: %[[DATA:[^:]+]]: tensor<4xf32>) -> tensor<4xf32> { +// CHECK-NEXT: %[[NEGATIVE:[^ ]+]] = stablehlo.negate %[[DATA]] : tensor<4xf32> +// CHECK-NEXT: return %[[NEGATIVE]] : tensor<4xf32> +// CHECK-NEXT: } +// CHECK-NEXT: } + +// ----- + +// One wrapper is not the only symbol use: an ordinary call also needs the +// private function. Keep that call and the original signature unchanged. +// Before: wrapper @negate(data, spare) +// After: wrapper @negate_without_unused(data) +module { + func.func @wrapper_beside_direct_call(%data: memref<4xf32>, %spare: memref<4xf32>) { + enzymexla.xla_wrapper @negate (%data, %spare) : + (memref<4xf32>, memref<4xf32>) -> () + return + } + func.func @keep_direct_call(%data: tensor<4xf32>, %spare: tensor<4xf32>) + -> (tensor<4xf32>, tensor<4xf32>) { + %results:2 = call @negate(%data, %spare) : + (tensor<4xf32>, tensor<4xf32>) -> (tensor<4xf32>, tensor<4xf32>) + return %results#0, %results#1 : tensor<4xf32>, tensor<4xf32> + } + func.func private @negate(%data: tensor<4xf32>, %spare: tensor<4xf32>) + -> (tensor<4xf32>, tensor<4xf32>) { + %negative = stablehlo.negate %data : tensor<4xf32> + return %negative, %spare : tensor<4xf32>, tensor<4xf32> + } +} + +// CHECK-LABEL: func.func @wrapper_beside_direct_call( +// CHECK-SAME: %[[DATA:[^:]+]]: memref<4xf32>, %[[SPARE:[^:]+]]: memref<4xf32>) { +// CHECK-NEXT: enzymexla.xla_wrapper @negate_without_unused (%[[DATA]]) : (memref<4xf32>) -> () +// CHECK-NEXT: return +// CHECK-NEXT: } +// CHECK-LABEL: func.func @keep_direct_call( +// CHECK-SAME: %[[DATA:[^:]+]]: tensor<4xf32>, %[[SPARE:[^:]+]]: tensor<4xf32>) -> (tensor<4xf32>, tensor<4xf32>) { +// CHECK-NEXT: %[[RESULTS:[^:]+]]:2 = call @negate(%[[DATA]], %[[SPARE]]) : (tensor<4xf32>, tensor<4xf32>) -> (tensor<4xf32>, tensor<4xf32>) +// CHECK-NEXT: return %[[RESULTS]]#0, %[[RESULTS]]#1 : tensor<4xf32>, tensor<4xf32> +// CHECK-NEXT: } +// CHECK-LABEL: func.func private @negate( +// CHECK-SAME: %[[DATA:[^:]+]]: tensor<4xf32>, %[[SPARE:[^:]+]]: tensor<4xf32>) -> (tensor<4xf32>, tensor<4xf32>) { +// CHECK-NEXT: %[[NEGATIVE:[^ ]+]] = stablehlo.negate %[[DATA]] : tensor<4xf32> +// CHECK-NEXT: return %[[NEGATIVE]], %[[SPARE]] : tensor<4xf32>, tensor<4xf32> +// CHECK-NEXT: } +// CHECK-LABEL: func.func private @negate_without_unused( +// CHECK-SAME: %[[DATA:[^:]+]]: tensor<4xf32>) -> tensor<4xf32> { +// CHECK-NEXT: %[[NEGATIVE:[^ ]+]] = stablehlo.negate %[[DATA]] : tensor<4xf32> +// CHECK-NEXT: return %[[NEGATIVE]] : tensor<4xf32> +// CHECK-NEXT: } +// CHECK-NEXT: } + +// ----- + +// Remove the middle buffer: no operation reads or changes spare. +// Before: wrapper @scale_data(data, spare, scale) +// scale_data(data, spare, scale) -> (data * scale, spare, scale) +// After: wrapper @scale_data_without_unused(data, scale) +// scale_data_without_unused(data, scale) -> (data * scale, scale) +// +// Keep scale: the multiplication reads it, even though it is returned unchanged. +// Use one clone for all three compatible wrappers in the two host functions. +// The wrapper with metadata needs the original signature, so cloning is required. +module { + func.func @drop_unused_middle_buffer(%data: memref<4xf32>, + %spare: memref<4xf32>, + %scale: memref<4xf32>) { + enzymexla.xla_wrapper @scale_data (%data, %spare, %scale) : + (memref<4xf32>, memref<4xf32>, memref<4xf32>) -> () + return + } + func.func @share_reduced_function(%data: memref<4xf32>, + %spare: memref<4xf32>, + %scale: memref<4xf32>) { + enzymexla.xla_wrapper @scale_data (%data, %spare, %scale) : + (memref<4xf32>, memref<4xf32>, memref<4xf32>) -> () + enzymexla.xla_wrapper @scale_data (%data, %spare, %scale) : + (memref<4xf32>, memref<4xf32>, memref<4xf32>) -> () + enzymexla.xla_wrapper @scale_data (%data, %spare, %scale) + {arg_attrs = [{}, {test.keep}, {}]} : + (memref<4xf32>, memref<4xf32>, memref<4xf32>) -> () + return + } + func.func private @scale_data(%data: tensor<4xf32>, + %spare: tensor<4xf32>, + %scale: tensor<4xf32>) + -> (tensor<4xf32>, tensor<4xf32>, tensor<4xf32>) { + %scaled = stablehlo.multiply %data, %scale : tensor<4xf32> + return %scaled, %spare, %scale : tensor<4xf32>, tensor<4xf32>, tensor<4xf32> + } +} + +// The wrapper passes data and scale, in that order. It no longer passes spare. +// CHECK-LABEL: func.func @drop_unused_middle_buffer( +// CHECK-SAME: %[[DATA:[^:]+]]: memref<4xf32>, +// CHECK-SAME: %[[SPARE:[^:]+]]: memref<4xf32>, +// CHECK-SAME: %[[SCALE:[^:]+]]: memref<4xf32>) { +// CHECK-NEXT: enzymexla.xla_wrapper @scale_data_without_unused (%[[DATA]], %[[SCALE]]) : (memref<4xf32>, memref<4xf32>) -> () +// CHECK-NEXT: return +// CHECK-NEXT: } + +// These wrappers share the same clone. The wrapper with metadata stays intact. +// CHECK-LABEL: func.func @share_reduced_function( +// CHECK-SAME: %[[DATA:[^:]+]]: memref<4xf32>, +// CHECK-SAME: %[[SPARE:[^:]+]]: memref<4xf32>, +// CHECK-SAME: %[[SCALE:[^:]+]]: memref<4xf32>) { +// CHECK-NEXT: enzymexla.xla_wrapper @scale_data_without_unused (%[[DATA]], %[[SCALE]]) : (memref<4xf32>, memref<4xf32>) -> () +// CHECK-NEXT: enzymexla.xla_wrapper @scale_data_without_unused (%[[DATA]], %[[SCALE]]) : (memref<4xf32>, memref<4xf32>) -> () +// CHECK-NEXT: enzymexla.xla_wrapper @scale_data (%[[DATA]], %[[SPARE]], %[[SCALE]]) {arg_attrs = [{}, {test.keep}, {}]} : (memref<4xf32>, memref<4xf32>, memref<4xf32>) -> () +// CHECK-NEXT: return +// CHECK-NEXT: } + +// The original function retains the spare argument and its matching result. +// CHECK-LABEL: func.func private @scale_data( +// CHECK-SAME: %[[DATA:[^:]+]]: tensor<4xf32>, +// CHECK-SAME: %[[SPARE:[^:]+]]: tensor<4xf32>, +// CHECK-SAME: %[[SCALE:[^:]+]]: tensor<4xf32>) -> (tensor<4xf32>, tensor<4xf32>, tensor<4xf32>) { +// CHECK-NEXT: %[[SCALED:[^ ]+]] = stablehlo.multiply %[[DATA]], %[[SCALE]] : tensor<4xf32> +// CHECK-NEXT: return %[[SCALED]], %[[SPARE]], %[[SCALE]] : tensor<4xf32>, tensor<4xf32>, tensor<4xf32> +// CHECK-NEXT: } + +// The clone has two arguments and two results. The multiplication is unchanged. +// CHECK-LABEL: func.func private @scale_data_without_unused( +// CHECK-SAME: %[[DATA:[^:]+]]: tensor<4xf32>, +// CHECK-SAME: %[[SCALE:[^:]+]]: tensor<4xf32>) -> (tensor<4xf32>, tensor<4xf32>) { +// CHECK-NEXT: %[[SCALED:[^ ]+]] = stablehlo.multiply %[[DATA]], %[[SCALE]] : tensor<4xf32> +// CHECK-NEXT: return %[[SCALED]], %[[SCALE]] : tensor<4xf32>, tensor<4xf32> +// CHECK-NEXT: } +// CHECK-NEXT: } + +// ----- + +// Keep a read-only input. Returning source unchanged does not make it unused: +// the addition reads source to compute the new destination. +// Before and after: add_to_destination(source, destination) +// -> (source, source + destination) +module { + func.func @keep_readonly_source(%source: memref<4xf32>, %destination: memref<4xf32>) { + enzymexla.xla_wrapper @add_to_destination (%source, %destination) : + (memref<4xf32>, memref<4xf32>) -> () + return + } + func.func private @add_to_destination(%source: tensor<4xf32>, %destination: tensor<4xf32>) + -> (tensor<4xf32>, tensor<4xf32>) { + %sum = stablehlo.add %source, %destination : tensor<4xf32> + return %source, %sum : tensor<4xf32>, tensor<4xf32> + } +} + +// CHECK-LABEL: func.func @keep_readonly_source( +// CHECK-SAME: %[[SOURCE:[^:]+]]: memref<4xf32>, %[[DEST:[^:]+]]: memref<4xf32>) { +// CHECK-NEXT: enzymexla.xla_wrapper @add_to_destination (%[[SOURCE]], %[[DEST]]) : (memref<4xf32>, memref<4xf32>) -> () +// CHECK-NEXT: return +// CHECK-NEXT: } +// CHECK-LABEL: func.func private @add_to_destination( +// CHECK-SAME: %[[SOURCE:[^:]+]]: tensor<4xf32>, %[[DEST:[^:]+]]: tensor<4xf32>) -> (tensor<4xf32>, tensor<4xf32>) { +// CHECK-NEXT: %[[SUM:[^ ]+]] = stablehlo.add %[[SOURCE]], %[[DEST]] : tensor<4xf32> +// CHECK-NEXT: return %[[SOURCE]], %[[SUM]] : tensor<4xf32>, tensor<4xf32> +// CHECK-NEXT: } + +// ----- + +// Keep a buffer that receives a new value. The old destination is never read, +// but the second result writes source into it. Removing this position loses +// that write. The first input also remains: two result positions use it. +// Before and after: copy_to_destination(source, destination) -> (source, source) +module { + func.func @keep_written_destination(%source: memref<4xf32>, %destination: memref<4xf32>) { + enzymexla.xla_wrapper @copy_to_destination (%source, %destination) : + (memref<4xf32>, memref<4xf32>) -> () + return + } + func.func private @copy_to_destination(%source: tensor<4xf32>, %destination: tensor<4xf32>) + -> (tensor<4xf32>, tensor<4xf32>) { + return %source, %source : tensor<4xf32>, tensor<4xf32> + } +} + +// CHECK-LABEL: func.func @keep_written_destination( +// CHECK-SAME: %[[SOURCE:[^:]+]]: memref<4xf32>, %[[DEST:[^:]+]]: memref<4xf32>) { +// CHECK-NEXT: enzymexla.xla_wrapper @copy_to_destination (%[[SOURCE]], %[[DEST]]) : (memref<4xf32>, memref<4xf32>) -> () +// CHECK-NEXT: return +// CHECK-NEXT: } +// CHECK-LABEL: func.func private @copy_to_destination( +// CHECK-SAME: %[[SOURCE:[^:]+]]: tensor<4xf32>, %[[DEST:[^:]+]]: tensor<4xf32>) -> (tensor<4xf32>, tensor<4xf32>) { +// CHECK-NEXT: return %[[SOURCE]], %[[SOURCE]] : tensor<4xf32>, tensor<4xf32> +// CHECK-NEXT: } + +// ----- + +// Remove every buffer, but keep the call and its side effect. +// Before: wrapper @tick(unused); tick(unused) runs side_effect and returns unused. +// After: wrapper @tick(); tick() runs side_effect. +module { + func.func @keep_effect(%unused: memref<4xf32>) { + enzymexla.xla_wrapper @tick (%unused) : (memref<4xf32>) -> () + return + } + func.func private @tick(%a: tensor<4xf32>) -> tensor<4xf32> { + stablehlo.custom_call @side_effect() {has_side_effect = true} : () -> () + return %a : tensor<4xf32> + } +} + +// CHECK-LABEL: func.func @keep_effect( +// CHECK-SAME: %[[UNUSED:[^:]+]]: memref<4xf32>) { +// CHECK-NEXT: enzymexla.xla_wrapper @tick () : () -> () +// CHECK-NEXT: return +// CHECK-NEXT: } +// CHECK-LABEL: func.func private @tick() { +// CHECK-NEXT: stablehlo.custom_call @side_effect() {has_side_effect = true} : () -> () +// CHECK-NEXT: return +// CHECK-NEXT: } + +// ----- + +// Keep these calls unchanged. Removing a position would also require changes +// to argument/result attributes or the specialized scalar input mapping. +module { + func.func @keep_metadata(%buffer: memref<4xf32>, %bound: i64) { + enzymexla.xla_wrapper @identity (%buffer) {arg_attrs = [{test.keep}]} : (memref<4xf32>) -> () + enzymexla.xla_wrapper @identity (%buffer) {res_attrs = [{test.keep}]} : (memref<4xf32>) -> () + enzymexla.xla_wrapper @arg_metadata (%buffer) : (memref<4xf32>) -> () + enzymexla.xla_wrapper @result_metadata (%buffer) : (memref<4xf32>) -> () + enzymexla.xla_wrapper @specialized (%buffer, %bound) {num_specialized = 1 : i64} : (memref<4xf32>, i64) -> () + return + } + func.func private @identity(%a: tensor<4xf32>) -> tensor<4xf32> { + return %a : tensor<4xf32> + } + func.func private @arg_metadata(%a: tensor<4xf32> {test.keep}) -> tensor<4xf32> { + return %a : tensor<4xf32> + } + func.func private @result_metadata(%a: tensor<4xf32>) -> (tensor<4xf32> {test.keep}) { + return %a : tensor<4xf32> + } + func.func private @specialized(%a: tensor<4xf32>, %bound: tensor) -> tensor<4xf32> { + return %a : tensor<4xf32> + } +} + +// CHECK-LABEL: func.func @keep_metadata( +// CHECK-SAME: %[[BUFFER:[^:]+]]: memref<4xf32>, %[[BOUND:[^:]+]]: i64) { +// CHECK-NEXT: enzymexla.xla_wrapper @identity (%[[BUFFER]]) {arg_attrs = [{test.keep}]} : (memref<4xf32>) -> () +// CHECK-NEXT: enzymexla.xla_wrapper @identity (%[[BUFFER]]) {res_attrs = [{test.keep}]} : (memref<4xf32>) -> () +// CHECK-NEXT: enzymexla.xla_wrapper @arg_metadata (%[[BUFFER]]) : (memref<4xf32>) -> () +// CHECK-NEXT: enzymexla.xla_wrapper @result_metadata (%[[BUFFER]]) : (memref<4xf32>) -> () +// CHECK-NEXT: enzymexla.xla_wrapper @specialized (%[[BUFFER]], %[[BOUND]]) {num_specialized = 1 : i64} : (memref<4xf32>, i64) -> () +// CHECK-NEXT: return +// CHECK-NEXT: } + +// ----- + +// Keep the nested callee unchanged. Moving its clone into another symbol table +// could change which symbols its body references. +module { + func.func @keep_nested_callee(%buffer: memref<4xf32>) { + enzymexla.xla_wrapper @nested::@identity (%buffer) : (memref<4xf32>) -> () + return + } + module @nested { + func.func @identity(%a: tensor<4xf32>) -> tensor<4xf32> { + return %a : tensor<4xf32> + } + } +} + +// CHECK-LABEL: func.func @keep_nested_callee( +// CHECK-SAME: %[[BUFFER:[^:]+]]: memref<4xf32>) { +// CHECK-NEXT: enzymexla.xla_wrapper @nested::@identity (%[[BUFFER]]) : (memref<4xf32>) -> () +// CHECK-NEXT: return +// CHECK-NEXT: } +// CHECK-LABEL: module @nested { +// CHECK-NEXT: func.func @identity(%[[BUFFER:[^:]+]]: tensor<4xf32>) -> tensor<4xf32> { +// CHECK-NEXT: return %[[BUFFER]] : tensor<4xf32> +// CHECK-NEXT: } + +// ----- + +// Both inputs point to the same buffer. Both results write the same value. +// Keep one input and one result. Use its tensor twice in the addition. +// Before: wrapper @pair(a, a); pair(a, b) -> (a + b, a + b) +// After: wrapper @pair_without_duplicates(a) +// pair_without_duplicates(a) -> a + a +// Keep the original function for the caller that passes distinct buffers. +module { + llvm.func @deduplicate_inputs(%a: !llvm.ptr) { + %first = "enzymexla.pointer2memref"(%a) : (!llvm.ptr) -> memref + %second = "enzymexla.pointer2memref"(%a) : (!llvm.ptr) -> memref + enzymexla.xla_wrapper @pair (%first, %second) : + (memref, memref) -> () + llvm.return + } + llvm.func @keep_other_call(%a: !llvm.ptr, %b: !llvm.ptr) { + %first = "enzymexla.pointer2memref"(%a) : (!llvm.ptr) -> memref + %second = "enzymexla.pointer2memref"(%b) : (!llvm.ptr) -> memref + enzymexla.xla_wrapper @pair (%first, %second) : + (memref, memref) -> () + llvm.return + } + func.func private @pair(%a: tensor, %b: tensor) + -> (tensor, tensor) { + %sum = stablehlo.add %a, %b : tensor + return %sum, %sum : tensor, tensor + } +} + +// CHECK-LABEL: llvm.func @deduplicate_inputs( +// CHECK-SAME: %[[A:[^:]+]]: !llvm.ptr) { +// CHECK-NEXT: %[[BUFFER:[^ ]+]] = "enzymexla.pointer2memref"(%[[A]]) : (!llvm.ptr) -> memref +// CHECK-NEXT: enzymexla.xla_wrapper @pair_without_duplicates (%[[BUFFER]]) : (memref) -> () +// CHECK-NEXT: llvm.return +// CHECK-NEXT: } + +// CHECK-LABEL: llvm.func @keep_other_call( +// CHECK-SAME: %[[A:[^:]+]]: !llvm.ptr, %[[B:[^:]+]]: !llvm.ptr) { +// CHECK-NEXT: %[[FIRST:[^ ]+]] = "enzymexla.pointer2memref"(%[[A]]) : (!llvm.ptr) -> memref +// CHECK-NEXT: %[[SECOND:[^ ]+]] = "enzymexla.pointer2memref"(%[[B]]) : (!llvm.ptr) -> memref +// CHECK-NEXT: enzymexla.xla_wrapper @pair (%[[FIRST]], %[[SECOND]]) : (memref, memref) -> () +// CHECK-NEXT: llvm.return +// CHECK-NEXT: } + +// CHECK-LABEL: func.func private @pair( +// CHECK-SAME: %[[A:[^:]+]]: tensor, %[[B:[^:]+]]: tensor) -> (tensor, tensor) { +// CHECK-NEXT: %[[SUM:[^ ]+]] = stablehlo.add %[[A]], %[[B]] : tensor +// CHECK-NEXT: return %[[SUM]], %[[SUM]] : tensor, tensor +// CHECK-NEXT: } + +// CHECK-LABEL: func.func private @pair_without_duplicates( +// CHECK-SAME: %[[A:[^:]+]]: tensor) -> tensor { +// CHECK-NEXT: %[[SUM:[^ ]+]] = stablehlo.add %[[A]], %[[A]] : tensor +// CHECK-NEXT: return %[[SUM]] : tensor +// CHECK-NEXT: } + +// ----- + +// Keep duplicate inputs when their results differ. Removing either result +// would remove one of the two different writes to the buffer. +// Before and after: wrapper @different_results(a, a) +// different_results(a, b) -> (a + b, -(a + b)) +module { + func.func @keep_conflicting_results(%buffer: memref<4xf32>) { + enzymexla.xla_wrapper @different_results (%buffer, %buffer) : + (memref<4xf32>, memref<4xf32>) -> () + return + } + func.func private @different_results(%a: tensor<4xf32>, %b: tensor<4xf32>) + -> (tensor<4xf32>, tensor<4xf32>) { + %sum = stablehlo.add %a, %b : tensor<4xf32> + %negative = stablehlo.negate %sum : tensor<4xf32> + return %sum, %negative : tensor<4xf32>, tensor<4xf32> + } +} + +// CHECK-LABEL: func.func @keep_conflicting_results( +// CHECK-SAME: %[[BUFFER:[^:]+]]: memref<4xf32>) { +// CHECK-NEXT: enzymexla.xla_wrapper @different_results (%[[BUFFER]], %[[BUFFER]]) : (memref<4xf32>, memref<4xf32>) -> () +// CHECK-NEXT: return +// CHECK-NEXT: } +// CHECK-LABEL: func.func private @different_results( +// CHECK-SAME: %[[A:[^:]+]]: tensor<4xf32>, %[[B:[^:]+]]: tensor<4xf32>) -> (tensor<4xf32>, tensor<4xf32>) { +// CHECK-NEXT: %[[SUM:[^ ]+]] = stablehlo.add %[[A]], %[[B]] : tensor<4xf32> +// CHECK-NEXT: %[[NEGATIVE:[^ ]+]] = stablehlo.negate %[[SUM]] : tensor<4xf32> +// CHECK-NEXT: return %[[SUM]], %[[NEGATIVE]] : tensor<4xf32>, tensor<4xf32> +// CHECK-NEXT: } +// CHECK-NEXT: } + +// ----- + +// The first two results are different arguments for the same input buffer. +// They become the same value when those arguments are merged. The addition +// reads both arguments, so unused-input removal cannot remove either one. +// Before: wrapper @read_pair(a, a, destination) +// read_pair(a, b, destination) -> (a, b, a + b) +// After: wrapper @read_pair_without_duplicates(a, destination) +// read_pair_without_duplicates(a, destination) -> (a, a + a) +module { + func.func @merge_returned_arguments(%source: memref<4xf32>, + %destination: memref<4xf32>) { + enzymexla.xla_wrapper @read_pair (%source, %source, %destination) : + (memref<4xf32>, memref<4xf32>, memref<4xf32>) -> () + return + } + func.func private @read_pair(%a: tensor<4xf32>, %b: tensor<4xf32>, + %destination: tensor<4xf32>) + -> (tensor<4xf32>, tensor<4xf32>, tensor<4xf32>) { + %sum = stablehlo.add %a, %b : tensor<4xf32> + return %a, %b, %sum : tensor<4xf32>, tensor<4xf32>, tensor<4xf32> + } +} + +// CHECK-LABEL: func.func @merge_returned_arguments( +// CHECK-SAME: %[[SOURCE:[^:]+]]: memref<4xf32>, %[[DESTINATION:[^:]+]]: memref<4xf32>) { +// CHECK-NEXT: enzymexla.xla_wrapper @read_pair_without_duplicates (%[[SOURCE]], %[[DESTINATION]]) : (memref<4xf32>, memref<4xf32>) -> () +// CHECK-NEXT: return +// CHECK-NEXT: } +// CHECK-LABEL: func.func private @read_pair_without_duplicates( +// CHECK-SAME: %[[SOURCE:[^:]+]]: tensor<4xf32>, %[[DESTINATION:[^:]+]]: tensor<4xf32>) -> (tensor<4xf32>, tensor<4xf32>) { +// CHECK-NEXT: %[[SUM:[^ ]+]] = stablehlo.add %[[SOURCE]], %[[SOURCE]] : tensor<4xf32> +// CHECK-NEXT: return %[[SOURCE]], %[[SUM]] : tensor<4xf32>, tensor<4xf32> +// CHECK-NEXT: } diff --git a/test/lit_tests/xla_wrapper_canonicalize_unknown_uses.mlir b/test/lit_tests/xla_wrapper_canonicalize_unknown_uses.mlir new file mode 100644 index 0000000000..721058d9d1 --- /dev/null +++ b/test/lit_tests/xla_wrapper_canonicalize_unknown_uses.mlir @@ -0,0 +1,38 @@ +// RUN: enzymexlamlir-opt %s --allow-unregistered-dialect --canonicalize | FileCheck %s + +// The unknown region can contain symbol uses that cannot be enumerated. +// Keep the original private function. Rewrite the known wrapper to a clone. +// Before: wrapper @negate(data, spare) +// After: wrapper @negate_without_unused(data) +module { + func.func @call_with_unknown_uses(%data: memref<4xf32>, %spare: memref<4xf32>) { + enzymexla.xla_wrapper @negate (%data, %spare) : + (memref<4xf32>, memref<4xf32>) -> () + return + } + func.func private @negate(%data: tensor<4xf32>, %spare: tensor<4xf32>) + -> (tensor<4xf32>, tensor<4xf32>) { + %negative = stablehlo.negate %data : tensor<4xf32> + return %negative, %spare : tensor<4xf32>, tensor<4xf32> + } + "test.unknown_scope"() ({}) : () -> () +} + +// CHECK-LABEL: func.func @call_with_unknown_uses( +// CHECK-SAME: %[[DATA:[^:]+]]: memref<4xf32>, %[[SPARE:[^:]+]]: memref<4xf32>) { +// CHECK-NEXT: enzymexla.xla_wrapper @negate_without_unused (%[[DATA]]) : (memref<4xf32>) -> () +// CHECK-NEXT: return +// CHECK-NEXT: } +// CHECK-LABEL: func.func private @negate( +// CHECK-SAME: %[[DATA:[^:]+]]: tensor<4xf32>, %[[SPARE:[^:]+]]: tensor<4xf32>) -> (tensor<4xf32>, tensor<4xf32>) { +// CHECK-NEXT: %[[NEGATIVE:[^ ]+]] = stablehlo.negate %[[DATA]] : tensor<4xf32> +// CHECK-NEXT: return %[[NEGATIVE]], %[[SPARE]] : tensor<4xf32>, tensor<4xf32> +// CHECK-NEXT: } +// CHECK-LABEL: func.func private @negate_without_unused( +// CHECK-SAME: %[[DATA:[^:]+]]: tensor<4xf32>) -> tensor<4xf32> { +// CHECK-NEXT: %[[NEGATIVE:[^ ]+]] = stablehlo.negate %[[DATA]] : tensor<4xf32> +// CHECK-NEXT: return %[[NEGATIVE]] : tensor<4xf32> +// CHECK-NEXT: } +// CHECK-NEXT: "test.unknown_scope"() ({ +// CHECK-NEXT: }) : () -> () +// CHECK-NEXT: } From 0bc08d48374527cb24c74153b118f79264235316 Mon Sep 17 00:00:00 2001 From: Vimarsh Sathia Date: Mon, 5 Oct 2026 09:44:42 -0500 Subject: [PATCH 2/5] Use NoOffsetViewInterface for wrapper canonicalization Mark the pointer conversion operations with NoOffsetViewInterface and require a zero-offset result layout for Pointer2MemrefOp. Add the interface headers and TableGen dependency. Use getBaseObject with offsets disabled to identify duplicate wrapper inputs. Remove the local helper and pin Enzyme to 95acfd4c from EnzymeAD/Enzyme#3420. Assisted-By: OpenAI Codex --- src/enzyme_ad/jax/Dialect/Ops.cpp | 21 +++------------------ 1 file changed, 3 insertions(+), 18 deletions(-) diff --git a/src/enzyme_ad/jax/Dialect/Ops.cpp b/src/enzyme_ad/jax/Dialect/Ops.cpp index c37798686d..34f15cffcf 100644 --- a/src/enzyme_ad/jax/Dialect/Ops.cpp +++ b/src/enzyme_ad/jax/Dialect/Ops.cpp @@ -9,6 +9,7 @@ #include "Enzyme/MLIR/Dialect/Ops.h" #include "Dialect.h" #include "Interfaces/AutoDiffTypeInterface.h" +#include "Interfaces/Utils.h" #include "Ops.h" #include "src/enzyme_ad/jax/Dialect/Canonicalizers.h" #include "src/enzyme_ad/jax/Dialect/Utils.h" @@ -2308,22 +2309,6 @@ class RemoveUnusedXLAWrapperInputs final } }; -/// Remove casts that preserve the buffer address. Keep views that change it. -static Value getXLAWrapperBufferIdentity(Value value) { - while (Operation *definingOp = value.getDefiningOp()) { - if (auto pointerToMemref = dyn_cast(definingOp)) { - value = pointerToMemref.getSource(); - continue; - } - if (auto cast = dyn_cast(definingOp)) { - value = cast.getSource(); - continue; - } - break; - } - return value; -} - /// Merge equal inputs only when their returned values also agree. /// Clone the function because other calls can pass distinct buffers. class DeduplicateXLAWrapperInputs final @@ -2364,8 +2349,8 @@ class DeduplicateXLAWrapperInputs final SmallVector representatives; llvm::BitVector duplicates(function.getNumArguments()); for (auto [index, input] : llvm::enumerate(op.getInputs())) { - auto [entry, inserted] = - identities.try_emplace(getXLAWrapperBufferIdentity(input), index); + auto [entry, inserted] = identities.try_emplace( + enzyme::oputils::getBaseObject(input, /*offsetAllowed=*/false), index); unsigned representative = entry->second; representatives.push_back(representative); if (inserted) From 8a9d2669bf7965cb393ea5bdd04e4f239bf8b29b Mon Sep 17 00:00:00 2001 From: Vimarsh Sathia Date: Mon, 5 Oct 2026 18:50:18 -0500 Subject: [PATCH 3/5] Use OffsetViewInterface for pointer conversions Use OffsetViewInterface for pointer2memref and memref2pointer. Keep the pointer2memref zero-offset verifier. The Enzyme dependency and shared pointer utility cleanup now come from main through Enzyme-JAX#3373. Update wrapper metadata in the LIT tests to use the current property syntax. Assisted-By: OpenAI Codex --- test/lit_tests/xla_wrapper_canonicalize.mlir | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/test/lit_tests/xla_wrapper_canonicalize.mlir b/test/lit_tests/xla_wrapper_canonicalize.mlir index 29adeba845..d176ccc032 100644 --- a/test/lit_tests/xla_wrapper_canonicalize.mlir +++ b/test/lit_tests/xla_wrapper_canonicalize.mlir @@ -184,7 +184,7 @@ module { enzymexla.xla_wrapper @scale_data (%data, %spare, %scale) : (memref<4xf32>, memref<4xf32>, memref<4xf32>) -> () enzymexla.xla_wrapper @scale_data (%data, %spare, %scale) - {arg_attrs = [{}, {test.keep}, {}]} : + <{arg_attrs = [{}, {test.keep}, {}]}> : (memref<4xf32>, memref<4xf32>, memref<4xf32>) -> () return } @@ -213,7 +213,7 @@ module { // CHECK-SAME: %[[SCALE:[^:]+]]: memref<4xf32>) { // CHECK-NEXT: enzymexla.xla_wrapper @scale_data_without_unused (%[[DATA]], %[[SCALE]]) : (memref<4xf32>, memref<4xf32>) -> () // CHECK-NEXT: enzymexla.xla_wrapper @scale_data_without_unused (%[[DATA]], %[[SCALE]]) : (memref<4xf32>, memref<4xf32>) -> () -// CHECK-NEXT: enzymexla.xla_wrapper @scale_data (%[[DATA]], %[[SPARE]], %[[SCALE]]) {arg_attrs = [{}, {test.keep}, {}]} : (memref<4xf32>, memref<4xf32>, memref<4xf32>) -> () +// CHECK-NEXT: enzymexla.xla_wrapper @scale_data (%[[DATA]], %[[SPARE]], %[[SCALE]]) <{arg_attrs = [{}, {test.keep}, {}]}> : (memref<4xf32>, memref<4xf32>, memref<4xf32>) -> () // CHECK-NEXT: return // CHECK-NEXT: } @@ -325,11 +325,11 @@ module { // to argument/result attributes or the specialized scalar input mapping. module { func.func @keep_metadata(%buffer: memref<4xf32>, %bound: i64) { - enzymexla.xla_wrapper @identity (%buffer) {arg_attrs = [{test.keep}]} : (memref<4xf32>) -> () - enzymexla.xla_wrapper @identity (%buffer) {res_attrs = [{test.keep}]} : (memref<4xf32>) -> () + enzymexla.xla_wrapper @identity (%buffer) <{arg_attrs = [{test.keep}]}> : (memref<4xf32>) -> () + enzymexla.xla_wrapper @identity (%buffer) <{res_attrs = [{test.keep}]}> : (memref<4xf32>) -> () enzymexla.xla_wrapper @arg_metadata (%buffer) : (memref<4xf32>) -> () enzymexla.xla_wrapper @result_metadata (%buffer) : (memref<4xf32>) -> () - enzymexla.xla_wrapper @specialized (%buffer, %bound) {num_specialized = 1 : i64} : (memref<4xf32>, i64) -> () + enzymexla.xla_wrapper @specialized (%buffer, %bound) <{num_specialized = 1 : i64}> : (memref<4xf32>, i64) -> () return } func.func private @identity(%a: tensor<4xf32>) -> tensor<4xf32> { @@ -348,11 +348,11 @@ module { // CHECK-LABEL: func.func @keep_metadata( // CHECK-SAME: %[[BUFFER:[^:]+]]: memref<4xf32>, %[[BOUND:[^:]+]]: i64) { -// CHECK-NEXT: enzymexla.xla_wrapper @identity (%[[BUFFER]]) {arg_attrs = [{test.keep}]} : (memref<4xf32>) -> () -// CHECK-NEXT: enzymexla.xla_wrapper @identity (%[[BUFFER]]) {res_attrs = [{test.keep}]} : (memref<4xf32>) -> () +// CHECK-NEXT: enzymexla.xla_wrapper @identity (%[[BUFFER]]) <{arg_attrs = [{test.keep}]}> : (memref<4xf32>) -> () +// CHECK-NEXT: enzymexla.xla_wrapper @identity (%[[BUFFER]]) <{res_attrs = [{test.keep}]}> : (memref<4xf32>) -> () // CHECK-NEXT: enzymexla.xla_wrapper @arg_metadata (%[[BUFFER]]) : (memref<4xf32>) -> () // CHECK-NEXT: enzymexla.xla_wrapper @result_metadata (%[[BUFFER]]) : (memref<4xf32>) -> () -// CHECK-NEXT: enzymexla.xla_wrapper @specialized (%[[BUFFER]], %[[BOUND]]) {num_specialized = 1 : i64} : (memref<4xf32>, i64) -> () +// CHECK-NEXT: enzymexla.xla_wrapper @specialized (%[[BUFFER]], %[[BOUND]]) <{num_specialized = 1 : i64}> : (memref<4xf32>, i64) -> () // CHECK-NEXT: return // CHECK-NEXT: } From bd8cad52a7a5bff779353d0cfacce48dc0131104 Mon Sep 17 00:00:00 2001 From: Vimarsh Sathia Date: Mon, 5 Oct 2026 22:00:11 -0500 Subject: [PATCH 4/5] Support specialized inputs and metadata in wrapper canonicalization Remove unused specialized scalars and update their count. Remove argument and result attributes with their slots, and preserve attributes on retained inputs and results. Keep duplicate buffers separate when their attributes differ. Use static helpers instead of closures. Assisted-By: OpenAI Codex --- src/enzyme_ad/jax/Dialect/Ops.cpp | 186 +++++++++++------- test/lit_tests/xla_wrapper_canonicalize.mlir | 47 ++--- .../xla_wrapper_canonicalize_specialized.mlir | 141 +++++++++++++ 3 files changed, 279 insertions(+), 95 deletions(-) create mode 100644 test/lit_tests/xla_wrapper_canonicalize_specialized.mlir diff --git a/src/enzyme_ad/jax/Dialect/Ops.cpp b/src/enzyme_ad/jax/Dialect/Ops.cpp index 34f15cffcf..6dbf314119 100644 --- a/src/enzyme_ad/jax/Dialect/Ops.cpp +++ b/src/enzyme_ad/jax/Dialect/Ops.cpp @@ -2201,30 +2201,60 @@ void XLAWrapperOp::getEffects( effects.emplace_back(MemoryEffects::Effect::get()); } -/// Remove a state slot if the function only returns its input unchanged. +static bool hasXLAWrapperSignature(XLAWrapperOp wrapper, + FunctionOpInterface function) { + int64_t numSpecialized = wrapper.getNumSpecialized(); + if (numSpecialized < 0 || + numSpecialized > static_cast(wrapper.getInputs().size())) + return false; + unsigned numBuffers = wrapper.getInputs().size() - numSpecialized; + return function.getNumArguments() == wrapper.getInputs().size() && + function.getNumResults() == numBuffers && + function.getArgumentTypes().take_front(numBuffers) == + function.getResultTypes() && + (!wrapper.getArgAttrsAttr() || + wrapper.getArgAttrsAttr().size() == wrapper.getInputs().size()) && + (!wrapper.getResAttrsAttr() || + wrapper.getResAttrsAttr().size() == numBuffers); +} + +static ArrayAttr removeWrapperAttrs(ArrayAttr attributes, + const llvm::BitVector &removed) { + SmallVector remaining; + for (auto [index, attribute] : llvm::enumerate(attributes)) + if (!removed.test(index)) + remaining.push_back(attribute); + return ArrayAttr::get(attributes.getContext(), remaining); +} + +static void removeWrapperInputs(XLAWrapperOp wrapper, + const llvm::BitVector &arguments, + const llvm::BitVector &results) { + if (auto attributes = wrapper.getArgAttrsAttr()) + wrapper.setArgAttrsAttr(removeWrapperAttrs(attributes, arguments)); + if (auto attributes = wrapper.getResAttrsAttr()) + wrapper.setResAttrsAttr(removeWrapperAttrs(attributes, results)); + unsigned numSpecialized = + wrapper.getNumSpecialized() - (arguments.count() - results.count()); + wrapper->eraseOperands(arguments); + wrapper.setNumSpecialized(numSpecialized); +} + +/// Remove a buffer slot if the function only returns its input unchanged. +/// Remove a specialized scalar input if the function does not use it. /// Change a private function in place if every use supports the rewrite. /// Otherwise, use one clone for compatible wrappers and keep other calls. class RemoveUnusedXLAWrapperInputs final : public OpRewritePattern { - static bool canRewrite(XLAWrapperOp wrapper) { - return !wrapper.getArgAttrsAttr() && !wrapper.getResAttrsAttr() && - !wrapper.getNumSpecialized(); - } - public: using OpRewritePattern::OpRewritePattern; LogicalResult matchAndRewrite(XLAWrapperOp op, PatternRewriter &rewriter) const override { - if (!canRewrite(op)) - return failure(); - auto function = dyn_cast_or_null( SymbolTable::lookupNearestSymbolFrom(op, op.getFnAttr())); if (!function || !function.getFunctionBody().hasOneBlock() || - function.getNumArguments() != op.getInputs().size() || - function.getArgumentTypes() != function.getResultTypes() || - function.getAllArgAttrs() || function.getAllResultAttrs()) + !hasXLAWrapperSignature(op, function)) return failure(); // Keep relative symbol references in the cloned body in the same scope. @@ -2234,21 +2264,31 @@ class RemoveUnusedXLAWrapperInputs final Block &body = function.getFunctionBody().front(); Operation *returnOp = body.getTerminator(); + unsigned numBuffers = function.getNumResults(); if (!returnOp->hasTrait() || - returnOp->getNumOperands() != function.getNumArguments()) + returnOp->getNumOperands() != numBuffers) return failure(); - llvm::BitVector unused(function.getNumArguments()); + llvm::BitVector unusedArguments(function.getNumArguments()); + llvm::BitVector unusedResults(numBuffers); for (BlockArgument argument : body.getArguments()) { unsigned index = argument.getArgNumber(); - if (argument.hasOneUse() && returnOp->getOperand(index) == argument) - unused.set(index); + if (index < numBuffers) { + if (argument.hasOneUse() && returnOp->getOperand(index) == argument) { + unusedArguments.set(index); + unusedResults.set(index); + } + } else if (argument.use_empty()) { + unusedArguments.set(index); + } } - if (unused.none()) + if (unusedArguments.none()) return failure(); - Type type = function.getTypeWithoutArgsAndResults(unused, unused); - if (!type) + Type withoutArguments = function.getTypeWithoutArgs(unusedArguments); + Type type = + function.getTypeWithoutArgsAndResults(unusedArguments, unusedResults); + if (!withoutArguments || !type) return failure(); llvm::SmallSetVector callers; @@ -2260,9 +2300,8 @@ class RemoveUnusedXLAWrapperInputs final if (uses) { for (const auto &use : *uses) { auto wrapper = dyn_cast(use.getUser()); - if (wrapper && canRewrite(wrapper) && - wrapper.getFnAttr() == op.getFnAttr() && - wrapper.getInputs().size() == function.getNumArguments() && + if (wrapper && wrapper.getFnAttr() == op.getFnAttr() && + hasXLAWrapperSignature(wrapper, function) && SymbolTable::getNearestSymbolTable(wrapper) == scope) callers.insert(wrapper); else @@ -2289,19 +2328,21 @@ class RemoveUnusedXLAWrapperInputs final } rewriter.startOpModification(specialized); - Block &specializedBody = specialized.getFunctionBody().front(); - Operation *specializedReturn = specializedBody.getTerminator(); + Operation *specializedReturn = + specialized.getFunctionBody().front().getTerminator(); rewriter.startOpModification(specializedReturn); - specializedReturn->eraseOperands(unused); - specializedBody.eraseArguments(unused); - specialized.setType(type); + specializedReturn->eraseOperands(unusedResults); + function_interface_impl::eraseFunctionArguments( + specialized, unusedArguments, withoutArguments); + function_interface_impl::eraseFunctionResults(specialized, unusedResults, + type); rewriter.finalizeOpModification(specializedReturn); rewriter.finalizeOpModification(specialized); for (Operation *caller : callers) { auto wrapper = cast(caller); rewriter.startOpModification(wrapper); - wrapper->eraseOperands(unused); + removeWrapperInputs(wrapper, unusedArguments, unusedResults); wrapper.setFnAttr(FlatSymbolRefAttr::get(op.getContext(), name)); rewriter.finalizeOpModification(wrapper); } @@ -2309,30 +2350,34 @@ class RemoveUnusedXLAWrapperInputs final } }; -/// Merge equal inputs only when their returned values also agree. +/// Merge equal buffer inputs only when their returned values and attributes +/// also agree. Keep specialized scalar inputs at the end of the input list. /// Clone the function because other calls can pass distinct buffers. class DeduplicateXLAWrapperInputs final : public OpRewritePattern { + static Value canonicalResult(Value result, Block &body, + ArrayRef representatives) { + if (auto argument = dyn_cast(result); + argument && argument.getOwner() == &body && + argument.getArgNumber() < representatives.size()) + return body.getArgument(representatives[argument.getArgNumber()]); + return result; + } + + static bool sameAttributes(ArrayAttr attributes, unsigned first, + unsigned second) { + return !attributes || attributes[first] == attributes[second]; + } + public: using OpRewritePattern::OpRewritePattern; LogicalResult matchAndRewrite(XLAWrapperOp op, PatternRewriter &rewriter) const override { - if (op.getArgAttrsAttr() || op.getResAttrsAttr() || op.getNumSpecialized()) - return failure(); - auto function = dyn_cast_or_null( SymbolTable::lookupNearestSymbolFrom(op, op.getFnAttr())); - auto hasMetadata = [](ArrayAttr attributes) { - return attributes && llvm::any_of(attributes, [](Attribute attribute) { - return !cast(attribute).empty(); - }); - }; if (!function || !function.getFunctionBody().hasOneBlock() || - function.getNumArguments() != op.getInputs().size() || - function.getArgumentTypes() != function.getResultTypes() || - hasMetadata(function.getAllArgAttrs()) || - hasMetadata(function.getAllResultAttrs())) + !hasXLAWrapperSignature(op, function)) return failure(); Operation *scope = SymbolTable::getNearestSymbolTable(op); @@ -2341,43 +2386,48 @@ class DeduplicateXLAWrapperInputs final Block &body = function.getFunctionBody().front(); Operation *returnOp = body.getTerminator(); + unsigned numBuffers = function.getNumResults(); if (!returnOp->hasTrait() || - returnOp->getNumOperands() != function.getNumArguments()) + returnOp->getNumOperands() != numBuffers) return failure(); DenseMap identities; SmallVector representatives; - llvm::BitVector duplicates(function.getNumArguments()); - for (auto [index, input] : llvm::enumerate(op.getInputs())) { + llvm::BitVector duplicateArguments(function.getNumArguments()); + llvm::BitVector duplicateResults(numBuffers); + for (auto [index, input] : + llvm::enumerate(op.getInputs().take_front(numBuffers))) { auto [entry, inserted] = identities.try_emplace( - enzyme::oputils::getBaseObject(input, /*offsetAllowed=*/false), index); + enzyme::oputils::getBaseObject(input, /*offsetAllowed=*/false), + index); unsigned representative = entry->second; representatives.push_back(representative); if (inserted) continue; if (input.getType() != op.getInputs()[representative].getType() || function.getArgumentTypes()[index] != - function.getArgumentTypes()[representative]) + function.getArgumentTypes()[representative] || + !sameAttributes(op.getArgAttrsAttr(), index, representative) || + !sameAttributes(op.getResAttrsAttr(), index, representative) || + !sameAttributes(function.getAllArgAttrs(), index, representative) || + !sameAttributes(function.getAllResultAttrs(), index, representative)) return failure(); - duplicates.set(index); + duplicateArguments.set(index); + duplicateResults.set(index); } - if (duplicates.none()) + if (duplicateArguments.none()) return failure(); - auto canonicalResult = [&](Value result) { - if (auto argument = dyn_cast(result); - argument && argument.getOwner() == &body) - return Value( - function.getArgument(representatives[argument.getArgNumber()])); - return result; - }; for (auto [index, representative] : llvm::enumerate(representatives)) - if (canonicalResult(returnOp->getOperand(index)) != - canonicalResult(returnOp->getOperand(representative))) + if (canonicalResult(returnOp->getOperand(index), body, representatives) != + canonicalResult(returnOp->getOperand(representative), body, + representatives)) return failure(); - Type type = function.getTypeWithoutArgsAndResults(duplicates, duplicates); - if (!type) + Type withoutArguments = function.getTypeWithoutArgs(duplicateArguments); + Type type = function.getTypeWithoutArgsAndResults(duplicateArguments, + duplicateResults); + if (!withoutArguments || !type) return failure(); SymbolTable symbols(scope); @@ -2390,12 +2440,14 @@ class DeduplicateXLAWrapperInputs final auto specialized = cast(clone); Block &specializedBody = specialized.getFunctionBody().front(); for (auto [index, representative] : llvm::enumerate(representatives)) - if (duplicates.test(index)) + if (duplicateArguments.test(index)) specializedBody.getArgument(index).replaceAllUsesWith( specializedBody.getArgument(representative)); - specializedBody.getTerminator()->eraseOperands(duplicates); - specializedBody.eraseArguments(duplicates); - specialized.setType(type); + specializedBody.getTerminator()->eraseOperands(duplicateResults); + function_interface_impl::eraseFunctionArguments( + specialized, duplicateArguments, withoutArguments); + function_interface_impl::eraseFunctionResults(specialized, duplicateResults, + type); SymbolTable::setSymbolName(clone, name); SymbolTable::setSymbolVisibility(clone, SymbolTable::Visibility::Private); { @@ -2404,10 +2456,10 @@ class DeduplicateXLAWrapperInputs final rewriter.insert(clone); } - rewriter.modifyOpInPlace(op, [&] { - op->eraseOperands(duplicates); - op.setFnAttr(FlatSymbolRefAttr::get(op.getContext(), name)); - }); + rewriter.startOpModification(op); + removeWrapperInputs(op, duplicateArguments, duplicateResults); + op.setFnAttr(FlatSymbolRefAttr::get(op.getContext(), name)); + rewriter.finalizeOpModification(op); return success(); } }; diff --git a/test/lit_tests/xla_wrapper_canonicalize.mlir b/test/lit_tests/xla_wrapper_canonicalize.mlir index d176ccc032..5649a8e5b4 100644 --- a/test/lit_tests/xla_wrapper_canonicalize.mlir +++ b/test/lit_tests/xla_wrapper_canonicalize.mlir @@ -162,12 +162,12 @@ module { // Remove the middle buffer: no operation reads or changes spare. // Before: wrapper @scale_data(data, spare, scale) // scale_data(data, spare, scale) -> (data * scale, spare, scale) -// After: wrapper @scale_data_without_unused(data, scale) -// scale_data_without_unused(data, scale) -> (data * scale, scale) +// After: wrapper @scale_data(data, scale) +// scale_data(data, scale) -> (data * scale, scale) // // Keep scale: the multiplication reads it, even though it is returned unchanged. -// Use one clone for all three compatible wrappers in the two host functions. -// The wrapper with metadata needs the original signature, so cloning is required. +// Update all four wrappers and the private function. Remove the attributes +// for spare from the wrapper that has argument attributes. module { func.func @drop_unused_middle_buffer(%data: memref<4xf32>, %spare: memref<4xf32>, @@ -202,33 +202,24 @@ module { // CHECK-SAME: %[[DATA:[^:]+]]: memref<4xf32>, // CHECK-SAME: %[[SPARE:[^:]+]]: memref<4xf32>, // CHECK-SAME: %[[SCALE:[^:]+]]: memref<4xf32>) { -// CHECK-NEXT: enzymexla.xla_wrapper @scale_data_without_unused (%[[DATA]], %[[SCALE]]) : (memref<4xf32>, memref<4xf32>) -> () +// CHECK-NEXT: enzymexla.xla_wrapper @scale_data (%[[DATA]], %[[SCALE]]) : (memref<4xf32>, memref<4xf32>) -> () // CHECK-NEXT: return // CHECK-NEXT: } -// These wrappers share the same clone. The wrapper with metadata stays intact. +// All wrappers now use the same reduced function. // CHECK-LABEL: func.func @share_reduced_function( // CHECK-SAME: %[[DATA:[^:]+]]: memref<4xf32>, // CHECK-SAME: %[[SPARE:[^:]+]]: memref<4xf32>, // CHECK-SAME: %[[SCALE:[^:]+]]: memref<4xf32>) { -// CHECK-NEXT: enzymexla.xla_wrapper @scale_data_without_unused (%[[DATA]], %[[SCALE]]) : (memref<4xf32>, memref<4xf32>) -> () -// CHECK-NEXT: enzymexla.xla_wrapper @scale_data_without_unused (%[[DATA]], %[[SCALE]]) : (memref<4xf32>, memref<4xf32>) -> () -// CHECK-NEXT: enzymexla.xla_wrapper @scale_data (%[[DATA]], %[[SPARE]], %[[SCALE]]) <{arg_attrs = [{}, {test.keep}, {}]}> : (memref<4xf32>, memref<4xf32>, memref<4xf32>) -> () +// CHECK-NEXT: enzymexla.xla_wrapper @scale_data (%[[DATA]], %[[SCALE]]) : (memref<4xf32>, memref<4xf32>) -> () +// CHECK-NEXT: enzymexla.xla_wrapper @scale_data (%[[DATA]], %[[SCALE]]) : (memref<4xf32>, memref<4xf32>) -> () +// CHECK-NEXT: enzymexla.xla_wrapper @scale_data (%[[DATA]], %[[SCALE]]) : (memref<4xf32>, memref<4xf32>) -> () // CHECK-NEXT: return // CHECK-NEXT: } -// The original function retains the spare argument and its matching result. +// The private function has two arguments and two results. // CHECK-LABEL: func.func private @scale_data( // CHECK-SAME: %[[DATA:[^:]+]]: tensor<4xf32>, -// CHECK-SAME: %[[SPARE:[^:]+]]: tensor<4xf32>, -// CHECK-SAME: %[[SCALE:[^:]+]]: tensor<4xf32>) -> (tensor<4xf32>, tensor<4xf32>, tensor<4xf32>) { -// CHECK-NEXT: %[[SCALED:[^ ]+]] = stablehlo.multiply %[[DATA]], %[[SCALE]] : tensor<4xf32> -// CHECK-NEXT: return %[[SCALED]], %[[SPARE]], %[[SCALE]] : tensor<4xf32>, tensor<4xf32>, tensor<4xf32> -// CHECK-NEXT: } - -// The clone has two arguments and two results. The multiplication is unchanged. -// CHECK-LABEL: func.func private @scale_data_without_unused( -// CHECK-SAME: %[[DATA:[^:]+]]: tensor<4xf32>, // CHECK-SAME: %[[SCALE:[^:]+]]: tensor<4xf32>) -> (tensor<4xf32>, tensor<4xf32>) { // CHECK-NEXT: %[[SCALED:[^ ]+]] = stablehlo.multiply %[[DATA]], %[[SCALE]] : tensor<4xf32> // CHECK-NEXT: return %[[SCALED]], %[[SCALE]] : tensor<4xf32>, tensor<4xf32> @@ -321,10 +312,10 @@ module { // ----- -// Keep these calls unchanged. Removing a position would also require changes -// to argument/result attributes or the specialized scalar input mapping. +// Remove the only buffer and its attributes from each call and function. +// The specialized function does not use its scalar input, so remove it too. module { - func.func @keep_metadata(%buffer: memref<4xf32>, %bound: i64) { + func.func @remove_unused_metadata(%buffer: memref<4xf32>, %bound: i64) { enzymexla.xla_wrapper @identity (%buffer) <{arg_attrs = [{test.keep}]}> : (memref<4xf32>) -> () enzymexla.xla_wrapper @identity (%buffer) <{res_attrs = [{test.keep}]}> : (memref<4xf32>) -> () enzymexla.xla_wrapper @arg_metadata (%buffer) : (memref<4xf32>) -> () @@ -346,13 +337,13 @@ module { } } -// CHECK-LABEL: func.func @keep_metadata( +// CHECK-LABEL: func.func @remove_unused_metadata( // CHECK-SAME: %[[BUFFER:[^:]+]]: memref<4xf32>, %[[BOUND:[^:]+]]: i64) { -// CHECK-NEXT: enzymexla.xla_wrapper @identity (%[[BUFFER]]) <{arg_attrs = [{test.keep}]}> : (memref<4xf32>) -> () -// CHECK-NEXT: enzymexla.xla_wrapper @identity (%[[BUFFER]]) <{res_attrs = [{test.keep}]}> : (memref<4xf32>) -> () -// CHECK-NEXT: enzymexla.xla_wrapper @arg_metadata (%[[BUFFER]]) : (memref<4xf32>) -> () -// CHECK-NEXT: enzymexla.xla_wrapper @result_metadata (%[[BUFFER]]) : (memref<4xf32>) -> () -// CHECK-NEXT: enzymexla.xla_wrapper @specialized (%[[BUFFER]], %[[BOUND]]) <{num_specialized = 1 : i64}> : (memref<4xf32>, i64) -> () +// CHECK-NEXT: enzymexla.xla_wrapper @identity () : () -> () +// CHECK-NEXT: enzymexla.xla_wrapper @identity () : () -> () +// CHECK-NEXT: enzymexla.xla_wrapper @arg_metadata () : () -> () +// CHECK-NEXT: enzymexla.xla_wrapper @result_metadata () : () -> () +// CHECK-NEXT: enzymexla.xla_wrapper @specialized () : () -> () // CHECK-NEXT: return // CHECK-NEXT: } diff --git a/test/lit_tests/xla_wrapper_canonicalize_specialized.mlir b/test/lit_tests/xla_wrapper_canonicalize_specialized.mlir new file mode 100644 index 0000000000..1907386675 --- /dev/null +++ b/test/lit_tests/xla_wrapper_canonicalize_specialized.mlir @@ -0,0 +1,141 @@ +// RUN: enzymexlamlir-opt %s --split-input-file --canonicalize --symbol-dce | FileCheck %s + +// Remove the unused buffer and the first and last specialized scalars. +// Keep data and increment, in that order. Their attributes must follow them. +// Both wrappers can use the reduced private function, although their attributes +// differ. The function has one result because increment is a scalar input. +// Before: update(spare, data, unused, increment, unused_last) -> (spare, data + increment) +// After: update(data, increment) -> data + increment +module { + func.func @remove_unused_specialized(%spare: memref, %data: memref, + %unused: i32, %increment: i32, + %unused_last: i32) { + enzymexla.xla_wrapper @update (%spare, %data, %unused, %increment, %unused_last) + <{arg_attrs = [{test.drop}, {test.data}, {test.drop}, {test.increment}, {test.drop}], + res_attrs = [{test.drop}, {test.result}], num_specialized = 3 : i64}> : + (memref, memref, i32, i32, i32) -> () + enzymexla.xla_wrapper @update (%spare, %data, %unused, %increment, %unused_last) + <{arg_attrs = [{test.drop}, {test.other_data}, {test.drop}, {test.other_increment}, {test.drop}], + res_attrs = [{test.drop}, {test.other_result}], num_specialized = 3 : i64}> : + (memref, memref, i32, i32, i32) -> () + return + } + func.func private @update(%spare: tensor {test.drop}, + %data: tensor {test.function_data}, + %unused: tensor {test.drop}, + %increment: tensor {test.function_increment}, + %unused_last: tensor {test.drop}) + -> (tensor {test.drop}, tensor {test.function_result}) { + %sum = stablehlo.add %data, %increment : tensor + return %spare, %sum : tensor, tensor + } +} + +// CHECK-LABEL: func.func @remove_unused_specialized( +// CHECK-SAME: %[[SPARE:[^:]+]]: memref, %[[DATA:[^:]+]]: memref, %[[UNUSED:[^:]+]]: i32, %[[INC:[^:]+]]: i32, +// CHECK-NEXT: enzymexla.xla_wrapper @update (%[[DATA]], %[[INC]]) : (memref, i32) -> () +// CHECK-NEXT: enzymexla.xla_wrapper @update (%[[DATA]], %[[INC]]) : (memref, i32) -> () +// CHECK-LABEL: func.func private @update( +// CHECK-SAME: %[[DATA:[^:]+]]: tensor {test.function_data}, %[[INC:[^:]+]]: tensor {test.function_increment}) -> (tensor {test.function_result}) { +// CHECK-NEXT: %[[SUM:[^ ]+]] = stablehlo.add %[[DATA]], %[[INC]] : tensor +// CHECK-NEXT: return %[[SUM]] : tensor + +// ----- + +// Remove a specialized scalar without removing a buffer. The public function +// can have unknown callers, so keep it and give the wrapper a reduced clone. +module { + func.func @remove_only_scalar(%data: memref, %unused: i32) { + enzymexla.xla_wrapper @negate (%data, %unused) + <{num_specialized = 1 : i64}> : (memref, i32) -> () + return + } + func.func @negate(%data: tensor, %unused: tensor) -> tensor { + %negative = stablehlo.negate %data : tensor + return %negative : tensor + } +} + +// CHECK-LABEL: func.func @remove_only_scalar( +// CHECK-SAME: %[[DATA:[^:]+]]: memref, +// CHECK-NEXT: enzymexla.xla_wrapper @negate_without_unused (%[[DATA]]) : (memref) -> () +// CHECK-LABEL: func.func @negate( +// CHECK-SAME: tensor, {{.*}}: tensor) -> tensor +// CHECK-LABEL: func.func private @negate_without_unused( +// CHECK-SAME: %[[DATA:[^:]+]]: tensor) -> tensor { +// CHECK-NEXT: %[[NEG:[^ ]+]] = stablehlo.negate %[[DATA]] : tensor +// CHECK-NEXT: return %[[NEG]] : tensor + +// ----- + +// Both buffer positions write the same specialized scalar to the same buffer. +// Remove the second buffer and its attributes. Keep both scalar positions and +// their attributes, even though the wrapper passes the same scalar to both. +// The first scalar supplies the result; the side effect reads the second. +// Before: write(data, data, value, value) -> (value, value) +// After: write_without_duplicates(data, value, value) -> value +module { + func.func @deduplicate_with_specialized(%data: memref, %value: i32) { + enzymexla.xla_wrapper @write (%data, %data, %value, %value) + <{arg_attrs = [{test.buffer}, {test.buffer}, {test.value}, {test.effect}], + res_attrs = [{test.result}, {test.result}], num_specialized = 2 : i64}> : + (memref, memref, i32, i32) -> () + return + } + func.func private @write(%first: tensor {test.function_buffer}, + %second: tensor {test.function_buffer}, + %value: tensor {test.function_value}, + %effect: tensor {test.function_effect}) + -> (tensor {test.function_result}, tensor {test.function_result}) { + stablehlo.custom_call @record(%effect) {has_side_effect = true} : (tensor) -> () + return %value, %value : tensor, tensor + } +} + +// CHECK-LABEL: func.func @deduplicate_with_specialized( +// CHECK-SAME: %[[DATA:[^:]+]]: memref, %[[VALUE:[^:]+]]: i32) { +// CHECK-NEXT: enzymexla.xla_wrapper @write_without_duplicates (%[[DATA]], %[[VALUE]], %[[VALUE]]) : (memref, i32, i32) -> () +// CHECK-LABEL: func.func private @write_without_duplicates( +// CHECK-SAME: %[[DATA:[^:]+]]: tensor {test.function_buffer}, %[[VALUE:[^:]+]]: tensor {test.function_value}, %[[EFFECT:[^:]+]]: tensor {test.function_effect}) -> (tensor {test.function_result}) { +// CHECK-NEXT: stablehlo.custom_call @record(%[[EFFECT]]) {has_side_effect = true} : (tensor) -> () +// CHECK-NEXT: return %[[VALUE]] : tensor + +// ----- + +// Keep both buffer positions when their attributes differ. Do not choose one +// position's attributes and discard the other position's attributes. +module { + func.func @keep_different_attributes(%data: memref) { + enzymexla.xla_wrapper @pair (%data, %data) + <{arg_attrs = [{test.first}, {test.second}]}> : (memref, memref) -> () + enzymexla.xla_wrapper @pair (%data, %data) + <{res_attrs = [{test.first}, {test.second}]}> : (memref, memref) -> () + enzymexla.xla_wrapper @arg_attributes (%data, %data) : (memref, memref) -> () + enzymexla.xla_wrapper @result_attributes (%data, %data) : (memref, memref) -> () + return + } + func.func private @pair(%a: tensor, %b: tensor) + -> (tensor, tensor) { + %sum = stablehlo.add %a, %b : tensor + return %sum, %sum : tensor, tensor + } + func.func private @arg_attributes(%a: tensor {test.first}, + %b: tensor {test.second}) + -> (tensor, tensor) { + %sum = stablehlo.add %a, %b : tensor + return %sum, %sum : tensor, tensor + } + func.func private @result_attributes(%a: tensor, %b: tensor) + -> (tensor {test.first}, tensor {test.second}) { + %sum = stablehlo.add %a, %b : tensor + return %sum, %sum : tensor, tensor + } +} + +// CHECK-LABEL: func.func @keep_different_attributes( +// CHECK-SAME: %[[DATA:[^:]+]]: memref) { +// CHECK-NEXT: enzymexla.xla_wrapper @pair (%[[DATA]], %[[DATA]]) : (memref, memref) -> () +// CHECK-NEXT: enzymexla.xla_wrapper @pair (%[[DATA]], %[[DATA]]) : (memref, memref) -> () +// CHECK-NEXT: enzymexla.xla_wrapper @arg_attributes (%[[DATA]], %[[DATA]]) : (memref, memref) -> () +// CHECK-NEXT: enzymexla.xla_wrapper @result_attributes (%[[DATA]], %[[DATA]]) : (memref, memref) -> () +// CHECK-NEXT: return From 59a440e1a04c5c68380ce440bc57e42b6fe3338a Mon Sep 17 00:00:00 2001 From: Vimarsh Sathia Date: Tue, 6 Oct 2026 00:00:58 -0500 Subject: [PATCH 5/5] Deduplicate specialized XLA wrapper inputs Merge repeated specialized scalar arguments when their types and attributes agree. Keep buffer results and update num_specialized. Keep the original function for callers that pass distinct values. Add direct tests for specialized input removal and deduplication. Check scalar order, distinct callers, combined buffer and scalar deduplication, and the result after wrapper fusion. Validation: optimized build passed; MLIR LIT suite had 1,424 passes and 9 expected failures. Assisted-By: OpenAI Codex --- src/enzyme_ad/jax/Dialect/Ops.cpp | 43 ++++--- .../xla_megakernelize_specialized.mlir | 13 +- .../xla_wrapper_canonicalize_specialized.mlir | 121 +++++++++++++++++- 3 files changed, 154 insertions(+), 23 deletions(-) diff --git a/src/enzyme_ad/jax/Dialect/Ops.cpp b/src/enzyme_ad/jax/Dialect/Ops.cpp index 6dbf314119..910a8e6843 100644 --- a/src/enzyme_ad/jax/Dialect/Ops.cpp +++ b/src/enzyme_ad/jax/Dialect/Ops.cpp @@ -2351,15 +2351,14 @@ class RemoveUnusedXLAWrapperInputs final }; /// Merge equal buffer inputs only when their returned values and attributes -/// also agree. Keep specialized scalar inputs at the end of the input list. -/// Clone the function because other calls can pass distinct buffers. +/// also agree. Merge equal specialized scalar inputs if their attributes agree. +/// Clone the function because other calls can pass distinct inputs. class DeduplicateXLAWrapperInputs final : public OpRewritePattern { static Value canonicalResult(Value result, Block &body, ArrayRef representatives) { if (auto argument = dyn_cast(result); - argument && argument.getOwner() == &body && - argument.getArgNumber() < representatives.size()) + argument && argument.getOwner() == &body) return body.getArgument(representatives[argument.getArgNumber()]); return result; } @@ -2395,30 +2394,44 @@ class DeduplicateXLAWrapperInputs final SmallVector representatives; llvm::BitVector duplicateArguments(function.getNumArguments()); llvm::BitVector duplicateResults(numBuffers); - for (auto [index, input] : - llvm::enumerate(op.getInputs().take_front(numBuffers))) { - auto [entry, inserted] = identities.try_emplace( - enzyme::oputils::getBaseObject(input, /*offsetAllowed=*/false), - index); + for (auto [index, input] : llvm::enumerate(op.getInputs())) { + bool isBuffer = index < numBuffers; + // Match scalars by SSA value. Keep buffer addresses separate. + if (index == numBuffers) + identities.clear(); + Value identity = + isBuffer + ? enzyme::oputils::getBaseObject(input, /*offsetAllowed=*/false) + : input; + auto [entry, inserted] = identities.try_emplace(identity, index); unsigned representative = entry->second; - representatives.push_back(representative); + representatives.push_back(index); if (inserted) continue; if (input.getType() != op.getInputs()[representative].getType() || function.getArgumentTypes()[index] != function.getArgumentTypes()[representative] || !sameAttributes(op.getArgAttrsAttr(), index, representative) || - !sameAttributes(op.getResAttrsAttr(), index, representative) || - !sameAttributes(function.getAllArgAttrs(), index, representative) || - !sameAttributes(function.getAllResultAttrs(), index, representative)) + !sameAttributes(function.getAllArgAttrs(), index, representative)) { + if (isBuffer) + return failure(); + continue; + } + if (isBuffer && + (!sameAttributes(op.getResAttrsAttr(), index, representative) || + !sameAttributes(function.getAllResultAttrs(), index, + representative))) return failure(); + representatives.back() = representative; duplicateArguments.set(index); - duplicateResults.set(index); + if (isBuffer) + duplicateResults.set(index); } if (duplicateArguments.none()) return failure(); - for (auto [index, representative] : llvm::enumerate(representatives)) + for (auto [index, representative] : + llvm::enumerate(ArrayRef(representatives).take_front(numBuffers))) if (canonicalResult(returnOp->getOperand(index), body, representatives) != canonicalResult(returnOp->getOperand(representative), body, representatives)) diff --git a/test/lit_tests/xla_megakernelize_specialized.mlir b/test/lit_tests/xla_megakernelize_specialized.mlir index 53ab0175e0..293d9b495e 100644 --- a/test/lit_tests/xla_megakernelize_specialized.mlir +++ b/test/lit_tests/xla_megakernelize_specialized.mlir @@ -61,8 +61,9 @@ module { // ----- // first(A, n), second(B, A, m, k) becomes first(A, B, n, m, k). -// Here n and m receive the same host SSA value, but remain separate specialized -// arguments. Only A and B have results, in the union's buffer order. +// Here n and m receive the same host SSA value. Canonicalization merges those +// arguments, so the final call is first_without_duplicates(A, B, n, k). +// Only A and B have results, in the union's buffer order. module { func.func @different_specialized_counts(%a: memref, %b: memref, %n: i32, %k: i64) { @@ -85,12 +86,12 @@ module { } // CHECK-LABEL: func.func @different_specialized_counts( // CHECK-SAME: %[[HOST_A:.*]]: memref, %[[HOST_B:.*]]: memref, %[[HOST_N:.*]]: i32, %[[HOST_K:.*]]: i64) -// CHECK-NEXT: enzymexla.xla_wrapper @first (%[[HOST_A]], %[[HOST_B]], %[[HOST_N]], %[[HOST_N]], %[[HOST_K]]) : (memref, memref, i32, i32, i64) -> () +// CHECK-NEXT: enzymexla.xla_wrapper @first_without_duplicates (%[[HOST_A]], %[[HOST_B]], %[[HOST_N]], %[[HOST_K]]) : (memref, memref, i32, i64) -> () // CHECK-NEXT: return -// CHECK-LABEL: func.func private @first( -// CHECK-SAME: %[[A:.*]]: tensor, %[[B:.*]]: tensor, %[[N:.*]]: tensor, %[[M:.*]]: tensor, %[[K:.*]]: tensor) -> (tensor, tensor) +// CHECK-LABEL: func.func private @first_without_duplicates( +// CHECK-SAME: %[[A:.*]]: tensor, %[[B:.*]]: tensor, %[[N:.*]]: tensor, %[[K:.*]]: tensor) -> (tensor, tensor) // CHECK-NEXT: %[[FIRST_SUM:.*]] = stablehlo.add %[[A]], %[[N]] : tensor -// CHECK-NEXT: %[[SCALED:.*]] = stablehlo.multiply %[[B]], %[[M]] : tensor +// CHECK-NEXT: %[[SCALED:.*]] = stablehlo.multiply %[[B]], %[[N]] : tensor // CHECK-NEXT: %[[K32:.*]] = stablehlo.convert %[[K]] : (tensor) -> tensor // CHECK-NEXT: %[[SECOND_SUM:.*]] = stablehlo.add %[[FIRST_SUM]], %[[K32]] : tensor // CHECK-NEXT: return %[[SECOND_SUM]], %[[SCALED]] : tensor, tensor diff --git a/test/lit_tests/xla_wrapper_canonicalize_specialized.mlir b/test/lit_tests/xla_wrapper_canonicalize_specialized.mlir index 1907386675..9a1c9f654d 100644 --- a/test/lit_tests/xla_wrapper_canonicalize_specialized.mlir +++ b/test/lit_tests/xla_wrapper_canonicalize_specialized.mlir @@ -68,9 +68,126 @@ module { // ----- +// Only the middle input is unused. It is a specialized scalar, so remove its +// argument without removing a buffer result. Keep increment at the end. +// Before: add_increment(data, unused, increment) -> data + increment +// After: add_increment(data, increment) -> data + increment +// num_specialized changes from 2 to 1. There is still one buffer result. +module { + func.func @remove_one_specialized_scalar(%data: memref, %unused: i32, + %increment: i32) { + enzymexla.xla_wrapper @add_increment (%data, %unused, %increment) + <{num_specialized = 2 : i64}> : (memref, i32, i32) -> () + return + } + func.func private @add_increment(%data: tensor, %unused: tensor, + %increment: tensor) -> tensor { + %sum = stablehlo.add %data, %increment : tensor + return %sum : tensor + } +} + +// CHECK-LABEL: func.func @remove_one_specialized_scalar( +// CHECK-SAME: %[[DATA:[^:]+]]: memref, %[[UNUSED:[^:]+]]: i32, %[[INC:[^:]+]]: i32) { +// CHECK-NEXT: enzymexla.xla_wrapper @add_increment (%[[DATA]], %[[INC]]) : (memref, i32) -> () +// CHECK-NEXT: return +// CHECK-LABEL: func.func private @add_increment( +// CHECK-SAME: %[[DATA:[^:]+]]: tensor, %[[INC:[^:]+]]: tensor) -> tensor { +// CHECK-NEXT: %[[SUM:[^ ]+]] = stablehlo.add %[[DATA]], %[[INC]] : tensor +// CHECK-NEXT: return %[[SUM]] : tensor + +// ----- + +// Both increment arguments are used, but the wrapper passes the same scalar +// to both. Merge those arguments and use the remaining argument twice. +// Keep the buffer and the later scale argument in their original order. +// Before: advance(data, increment, increment, scale) +// After: advance_without_duplicates(data, increment, scale) +// Both compute (data + increment + increment) * scale. +// num_specialized changes from 3 to 2. There is still one buffer result. +// A second caller passes distinct scalars. Keep that call and the original +// function so that it still computes (data + first + second) * scale. +module { + func.func @deduplicate_specialized_scalar(%data: memref, %increment: i32, + %scale: i32) { + enzymexla.xla_wrapper @advance (%data, %increment, %increment, %scale) + <{num_specialized = 3 : i64}> : (memref, i32, i32, i32) -> () + return + } + func.func @keep_distinct_specialized_scalars(%data: memref, %first: i32, + %second: i32, %scale: i32) { + enzymexla.xla_wrapper @advance (%data, %first, %second, %scale) + <{num_specialized = 3 : i64}> : (memref, i32, i32, i32) -> () + return + } + func.func private @advance(%data: tensor, %first: tensor, + %second: tensor, %scale: tensor) + -> tensor { + %once = stablehlo.add %data, %first : tensor + %twice = stablehlo.add %once, %second : tensor + %scaled = stablehlo.multiply %twice, %scale : tensor + return %scaled : tensor + } +} + +// CHECK-LABEL: func.func @deduplicate_specialized_scalar( +// CHECK-SAME: %[[DATA:[^:]+]]: memref, %[[INC:[^:]+]]: i32, %[[SCALE:[^:]+]]: i32) { +// CHECK-NEXT: enzymexla.xla_wrapper @advance_without_duplicates (%[[DATA]], %[[INC]], %[[SCALE]]) : (memref, i32, i32) -> () +// CHECK-NEXT: return +// CHECK-LABEL: func.func @keep_distinct_specialized_scalars( +// CHECK-SAME: %[[DATA:[^:]+]]: memref, %[[FIRST:[^:]+]]: i32, %[[SECOND:[^:]+]]: i32, %[[SCALE:[^:]+]]: i32) { +// CHECK-NEXT: enzymexla.xla_wrapper @advance (%[[DATA]], %[[FIRST]], %[[SECOND]], %[[SCALE]]) : (memref, i32, i32, i32) -> () +// CHECK-NEXT: return +// CHECK-LABEL: func.func private @advance( +// CHECK-SAME: %[[DATA:[^:]+]]: tensor, %[[FIRST:[^:]+]]: tensor, %[[SECOND:[^:]+]]: tensor, %[[SCALE:[^:]+]]: tensor) -> tensor { +// CHECK-NEXT: %[[ONCE:[^ ]+]] = stablehlo.add %[[DATA]], %[[FIRST]] : tensor +// CHECK-NEXT: %[[TWICE:[^ ]+]] = stablehlo.add %[[ONCE]], %[[SECOND]] : tensor +// CHECK-NEXT: %[[SCALED:[^ ]+]] = stablehlo.multiply %[[TWICE]], %[[SCALE]] : tensor +// CHECK-NEXT: return %[[SCALED]] : tensor +// CHECK-LABEL: func.func private @advance_without_duplicates( +// CHECK-SAME: %[[DATA:[^:]+]]: tensor, %[[INC:[^:]+]]: tensor, %[[SCALE:[^:]+]]: tensor) -> tensor { +// CHECK-NEXT: %[[ONCE:[^ ]+]] = stablehlo.add %[[DATA]], %[[INC]] : tensor +// CHECK-NEXT: %[[TWICE:[^ ]+]] = stablehlo.add %[[ONCE]], %[[INC]] : tensor +// CHECK-NEXT: %[[SCALED:[^ ]+]] = stablehlo.multiply %[[TWICE]], %[[SCALE]] : tensor +// CHECK-NEXT: return %[[SCALED]] : tensor + +// ----- + +// Both buffer positions refer to data. Their results are different scalar +// arguments, but the wrapper supplies the same value for those arguments. +// Merge the scalar arguments first when comparing the buffer results. Then +// both results agree, so one buffer input and one buffer result can be removed. +// Before: write_pair(data, data, value, value) -> (value, value) +// After: write_pair_without_duplicates(data, value) -> value +// num_specialized changes from 2 to 1. Two buffer results become one. +module { + func.func @deduplicate_buffers_and_scalars(%data: memref, %value: i32) { + enzymexla.xla_wrapper @write_pair (%data, %data, %value, %value) + <{num_specialized = 2 : i64}> : (memref, memref, i32, i32) -> () + return + } + func.func private @write_pair(%first_buffer: tensor, + %second_buffer: tensor, + %first_value: tensor, + %second_value: tensor) + -> (tensor, tensor) { + return %first_value, %second_value : tensor, tensor + } +} + +// CHECK-LABEL: func.func @deduplicate_buffers_and_scalars( +// CHECK-SAME: %[[DATA:[^:]+]]: memref, %[[VALUE:[^:]+]]: i32) { +// CHECK-NEXT: enzymexla.xla_wrapper @write_pair_without_duplicates (%[[DATA]], %[[VALUE]]) : (memref, i32) -> () +// CHECK-NEXT: return +// CHECK-LABEL: func.func private @write_pair_without_duplicates( +// CHECK-SAME: %[[DATA:[^:]+]]: tensor, %[[VALUE:[^:]+]]: tensor) -> tensor { +// CHECK-NEXT: return %[[VALUE]] : tensor + +// ----- + // Both buffer positions write the same specialized scalar to the same buffer. -// Remove the second buffer and its attributes. Keep both scalar positions and -// their attributes, even though the wrapper passes the same scalar to both. +// Remove the second buffer and its attributes. The scalar positions have +// different attributes, so keep both even though their values are equal. // The first scalar supplies the result; the side effect reads the second. // Before: write(data, data, value, value) -> (value, value) // After: write_without_duplicates(data, value, value) -> value