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<MemoryEffectsOpInterface>
 ]> {
   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<AnyType>:$inputs,
       OptionalAttr<DictArrayAttr>:$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..910a8e6843 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"
@@ -2200,6 +2201,288 @@ void XLAWrapperOp::getEffects(
   effects.emplace_back(MemoryEffects::Effect::get<MemoryEffects::Write>());
 }
 
+static bool hasXLAWrapperSignature(XLAWrapperOp wrapper,
+                                   FunctionOpInterface function) {
+  int64_t numSpecialized = wrapper.getNumSpecialized();
+  if (numSpecialized < 0 ||
+      numSpecialized > static_cast<int64_t>(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<Attribute> 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<XLAWrapperOp> {
+public:
+  using OpRewritePattern<XLAWrapperOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(XLAWrapperOp op,
+                                PatternRewriter &rewriter) const override {
+    auto function = dyn_cast_or_null<FunctionOpInterface>(
+        SymbolTable::lookupNearestSymbolFrom(op, op.getFnAttr()));
+    if (!function || !function.getFunctionBody().hasOneBlock() ||
+        !hasXLAWrapperSignature(op, function))
+      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();
+    unsigned numBuffers = function.getNumResults();
+    if (!returnOp->hasTrait<OpTrait::ReturnLike>() ||
+        returnOp->getNumOperands() != numBuffers)
+      return failure();
+
+    llvm::BitVector unusedArguments(function.getNumArguments());
+    llvm::BitVector unusedResults(numBuffers);
+    for (BlockArgument argument : body.getArguments()) {
+      unsigned index = argument.getArgNumber();
+      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 (unusedArguments.none())
+      return failure();
+
+    Type withoutArguments = function.getTypeWithoutArgs(unusedArguments);
+    Type type =
+        function.getTypeWithoutArgsAndResults(unusedArguments, unusedResults);
+    if (!withoutArguments || !type)
+      return failure();
+
+    llvm::SmallSetVector<Operation *, 4> 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<XLAWrapperOp>(use.getUser());
+        if (wrapper && wrapper.getFnAttr() == op.getFnAttr() &&
+            hasXLAWrapperSignature(wrapper, function) &&
+            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<FunctionOpInterface>(clone);
+      SymbolTable::setSymbolName(clone, name);
+      SymbolTable::setSymbolVisibility(clone, SymbolTable::Visibility::Private);
+      OpBuilder::InsertionGuard guard(rewriter);
+      rewriter.setInsertionPointAfter(function);
+      rewriter.insert(clone);
+    }
+
+    rewriter.startOpModification(specialized);
+    Operation *specializedReturn =
+        specialized.getFunctionBody().front().getTerminator();
+    rewriter.startOpModification(specializedReturn);
+    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<XLAWrapperOp>(caller);
+      rewriter.startOpModification(wrapper);
+      removeWrapperInputs(wrapper, unusedArguments, unusedResults);
+      wrapper.setFnAttr(FlatSymbolRefAttr::get(op.getContext(), name));
+      rewriter.finalizeOpModification(wrapper);
+    }
+    return success();
+  }
+};
+
+/// Merge equal buffer inputs only when their returned values and attributes
+/// 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<XLAWrapperOp> {
+  static Value canonicalResult(Value result, Block &body,
+                               ArrayRef<unsigned> representatives) {
+    if (auto argument = dyn_cast<BlockArgument>(result);
+        argument && argument.getOwner() == &body)
+      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<XLAWrapperOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(XLAWrapperOp op,
+                                PatternRewriter &rewriter) const override {
+    auto function = dyn_cast_or_null<FunctionOpInterface>(
+        SymbolTable::lookupNearestSymbolFrom(op, op.getFnAttr()));
+    if (!function || !function.getFunctionBody().hasOneBlock() ||
+        !hasXLAWrapperSignature(op, function))
+      return failure();
+
+    Operation *scope = SymbolTable::getNearestSymbolTable(op);
+    if (function->getParentOp() != scope)
+      return failure();
+
+    Block &body = function.getFunctionBody().front();
+    Operation *returnOp = body.getTerminator();
+    unsigned numBuffers = function.getNumResults();
+    if (!returnOp->hasTrait<OpTrait::ReturnLike>() ||
+        returnOp->getNumOperands() != numBuffers)
+      return failure();
+
+    DenseMap<Value, unsigned> identities;
+    SmallVector<unsigned> representatives;
+    llvm::BitVector duplicateArguments(function.getNumArguments());
+    llvm::BitVector duplicateResults(numBuffers);
+    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(index);
+      if (inserted)
+        continue;
+      if (input.getType() != op.getInputs()[representative].getType() ||
+          function.getArgumentTypes()[index] !=
+              function.getArgumentTypes()[representative] ||
+          !sameAttributes(op.getArgAttrsAttr(), 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);
+      if (isBuffer)
+        duplicateResults.set(index);
+    }
+    if (duplicateArguments.none())
+      return failure();
+
+    for (auto [index, representative] :
+         llvm::enumerate(ArrayRef(representatives).take_front(numBuffers)))
+      if (canonicalResult(returnOp->getOperand(index), body, representatives) !=
+          canonicalResult(returnOp->getOperand(representative), body,
+                          representatives))
+        return failure();
+
+    Type withoutArguments = function.getTypeWithoutArgs(duplicateArguments);
+    Type type = function.getTypeWithoutArgsAndResults(duplicateArguments,
+                                                      duplicateResults);
+    if (!withoutArguments || !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<FunctionOpInterface>(clone);
+    Block &specializedBody = specialized.getFunctionBody().front();
+    for (auto [index, representative] : llvm::enumerate(representatives))
+      if (duplicateArguments.test(index))
+        specializedBody.getArgument(index).replaceAllUsesWith(
+            specializedBody.getArgument(representative));
+    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);
+    {
+      OpBuilder::InsertionGuard guard(rewriter);
+      rewriter.setInsertionPointAfter(function);
+      rewriter.insert(clone);
+    }
+
+    rewriter.startOpModification(op);
+    removeWrapperInputs(op, duplicateArguments, duplicateResults);
+    op.setFnAttr(FlatSymbolRefAttr::get(op.getContext(), name));
+    rewriter.finalizeOpModification(op);
+    return success();
+  }
+};
+
+void XLAWrapperOp::getCanonicalizationPatterns(RewritePatternSet &results,
+                                               MLIRContext *context) {
+  results.insert<RemoveUnusedXLAWrapperInputs, DeduplicateXLAWrapperInputs>(
+      context);
+}
+
 //===----------------------------------------------------------------------===//
 // AlternativesOp
 //===----------------------------------------------------------------------===//
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<i32>, %b: memref<i32>,
                                           %n: i32, %k: i64) {
@@ -85,12 +86,12 @@ module {
 }
 // CHECK-LABEL: func.func @different_specialized_counts(
 // CHECK-SAME: %[[HOST_A:.*]]: memref<i32>, %[[HOST_B:.*]]: memref<i32>, %[[HOST_N:.*]]: i32, %[[HOST_K:.*]]: i64)
-// CHECK-NEXT: enzymexla.xla_wrapper @first (%[[HOST_A]], %[[HOST_B]], %[[HOST_N]], %[[HOST_N]], %[[HOST_K]]) <num_specialized = 3> : (memref<i32>, memref<i32>, i32, i32, i64) -> ()
+// CHECK-NEXT: enzymexla.xla_wrapper @first_without_duplicates (%[[HOST_A]], %[[HOST_B]], %[[HOST_N]], %[[HOST_K]]) <num_specialized = 2> : (memref<i32>, memref<i32>, i32, i64) -> ()
 // CHECK-NEXT: return
-// CHECK-LABEL: func.func private @first(
-// CHECK-SAME: %[[A:.*]]: tensor<i32>, %[[B:.*]]: tensor<i32>, %[[N:.*]]: tensor<i32>, %[[M:.*]]: tensor<i32>, %[[K:.*]]: tensor<i64>) -> (tensor<i32>, tensor<i32>)
+// CHECK-LABEL: func.func private @first_without_duplicates(
+// CHECK-SAME: %[[A:.*]]: tensor<i32>, %[[B:.*]]: tensor<i32>, %[[N:.*]]: tensor<i32>, %[[K:.*]]: tensor<i64>) -> (tensor<i32>, tensor<i32>)
 // CHECK-NEXT: %[[FIRST_SUM:.*]] = stablehlo.add %[[A]], %[[N]] : tensor<i32>
-// CHECK-NEXT: %[[SCALED:.*]] = stablehlo.multiply %[[B]], %[[M]] : tensor<i32>
+// CHECK-NEXT: %[[SCALED:.*]] = stablehlo.multiply %[[B]], %[[N]] : tensor<i32>
 // CHECK-NEXT: %[[K32:.*]] = stablehlo.convert %[[K]] : (tensor<i64>) -> tensor<i32>
 // CHECK-NEXT: %[[SECOND_SUM:.*]] = stablehlo.add %[[FIRST_SUM]], %[[K32]] : tensor<i32>
 // CHECK-NEXT: return %[[SECOND_SUM]], %[[SCALED]] : tensor<i32>, tensor<i32>
diff --git a/test/lit_tests/xla_wrapper_canonicalize.mlir b/test/lit_tests/xla_wrapper_canonicalize.mlir
new file mode 100644
index 0000000000..5649a8e5b4
--- /dev/null
+++ b/test/lit_tests/xla_wrapper_canonicalize.mlir
@@ -0,0 +1,499 @@
+// 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(data, scale)
+//         scale_data(data, scale) -> (data * scale, scale)
+//
+// Keep scale: the multiplication reads it, even though it is returned unchanged.
+// 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>,
+                                      %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 (%[[DATA]], %[[SCALE]]) : (memref<4xf32>, memref<4xf32>) -> ()
+// CHECK-NEXT: return
+// CHECK-NEXT: }
+
+// 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 (%[[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]]) <arg_attrs = [{}, {}]> : (memref<4xf32>, memref<4xf32>) -> ()
+// CHECK-NEXT: return
+// CHECK-NEXT: }
+
+// The private function has two arguments and two results.
+// CHECK-LABEL: func.func private @scale_data(
+// 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: }
+
+// -----
+
+// 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 @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>) -> ()
+    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<i64>) -> tensor<4xf32> {
+    return %a : tensor<4xf32>
+  }
+}
+
+// CHECK-LABEL: func.func @remove_unused_metadata(
+// CHECK-SAME: %[[BUFFER:[^:]+]]: memref<4xf32>, %[[BOUND:[^:]+]]: i64) {
+// CHECK-NEXT: enzymexla.xla_wrapper @identity () <arg_attrs = []> : () -> ()
+// CHECK-NEXT: enzymexla.xla_wrapper @identity () <res_attrs = []> : () -> ()
+// 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: }
+
+// -----
+
+// 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<?xf32>
+    %second = "enzymexla.pointer2memref"(%a) : (!llvm.ptr) -> memref<?xf32>
+    enzymexla.xla_wrapper @pair (%first, %second) :
+        (memref<?xf32>, memref<?xf32>) -> ()
+    llvm.return
+  }
+  llvm.func @keep_other_call(%a: !llvm.ptr, %b: !llvm.ptr) {
+    %first = "enzymexla.pointer2memref"(%a) : (!llvm.ptr) -> memref<?xf32>
+    %second = "enzymexla.pointer2memref"(%b) : (!llvm.ptr) -> memref<?xf32>
+    enzymexla.xla_wrapper @pair (%first, %second) :
+        (memref<?xf32>, memref<?xf32>) -> ()
+    llvm.return
+  }
+  func.func private @pair(%a: tensor<?xf32>, %b: tensor<?xf32>)
+      -> (tensor<?xf32>, tensor<?xf32>) {
+    %sum = stablehlo.add %a, %b : tensor<?xf32>
+    return %sum, %sum : tensor<?xf32>, tensor<?xf32>
+  }
+}
+
+// CHECK-LABEL: llvm.func @deduplicate_inputs(
+// CHECK-SAME: %[[A:[^:]+]]: !llvm.ptr) {
+// CHECK-NEXT: %[[BUFFER:[^ ]+]] = "enzymexla.pointer2memref"(%[[A]]) : (!llvm.ptr) -> memref<?xf32>
+// CHECK-NEXT: enzymexla.xla_wrapper @pair_without_duplicates (%[[BUFFER]]) : (memref<?xf32>) -> ()
+// 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<?xf32>
+// CHECK-NEXT: %[[SECOND:[^ ]+]] = "enzymexla.pointer2memref"(%[[B]]) : (!llvm.ptr) -> memref<?xf32>
+// CHECK-NEXT: enzymexla.xla_wrapper @pair (%[[FIRST]], %[[SECOND]]) : (memref<?xf32>, memref<?xf32>) -> ()
+// CHECK-NEXT: llvm.return
+// CHECK-NEXT: }
+
+// CHECK-LABEL: func.func private @pair(
+// CHECK-SAME: %[[A:[^:]+]]: tensor<?xf32>, %[[B:[^:]+]]: tensor<?xf32>) -> (tensor<?xf32>, tensor<?xf32>) {
+// CHECK-NEXT: %[[SUM:[^ ]+]] = stablehlo.add %[[A]], %[[B]] : tensor<?xf32>
+// CHECK-NEXT: return %[[SUM]], %[[SUM]] : tensor<?xf32>, tensor<?xf32>
+// CHECK-NEXT: }
+
+// CHECK-LABEL: func.func private @pair_without_duplicates(
+// CHECK-SAME: %[[A:[^:]+]]: tensor<?xf32>) -> tensor<?xf32> {
+// CHECK-NEXT: %[[SUM:[^ ]+]] = stablehlo.add %[[A]], %[[A]] : tensor<?xf32>
+// CHECK-NEXT: return %[[SUM]] : tensor<?xf32>
+// 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_specialized.mlir b/test/lit_tests/xla_wrapper_canonicalize_specialized.mlir
new file mode 100644
index 0000000000..9a1c9f654d
--- /dev/null
+++ b/test/lit_tests/xla_wrapper_canonicalize_specialized.mlir
@@ -0,0 +1,258 @@
+// 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<i32>, %data: memref<i32>,
+                                       %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<i32>, memref<i32>, 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<i32>, memref<i32>, i32, i32, i32) -> ()
+    return
+  }
+  func.func private @update(%spare: tensor<i32> {test.drop},
+                            %data: tensor<i32> {test.function_data},
+                            %unused: tensor<i32> {test.drop},
+                            %increment: tensor<i32> {test.function_increment},
+                            %unused_last: tensor<i32> {test.drop})
+      -> (tensor<i32> {test.drop}, tensor<i32> {test.function_result}) {
+    %sum = stablehlo.add %data, %increment : tensor<i32>
+    return %spare, %sum : tensor<i32>, tensor<i32>
+  }
+}
+
+// CHECK-LABEL: func.func @remove_unused_specialized(
+// CHECK-SAME: %[[SPARE:[^:]+]]: memref<i32>, %[[DATA:[^:]+]]: memref<i32>, %[[UNUSED:[^:]+]]: i32, %[[INC:[^:]+]]: i32,
+// CHECK-NEXT: enzymexla.xla_wrapper @update (%[[DATA]], %[[INC]]) <arg_attrs = [{test.data}, {test.increment}], res_attrs = [{test.result}], num_specialized = 1> : (memref<i32>, i32) -> ()
+// CHECK-NEXT: enzymexla.xla_wrapper @update (%[[DATA]], %[[INC]]) <arg_attrs = [{test.other_data}, {test.other_increment}], res_attrs = [{test.other_result}], num_specialized = 1> : (memref<i32>, i32) -> ()
+// CHECK-LABEL: func.func private @update(
+// CHECK-SAME: %[[DATA:[^:]+]]: tensor<i32> {test.function_data}, %[[INC:[^:]+]]: tensor<i32> {test.function_increment}) -> (tensor<i32> {test.function_result}) {
+// CHECK-NEXT: %[[SUM:[^ ]+]] = stablehlo.add %[[DATA]], %[[INC]] : tensor<i32>
+// CHECK-NEXT: return %[[SUM]] : tensor<i32>
+
+// -----
+
+// 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<i32>, %unused: i32) {
+    enzymexla.xla_wrapper @negate (%data, %unused)
+        <{num_specialized = 1 : i64}> : (memref<i32>, i32) -> ()
+    return
+  }
+  func.func @negate(%data: tensor<i32>, %unused: tensor<i32>) -> tensor<i32> {
+    %negative = stablehlo.negate %data : tensor<i32>
+    return %negative : tensor<i32>
+  }
+}
+
+// CHECK-LABEL: func.func @remove_only_scalar(
+// CHECK-SAME: %[[DATA:[^:]+]]: memref<i32>,
+// CHECK-NEXT: enzymexla.xla_wrapper @negate_without_unused (%[[DATA]]) : (memref<i32>) -> ()
+// CHECK-LABEL: func.func @negate(
+// CHECK-SAME: tensor<i32>, {{.*}}: tensor<i32>) -> tensor<i32>
+// CHECK-LABEL: func.func private @negate_without_unused(
+// CHECK-SAME: %[[DATA:[^:]+]]: tensor<i32>) -> tensor<i32> {
+// CHECK-NEXT: %[[NEG:[^ ]+]] = stablehlo.negate %[[DATA]] : tensor<i32>
+// CHECK-NEXT: return %[[NEG]] : tensor<i32>
+
+// -----
+
+// 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<i32>, %unused: i32,
+                                           %increment: i32) {
+    enzymexla.xla_wrapper @add_increment (%data, %unused, %increment)
+        <{num_specialized = 2 : i64}> : (memref<i32>, i32, i32) -> ()
+    return
+  }
+  func.func private @add_increment(%data: tensor<i32>, %unused: tensor<i32>,
+                                   %increment: tensor<i32>) -> tensor<i32> {
+    %sum = stablehlo.add %data, %increment : tensor<i32>
+    return %sum : tensor<i32>
+  }
+}
+
+// CHECK-LABEL: func.func @remove_one_specialized_scalar(
+// CHECK-SAME: %[[DATA:[^:]+]]: memref<i32>, %[[UNUSED:[^:]+]]: i32, %[[INC:[^:]+]]: i32) {
+// CHECK-NEXT: enzymexla.xla_wrapper @add_increment (%[[DATA]], %[[INC]]) <num_specialized = 1> : (memref<i32>, i32) -> ()
+// CHECK-NEXT: return
+// CHECK-LABEL: func.func private @add_increment(
+// CHECK-SAME: %[[DATA:[^:]+]]: tensor<i32>, %[[INC:[^:]+]]: tensor<i32>) -> tensor<i32> {
+// CHECK-NEXT: %[[SUM:[^ ]+]] = stablehlo.add %[[DATA]], %[[INC]] : tensor<i32>
+// CHECK-NEXT: return %[[SUM]] : tensor<i32>
+
+// -----
+
+// 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<i32>, %increment: i32,
+                                            %scale: i32) {
+    enzymexla.xla_wrapper @advance (%data, %increment, %increment, %scale)
+        <{num_specialized = 3 : i64}> : (memref<i32>, i32, i32, i32) -> ()
+    return
+  }
+  func.func @keep_distinct_specialized_scalars(%data: memref<i32>, %first: i32,
+                                               %second: i32, %scale: i32) {
+    enzymexla.xla_wrapper @advance (%data, %first, %second, %scale)
+        <{num_specialized = 3 : i64}> : (memref<i32>, i32, i32, i32) -> ()
+    return
+  }
+  func.func private @advance(%data: tensor<i32>, %first: tensor<i32>,
+                             %second: tensor<i32>, %scale: tensor<i32>)
+      -> tensor<i32> {
+    %once = stablehlo.add %data, %first : tensor<i32>
+    %twice = stablehlo.add %once, %second : tensor<i32>
+    %scaled = stablehlo.multiply %twice, %scale : tensor<i32>
+    return %scaled : tensor<i32>
+  }
+}
+
+// CHECK-LABEL: func.func @deduplicate_specialized_scalar(
+// CHECK-SAME: %[[DATA:[^:]+]]: memref<i32>, %[[INC:[^:]+]]: i32, %[[SCALE:[^:]+]]: i32) {
+// CHECK-NEXT: enzymexla.xla_wrapper @advance_without_duplicates (%[[DATA]], %[[INC]], %[[SCALE]]) <num_specialized = 2> : (memref<i32>, i32, i32) -> ()
+// CHECK-NEXT: return
+// CHECK-LABEL: func.func @keep_distinct_specialized_scalars(
+// CHECK-SAME: %[[DATA:[^:]+]]: memref<i32>, %[[FIRST:[^:]+]]: i32, %[[SECOND:[^:]+]]: i32, %[[SCALE:[^:]+]]: i32) {
+// CHECK-NEXT: enzymexla.xla_wrapper @advance (%[[DATA]], %[[FIRST]], %[[SECOND]], %[[SCALE]]) <num_specialized = 3> : (memref<i32>, i32, i32, i32) -> ()
+// CHECK-NEXT: return
+// CHECK-LABEL: func.func private @advance(
+// CHECK-SAME: %[[DATA:[^:]+]]: tensor<i32>, %[[FIRST:[^:]+]]: tensor<i32>, %[[SECOND:[^:]+]]: tensor<i32>, %[[SCALE:[^:]+]]: tensor<i32>) -> tensor<i32> {
+// CHECK-NEXT: %[[ONCE:[^ ]+]] = stablehlo.add %[[DATA]], %[[FIRST]] : tensor<i32>
+// CHECK-NEXT: %[[TWICE:[^ ]+]] = stablehlo.add %[[ONCE]], %[[SECOND]] : tensor<i32>
+// CHECK-NEXT: %[[SCALED:[^ ]+]] = stablehlo.multiply %[[TWICE]], %[[SCALE]] : tensor<i32>
+// CHECK-NEXT: return %[[SCALED]] : tensor<i32>
+// CHECK-LABEL: func.func private @advance_without_duplicates(
+// CHECK-SAME: %[[DATA:[^:]+]]: tensor<i32>, %[[INC:[^:]+]]: tensor<i32>, %[[SCALE:[^:]+]]: tensor<i32>) -> tensor<i32> {
+// CHECK-NEXT: %[[ONCE:[^ ]+]] = stablehlo.add %[[DATA]], %[[INC]] : tensor<i32>
+// CHECK-NEXT: %[[TWICE:[^ ]+]] = stablehlo.add %[[ONCE]], %[[INC]] : tensor<i32>
+// CHECK-NEXT: %[[SCALED:[^ ]+]] = stablehlo.multiply %[[TWICE]], %[[SCALE]] : tensor<i32>
+// CHECK-NEXT: return %[[SCALED]] : tensor<i32>
+
+// -----
+
+// 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<i32>, %value: i32) {
+    enzymexla.xla_wrapper @write_pair (%data, %data, %value, %value)
+        <{num_specialized = 2 : i64}> : (memref<i32>, memref<i32>, i32, i32) -> ()
+    return
+  }
+  func.func private @write_pair(%first_buffer: tensor<i32>,
+                                %second_buffer: tensor<i32>,
+                                %first_value: tensor<i32>,
+                                %second_value: tensor<i32>)
+      -> (tensor<i32>, tensor<i32>) {
+    return %first_value, %second_value : tensor<i32>, tensor<i32>
+  }
+}
+
+// CHECK-LABEL: func.func @deduplicate_buffers_and_scalars(
+// CHECK-SAME: %[[DATA:[^:]+]]: memref<i32>, %[[VALUE:[^:]+]]: i32) {
+// CHECK-NEXT: enzymexla.xla_wrapper @write_pair_without_duplicates (%[[DATA]], %[[VALUE]]) <num_specialized = 1> : (memref<i32>, i32) -> ()
+// CHECK-NEXT: return
+// CHECK-LABEL: func.func private @write_pair_without_duplicates(
+// CHECK-SAME: %[[DATA:[^:]+]]: tensor<i32>, %[[VALUE:[^:]+]]: tensor<i32>) -> tensor<i32> {
+// CHECK-NEXT: return %[[VALUE]] : tensor<i32>
+
+// -----
+
+// Both buffer positions write the same specialized scalar to the same buffer.
+// 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
+module {
+  func.func @deduplicate_with_specialized(%data: memref<i32>, %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<i32>, memref<i32>, i32, i32) -> ()
+    return
+  }
+  func.func private @write(%first: tensor<i32> {test.function_buffer},
+                           %second: tensor<i32> {test.function_buffer},
+                           %value: tensor<i32> {test.function_value},
+                           %effect: tensor<i32> {test.function_effect})
+      -> (tensor<i32> {test.function_result}, tensor<i32> {test.function_result}) {
+    stablehlo.custom_call @record(%effect) {has_side_effect = true} : (tensor<i32>) -> ()
+    return %value, %value : tensor<i32>, tensor<i32>
+  }
+}
+
+// CHECK-LABEL: func.func @deduplicate_with_specialized(
+// CHECK-SAME: %[[DATA:[^:]+]]: memref<i32>, %[[VALUE:[^:]+]]: i32) {
+// CHECK-NEXT: enzymexla.xla_wrapper @write_without_duplicates (%[[DATA]], %[[VALUE]], %[[VALUE]]) <arg_attrs = [{test.buffer}, {test.value}, {test.effect}], res_attrs = [{test.result}], num_specialized = 2> : (memref<i32>, i32, i32) -> ()
+// CHECK-LABEL: func.func private @write_without_duplicates(
+// CHECK-SAME: %[[DATA:[^:]+]]: tensor<i32> {test.function_buffer}, %[[VALUE:[^:]+]]: tensor<i32> {test.function_value}, %[[EFFECT:[^:]+]]: tensor<i32> {test.function_effect}) -> (tensor<i32> {test.function_result}) {
+// CHECK-NEXT: stablehlo.custom_call @record(%[[EFFECT]]) {has_side_effect = true} : (tensor<i32>) -> ()
+// CHECK-NEXT: return %[[VALUE]] : tensor<i32>
+
+// -----
+
+// 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<i32>) {
+    enzymexla.xla_wrapper @pair (%data, %data)
+        <{arg_attrs = [{test.first}, {test.second}]}> : (memref<i32>, memref<i32>) -> ()
+    enzymexla.xla_wrapper @pair (%data, %data)
+        <{res_attrs = [{test.first}, {test.second}]}> : (memref<i32>, memref<i32>) -> ()
+    enzymexla.xla_wrapper @arg_attributes (%data, %data) : (memref<i32>, memref<i32>) -> ()
+    enzymexla.xla_wrapper @result_attributes (%data, %data) : (memref<i32>, memref<i32>) -> ()
+    return
+  }
+  func.func private @pair(%a: tensor<i32>, %b: tensor<i32>)
+      -> (tensor<i32>, tensor<i32>) {
+    %sum = stablehlo.add %a, %b : tensor<i32>
+    return %sum, %sum : tensor<i32>, tensor<i32>
+  }
+  func.func private @arg_attributes(%a: tensor<i32> {test.first},
+                                    %b: tensor<i32> {test.second})
+      -> (tensor<i32>, tensor<i32>) {
+    %sum = stablehlo.add %a, %b : tensor<i32>
+    return %sum, %sum : tensor<i32>, tensor<i32>
+  }
+  func.func private @result_attributes(%a: tensor<i32>, %b: tensor<i32>)
+      -> (tensor<i32> {test.first}, tensor<i32> {test.second}) {
+    %sum = stablehlo.add %a, %b : tensor<i32>
+    return %sum, %sum : tensor<i32>, tensor<i32>
+  }
+}
+
+// CHECK-LABEL: func.func @keep_different_attributes(
+// CHECK-SAME: %[[DATA:[^:]+]]: memref<i32>) {
+// CHECK-NEXT: enzymexla.xla_wrapper @pair (%[[DATA]], %[[DATA]]) <arg_attrs = [{test.first}, {test.second}]> : (memref<i32>, memref<i32>) -> ()
+// CHECK-NEXT: enzymexla.xla_wrapper @pair (%[[DATA]], %[[DATA]]) <res_attrs = [{test.first}, {test.second}]> : (memref<i32>, memref<i32>) -> ()
+// CHECK-NEXT: enzymexla.xla_wrapper @arg_attributes (%[[DATA]], %[[DATA]]) : (memref<i32>, memref<i32>) -> ()
+// CHECK-NEXT: enzymexla.xla_wrapper @result_attributes (%[[DATA]], %[[DATA]]) : (memref<i32>, memref<i32>) -> ()
+// CHECK-NEXT: return
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: }