# MLIR HLO mlir_bisect

This is a test case reduction tool, similar in purpose to `mlir-reduce`, but
specific to the `mlir-interpreter` infrastructure. In particular, reductions can
depend on concrete values encountered during execution, and reductions can (and
usually do) generate multiple candidates.

For example, the `ReplaceOpWithConstant` reduction will attempt to replace each
op with each of its results. If the op is in a loop, each execution will be a
candidate for replacement.

## Using this tool

1.  Run a JAX test with snapshots enabled:

    ```
    bazel test some-jax-test
      --test_env=XLA_FLAGS="--xla_cpu_use_xla_runtime --xla_dump_to=/tmp/dump
      --xla_dump_hlo_snapshots" --test_filter=SomeSpecific.Test
      --test_sharding_strategy=disabled --test_strategy=local
    ```

1.  Figure out the culprit module and pass (sorry, no automation yet):

    ```
    bazel run tensorflow/compiler/xla/mlir/tools/mlir_replay:mlir_replay -- \
      --mlir-compilation-trace=/tmp/dump/module_0000.jit__something.mlir-trace.pb \
      --hlo-snapshot=/tmp/dump/module_0000.jit__something.snapshot.0.pb \
      --print-changes-only \
      --execution-trace-dir=/tmp/execution
    ```

    You should see a pass after which results change. You'll want to use the
    .mlir file in `/tmp/execution` corresponding to the pass *before* that with
    the bisect tool.

    Note: If the failing pass is bufferization, you may have to use an earlier
    snapshot, e.g. before EmptyTensorToAllocTensor.
1.  Run bisect:

    ```
    bazel run tensorflow/compiler/xla/mlir/tools/mlir_bisect:mlir-bisect -- \
      --hlo-snapshot=/tmp/dump/module_0000.jit_something.snapshot.0.pb \
      --pass-pipeline="builtin.module(empty-tensor-to-alloc-tensor,one-shot-bufferize{allow-return-allocs bufferize-function-boundaries create-deallocs=0})" \
      /tmp/execution/0052.ScalarizationPass.mlir
    ```

## Adding a reduction

To add a reduction, create a function that generates the candidates and register
it:

```
SmallVector<OwningOpRef<ModuleOp>>
FrobulateAndDefenestrate(BisectState&, dialect::SomeOp some_op) {
  auto [cloned_module_1, cloned_op_1] = CloneModuleFor(some_op);
  Frobulate(cloned_op_1);

  auto [cloned_module_2, cloned_op_2] = CloneModuleFor(some_op);
  Defenestrate(cloned_op_2);

  return {cloned_module_1, cloned_module_2};
}

REGISTER_MLIR_REDUCE_STRATEGY(FrobulateAndDefenestrate);
```

Then, add a test for the strategy. Make sure your strategy is linked into
mlir-bisect and has `alwayslink` set.

```
// RUN: mlir-bisect %s --debug-strategy=FrobulateAndDefenestrate | FileCheck %s

func.func @main() {
  dialect.some_op()
}

// CHECK: func @main()
// CHECK-NEXT: frobulated

// CHECK: func @main()
// CHECK-NEXT: defenestrated
```

`--debug-strategy` will print all candidates generated by the given strategy.
