Add test for EmitReducePrecisionIR

[apivovarov]

I noticed that the EmitReducePrecisionIR function from xla/service/elemental_ir_emitter.h is not covered by unit tests.

Given its non-trivial logic, I believe it should be thoroughly tested, particularly for corner cases.

Changes in this PR:

• Declare EmitReducePrecisionIR function in xla/service/elemental_ir_emitter.h
• Add EmitReducePrecisionIR_F16ToF8e5m2 test
• Add EmitReducePrecisionIR_F16ToF8e4m3fn test

Related PR:

• PR-16585 Add support for float8_e4m3

[reedwm]

Maybe such tests should be in xla/tests/reduce_precision_test.cc. That way it will be tested on backends which don't use element_ir_emitter.

[apivovarov]

Maybe such tests should be in xla/tests/reduce_precision_test.cc. That way it will be tested on backends which don't use element_ir_emitter.

The function EmitReducePrecisionIR is declared in xla/service/elemental_ir_emitter.h and implemented in xla/service/elemental_ir_emitter.cc.

To test the functions declared in elemental_ir_emitter.h, XLA provides a corresponding file, xla/service/elemental_ir_emitter_test.cc.

EmitReducePrecisionIR utilizes LLVM and returns an llvm::Value (LLVM IR), which is then converted to a string representation and verified in the related tests.

In contrast, xla/tests/reduce_precision_test.cc does not rely on LLVM.

Given this, I believe that xla/service/elemental_ir_emitter_test.cc is the correct location for the EmitReducePrecisionIR tests.

[akuegel]

Maybe such tests should be in xla/tests/reduce_precision_test.cc. That way it will be tested on backends which don't use element_ir_emitter.

The function EmitReducePrecisionIR is declared in xla/service/elemental_ir_emitter.h and implemented in xla/service/elemental_ir_emitter.cc.

To test the functions declared in elemental_ir_emitter.h, XLA provides a corresponding file, xla/service/elemental_ir_emitter_test.cc.

EmitReducePrecisionIR utilizes LLVM and returns an llvm::Value (LLVM IR), which is then converted to a string representation and verified in the related tests.

In contrast, xla/tests/reduce_precision_test.cc does not rely on LLVM.

Given this, I believe that xla/service/elemental_ir_emitter_test.cc is the correct location for the EmitReducePrecisionIR tests.

So we have different kind of tests:

• End2End tests for all backends, those are in xla/tests
• GPU specific End2End tests, those are in xla/service/gpu/tests
• Filecheck tests, those are also in xla/service/gpu/tests

For the tests you are adding here, I think filecheck based tests that use the hlo-opt tool would be the best fit. Can you please rewrite the tests based on that? You can take a look at .hlo files in xla/service/gpu/tests to see how to use it.

[apivovarov]

Thank you, Adrian and Reed, for your feedback. I’d like to provide some additional context on why I’m using the GTest framework for this test and the associated code.

The EmitReducePrecisionIR function consists of 138 lines of code and handles several specific cases:

• When the destination mantissa is smaller than the source mantissa (adds 6 LLVM ops).
• When the destination exponent is smaller than the source exponent (adds 7 LLVM ops).
• NaN handling (adds 5 LLVM ops).
• Prolog/epilog setup (adds 5 LLVM ops).

The function generates a sequence of LLVM IR operations that convert an input number in f16 format to an output in a reduced-precision f16-like format, with fewer bits for the exponent and mantissa. The result is returned as an llvm::Value*.

When working on f8E4M3 type support, it wasn’t immediately clear what LLVM IR or output I would get when reducing f16 to a 4-bit exponent and 3-bit mantissa. To explore this, I created a test that generates the LLVM IR and uses specific f16 constants as inputs.

Since the input is a constant, the entire LLVM IR can be evaluated and simplified to a final result. By calling llvm::Value*->print(), the result is folded into a final f16-like LLVM IR constant, represented as specific bits in hexadecimal format.

Note that the GPU compiler has its own separate EmitReducePrecision and EmitF16ToF8e5m2 functions, which use mlir::Value instead of llvm::Value.

The EmitReducePrecisionIR function’s test:

• Is not an end-to-end test.
• Is not an HLO test.
• Is not a GPU test.

It’s a unit test for a specific C++ function (EmitReducePrecisionIR), which uses the LLVM API internally to emit LLVM IR for the CPU compiler.

Given this, I believe the test should use GTest along with the LLVM API to properly validate the result of EmitReducePrecisionIR.

@akuegel @reedwm

[akuegel]

Thanks for the explanation. If it is for the CPU backend, then I think @ezhulenev might be a better reviewer.

[apivovarov]

Hi Eugene,

Could you please review this PR? I've included details above explaining the necessity of the test and the reasoning behind using GTest with the LLVM API.

@ezhulenev

[ezhulenev]

can you split struct declaration from variable declaration, this is too exotic for XLA codebase

[apivovarov]

This approach is used in

• xla/tests/convert_test.cc (in 14 tests)
• xla/util_test.cc (in CommonFactors)

BTW, struct name "TestCase" can be omitted (example xla/util_test.cc CommonFactors)

[dimitar-asenov]

I agree with Eugene's comment. Despite this approach being used else where, it would better to split the declarations.

[apivovarov]

Sure, I refactored the test file and added

struct EmitReducePrecisionIrTestCase {
  float input;
  std::string expected_res;
};

I also rebased the branch and switched to llvm::IRBuilderBase* b.

[dimitar-asenov]

Thanks!

[ezhulenev]

We have https://github.com/openxla/xla/blob/main/xla/service/llvm_ir/llvm_util.h#L61

[apivovarov]

replaced llvmValueToString with llvm_ir::DumpToString

[apivovarov]

Maybe such tests should be in xla/tests/reduce_precision_test.cc. That way it will be tested on backends which don't use element_ir_emitter.

Some argument in favor of using reduce_precision_test.cc to test elemental_ir_emitter.h functions:

• Typically, .h/.cc files have corresponding _test.cc files. In this case, elemental_ir_emitter.cc aligns with reduce_precision_test.cc.
• reduce_precision_test.cc doesn't include the elemental_ir_emitter.h header, nor does it directly test functions from elemental_ir_emitter.h.
• The tests in reduce_precision_test.cc focus on using XlaBuilder::ReducePrecision.

[dimitar-asenov]

@ezhulenev I believe Alexander has addressed your original comments. Do you have any other concerns or can we merge this?