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:

   shCopy

   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
   

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

   shCopy

   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.

3. Run bisect:

   shCopy

   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.