[all-commits] [llvm/llvm-project] 51b65d: [mlir][AMDGPU] Improve BF16 handling through AMDGP...

[Krzysztof Drewniak via All-commits]

  Branch: refs/heads/main
  Home:   https://github.com/llvm/llvm-project
  Commit: 51b65d0895af613f41903faf03dcffe700248d03
      https://github.com/llvm/llvm-project/commit/51b65d0895af613f41903faf03dcffe700248d03
  Author: Krzysztof Drewniak <Krzysztof.Drewniak at amd.com>
  Date:   2023-08-17 (Thu, 17 Aug 2023)

  Changed paths:
    M mlir/include/mlir/Conversion/AMDGPUToROCDL/AMDGPUToROCDL.h
    M mlir/lib/Conversion/AMDGPUToROCDL/AMDGPUToROCDL.cpp
    M mlir/lib/Conversion/GPUToROCDL/CMakeLists.txt
    M mlir/lib/Conversion/GPUToROCDL/LowerGpuOpsToROCDLOps.cpp
    M mlir/lib/Dialect/Arith/Transforms/EmulateUnsupportedFloats.cpp
    M mlir/lib/Target/LLVMIR/ModuleTranslation.cpp
    M mlir/test/Conversion/AMDGPUToROCDL/mfma.mlir
    M mlir/test/Conversion/AMDGPUToROCDL/wmma.mlir
    M mlir/test/Conversion/GPUToROCDL/gpu-to-rocdl.mlir
    M mlir/test/Target/LLVMIR/llvmir.mlir

  Log Message:
  -----------
  [mlir][AMDGPU] Improve BF16 handling through AMDGPU compilation

Many previous sets of AMDGPU dialect code have been incorrect in the
presence of the bf16 type (when lowered to LLVM's bfloat) as they were
developed in a setting that run a custom bf16-to-i16 pass before LLVM
lowering.

An overall effect of this patch is that you should run
--arith-emulate-unsupported-floats="source-types=bf16 target-type=f32"
on your GPU module before calling --convert-gpu-to-rocdl if your code
performs bf16 arithmetic.

While LLVM now supports software bfloat, initial experiments showed
that using this support on AMDGPU inserted a large number of
conversions around loads and stores which had substantial performance
imparts. Furthermore, all of the native AMDGPU operations on bf16
types (like the WMMA operations) operate on 16-bit integers instead of
the bfloat type.

First, we make the following changes to preserve compatibility once
the LLVM bfloat type is reenabled.
1. The matrix multiplication operations (MFMA and WMMA) will bitcast
bfloat vectors to i16 vectors.
2. Buffer loads and stores will operate on the relevant integer
datatype and then cast to bfloat if needed.

Second, we add type conversions to convert bf16 and vectors of it to
equivalent i16 types.

Third, we add the bfloat <-> f32 expansion patterns to the set of
operations run before the main LLVM conversion so that MLIR's
implementation of these conversion routines is used.

Finally, we extend the "floats treated as integers" support in the
LLVM exporter to handle types other than fp8.

We also fix a bug in the unsupported floats emulation where it tried
to operate on `arith.bitcast` due to an oversight.

Reviewed By: rsuderman

Differential Revision: https://reviews.llvm.org/D156361