#include "mlir/Dialect/LLVMIR/BasicPtxBuilderInterface.h"
#include "mlir/IR/BuiltinTypes.h"
#include "mlir/IR/Diagnostics.h"
#include "mlir/IR/Location.h"
#include "mlir/IR/MLIRContext.h"
#include "mlir/Support/LLVM.h"
#include "llvm/ADT/StringExtras.h"
#include "llvm/ADT/TypeSwitch.h"
#include "llvm/Support/DebugLog.h"
#include "llvm/Support/FormatVariadic.h"
#include "llvm/Support/LogicalResult.h"
#include "llvm/Support/Regex.h"
#include "mlir/Dialect/LLVMIR/BasicPtxBuilderInterface.cpp.inc"

Macros
#define	DEBUG_TYPE "ptx-builder"

Functions
static FailureOr< char >	getRegisterType (Type type, Location loc)

static FailureOr< char >	getRegisterType (Value v, Location loc)

static SmallVector< Value >	extractStructElements (PatternRewriter &rewriter, Location loc, Value structVal)
	Extract every element of a struct value. More...

static bool	needsPackUnpack (BasicPtxBuilderInterface interfaceOp, bool needsManualRegisterMapping, SmallVectorImpl< PTXRegisterMod > &registerModifiers)
	Check if the operation needs to pack and unpack results. More...

static SmallVector< Type >	packResultTypes (BasicPtxBuilderInterface interfaceOp, bool needsManualRegisterMapping, SmallVectorImpl< PTXRegisterMod > &registerModifiers, SmallVectorImpl< Value > &ptxOperands)
	Pack the result types of the interface operation. More...

static std::string	canonicalizeRegisterConstraints (llvm::StringRef csv)
	Canonicalize the register constraints: More...

static llvm::Regex	getPredicateMappingRegex ()
	Returns a regex that matches {$rwN}, {$wN}, {$rN}. More...

static std::string	rewriteAsmPlaceholders (llvm::StringRef ptxCode)
	Rewrites `{$rwN}`, `{$wN}`, and `{$rN}` placeholders in `ptxCode` into compact `$K` indices: More...

Variables
static constexpr int64_t	kSharedMemorySpace = 3

constexpr llvm::StringLiteral	kReadWritePrefix {"rw"}

constexpr llvm::StringLiteral	kWriteOnlyPrefix {"w"}

constexpr llvm::StringLiteral	kReadOnlyPrefix {"r"}

Macro Definition Documentation

◆ DEBUG_TYPE

#define DEBUG_TYPE "ptx-builder"

Definition at line 28 of file BasicPtxBuilderInterface.cpp.

Function Documentation

◆ canonicalizeRegisterConstraints()

static std::string canonicalizeRegisterConstraints ( llvm::StringRef csv )

static

Canonicalize the register constraints:

Turn every "+X" into "=X"
Append (at the very end) the 0-based indices of tokens that were "+X" Examples: "+f,+f,+r,=r,=r,r,r" -> "=f,=f,=r,=r,=r,r,r,0,1,2" "+f,+f,+r,=r,=r" -> "=f,=f,=r,=r,=r,0,1,2"

Definition at line 258 of file BasicPtxBuilderInterface.cpp.

Referenced by mlir::NVVM::PtxBuilder::build().

◆ extractStructElements()

static SmallVector<Value> extractStructElements	(	PatternRewriter &	rewriter,
		Location	loc,
		Value	structVal
	)

static

Extract every element of a struct value.

Definition at line 126 of file BasicPtxBuilderInterface.cpp.

References mlir::Value::getType().

Referenced by mlir::NVVM::PtxBuilder::buildAndReplaceOp().

◆ getPredicateMappingRegex()

static llvm::Regex getPredicateMappingRegex ( )

static

Returns a regex that matches {$rwN}, {$wN}, {$rN}.

Definition at line 297 of file BasicPtxBuilderInterface.cpp.

References kReadOnlyPrefix, kReadWritePrefix, and kWriteOnlyPrefix.

Referenced by rewriteAsmPlaceholders().

◆ getRegisterType() [1/2]

static FailureOr<char> getRegisterType	(	Type	type,
		Location	loc
	)

static

Definition at line 41 of file BasicPtxBuilderInterface.cpp.

References mlir::emitError(), mlir::get(), mlir::Type::getContext(), mlir::Type::isBF16(), mlir::Type::isF16(), mlir::Type::isF32(), mlir::Type::isF64(), mlir::Type::isFloat(), mlir::Type::isInteger(), and kSharedMemorySpace.

Referenced by getRegisterType(), and mlir::NVVM::PtxBuilder::insertValue().

◆ getRegisterType() [2/2]

static FailureOr<char> getRegisterType	(	Value	v,
		Location	loc
	)

static

Definition at line 119 of file BasicPtxBuilderInterface.cpp.

References mlir::Value::getDefiningOp(), getRegisterType(), and mlir::Value::getType().

◆ needsPackUnpack()

static bool needsPackUnpack	(	BasicPtxBuilderInterface	interfaceOp,
		bool	needsManualRegisterMapping,
		SmallVectorImpl< PTXRegisterMod > &	registerModifiers
	)

static

Check if the operation needs to pack and unpack results.

Definition at line 202 of file BasicPtxBuilderInterface.cpp.

References mlir::NVVM::ReadWrite.

Referenced by mlir::NVVM::PtxBuilder::buildAndReplaceOp(), and packResultTypes().

◆ packResultTypes()

static SmallVector<Type> packResultTypes	(	BasicPtxBuilderInterface	interfaceOp,
		bool	needsManualRegisterMapping,
		SmallVectorImpl< PTXRegisterMod > &	registerModifiers,
		SmallVectorImpl< Value > &	ptxOperands
	)

static

Pack the result types of the interface operation.

If the operation has multiple results, it packs them into a struct type. Otherwise, it returns the original result types.

Definition at line 219 of file BasicPtxBuilderInterface.cpp.

References needsPackUnpack(), and mlir::NVVM::ReadWrite.

Referenced by mlir::NVVM::PtxBuilder::build().

◆ rewriteAsmPlaceholders()

static std::string rewriteAsmPlaceholders ( llvm::StringRef ptxCode )

static

Rewrites {$rwN}, {$wN}, and {$rN} placeholders in ptxCode into compact $K indices:

All rw* first (sorted by N),
Then w*,
Then r*. If there a predicate, it comes always in the end. Each number is assigned once; duplicates are ignored.

Example Input: "{ reg .pred p; setp.ge.s32 p, {$r0}, {$r1};" selp.s32 {$rw0}, {$r0}, {$r1}, p; selp.s32 {$rw1}, {$r0}, {$r1}, p; selp.s32 {$w0}, {$r0}, {$r1}, p; selp.s32 {$w1}, {$r0}, {$r1}, p; }
" Example Output: "{ reg .pred p; setp.ge.s32 p, $4, $5;" selp.s32 $0, $4, $5, p; selp.s32 $1, $4, $5, p; selp.s32 $2, $4, $5, p; selp.s32 $3, $4, $5, p; }
"

Definition at line 363 of file BasicPtxBuilderInterface.cpp.

References mlir::NVVM::countPlaceholderNumbers(), getPredicateMappingRegex(), kReadWritePrefix, and kWriteOnlyPrefix.

Referenced by mlir::NVVM::PtxBuilder::build().