[all-commits] [llvm/llvm-project] 12c241: [MLIR][NVVM] Explicit Data Type for Output in `wgm...

[Guray Ozen via All-commits]

  Branch: refs/heads/main
  Home:   https://github.com/llvm/llvm-project
  Commit: 12c241b3654800ab708607dbc1998975c893fc14
      https://github.com/llvm/llvm-project/commit/12c241b3654800ab708607dbc1998975c893fc14
  Author: Guray Ozen <guray.ozen at gmail.com>
  Date:   2024-01-22 (Mon, 22 Jan 2024)

  Changed paths:
    M mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
    M mlir/lib/Conversion/NVGPUToNVVM/NVGPUToNVVM.cpp
    M mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
    M mlir/test/Conversion/NVGPUToNVVM/nvgpu-to-nvvm.mlir
    M mlir/test/Conversion/NVVMToLLVM/invalid.mlir
    M mlir/test/Conversion/NVVMToLLVM/nvvm-to-llvm.mlir
    M mlir/test/python/dialects/nvvm.py

  Log Message:
  -----------
  [MLIR][NVVM] Explicit Data Type for Output in `wgmma.mma_async` (#78713)

The current implementation of `nvvm.wgmma.mma_async` Op deduces the data
type of the output matrix from the data type of struct member, which can be
non-intuitive, especially in cases where types like `2xf16` are packed
into `i32`.

This PR addresses this issue by improving the Op to include an explicit
data type for the output matrix.

The modified Op now includes an explicit data type for Matrix-D (<f16>),
and looks as follows:

```
%result = llvm.mlir.undef : !llvm.struct<(struct<(i32, i32, ...
nvvm.wgmma.mma_async
    %descA, %descB, %result,
    #nvvm.shape<m = 64, n = 32, k = 16>,
    D [<f16>, #nvvm.wgmma_scale_out<zero>],
    A [<f16>, #nvvm.wgmma_scale_in<neg>, <col>],
    B [<f16>, #nvvm.wgmma_scale_in<neg>, <col>]
```