1710 lines · cpp
1//===-- CUDAIntrinsicCall.cpp ---------------------------------------------===//2//3// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.4// See https://llvm.org/LICENSE.txt for license information.5// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception6//7//===----------------------------------------------------------------------===//8//9// Helper routines for constructing the FIR dialect of MLIR for PowerPC10// intrinsics. Extensive use of MLIR interfaces and MLIR's coding style11// (https://mlir.llvm.org/getting_started/DeveloperGuide/) is used in this12// module.13//14//===----------------------------------------------------------------------===//15 16#include "flang/Optimizer/Builder/CUDAIntrinsicCall.h"17#include "flang/Evaluate/common.h"18#include "flang/Optimizer/Builder/FIRBuilder.h"19#include "flang/Optimizer/Builder/MutableBox.h"20#include "mlir/Dialect/Index/IR/IndexOps.h"21#include "mlir/Dialect/SCF/IR/SCF.h"22#include "mlir/Dialect/Vector/IR/VectorOps.h"23 24namespace fir {25 26using CI = CUDAIntrinsicLibrary;27 28static const char __ldca_i4x4[] = "__ldca_i4x4_";29static const char __ldca_i8x2[] = "__ldca_i8x2_";30static const char __ldca_r2x2[] = "__ldca_r2x2_";31static const char __ldca_r4x4[] = "__ldca_r4x4_";32static const char __ldca_r8x2[] = "__ldca_r8x2_";33static const char __ldcg_i4x4[] = "__ldcg_i4x4_";34static const char __ldcg_i8x2[] = "__ldcg_i8x2_";35static const char __ldcg_r2x2[] = "__ldcg_r2x2_";36static const char __ldcg_r4x4[] = "__ldcg_r4x4_";37static const char __ldcg_r8x2[] = "__ldcg_r8x2_";38static const char __ldcs_i4x4[] = "__ldcs_i4x4_";39static const char __ldcs_i8x2[] = "__ldcs_i8x2_";40static const char __ldcs_r2x2[] = "__ldcs_r2x2_";41static const char __ldcs_r4x4[] = "__ldcs_r4x4_";42static const char __ldcs_r8x2[] = "__ldcs_r8x2_";43static const char __ldcv_i4x4[] = "__ldcv_i4x4_";44static const char __ldcv_i8x2[] = "__ldcv_i8x2_";45static const char __ldcv_r2x2[] = "__ldcv_r2x2_";46static const char __ldcv_r4x4[] = "__ldcv_r4x4_";47static const char __ldcv_r8x2[] = "__ldcv_r8x2_";48static const char __ldlu_i4x4[] = "__ldlu_i4x4_";49static const char __ldlu_i8x2[] = "__ldlu_i8x2_";50static const char __ldlu_r2x2[] = "__ldlu_r2x2_";51static const char __ldlu_r4x4[] = "__ldlu_r4x4_";52static const char __ldlu_r8x2[] = "__ldlu_r8x2_";53 54// CUDA specific intrinsic handlers.55static constexpr IntrinsicHandler cudaHandlers[]{56 {"__ldca_i4x4",57 static_cast<CUDAIntrinsicLibrary::ExtendedGenerator>(58 &CI::genLDXXFunc<__ldca_i4x4, 4>),59 {{{"a", asAddr}}},60 /*isElemental=*/false},61 {"__ldca_i8x2",62 static_cast<CUDAIntrinsicLibrary::ExtendedGenerator>(63 &CI::genLDXXFunc<__ldca_i8x2, 2>),64 {{{"a", asAddr}}},65 /*isElemental=*/false},66 {"__ldca_r2x2",67 static_cast<CUDAIntrinsicLibrary::ExtendedGenerator>(68 &CI::genLDXXFunc<__ldca_r2x2, 2>),69 {{{"a", asAddr}}},70 /*isElemental=*/false},71 {"__ldca_r4x4",72 static_cast<CUDAIntrinsicLibrary::ExtendedGenerator>(73 &CI::genLDXXFunc<__ldca_r4x4, 4>),74 {{{"a", asAddr}}},75 /*isElemental=*/false},76 {"__ldca_r8x2",77 static_cast<CUDAIntrinsicLibrary::ExtendedGenerator>(78 &CI::genLDXXFunc<__ldca_r8x2, 2>),79 {{{"a", asAddr}}},80 /*isElemental=*/false},81 {"__ldcg_i4x4",82 static_cast<CUDAIntrinsicLibrary::ExtendedGenerator>(83 &CI::genLDXXFunc<__ldcg_i4x4, 4>),84 {{{"a", asAddr}}},85 /*isElemental=*/false},86 {"__ldcg_i8x2",87 static_cast<CUDAIntrinsicLibrary::ExtendedGenerator>(88 &CI::genLDXXFunc<__ldcg_i8x2, 2>),89 {{{"a", asAddr}}},90 /*isElemental=*/false},91 {"__ldcg_r2x2",92 static_cast<CUDAIntrinsicLibrary::ExtendedGenerator>(93 &CI::genLDXXFunc<__ldcg_r2x2, 2>),94 {{{"a", asAddr}}},95 /*isElemental=*/false},96 {"__ldcg_r4x4",97 static_cast<CUDAIntrinsicLibrary::ExtendedGenerator>(98 &CI::genLDXXFunc<__ldcg_r4x4, 4>),99 {{{"a", asAddr}}},100 /*isElemental=*/false},101 {"__ldcg_r8x2",102 static_cast<CUDAIntrinsicLibrary::ExtendedGenerator>(103 &CI::genLDXXFunc<__ldcg_r8x2, 2>),104 {{{"a", asAddr}}},105 /*isElemental=*/false},106 {"__ldcs_i4x4",107 static_cast<CUDAIntrinsicLibrary::ExtendedGenerator>(108 &CI::genLDXXFunc<__ldcs_i4x4, 4>),109 {{{"a", asAddr}}},110 /*isElemental=*/false},111 {"__ldcs_i8x2",112 static_cast<CUDAIntrinsicLibrary::ExtendedGenerator>(113 &CI::genLDXXFunc<__ldcs_i8x2, 2>),114 {{{"a", asAddr}}},115 /*isElemental=*/false},116 {"__ldcs_r2x2",117 static_cast<CUDAIntrinsicLibrary::ExtendedGenerator>(118 &CI::genLDXXFunc<__ldcs_r2x2, 2>),119 {{{"a", asAddr}}},120 /*isElemental=*/false},121 {"__ldcs_r4x4",122 static_cast<CUDAIntrinsicLibrary::ExtendedGenerator>(123 &CI::genLDXXFunc<__ldcs_r4x4, 4>),124 {{{"a", asAddr}}},125 /*isElemental=*/false},126 {"__ldcs_r8x2",127 static_cast<CUDAIntrinsicLibrary::ExtendedGenerator>(128 &CI::genLDXXFunc<__ldcs_r8x2, 2>),129 {{{"a", asAddr}}},130 /*isElemental=*/false},131 {"__ldcv_i4x4",132 static_cast<CUDAIntrinsicLibrary::ExtendedGenerator>(133 &CI::genLDXXFunc<__ldcv_i4x4, 4>),134 {{{"a", asAddr}}},135 /*isElemental=*/false},136 {"__ldcv_i8x2",137 static_cast<CUDAIntrinsicLibrary::ExtendedGenerator>(138 &CI::genLDXXFunc<__ldcv_i8x2, 2>),139 {{{"a", asAddr}}},140 /*isElemental=*/false},141 {"__ldcv_r2x2",142 static_cast<CUDAIntrinsicLibrary::ExtendedGenerator>(143 &CI::genLDXXFunc<__ldcv_r2x2, 2>),144 {{{"a", asAddr}}},145 /*isElemental=*/false},146 {"__ldcv_r4x4",147 static_cast<CUDAIntrinsicLibrary::ExtendedGenerator>(148 &CI::genLDXXFunc<__ldcv_r4x4, 4>),149 {{{"a", asAddr}}},150 /*isElemental=*/false},151 {"__ldcv_r8x2",152 static_cast<CUDAIntrinsicLibrary::ExtendedGenerator>(153 &CI::genLDXXFunc<__ldcv_r8x2, 2>),154 {{{"a", asAddr}}},155 /*isElemental=*/false},156 {"__ldlu_i4x4",157 static_cast<CUDAIntrinsicLibrary::ExtendedGenerator>(158 &CI::genLDXXFunc<__ldlu_i4x4, 4>),159 {{{"a", asAddr}}},160 /*isElemental=*/false},161 {"__ldlu_i8x2",162 static_cast<CUDAIntrinsicLibrary::ExtendedGenerator>(163 &CI::genLDXXFunc<__ldlu_i8x2, 2>),164 {{{"a", asAddr}}},165 /*isElemental=*/false},166 {"__ldlu_r2x2",167 static_cast<CUDAIntrinsicLibrary::ExtendedGenerator>(168 &CI::genLDXXFunc<__ldlu_r2x2, 2>),169 {{{"a", asAddr}}},170 /*isElemental=*/false},171 {"__ldlu_r4x4",172 static_cast<CUDAIntrinsicLibrary::ExtendedGenerator>(173 &CI::genLDXXFunc<__ldlu_r4x4, 4>),174 {{{"a", asAddr}}},175 /*isElemental=*/false},176 {"__ldlu_r8x2",177 static_cast<CUDAIntrinsicLibrary::ExtendedGenerator>(178 &CI::genLDXXFunc<__ldlu_r8x2, 2>),179 {{{"a", asAddr}}},180 /*isElemental=*/false},181 {"all_sync",182 static_cast<CUDAIntrinsicLibrary::ElementalGenerator>(183 &CI::genVoteSync<mlir::NVVM::VoteSyncKind::all>),184 {{{"mask", asValue}, {"pred", asValue}}},185 /*isElemental=*/false},186 {"any_sync",187 static_cast<CUDAIntrinsicLibrary::ElementalGenerator>(188 &CI::genVoteSync<mlir::NVVM::VoteSyncKind::any>),189 {{{"mask", asValue}, {"pred", asValue}}},190 /*isElemental=*/false},191 {"atomicadd_r4x2",192 static_cast<CUDAIntrinsicLibrary::ExtendedGenerator>(193 &CI::genAtomicAddVector<2>),194 {{{"a", asAddr}, {"v", asAddr}}},195 false},196 {"atomicadd_r4x4",197 static_cast<CUDAIntrinsicLibrary::ExtendedGenerator>(198 &CI::genAtomicAddVector4x4),199 {{{"a", asAddr}, {"v", asAddr}}},200 false},201 {"atomicaddd",202 static_cast<CUDAIntrinsicLibrary::ElementalGenerator>(&CI::genAtomicAdd),203 {{{"a", asAddr}, {"v", asValue}}},204 false},205 {"atomicaddf",206 static_cast<CUDAIntrinsicLibrary::ElementalGenerator>(&CI::genAtomicAdd),207 {{{"a", asAddr}, {"v", asValue}}},208 false},209 {"atomicaddi",210 static_cast<CUDAIntrinsicLibrary::ElementalGenerator>(&CI::genAtomicAdd),211 {{{"a", asAddr}, {"v", asValue}}},212 false},213 {"atomicaddl",214 static_cast<CUDAIntrinsicLibrary::ElementalGenerator>(&CI::genAtomicAdd),215 {{{"a", asAddr}, {"v", asValue}}},216 false},217 {"atomicaddr2",218 static_cast<CUDAIntrinsicLibrary::ExtendedGenerator>(&CI::genAtomicAddR2),219 {{{"a", asAddr}, {"v", asAddr}}},220 false},221 {"atomicaddvector_r2x2",222 static_cast<CUDAIntrinsicLibrary::ExtendedGenerator>(223 &CI::genAtomicAddVector<2>),224 {{{"a", asAddr}, {"v", asAddr}}},225 false},226 {"atomicaddvector_r4x2",227 static_cast<CUDAIntrinsicLibrary::ExtendedGenerator>(228 &CI::genAtomicAddVector<2>),229 {{{"a", asAddr}, {"v", asAddr}}},230 false},231 {"atomicandi",232 static_cast<CUDAIntrinsicLibrary::ElementalGenerator>(&CI::genAtomicAnd),233 {{{"a", asAddr}, {"v", asValue}}},234 false},235 {"atomiccasd",236 static_cast<CUDAIntrinsicLibrary::ExtendedGenerator>(&CI::genAtomicCas),237 {{{"a", asAddr}, {"v1", asValue}, {"v2", asValue}}},238 false},239 {"atomiccasf",240 static_cast<CUDAIntrinsicLibrary::ExtendedGenerator>(&CI::genAtomicCas),241 {{{"a", asAddr}, {"v1", asValue}, {"v2", asValue}}},242 false},243 {"atomiccasi",244 static_cast<CUDAIntrinsicLibrary::ExtendedGenerator>(&CI::genAtomicCas),245 {{{"a", asAddr}, {"v1", asValue}, {"v2", asValue}}},246 false},247 {"atomiccasul",248 static_cast<CUDAIntrinsicLibrary::ExtendedGenerator>(&CI::genAtomicCas),249 {{{"a", asAddr}, {"v1", asValue}, {"v2", asValue}}},250 false},251 {"atomicdeci",252 static_cast<CUDAIntrinsicLibrary::ElementalGenerator>(&CI::genAtomicDec),253 {{{"a", asAddr}, {"v", asValue}}},254 false},255 {"atomicexchd",256 static_cast<CUDAIntrinsicLibrary::ExtendedGenerator>(&CI::genAtomicExch),257 {{{"a", asAddr}, {"v", asValue}}},258 false},259 {"atomicexchf",260 static_cast<CUDAIntrinsicLibrary::ExtendedGenerator>(&CI::genAtomicExch),261 {{{"a", asAddr}, {"v", asValue}}},262 false},263 {"atomicexchi",264 static_cast<CUDAIntrinsicLibrary::ExtendedGenerator>(&CI::genAtomicExch),265 {{{"a", asAddr}, {"v", asValue}}},266 false},267 {"atomicexchul",268 static_cast<CUDAIntrinsicLibrary::ExtendedGenerator>(&CI::genAtomicExch),269 {{{"a", asAddr}, {"v", asValue}}},270 false},271 {"atomicinci",272 static_cast<CUDAIntrinsicLibrary::ElementalGenerator>(&CI::genAtomicInc),273 {{{"a", asAddr}, {"v", asValue}}},274 false},275 {"atomicmaxd",276 static_cast<CUDAIntrinsicLibrary::ElementalGenerator>(&CI::genAtomicMax),277 {{{"a", asAddr}, {"v", asValue}}},278 false},279 {"atomicmaxf",280 static_cast<CUDAIntrinsicLibrary::ElementalGenerator>(&CI::genAtomicMax),281 {{{"a", asAddr}, {"v", asValue}}},282 false},283 {"atomicmaxi",284 static_cast<CUDAIntrinsicLibrary::ElementalGenerator>(&CI::genAtomicMax),285 {{{"a", asAddr}, {"v", asValue}}},286 false},287 {"atomicmaxl",288 static_cast<CUDAIntrinsicLibrary::ElementalGenerator>(&CI::genAtomicMax),289 {{{"a", asAddr}, {"v", asValue}}},290 false},291 {"atomicmind",292 static_cast<CUDAIntrinsicLibrary::ElementalGenerator>(&CI::genAtomicMin),293 {{{"a", asAddr}, {"v", asValue}}},294 false},295 {"atomicminf",296 static_cast<CUDAIntrinsicLibrary::ElementalGenerator>(&CI::genAtomicMin),297 {{{"a", asAddr}, {"v", asValue}}},298 false},299 {"atomicmini",300 static_cast<CUDAIntrinsicLibrary::ElementalGenerator>(&CI::genAtomicMin),301 {{{"a", asAddr}, {"v", asValue}}},302 false},303 {"atomicminl",304 static_cast<CUDAIntrinsicLibrary::ElementalGenerator>(&CI::genAtomicMin),305 {{{"a", asAddr}, {"v", asValue}}},306 false},307 {"atomicori",308 static_cast<CUDAIntrinsicLibrary::ElementalGenerator>(&CI::genAtomicOr),309 {{{"a", asAddr}, {"v", asValue}}},310 false},311 {"atomicsubd",312 static_cast<CUDAIntrinsicLibrary::ElementalGenerator>(&CI::genAtomicSub),313 {{{"a", asAddr}, {"v", asValue}}},314 false},315 {"atomicsubf",316 static_cast<CUDAIntrinsicLibrary::ElementalGenerator>(&CI::genAtomicSub),317 {{{"a", asAddr}, {"v", asValue}}},318 false},319 {"atomicsubi",320 static_cast<CUDAIntrinsicLibrary::ElementalGenerator>(&CI::genAtomicSub),321 {{{"a", asAddr}, {"v", asValue}}},322 false},323 {"atomicsubl",324 static_cast<CUDAIntrinsicLibrary::ElementalGenerator>(&CI::genAtomicSub),325 {{{"a", asAddr}, {"v", asValue}}},326 false},327 {"atomicxori",328 static_cast<CUDAIntrinsicLibrary::ExtendedGenerator>(&CI::genAtomicXor),329 {{{"a", asAddr}, {"v", asValue}}},330 false},331 {"ballot_sync",332 static_cast<CUDAIntrinsicLibrary::ElementalGenerator>(333 &CI::genVoteSync<mlir::NVVM::VoteSyncKind::ballot>),334 {{{"mask", asValue}, {"pred", asValue}}},335 /*isElemental=*/false},336 {"barrier_arrive",337 static_cast<CUDAIntrinsicLibrary::ElementalGenerator>(338 &CI::genBarrierArrive),339 {{{"barrier", asAddr}}},340 /*isElemental=*/false},341 {"barrier_arrive_cnt",342 static_cast<CUDAIntrinsicLibrary::ElementalGenerator>(343 &CI::genBarrierArriveCnt),344 {{{"barrier", asAddr}, {"count", asValue}}},345 /*isElemental=*/false},346 {"barrier_init",347 static_cast<CUDAIntrinsicLibrary::SubroutineGenerator>(348 &CI::genBarrierInit),349 {{{"barrier", asAddr}, {"count", asValue}}},350 /*isElemental=*/false},351 {"barrier_try_wait",352 static_cast<CUDAIntrinsicLibrary::ElementalGenerator>(353 &CI::genBarrierTryWait),354 {{{"barrier", asAddr}, {"token", asValue}}},355 /*isElemental=*/false},356 {"barrier_try_wait_sleep",357 static_cast<CUDAIntrinsicLibrary::ElementalGenerator>(358 &CI::genBarrierTryWaitSleep),359 {{{"barrier", asAddr}, {"token", asValue}, {"ns", asValue}}},360 /*isElemental=*/false},361 {"clock",362 static_cast<CUDAIntrinsicLibrary::ElementalGenerator>(363 &CI::genNVVMTime<mlir::NVVM::ClockOp>),364 {},365 /*isElemental=*/false},366 {"clock64",367 static_cast<CUDAIntrinsicLibrary::ElementalGenerator>(368 &CI::genNVVMTime<mlir::NVVM::Clock64Op>),369 {},370 /*isElemental=*/false},371 {"cluster_block_index",372 static_cast<CUDAIntrinsicLibrary::ElementalGenerator>(373 &CI::genClusterBlockIndex),374 {},375 /*isElemental=*/false},376 {"cluster_dim_blocks",377 static_cast<CUDAIntrinsicLibrary::ElementalGenerator>(378 &CI::genClusterDimBlocks),379 {},380 /*isElemental=*/false},381 {"fence_proxy_async",382 static_cast<CUDAIntrinsicLibrary::SubroutineGenerator>(383 &CI::genFenceProxyAsync),384 {},385 /*isElemental=*/false},386 {"globaltimer",387 static_cast<CUDAIntrinsicLibrary::ElementalGenerator>(388 &CI::genNVVMTime<mlir::NVVM::GlobalTimerOp>),389 {},390 /*isElemental=*/false},391 {"match_all_syncjd",392 static_cast<CUDAIntrinsicLibrary::ElementalGenerator>(393 &CI::genMatchAllSync),394 {{{"mask", asValue}, {"value", asValue}, {"pred", asAddr}}},395 /*isElemental=*/false},396 {"match_all_syncjf",397 static_cast<CUDAIntrinsicLibrary::ElementalGenerator>(398 &CI::genMatchAllSync),399 {{{"mask", asValue}, {"value", asValue}, {"pred", asAddr}}},400 /*isElemental=*/false},401 {"match_all_syncjj",402 static_cast<CUDAIntrinsicLibrary::ElementalGenerator>(403 &CI::genMatchAllSync),404 {{{"mask", asValue}, {"value", asValue}, {"pred", asAddr}}},405 /*isElemental=*/false},406 {"match_all_syncjx",407 static_cast<CUDAIntrinsicLibrary::ElementalGenerator>(408 &CI::genMatchAllSync),409 {{{"mask", asValue}, {"value", asValue}, {"pred", asAddr}}},410 /*isElemental=*/false},411 {"match_any_syncjd",412 static_cast<CUDAIntrinsicLibrary::ElementalGenerator>(413 &CI::genMatchAnySync),414 {{{"mask", asValue}, {"value", asValue}}},415 /*isElemental=*/false},416 {"match_any_syncjf",417 static_cast<CUDAIntrinsicLibrary::ElementalGenerator>(418 &CI::genMatchAnySync),419 {{{"mask", asValue}, {"value", asValue}}},420 /*isElemental=*/false},421 {"match_any_syncjj",422 static_cast<CUDAIntrinsicLibrary::ElementalGenerator>(423 &CI::genMatchAnySync),424 {{{"mask", asValue}, {"value", asValue}}},425 /*isElemental=*/false},426 {"match_any_syncjx",427 static_cast<CUDAIntrinsicLibrary::ElementalGenerator>(428 &CI::genMatchAnySync),429 {{{"mask", asValue}, {"value", asValue}}},430 /*isElemental=*/false},431 {"syncthreads",432 static_cast<CUDAIntrinsicLibrary::SubroutineGenerator>(433 &CI::genSyncThreads),434 {},435 /*isElemental=*/false},436 {"syncthreads_and_i4",437 static_cast<CUDAIntrinsicLibrary::ElementalGenerator>(438 &CI::genSyncThreadsAnd),439 {},440 /*isElemental=*/false},441 {"syncthreads_and_l4",442 static_cast<CUDAIntrinsicLibrary::ElementalGenerator>(443 &CI::genSyncThreadsAnd),444 {},445 /*isElemental=*/false},446 {"syncthreads_count_i4",447 static_cast<CUDAIntrinsicLibrary::ElementalGenerator>(448 &CI::genSyncThreadsCount),449 {},450 /*isElemental=*/false},451 {"syncthreads_count_l4",452 static_cast<CUDAIntrinsicLibrary::ElementalGenerator>(453 &CI::genSyncThreadsCount),454 {},455 /*isElemental=*/false},456 {"syncthreads_or_i4",457 static_cast<CUDAIntrinsicLibrary::ElementalGenerator>(458 &CI::genSyncThreadsOr),459 {},460 /*isElemental=*/false},461 {"syncthreads_or_l4",462 static_cast<CUDAIntrinsicLibrary::ElementalGenerator>(463 &CI::genSyncThreadsOr),464 {},465 /*isElemental=*/false},466 {"syncwarp",467 static_cast<CUDAIntrinsicLibrary::SubroutineGenerator>(&CI::genSyncWarp),468 {},469 /*isElemental=*/false},470 {"this_cluster",471 static_cast<CUDAIntrinsicLibrary::ElementalGenerator>(&CI::genThisCluster),472 {},473 /*isElemental=*/false},474 {"this_grid",475 static_cast<CUDAIntrinsicLibrary::ElementalGenerator>(&CI::genThisGrid),476 {},477 /*isElemental=*/false},478 {"this_thread_block",479 static_cast<CUDAIntrinsicLibrary::ElementalGenerator>(480 &CI::genThisThreadBlock),481 {},482 /*isElemental=*/false},483 {"this_warp",484 static_cast<CUDAIntrinsicLibrary::ElementalGenerator>(&CI::genThisWarp),485 {},486 /*isElemental=*/false},487 {"threadfence",488 static_cast<CUDAIntrinsicLibrary::SubroutineGenerator>(489 &CI::genThreadFence<mlir::NVVM::MemScopeKind::GPU>),490 {},491 /*isElemental=*/false},492 {"threadfence_block",493 static_cast<CUDAIntrinsicLibrary::SubroutineGenerator>(494 &CI::genThreadFence<mlir::NVVM::MemScopeKind::CTA>),495 {},496 /*isElemental=*/false},497 {"threadfence_system",498 static_cast<CUDAIntrinsicLibrary::SubroutineGenerator>(499 &CI::genThreadFence<mlir::NVVM::MemScopeKind::SYS>),500 {},501 /*isElemental=*/false},502 {"tma_bulk_commit_group",503 static_cast<CUDAIntrinsicLibrary::SubroutineGenerator>(504 &CI::genTMABulkCommitGroup),505 {{}},506 /*isElemental=*/false},507 {"tma_bulk_g2s",508 static_cast<CUDAIntrinsicLibrary::SubroutineGenerator>(&CI::genTMABulkG2S),509 {{{"barrier", asAddr},510 {"src", asAddr},511 {"dst", asAddr},512 {"nbytes", asValue}}},513 /*isElemental=*/false},514 {"tma_bulk_ldc4",515 static_cast<CUDAIntrinsicLibrary::SubroutineGenerator>(516 &CI::genTMABulkLoadC4),517 {{{"barrier", asAddr},518 {"src", asAddr},519 {"dst", asAddr},520 {"nelems", asValue}}},521 /*isElemental=*/false},522 {"tma_bulk_ldc8",523 static_cast<CUDAIntrinsicLibrary::SubroutineGenerator>(524 &CI::genTMABulkLoadC8),525 {{{"barrier", asAddr},526 {"src", asAddr},527 {"dst", asAddr},528 {"nelems", asValue}}},529 /*isElemental=*/false},530 {"tma_bulk_ldi4",531 static_cast<CUDAIntrinsicLibrary::SubroutineGenerator>(532 &CI::genTMABulkLoadI4),533 {{{"barrier", asAddr},534 {"src", asAddr},535 {"dst", asAddr},536 {"nelems", asValue}}},537 /*isElemental=*/false},538 {"tma_bulk_ldi8",539 static_cast<CUDAIntrinsicLibrary::SubroutineGenerator>(540 &CI::genTMABulkLoadI8),541 {{{"barrier", asAddr},542 {"src", asAddr},543 {"dst", asAddr},544 {"nelems", asValue}}},545 /*isElemental=*/false},546 {"tma_bulk_ldr2",547 static_cast<CUDAIntrinsicLibrary::SubroutineGenerator>(548 &CI::genTMABulkLoadR2),549 {{{"barrier", asAddr},550 {"src", asAddr},551 {"dst", asAddr},552 {"nelems", asValue}}},553 /*isElemental=*/false},554 {"tma_bulk_ldr4",555 static_cast<CUDAIntrinsicLibrary::SubroutineGenerator>(556 &CI::genTMABulkLoadR4),557 {{{"barrier", asAddr},558 {"src", asAddr},559 {"dst", asAddr},560 {"nelems", asValue}}},561 /*isElemental=*/false},562 {"tma_bulk_ldr8",563 static_cast<CUDAIntrinsicLibrary::SubroutineGenerator>(564 &CI::genTMABulkLoadR8),565 {{{"barrier", asAddr},566 {"src", asAddr},567 {"dst", asAddr},568 {"nelems", asValue}}},569 /*isElemental=*/false},570 {"tma_bulk_s2g",571 static_cast<CUDAIntrinsicLibrary::SubroutineGenerator>(&CI::genTMABulkS2G),572 {{{"src", asAddr}, {"dst", asAddr}, {"nbytes", asValue}}},573 /*isElemental=*/false},574 {"tma_bulk_store_c4",575 static_cast<CUDAIntrinsicLibrary::SubroutineGenerator>(576 &CI::genTMABulkStoreC4),577 {{{"src", asAddr}, {"dst", asAddr}, {"count", asValue}}},578 /*isElemental=*/false},579 {"tma_bulk_store_c8",580 static_cast<CUDAIntrinsicLibrary::SubroutineGenerator>(581 &CI::genTMABulkStoreC8),582 {{{"src", asAddr}, {"dst", asAddr}, {"count", asValue}}},583 /*isElemental=*/false},584 {"tma_bulk_store_i4",585 static_cast<CUDAIntrinsicLibrary::SubroutineGenerator>(586 &CI::genTMABulkStoreI4),587 {{{"src", asAddr}, {"dst", asAddr}, {"count", asValue}}},588 /*isElemental=*/false},589 {"tma_bulk_store_i8",590 static_cast<CUDAIntrinsicLibrary::SubroutineGenerator>(591 &CI::genTMABulkStoreI8),592 {{{"src", asAddr}, {"dst", asAddr}, {"count", asValue}}},593 /*isElemental=*/false},594 {"tma_bulk_store_r2",595 static_cast<CUDAIntrinsicLibrary::SubroutineGenerator>(596 &CI::genTMABulkStoreR2),597 {{{"src", asAddr}, {"dst", asAddr}, {"count", asValue}}},598 /*isElemental=*/false},599 {"tma_bulk_store_r4",600 static_cast<CUDAIntrinsicLibrary::SubroutineGenerator>(601 &CI::genTMABulkStoreR4),602 {{{"src", asAddr}, {"dst", asAddr}, {"count", asValue}}},603 /*isElemental=*/false},604 {"tma_bulk_store_r8",605 static_cast<CUDAIntrinsicLibrary::SubroutineGenerator>(606 &CI::genTMABulkStoreR8),607 {{{"src", asAddr}, {"dst", asAddr}, {"count", asValue}}},608 /*isElemental=*/false},609 {"tma_bulk_wait_group",610 static_cast<CUDAIntrinsicLibrary::SubroutineGenerator>(611 &CI::genTMABulkWaitGroup),612 {{}},613 /*isElemental=*/false},614};615 616template <std::size_t N>617static constexpr bool isSorted(const IntrinsicHandler (&array)[N]) {618 // Replace by std::sorted when C++20 is default (will be constexpr).619 const IntrinsicHandler *lastSeen{nullptr};620 bool isSorted{true};621 for (const auto &x : array) {622 if (lastSeen)623 isSorted &= std::string_view{lastSeen->name} < std::string_view{x.name};624 lastSeen = &x;625 }626 return isSorted;627}628static_assert(isSorted(cudaHandlers) && "map must be sorted");629 630const IntrinsicHandler *findCUDAIntrinsicHandler(llvm::StringRef name) {631 auto compare = [](const IntrinsicHandler &cudaHandler, llvm::StringRef name) {632 return name.compare(cudaHandler.name) > 0;633 };634 auto result = llvm::lower_bound(cudaHandlers, name, compare);635 return result != std::end(cudaHandlers) && result->name == name ? result636 : nullptr;637}638 639static mlir::Value convertPtrToNVVMSpace(fir::FirOpBuilder &builder,640 mlir::Location loc,641 mlir::Value barrier,642 mlir::NVVM::NVVMMemorySpace space) {643 mlir::Value llvmPtr = fir::ConvertOp::create(644 builder, loc, mlir::LLVM::LLVMPointerType::get(builder.getContext()),645 barrier);646 mlir::Value addrCast = mlir::LLVM::AddrSpaceCastOp::create(647 builder, loc,648 mlir::LLVM::LLVMPointerType::get(builder.getContext(),649 static_cast<unsigned>(space)),650 llvmPtr);651 return addrCast;652}653 654static mlir::Value genAtomBinOp(fir::FirOpBuilder &builder, mlir::Location &loc,655 mlir::LLVM::AtomicBinOp binOp, mlir::Value arg0,656 mlir::Value arg1) {657 auto llvmPointerType = mlir::LLVM::LLVMPointerType::get(builder.getContext());658 arg0 = builder.createConvert(loc, llvmPointerType, arg0);659 return mlir::LLVM::AtomicRMWOp::create(builder, loc, binOp, arg0, arg1,660 mlir::LLVM::AtomicOrdering::seq_cst);661}662 663// ATOMICADD664mlir::Value665CUDAIntrinsicLibrary::genAtomicAdd(mlir::Type resultType,666 llvm::ArrayRef<mlir::Value> args) {667 assert(args.size() == 2);668 mlir::LLVM::AtomicBinOp binOp =669 mlir::isa<mlir::IntegerType>(args[1].getType())670 ? mlir::LLVM::AtomicBinOp::add671 : mlir::LLVM::AtomicBinOp::fadd;672 return genAtomBinOp(builder, loc, binOp, args[0], args[1]);673}674 675fir::ExtendedValue676CUDAIntrinsicLibrary::genAtomicAddR2(mlir::Type resultType,677 llvm::ArrayRef<fir::ExtendedValue> args) {678 assert(args.size() == 2);679 680 mlir::Value a = fir::getBase(args[0]);681 682 if (mlir::isa<fir::BaseBoxType>(a.getType())) {683 a = fir::BoxAddrOp::create(builder, loc, a);684 }685 686 auto loc = builder.getUnknownLoc();687 auto f16Ty = builder.getF16Type();688 auto i32Ty = builder.getI32Type();689 auto vecF16Ty = mlir::VectorType::get({2}, f16Ty);690 mlir::Type idxTy = builder.getIndexType();691 auto f16RefTy = fir::ReferenceType::get(f16Ty);692 auto zero = builder.createIntegerConstant(loc, idxTy, 0);693 auto one = builder.createIntegerConstant(loc, idxTy, 1);694 auto v1Coord = fir::CoordinateOp::create(builder, loc, f16RefTy,695 fir::getBase(args[1]), zero);696 auto v2Coord = fir::CoordinateOp::create(builder, loc, f16RefTy,697 fir::getBase(args[1]), one);698 auto v1 = fir::LoadOp::create(builder, loc, v1Coord);699 auto v2 = fir::LoadOp::create(builder, loc, v2Coord);700 mlir::Value undef = mlir::LLVM::UndefOp::create(builder, loc, vecF16Ty);701 mlir::Value vec1 = mlir::LLVM::InsertElementOp::create(702 builder, loc, undef, v1, builder.createIntegerConstant(loc, i32Ty, 0));703 mlir::Value vec2 = mlir::LLVM::InsertElementOp::create(704 builder, loc, vec1, v2, builder.createIntegerConstant(loc, i32Ty, 1));705 auto res = genAtomBinOp(builder, loc, mlir::LLVM::AtomicBinOp::fadd, a, vec2);706 auto i32VecTy = mlir::VectorType::get({1}, i32Ty);707 mlir::Value vecI32 =708 mlir::vector::BitCastOp::create(builder, loc, i32VecTy, res);709 return mlir::vector::ExtractOp::create(builder, loc, vecI32,710 mlir::ArrayRef<int64_t>{0});711}712 713// ATOMICADDVECTOR714template <int extent>715fir::ExtendedValue CUDAIntrinsicLibrary::genAtomicAddVector(716 mlir::Type resultType, llvm::ArrayRef<fir::ExtendedValue> args) {717 assert(args.size() == 2);718 mlir::Value res = fir::AllocaOp::create(719 builder, loc, fir::SequenceType::get({extent}, resultType));720 mlir::Value a = fir::getBase(args[0]);721 if (mlir::isa<fir::BaseBoxType>(a.getType())) {722 a = fir::BoxAddrOp::create(builder, loc, a);723 }724 auto vecTy = mlir::VectorType::get({extent}, resultType);725 auto refTy = fir::ReferenceType::get(resultType);726 mlir::Type i32Ty = builder.getI32Type();727 mlir::Type idxTy = builder.getIndexType();728 729 // Extract the values from the array.730 llvm::SmallVector<mlir::Value> values;731 for (unsigned i = 0; i < extent; ++i) {732 mlir::Value pos = builder.createIntegerConstant(loc, idxTy, i);733 mlir::Value coord = fir::CoordinateOp::create(builder, loc, refTy,734 fir::getBase(args[1]), pos);735 mlir::Value value = fir::LoadOp::create(builder, loc, coord);736 values.push_back(value);737 }738 // Pack extracted values into a vector to call the atomic add.739 mlir::Value undef = mlir::LLVM::UndefOp::create(builder, loc, vecTy);740 for (unsigned i = 0; i < extent; ++i) {741 mlir::Value insert = mlir::LLVM::InsertElementOp::create(742 builder, loc, undef, values[i],743 builder.createIntegerConstant(loc, i32Ty, i));744 undef = insert;745 }746 // Atomic operation with a vector of values.747 mlir::Value add =748 genAtomBinOp(builder, loc, mlir::LLVM::AtomicBinOp::fadd, a, undef);749 // Store results in the result array.750 for (unsigned i = 0; i < extent; ++i) {751 mlir::Value r = mlir::LLVM::ExtractElementOp::create(752 builder, loc, add, builder.createIntegerConstant(loc, i32Ty, i));753 mlir::Value c = fir::CoordinateOp::create(754 builder, loc, refTy, res, builder.createIntegerConstant(loc, idxTy, i));755 fir::StoreOp::create(builder, loc, r, c);756 }757 mlir::Value ext = builder.createIntegerConstant(loc, idxTy, extent);758 return fir::ArrayBoxValue(res, {ext});759}760 761// ATOMICADDVECTOR4x4762fir::ExtendedValue CUDAIntrinsicLibrary::genAtomicAddVector4x4(763 mlir::Type resultType, llvm::ArrayRef<fir::ExtendedValue> args) {764 assert(args.size() == 2);765 mlir::Value a = fir::getBase(args[0]);766 if (mlir::isa<fir::BaseBoxType>(a.getType()))767 a = fir::BoxAddrOp::create(builder, loc, a);768 769 const unsigned extent = 4;770 auto llvmPtrTy = mlir::LLVM::LLVMPointerType::get(builder.getContext());771 mlir::Value ptr = builder.createConvert(loc, llvmPtrTy, a);772 mlir::Type f32Ty = builder.getF32Type();773 mlir::Type idxTy = builder.getIndexType();774 mlir::Type refTy = fir::ReferenceType::get(f32Ty);775 llvm::SmallVector<mlir::Value> values;776 for (unsigned i = 0; i < extent; ++i) {777 mlir::Value pos = builder.createIntegerConstant(loc, idxTy, i);778 mlir::Value coord = fir::CoordinateOp::create(builder, loc, refTy,779 fir::getBase(args[1]), pos);780 mlir::Value value = fir::LoadOp::create(builder, loc, coord);781 values.push_back(value);782 }783 784 auto inlinePtx = mlir::NVVM::InlinePtxOp::create(785 builder, loc, {f32Ty, f32Ty, f32Ty, f32Ty},786 {ptr, values[0], values[1], values[2], values[3]}, {},787 "atom.add.v4.f32 {%0, %1, %2, %3}, [%4], {%5, %6, %7, %8};", {});788 789 llvm::SmallVector<mlir::Value> results;790 results.push_back(inlinePtx.getResult(0));791 results.push_back(inlinePtx.getResult(1));792 results.push_back(inlinePtx.getResult(2));793 results.push_back(inlinePtx.getResult(3));794 795 mlir::Type vecF32Ty = mlir::VectorType::get({extent}, f32Ty);796 mlir::Value undef = mlir::LLVM::UndefOp::create(builder, loc, vecF32Ty);797 mlir::Type i32Ty = builder.getI32Type();798 for (unsigned i = 0; i < extent; ++i)799 undef = mlir::LLVM::InsertElementOp::create(800 builder, loc, undef, results[i],801 builder.createIntegerConstant(loc, i32Ty, i));802 803 auto i128Ty = builder.getIntegerType(128);804 auto i128VecTy = mlir::VectorType::get({1}, i128Ty);805 mlir::Value vec128 =806 mlir::vector::BitCastOp::create(builder, loc, i128VecTy, undef);807 return mlir::vector::ExtractOp::create(builder, loc, vec128,808 mlir::ArrayRef<int64_t>{0});809}810 811mlir::Value812CUDAIntrinsicLibrary::genAtomicAnd(mlir::Type resultType,813 llvm::ArrayRef<mlir::Value> args) {814 assert(args.size() == 2);815 assert(mlir::isa<mlir::IntegerType>(args[1].getType()));816 817 mlir::LLVM::AtomicBinOp binOp = mlir::LLVM::AtomicBinOp::_and;818 return genAtomBinOp(builder, loc, binOp, args[0], args[1]);819}820 821mlir::Value822CUDAIntrinsicLibrary::genAtomicOr(mlir::Type resultType,823 llvm::ArrayRef<mlir::Value> args) {824 assert(args.size() == 2);825 assert(mlir::isa<mlir::IntegerType>(args[1].getType()));826 827 mlir::LLVM::AtomicBinOp binOp = mlir::LLVM::AtomicBinOp::_or;828 return genAtomBinOp(builder, loc, binOp, args[0], args[1]);829}830 831// ATOMICCAS832fir::ExtendedValue833CUDAIntrinsicLibrary::genAtomicCas(mlir::Type resultType,834 llvm::ArrayRef<fir::ExtendedValue> args) {835 assert(args.size() == 3);836 auto successOrdering = mlir::LLVM::AtomicOrdering::acq_rel;837 auto failureOrdering = mlir::LLVM::AtomicOrdering::monotonic;838 auto llvmPtrTy = mlir::LLVM::LLVMPointerType::get(resultType.getContext());839 840 mlir::Value arg0 = fir::getBase(args[0]);841 mlir::Value arg1 = fir::getBase(args[1]);842 mlir::Value arg2 = fir::getBase(args[2]);843 844 auto bitCastFloat = [&](mlir::Value arg) -> mlir::Value {845 if (mlir::isa<mlir::Float32Type>(arg.getType()))846 return mlir::LLVM::BitcastOp::create(builder, loc, builder.getI32Type(),847 arg);848 if (mlir::isa<mlir::Float64Type>(arg.getType()))849 return mlir::LLVM::BitcastOp::create(builder, loc, builder.getI64Type(),850 arg);851 return arg;852 };853 854 arg1 = bitCastFloat(arg1);855 arg2 = bitCastFloat(arg2);856 857 if (arg1.getType() != arg2.getType()) {858 // arg1 and arg2 need to have the same type in AtomicCmpXchgOp.859 arg2 = builder.createConvert(loc, arg1.getType(), arg2);860 }861 862 auto address =863 mlir::UnrealizedConversionCastOp::create(builder, loc, llvmPtrTy, arg0)864 .getResult(0);865 auto cmpxchg = mlir::LLVM::AtomicCmpXchgOp::create(866 builder, loc, address, arg1, arg2, successOrdering, failureOrdering);867 mlir::Value boolResult =868 mlir::LLVM::ExtractValueOp::create(builder, loc, cmpxchg, 1);869 return builder.createConvert(loc, resultType, boolResult);870}871 872mlir::Value873CUDAIntrinsicLibrary::genAtomicDec(mlir::Type resultType,874 llvm::ArrayRef<mlir::Value> args) {875 assert(args.size() == 2);876 assert(mlir::isa<mlir::IntegerType>(args[1].getType()));877 878 mlir::LLVM::AtomicBinOp binOp = mlir::LLVM::AtomicBinOp::udec_wrap;879 return genAtomBinOp(builder, loc, binOp, args[0], args[1]);880}881 882// ATOMICEXCH883fir::ExtendedValue884CUDAIntrinsicLibrary::genAtomicExch(mlir::Type resultType,885 llvm::ArrayRef<fir::ExtendedValue> args) {886 assert(args.size() == 2);887 mlir::Value arg0 = fir::getBase(args[0]);888 mlir::Value arg1 = fir::getBase(args[1]);889 assert(arg1.getType().isIntOrFloat());890 891 mlir::LLVM::AtomicBinOp binOp = mlir::LLVM::AtomicBinOp::xchg;892 return genAtomBinOp(builder, loc, binOp, arg0, arg1);893}894 895mlir::Value896CUDAIntrinsicLibrary::genAtomicInc(mlir::Type resultType,897 llvm::ArrayRef<mlir::Value> args) {898 assert(args.size() == 2);899 assert(mlir::isa<mlir::IntegerType>(args[1].getType()));900 901 mlir::LLVM::AtomicBinOp binOp = mlir::LLVM::AtomicBinOp::uinc_wrap;902 return genAtomBinOp(builder, loc, binOp, args[0], args[1]);903}904 905mlir::Value906CUDAIntrinsicLibrary::genAtomicMax(mlir::Type resultType,907 llvm::ArrayRef<mlir::Value> args) {908 assert(args.size() == 2);909 910 mlir::LLVM::AtomicBinOp binOp =911 mlir::isa<mlir::IntegerType>(args[1].getType())912 ? mlir::LLVM::AtomicBinOp::max913 : mlir::LLVM::AtomicBinOp::fmax;914 return genAtomBinOp(builder, loc, binOp, args[0], args[1]);915}916 917mlir::Value918CUDAIntrinsicLibrary::genAtomicMin(mlir::Type resultType,919 llvm::ArrayRef<mlir::Value> args) {920 assert(args.size() == 2);921 922 mlir::LLVM::AtomicBinOp binOp =923 mlir::isa<mlir::IntegerType>(args[1].getType())924 ? mlir::LLVM::AtomicBinOp::min925 : mlir::LLVM::AtomicBinOp::fmin;926 return genAtomBinOp(builder, loc, binOp, args[0], args[1]);927}928 929// ATOMICSUB930mlir::Value931CUDAIntrinsicLibrary::genAtomicSub(mlir::Type resultType,932 llvm::ArrayRef<mlir::Value> args) {933 assert(args.size() == 2);934 mlir::LLVM::AtomicBinOp binOp =935 mlir::isa<mlir::IntegerType>(args[1].getType())936 ? mlir::LLVM::AtomicBinOp::sub937 : mlir::LLVM::AtomicBinOp::fsub;938 return genAtomBinOp(builder, loc, binOp, args[0], args[1]);939}940 941// ATOMICXOR942fir::ExtendedValue943CUDAIntrinsicLibrary::genAtomicXor(mlir::Type resultType,944 llvm::ArrayRef<fir::ExtendedValue> args) {945 assert(args.size() == 2);946 mlir::Value arg0 = fir::getBase(args[0]);947 mlir::Value arg1 = fir::getBase(args[1]);948 return genAtomBinOp(builder, loc, mlir::LLVM::AtomicBinOp::_xor, arg0, arg1);949}950 951// BARRIER_ARRIVE952mlir::Value953CUDAIntrinsicLibrary::genBarrierArrive(mlir::Type resultType,954 llvm::ArrayRef<mlir::Value> args) {955 assert(args.size() == 1);956 mlir::Value barrier = convertPtrToNVVMSpace(957 builder, loc, args[0], mlir::NVVM::NVVMMemorySpace::Shared);958 return mlir::NVVM::MBarrierArriveOp::create(builder, loc, resultType, barrier)959 .getResult(0);960}961 962// BARRIER_ARRIBVE_CNT963mlir::Value964CUDAIntrinsicLibrary::genBarrierArriveCnt(mlir::Type resultType,965 llvm::ArrayRef<mlir::Value> args) {966 assert(args.size() == 2);967 mlir::Value barrier = convertPtrToNVVMSpace(968 builder, loc, args[0], mlir::NVVM::NVVMMemorySpace::Shared);969 return mlir::NVVM::InlinePtxOp::create(builder, loc, {resultType},970 {barrier, args[1]}, {},971 "mbarrier.arrive.expect_tx.release."972 "cta.shared::cta.b64 %0, [%1], %2;",973 {})974 .getResult(0);975}976 977// BARRIER_INIT978void CUDAIntrinsicLibrary::genBarrierInit(979 llvm::ArrayRef<fir::ExtendedValue> args) {980 assert(args.size() == 2);981 mlir::Value barrier = convertPtrToNVVMSpace(982 builder, loc, fir::getBase(args[0]), mlir::NVVM::NVVMMemorySpace::Shared);983 mlir::NVVM::MBarrierInitOp::create(builder, loc, barrier,984 fir::getBase(args[1]), {});985 auto kind = mlir::NVVM::ProxyKindAttr::get(986 builder.getContext(), mlir::NVVM::ProxyKind::async_shared);987 auto space = mlir::NVVM::SharedSpaceAttr::get(988 builder.getContext(), mlir::NVVM::SharedSpace::shared_cta);989 mlir::NVVM::FenceProxyOp::create(builder, loc, kind, space);990}991 992// BARRIER_TRY_WAIT993mlir::Value994CUDAIntrinsicLibrary::genBarrierTryWait(mlir::Type resultType,995 llvm::ArrayRef<mlir::Value> args) {996 assert(args.size() == 2);997 mlir::Value res = fir::AllocaOp::create(builder, loc, resultType);998 mlir::Value zero = builder.createIntegerConstant(loc, resultType, 0);999 fir::StoreOp::create(builder, loc, zero, res);1000 mlir::Value ns =1001 builder.createIntegerConstant(loc, builder.getI32Type(), 1000000);1002 mlir::Value load = fir::LoadOp::create(builder, loc, res);1003 auto whileOp = mlir::scf::WhileOp::create(1004 builder, loc, mlir::TypeRange{resultType}, mlir::ValueRange{load});1005 mlir::Block *beforeBlock = builder.createBlock(&whileOp.getBefore());1006 mlir::Value beforeArg = beforeBlock->addArgument(resultType, loc);1007 builder.setInsertionPointToStart(beforeBlock);1008 mlir::Value condition = mlir::arith::CmpIOp::create(1009 builder, loc, mlir::arith::CmpIPredicate::ne, beforeArg, zero);1010 mlir::scf::ConditionOp::create(builder, loc, condition, beforeArg);1011 mlir::Block *afterBlock = builder.createBlock(&whileOp.getAfter());1012 afterBlock->addArgument(resultType, loc);1013 builder.setInsertionPointToStart(afterBlock);1014 auto llvmPtrTy = mlir::LLVM::LLVMPointerType::get(builder.getContext());1015 auto barrier = builder.createConvert(loc, llvmPtrTy, args[0]);1016 mlir::Value ret = mlir::NVVM::InlinePtxOp::create(1017 builder, loc, {resultType}, {barrier, args[1], ns}, {},1018 "{\n"1019 " .reg .pred p;\n"1020 " mbarrier.try_wait.shared.b64 p, [%1], %2, %3;\n"1021 " selp.b32 %0, 1, 0, p;\n"1022 "}",1023 {})1024 .getResult(0);1025 mlir::scf::YieldOp::create(builder, loc, ret);1026 builder.setInsertionPointAfter(whileOp);1027 return whileOp.getResult(0);1028}1029 1030// BARRIER_TRY_WAIT_SLEEP1031mlir::Value1032CUDAIntrinsicLibrary::genBarrierTryWaitSleep(mlir::Type resultType,1033 llvm::ArrayRef<mlir::Value> args) {1034 assert(args.size() == 3);1035 auto llvmPtrTy = mlir::LLVM::LLVMPointerType::get(builder.getContext());1036 auto barrier = builder.createConvert(loc, llvmPtrTy, args[0]);1037 return mlir::NVVM::InlinePtxOp::create(1038 builder, loc, {resultType}, {barrier, args[1], args[2]}, {},1039 "{\n"1040 " .reg .pred p;\n"1041 " mbarrier.try_wait.shared.b64 p, [%1], %2, %3;\n"1042 " selp.b32 %0, 1, 0, p;\n"1043 "}",1044 {})1045 .getResult(0);1046}1047 1048static void insertValueAtPos(fir::FirOpBuilder &builder, mlir::Location loc,1049 fir::RecordType recTy, mlir::Value base,1050 mlir::Value dim, unsigned fieldPos) {1051 auto fieldName = recTy.getTypeList()[fieldPos].first;1052 mlir::Type fieldTy = recTy.getTypeList()[fieldPos].second;1053 mlir::Type fieldIndexType = fir::FieldType::get(base.getContext());1054 mlir::Value fieldIndex =1055 fir::FieldIndexOp::create(builder, loc, fieldIndexType, fieldName, recTy,1056 /*typeParams=*/mlir::ValueRange{});1057 mlir::Value coord = fir::CoordinateOp::create(1058 builder, loc, builder.getRefType(fieldTy), base, fieldIndex);1059 fir::StoreOp::create(builder, loc, dim, coord);1060}1061 1062// CLUSTER_BLOCK_INDEX1063mlir::Value1064CUDAIntrinsicLibrary::genClusterBlockIndex(mlir::Type resultType,1065 llvm::ArrayRef<mlir::Value> args) {1066 assert(args.size() == 0);1067 auto recTy = mlir::cast<fir::RecordType>(resultType);1068 assert(recTy && "RecordType expepected");1069 mlir::Value res = fir::AllocaOp::create(builder, loc, resultType);1070 mlir::Type i32Ty = builder.getI32Type();1071 mlir::Value x = mlir::NVVM::BlockInClusterIdXOp::create(builder, loc, i32Ty);1072 mlir::Value one = builder.createIntegerConstant(loc, i32Ty, 1);1073 x = mlir::arith::AddIOp::create(builder, loc, x, one);1074 insertValueAtPos(builder, loc, recTy, res, x, 0);1075 mlir::Value y = mlir::NVVM::BlockInClusterIdYOp::create(builder, loc, i32Ty);1076 y = mlir::arith::AddIOp::create(builder, loc, y, one);1077 insertValueAtPos(builder, loc, recTy, res, y, 1);1078 mlir::Value z = mlir::NVVM::BlockInClusterIdZOp::create(builder, loc, i32Ty);1079 z = mlir::arith::AddIOp::create(builder, loc, z, one);1080 insertValueAtPos(builder, loc, recTy, res, z, 2);1081 return res;1082}1083 1084// CLUSTER_DIM_BLOCKS1085mlir::Value1086CUDAIntrinsicLibrary::genClusterDimBlocks(mlir::Type resultType,1087 llvm::ArrayRef<mlir::Value> args) {1088 assert(args.size() == 0);1089 auto recTy = mlir::cast<fir::RecordType>(resultType);1090 assert(recTy && "RecordType expepected");1091 mlir::Value res = fir::AllocaOp::create(builder, loc, resultType);1092 mlir::Type i32Ty = builder.getI32Type();1093 mlir::Value x = mlir::NVVM::ClusterDimBlocksXOp::create(builder, loc, i32Ty);1094 insertValueAtPos(builder, loc, recTy, res, x, 0);1095 mlir::Value y = mlir::NVVM::ClusterDimBlocksYOp::create(builder, loc, i32Ty);1096 insertValueAtPos(builder, loc, recTy, res, y, 1);1097 mlir::Value z = mlir::NVVM::ClusterDimBlocksZOp::create(builder, loc, i32Ty);1098 insertValueAtPos(builder, loc, recTy, res, z, 2);1099 return res;1100}1101 1102// FENCE_PROXY_ASYNC1103void CUDAIntrinsicLibrary::genFenceProxyAsync(1104 llvm::ArrayRef<fir::ExtendedValue> args) {1105 assert(args.size() == 0);1106 auto kind = mlir::NVVM::ProxyKindAttr::get(1107 builder.getContext(), mlir::NVVM::ProxyKind::async_shared);1108 auto space = mlir::NVVM::SharedSpaceAttr::get(1109 builder.getContext(), mlir::NVVM::SharedSpace::shared_cta);1110 mlir::NVVM::FenceProxyOp::create(builder, loc, kind, space);1111}1112 1113// __LDCA, __LDCS, __LDLU, __LDCV1114template <const char *fctName, int extent>1115fir::ExtendedValue1116CUDAIntrinsicLibrary::genLDXXFunc(mlir::Type resultType,1117 llvm::ArrayRef<fir::ExtendedValue> args) {1118 assert(args.size() == 1);1119 mlir::Type resTy = fir::SequenceType::get(extent, resultType);1120 mlir::Value arg = fir::getBase(args[0]);1121 mlir::Value res = fir::AllocaOp::create(builder, loc, resTy);1122 if (mlir::isa<fir::BaseBoxType>(arg.getType()))1123 arg = fir::BoxAddrOp::create(builder, loc, arg);1124 mlir::Type refResTy = fir::ReferenceType::get(resTy);1125 mlir::FunctionType ftype =1126 mlir::FunctionType::get(arg.getContext(), {refResTy, refResTy}, {});1127 auto funcOp = builder.createFunction(loc, fctName, ftype);1128 llvm::SmallVector<mlir::Value> funcArgs;1129 funcArgs.push_back(res);1130 funcArgs.push_back(arg);1131 fir::CallOp::create(builder, loc, funcOp, funcArgs);1132 mlir::Value ext =1133 builder.createIntegerConstant(loc, builder.getIndexType(), extent);1134 return fir::ArrayBoxValue(res, {ext});1135}1136 1137// CLOCK, CLOCK64, GLOBALTIMER1138template <typename OpTy>1139mlir::Value1140CUDAIntrinsicLibrary::genNVVMTime(mlir::Type resultType,1141 llvm::ArrayRef<mlir::Value> args) {1142 assert(args.size() == 0 && "expect no arguments");1143 return OpTy::create(builder, loc, resultType).getResult();1144}1145 1146// MATCH_ALL_SYNC1147mlir::Value1148CUDAIntrinsicLibrary::genMatchAllSync(mlir::Type resultType,1149 llvm::ArrayRef<mlir::Value> args) {1150 assert(args.size() == 3);1151 bool is32 = args[1].getType().isInteger(32) || args[1].getType().isF32();1152 1153 mlir::Type i1Ty = builder.getI1Type();1154 mlir::MLIRContext *context = builder.getContext();1155 1156 mlir::Value arg1 = args[1];1157 if (arg1.getType().isF32() || arg1.getType().isF64())1158 arg1 = fir::ConvertOp::create(1159 builder, loc, is32 ? builder.getI32Type() : builder.getI64Type(), arg1);1160 1161 mlir::Type retTy =1162 mlir::LLVM::LLVMStructType::getLiteral(context, {resultType, i1Ty});1163 auto match =1164 mlir::NVVM::MatchSyncOp::create(builder, loc, retTy, args[0], arg1,1165 mlir::NVVM::MatchSyncKind::all)1166 .getResult();1167 auto value = mlir::LLVM::ExtractValueOp::create(builder, loc, match, 0);1168 auto pred = mlir::LLVM::ExtractValueOp::create(builder, loc, match, 1);1169 auto conv = mlir::LLVM::ZExtOp::create(builder, loc, resultType, pred);1170 fir::StoreOp::create(builder, loc, conv, args[2]);1171 return value;1172}1173 1174// MATCH_ANY_SYNC1175mlir::Value1176CUDAIntrinsicLibrary::genMatchAnySync(mlir::Type resultType,1177 llvm::ArrayRef<mlir::Value> args) {1178 assert(args.size() == 2);1179 bool is32 = args[1].getType().isInteger(32) || args[1].getType().isF32();1180 1181 mlir::Value arg1 = args[1];1182 if (arg1.getType().isF32() || arg1.getType().isF64())1183 arg1 = fir::ConvertOp::create(1184 builder, loc, is32 ? builder.getI32Type() : builder.getI64Type(), arg1);1185 1186 return mlir::NVVM::MatchSyncOp::create(builder, loc, resultType, args[0],1187 arg1, mlir::NVVM::MatchSyncKind::any)1188 .getResult();1189}1190 1191// SYNCTHREADS1192void CUDAIntrinsicLibrary::genSyncThreads(1193 llvm::ArrayRef<fir::ExtendedValue> args) {1194 mlir::NVVM::Barrier0Op::create(builder, loc);1195}1196 1197// SYNCTHREADS_AND1198mlir::Value1199CUDAIntrinsicLibrary::genSyncThreadsAnd(mlir::Type resultType,1200 llvm::ArrayRef<mlir::Value> args) {1201 mlir::Value arg = builder.createConvert(loc, builder.getI32Type(), args[0]);1202 return mlir::NVVM::BarrierOp::create(1203 builder, loc, resultType, {}, {},1204 mlir::NVVM::BarrierReductionAttr::get(1205 builder.getContext(), mlir::NVVM::BarrierReduction::AND),1206 arg)1207 .getResult(0);1208}1209 1210// SYNCTHREADS_COUNT1211mlir::Value1212CUDAIntrinsicLibrary::genSyncThreadsCount(mlir::Type resultType,1213 llvm::ArrayRef<mlir::Value> args) {1214 mlir::Value arg = builder.createConvert(loc, builder.getI32Type(), args[0]);1215 return mlir::NVVM::BarrierOp::create(1216 builder, loc, resultType, {}, {},1217 mlir::NVVM::BarrierReductionAttr::get(1218 builder.getContext(), mlir::NVVM::BarrierReduction::POPC),1219 arg)1220 .getResult(0);1221}1222 1223// SYNCTHREADS_OR1224mlir::Value1225CUDAIntrinsicLibrary::genSyncThreadsOr(mlir::Type resultType,1226 llvm::ArrayRef<mlir::Value> args) {1227 mlir::Value arg = builder.createConvert(loc, builder.getI32Type(), args[0]);1228 return mlir::NVVM::BarrierOp::create(1229 builder, loc, resultType, {}, {},1230 mlir::NVVM::BarrierReductionAttr::get(1231 builder.getContext(), mlir::NVVM::BarrierReduction::OR),1232 arg)1233 .getResult(0);1234}1235 1236// SYNCWARP1237void CUDAIntrinsicLibrary::genSyncWarp(1238 llvm::ArrayRef<fir::ExtendedValue> args) {1239 assert(args.size() == 1);1240 mlir::NVVM::SyncWarpOp::create(builder, loc, fir::getBase(args[0]));1241}1242 1243// THIS_CLUSTER1244mlir::Value1245CUDAIntrinsicLibrary::genThisCluster(mlir::Type resultType,1246 llvm::ArrayRef<mlir::Value> args) {1247 assert(args.size() == 0);1248 auto recTy = mlir::cast<fir::RecordType>(resultType);1249 assert(recTy && "RecordType expepected");1250 mlir::Value res = fir::AllocaOp::create(builder, loc, resultType);1251 mlir::Type i32Ty = builder.getI32Type();1252 1253 // SIZE1254 mlir::Value size = mlir::NVVM::ClusterDim::create(builder, loc, i32Ty);1255 auto sizeFieldName = recTy.getTypeList()[1].first;1256 mlir::Type sizeFieldTy = recTy.getTypeList()[1].second;1257 mlir::Type fieldIndexType = fir::FieldType::get(resultType.getContext());1258 mlir::Value sizeFieldIndex = fir::FieldIndexOp::create(1259 builder, loc, fieldIndexType, sizeFieldName, recTy,1260 /*typeParams=*/mlir::ValueRange{});1261 mlir::Value sizeCoord = fir::CoordinateOp::create(1262 builder, loc, builder.getRefType(sizeFieldTy), res, sizeFieldIndex);1263 fir::StoreOp::create(builder, loc, size, sizeCoord);1264 1265 // RANK1266 mlir::Value rank = mlir::NVVM::ClusterId::create(builder, loc, i32Ty);1267 mlir::Value one = builder.createIntegerConstant(loc, i32Ty, 1);1268 rank = mlir::arith::AddIOp::create(builder, loc, rank, one);1269 auto rankFieldName = recTy.getTypeList()[2].first;1270 mlir::Type rankFieldTy = recTy.getTypeList()[2].second;1271 mlir::Value rankFieldIndex = fir::FieldIndexOp::create(1272 builder, loc, fieldIndexType, rankFieldName, recTy,1273 /*typeParams=*/mlir::ValueRange{});1274 mlir::Value rankCoord = fir::CoordinateOp::create(1275 builder, loc, builder.getRefType(rankFieldTy), res, rankFieldIndex);1276 fir::StoreOp::create(builder, loc, rank, rankCoord);1277 1278 return res;1279}1280 1281// THIS_GRID1282mlir::Value1283CUDAIntrinsicLibrary::genThisGrid(mlir::Type resultType,1284 llvm::ArrayRef<mlir::Value> args) {1285 assert(args.size() == 0);1286 auto recTy = mlir::cast<fir::RecordType>(resultType);1287 assert(recTy && "RecordType expepected");1288 mlir::Value res = fir::AllocaOp::create(builder, loc, resultType);1289 mlir::Type i32Ty = builder.getI32Type();1290 1291 mlir::Value threadIdX = mlir::NVVM::ThreadIdXOp::create(builder, loc, i32Ty);1292 mlir::Value threadIdY = mlir::NVVM::ThreadIdYOp::create(builder, loc, i32Ty);1293 mlir::Value threadIdZ = mlir::NVVM::ThreadIdZOp::create(builder, loc, i32Ty);1294 1295 mlir::Value blockIdX = mlir::NVVM::BlockIdXOp::create(builder, loc, i32Ty);1296 mlir::Value blockIdY = mlir::NVVM::BlockIdYOp::create(builder, loc, i32Ty);1297 mlir::Value blockIdZ = mlir::NVVM::BlockIdZOp::create(builder, loc, i32Ty);1298 1299 mlir::Value blockDimX = mlir::NVVM::BlockDimXOp::create(builder, loc, i32Ty);1300 mlir::Value blockDimY = mlir::NVVM::BlockDimYOp::create(builder, loc, i32Ty);1301 mlir::Value blockDimZ = mlir::NVVM::BlockDimZOp::create(builder, loc, i32Ty);1302 mlir::Value gridDimX = mlir::NVVM::GridDimXOp::create(builder, loc, i32Ty);1303 mlir::Value gridDimY = mlir::NVVM::GridDimYOp::create(builder, loc, i32Ty);1304 mlir::Value gridDimZ = mlir::NVVM::GridDimZOp::create(builder, loc, i32Ty);1305 1306 // this_grid.size = ((blockDim.z * gridDim.z) * (blockDim.y * gridDim.y)) *1307 // (blockDim.x * gridDim.x);1308 mlir::Value resZ =1309 mlir::arith::MulIOp::create(builder, loc, blockDimZ, gridDimZ);1310 mlir::Value resY =1311 mlir::arith::MulIOp::create(builder, loc, blockDimY, gridDimY);1312 mlir::Value resX =1313 mlir::arith::MulIOp::create(builder, loc, blockDimX, gridDimX);1314 mlir::Value resZY = mlir::arith::MulIOp::create(builder, loc, resZ, resY);1315 mlir::Value size = mlir::arith::MulIOp::create(builder, loc, resZY, resX);1316 1317 // tmp = ((blockIdx.z * gridDim.y * gridDim.x) + (blockIdx.y * gridDim.x)) +1318 // blockIdx.x;1319 // this_group.rank = tmp * ((blockDim.x * blockDim.y) * blockDim.z) +1320 // ((threadIdx.z * blockDim.y) * blockDim.x) +1321 // (threadIdx.y * blockDim.x) + threadIdx.x + 1;1322 mlir::Value r1 =1323 mlir::arith::MulIOp::create(builder, loc, blockIdZ, gridDimY);1324 mlir::Value r2 = mlir::arith::MulIOp::create(builder, loc, r1, gridDimX);1325 mlir::Value r3 =1326 mlir::arith::MulIOp::create(builder, loc, blockIdY, gridDimX);1327 mlir::Value r2r3 = mlir::arith::AddIOp::create(builder, loc, r2, r3);1328 mlir::Value tmp = mlir::arith::AddIOp::create(builder, loc, r2r3, blockIdX);1329 1330 mlir::Value bXbY =1331 mlir::arith::MulIOp::create(builder, loc, blockDimX, blockDimY);1332 mlir::Value bXbYbZ =1333 mlir::arith::MulIOp::create(builder, loc, bXbY, blockDimZ);1334 mlir::Value tZbY =1335 mlir::arith::MulIOp::create(builder, loc, threadIdZ, blockDimY);1336 mlir::Value tZbYbX =1337 mlir::arith::MulIOp::create(builder, loc, tZbY, blockDimX);1338 mlir::Value tYbX =1339 mlir::arith::MulIOp::create(builder, loc, threadIdY, blockDimX);1340 mlir::Value rank = mlir::arith::MulIOp::create(builder, loc, tmp, bXbYbZ);1341 rank = mlir::arith::AddIOp::create(builder, loc, rank, tZbYbX);1342 rank = mlir::arith::AddIOp::create(builder, loc, rank, tYbX);1343 rank = mlir::arith::AddIOp::create(builder, loc, rank, threadIdX);1344 mlir::Value one = builder.createIntegerConstant(loc, i32Ty, 1);1345 rank = mlir::arith::AddIOp::create(builder, loc, rank, one);1346 1347 auto sizeFieldName = recTy.getTypeList()[1].first;1348 mlir::Type sizeFieldTy = recTy.getTypeList()[1].second;1349 mlir::Type fieldIndexType = fir::FieldType::get(resultType.getContext());1350 mlir::Value sizeFieldIndex = fir::FieldIndexOp::create(1351 builder, loc, fieldIndexType, sizeFieldName, recTy,1352 /*typeParams=*/mlir::ValueRange{});1353 mlir::Value sizeCoord = fir::CoordinateOp::create(1354 builder, loc, builder.getRefType(sizeFieldTy), res, sizeFieldIndex);1355 fir::StoreOp::create(builder, loc, size, sizeCoord);1356 1357 auto rankFieldName = recTy.getTypeList()[2].first;1358 mlir::Type rankFieldTy = recTy.getTypeList()[2].second;1359 mlir::Value rankFieldIndex = fir::FieldIndexOp::create(1360 builder, loc, fieldIndexType, rankFieldName, recTy,1361 /*typeParams=*/mlir::ValueRange{});1362 mlir::Value rankCoord = fir::CoordinateOp::create(1363 builder, loc, builder.getRefType(rankFieldTy), res, rankFieldIndex);1364 fir::StoreOp::create(builder, loc, rank, rankCoord);1365 return res;1366}1367 1368// THIS_THREAD_BLOCK1369mlir::Value1370CUDAIntrinsicLibrary::genThisThreadBlock(mlir::Type resultType,1371 llvm::ArrayRef<mlir::Value> args) {1372 assert(args.size() == 0);1373 auto recTy = mlir::cast<fir::RecordType>(resultType);1374 assert(recTy && "RecordType expepected");1375 mlir::Value res = fir::AllocaOp::create(builder, loc, resultType);1376 mlir::Type i32Ty = builder.getI32Type();1377 1378 // this_thread_block%size = blockDim.z * blockDim.y * blockDim.x;1379 mlir::Value blockDimX = mlir::NVVM::BlockDimXOp::create(builder, loc, i32Ty);1380 mlir::Value blockDimY = mlir::NVVM::BlockDimYOp::create(builder, loc, i32Ty);1381 mlir::Value blockDimZ = mlir::NVVM::BlockDimZOp::create(builder, loc, i32Ty);1382 mlir::Value size =1383 mlir::arith::MulIOp::create(builder, loc, blockDimZ, blockDimY);1384 size = mlir::arith::MulIOp::create(builder, loc, size, blockDimX);1385 1386 // this_thread_block%rank = ((threadIdx.z * blockDim.y) * blockDim.x) +1387 // (threadIdx.y * blockDim.x) + threadIdx.x + 1;1388 mlir::Value threadIdX = mlir::NVVM::ThreadIdXOp::create(builder, loc, i32Ty);1389 mlir::Value threadIdY = mlir::NVVM::ThreadIdYOp::create(builder, loc, i32Ty);1390 mlir::Value threadIdZ = mlir::NVVM::ThreadIdZOp::create(builder, loc, i32Ty);1391 mlir::Value r1 =1392 mlir::arith::MulIOp::create(builder, loc, threadIdZ, blockDimY);1393 mlir::Value r2 = mlir::arith::MulIOp::create(builder, loc, r1, blockDimX);1394 mlir::Value r3 =1395 mlir::arith::MulIOp::create(builder, loc, threadIdY, blockDimX);1396 mlir::Value r2r3 = mlir::arith::AddIOp::create(builder, loc, r2, r3);1397 mlir::Value rank = mlir::arith::AddIOp::create(builder, loc, r2r3, threadIdX);1398 mlir::Value one = builder.createIntegerConstant(loc, i32Ty, 1);1399 rank = mlir::arith::AddIOp::create(builder, loc, rank, one);1400 1401 auto sizeFieldName = recTy.getTypeList()[1].first;1402 mlir::Type sizeFieldTy = recTy.getTypeList()[1].second;1403 mlir::Type fieldIndexType = fir::FieldType::get(resultType.getContext());1404 mlir::Value sizeFieldIndex = fir::FieldIndexOp::create(1405 builder, loc, fieldIndexType, sizeFieldName, recTy,1406 /*typeParams=*/mlir::ValueRange{});1407 mlir::Value sizeCoord = fir::CoordinateOp::create(1408 builder, loc, builder.getRefType(sizeFieldTy), res, sizeFieldIndex);1409 fir::StoreOp::create(builder, loc, size, sizeCoord);1410 1411 auto rankFieldName = recTy.getTypeList()[2].first;1412 mlir::Type rankFieldTy = recTy.getTypeList()[2].second;1413 mlir::Value rankFieldIndex = fir::FieldIndexOp::create(1414 builder, loc, fieldIndexType, rankFieldName, recTy,1415 /*typeParams=*/mlir::ValueRange{});1416 mlir::Value rankCoord = fir::CoordinateOp::create(1417 builder, loc, builder.getRefType(rankFieldTy), res, rankFieldIndex);1418 fir::StoreOp::create(builder, loc, rank, rankCoord);1419 return res;1420}1421 1422// THIS_WARP1423mlir::Value1424CUDAIntrinsicLibrary::genThisWarp(mlir::Type resultType,1425 llvm::ArrayRef<mlir::Value> args) {1426 assert(args.size() == 0);1427 auto recTy = mlir::cast<fir::RecordType>(resultType);1428 assert(recTy && "RecordType expepected");1429 mlir::Value res = fir::AllocaOp::create(builder, loc, resultType);1430 mlir::Type i32Ty = builder.getI32Type();1431 1432 // coalesced_group%size = 321433 mlir::Value size = builder.createIntegerConstant(loc, i32Ty, 32);1434 auto sizeFieldName = recTy.getTypeList()[1].first;1435 mlir::Type sizeFieldTy = recTy.getTypeList()[1].second;1436 mlir::Type fieldIndexType = fir::FieldType::get(resultType.getContext());1437 mlir::Value sizeFieldIndex = fir::FieldIndexOp::create(1438 builder, loc, fieldIndexType, sizeFieldName, recTy,1439 /*typeParams=*/mlir::ValueRange{});1440 mlir::Value sizeCoord = fir::CoordinateOp::create(1441 builder, loc, builder.getRefType(sizeFieldTy), res, sizeFieldIndex);1442 fir::StoreOp::create(builder, loc, size, sizeCoord);1443 1444 // coalesced_group%rank = threadIdx.x & 31 + 11445 mlir::Value threadIdX = mlir::NVVM::ThreadIdXOp::create(builder, loc, i32Ty);1446 mlir::Value mask = builder.createIntegerConstant(loc, i32Ty, 31);1447 mlir::Value one = builder.createIntegerConstant(loc, i32Ty, 1);1448 mlir::Value masked =1449 mlir::arith::AndIOp::create(builder, loc, threadIdX, mask);1450 mlir::Value rank = mlir::arith::AddIOp::create(builder, loc, masked, one);1451 auto rankFieldName = recTy.getTypeList()[2].first;1452 mlir::Type rankFieldTy = recTy.getTypeList()[2].second;1453 mlir::Value rankFieldIndex = fir::FieldIndexOp::create(1454 builder, loc, fieldIndexType, rankFieldName, recTy,1455 /*typeParams=*/mlir::ValueRange{});1456 mlir::Value rankCoord = fir::CoordinateOp::create(1457 builder, loc, builder.getRefType(rankFieldTy), res, rankFieldIndex);1458 fir::StoreOp::create(builder, loc, rank, rankCoord);1459 return res;1460}1461 1462// THREADFENCE, THREADFENCE_BLOCK, THREADFENCE_SYSTEM1463template <mlir::NVVM::MemScopeKind scope>1464void CUDAIntrinsicLibrary::genThreadFence(1465 llvm::ArrayRef<fir::ExtendedValue> args) {1466 assert(args.size() == 0);1467 mlir::NVVM::MembarOp::create(builder, loc, scope);1468}1469 1470// TMA_BULK_COMMIT_GROUP1471void CUDAIntrinsicLibrary::genTMABulkCommitGroup(1472 llvm::ArrayRef<fir::ExtendedValue> args) {1473 assert(args.size() == 0);1474 mlir::NVVM::CpAsyncBulkCommitGroupOp::create(builder, loc);1475}1476 1477// TMA_BULK_G2S1478void CUDAIntrinsicLibrary::genTMABulkG2S(1479 llvm::ArrayRef<fir::ExtendedValue> args) {1480 assert(args.size() == 4);1481 mlir::Value barrier = convertPtrToNVVMSpace(1482 builder, loc, fir::getBase(args[0]), mlir::NVVM::NVVMMemorySpace::Shared);1483 mlir::Value dst =1484 convertPtrToNVVMSpace(builder, loc, fir::getBase(args[2]),1485 mlir::NVVM::NVVMMemorySpace::SharedCluster);1486 mlir::Value src = convertPtrToNVVMSpace(builder, loc, fir::getBase(args[1]),1487 mlir::NVVM::NVVMMemorySpace::Global);1488 mlir::NVVM::CpAsyncBulkGlobalToSharedClusterOp::create(1489 builder, loc, dst, src, barrier, fir::getBase(args[3]), {}, {});1490}1491 1492static void genTMABulkLoad(fir::FirOpBuilder &builder, mlir::Location loc,1493 mlir::Value barrier, mlir::Value src,1494 mlir::Value dst, mlir::Value nelem,1495 mlir::Value eleSize) {1496 mlir::Value size = mlir::arith::MulIOp::create(builder, loc, nelem, eleSize);1497 auto llvmPtrTy = mlir::LLVM::LLVMPointerType::get(builder.getContext());1498 barrier = builder.createConvert(loc, llvmPtrTy, barrier);1499 dst = builder.createConvert(loc, llvmPtrTy, dst);1500 src = builder.createConvert(loc, llvmPtrTy, src);1501 mlir::NVVM::InlinePtxOp::create(1502 builder, loc, mlir::TypeRange{}, {dst, src, size, barrier}, {},1503 "cp.async.bulk.shared::cluster.global.mbarrier::complete_tx::bytes [%0], "1504 "[%1], %2, [%3];",1505 {});1506 mlir::NVVM::InlinePtxOp::create(1507 builder, loc, mlir::TypeRange{}, {barrier, size}, {},1508 "mbarrier.expect_tx.relaxed.cta.shared::cta.b64 [%0], %1;", {});1509}1510 1511// TMA_BULK_LOADC41512void CUDAIntrinsicLibrary::genTMABulkLoadC4(1513 llvm::ArrayRef<fir::ExtendedValue> args) {1514 assert(args.size() == 4);1515 mlir::Value eleSize =1516 builder.createIntegerConstant(loc, builder.getI32Type(), 8);1517 genTMABulkLoad(builder, loc, fir::getBase(args[0]), fir::getBase(args[1]),1518 fir::getBase(args[2]), fir::getBase(args[3]), eleSize);1519}1520 1521// TMA_BULK_LOADC81522void CUDAIntrinsicLibrary::genTMABulkLoadC8(1523 llvm::ArrayRef<fir::ExtendedValue> args) {1524 assert(args.size() == 4);1525 mlir::Value eleSize =1526 builder.createIntegerConstant(loc, builder.getI32Type(), 16);1527 genTMABulkLoad(builder, loc, fir::getBase(args[0]), fir::getBase(args[1]),1528 fir::getBase(args[2]), fir::getBase(args[3]), eleSize);1529}1530 1531// TMA_BULK_LOADI41532void CUDAIntrinsicLibrary::genTMABulkLoadI4(1533 llvm::ArrayRef<fir::ExtendedValue> args) {1534 assert(args.size() == 4);1535 mlir::Value eleSize =1536 builder.createIntegerConstant(loc, builder.getI32Type(), 4);1537 genTMABulkLoad(builder, loc, fir::getBase(args[0]), fir::getBase(args[1]),1538 fir::getBase(args[2]), fir::getBase(args[3]), eleSize);1539}1540 1541// TMA_BULK_LOADI81542void CUDAIntrinsicLibrary::genTMABulkLoadI8(1543 llvm::ArrayRef<fir::ExtendedValue> args) {1544 assert(args.size() == 4);1545 mlir::Value eleSize =1546 builder.createIntegerConstant(loc, builder.getI32Type(), 8);1547 genTMABulkLoad(builder, loc, fir::getBase(args[0]), fir::getBase(args[1]),1548 fir::getBase(args[2]), fir::getBase(args[3]), eleSize);1549}1550 1551// TMA_BULK_LOADR21552void CUDAIntrinsicLibrary::genTMABulkLoadR2(1553 llvm::ArrayRef<fir::ExtendedValue> args) {1554 assert(args.size() == 4);1555 mlir::Value eleSize =1556 builder.createIntegerConstant(loc, builder.getI32Type(), 2);1557 genTMABulkLoad(builder, loc, fir::getBase(args[0]), fir::getBase(args[1]),1558 fir::getBase(args[2]), fir::getBase(args[3]), eleSize);1559}1560 1561// TMA_BULK_LOADR41562void CUDAIntrinsicLibrary::genTMABulkLoadR4(1563 llvm::ArrayRef<fir::ExtendedValue> args) {1564 assert(args.size() == 4);1565 mlir::Value eleSize =1566 builder.createIntegerConstant(loc, builder.getI32Type(), 4);1567 genTMABulkLoad(builder, loc, fir::getBase(args[0]), fir::getBase(args[1]),1568 fir::getBase(args[2]), fir::getBase(args[3]), eleSize);1569}1570 1571// TMA_BULK_LOADR81572void CUDAIntrinsicLibrary::genTMABulkLoadR8(1573 llvm::ArrayRef<fir::ExtendedValue> args) {1574 assert(args.size() == 4);1575 mlir::Value eleSize =1576 builder.createIntegerConstant(loc, builder.getI32Type(), 8);1577 genTMABulkLoad(builder, loc, fir::getBase(args[0]), fir::getBase(args[1]),1578 fir::getBase(args[2]), fir::getBase(args[3]), eleSize);1579}1580 1581// TMA_BULK_S2G1582void CUDAIntrinsicLibrary::genTMABulkS2G(1583 llvm::ArrayRef<fir::ExtendedValue> args) {1584 assert(args.size() == 3);1585 mlir::Value src = convertPtrToNVVMSpace(builder, loc, fir::getBase(args[0]),1586 mlir::NVVM::NVVMMemorySpace::Shared);1587 mlir::Value dst = convertPtrToNVVMSpace(builder, loc, fir::getBase(args[1]),1588 mlir::NVVM::NVVMMemorySpace::Global);1589 mlir::NVVM::CpAsyncBulkSharedCTAToGlobalOp::create(1590 builder, loc, dst, src, fir::getBase(args[2]), {}, {});1591 1592 mlir::NVVM::InlinePtxOp::create(builder, loc, mlir::TypeRange{}, {}, {},1593 "cp.async.bulk.commit_group;", {});1594 mlir::NVVM::CpAsyncBulkWaitGroupOp::create(builder, loc,1595 builder.getI32IntegerAttr(0), {});1596}1597 1598static void genTMABulkStore(fir::FirOpBuilder &builder, mlir::Location loc,1599 mlir::Value src, mlir::Value dst, mlir::Value count,1600 mlir::Value eleSize) {1601 mlir::Value size = mlir::arith::MulIOp::create(builder, loc, eleSize, count);1602 src = convertPtrToNVVMSpace(builder, loc, src,1603 mlir::NVVM::NVVMMemorySpace::Shared);1604 dst = convertPtrToNVVMSpace(builder, loc, dst,1605 mlir::NVVM::NVVMMemorySpace::Global);1606 mlir::NVVM::CpAsyncBulkSharedCTAToGlobalOp::create(builder, loc, dst, src,1607 size, {}, {});1608 mlir::NVVM::InlinePtxOp::create(builder, loc, mlir::TypeRange{}, {}, {},1609 "cp.async.bulk.commit_group;", {});1610 mlir::NVVM::CpAsyncBulkWaitGroupOp::create(builder, loc,1611 builder.getI32IntegerAttr(0), {});1612}1613 1614// TMA_BULK_STORE_C41615void CUDAIntrinsicLibrary::genTMABulkStoreC4(1616 llvm::ArrayRef<fir::ExtendedValue> args) {1617 assert(args.size() == 3);1618 mlir::Value eleSize =1619 builder.createIntegerConstant(loc, builder.getI32Type(), 8);1620 genTMABulkStore(builder, loc, fir::getBase(args[0]), fir::getBase(args[1]),1621 fir::getBase(args[2]), eleSize);1622}1623 1624// TMA_BULK_STORE_C81625void CUDAIntrinsicLibrary::genTMABulkStoreC8(1626 llvm::ArrayRef<fir::ExtendedValue> args) {1627 assert(args.size() == 3);1628 mlir::Value eleSize =1629 builder.createIntegerConstant(loc, builder.getI32Type(), 16);1630 genTMABulkStore(builder, loc, fir::getBase(args[0]), fir::getBase(args[1]),1631 fir::getBase(args[2]), eleSize);1632}1633 1634// TMA_BULK_STORE_I41635void CUDAIntrinsicLibrary::genTMABulkStoreI4(1636 llvm::ArrayRef<fir::ExtendedValue> args) {1637 assert(args.size() == 3);1638 mlir::Value eleSize =1639 builder.createIntegerConstant(loc, builder.getI32Type(), 4);1640 genTMABulkStore(builder, loc, fir::getBase(args[0]), fir::getBase(args[1]),1641 fir::getBase(args[2]), eleSize);1642}1643 1644// TMA_BULK_STORE_I81645void CUDAIntrinsicLibrary::genTMABulkStoreI8(1646 llvm::ArrayRef<fir::ExtendedValue> args) {1647 assert(args.size() == 3);1648 mlir::Value eleSize =1649 builder.createIntegerConstant(loc, builder.getI32Type(), 8);1650 genTMABulkStore(builder, loc, fir::getBase(args[0]), fir::getBase(args[1]),1651 fir::getBase(args[2]), eleSize);1652}1653 1654// TMA_BULK_STORE_R21655void CUDAIntrinsicLibrary::genTMABulkStoreR2(1656 llvm::ArrayRef<fir::ExtendedValue> args) {1657 assert(args.size() == 3);1658 mlir::Value eleSize =1659 builder.createIntegerConstant(loc, builder.getI32Type(), 2);1660 genTMABulkStore(builder, loc, fir::getBase(args[0]), fir::getBase(args[1]),1661 fir::getBase(args[2]), eleSize);1662}1663 1664// TMA_BULK_STORE_R41665void CUDAIntrinsicLibrary::genTMABulkStoreR4(1666 llvm::ArrayRef<fir::ExtendedValue> args) {1667 assert(args.size() == 3);1668 mlir::Value eleSize =1669 builder.createIntegerConstant(loc, builder.getI32Type(), 4);1670 genTMABulkStore(builder, loc, fir::getBase(args[0]), fir::getBase(args[1]),1671 fir::getBase(args[2]), eleSize);1672}1673 1674// TMA_BULK_STORE_R81675void CUDAIntrinsicLibrary::genTMABulkStoreR8(1676 llvm::ArrayRef<fir::ExtendedValue> args) {1677 assert(args.size() == 3);1678 mlir::Value eleSize =1679 builder.createIntegerConstant(loc, builder.getI32Type(), 8);1680 genTMABulkStore(builder, loc, fir::getBase(args[0]), fir::getBase(args[1]),1681 fir::getBase(args[2]), eleSize);1682}1683 1684// TMA_BULK_WAIT_GROUP1685void CUDAIntrinsicLibrary::genTMABulkWaitGroup(1686 llvm::ArrayRef<fir::ExtendedValue> args) {1687 assert(args.size() == 0);1688 auto group = builder.getIntegerAttr(builder.getI32Type(), 0);1689 mlir::NVVM::CpAsyncBulkWaitGroupOp::create(builder, loc, group, {});1690}1691 1692// ALL_SYNC, ANY_SYNC, BALLOT_SYNC1693template <mlir::NVVM::VoteSyncKind kind>1694mlir::Value1695CUDAIntrinsicLibrary::genVoteSync(mlir::Type resultType,1696 llvm::ArrayRef<mlir::Value> args) {1697 assert(args.size() == 2);1698 mlir::Value arg1 =1699 fir::ConvertOp::create(builder, loc, builder.getI1Type(), args[1]);1700 mlir::Type resTy = kind == mlir::NVVM::VoteSyncKind::ballot1701 ? builder.getI32Type()1702 : builder.getI1Type();1703 auto voteRes =1704 mlir::NVVM::VoteSyncOp::create(builder, loc, resTy, args[0], arg1, kind)1705 .getResult();1706 return fir::ConvertOp::create(builder, loc, resultType, voteRes);1707}1708 1709} // namespace fir1710