derek.holloway
7502bf64a4
CodeQL / Analyze (actions) (push) Canceled after 0s
CodeQL / Analyze (python) (push) Canceled after 0s
pre-commit / pre-commit (push) Canceled after 0s
Multi-Arch CI / setup (push) Canceled after 0s
Publish build_manylinux_x86_64 images / publish_build_manylinux_x86_64 (push) Canceled after 0s
Setup (Multi-Arch) / setup (push) Canceled after 0s
Multi-Arch CI / Linux::${{ fromJSON(needs.setup.outputs.linux_build_config || '{}').build_variant_label || 'skip' }} (push) Canceled after 0s
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Publish no_rocm_image_ubi10 image / publish_no_rocm_image_ubi10 (push) Canceled after 0s
Multi-Arch CI / Windows::${{ fromJSON(needs.setup.outputs.windows_build_config || '{}').build_variant_label || 'skip' }} (push) Canceled after 0s
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Multi-Arch CI / Manifest Diff (push) Canceled after 0s
TheRock Manifest Diff Report / Generate Manifest Diff Report (push) Canceled after 0s
Multi-Arch CI / CI Summary (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04 (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04_ocl_rt images / publish_no_rocm_image_ubuntu24_04_ocl_rt (push) Canceled after 0s
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04_rocgdb (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04_rocgdb images / publish_no_rocm_image_ubuntu24_04_rocgdb (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04_ocl_rt (push) Canceled after 0s
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Unit Tests / Unit Tests :: ubuntu-24.04 (push) Canceled after 0s
Unit Tests / Unit Tests :: windows-2022 (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04_media (push) Canceled after 0s
Unit Tests / Unit Tests Summary (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04_openmpi (push) Canceled after 0s
815 lines
33 KiB
C++
815 lines
33 KiB
C++
//===- NVVMToLLVMIRTranslation.cpp - Translate NVVM to LLVM IR ------------===//
|
|
//
|
|
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
|
|
// See https://llvm.org/LICENSE.txt for license information.
|
|
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
|
|
//
|
|
//===----------------------------------------------------------------------===//
|
|
//
|
|
// This file implements a translation between the MLIR NVVM dialect and
|
|
// LLVM IR.
|
|
//
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
#include "mlir/Target/LLVMIR/Dialect/NVVM/NVVMToLLVMIRTranslation.h"
|
|
#include "mlir/Dialect/LLVMIR/NVVMDialect.h"
|
|
#include "mlir/IR/Operation.h"
|
|
#include "mlir/Target/LLVMIR/ModuleTranslation.h"
|
|
|
|
#include "llvm/ADT/StringExtras.h"
|
|
#include "llvm/ADT/iterator_range.h"
|
|
#include "llvm/IR/IRBuilder.h"
|
|
#include "llvm/IR/IntrinsicsNVPTX.h"
|
|
#include "llvm/Support/FormatVariadic.h"
|
|
#include "llvm/Support/NVVMAttributes.h"
|
|
|
|
using namespace mlir;
|
|
using namespace mlir::LLVM;
|
|
using mlir::LLVM::detail::createIntrinsicCall;
|
|
|
|
#define REDUX_F32_ID_IMPL(op, abs, hasNaN) \
|
|
hasNaN ? llvm::Intrinsic::nvvm_redux_sync_f##op##abs##_NaN \
|
|
: llvm::Intrinsic::nvvm_redux_sync_f##op##abs
|
|
|
|
#define GET_REDUX_F32_ID(op, hasAbs, hasNaN) \
|
|
hasAbs ? REDUX_F32_ID_IMPL(op, _abs, hasNaN) : REDUX_F32_ID_IMPL(op, , hasNaN)
|
|
|
|
static llvm::Intrinsic::ID getReduxIntrinsicId(llvm::Type *resultType,
|
|
NVVM::ReductionKind kind,
|
|
bool hasAbs, bool hasNaN) {
|
|
switch (kind) {
|
|
case NVVM::ReductionKind::ADD:
|
|
return llvm::Intrinsic::nvvm_redux_sync_add;
|
|
case NVVM::ReductionKind::UMAX:
|
|
return llvm::Intrinsic::nvvm_redux_sync_umax;
|
|
case NVVM::ReductionKind::UMIN:
|
|
return llvm::Intrinsic::nvvm_redux_sync_umin;
|
|
case NVVM::ReductionKind::AND:
|
|
return llvm::Intrinsic::nvvm_redux_sync_and;
|
|
case NVVM::ReductionKind::OR:
|
|
return llvm::Intrinsic::nvvm_redux_sync_or;
|
|
case NVVM::ReductionKind::XOR:
|
|
return llvm::Intrinsic::nvvm_redux_sync_xor;
|
|
case NVVM::ReductionKind::MAX:
|
|
return llvm::Intrinsic::nvvm_redux_sync_max;
|
|
case NVVM::ReductionKind::MIN:
|
|
return llvm::Intrinsic::nvvm_redux_sync_min;
|
|
case NVVM::ReductionKind::FMIN:
|
|
return GET_REDUX_F32_ID(min, hasAbs, hasNaN);
|
|
case NVVM::ReductionKind::FMAX:
|
|
return GET_REDUX_F32_ID(max, hasAbs, hasNaN);
|
|
}
|
|
llvm_unreachable("unknown reduction kind");
|
|
}
|
|
|
|
static llvm::Intrinsic::ID getShflIntrinsicId(llvm::Type *resultType,
|
|
NVVM::ShflKind kind,
|
|
bool withPredicate) {
|
|
|
|
if (withPredicate) {
|
|
resultType = cast<llvm::StructType>(resultType)->getElementType(0);
|
|
switch (kind) {
|
|
case NVVM::ShflKind::bfly:
|
|
return resultType->isFloatTy()
|
|
? llvm::Intrinsic::nvvm_shfl_sync_bfly_f32p
|
|
: llvm::Intrinsic::nvvm_shfl_sync_bfly_i32p;
|
|
case NVVM::ShflKind::up:
|
|
return resultType->isFloatTy() ? llvm::Intrinsic::nvvm_shfl_sync_up_f32p
|
|
: llvm::Intrinsic::nvvm_shfl_sync_up_i32p;
|
|
case NVVM::ShflKind::down:
|
|
return resultType->isFloatTy()
|
|
? llvm::Intrinsic::nvvm_shfl_sync_down_f32p
|
|
: llvm::Intrinsic::nvvm_shfl_sync_down_i32p;
|
|
case NVVM::ShflKind::idx:
|
|
return resultType->isFloatTy() ? llvm::Intrinsic::nvvm_shfl_sync_idx_f32p
|
|
: llvm::Intrinsic::nvvm_shfl_sync_idx_i32p;
|
|
}
|
|
} else {
|
|
switch (kind) {
|
|
case NVVM::ShflKind::bfly:
|
|
return resultType->isFloatTy() ? llvm::Intrinsic::nvvm_shfl_sync_bfly_f32
|
|
: llvm::Intrinsic::nvvm_shfl_sync_bfly_i32;
|
|
case NVVM::ShflKind::up:
|
|
return resultType->isFloatTy() ? llvm::Intrinsic::nvvm_shfl_sync_up_f32
|
|
: llvm::Intrinsic::nvvm_shfl_sync_up_i32;
|
|
case NVVM::ShflKind::down:
|
|
return resultType->isFloatTy() ? llvm::Intrinsic::nvvm_shfl_sync_down_f32
|
|
: llvm::Intrinsic::nvvm_shfl_sync_down_i32;
|
|
case NVVM::ShflKind::idx:
|
|
return resultType->isFloatTy() ? llvm::Intrinsic::nvvm_shfl_sync_idx_f32
|
|
: llvm::Intrinsic::nvvm_shfl_sync_idx_i32;
|
|
}
|
|
}
|
|
llvm_unreachable("unknown shuffle kind");
|
|
}
|
|
|
|
static llvm::Intrinsic::ID getMatchSyncIntrinsicId(Type valType,
|
|
NVVM::MatchSyncKind kind) {
|
|
switch (kind) {
|
|
case NVVM::MatchSyncKind::any:
|
|
return valType.isInteger(32) ? llvm::Intrinsic::nvvm_match_any_sync_i32
|
|
: llvm::Intrinsic::nvvm_match_any_sync_i64;
|
|
case NVVM::MatchSyncKind::all:
|
|
// match.all instruction has two variants -- one returns a single value,
|
|
// another returns a pair {value, predicate}. We currently only implement
|
|
// the latter as that's the variant exposed by CUDA API.
|
|
return valType.isInteger(32) ? llvm::Intrinsic::nvvm_match_all_sync_i32p
|
|
: llvm::Intrinsic::nvvm_match_all_sync_i64p;
|
|
}
|
|
llvm_unreachable("unsupported match sync kind");
|
|
}
|
|
|
|
static llvm::Intrinsic::ID getVoteSyncIntrinsicId(NVVM::VoteSyncKind kind) {
|
|
switch (kind) {
|
|
case NVVM::VoteSyncKind::any:
|
|
return llvm::Intrinsic::nvvm_vote_any_sync;
|
|
case NVVM::VoteSyncKind::all:
|
|
return llvm::Intrinsic::nvvm_vote_all_sync;
|
|
case NVVM::VoteSyncKind::ballot:
|
|
return llvm::Intrinsic::nvvm_vote_ballot_sync;
|
|
case NVVM::VoteSyncKind::uni:
|
|
return llvm::Intrinsic::nvvm_vote_uni_sync;
|
|
}
|
|
llvm_unreachable("unsupported vote kind");
|
|
}
|
|
|
|
static llvm::Intrinsic::ID
|
|
getLdMatrixIntrinsicId(NVVM::MMALayout layout, int32_t num,
|
|
NVVM::LdStMatrixShapeAttr shape,
|
|
NVVM::LdStMatrixEltType eltType) {
|
|
if (shape.getM() == 8 && shape.getN() == 8) {
|
|
switch (num) {
|
|
case 1:
|
|
return (layout == NVVM::MMALayout::row)
|
|
? llvm::Intrinsic::nvvm_ldmatrix_sync_aligned_m8n8_x1_b16
|
|
: llvm::Intrinsic::
|
|
nvvm_ldmatrix_sync_aligned_m8n8_x1_trans_b16;
|
|
case 2:
|
|
return (layout == NVVM::MMALayout::row)
|
|
? llvm::Intrinsic::nvvm_ldmatrix_sync_aligned_m8n8_x2_b16
|
|
: llvm::Intrinsic::
|
|
nvvm_ldmatrix_sync_aligned_m8n8_x2_trans_b16;
|
|
case 4:
|
|
return (layout == NVVM::MMALayout::row)
|
|
? llvm::Intrinsic::nvvm_ldmatrix_sync_aligned_m8n8_x4_b16
|
|
: llvm::Intrinsic::
|
|
nvvm_ldmatrix_sync_aligned_m8n8_x4_trans_b16;
|
|
}
|
|
} else if (shape.getM() == 8 && shape.getN() == 16) {
|
|
if (eltType == NVVM::LdStMatrixEltType::B8X16_B6X16_P32) {
|
|
switch (num) {
|
|
case 1:
|
|
return llvm::Intrinsic::
|
|
nvvm_ldmatrix_sync_aligned_m8n16_x1_b8x16_b6x16_p32;
|
|
case 2:
|
|
return llvm::Intrinsic::
|
|
nvvm_ldmatrix_sync_aligned_m8n16_x2_b8x16_b6x16_p32;
|
|
case 4:
|
|
return llvm::Intrinsic::
|
|
nvvm_ldmatrix_sync_aligned_m8n16_x4_b8x16_b6x16_p32;
|
|
}
|
|
} else if (eltType == NVVM::LdStMatrixEltType::B8X16_B4X16_P64) {
|
|
switch (num) {
|
|
case 1:
|
|
return llvm::Intrinsic::
|
|
nvvm_ldmatrix_sync_aligned_m8n16_x1_b8x16_b4x16_p64;
|
|
case 2:
|
|
return llvm::Intrinsic::
|
|
nvvm_ldmatrix_sync_aligned_m8n16_x2_b8x16_b4x16_p64;
|
|
case 4:
|
|
return llvm::Intrinsic::
|
|
nvvm_ldmatrix_sync_aligned_m8n16_x4_b8x16_b4x16_p64;
|
|
}
|
|
}
|
|
} else if (shape.getM() == 16 && shape.getN() == 16) {
|
|
if (eltType == NVVM::LdStMatrixEltType::B8) {
|
|
switch (num) {
|
|
case 1:
|
|
return llvm::Intrinsic::nvvm_ldmatrix_sync_aligned_m16n16_x1_trans_b8;
|
|
case 2:
|
|
return llvm::Intrinsic::nvvm_ldmatrix_sync_aligned_m16n16_x2_trans_b8;
|
|
}
|
|
} else if (eltType == NVVM::LdStMatrixEltType::B8X16_B6X16_P32) {
|
|
switch (num) {
|
|
case 1:
|
|
return llvm::Intrinsic::
|
|
nvvm_ldmatrix_sync_aligned_m16n16_x1_trans_b8x16_b6x16_p32;
|
|
case 2:
|
|
return llvm::Intrinsic::
|
|
nvvm_ldmatrix_sync_aligned_m16n16_x2_trans_b8x16_b6x16_p32;
|
|
}
|
|
} else if (eltType == NVVM::LdStMatrixEltType::B8X16_B4X16_P64) {
|
|
switch (num) {
|
|
case 1:
|
|
return llvm::Intrinsic::
|
|
nvvm_ldmatrix_sync_aligned_m16n16_x1_trans_b8x16_b4x16_p64;
|
|
case 2:
|
|
return llvm::Intrinsic::
|
|
nvvm_ldmatrix_sync_aligned_m16n16_x2_trans_b8x16_b4x16_p64;
|
|
}
|
|
}
|
|
}
|
|
llvm_unreachable("unknown ldmatrix kind");
|
|
}
|
|
|
|
/// Return the intrinsic ID associated with stmatrix for the given paramters.
|
|
static llvm::Intrinsic::ID
|
|
getStMatrixIntrinsicId(NVVM::MMALayout layout, int32_t num,
|
|
NVVM::LdStMatrixShapeAttr shape,
|
|
NVVM::LdStMatrixEltType eltType) {
|
|
if (shape.getM() == 8 && shape.getN() == 8) {
|
|
switch (num) {
|
|
case 1:
|
|
return (layout == NVVM::MMALayout::row)
|
|
? llvm::Intrinsic::nvvm_stmatrix_sync_aligned_m8n8_x1_b16
|
|
: llvm::Intrinsic::
|
|
nvvm_stmatrix_sync_aligned_m8n8_x1_trans_b16;
|
|
case 2:
|
|
return (layout == NVVM::MMALayout::row)
|
|
? llvm::Intrinsic::nvvm_stmatrix_sync_aligned_m8n8_x2_b16
|
|
: llvm::Intrinsic::
|
|
nvvm_stmatrix_sync_aligned_m8n8_x2_trans_b16;
|
|
case 4:
|
|
return (layout == NVVM::MMALayout::row)
|
|
? llvm::Intrinsic::nvvm_stmatrix_sync_aligned_m8n8_x4_b16
|
|
: llvm::Intrinsic::
|
|
nvvm_stmatrix_sync_aligned_m8n8_x4_trans_b16;
|
|
}
|
|
} else if (shape.getM() == 16 && shape.getN() == 8) {
|
|
switch (num) {
|
|
case 1:
|
|
return llvm::Intrinsic::nvvm_stmatrix_sync_aligned_m16n8_x1_trans_b8;
|
|
case 2:
|
|
return llvm::Intrinsic::nvvm_stmatrix_sync_aligned_m16n8_x2_trans_b8;
|
|
case 4:
|
|
return llvm::Intrinsic::nvvm_stmatrix_sync_aligned_m16n8_x4_trans_b8;
|
|
}
|
|
}
|
|
llvm_unreachable("unknown stmatrix kind");
|
|
}
|
|
|
|
/// Return the intrinsic ID associated with st.bulk for the given address type.
|
|
static llvm::Intrinsic::ID
|
|
getStBulkIntrinsicId(LLVM::LLVMPointerType addrType) {
|
|
bool isSharedMemory = addrType.getAddressSpace() ==
|
|
static_cast<unsigned>(NVVM::NVVMMemorySpace::Shared);
|
|
return isSharedMemory ? llvm::Intrinsic::nvvm_st_bulk_shared_cta
|
|
: llvm::Intrinsic::nvvm_st_bulk;
|
|
}
|
|
|
|
static unsigned getUnidirectionalFenceProxyID(NVVM::ProxyKind fromProxy,
|
|
NVVM::ProxyKind toProxy,
|
|
NVVM::MemScopeKind scope,
|
|
bool isRelease) {
|
|
if (fromProxy == NVVM::ProxyKind::GENERIC &&
|
|
toProxy == NVVM::ProxyKind::TENSORMAP) {
|
|
switch (scope) {
|
|
case NVVM::MemScopeKind::CTA: {
|
|
if (isRelease)
|
|
return llvm::Intrinsic::nvvm_fence_proxy_tensormap_generic_release_cta;
|
|
return llvm::Intrinsic::nvvm_fence_proxy_tensormap_generic_acquire_cta;
|
|
}
|
|
case NVVM::MemScopeKind::CLUSTER: {
|
|
if (isRelease)
|
|
return llvm::Intrinsic::
|
|
nvvm_fence_proxy_tensormap_generic_release_cluster;
|
|
return llvm::Intrinsic::
|
|
nvvm_fence_proxy_tensormap_generic_acquire_cluster;
|
|
}
|
|
case NVVM::MemScopeKind::GPU: {
|
|
if (isRelease)
|
|
return llvm::Intrinsic::nvvm_fence_proxy_tensormap_generic_release_gpu;
|
|
return llvm::Intrinsic::nvvm_fence_proxy_tensormap_generic_acquire_gpu;
|
|
}
|
|
case NVVM::MemScopeKind::SYS: {
|
|
if (isRelease)
|
|
return llvm::Intrinsic::nvvm_fence_proxy_tensormap_generic_release_sys;
|
|
return llvm::Intrinsic::nvvm_fence_proxy_tensormap_generic_acquire_sys;
|
|
}
|
|
}
|
|
llvm_unreachable("Unknown scope for uni-directional fence.proxy operation");
|
|
}
|
|
llvm_unreachable("Unsupported proxy kinds");
|
|
}
|
|
|
|
static unsigned getMembarIntrinsicID(NVVM::MemScopeKind scope) {
|
|
switch (scope) {
|
|
case NVVM::MemScopeKind::CTA:
|
|
return llvm::Intrinsic::nvvm_membar_cta;
|
|
case NVVM::MemScopeKind::CLUSTER:
|
|
return llvm::Intrinsic::nvvm_fence_sc_cluster;
|
|
case NVVM::MemScopeKind::GPU:
|
|
return llvm::Intrinsic::nvvm_membar_gl;
|
|
case NVVM::MemScopeKind::SYS:
|
|
return llvm::Intrinsic::nvvm_membar_sys;
|
|
}
|
|
llvm_unreachable("Unknown scope for memory barrier");
|
|
}
|
|
|
|
#define TCGEN05LD(SHAPE, NUM) llvm::Intrinsic::nvvm_tcgen05_ld_##SHAPE##_##NUM
|
|
|
|
static llvm::Intrinsic::ID
|
|
getTcgen05LdIntrinsicID(mlir::NVVM::Tcgen05LdStShape shape, uint32_t num) {
|
|
llvm::Intrinsic::ID Shape16x64b[] = {
|
|
TCGEN05LD(16x64b, x1), TCGEN05LD(16x64b, x2), TCGEN05LD(16x64b, x4),
|
|
TCGEN05LD(16x64b, x8), TCGEN05LD(16x64b, x16), TCGEN05LD(16x64b, x32),
|
|
TCGEN05LD(16x64b, x64), TCGEN05LD(16x64b, x128),
|
|
};
|
|
|
|
llvm::Intrinsic::ID Shape16x128b[] = {
|
|
TCGEN05LD(16x128b, x1), TCGEN05LD(16x128b, x2), TCGEN05LD(16x128b, x4),
|
|
TCGEN05LD(16x128b, x8), TCGEN05LD(16x128b, x16), TCGEN05LD(16x128b, x32),
|
|
TCGEN05LD(16x128b, x64),
|
|
};
|
|
|
|
llvm::Intrinsic::ID Shape16x256b[] = {
|
|
TCGEN05LD(16x256b, x1), TCGEN05LD(16x256b, x2), TCGEN05LD(16x256b, x4),
|
|
TCGEN05LD(16x256b, x8), TCGEN05LD(16x256b, x16), TCGEN05LD(16x256b, x32),
|
|
};
|
|
|
|
llvm::Intrinsic::ID Shape16x32bx2[] = {
|
|
TCGEN05LD(16x32bx2, x1), TCGEN05LD(16x32bx2, x2),
|
|
TCGEN05LD(16x32bx2, x4), TCGEN05LD(16x32bx2, x8),
|
|
TCGEN05LD(16x32bx2, x16), TCGEN05LD(16x32bx2, x32),
|
|
TCGEN05LD(16x32bx2, x64), TCGEN05LD(16x32bx2, x128),
|
|
};
|
|
|
|
llvm::Intrinsic::ID Shape32x32b[] = {
|
|
TCGEN05LD(32x32b, x1), TCGEN05LD(32x32b, x2), TCGEN05LD(32x32b, x4),
|
|
TCGEN05LD(32x32b, x8), TCGEN05LD(32x32b, x16), TCGEN05LD(32x32b, x32),
|
|
TCGEN05LD(32x32b, x64), TCGEN05LD(32x32b, x128),
|
|
};
|
|
|
|
// `num` contains the length of vector and log2 of `num` returns the index
|
|
// into the shape array
|
|
unsigned Idx = std::log2(num);
|
|
|
|
switch (shape) {
|
|
case NVVM::Tcgen05LdStShape::SHAPE_16X64B:
|
|
return Shape16x64b[Idx];
|
|
case NVVM::Tcgen05LdStShape::SHAPE_16X128B:
|
|
return Shape16x128b[Idx - 1];
|
|
case NVVM::Tcgen05LdStShape::SHAPE_16X256B:
|
|
return Shape16x256b[Idx - 2];
|
|
case NVVM::Tcgen05LdStShape::SHAPE_32X32B:
|
|
return Shape32x32b[Idx];
|
|
case NVVM::Tcgen05LdStShape::SHAPE_16X32BX2:
|
|
return Shape16x32bx2[Idx];
|
|
}
|
|
llvm_unreachable("unhandled tcgen05.ld lowering");
|
|
}
|
|
|
|
#define TCGEN05ST(SHAPE, NUM) llvm::Intrinsic::nvvm_tcgen05_st_##SHAPE##_##NUM
|
|
|
|
static llvm::Intrinsic::ID
|
|
getTcgen05StIntrinsicID(mlir::NVVM::Tcgen05LdStShape shape, uint32_t num) {
|
|
llvm::Intrinsic::ID Shape16x64b[] = {
|
|
TCGEN05ST(16x64b, x1), TCGEN05ST(16x64b, x2), TCGEN05ST(16x64b, x4),
|
|
TCGEN05ST(16x64b, x8), TCGEN05ST(16x64b, x16), TCGEN05ST(16x64b, x32),
|
|
TCGEN05ST(16x64b, x64), TCGEN05ST(16x64b, x128),
|
|
};
|
|
|
|
llvm::Intrinsic::ID Shape16x128b[] = {
|
|
TCGEN05ST(16x128b, x1), TCGEN05ST(16x128b, x2), TCGEN05ST(16x128b, x4),
|
|
TCGEN05ST(16x128b, x8), TCGEN05ST(16x128b, x16), TCGEN05ST(16x128b, x32),
|
|
TCGEN05ST(16x128b, x64),
|
|
};
|
|
|
|
llvm::Intrinsic::ID Shape16x256b[] = {
|
|
TCGEN05ST(16x256b, x1), TCGEN05ST(16x256b, x2), TCGEN05ST(16x256b, x4),
|
|
TCGEN05ST(16x256b, x8), TCGEN05ST(16x256b, x16), TCGEN05ST(16x256b, x32),
|
|
};
|
|
|
|
llvm::Intrinsic::ID Shape16x32bx2[] = {
|
|
TCGEN05ST(16x32bx2, x1), TCGEN05ST(16x32bx2, x2),
|
|
TCGEN05ST(16x32bx2, x4), TCGEN05ST(16x32bx2, x8),
|
|
TCGEN05ST(16x32bx2, x16), TCGEN05ST(16x32bx2, x32),
|
|
TCGEN05ST(16x32bx2, x64), TCGEN05ST(16x32bx2, x128),
|
|
};
|
|
|
|
llvm::Intrinsic::ID Shape32x32b[] = {
|
|
TCGEN05ST(32x32b, x1), TCGEN05ST(32x32b, x2), TCGEN05ST(32x32b, x4),
|
|
TCGEN05ST(32x32b, x8), TCGEN05ST(32x32b, x16), TCGEN05ST(32x32b, x32),
|
|
TCGEN05ST(32x32b, x64), TCGEN05ST(32x32b, x128),
|
|
};
|
|
|
|
// `num` contains the length of vector and log2 of `num` returns the index
|
|
// into the shape array
|
|
unsigned Idx = std::log2(num);
|
|
|
|
switch (shape) {
|
|
case NVVM::Tcgen05LdStShape::SHAPE_16X64B:
|
|
return Shape16x64b[Idx];
|
|
case NVVM::Tcgen05LdStShape::SHAPE_16X128B:
|
|
return Shape16x128b[Idx - 1];
|
|
case NVVM::Tcgen05LdStShape::SHAPE_16X256B:
|
|
return Shape16x256b[Idx - 2];
|
|
case NVVM::Tcgen05LdStShape::SHAPE_32X32B:
|
|
return Shape32x32b[Idx];
|
|
case NVVM::Tcgen05LdStShape::SHAPE_16X32BX2:
|
|
return Shape16x32bx2[Idx];
|
|
}
|
|
llvm_unreachable("unhandled tcgen05.st lowering");
|
|
}
|
|
|
|
static llvm::Intrinsic::ID getFenceSyncRestrictID(NVVM::MemOrderKind order) {
|
|
return order == NVVM::MemOrderKind::ACQUIRE
|
|
? llvm::Intrinsic::
|
|
nvvm_fence_acquire_sync_restrict_space_cluster_scope_cluster
|
|
: llvm::Intrinsic::
|
|
nvvm_fence_release_sync_restrict_space_cta_scope_cluster;
|
|
}
|
|
|
|
static llvm::Intrinsic::ID
|
|
getFenceProxyID(NVVM::ProxyKind kind, std::optional<NVVM::SharedSpace> space) {
|
|
switch (kind) {
|
|
case NVVM::ProxyKind::alias:
|
|
return llvm::Intrinsic::nvvm_fence_proxy_alias;
|
|
case NVVM::ProxyKind::async:
|
|
return llvm::Intrinsic::nvvm_fence_proxy_async;
|
|
case NVVM::ProxyKind::async_global:
|
|
return llvm::Intrinsic::nvvm_fence_proxy_async_global;
|
|
case NVVM::ProxyKind::async_shared:
|
|
return *space == NVVM::SharedSpace::shared_cta
|
|
? llvm::Intrinsic::nvvm_fence_proxy_async_shared_cta
|
|
: llvm::Intrinsic::nvvm_fence_proxy_async_shared_cluster;
|
|
default:
|
|
llvm_unreachable("unsupported proxy kind");
|
|
}
|
|
}
|
|
|
|
static llvm::Intrinsic::ID
|
|
getFenceProxySyncRestrictID(NVVM::MemOrderKind order) {
|
|
return order == NVVM::MemOrderKind::ACQUIRE
|
|
? llvm::Intrinsic::
|
|
nvvm_fence_proxy_async_generic_acquire_sync_restrict_space_cluster_scope_cluster
|
|
: llvm::Intrinsic::
|
|
nvvm_fence_proxy_async_generic_release_sync_restrict_space_cta_scope_cluster;
|
|
}
|
|
|
|
// Calls an LLVM intrinsic on the given operands. For f32/f64 vector types,
|
|
// the intrinsic is called per-element and the results are packed back into a
|
|
// vector. If retType is non-null, it is forwarded as the return-type
|
|
// overload to `createIntrinsicCall`.
|
|
static llvm::Value *
|
|
createScalarizedIntrinsicCall(llvm::IRBuilderBase &builder,
|
|
llvm::Intrinsic::ID IID, llvm::Type *opTypeLLVM,
|
|
ArrayRef<llvm::Value *> operands,
|
|
llvm::Type *retType) {
|
|
if (opTypeLLVM->isVectorTy() && (opTypeLLVM->getScalarType()->isFloatTy() ||
|
|
opTypeLLVM->getScalarType()->isDoubleTy())) {
|
|
llvm::Value *result = llvm::PoisonValue::get(
|
|
llvm::FixedVectorType::get(opTypeLLVM->getScalarType(), 2));
|
|
for (int64_t i = 0; i < 2; ++i) {
|
|
llvm::SmallVector<llvm::Value *> scalarArgs;
|
|
for (llvm::Value *op : operands)
|
|
scalarArgs.push_back(
|
|
builder.CreateExtractElement(op, builder.getInt32(i)));
|
|
llvm::Value *res = createIntrinsicCall(builder, IID, retType, scalarArgs);
|
|
result = builder.CreateInsertElement(result, res, builder.getInt32(i));
|
|
}
|
|
return result;
|
|
}
|
|
|
|
return createIntrinsicCall(builder, IID, retType, operands);
|
|
}
|
|
|
|
void NVVM::AddFOp::lowerAddFToLLVMIR(llvm::Value *argLHS, llvm::Value *argRHS,
|
|
Value res, NVVM::FPRoundingMode rndMode,
|
|
NVVM::SaturationMode satMode, bool isFTZ,
|
|
LLVM::ModuleTranslation &mt,
|
|
llvm::IRBuilderBase &builder) {
|
|
llvm::Type *opTypeLLVM = argLHS->getType();
|
|
bool isVectorOp = opTypeLLVM->isVectorTy();
|
|
bool isSat = satMode != NVVM::SaturationMode::NONE;
|
|
|
|
// FIXME: Add intrinsics for add.rn.ftz.f16x2 and add.rn.ftz.f16 here when
|
|
// they are available.
|
|
static constexpr llvm::Intrinsic::ID f16IDs[] = {
|
|
llvm::Intrinsic::nvvm_add_rn_sat_f16,
|
|
llvm::Intrinsic::nvvm_add_rn_ftz_sat_f16,
|
|
llvm::Intrinsic::nvvm_add_rn_sat_v2f16,
|
|
llvm::Intrinsic::nvvm_add_rn_ftz_sat_v2f16,
|
|
};
|
|
|
|
static constexpr llvm::Intrinsic::ID f32IDs[] = {
|
|
llvm::Intrinsic::nvvm_add_rn_f, // default rounding mode RN
|
|
llvm::Intrinsic::nvvm_add_rn_f,
|
|
llvm::Intrinsic::nvvm_add_rm_f,
|
|
llvm::Intrinsic::nvvm_add_rp_f,
|
|
llvm::Intrinsic::nvvm_add_rz_f,
|
|
llvm::Intrinsic::nvvm_add_rn_sat_f, // default rounding mode RN
|
|
llvm::Intrinsic::nvvm_add_rn_sat_f,
|
|
llvm::Intrinsic::nvvm_add_rm_sat_f,
|
|
llvm::Intrinsic::nvvm_add_rp_sat_f,
|
|
llvm::Intrinsic::nvvm_add_rz_sat_f,
|
|
llvm::Intrinsic::nvvm_add_rn_ftz_f, // default rounding mode RN
|
|
llvm::Intrinsic::nvvm_add_rn_ftz_f,
|
|
llvm::Intrinsic::nvvm_add_rm_ftz_f,
|
|
llvm::Intrinsic::nvvm_add_rp_ftz_f,
|
|
llvm::Intrinsic::nvvm_add_rz_ftz_f,
|
|
llvm::Intrinsic::nvvm_add_rn_ftz_sat_f, // default rounding mode RN
|
|
llvm::Intrinsic::nvvm_add_rn_ftz_sat_f,
|
|
llvm::Intrinsic::nvvm_add_rm_ftz_sat_f,
|
|
llvm::Intrinsic::nvvm_add_rp_ftz_sat_f,
|
|
llvm::Intrinsic::nvvm_add_rz_ftz_sat_f,
|
|
};
|
|
|
|
static constexpr llvm::Intrinsic::ID f64IDs[] = {
|
|
llvm::Intrinsic::nvvm_add_rn_d, // default rounding mode RN
|
|
llvm::Intrinsic::nvvm_add_rn_d, llvm::Intrinsic::nvvm_add_rm_d,
|
|
llvm::Intrinsic::nvvm_add_rp_d, llvm::Intrinsic::nvvm_add_rz_d};
|
|
|
|
auto addIntrinsic = [&](llvm::Intrinsic::ID IID) -> llvm::Value * {
|
|
return createScalarizedIntrinsicCall(builder, IID, opTypeLLVM,
|
|
{argLHS, argRHS}, opTypeLLVM);
|
|
};
|
|
|
|
// f16 + f16 -> f16 / vector<2xf16> + vector<2xf16> -> vector<2xf16>
|
|
// FIXME: Allow lowering to add.rn.ftz.f16x2 and add.rn.ftz.f16 here when the
|
|
// intrinsics are available.
|
|
if (opTypeLLVM->getScalarType()->isHalfTy()) {
|
|
llvm::Value *result;
|
|
if (isSat) {
|
|
unsigned index = (isVectorOp << 1) | isFTZ;
|
|
result = addIntrinsic(f16IDs[index]);
|
|
} else {
|
|
result = builder.CreateFAdd(argLHS, argRHS);
|
|
}
|
|
mt.mapValue(res, result);
|
|
return;
|
|
}
|
|
|
|
// bf16 + bf16 -> bf16 / vector<2xbf16> + vector<2xbf16> -> vector<2xbf16>
|
|
if (opTypeLLVM->getScalarType()->isBFloatTy()) {
|
|
mt.mapValue(res, builder.CreateFAdd(argLHS, argRHS));
|
|
return;
|
|
}
|
|
|
|
// f64 + f64 -> f64 / vector<2xf64> + vector<2xf64> -> vector<2xf64>
|
|
if (opTypeLLVM->getScalarType()->isDoubleTy()) {
|
|
unsigned index = static_cast<unsigned>(rndMode);
|
|
mt.mapValue(res, addIntrinsic(f64IDs[index]));
|
|
return;
|
|
}
|
|
|
|
// f32 + f32 -> f32 / vector<2xf32> + vector<2xf32> -> vector<2xf32>
|
|
const unsigned numRndModes = 5; // NONE, RM, RN, RP, RZ
|
|
if (opTypeLLVM->getScalarType()->isFloatTy()) {
|
|
unsigned index =
|
|
((isFTZ << 1) | isSat) * numRndModes + static_cast<unsigned>(rndMode);
|
|
mt.mapValue(res, addIntrinsic(f32IDs[index]));
|
|
return;
|
|
}
|
|
}
|
|
|
|
void NVVM::FmaOp::lowerFmaToLLVMIR(Operation &op, LLVM::ModuleTranslation &mt,
|
|
llvm::IRBuilderBase &builder) {
|
|
auto thisOp = cast<NVVM::FmaOp>(op);
|
|
mlir::NVVM::FPRoundingMode rndMode = thisOp.getRnd();
|
|
unsigned rndIndex = static_cast<unsigned>(rndMode) - 1; // 1-4 mapped to 0-3
|
|
mlir::NVVM::SaturationMode satMode = thisOp.getSat();
|
|
bool isFTZ = thisOp.getFtz();
|
|
bool isRelu = thisOp.getRelu();
|
|
bool isSat = satMode == NVVM::SaturationMode::SAT;
|
|
bool isOOB = thisOp.getOob();
|
|
|
|
mlir::Type opType = thisOp.getRes().getType();
|
|
llvm::Type *opTypeLLVM = mt.convertType(opType);
|
|
bool isVectorFma = opTypeLLVM->isVectorTy();
|
|
|
|
llvm::Value *argA = mt.lookupValue(thisOp.getA());
|
|
llvm::Value *argB = mt.lookupValue(thisOp.getB());
|
|
llvm::Value *argC = mt.lookupValue(thisOp.getC());
|
|
|
|
static constexpr llvm::Intrinsic::ID f16IDs[] = {
|
|
llvm::Intrinsic::nvvm_fma_rn_f16,
|
|
llvm::Intrinsic::nvvm_fma_rn_f16x2,
|
|
llvm::Intrinsic::nvvm_fma_rn_ftz_f16,
|
|
llvm::Intrinsic::nvvm_fma_rn_ftz_f16x2,
|
|
llvm::Intrinsic::nvvm_fma_rn_sat_f16,
|
|
llvm::Intrinsic::nvvm_fma_rn_sat_f16x2,
|
|
llvm::Intrinsic::nvvm_fma_rn_ftz_sat_f16,
|
|
llvm::Intrinsic::nvvm_fma_rn_ftz_sat_f16x2,
|
|
llvm::Intrinsic::nvvm_fma_rn_relu_f16,
|
|
llvm::Intrinsic::nvvm_fma_rn_relu_f16x2,
|
|
llvm::Intrinsic::nvvm_fma_rn_ftz_relu_f16,
|
|
llvm::Intrinsic::nvvm_fma_rn_ftz_relu_f16x2};
|
|
|
|
static constexpr llvm::Intrinsic::ID bf16IDs[] = {
|
|
llvm::Intrinsic::nvvm_fma_rn_bf16, llvm::Intrinsic::nvvm_fma_rn_bf16x2,
|
|
llvm::Intrinsic::nvvm_fma_rn_relu_bf16,
|
|
llvm::Intrinsic::nvvm_fma_rn_relu_bf16x2};
|
|
|
|
static constexpr llvm::Intrinsic::ID f32IDs[] = {
|
|
llvm::Intrinsic::nvvm_fma_rn_f,
|
|
llvm::Intrinsic::nvvm_fma_rm_f,
|
|
llvm::Intrinsic::nvvm_fma_rp_f,
|
|
llvm::Intrinsic::nvvm_fma_rz_f,
|
|
llvm::Intrinsic::nvvm_fma_rn_sat_f,
|
|
llvm::Intrinsic::nvvm_fma_rm_sat_f,
|
|
llvm::Intrinsic::nvvm_fma_rp_sat_f,
|
|
llvm::Intrinsic::nvvm_fma_rz_sat_f,
|
|
llvm::Intrinsic::nvvm_fma_rn_ftz_f,
|
|
llvm::Intrinsic::nvvm_fma_rm_ftz_f,
|
|
llvm::Intrinsic::nvvm_fma_rp_ftz_f,
|
|
llvm::Intrinsic::nvvm_fma_rz_ftz_f,
|
|
llvm::Intrinsic::nvvm_fma_rn_ftz_sat_f,
|
|
llvm::Intrinsic::nvvm_fma_rm_ftz_sat_f,
|
|
llvm::Intrinsic::nvvm_fma_rp_ftz_sat_f,
|
|
llvm::Intrinsic::nvvm_fma_rz_ftz_sat_f,
|
|
};
|
|
|
|
static constexpr llvm::Intrinsic::ID f64IDs[] = {
|
|
llvm::Intrinsic::nvvm_fma_rn_d, llvm::Intrinsic::nvvm_fma_rm_d,
|
|
llvm::Intrinsic::nvvm_fma_rp_d, llvm::Intrinsic::nvvm_fma_rz_d};
|
|
|
|
auto fmaIntrinsic = [&](llvm::Intrinsic::ID IID,
|
|
llvm::Type *retType) -> llvm::Value * {
|
|
return createScalarizedIntrinsicCall(
|
|
builder, IID, opTypeLLVM, {argA, argB, argC}, /*retType=*/retType);
|
|
};
|
|
|
|
// f16 + f16 -> f16 / vector<2xf16> + vector<2xf16> -> vector<2xf16>
|
|
if (opTypeLLVM->getScalarType()->isHalfTy()) {
|
|
llvm::Value *result;
|
|
if (isOOB) {
|
|
result = fmaIntrinsic(isRelu ? llvm::Intrinsic::nvvm_fma_rn_oob_relu
|
|
: llvm::Intrinsic::nvvm_fma_rn_oob,
|
|
opTypeLLVM);
|
|
} else {
|
|
unsigned index =
|
|
(isRelu << 3) | (isSat << 2) | (isFTZ << 1) |
|
|
isVectorFma; // Op verifier ensures that this index is valid
|
|
result = fmaIntrinsic(f16IDs[index], opTypeLLVM);
|
|
}
|
|
mt.mapValue(thisOp.getRes(), result);
|
|
return;
|
|
}
|
|
|
|
// bf16 + bf16 -> bf16 / vector<2xbf16> + vector<2xbf16> -> vector<2xbf16>
|
|
if (opTypeLLVM->getScalarType()->isBFloatTy()) {
|
|
llvm::Value *result;
|
|
if (isOOB) {
|
|
result = fmaIntrinsic(isRelu ? llvm::Intrinsic::nvvm_fma_rn_oob_relu
|
|
: llvm::Intrinsic::nvvm_fma_rn_oob,
|
|
opTypeLLVM);
|
|
} else {
|
|
unsigned index = (isRelu << 1) | isVectorFma;
|
|
result = fmaIntrinsic(bf16IDs[index], opTypeLLVM);
|
|
}
|
|
mt.mapValue(thisOp.getRes(), result);
|
|
return;
|
|
}
|
|
|
|
// f64 + f64 -> f64 / vector<2xf64> + vector<2xf64> -> vector<2xf64>
|
|
if (opTypeLLVM->getScalarType()->isDoubleTy()) {
|
|
mt.mapValue(thisOp.getRes(),
|
|
fmaIntrinsic(f64IDs[rndIndex], opTypeLLVM->getScalarType()));
|
|
return;
|
|
}
|
|
|
|
// f32 + f32 -> f32 / vector<2xf32> + vector<2xf32> -> vector<2xf32>
|
|
const unsigned numRndModes = 4; // RN, RM, RP, RZ
|
|
if (opTypeLLVM->getScalarType()->isFloatTy()) {
|
|
unsigned index = ((isFTZ << 1) | isSat) * numRndModes + rndIndex;
|
|
mt.mapValue(thisOp.getRes(),
|
|
fmaIntrinsic(f32IDs[index], opTypeLLVM->getScalarType()));
|
|
return;
|
|
}
|
|
}
|
|
|
|
namespace {
|
|
/// Implementation of the dialect interface that converts operations belonging
|
|
/// to the NVVM dialect to LLVM IR.
|
|
class NVVMDialectLLVMIRTranslationInterface
|
|
: public LLVMTranslationDialectInterface {
|
|
public:
|
|
using LLVMTranslationDialectInterface::LLVMTranslationDialectInterface;
|
|
|
|
/// Translates the given operation to LLVM IR using the provided IR builder
|
|
/// and saving the state in `moduleTranslation`.
|
|
LogicalResult
|
|
convertOperation(Operation *op, llvm::IRBuilderBase &builder,
|
|
LLVM::ModuleTranslation &moduleTranslation) const final {
|
|
// All NVVM ops are instruction-level and require an active insertion point.
|
|
// A null insert block means the op is misplaced (e.g., at module scope),
|
|
// which would otherwise cause a null dereference in createIntrinsicCall.
|
|
if (!builder.GetInsertBlock())
|
|
return op->emitOpError(
|
|
"cannot be translated to LLVM IR without an active insertion "
|
|
"point; make sure the op is inside a function");
|
|
Operation &opInst = *op;
|
|
#include "mlir/Dialect/LLVMIR/NVVMConversions.inc"
|
|
|
|
return failure();
|
|
}
|
|
|
|
/// Attaches module-level metadata for functions marked as kernels
|
|
/// and managed annotations for global variables.
|
|
LogicalResult
|
|
amendOperation(Operation *op, ArrayRef<llvm::Instruction *> instructions,
|
|
NamedAttribute attribute,
|
|
LLVM::ModuleTranslation &moduleTranslation) const final {
|
|
if (auto globalOp = dyn_cast<LLVM::GlobalOp>(op)) {
|
|
if (attribute.getName() == NVVM::NVVMDialect::getManagedAttrName()) {
|
|
auto *gv = cast<llvm::GlobalVariable>(
|
|
moduleTranslation.lookupGlobal(globalOp));
|
|
llvm::Module *m = gv->getParent();
|
|
llvm::LLVMContext &ctx = m->getContext();
|
|
llvm::NamedMDNode *md = m->getOrInsertNamedMetadata("nvvm.annotations");
|
|
md->addOperand(llvm::MDNode::get(
|
|
ctx, {llvm::ConstantAsMetadata::get(gv),
|
|
llvm::MDString::get(ctx, "managed"),
|
|
llvm::ConstantAsMetadata::get(llvm::ConstantInt::get(
|
|
llvm::Type::getInt32Ty(ctx), 1))}));
|
|
}
|
|
return success();
|
|
}
|
|
|
|
auto func = dyn_cast<LLVM::LLVMFuncOp>(op);
|
|
if (!func)
|
|
return failure();
|
|
llvm::Function *llvmFunc = moduleTranslation.lookupFunction(func.getName());
|
|
|
|
if (attribute.getName() == NVVM::NVVMDialect::getMaxntidAttrName()) {
|
|
if (!isa<DenseI32ArrayAttr>(attribute.getValue()))
|
|
return failure();
|
|
auto values = cast<DenseI32ArrayAttr>(attribute.getValue());
|
|
const std::string attr = llvm::formatv(
|
|
"{0:$[,]}", llvm::make_range(values.asArrayRef().begin(),
|
|
values.asArrayRef().end()));
|
|
llvmFunc->addFnAttr(llvm::NVVMAttr::MaxNTID, attr);
|
|
} else if (attribute.getName() == NVVM::NVVMDialect::getReqntidAttrName()) {
|
|
if (!isa<DenseI32ArrayAttr>(attribute.getValue()))
|
|
return failure();
|
|
auto values = cast<DenseI32ArrayAttr>(attribute.getValue());
|
|
const std::string attr = llvm::formatv(
|
|
"{0:$[,]}", llvm::make_range(values.asArrayRef().begin(),
|
|
values.asArrayRef().end()));
|
|
llvmFunc->addFnAttr(llvm::NVVMAttr::ReqNTID, attr);
|
|
} else if (attribute.getName() ==
|
|
NVVM::NVVMDialect::getClusterDimAttrName()) {
|
|
if (!isa<DenseI32ArrayAttr>(attribute.getValue()))
|
|
return failure();
|
|
auto values = cast<DenseI32ArrayAttr>(attribute.getValue());
|
|
const std::string attr = llvm::formatv(
|
|
"{0:$[,]}", llvm::make_range(values.asArrayRef().begin(),
|
|
values.asArrayRef().end()));
|
|
llvmFunc->addFnAttr(llvm::NVVMAttr::ClusterDim, attr);
|
|
} else if (attribute.getName() ==
|
|
NVVM::NVVMDialect::getClusterMaxBlocksAttrName()) {
|
|
auto value = dyn_cast<IntegerAttr>(attribute.getValue());
|
|
llvmFunc->addFnAttr(llvm::NVVMAttr::MaxClusterRank,
|
|
llvm::utostr(value.getInt()));
|
|
} else if (attribute.getName() ==
|
|
NVVM::NVVMDialect::getMinctasmAttrName()) {
|
|
auto value = dyn_cast<IntegerAttr>(attribute.getValue());
|
|
llvmFunc->addFnAttr(llvm::NVVMAttr::MinCTASm,
|
|
llvm::utostr(value.getInt()));
|
|
} else if (attribute.getName() == NVVM::NVVMDialect::getMaxnregAttrName()) {
|
|
auto value = dyn_cast<IntegerAttr>(attribute.getValue());
|
|
llvmFunc->addFnAttr(llvm::NVVMAttr::MaxNReg,
|
|
llvm::utostr(value.getInt()));
|
|
} else if (attribute.getName() ==
|
|
NVVM::NVVMDialect::getKernelFuncAttrName()) {
|
|
llvmFunc->setCallingConv(llvm::CallingConv::PTX_Kernel);
|
|
} else if (attribute.getName() ==
|
|
NVVM::NVVMDialect::getBlocksAreClustersAttrName()) {
|
|
llvmFunc->addFnAttr(llvm::NVVMAttr::BlocksAreClusters);
|
|
}
|
|
|
|
return success();
|
|
}
|
|
|
|
LogicalResult
|
|
convertParameterAttr(LLVMFuncOp funcOp, int argIdx, NamedAttribute attribute,
|
|
LLVM::ModuleTranslation &moduleTranslation) const final {
|
|
|
|
llvm::LLVMContext &llvmContext = moduleTranslation.getLLVMContext();
|
|
llvm::Function *llvmFunc =
|
|
moduleTranslation.lookupFunction(funcOp.getName());
|
|
|
|
if (attribute.getName() == NVVM::NVVMDialect::getGridConstantAttrName()) {
|
|
llvmFunc->addParamAttr(
|
|
argIdx,
|
|
llvm::Attribute::get(llvmContext, llvm::NVVMAttr::GridConstant));
|
|
}
|
|
return success();
|
|
}
|
|
};
|
|
} // namespace
|
|
|
|
void mlir::registerNVVMDialectTranslation(DialectRegistry ®istry) {
|
|
registry.insert<NVVM::NVVMDialect>();
|
|
registry.addExtension(+[](MLIRContext *ctx, NVVM::NVVMDialect *dialect) {
|
|
dialect->addInterfaces<NVVMDialectLLVMIRTranslationInterface>();
|
|
});
|
|
}
|
|
|
|
void mlir::registerNVVMDialectTranslation(MLIRContext &context) {
|
|
DialectRegistry registry;
|
|
registerNVVMDialectTranslation(registry);
|
|
context.appendDialectRegistry(registry);
|
|
}
|